Merge pull request #183 from OpenBioSim/sync_exscientia

Synchronise with Exscientia remote

.github/workflows/Sandpit_exs.yml

@@ -19,7 +19,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest", "macOS-latest",]
-        python-version: ["3.8", ]
+        python-version: ["3.10",]
 
     steps:
       - uses: actions/checkout@v2
@@ -35,12 +35,10 @@ jobs:
 
       - name: Install dependency
         run: |
-          mamba install -c conda-forge -c openbiosim/label/main biosimspace python=3.8 ambertools gromacs "sire=2023.2.2" alchemlyb pytest pyarrow "typing-extensions!=4.6.0" openff-interchange pint=0.21 "rdkit<2023" "jaxlib>0.3.7" tqdm
+          mamba install -c conda-forge -c openbiosim/label/main biosimspace python=3.10 ambertools gromacs "sire=2023.3" "alchemlyb>=2.1" pytest openff-interchange pint=0.21 rdkit "jaxlib>0.3.7" tqdm
           python -m pip install git+https://github.com/Exscientia/MDRestraintsGenerator.git
           # For the testing of BSS.FreeEnergy.AlchemicalFreeEnergy.analysis
           python -m pip install https://github.com/alchemistry/alchemtest/archive/master.zip
-          # Before the new alchemlyb is released
-          python -m pip install git+https://github.com/alchemistry/alchemlyb.git
 
       - name: Install the dev version
         run: |

python/BioSimSpace/Sandpit/Exscientia/Align/_alch_ion.py

@@ -0,0 +1,52 @@
+import warnings
+
+from .._SireWrappers import Molecule as _Molecule
+
+
+def _mark_alchemical_ion(molecule):
+    """
+    Mark the ion molecule as being alchemical ion.
+
+    This enables one to use
+
+        * :meth:`~BioSimSpace.Sandpit.Exscientia._SireWrappers._system.System.getAlchemicalIon` to get the alchemical ion.
+        * :meth:`~BioSimSpace.Sandpit.Exscientia._SireWrappers._system.System.getAlchemicalIonIdx` to get the index of alchemical ion.
+        * :meth:`~BioSimSpace.Sandpit.Exscientia._SireWrappers._molecule.Molecule.isAlchemicalIon` to check if a molecule is an alchemical ion.
+
+
+    Parameters
+    ----------
+
+    molecule : BioSimSpace._SireWrappers.Molecule
+        The molecule to be marked as alchemical ion.
+
+    Returns
+    -------
+
+    alchemical_ion : BSS._SireWrappers.Molecule
+        The molecule marked as being alchemical ion.
+    """
+    # Validate input.
+
+    if not isinstance(molecule, _Molecule):
+        raise TypeError(
+            "'molecule' must be of type 'BioSimSpace._SireWrappers.Molecule'"
+        )
+
+    # Cannot decouple a perturbable molecule.
+    if molecule.isAlchemicalIon():
+        warnings.warn("'molecule' has already been marked as alchemical ion!")
+
+    # Create a copy of this molecule.
+    mol = _Molecule(molecule)
+    mol_sire = mol._sire_object
+
+    # Edit the molecule
+    mol_edit = mol_sire.edit()
+
+    mol_edit.setProperty("AlchemicalIon", True)
+
+    # Update the Sire molecule object of the new molecule.
+    mol._sire_object = mol_edit.commit()
+
+    return mol

python/BioSimSpace/Sandpit/Exscientia/Align/_squash.py

@@ -113,9 +113,30 @@ def _squash_molecule(molecule, explicit_dummies=False):
         return molecule
 
     if explicit_dummies:
-        # We make sure we use the same coordinates at both endstates.
+        # Get the common core atoms
+        atom_mapping0_common = _squashed_atom_mapping(
+            molecule,
+            is_lambda1=False,
+            explicit_dummies=explicit_dummies,
+            environment=False,
+            dummies=False,
+        )
+        atom_mapping1_common = _squashed_atom_mapping(
+            molecule,
+            is_lambda1=True,
+            explicit_dummies=explicit_dummies,
+            environment=False,
+            dummies=False,
+        )
+        if set(atom_mapping0_common) != set(atom_mapping1_common):
+            raise RuntimeError("The MCS atoms don't match between the two endstates")
+        common_atoms = set(atom_mapping0_common)
+
+        # We make sure we use the same coordinates for the common core at both endstates.
         c = molecule.copy()._sire_object.cursor()
