Avoid setting all end-state properties when decoupling

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

@@ -76,7 +76,7 @@ def decouple(
             if not isinstance(value, bool):
                 raise ValueError(f"{value} in {name} must be bool.")
 
-    # Change names of charge and LJ tuples to avoid clashes with properties
+    # Change names of charge and LJ tuples to avoid clashes with properties.
     charge_tuple = charge
     LJ_tuple = LJ
 
@@ -86,7 +86,7 @@ def decouple(
     # Invert the user property mappings.
     inv_property_map = {v: k for k, v in property_map.items()}
 
-    # Create a copy of this molecule and Sire object to check properties
+    # Create a copy of this molecule and Sire object to check properties.
     mol = _Molecule(molecule)
     mol_sire = mol._sire_object
 
@@ -97,7 +97,7 @@ def decouple(
     element = inv_property_map.get("element", "element")
     ambertype = inv_property_map.get("ambertype", "ambertype")
 
-    # Check for missing information
+    # Check for missing information.
     if not mol_sire.hasProperty(ff):
         raise _IncompatibleError("Cannot determine 'forcefield' of 'molecule'!")
     if not mol_sire.hasProperty(LJ):
@@ -107,54 +107,32 @@ def decouple(
     if not mol_sire.hasProperty(element):
         raise _IncompatibleError("Cannot determine elements in molecule")
 
-    # Check for ambertype property (optional)
+    # Check for ambertype property (optional).
     has_ambertype = True
     if not mol_sire.hasProperty(ambertype):
         has_ambertype = False
 
     if not isinstance(intramol, bool):
         raise TypeError("'intramol' must be of type 'bool'")
 
-    # Edit the molecule
+    # Edit the molecule.
     mol_edit = mol_sire.edit()
 
-    # Create dictionary to store charge and LJ tuples
+    # Create dictionary to store charge and LJ tuples.
     mol_edit.setProperty(
         "decouple", {"charge": charge_tuple, "LJ": LJ_tuple, "intramol": intramol}
     )
 
     # Set the "forcefield0" property.
     mol_edit.setProperty("forcefield0", molecule._sire_object.property(ff))
 
-    # Set starting properties based on fully-interacting molecule
+    # Set starting properties based on fully-interacting molecule.
     mol_edit.setProperty("charge0", molecule._sire_object.property(charge))
     mol_edit.setProperty("LJ0", molecule._sire_object.property(LJ))
     mol_edit.setProperty("element0", molecule._sire_object.property(element))
     if has_ambertype:
         mol_edit.setProperty("ambertype0", molecule._sire_object.property(ambertype))
 
-    # Set final charges and LJ terms to 0, elements to "X" and (if required) ambertypes to du
-    for atom in mol_sire.atoms():
-        mol_edit = (
-            mol_edit.atom(atom.index())
-            .setProperty("charge1", 0 * _SireUnits.e_charge)
-            .molecule()
-        )
-        mol_edit = (
-            mol_edit.atom(atom.index())
-            .setProperty("LJ1", _SireMM.LJParameter())
-            .molecule()
-        )
-        mol_edit = (
-            mol_edit.atom(atom.index())
-            .setProperty("element1", _SireMol.Element(0))
-            .molecule()
-        )
-        if has_ambertype:
-            mol_edit = (
-                mol_edit.atom(atom.index()).setProperty("ambertype1", "du").molecule()
-            )
-
     mol_edit.setProperty("annihilated", _SireBase.wrap(intramol))
 
     # Flag that this molecule is decoupled.

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

@@ -26,10 +26,10 @@
 
 __all__ = ["Somd"]
 
-from .._Utils import _try_import
-
 import os as _os
 
+from .._Utils import _try_import
+
 _pygtail = _try_import("pygtail")
 import glob as _glob
 import random as _random
@@ -38,23 +38,21 @@
 import timeit as _timeit
 import warnings as _warnings
 
-from sire.legacy import Base as _SireBase
 from sire.legacy import CAS as _SireCAS
 from sire.legacy import IO as _SireIO
 from sire.legacy import MM as _SireMM
+from sire.legacy import Base as _SireBase
 from sire.legacy import Mol as _SireMol
+from sire.legacy import Units as _SireUnits
 
-from .. import _isVerbose
+from .. import IO as _IO
+from .. import Protocol as _Protocol
+from .. import Trajectory as _Trajectory
+from .. import _isVerbose, _Utils
 from .._Exceptions import IncompatibleError as _IncompatibleError
 from .._Exceptions import MissingSoftwareError as _MissingSoftwareError
 from .._SireWrappers import Molecule as _Molecule
 from .._SireWrappers import System as _System
-
-from .. import IO as _IO
-from .. import Protocol as _Protocol
-from .. import Trajectory as _Trajectory
-from .. import _Utils
-
 from . import _process
 
 
@@ -1001,6 +999,9 @@ def _to_pert_file(
     if not isinstance(perturbation_type, str):
         raise TypeError("'perturbation_type' must be of type 'str'")
 
+    if not isinstance(property_map, dict):
+        raise TypeError("'property_map' must be of type 'dict'")
+
     # Convert to lower case and strip whitespace.
     perturbation_type = perturbation_type.lower().replace(" ", "")
 
@@ -1027,6 +1028,60 @@ def _to_pert_file(
     # Extract and copy the Sire molecule.
     mol = molecule._sire_object.__deepcopy__()
 
+    # If the molecule is decoupled (for an ABFE calculation), then we need to
+    # set the end-state properties of the molecule.
+    if molecule.isDecoupled():
+        # Invert the user property mappings.
+        inv_property_map = {v: k for k, v in property_map.items()}
+
+        # Get required properties.
+        lj = inv_property_map.get("LJ", "LJ")
+        charge = inv_property_map.get("charge", "charge")
+        element = inv_property_map.get("element", "element")
+        ambertype = inv_property_map.get("ambertype", "ambertype")
+
+        # Check for missing information.
+        if not mol.hasProperty(lj):
+            raise _IncompatibleError("Cannot determine LJ terms for molecule")
+        if not mol.hasProperty(charge):
+            raise _IncompatibleError("Cannot determine charges for molecule")
+        if not mol.hasProperty(element):
+            raise _IncompatibleError("Cannot determine elements in molecule")
+
+        # Check for ambertype property.
+        has_ambertype = True
+        if not mol.hasProperty(ambertype):
+            has_ambertype = False
+
+        mol_edit = mol.edit()
+
+        # Set final charges and LJ terms to 0, elements to "X" and (if required) ambertypes to du
+        for atom in mol.atoms():
+            mol_edit = (
+                mol_edit.atom(atom.index())
+                .setProperty("charge1", 0 * _SireUnits.e_charge)
+                .molecule()
+            )
+            mol_edit = (
+                mol_edit.atom(atom.index())
+                .setProperty("LJ1", _SireMM.LJParameter())
+                .molecule()
+            )
+            mol_edit = (
+                mol_edit.atom(atom.index())
+                .setProperty("element1", _SireMol.Element(0))
+                .molecule()
+            )
+            if has_ambertype:
+                mol_edit = (
+                    mol_edit.atom(atom.index())
+                    .setProperty("ambertype1", "du")
+                    .molecule()
+                )
+
+        # Update the Sire molecule object of the new molecule.
+        mol = mol_edit.commit()
+
     # First work out the indices of atoms that are perturbed.
     pert_idxs = []
 

tests/Sandpit/Exscientia/Align/test_decouple.py

@@ -81,8 +81,11 @@ def test_topology(mol, tmp_path):
 
 
 def test_end_types(mol):
-    """Check that the correct properties have been set at either
-    end of the perturbation."""
+    """
+    Check that the correct properties have been set for the 0
+    end of the perturbation. Note that for SOMD, the 1 end state
+    properties are set in Process.Somd._to_pert_file.
+    """
 
     decoupled_mol = decouple(mol)
     assert decoupled_mol._sire_object.property("charge0") == mol._sire_object.property(
@@ -95,8 +98,3 @@ def test_end_types(mol):
     assert decoupled_mol._sire_object.property(
         "ambertype0"
     ) == mol._sire_object.property("ambertype")
-    for atom in decoupled_mol._sire_object.atoms():
-        assert atom.property("charge1") == 0 * _SireUnits.e_charge
-        assert atom.property("LJ1") == _SireMM.LJParameter()
-        assert atom.property("element1") == _SireMol.Element(0)
-        assert atom.property("ambertype1") == "du"

tests/Sandpit/Exscientia/FreeEnergy/test_alchemical_free_energy.py

@@ -1,4 +1,5 @@
 import bz2
+from math import exp
 import pandas as pd
 import pathlib
 import pytest
@@ -270,9 +271,11 @@ def freenrg():
         return freenrg
 
     def test_files_exist(self, freenrg):
-        """Test if the files have been created. Note that e.g. gradients.dat
+        """
+        Test if the files have been created. Note that e.g. gradients.dat
         are not created until later in the simulation, so their presence is
-        not tested for."""
+        not tested for.
+        """
         path = pathlib.Path(freenrg.workDir())
         for lam in ["0.0000", "0.5000", "1.0000"]:
             assert (path / f"lambda_{lam}" / "simfile.dat").is_file()
@@ -282,6 +285,31 @@ def test_files_exist(self, freenrg):
             assert (path / f"lambda_{lam}" / "somd.prm7").is_file()
             assert (path / f"lambda_{lam}" / "somd.err").is_file()
             assert (path / f"lambda_{lam}" / "somd.out").is_file()
+            assert (path / f"lambda_{lam}" / "somd.pert").is_file()
+
+    def test_correct_pert_file(self, freenrg):
+        """Check that pert file is correct."""
+        path = pathlib.Path(freenrg.workDir()) / "lambda_0.0000"
+        with open(os.path.join(path, "somd.pert"), "rt") as f:
+            lines = f.readlines()
+
+            for i, line in enumerate(lines):
+                # Check that the end-state properties are correct.
+                if "final_type" in line:
+                    assert "final_type     du" in line
+                if "final_LJ" in line:
+                    assert "final_LJ       0.00000 0.00000" in line
+                if "final_charge" in line:
+                    assert "final_charge   0.00000" in line
+                # Check that the initial state properties are correct for the first and last atoms.
+                if line == "        name           C1\n":
+                    assert "initial_type   c" in lines[i + 1]
+                    assert "initial_LJ     3.39967 0.08600" in lines[i + 3]
+                    assert "initial_charge 0.67120" in lines[i + 5]
+                if "name           O3" in line:
+                    assert "initial_type   o" in lines[i + 1]
+                    assert "initial_LJ     2.95992 0.21000" in lines[i + 3]
+                    assert "initial_charge -0.52110" in lines[i + 5]
 
     def test_correct_conf_file(self, freenrg):
         """Check that lambda data is correct in somd.cfg"""