From 5d56c7c4ac8d56ab952d115d1df538f616e65435 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Fri, 3 Jul 2026 17:06:36 +0100 Subject: [PATCH] Backport fix from PR #169. [ci skip] --- src/somd2/runner/_base.py | 67 +++++++++++++++++++--------- tests/runner/test_alchemical_ions.py | 67 +++++++++++++++++++++++++++- 2 files changed, 112 insertions(+), 22 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index c9734001..589da8df 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -360,11 +360,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. @@ -387,24 +390,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 @@ -414,6 +403,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 @@ -1061,6 +1064,15 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None): system: :class: `System ` The perturbed system with alchemical ions added. + + restraints: :class: `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 `. """ from sire.legacy.IO import createChlorineIon as _createChlorineIon @@ -1126,6 +1138,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 @@ -1241,6 +1256,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()) @@ -1256,7 +1274,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): diff --git a/tests/runner/test_alchemical_ions.py b/tests/runner/test_alchemical_ions.py index 6024d1a3..f4a4bfe7 100644 --- a/tests/runner/test_alchemical_ions.py +++ b/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()