diff --git a/python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py b/python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py index 0cb99ef12..b22d18dc5 100644 --- a/python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py +++ b/python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py @@ -628,7 +628,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" @@ -637,6 +637,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( @@ -650,6 +663,18 @@ def addMolecules(self, molecules): _warnings.warn( "Failed to remove 'velocity' property from all molecules!" ) + if num_pert_vels > 0: + try: + self._sire_object = _SireIO.removeProperty( + self._sire_object, "velocity0" + ) + self._sire_object = _SireIO.removeProperty( + self._sire_object, "velocity1" + ) + except: + _warnings.warn( + "Failed to remove 'velocity0' and 'velocity1' property from molecules!" + ) def removeMolecules(self, molecules): """ diff --git a/python/BioSimSpace/_SireWrappers/_system.py b/python/BioSimSpace/_SireWrappers/_system.py index d8200e55e..8137a1478 100644 --- a/python/BioSimSpace/_SireWrappers/_system.py +++ b/python/BioSimSpace/_SireWrappers/_system.py @@ -628,7 +628,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" @@ -637,6 +637,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( @@ -650,6 +663,18 @@ def addMolecules(self, molecules): _warnings.warn( "Failed to remove 'velocity' property from all molecules!" ) + if num_pert_vels > 0: + try: + self._sire_object = _SireIO.removeProperty( + self._sire_object, "velocity0" + ) + self._sire_object = _SireIO.removeProperty( + self._sire_object, "velocity1" + ) + except: + _warnings.warn( + "Failed to remove 'velocity0' and 'velocity1' property from molecules!" + ) def removeMolecules(self, molecules): """