Add alchemical ion support

.github/workflows/main.yaml

@@ -28,7 +28,7 @@ jobs:
       max-parallel: 9
       fail-fast: false
       matrix:
-        python-version: ["3.11"]
+        python-version: ["3.12"]
         platform:
           - { name: "windows", os: "windows-latest", shell: "pwsh" }
           - { name: "linux", os: "ubuntu-latest", shell: "bash -l {0}" }
@@ -55,12 +55,10 @@ jobs:
           activate-environment: somd2
           environment-file: environment.yaml
           miniforge-version: latest
-          miniforge-variant: Mambaforge
-          use-mamba: true
           run-post: ${{ matrix.platform.name != 'windows' }}
 #
       - name: Install pytest
-        run: mamba install pytest
+        run: conda install pytest
 #
       - name: Install the package
         run: pip install .

src/somd2/config/_config.py

@@ -104,6 +104,7 @@ def __init__(
         perturbable_constraint="h_bonds_not_heavy_perturbed",
         include_constrained_energies=False,
         dynamic_constraints=True,
+        charge_difference=0,
         com_reset_frequency=10,
         minimise=True,
         equilibration_time="0ps",
@@ -211,6 +212,12 @@ def __init__(
             to the value of r0 at that lambda value. If this is False, then
             the constraint is set based on the current length.
 
+        charge_difference: int
+            The charge difference between the two end states. (Perturbed minus
+            reference.) If specified, then a number of alchemical ions will be
+            added to the system to keep the charge constant, i.e. make the
+            perturbed state charge the same as the reference state.
+
         com_reset_frequency: int
             Frequency at which to reset the centre of mass of the system.
 
@@ -315,6 +322,7 @@ def __init__(
         self.perturbable_constraint = perturbable_constraint
         self.include_constrained_energies = include_constrained_energies
         self.dynamic_constraints = dynamic_constraints
+        self.charge_difference = charge_difference
         self.com_reset_frequency = com_reset_frequency
         self.minimise = minimise
         self.equilibration_time = equilibration_time
@@ -861,6 +869,17 @@ def dynamic_constraints(self, dynamic_constraints):
             raise ValueError("'dynamic_constraints' must be of type 'bool'")
         self._dynamic_constraints = dynamic_constraints
 
+    @property
+    def charge_difference(self):
+        return self._charge_difference
+
+    @charge_difference.setter
+    def charge_difference(self, charge_difference):
+        if charge_difference is not None:
+            if not isinstance(charge_difference, int):
+                raise ValueError("'charge_difference' must be an integer")
+        self._charge_difference = charge_difference
+
     @property
     def com_reset_frequency(self):
         return self._com_reset_frequency

src/somd2/runner/_dynamics.py

@@ -151,7 +151,8 @@ def increment_filename(base_filename, suffix):
             raise ValueError("lambda_value not in lambda_array")
         lam = f"{lambda_value:.5f}"
         filenames = {}
-        filenames["topology"] = "system.prm7"
+        filenames["topology0"] = "system0.prm7"
+        filenames["topology1"] = "system1.prm7"
         filenames["checkpoint"] = f"checkpoint_{lam}.s3"
         filenames["energy_traj"] = f"energy_traj_{lam}.parquet"
         filenames["trajectory"] = f"traj_{lam}.dcd"
@@ -307,8 +308,12 @@ def _run(self, lambda_minimisation=None, is_restart=False):
 
         # Save the system topology to a PRM7 file that can be used to load the
         # trajectory.
-        topology = str(self._config.output_directory / self._filenames["topology"])
-        sr.save(self._system, topology)
+        topology0 = str(self._config.output_directory / self._filenames["topology0"])
+        topology1 = str(self._config.output_directory / self._filenames["topology1"])
+        mols0 = sr.morph.link_to_reference(self._system)
+        mols1 = sr.morph.link_to_perturbed(self._system)
+        sr.save(mols0, topology0)
+        sr.save(mols1, topology1)
 
         def generate_lam_vals(lambda_base, increment):
             """Generate lambda values for a given lambda_base and increment"""
@@ -560,7 +565,7 @@ def generate_lam_vals(lambda_base, increment):
                     traj_chunks = [f"{traj_filename}.bak"] + traj_chunks
 
             # Load the topology and chunked trajectory files.
-            system = sr.load([topology] + traj_chunks)
+            system = sr.load([topology0] + traj_chunks)
 
             # Save the final trajectory to a single file.
             traj_filename = str(

src/somd2/runner/_runner.py

@@ -173,6 +173,30 @@ def __init__(self, system, config):
         # Check the end state contraints.
         self._check_end_state_constraints()
 
+        # Get the charge difference between the two end states.
+        charge_diff = self._get_charge_difference(self._system)
+
+        # Make sure the difference is integer valued to 5 decimal places.
+        if not round(charge_diff, 4).is_integer():
+            _logger.warning("Charge difference between end states is not an integer.")
+        charge_diff = int(round(charge_diff, 4))
+
+        # Make sure the charge difference matches the expected value
+        # from the config.
+        if self._config.charge_difference != charge_diff:
+            _logger.warning(
+                f"The charge difference of {charge_diff} between the end states "
+                f"does not match the expected value of {self._config.charge_difference}. "
+                "Please specify the 'charge_difference' if you wish to keep the charge "
+                "constant."
+            )
+
+        # Create alchemical ions.
+        if self._config.charge_difference != 0:
+            self._system = self._create_alchemical_ions(
+                self._system, self._config.charge_difference
+            )
+
         # Set the lambda values.
         if self._config.lambda_values:
             self._lambda_values = self._config.lambda_values
@@ -318,6 +342,106 @@ def _check_end_state_constraints(self):
                         )
                         break
 
+    @staticmethod
+    def _get_charge_difference(system):
+        """
+        Internal function to check the charge difference between the two end states.
+
+        Parameters
+        ----------
+
+        system: :class: `System <sire.system.System>`
+            The system to be perturbed.
+
+        Returns
+        -------
+
+        charge_diff: int
+            The charge difference between the perturbed and reference states.
+        """
+
+        reference = _morph.link_to_reference(system).charge().value()
+        perturbed = _morph.link_to_perturbed(system).charge().value()
+
+        return perturbed - reference
+
+    @staticmethod
+    def _create_alchemical_ions(system, charge_diff):
+        """
+        Internal function to create alchemical ions to maintain a constant charge.
+
+        Parameters
+        ----------
+
+        system: :class: `System <sire.system.System>`
+            The system to be perturbed.
+
+        charge_diff: int
+            The charge difference between perturbed and reference states.
+
+        Returns
+        -------
+
+        system: :class: `System <sire.system.System>`
+            The perturbed system with alchemical ions added.
+        """
+
+        from sire.legacy.IO import createChlorineIon as _createChlorineIon
+        from sire.legacy.IO import createSodiumIon as _createSodiumIon
+
+        # Clone the system.
+        system = system.clone()
+
+        # The number of waters to convert is the absolute charge difference.
+        num_waters = abs(charge_diff)
+
+        # Reference coordinates.
+        coords = system.molecules("property is_perturbable").coordinates()
+        coord_string = f"{coords[0].value()}, {coords[1].value()}, {coords[2].value()}"
+
+        # Find the furthest N waters from the perturbable molecule.
+        waters = system[f"furthest {num_waters} waters from {coord_string}"].molecules()
+
+        # Determine the water model.
+        if waters[0].num_atoms() == 3:
+            model = "tip3p"
+        elif waters[0].num_atoms() == 4:
+            model = "tip4p"
+        elif waters[0].num_atoms() == 5:
+            # Note that AMBER has no ion model for tip5p.
+            model = "tip4p"
+
+        # Store the molecule numbers for the system.
+        numbers = system.numbers()
+
+        # Create the ions.
+        for water in waters:
+            # Create an ion. This is to adjust the charge of the perturbed state
+            # to match that of the reference.
+            if charge_diff > 0:
+                ion = _createChlorineIon(water["element O"].coordinates(), model)
+                ion_str = "Cl-"
+            else:
+                ion = _createSodiumIon(water["element O"].coordinates(), model)
+                ion_str = "Na+"
+
+            # Create a perturbable molecule: water --> ion.
+            merged = _morph.merge(water, ion, map={"as_new_molecule": False})
+
+            # Update the system.
+            system.update(merged)
+
+            # Get the index of the perturbed water.
+            index = numbers.index(water.number())
+
+            # Log that the water was perturbed.
+            _logger.info(
+                f"Water at molecule index {index} will be perturbed to "
+                f"{ion_str} to keep charge constant."
+            )
+
+        return system
+
     def _check_directory(self):
         """
         Find the name of the file from which simulations will be started.

tests/runner/test_alchemical_ions.py

@@ -0,0 +1,22 @@
+import math
+
+from somd2.runner import Runner
+
+
+def test_alchemical_ions(ethane_methanol):
+    """Ensure that alchemical ions are added correctly."""
+
+    # Clone the system.
+    mols = ethane_methanol.clone()
+
+    # Add 10 Cl- ions.
+    new_mols = Runner._create_alchemical_ions(mols, 10)
+
+    # Make sure the charge difference is correct.
+    assert math.isclose(Runner._get_charge_difference(new_mols), -10.0, rel_tol=1e-6)
+
+    # Add 10 Na+ ions.
+    new_mols = Runner._create_alchemical_ions(mols, -10)
+
+    # Make sure the charge difference is correct.
+    assert math.isclose(Runner._get_charge_difference(new_mols), 10.0, rel_tol=1e-6)

tests/runner/test_restart.py

@@ -45,7 +45,7 @@ def test_restart(mols, request):
 
         # Load the trajectory.
         traj_1 = sr.load(
-            [str(Path(tmpdir) / "system.prm7"), str(Path(tmpdir) / "traj_0.00000.dcd")]
+            [str(Path(tmpdir) / "system0.prm7"), str(Path(tmpdir) / "traj_0.00000.dcd")]
         )
 
         # Check that both config and lambda have been written
@@ -89,7 +89,7 @@ def test_restart(mols, request):
 
         # Reload the trajectory.
         traj_2 = sr.load(
-            [str(Path(tmpdir) / "system.prm7"), str(Path(tmpdir) / "traj_0.00000.dcd")]
+            [str(Path(tmpdir) / "system0.prm7"), str(Path(tmpdir) / "traj_0.00000.dcd")]
         )
 
         # Check that the trajectory is twice as long as the first.