Check perturbable molecules for velocity property. [closes #192]

michellab · Oct 27, 2023 · f234669 · f234669
1 parent cb92227
commit f234669
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py b/python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py
@@ -629,7 +629,7 @@ def addMolecules(self, molecules):
 
         # Remove velocities if any molecules are missing them.
         if self.nMolecules() > 1:
-            # Search for water molecules in the system.
+            # Search for molecules with a velocity property.
             try:
                 mols_with_velocities = self.search(
                     f"mols with property velocity"
@@ -638,6 +638,19 @@ def addMolecules(self, molecules):
             except:
                 num_vels = 0
 
+            # Search for perturbable molecules with a velocity property.
+            # Only consider the lambda = 0 end state.
+            try:
+                pert_mols_with_velocities = self.search(
+                    f"mols with property velocity0"
+                ).molecules()
+                num_pert_vels = len(pert_mols_with_velocities)
+            except:
+                num_pert_vels = 0
+
+            # Compute the total number of molecules with velocities.
+            num_vels = num_vels + num_pert_vels
+
             # Not all molecules have velocities.
             if num_vels > 0 and num_vels != self.nMolecules():
                 _warnings.warn(

diff --git a/python/BioSimSpace/_SireWrappers/_system.py b/python/BioSimSpace/_SireWrappers/_system.py
@@ -629,7 +629,7 @@ def addMolecules(self, molecules):
 
         # Remove velocities if any molecules are missing them.
         if self.nMolecules() > 1:
-            # Search for water molecules in the system.
+            # Search for molecules with a velocity property.
             try:
                 mols_with_velocities = self.search(
                     f"mols with property velocity"
@@ -638,6 +638,19 @@ def addMolecules(self, molecules):
             except:
                 num_vels = 0
 
+            # Search for perturbable molecules with a velocity property.
+            # Only consider the lambda = 0 end state.
+            try:
+                pert_mols_with_velocities = self.search(
+                    f"mols with property velocity0"
+                ).molecules()
+                num_pert_vels = len(pert_mols_with_velocities)
+            except:
+                num_pert_vels = 0
+
+            # Compute the total number of molecules with velocities.
+            num_vels = num_vels + num_pert_vels
+
             # Not all molecules have velocities.
             if num_vels > 0 and num_vels != self.nMolecules():
                 _warnings.warn(