diff --git a/CHANGELOG.md b/CHANGELOG.md index cc8f5c8f..159f3fcf 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. 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 -------------------------------------------------------------------------------------- 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..b0e066e1 100644 --- a/src/somd2/runner/_base.py +++ b/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 diff --git a/tests/runner/test_alchemical_ions.py b/tests/runner/test_alchemical_ions.py index b6707968..64e06e31 100644 --- a/tests/runner/test_alchemical_ions.py +++ b/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): 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",