-        c["coordinates1"] = c["coordinates0"]
+        for i, atom in enumerate(c.atoms()):
+            if i in common_atoms:
+                atom["coordinates1"] = atom["coordinates0"]
         molecule = _Molecule(c.commit())
 
     # Generate a "system" from the molecule at lambda = 0 and another copy at lambda = 1.
@@ -269,20 +290,38 @@ def _unsquash_molecule(molecule, squashed_molecules, explicit_dummies=False):
     molecule : BioSimSpace._SireWrappers.Molecule
          The output updated merged molecule.
     """
-    # Get the atom mapping and combine it with the lambda=0 molecule being prioritised
+    # Get the common core atoms
+    atom_mapping0_common = _squashed_atom_mapping(
+        molecule,
+        is_lambda1=False,
+        explicit_dummies=explicit_dummies,
+        environment=False,
+        dummies=False,
+    )
+    atom_mapping1_common = _squashed_atom_mapping(
+        molecule,
+        is_lambda1=True,
+        explicit_dummies=explicit_dummies,
+        environment=False,
+        dummies=False,
+    )
+    if set(atom_mapping0_common) != set(atom_mapping1_common):
+        raise RuntimeError("The MCS atoms don't match between the two endstates")
+    common_atoms = set(atom_mapping0_common)
+
+    # Get the atom mapping from both endstates
     atom_mapping0 = _squashed_atom_mapping(
         molecule, is_lambda1=False, explicit_dummies=explicit_dummies
     )
     atom_mapping1 = _squashed_atom_mapping(
         molecule, is_lambda1=True, explicit_dummies=explicit_dummies
     )
-    atom_mapping = {**atom_mapping1, **atom_mapping0}
     update_velocity = squashed_molecules[0]._sire_object.hasProperty("velocity")
 
-    # Even though the two molecules should have the same coordinates, they might be PBC wrapped differently.
+    # Even though the common core of the two molecules should have the same coordinates,
+    # they might be PBC wrapped differently.
     # Here we take the first common core atom and translate the second molecule.
     if len(squashed_molecules) == 2:
-        common_atoms = set(atom_mapping0.keys()) & set(atom_mapping1.keys())
         first_common_atom = list(sorted(common_atoms))[0]
         pertatom0 = squashed_molecules.getAtom(atom_mapping0[first_common_atom])
         pertatom1 = squashed_molecules.getAtom(atom_mapping1[first_common_atom])
@@ -292,25 +331,42 @@ def _unsquash_molecule(molecule, squashed_molecules, explicit_dummies=False):
 
     # Update the coordinates and velocities.
     siremol = molecule.copy()._sire_object.edit()
-    for merged_atom_idx, squashed_atom_idx in atom_mapping.items():
+    for merged_atom_idx in range(molecule.nAtoms()):
+        # Get the relevant atom indices
         merged_atom = siremol.atom(_SireMol.AtomIdx(merged_atom_idx))
-        squashed_atom = squashed_molecules.getAtom(squashed_atom_idx)
+        if merged_atom_idx in atom_mapping0:
+            squashed_atom_idx0 = atom_mapping0[merged_atom_idx]
+        else:
+            squashed_atom_idx0 = atom_mapping1[merged_atom_idx]
+        if merged_atom_idx in atom_mapping1:
+            squashed_atom_idx1 = atom_mapping1[merged_atom_idx]
+            apply_translation_vec = True
+        else:
+            squashed_atom_idx1 = atom_mapping0[merged_atom_idx]
+            apply_translation_vec = False
 
-        # Update the coordinates.
-        coordinates = squashed_atom._sire_object.property("coordinates")
+        # Get the coordinates.
+        squashed_atom0 = squashed_molecules.getAtom(squashed_atom_idx0)
+        squashed_atom1 = squashed_molecules.getAtom(squashed_atom_idx1)
+        coordinates0 = squashed_atom0._sire_object.property("coordinates")
+        coordinates1 = squashed_atom1._sire_object.property("coordinates")
 
         # Apply the translation if the atom is coming from the second molecule.
-        if len(squashed_molecules) == 2 and squashed_atom_idx in atom_mapping1.values():
-            coordinates -= translation_vec
+        if len(squashed_molecules) == 2 and apply_translation_vec:
+            # This is a dummy atom so we need to translate coordinates0 as well
+            if squashed_atom_idx0 == squashed_atom_idx1:
+                coordinates0 -= translation_vec
+            coordinates1 -= translation_vec
 
-        siremol = merged_atom.setProperty("coordinates0", coordinates).molecule()
-        siremol = merged_atom.setProperty("coordinates1", coordinates).molecule()
+        siremol = merged_atom.setProperty("coordinates0", coordinates0).molecule()
+        siremol = merged_atom.setProperty("coordinates1", coordinates1).molecule()
 
         # Update the velocities.
         if update_velocity:
-            velocities = squashed_atom._sire_object.property("velocity")
-            siremol = merged_atom.setProperty("velocity0", velocities).molecule()
-            siremol = merged_atom.setProperty("velocity1", velocities).molecule()
+            velocities0 = squashed_atom0._sire_object.property("velocity")
+            velocities1 = squashed_atom1._sire_object.property("velocity")
+            siremol = merged_atom.setProperty("velocity0", velocities0).molecule()
+            siremol = merged_atom.setProperty("velocity1", velocities1).molecule()
 
     return _Molecule(siremol.commit())
 

python/BioSimSpace/Sandpit/Exscientia/Process/_amber.py

@@ -27,6 +27,7 @@
 __all__ = ["Amber"]
 
 import os
+import shutil
 from pathlib import Path as _Path
 
 from .._Utils import _try_import
@@ -303,9 +304,11 @@ def _write_system(self, system, coord_file=None, topol_file=None, ref_file=None)
         if coord_file is not None:
             try:
                 file = _os.path.splitext(coord_file)[0]
-                _IO.saveMolecules(file, system, "rst7", property_map=self._property_map)
+                _IO.saveMolecules(file, system, "rst", property_map=self._property_map)
+                # To keep the file extension the same.
+                shutil.move(f"{file}.rst", f"{file}.rst7")
             except Exception as e:
-                msg = "Failed to write system to 'RST7' format."
+                msg = "Failed to write system to 'rst7' format."
                 if _isVerbose():
                     raise IOError(msg) from e
                 else:
@@ -315,9 +318,11 @@ def _write_system(self, system, coord_file=None, topol_file=None, ref_file=None)
         if ref_file is not None:
             try:
                 file = _os.path.splitext(ref_file)[0]
-                _IO.saveMolecules(file, system, "rst7", property_map=self._property_map)
+                _IO.saveMolecules(file, system, "rst", property_map=self._property_map)
+                # To keep the file extension the same.
+                shutil.move(f"{file}.rst", f"{file}.rst7")
             except Exception as e:
-                msg = "Failed to write system to 'RST7' format."
+                msg = "Failed to write system to 'rst7' format."
                 if _isVerbose():
                     raise IOError(msg) from e
                 else:

python/BioSimSpace/Sandpit/Exscientia/Process/_process.py

@@ -29,6 +29,7 @@
 import collections as _collections
 import glob as _glob
 import os as _os
+import traceback
 
 import pandas as pd
 
@@ -910,7 +911,16 @@ def wait(self, max_time=None):
 
         # The process isn't running.
         if not self.isRunning():
-            self.saveMetric()
+            try:
+                self.saveMetric()
+            except Exception:
+                exception_info = traceback.format_exc()
+                with open(f"{self.workDir()}/{self._name}.err", "a+") as f:
+                    f.write("Exception Information during saveMetric():\n")
+                    f.write("======================\n")
+                    f.write(exception_info)
+                    f.write("\n\n")
+
             return
 
         if max_time is not None:

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

@@ -518,6 +518,18 @@ def isML(self):
         else:
             return False
 
+    def isAlchemicalIon(self):
+        """
+        Whether this molecule is marked as Alchemical Ion.
+
+        Returns
+        -------
+
+        isAlchemicalIon : bool
+            Whether the molecule is marked as Alchemical Ion.
+        """
+        return self._sire_object.hasProperty("AlchemicalIon")
+
     def isWater(self, property_map={}):
         """
         Whether this is a water molecule.

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

