improve cntr initialization

JecaTosovic · Oct 10, 2023 · eb06ce7 · eb06ce7
1 parent 49cd8e1
commit eb06ce7
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/ConservedWaterSearch/utils.py b/ConservedWaterSearch/utils.py
@@ -258,7 +258,14 @@ def _add_hydrogen_and_bond(wname, Hpos, Hname, resn, resi):
 def _make_water_objects(water_type, waterO, waterH1, waterH2, output_file):
     from pymol import cmd
 
-    cntr = {"FCW": 0, "WCW": 0, "HCW": 0}
+    cntr = {
+        name: (
+            len(cmd.get_names("objects", False, f"model {name}*"))
+            if cmd.get_names("objects", False, f"model {name}*")
+            else 0
+        )
+        for name in ["FCW", "WCW", "HCW", "onlyO"]
+    }
     ind = 0  # initialize index
     while ind < len(water_type):
         tip, Opos, H1pos, H2pos = (