Preserve molecule ordering when adding molecules. [ref #193]

michellab · Apr 28, 2021 · 4017a46 · 4017a46
1 parent cf104f5
commit 4017a46
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 4 deletions.
diff --git a/python/BioSimSpace/_SireWrappers/_molecule.py b/python/BioSimSpace/_SireWrappers/_molecule.py
@@ -131,6 +131,11 @@ def __add__(self, other):
 
         # A single Molecule object.
         elif type(other) is Molecule:
+            mol = self.copy()
+            mol._renumber(_SireMol.MolNum.getUniqueNumber())
+            molecules[0] = mol
+            other = other.copy()
+            other._renumber(_SireMol.MolNum.getUniqueNumber())
             molecules.append(other)
 
         # A Molecules object.
@@ -139,7 +144,13 @@ def __add__(self, other):
 
         # A list of Molecule objects.
         elif type(other) is list and all(isinstance(x, Molecule) for x in other):
-            molecules.extend(other)
+            mol = self.copy()
+            mol._renumber(_SireMol.MolNum.getUniqueNumber())
+            molecules[0] = mol
+            for mol in other:
+                mol = mol.copy()
+                mol._renumber(_SireMol.MolNum.getUniqueNumber())
+                molecules.append(mol)
 
         # Unsupported.
         else:
@@ -1111,6 +1122,22 @@ def makeCompatibleWith(self, molecule, property_map={}, overwrite=True,
             # Finally, commit the changes to the internal object.
             self._sire_object = edit_mol.commit()
 
+    def _renumber(self, mol_num):
+        """Renumber the molecule with a unique MolNum.
+
+           Parameters
+           ----------
+
+           mol_num : Sire.Mol.MolNum
+               The molecule number.
+        """
+        if type(mol_num) is not _SireMol.MolNum:
+            raise TypeError("'mol_num' must be of type 'Sire.Mol.MolNum'")
+
+        edit_mol = self._sire_object.edit()
+        edit_mol.renumber(mol_num)
+        self._sire_object = edit_mol.commit()
+
     def _getPropertyMap0(self):
         """Generate a property map for the lambda = 0 state of the merged molecule."""
 

diff --git a/python/BioSimSpace/_SireWrappers/_molecules.py b/python/BioSimSpace/_SireWrappers/_molecules.py
@@ -124,16 +124,20 @@ def __add__(self, other):
 
         # A Molecule object.
         elif type(other) is _Molecule:
-            molecules.add(other._sire_object)
+            mol = other.copy()
+            mol._renumber(_SireMol.MolNum.getUniqueNumber())
+            molecules.add(mol._sire_object)
 
         # A Molecules object.
         elif type(other) is Molecules:
             molecules.add(other._sire_object)
 
         # A list of Molecule objects.
-        elif type(other) is list and all(isinstance(x, Molecule) for x in other):
+        elif type(other) is list and all(isinstance(x, _Molecule) for x in other):
             for molecule in other:
-                molecules.add(molecule._sire_object)
+                mol = molecule.copy()
+                mol._renumber(_SireMol.MolNum.getUniqueNumber())
+                molecules.add(mol._sire_object)
 
         # Unsupported.
         else: