Backport fixes from PR #259

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.

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"

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

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()

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"])