From e806ab45f9ab9945e3c5ec0374338dc75bbd9c70 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 20:13:01 +0100 Subject: [PATCH] Add auto-generated Boresch and Morse restraints to user restraints. --- CHANGELOG.md | 1 + README.md | 21 ++++---- 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 +++++++++++-- 7 files changed, 147 insertions(+), 46 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bfb43c44..cc8f5c8f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,7 @@ Changelog * Reduce the setup time of replica exchange simulations that re-use OpenMM contexts across lambda values. Minimisation and equilibration now create their dynamics objects once per context, rather than once per replica [#255](https://github.com/OpenBioSim/somd2/pull/255). * 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). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 -------------------------------------------------------------------------------------- diff --git a/README.md b/README.md index 83c92f3d..04aa9bfe 100644 --- a/README.md +++ b/README.md @@ -648,10 +648,12 @@ somd2 bound.s3 --lambda-schedule decouple > bonded terms, so the ligand falls apart as λ increases unless every atom is > restrained. -The ligand must be restrained within the binding site. If no restraints are -passed, a Boresch restraint is generated automatically for the bound leg, i.e. -when the system contains a protein as well as the ligand and water. This is -done by minimising the system, running a short trajectory at λ = 0, then +The ligand must be restrained within the binding site. If no Boresch restraint +is passed, one is generated automatically for the bound leg, i.e. when the +system contains a protein as well as the ligand and water, and is added to any +other restraints that are passed. A Boresch restraint is therefore always used +for the bound leg, since the standard state correction depends on it. It is +generated by minimising the system, running a short trajectory at λ = 0, then choosing the anchor atoms and force constants from it. The length of this trajectory and the frequency at which frames are saved can be controlled with the `--restraint-search-time` and `--restraint-search-frequency` options. By @@ -682,11 +684,12 @@ somd2 perturbable_system.bss --lambda-schedule ring_break_morph ``` These perturbations require a pair of Morse restraints on the atoms of the bond -that is broken. If no restraints are passed, both are generated automatically. -A "hard" Morse potential replaces the harmonic bond, inheriting its force -constant and equilibrium length, and is switched off as a weaker "soft" Morse -restraint holds the fragment in place. Their well depths and the force constant -of the soft restraint can be controlled with the `--morse-hard-well-depth`, +that is broken. If no Morse restraints are passed, both are generated +automatically and added to any other restraints that are passed. A "hard" Morse +potential replaces the harmonic bond, inheriting its force constant and +equilibrium length, and is switched off as a weaker "soft" Morse restraint holds +the fragment in place. Their well depths and the force constant of the soft +restraint can be controlled with the `--morse-hard-well-depth`, `--morse-soft-well-depth`, and `--morse-soft-force-constant` options. Unlike the ABFE restraints, these are regenerated on each run rather than being diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index 98f985a7..237fd2a4 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -361,13 +361,17 @@ def _prepare_system(self): 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 " @@ -1335,11 +1339,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 @@ -1363,7 +1403,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." ) @@ -1426,8 +1466,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 @@ -1499,7 +1539,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 36d9ee17..103510de 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -207,19 +207,10 @@ def __init__(self, system, config): else None ) - # 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 7b34b3e9..0bfd0cbc 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -268,17 +268,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"])