Fix issue #262

CHANGELOG.md

@@ -31,6 +31,7 @@ Changelog
 * Further speed up checkpointing, and join the final trajectory without reloading the topology [#257](https://github.com/OpenBioSim/somd2/pull/257).
 * Record frame times in final trajectories, rather than zero for every frame [#257](https://github.com/OpenBioSim/somd2/pull/257).
 * Fixed the Boresch restraint for ABFE and the Morse restraints for ring-breaking perturbations not being generated when any other restraint was passed. They are now generated unless restraints of the same type are passed, and added to the others [#259](https://github.com/OpenBioSim/somd2/pull/259).
+* Fixed alchemical ions changing their charge linearly over the whole lambda range for the `annihilate`, `decouple`, `charge_scaled_morph`, `ring_break_morph` and `reverse_ring_break_morph` lambda schedules, rather than alongside the perturbable molecule. A warning is raised for custom lambda schedules, where alchemical ions still use the `standard_morph` schedule [#263](https://github.com/OpenBioSim/somd2/pull/263).
 
 [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026
 --------------------------------------------------------------------------------------

src/somd2/_utils/_schedules.py

@@ -20,6 +20,7 @@
 #####################################################################
 
 __all__ = [
+    "abfe_alchemical_ion",
     "annihilate",
     "decouple",
     "ring_break_morph",
@@ -232,6 +233,40 @@ def decouple(fix_epsilon=True, restraint_lever="split"):
     return s
 
 
+def abfe_alchemical_ion(schedule):
+    """
+    Build the lambda schedule for an alchemical ion used with an ABFE lambda
+    schedule.
+
+    The ion is morphed during the decharge stage, so that the system stays
+    neutral while the ligand is decharged, then held at its final state.
+
+    Parameters
+    ----------
+
+    schedule : sire.legacy.CAS.LambdaSchedule
+        The ABFE lambda schedule used for the ligand.
+
+    Returns
+    -------
+
+    schedule : sire.legacy.CAS.LambdaSchedule
+        The lambda schedule for the alchemical ion.
+    """
+    from sire.cas import LambdaSchedule as _LambdaSchedule
+
+    s = _LambdaSchedule()
+
+    for stage, weight in zip(schedule.get_stages(), schedule.get_stage_weights()):
+        if stage == "decharge":
+            equation = (1 - s.lam()) * s.initial() + s.lam() * s.final()
+        else:
+            equation = s.final()
+        s.add_stage(stage, equation, weight)
+
+    return s
+
+
 def ring_break_morph():
     """
     Build a lambda schedule for ring-breaking perturbations.

src/somd2/runner/_base.py

@@ -584,20 +584,44 @@ def _set_schedule(self, pert_idxs, ion_indices):
                 restraint_lever = self._config._boresch_restraint_lever() or "split"
                 _logger.info(f"Using the '{restraint_lever}' Boresch restraint lever.")
 
-        # 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).
+        # Alchemical ions change their charge alongside the ligand so that the
+        # system stays neutral. The morph schedules can be shared, but the ABFE
+        # schedules' ghost-atom equations, e.g. fixed epsilon, are meant for the
+        # ligand, so the ion gets its own schedule aligned with the decharge stage.
         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()
+            schedule = self._config._lambda_schedule
+            schedule_name = self._config._lambda_schedule_name
+
+            # Also match ABFE schedules passed as objects, e.g. from decouple().
+            is_abfe = schedule_name in ("annihilate", "decouple") or (
+                schedule_name is None
+                and schedule.get_stages()
+                in (["decharge", "annihilate"], ["decharge", "decouple"])
+            )
+
+            if is_abfe:
+                from .._utils._schedules import (
+                    abfe_alchemical_ion as _abfe_alchemical_ion,
                 )
 
+                ion_schedule = _abfe_alchemical_ion(schedule)
+            elif schedule_name is None and schedule != _LambdaSchedule.standard_morph():
+                _logger.warning(
+                    "Alchemical ions use the 'standard_morph' lambda schedule with a "
+                    "custom lambda schedule, so intermediate states will only be "
+                    "neutral if the perturbable molecule's charge also changes "
+                    "linearly with lambda."
+                )
+                ion_schedule = _LambdaSchedule.standard_morph()
+            else:
+                ion_schedule = None
+
+            if ion_schedule is not None:
+                for pert_idx in ion_indices:
+                    schedule.set_molecule_schedule(pert_idx, ion_schedule)
+
         # Set the lambda values.
         if self._config.lambda_values:
             self._lambda_values = self._config.lambda_values

tests/runner/test_alchemical_ions.py

@@ -194,16 +194,29 @@ def test_alchemical_ion_restart_reuses_same_ion(ethane_methanol_ions):
         assert ion_number_1 == ion_number_2
 
 
-@pytest.mark.parametrize("schedule_name", ["decouple", "annihilate"])
-def test_alchemical_ion_abfe_schedule(schedule_name, ethane_methanol_ions):
+@pytest.mark.parametrize(
+    "schedule_name",
+    [
+        "decouple",
+        "annihilate",
+        "standard_morph",
+        "charge_scaled_morph",
+        "custom_standard_morph",
+        "custom_decouple",
+        "custom_annihilate",
+        "custom",
+    ],
+)
+def test_alchemical_ion_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.
+    Ensure that an alchemical ion follows the ligand's schedule for the morph
+    schedules, gets a schedule aligned with the decharge stage for the ABFE
+    schedules, and falls back to a plain morph for other custom schedules.
     """
     from sire.cas import LambdaSchedule
 
+    from somd2._utils._schedules import annihilate, decouple
+
     # Clone the system (charge-neutral, but with plenty of waters and existing
     # free ions to match parameters against).
     mols = ethane_methanol_ions.clone()
@@ -212,7 +225,12 @@ def test_alchemical_ion_abfe_schedule(schedule_name, ethane_methanol_ions):
         config = Config(
             output_directory=tmpdir,
             platform="cpu",
-            lambda_schedule=schedule_name,
+            lambda_schedule={
+                "custom_standard_morph": LambdaSchedule.standard_morph(),
+                "custom_decouple": decouple(fix_epsilon=False),
+                "custom_annihilate": annihilate(fix_epsilon=False),
+                "custom": LambdaSchedule.charge_scaled_morph(0.5),
+            }.get(schedule_name, schedule_name),
             # Force a charge difference so that an alchemical ion is added,
             # even though the fixture itself is charge neutral.
             charge_difference=1,
@@ -238,16 +256,26 @@ def test_alchemical_ion_abfe_schedule(schedule_name, ethane_methanol_ions):
 
         schedule = runner._config.lambda_schedule
 
-        # The ligand follows the main ABFE schedule directly - no override.
+        # The ligand always follows the main schedule directly.
         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()
+
+        if schedule_name == "custom":
+            assert ion_schedule.to_string() == (
+                LambdaSchedule.standard_morph().to_string()
+            )
+            return
+
+        # The ion's charge changes at the same rate as the ligand's.
+        for lv in [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]:
+            assert ion_schedule.morph("clj", "charge", 0.0, 1.0, lv) == (
+                pytest.approx(schedule.morph("clj", "charge", 0.0, 1.0, lv))
+            )
+
+        # The ion's LJ parameters reach their final values.
+        for lever in ("sigma", "epsilon"):
+            assert ion_schedule.morph("clj", lever, 0.0, 1.0, 1.0) == pytest.approx(1.0)
 
 
 def test_charge_difference_with_virtual_sites(ethane_methanol):

tests/schedules/test_abfe.py

@@ -1,6 +1,6 @@
 import pytest
 
-from somd2._utils._schedules import annihilate, decouple
+from somd2._utils._schedules import abfe_alchemical_ion, annihilate, decouple
 
 # Lambda schedules are always symmetric between the two stages (decharge is
 # the first stage, annihilate/decouple is the second), so both builders share
@@ -85,6 +85,29 @@ def test_restraint_lever_split(builder):
     assert distance_angle_values == pytest.approx(_SPLIT_DISTANCE_ANGLE, abs=1e-6)
 
 
+@pytest.mark.parametrize("builder", BUILDERS)
+@pytest.mark.parametrize("fix_epsilon", [True, False])
+def test_alchemical_ion(builder, fix_epsilon):
+    """
+    The alchemical ion's charge changes at the same rate as the ligand's,
+    keeping the system neutral, and its LJ parameters reach their final values.
+    """
+    schedule = builder(fix_epsilon=fix_epsilon)
+    schedule.set_stage_weight("decharge", 2.0)
+    ion_schedule = abfe_alchemical_ion(schedule)
+
+    assert ion_schedule.get_stages() == schedule.get_stages()
+    assert ion_schedule.get_stage_weights() == schedule.get_stage_weights()
+
+    for lv in _LAMBDA_VALUES:
+        assert ion_schedule.morph("clj", "charge", 0.0, 1.0, lv) == pytest.approx(
+            schedule.morph("clj", "charge", 0.0, 1.0, lv)
+        )
+
+    for lever in ("charge", "sigma", "epsilon"):
+        assert ion_schedule.morph("clj", lever, 0.0, 1.0, 1.0) == pytest.approx(1.0)
+
+
 @pytest.mark.skipif(
     "openmm" not in __import__("sire").convert.supported_formats(),
     reason="openmm support is not available",