From 45e58975086c083b8c9042cecbd4735f4f57a899 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 31 Oct 2023 08:54:32 +0000
Subject: [PATCH] Backport fix from PR #193. [ci skip]

---
 .../Exscientia/_SireWrappers/_system.py       | 27 ++++++++++++++++++-
 python/BioSimSpace/_SireWrappers/_system.py   | 27 ++++++++++++++++++-
 2 files changed, 52 insertions(+), 2 deletions(-)

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):
         """