From 2d3f6ebd685f5cbd448931eb12fb4b27d7fe0be2 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 20:25:35 +0100 Subject: [PATCH] Backport fixes from PR #259. [ci skip] --- src/somd2/runner/_base.py | 62 ++++++++++++++++++---- src/somd2/runner/_repex.py | 13 +---- src/somd2/runner/_runner.py | 12 +---- tests/runner/test_boresch_search.py | 48 +++++++++++++++++ tests/runner/test_ring_break_restraints.py | 36 +++++++++++-- 5 files changed, 134 insertions(+), 37 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 7ed0183d..ffb7106e 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -296,13 +296,17 @@ def __init__(self, system, config): self._config._extra_args["num_gcmc_waters"] = self._config.gcmc_num_waters # Auto-generate Morse restraints for ring-breaking perturbations with no - # user-supplied restraint. This is done before any modification of the - # bonded terms, and before the reference system is stored and the restart - # checks are performed, since the hard restraint replaces a bond in the - # system, which must match the checkpoints. - if self._is_ring_break and self._config.restraints is None: + # user-supplied Morse restraints. This is done before any modification of + # the bonded terms, and before the reference system is stored and the + # restart checks are performed, since the hard restraint replaces a bond + # in the system, which must match the checkpoints. + if self._is_ring_break and not self._has_restraints( + _sr.mm.MorsePotentialRestraints + ): try: - self._config.restraints = self._generate_morse_restraints() + self._config.restraints = ( + self._config.restraints or [] + ) + self._generate_morse_restraints() except Exception as e: msg = ( "Unable to generate Morse restraints for ring-breaking " @@ -1103,11 +1107,47 @@ def _is_ring_break(self): "reverse_ring_break_morph", ) + def _has_restraints(self, restraint_type): + """ + Return whether any of the user-supplied restraints are of the given type. + """ + return any( + isinstance(restraint, restraint_type) + for restraint in self._config.restraints or [] + ) + + def _add_boresch_restraint(self, device=None): + """ + Generate a Boresch restraint for an ABFE bound leg with no user-supplied + Boresch restraint, and add it to the restraints passed to the dynamics. + It isn't added to the config, since it is generated after the config is + written, and is instead saved separately for reuse on restart. + + Parameters + ---------- + + device : int, optional + GPU device number to use for the restraint-search run. + """ + if not self._is_abfe_bound or self._has_restraints(_sr.mm.BoreschRestraints): + return + + try: + restraints = self._generate_boresch_restraint(device=device) + except Exception as e: + msg = f"Unable to generate Boresch restraint for ABFE simulation: {e}" + _logger.error(msg) + raise RuntimeError(msg) + + self._dynamics_kwargs["restraints"] = ( + self._dynamics_kwargs["restraints"] or [] + ) + [restraints] + def _generate_morse_restraints(self): """ Return the pair of Morse restraints required by the 'ring_break_morph' lambda schedule, or its reverse. Called automatically when running a - ring-breaking simulation with no user-supplied restraint. + ring-breaking simulation with no user-supplied Morse restraints. The "hard" restraint directly replaces the harmonic bond that is broken (or formed) by the perturbation, inheriting its force constant and @@ -1131,7 +1171,7 @@ def _generate_morse_restraints(self): from sire.restraints import morse_potential as _morse_potential _logger.info( - "No restraints supplied for ring-breaking perturbation. " + "No Morse restraints supplied for ring-breaking perturbation. " "Generating default Morse restraints." ) @@ -1194,8 +1234,8 @@ def _generate_boresch_restraint(self, device=None): Return a Boresch restraint for the ABFE simulation, either by loading one saved from a previous run or by running a short lambda=0 trajectory. Called automatically before minimisation/equilibration of the production - windows/replicas when the simulation is ABFE and no restraint has been - supplied. + windows/replicas when the simulation is ABFE and no Boresch restraint has + been supplied. The input system is assumed to already be equilibrated (SOMD2 does not run a separate equilibration stage for the restraint search); it is @@ -1267,7 +1307,7 @@ def _generate_boresch_restraint(self, device=None): return restraints _logger.info( - "No restraint supplied for ABFE. Running Boresch restraint search." + "No Boresch restraint supplied for ABFE. Running Boresch restraint search." ) # Both protocols need a minimum number of trajectory frames. diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 1f3ad7e5..db43f3f5 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -1264,19 +1264,10 @@ def __init__(self, system, config): self._set_num_slots() - # Auto-generate a Boresch restraint for ABFE runs with no user-supplied - # restraint. This must happen before the dynamics cache is built below, + # Add the Boresch restraint before the dynamics cache is built below, # since the per-replica OpenMM contexts it creates are fixed at # construction time and won't pick up a restraint added afterwards. - if self._is_abfe_bound and self._config.restraints is None: - try: - restraints = self._generate_boresch_restraint(device=0) - except Exception as e: - msg = f"Unable to generate Boresch restraint for ABFE simulation: {e}" - _logger.error(msg) - raise RuntimeError(msg) - self._config.restraints = restraints - self._dynamics_kwargs["restraints"] = restraints + self._add_boresch_restraint(device=0) # Store the name of the dynamics cache pickle file. self._repex_state = self._config.output_directory / "repex_state.pkl" diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 87f634b5..b9c1e06e 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -270,17 +270,7 @@ def run(self): else: self._max_workers = 1 - # Auto-generate a Boresch restraint for ABFE runs with no user-supplied restraint. - if self._is_abfe_bound and self._config.restraints is None: - device = self._gpu_pool[0] if self._is_gpu else None - try: - restraints = self._generate_boresch_restraint(device=device) - except Exception as e: - msg = f"Unable to generate Boresch restraint for ABFE simulation: {e}" - _logger.error(msg) - raise RuntimeError(msg) - self._config.restraints = restraints - self._dynamics_kwargs["restraints"] = restraints + self._add_boresch_restraint(device=self._gpu_pool[0] if self._is_gpu else None) import concurrent.futures as _futures import multiprocessing as _mp diff --git a/tests/runner/test_boresch_search.py b/tests/runner/test_boresch_search.py index d51342e0..896d574c 100644 --- a/tests/runner/test_boresch_search.py +++ b/tests/runner/test_boresch_search.py @@ -38,6 +38,54 @@ def test_too_few_frames_raises(abfe_charge_change_mols): runner._generate_boresch_restraint() +def test_restraint_added_to_user_restraints(abfe_charge_change_mols, monkeypatch): + """A generated Boresch restraint is added to the user's other restraints.""" + mols = abfe_charge_change_mols.clone() + distance = sr.restraints.distance(mols, atoms0=0, atoms1=1, name="loop") + boresch = sr.mm.BoreschRestraints() + + monkeypatch.setattr( + Runner, "_generate_boresch_restraint", lambda self, device=None: boresch + ) + + with tempfile.TemporaryDirectory() as tmpdir: + config = Config( + output_directory=tmpdir, + platform="cpu", + lambda_schedule="decouple", + restraints=distance, + ) + runner = Runner(mols, config) + runner._add_boresch_restraint() + + assert runner._dynamics_kwargs["restraints"] == [distance, boresch] + assert runner._config.restraints == [distance] + + +def test_user_boresch_restraint_is_kept(abfe_charge_change_mols, monkeypatch): + """No Boresch restraint is generated when the user supplies one.""" + mols = abfe_charge_change_mols.clone() + distance = sr.restraints.distance(mols, atoms0=0, atoms1=1, name="loop") + boresch = sr.mm.BoreschRestraints() + + def generate(self, device=None): + raise AssertionError("A Boresch restraint should not be generated.") + + monkeypatch.setattr(Runner, "_generate_boresch_restraint", generate) + + with tempfile.TemporaryDirectory() as tmpdir: + config = Config( + output_directory=tmpdir, + platform="cpu", + lambda_schedule="decouple", + restraints=[distance, boresch], + ) + runner = Runner(mols, config) + runner._add_boresch_restraint() + + assert runner._dynamics_kwargs["restraints"] == [distance, boresch] + + def test_rxrx_unsuitable_falls_back(abfe_charge_change_mols): """A ligand with no N/O atoms can't use RXRX, so Aldeghi is used instead.""" mols = abfe_charge_change_mols.clone() diff --git a/tests/runner/test_ring_break_restraints.py b/tests/runner/test_ring_break_restraints.py index 255931a4..42b05542 100644 --- a/tests/runner/test_ring_break_restraints.py +++ b/tests/runner/test_ring_break_restraints.py @@ -161,10 +161,10 @@ def test_reverse_schedule_generates_restraints(syk_ring_break_mols): assert len(runner._config.restraints) == 2 -def test_user_restraints_are_not_overridden(syk_ring_break_mols): +def test_restraints_added_to_user_restraints(syk_ring_break_mols): """ - Ensure that a user-supplied restraint is left alone, and that the system is - not modified behind their back. + Ensure that the generated Morse restraints are added to a user-supplied + restraint of a different type, rather than being skipped. """ import sire as sr @@ -175,13 +175,41 @@ def test_user_restraints_are_not_overridden(syk_ring_break_mols): atoms0=0, atoms1=1, k="10 kcal mol-1 A-2", + name="loop", ) with tempfile.TemporaryDirectory() as tmpdir: runner = Runner(mols, _config(tmpdir, restraints=restraints)) - assert len(runner._config.restraints) == 1 + assert len(runner._config.restraints) == 3 assert runner._config.restraints[0] == restraints + _restraint(runner._config.restraints, "morse_hard") + _restraint(runner._config.restraints, "morse_soft") + assert runner._dynamics_kwargs["restraints"] is runner._config.restraints + + +def test_user_morse_restraints_are_not_overridden(syk_ring_break_mols): + """ + Ensure that user-supplied Morse restraints are left alone, and that the + system is not modified behind their back. + """ + import sire as sr + + mols = syk_ring_break_mols.clone() + + hard_restraints, mols = sr.restraints.morse_potential( + mols, + de="150 kcal mol-1", + auto_parametrise=True, + direct_morse_replacement=True, + name="morse_hard", + ) + distance = sr.restraints.distance(mols, atoms0=0, atoms1=1, name="loop") + + with tempfile.TemporaryDirectory() as tmpdir: + runner = Runner(mols, _config(tmpdir, restraints=[distance, hard_restraints])) + + assert runner._config.restraints == [distance, hard_restraints] @pytest.mark.parametrize("schedule", ["standard_morph", "charge_scaled_morph"])