From b66a70459162ba279ef0235cf93cad9031f5ce75 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 22 Oct 2026 15:16:57 +0100 Subject: [PATCH 1/4] Align alchemical ion charge changes with the lambda schedule. --- CHANGELOG.md | 1 + src/somd2/_utils/_schedules.py | 35 +++++++++++++++++++++++++ src/somd2/runner/_base.py | 34 ++++++++++++++++++------- tests/runner/test_alchemical_ions.py | 38 ++++++++++++++++++---------- tests/schedules/test_abfe.py | 25 +++++++++++++++++- 5 files changed, 110 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cc8f5c8f..2e8f5735 100644 --- a/CHANGELOG.md +++ b/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, which left intermediate states with a net charge. A warning is raised for custom lambda schedules, where alchemical ions still use the `standard_morph` schedule [#262](https://github.com/OpenBioSim/somd2/pull/262). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/_utils/_schedules.py b/src/somd2/_utils/_schedules.py index 1f1f0e22..7c017932 100644 --- a/src/somd2/_utils/_schedules.py +++ b/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. diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 237fd2a4..572eb745 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -584,19 +584,35 @@ 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 would hold the ion's epsilon fixed, so need their own. 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_name = self._config._lambda_schedule_name + if schedule_name in ("annihilate", "decouple"): + from .._utils._schedules import ( + abfe_alchemical_ion as _abfe_alchemical_ion, + ) + + ion_schedule = _abfe_alchemical_ion(self._config._lambda_schedule) + elif schedule_name is None: + _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: + self._config._lambda_schedule.set_molecule_schedule( + pert_idx, ion_schedule + ) # Set the lambda values. if self._config.lambda_values: diff --git a/tests/runner/test_alchemical_ions.py b/tests/runner/test_alchemical_ions.py index b6707968..5f8813bb 100644 --- a/tests/runner/test_alchemical_ions.py +++ b/tests/runner/test_alchemical_ions.py @@ -194,16 +194,20 @@ 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"], +) +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 custom schedules. """ from sire.cas import LambdaSchedule + from somd2._utils._schedules import abfe_alchemical_ion + # 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 +216,11 @@ def test_alchemical_ion_abfe_schedule(schedule_name, ethane_methanol_ions): config = Config( output_directory=tmpdir, platform="cpu", - lambda_schedule=schedule_name, + lambda_schedule=( + LambdaSchedule.standard_morph() + if schedule_name == "custom" + else 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 +246,20 @@ 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). + if schedule_name in ("decouple", "annihilate"): + expected = abfe_alchemical_ion(schedule) + elif schedule_name == "custom": + expected = LambdaSchedule.standard_morph() + else: + assert not schedule.has_molecule_schedule(ion_idx) + return + 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() + assert ion_schedule.to_string() == expected.to_string() def test_charge_difference_with_virtual_sites(ethane_methanol): diff --git a/tests/schedules/test_abfe.py b/tests/schedules/test_abfe.py index e44a25a7..cc40bb2a 100644 --- a/tests/schedules/test_abfe.py +++ b/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", From 9a1f02afdb2873e95c205c6f75d80cb3940c1dd2 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 22 Oct 2026 15:19:59 +0100 Subject: [PATCH 2/4] Skip the alchemical ion warning for a custom standard_morph schedule. --- src/somd2/runner/_base.py | 5 ++++- tests/runner/test_alchemical_ions.py | 20 +++++++++++++------- 2 files changed, 17 insertions(+), 8 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 572eb745..5a0a1636 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -597,7 +597,10 @@ def _set_schedule(self, pert_idxs, ion_indices): ) ion_schedule = _abfe_alchemical_ion(self._config._lambda_schedule) - elif schedule_name is None: + elif ( + schedule_name is None + and self._config._lambda_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 " diff --git a/tests/runner/test_alchemical_ions.py b/tests/runner/test_alchemical_ions.py index 5f8813bb..c6e21882 100644 --- a/tests/runner/test_alchemical_ions.py +++ b/tests/runner/test_alchemical_ions.py @@ -196,13 +196,20 @@ def test_alchemical_ion_restart_reuses_same_ion(ethane_methanol_ions): @pytest.mark.parametrize( "schedule_name", - ["decouple", "annihilate", "standard_morph", "charge_scaled_morph", "custom"], + [ + "decouple", + "annihilate", + "standard_morph", + "charge_scaled_morph", + "custom_standard_morph", + "custom", + ], ) def test_alchemical_ion_schedule(schedule_name, ethane_methanol_ions): """ 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 custom schedules. + schedules, and falls back to a plain morph for other custom schedules. """ from sire.cas import LambdaSchedule @@ -216,11 +223,10 @@ def test_alchemical_ion_schedule(schedule_name, ethane_methanol_ions): config = Config( output_directory=tmpdir, platform="cpu", - lambda_schedule=( - LambdaSchedule.standard_morph() - if schedule_name == "custom" - else schedule_name - ), + lambda_schedule={ + "custom_standard_morph": LambdaSchedule.standard_morph(), + "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, From 1a4d0c0747d0cf09fa3cee6f90dee59293fa47f4 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 22 Oct 2026 15:26:09 +0100 Subject: [PATCH 3/4] Small fixes for alchemical ion schedule handling. --- CHANGELOG.md | 2 +- src/somd2/runner/_base.py | 25 +++++++++++++---------- tests/runner/test_alchemical_ions.py | 30 ++++++++++++++++++---------- 3 files changed, 36 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e8f5735..3f816ab6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,7 +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, which left intermediate states with a net charge. A warning is raised for custom lambda schedules, where alchemical ions still use the `standard_morph` schedule [#262](https://github.com/OpenBioSim/somd2/pull/262). +* 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 [#262](https://github.com/OpenBioSim/somd2/pull/262). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 5a0a1636..b0e066e1 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -586,21 +586,28 @@ def _set_schedule(self, pert_idxs, ion_indices): # Alchemical ions change their charge alongside the ligand so that the # system stays neutral. The morph schedules can be shared, but the ABFE - # schedules would hold the ion's epsilon fixed, so need their own. + # 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 + schedule = self._config._lambda_schedule schedule_name = self._config._lambda_schedule_name - if schedule_name in ("annihilate", "decouple"): + + # 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(self._config._lambda_schedule) - elif ( - schedule_name is None - and self._config._lambda_schedule != _LambdaSchedule.standard_morph() - ): + 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 " @@ -613,9 +620,7 @@ def _set_schedule(self, pert_idxs, ion_indices): if ion_schedule is not None: for pert_idx in ion_indices: - self._config._lambda_schedule.set_molecule_schedule( - pert_idx, ion_schedule - ) + schedule.set_molecule_schedule(pert_idx, ion_schedule) # Set the lambda values. if self._config.lambda_values: diff --git a/tests/runner/test_alchemical_ions.py b/tests/runner/test_alchemical_ions.py index c6e21882..64e06e31 100644 --- a/tests/runner/test_alchemical_ions.py +++ b/tests/runner/test_alchemical_ions.py @@ -202,6 +202,8 @@ def test_alchemical_ion_restart_reuses_same_ion(ethane_methanol_ions): "standard_morph", "charge_scaled_morph", "custom_standard_morph", + "custom_decouple", + "custom_annihilate", "custom", ], ) @@ -213,7 +215,7 @@ def test_alchemical_ion_schedule(schedule_name, ethane_methanol_ions): """ from sire.cas import LambdaSchedule - from somd2._utils._schedules import abfe_alchemical_ion + 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). @@ -225,6 +227,8 @@ def test_alchemical_ion_schedule(schedule_name, ethane_methanol_ions): platform="cpu", 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, @@ -255,17 +259,23 @@ def test_alchemical_ion_schedule(schedule_name, ethane_methanol_ions): # The ligand always follows the main schedule directly. assert not schedule.has_molecule_schedule(ligand_idx) - if schedule_name in ("decouple", "annihilate"): - expected = abfe_alchemical_ion(schedule) - elif schedule_name == "custom": - expected = LambdaSchedule.standard_morph() - else: - assert not schedule.has_molecule_schedule(ion_idx) + ion_schedule = schedule.get_molecule_schedule(ion_idx) + + if schedule_name == "custom": + assert ion_schedule.to_string() == ( + LambdaSchedule.standard_morph().to_string() + ) return - assert schedule.has_molecule_schedule(ion_idx) - ion_schedule = schedule.get_molecule_schedule(ion_idx) - assert ion_schedule.to_string() == expected.to_string() + # 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): From 89ac1db205591ab099258dcbf4ee9f9548451efd Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 22 Oct 2026 15:46:25 +0100 Subject: [PATCH 4/4] Fix CHANGELOG PR number. --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3f816ab6..159f3fcf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,7 +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 [#262](https://github.com/OpenBioSim/somd2/pull/262). +* 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 --------------------------------------------------------------------------------------