Fix issue #168

CHANGELOG.md

@@ -6,6 +6,7 @@ Changelog
 
 * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
 * Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166).
+* Give alchemical ions their own plain morph lambda schedule so they interpolate correctly under non-standard lambda schedules [#169](https://github.com/OpenBioSim/somd2/pull/169).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------

src/somd2/runner/_base.py

@@ -374,11 +374,14 @@ def __init__(self, system, config):
                 )
 
         # Create alchemical ions.
+        ion_indices = []
         if charge_diff != 0:
-            self._system, coalchemical_restraints = self._create_alchemical_ions(
-                self._system,
-                charge_diff,
-                restraint_distance=self._config.coalchemical_restraint_dist,
+            self._system, coalchemical_restraints, ion_indices = (
+                self._create_alchemical_ions(
+                    self._system,
+                    charge_diff,
+                    restraint_distance=self._config.coalchemical_restraint_dist,
+                )
             )
 
             # Add the coalchemical restraints to the extra args.
@@ -401,24 +404,10 @@ def __init__(self, system, config):
             self._config._extra_args["use_beutler_softening"] = True
             self._config._extra_args["beutler_alpha"] = self._config.beutler_alpha
 
-        # Build deferred schedules now that the softcore form is known. Epsilon is
-        # only held fixed (with LJ decay handled entirely by the Beutler soft-core
-        # prefactor) for molecules undergoing a ghost-atom decoupling/annihilation.
-        # An alchemical ion is a real (non-ghost) atom mutating identity (e.g. a
-        # water oxygen turning into Na+), so its LJ epsilon needs to interpolate
-        # normally; fixing it would leave the ion's persisting atom stuck at its
-        # initial LJ parameters for the whole stage. Disable fix_epsilon whenever
-        # an alchemical ion has been added, regardless of the configured value.
+        # Build deferred schedules now that the softcore form is known.
         fix_epsilon = (
             self._config.softcore_form == "beutler" and self._config.beutler_fix_epsilon
         )
-        if fix_epsilon and charge_diff != 0:
-            _logger.info(
-                "Disabling Beutler 'fix_epsilon' since an alchemical ion has been "
-                "added: the ion's persisting atom is a real (non-ghost) mutation "
-                "and needs its LJ epsilon to interpolate normally."
-            )
-            fix_epsilon = False
         if self._config._lambda_schedule_name == "annihilate":
             from .._utils._schedules import annihilate as _annihilate
 
@@ -428,6 +417,20 @@ def __init__(self, system, config):
 
             self._config._lambda_schedule = _decouple(fix_epsilon=fix_epsilon)
 
+        # Alchemical ions are real (non-ghost) atoms mutating identity (e.g. a
+        # water oxygen turning into Na+), not ghost-atom decoupling/annihilation
+        # like the main perturbable molecule(s). Give each one its own plain
+        # morph schedule so it always interpolates smoothly across the whole
+        # lambda range, regardless of the ligand's schedule (e.g. so it isn't
+        # held fixed outside of a "decharge" stage, or subject to fix_epsilon).
+        if ion_indices:
+            from sire.cas import LambdaSchedule as _LambdaSchedule
+
+            for pert_idx in ion_indices:
+                self._config._lambda_schedule.set_molecule_schedule(
+                    pert_idx, _LambdaSchedule.standard_morph()
+                )
+
         # Set the lambda values.
         if self._config.lambda_values:
             self._lambda_values = self._config.lambda_values
@@ -1203,6 +1206,15 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
 
         system: :class: `System <sire.system.System>`
             The perturbed system with alchemical ions added.
+
+        restraints: :class: `Restraints <sire.restraints.Restraints>`
+            The coalchemical restraints, or None if no restraint distance
+            was specified.
+
+        ion_indices: [int]
+            The perturbable-molecule index of each alchemical ion that was
+            added, suitable for use with
+            `LambdaSchedule.set_molecule_schedule <sire.cas.LambdaSchedule>`.
         """
 
         from sire.legacy.IO import createChlorineIon as _createChlorineIon
@@ -1268,6 +1280,9 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
         # Create a null set of coalchemical restraints.
         restraints = None
 
+        # Store the molecule numbers of the alchemical ions.
+        ion_numbers = []
+
         # Create the ions.
         for water in waters:
             # Flag to indicate whether we need to reverse the alchemical ion
@@ -1383,6 +1398,9 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
             # Update the system.
             system.update(merged)
 
+            # Record the molecule number of the alchemical ion.
+            ion_numbers.append(water.number())
+
             # Get the index of the perturbed water.
             index = numbers.index(water.number())
 
@@ -1398,7 +1416,16 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
                     f"{ion_str} ion to keep charge constant."
                 )
 
-        return system, restraints
+        # Work out the perturbable-molecule index of each alchemical ion, now
+        # that the full system (existing perturbable molecule(s) plus all of
+        # the newly merged ions) is in its final state. This index is what
+        # Sire's LambdaSchedule.set_molecule_schedule() expects, and depends on
+        # the relative order of *all* perturbable molecules in the system, so
+        # it can only be computed once every ion has been added.
+        perturbable_mols = system.molecules()["perturbable"].molecules()
+        ion_indices = [perturbable_mols.find(system[number]) for number in ion_numbers]
+
+        return system, restraints, ion_indices
 
     @staticmethod
     def _create_filenames(lambda_array, lambda_value, output_directory, restart=False):

tests/runner/test_alchemical_ions.py

@@ -1,6 +1,9 @@
 import math
+import tempfile
+
 import pytest
 
+from somd2.config import Config
 from somd2.runner import Runner
 
 
@@ -12,13 +15,73 @@ def test_alchemical_ions(mols, request):
     mols = request.getfixturevalue(mols).clone()
 
     # Add 10 Cl- ions.
-    new_mols, _ = Runner._create_alchemical_ions(mols, 10)
+    new_mols, _, ion_indices = 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)
 
+    # Make sure there is one perturbable-molecule index per ion.
+    assert len(ion_indices) == 10
+
     # Add 10 Na+ ions.
-    new_mols, _ = Runner._create_alchemical_ions(mols, -10)
+    new_mols, _, ion_indices = 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)
+    assert len(ion_indices) == 10
+
+
+@pytest.mark.parametrize("schedule_name", ["decouple", "annihilate"])
+def test_alchemical_ion_abfe_schedule(schedule_name, ethane_methanol_ions):
+    """
+    Ensure that an alchemical ion added alongside an ABFE (decouple/annihilate)
+    perturbable molecule gets its own plain morph schedule, rather than
+    inheriting the ghost-atom decoupling/annihilation lever equations meant
+    for the main perturbable molecule.
+    """
+    from sire.cas import LambdaSchedule
+
+    # Clone the system (charge-neutral, but with plenty of waters and existing
+    # free ions to match parameters against).
+    mols = ethane_methanol_ions.clone()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = Config(
+            output_directory=tmpdir,
+            platform="cpu",
+            lambda_schedule=schedule_name,
+            # Force a charge difference so that an alchemical ion is added,
+            # even though the fixture itself is charge neutral.
+            charge_difference=1,
+        )
+
+        runner = Runner(mols, config)
+
+        # There should be exactly two perturbable molecules now: the original
+        # ligand and the single alchemical ion that was added.
+        perturbable_mols = runner._system.molecules()["perturbable"].molecules()
+        assert len(perturbable_mols) == 2
+
+        ligand_idx = None
+        ion_idx = None
+        for i, mol in enumerate(perturbable_mols):
+            if mol.has_property("is_alchemical_ion"):
+                ion_idx = i
+            else:
+                ligand_idx = i
+
+        assert ligand_idx is not None
+        assert ion_idx is not None
+
+        schedule = runner._config.lambda_schedule
+
+        # The ligand follows the main ABFE schedule directly - no override.
+        assert not schedule.has_molecule_schedule(ligand_idx)
+
+        # The ion has its own molecule-specific schedule, and it is a plain
+        # morph (not the ligand's decharge/decouple or decharge/annihilate
+        # staging).
+        assert schedule.has_molecule_schedule(ion_idx)
+        ion_schedule = schedule.get_molecule_schedule(ion_idx)
+        assert ion_schedule.get_stages() == ["morph"]
+        assert ion_schedule.to_string() == LambdaSchedule.standard_morph().to_string()