@@ -1285,6 +1285,33 @@ def reduceBoxVectors(self, bias=0, property_map={}):
         # Update the space property in the sire object.
         self._sire_object.setProperty(space_prop, space)
 
+    def getAlchemicalIon(self):
+        """
+        Return the Alchemical Ion in the system.
+
+        Returns
+        -------
+
+        molecule : [:class:`Molecule <BioSimSpace._SireWrappers.Molecule>`]
+            The Alchemical Ion or None if there isn't any.
+        """
+        try:
+            return self.search("mols with property AlchemicalIon").molecules()[0]
+        except:
+            return None
+
+    def getAlchemicalIonIdx(self):
+        """
+        Return the index of Alchemical Ion in the system.
+
+        Returns
+        -------
+
+        index : int
+            The index of Alchemical Ion in the system.
+        """
+        return self.getIndex(self.getAlchemicalIon())
+
     def repartitionHydrogenMass(
         self, factor=4, water="no", use_coordinates=False, property_map={}
     ):

tests/Sandpit/Exscientia/Align/test_alchemical_ion.py

@@ -0,0 +1,52 @@
+import pytest
+
+import BioSimSpace.Sandpit.Exscientia as BSS
+from BioSimSpace.Sandpit.Exscientia.Align._alch_ion import _mark_alchemical_ion
+from BioSimSpace.Sandpit.Exscientia._SireWrappers import Molecule
+
+from tests.conftest import root_fp, has_gromacs
+
+
+@pytest.fixture
+def system():
+    mol = BSS.IO.readMolecules(
+        [f"{root_fp}/input/ala.top", f"{root_fp}/input/ala.crd"]
+    )[0]
+    system = BSS.Solvent.tip3p(mol, ion_conc=0.15, shell=2 * BSS.Units.Length.nanometer)
+    return system
+
+
+@pytest.fixture
+def alchemical_ion_system(system):
+    ion = system[-1]
+    ion = _mark_alchemical_ion(ion)
+    system.updateMolecules(ion)
+    return system
+
+
+@pytest.mark.skipif(has_gromacs is False, reason="Requires GROMACS to be installed.")
+@pytest.mark.parametrize(
+    "input_system,isalchem", [("system", False), ("alchemical_ion_system", True)]
+)
+def test_isAlchemicalIon(input_system, isalchem, request):
+    system = request.getfixturevalue(input_system)
+    assert system[-1].isAlchemicalIon() is isalchem
+
+
+@pytest.mark.skipif(has_gromacs is False, reason="Requires GROMACS to be installed.")
+@pytest.mark.parametrize(
+    "input_system,isalchem", [("system", None), ("alchemical_ion_system", True)]
+)
+def test_getAlchemicalIon(input_system, isalchem, request):
+    system = request.getfixturevalue(input_system)
+    ion = system.getAlchemicalIon()
+    if isalchem is None:
+        assert ion is None
+    else:
+        assert isinstance(ion, Molecule)
+
+
+@pytest.mark.skipif(has_gromacs is False, reason="Requires GROMACS to be installed.")
+def test_getAlchemicalIonIdx(alchemical_ion_system):
+    index = alchemical_ion_system.getAlchemicalIonIdx()
+    assert index == 680

tests/Sandpit/Exscientia/Align/test_decouple.py

@@ -6,6 +6,7 @@
 
 from BioSimSpace.Sandpit.Exscientia.Align._decouple import decouple
 import BioSimSpace.Sandpit.Exscientia as BSS
+from tests.conftest import root_fp
 
 # Store the tutorial URL.
 url = BSS.tutorialUrl()
@@ -14,7 +15,9 @@
 @pytest.fixture(scope="session")
 def mol():
     # Alanin-dipeptide.
-    return BSS.IO.readMolecules(["tests/input/ala.top", "tests/input/ala.crd"])[0]
+    return BSS.IO.readMolecules(
+        [f"{root_fp}/input/ala.top", f"{root_fp}/input/ala.crd"]
+    )[0]
 
 
 def test_sanity(mol):