Merge pull request #440 from OpenBioSim/feature_system_renumber

Add support for renumbering molecules when copying a system

Author: lohedges

python/BioSimSpace/Sandpit/Exscientia/_SireWrappers/_system.py

@@ -260,6 +260,47 @@ def __len__(self):
         """Return the number of molecules in the system."""
         return self.nMolecules()
 
+    def copy(self, renumber=False):
+        """
+        Return a copy of this System.
+
+        Parameters
+        ----------
+
+        renumber : bool
+            Whether to give the copied molecules unique molecule numbers.
+
+        Returns
+        -------
+
+        System : :class:`System <BioSimSpace._SireWrappers.System>`
+            A copy of the object.
+        """
+
+        if not isinstance(renumber, bool):
+            raise TypeError("'renumber' must be of type 'bool'")
+
+        if not renumber:
+            return super().copy()
+
+        # Create a molecules container.
+        mols = _SireMol.MoleculeGroup("all")
+
+        # Give each molecule a unique molecule number.
+        for mol in self._sire_object:
+            cursor = mol.cursor()
+            cursor.number = _SireMol.MolNum.getUniqueNumber()
+            mols.add(cursor.commit())
+
+        # Create a new system.
+        system = _Molecules(mols).toSystem()
+
+        # Copy over the system properties.
+        for prop in self._sire_object.propertyKeys():
+            system._sire_object.setProperty(prop, self._sire_object.property(prop))
+
+        return system
+
     def nMolecules(self):
         """
         Return the number of molecules in the system.

python/BioSimSpace/_SireWrappers/_system.py

@@ -260,6 +260,47 @@ def __len__(self):
         """Return the number of molecules in the system."""
         return self.nMolecules()
 
+    def copy(self, renumber=False):
+        """
+        Return a copy of this System.
+
+        Parameters
+        ----------
+
+        renumber : bool
+            Whether to give the copied molecules unique molecule numbers.
+
+        Returns
+        -------
+
+        System : :class:`System <BioSimSpace._SireWrappers.System>`
+            A copy of the object.
+        """
+
+        if not isinstance(renumber, bool):
+            raise TypeError("'renumber' must be of type 'bool'")
+
+        if not renumber:
+            return super().copy()
+
+        # Create a molecules container.
+        mols = _SireMol.MoleculeGroup("all")
+
+        # Give each molecule a unique molecule number.
+        for mol in self._sire_object:
+            cursor = mol.cursor()
+            cursor.number = _SireMol.MolNum.getUniqueNumber()
+            mols.add(cursor.commit())
+
+        # Create a new system.
+        system = _Molecules(mols).toSystem()
+
+        # Copy over the system properties.
+        for prop in self._sire_object.propertyKeys():
+            system._sire_object.setProperty(prop, self._sire_object.property(prop))
+
+        return system
+
     def nMolecules(self):
         """
         Return the number of molecules in the system.

tests/Sandpit/Exscientia/_SireWrappers/test_system.py

@@ -508,3 +508,26 @@ def test_remove_box(system):
 
     # Make sure the box is removed.
     assert not "space" in system._sire_object.propertyKeys()
+
+
+def test_renumber(system):
+    # Make sure all the molecules retain their original numbers
+    # using the default copy.
+    for mol0, mol1 in zip(system, system.copy()):
+        assert mol0.number() == mol1.number()
+
+    # Store the set of original molecule numbers.
+    original_numbers = set(mol.number() for mol in system)
+
+    # Create a renumbered copy of the system.
+    renumbered_system = system.copy(renumber=True)
+
+    # Make sure that all of the numbers differ when renumbering.
+    for mol0, mol1 in zip(system, renumbered_system):
+        assert mol0.number() != mol1.number()
+
+    # Store the set of renumbered molecule numbers.
+    renumbered_numbers = set(mol.number() for mol in renumbered_system)
+
+    # Make sure that no original numbers are present in the renumbered set.
+    assert original_numbers.isdisjoint(renumbered_numbers)

tests/_SireWrappers/test_system.py

@@ -498,3 +498,26 @@ def test_remove_box(system):
 
     # Make sure the box is removed.
     assert not "space" in system._sire_object.propertyKeys()
+
+
+def test_renumber(system):
+    # Make sure all the molecules retain their original numbers
+    # using the default copy.
+    for mol0, mol1 in zip(system, system.copy()):
+        assert mol0.number() == mol1.number()
+
+    # Store the set of original molecule numbers.
+    original_numbers = set(mol.number() for mol in system)
+
+    # Create a renumbered copy of the system.
+    renumbered_system = system.copy(renumber=True)
+
+    # Make sure that all of the numbers differ when renumbering.
+    for mol0, mol1 in zip(system, renumbered_system):
+        assert mol0.number() != mol1.number()
+
+    # Store the set of renumbered molecule numbers.
+    renumbered_numbers = set(mol.number() for mol in renumbered_system)
+
+    # Make sure that no original numbers are present in the renumbered set.
+    assert original_numbers.isdisjoint(renumbered_numbers)