From f58b8f11a741276c0d0439390b4b2a263a852cf0 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 29 Jun 2026 16:26:01 +0100
Subject: [PATCH 01/67] Update BioSimSpace development pin.

---
 pixi.toml                 | 4 ++--
 recipes/somd2/recipe.yaml | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pixi.toml b/pixi.toml
index e77834c..8a43389 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -6,9 +6,9 @@ platforms = ["linux-64", "osx-arm64"]
 [dependencies]
 python = ">=3.10"
 # main
-biosimspace = ">=2026.1.0,<2026.2.0"
+#biosimspace = ">=2026.1.0,<2026.2.0"
 # devel
-#biosimspace = "==2026.2.0.dev"
+biosimspace = "==2026.2.0.dev"
 filelock = "*"
 ghostly = "*"
 loch = "*"
diff --git a/recipes/somd2/recipe.yaml b/recipes/somd2/recipe.yaml
index 7303d8e..d8916d6 100644
--- a/recipes/somd2/recipe.yaml
+++ b/recipes/somd2/recipe.yaml
@@ -20,9 +20,9 @@ requirements:
     - versioningit
   run:
     # main
-    - biosimspace >=2026.1.0,<2026.2.0
+    #- biosimspace >=2026.1.0,<2026.2.0
     # devel
-    #- biosimspace ==2026.2.0.dev
+    - biosimspace ==2026.2.0.dev
     - filelock
     - ghostly
     - loch

From ebcc92c8bf085645495c65dd377c5fd5a65e0eec Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 29 Jun 2026 16:26:36 +0100
Subject: [PATCH 02/67] Update CHANGELOG for 2026.2.0 development.

---
 CHANGELOG.md | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1df7d6f..0ddf958 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,11 @@
 Changelog
 =========
 
+[2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - ********
+--------------------------------------------------------------------------------------
+
+* Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
+
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
 

From e5c79f4815e78da64b9afb02cc9b5098e96b37b8 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 29 Jun 2026 16:54:34 +0100
Subject: [PATCH 03/67] Add note regarding PyMBAR JAX issues.

---
 README.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/README.md b/README.md
index 6ec3b90..253709b 100644
--- a/README.md
+++ b/README.md
@@ -423,3 +423,7 @@ some systems, causing the simulation to hang. This can be resolved by either
 reducing the frequency at which frames are stored, or checkpointing more
 frequently. (Frames are written to disk and cleared from memory at each
 checkpoint.)
+
+PyMBAR uses JAX by default for GPU acceleration, which can cause issues in
+some environments. If you encounter issues when analysing simlation output,
+try setting the `PYMBAR_DISABLE_JAX` environment variable to `1`.

From c2f5660422f299d43d07bb47cb363911e2631668 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 1 Jul 2026 09:14:24 +0100
Subject: [PATCH 04/67] Remove pre-release warning. [ci skip]

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 253709b..cad485e 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,7 @@
 [![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0)
 
 Open-source GPU accelerated molecular dynamics engine for alchemical free-energy
-simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [OpenMM](https://github.com/openmm/openmm). The code is still under active development and is not yet ready for general use.
+simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [OpenMM](https://github.com/openmm/openmm).
 
 ## Installation
 

From 5eb74ec8217dd01eb2159b1c7bb4ffa679067449 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 1 Jul 2026 09:21:32 +0100
Subject: [PATCH 05/67] Clarify perturbed_system option. [ci skip]

---
 src/somd2/config/_config.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index d4482be..b350295 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -381,7 +381,8 @@ def __init__(
         perturbed_system: str
             The path to a stream file containing a Sire system for the equilibrated perturbed
             end state (lambda = 1). This will be used as the starting conformation all lambda
-            windows > 0.5 when performing a replica exchange simulation.
+            windows > 0.5 when performing a replica exchange simulation. (Note that this assumes
+            that the "coordinates1" property specifies the coordinates for perturbable molecules.)
 
         terminal_flip_frequency: str
             Frequency at which to attempt terminal ring flip Monte Carlo moves. If None

From 65ccd41761cfa5720444d99edd243fb56de21193 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 1 Jul 2026 12:02:10 +0100
Subject: [PATCH 06/67] Add support for generating Boresch restraints for ABFE.

---
 CHANGELOG.md                |   1 +
 src/somd2/config/_config.py | 110 ++++++++++++++++++++++++++++++-
 src/somd2/runner/_base.py   | 127 +++++++++++++++++++++++++++++++++++-
 src/somd2/runner/_repex.py  |  14 ++++
 src/somd2/runner/_runner.py |  12 ++++
 5 files changed, 261 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0ddf958..0190663 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@ Changelog
 --------------------------------------------------------------------------------------
 
 * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
+* Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index b350295..9efd79c 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -172,6 +172,9 @@ def __init__(
         save_xml=False,
         page_size=None,
         timeout="300 s",
+        restraint_search_time="1 ns",
+        restraint_search_frequency="10 ps",
+        restraint_search_receptor_selection=None,
     ):
         """
         Constructor.
@@ -557,6 +560,21 @@ def __init__(
         null_energy: str
             The energy value to use for lambda windows that are not
             being computed as part of the energy trajectory.
+
+        restraint_search_time: str
+            Length of the short pre-production trajectory used to auto-generate
+            a Boresch restraint when running an ABFE simulation without a
+            user-supplied restraint. Defaults to "1 ns".
+
+        restraint_search_frequency: str
+            Frame-saving frequency during the restraint-search trajectory.
+            Defaults to "10 ps". Should be small enough to yield at least 50
+            frames over ``restraint_search_time``.
+
+        restraint_search_receptor_selection: str
+            Sire selection string for receptor anchor atom candidates used
+            during automatic Boresch restraint generation. If None, the default
+            backbone selection is used (CA, C, N atoms in non-water molecules).
         """
 
         # Setup logger before doing anything else
@@ -645,9 +663,10 @@ def __init__(
         self.num_energy_neighbours = num_energy_neighbours
         self.null_energy = null_energy
         self.page_size = page_size
-
+        self.restraint_search_time = restraint_search_time
+        self.restraint_search_frequency = restraint_search_frequency
+        self.restraint_search_receptor_selection = restraint_search_receptor_selection
         self.write_config = write_config
-
         self.overwrite = overwrite
 
     def __str__(self):
@@ -2468,6 +2487,34 @@ def _from_hex(hex):
 
         return obj
 
+    def __getstate__(self):
+        """
+        Hex-encode the same fields that to_yaml()/from_yaml() already
+        hex-encode (currently 'restraints' and 'lambda_schedule'), since
+        these legacy Sire objects are not guaranteed to have native pickle
+        support. This is needed so that a Config holding these can be sent
+        to a spawned worker process, e.g. via
+        concurrent.futures.ProcessPoolExecutor.
+        """
+        state = self.__dict__.copy()
+        if state.get("_restraints") is not None:
+            state["_restraints"] = [
+                self._to_hex(restraint) for restraint in state["_restraints"]
+            ]
+        if state.get("_lambda_schedule") is not None:
+            state["_lambda_schedule"] = self._to_hex(state["_lambda_schedule"])
+        return state
+
+    def __setstate__(self, state):
+        """Reverse the hex-encoding performed in __getstate__."""
+        if state.get("_restraints") is not None:
+            state["_restraints"] = [
+                self._from_hex(restraint) for restraint in state["_restraints"]
+            ]
+        if state.get("_lambda_schedule") is not None:
+            state["_lambda_schedule"] = self._from_hex(state["_lambda_schedule"])
+        self.__dict__.update(state)
+
     @classmethod
     def _create_parser(cls):
         """
@@ -2572,6 +2619,65 @@ def _create_parser(cls):
 
         return parser
 
+    @property
+    def restraint_search_time(self):
+        return self._restraint_search_time
+
+    @restraint_search_time.setter
+    def restraint_search_time(self, restraint_search_time):
+        if not isinstance(restraint_search_time, str):
+            raise TypeError("'restraint_search_time' must be of type 'str'")
+
+        from sire.units import picosecond
+
+        try:
+            t = _sr.u(restraint_search_time)
+        except:
+            raise ValueError(
+                f"Unable to parse 'restraint_search_time' as a Sire GeneralUnit: {restraint_search_time}"
+            )
+
+        if not t.has_same_units(picosecond):
+            raise ValueError("'restraint_search_time' units are invalid.")
+
+        self._restraint_search_time = t
+
+    @property
+    def restraint_search_frequency(self):
+        return self._restraint_search_frequency
+
+    @restraint_search_frequency.setter
+    def restraint_search_frequency(self, restraint_search_frequency):
+        if not isinstance(restraint_search_frequency, str):
+            raise TypeError("'restraint_search_frequency' must be of type 'str'")
+
+        from sire.units import picosecond
+
+        try:
+            t = _sr.u(restraint_search_frequency)
+        except:
+            raise ValueError(
+                f"Unable to parse 'restraint_search_frequency' as a Sire GeneralUnit: {restraint_search_frequency}"
+            )
+
+        if not t.has_same_units(picosecond):
+            raise ValueError("'restraint_search_frequency' units are invalid.")
+
+        self._restraint_search_frequency = t
+
+    @property
+    def restraint_search_receptor_selection(self):
+        return self._restraint_search_receptor_selection
+
+    @restraint_search_receptor_selection.setter
+    def restraint_search_receptor_selection(self, restraint_search_receptor_selection):
+        if restraint_search_receptor_selection is not None:
+            if not isinstance(restraint_search_receptor_selection, str):
+                raise TypeError(
+                    "'restraint_search_receptor_selection' must be of type 'str'"
+                )
+        self._restraint_search_receptor_selection = restraint_search_receptor_selection
+
     def _reset_logger(self, logger):
         """
         Internal method to reset the logger.
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index c973400..d5101f2 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -223,6 +223,15 @@ def __init__(self, system, config):
         except:
             self._has_water = False
 
+        # Check for protein (distinguishes ABFE from AHFE). A protein+ligand
+        # system has more than one non-water molecule with at least 3 atoms;
+        # a solvated ligand alone has only one.
+        try:
+            non_water_mols = self._system["(not water) and (atomidx > 1)"].molecules()
+            self._has_protein = non_water_mols.num_molecules() > 1
+        except:
+            self._has_protein = False
+
         # Warn if dispersion correction is requested but can't be applied.
         if self._config.use_dispersion_correction and not self._has_water:
             msg = "Cannot use dispersion correction for vacuum simulations. Disabling!"
@@ -274,7 +283,7 @@ def __init__(self, system, config):
         elif self._config.ghost_modifications:
             from ghostly import modify
 
-            _logger.info("Applying modifications to ghost atom bonded terms")
+            _logger.info("Applying modifications to ghost atom bonded terms.")
             try:
                 self._system, self._modifications = modify(self._system)
             # Angle optimisation can sometimes fail.
@@ -954,6 +963,120 @@ def __init__(self, system, config):
             # Update the maximum number of threads.
             _sr.legacy.Base.set_max_num_threads(sire_threads)
 
+    @property
+    def _is_abfe_bound(self):
+        """
+        Whether this is the bound leg of an ABFE simulation: annihilate/decouple
+        schedule with a solvated protein present. False for the free leg (ligand
+        in solvent, no protein), even though it may use the same lambda schedule.
+        """
+        return (
+            self._config._lambda_schedule_name in ("annihilate", "decouple")
+            and self._has_protein
+            and self._has_water
+        )
+
+    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.
+
+        The input system is assumed to already be equilibrated (SOMD2 does not
+        run a separate equilibration stage for the restraint search); it is
+        minimised using the production minimisation settings, then a short
+        trajectory is run at lambda=0, matching the dynamics settings used for
+        production, to derive the restraint geometry and force constants.
+
+        Parameters
+        ----------
+
+        device : int, optional
+            GPU device number to use for the restraint-search run.
+
+        Returns
+        -------
+
+        restraints : sire.mm.BoreschRestraints
+        """
+        from sire.restraints import boresch_search
+
+        restraint_file = str(self._config.output_directory / "abfe_restraint.s3")
+
+        # On restart, load the restraint saved from the previous run.
+        if _Path(restraint_file).exists():
+            _logger.info(f"Loading existing Boresch restraint from {restraint_file}")
+            return _sr.stream.load(restraint_file)
+
+        _logger.info(
+            "No restraint supplied for ABFE. Running Boresch restraint search."
+        )
+
+        search_system = self._system
+
+        if self._config.minimise:
+            constraint = self._config.constraint
+            perturbable_constraint = self._config.perturbable_constraint
+
+            # Don't use constraints during minimisation.
+            if not self._config.minimisation_constraints:
+                constraint = "none"
+                perturbable_constraint = "none"
+
+            min_dynamics_kwargs = self._dynamics_kwargs.copy()
+            min_dynamics_kwargs.update(
+                {
+                    "device": device,
+                    "lambda_value": 0.0,
+                    "constraint": constraint,
+                    "perturbable_constraint": perturbable_constraint,
+                }
+            )
+
+            min_dynamics = search_system.dynamics(**min_dynamics_kwargs)
+            min_dynamics.minimise(timeout=self._config.timeout)
+            search_system = min_dynamics.commit()
+
+        dynamics_kwargs = self._dynamics_kwargs.copy()
+        dynamics_kwargs.update(
+            {
+                "device": device,
+                "lambda_value": 0.0,
+            }
+        )
+
+        dynamics = search_system.dynamics(**dynamics_kwargs)
+        dynamics.run(
+            self._config.restraint_search_time,
+            energy_frequency=0,
+            frame_frequency=self._config.restraint_search_frequency,
+            save_velocities=False,
+            auto_fix_minimise=self._config.auto_fix_minimise,
+            save_crash_report=self._config.save_crash_report,
+        )
+        search_system = dynamics.commit()
+
+        search_kwargs = {"temperature": self._config.temperature}
+        if self._config.restraint_search_receptor_selection is not None:
+            search_kwargs["receptor_selection"] = (
+                self._config.restraint_search_receptor_selection
+            )
+
+        restraints, correction = boresch_search(search_system, **search_kwargs)
+
+        correction_kcal_mol = float(correction.to(_sr.units.kcal_per_mol))
+        _logger.info(
+            f"Boresch restraint generated. Standard state correction: "
+            f"{correction_kcal_mol:.4f} kcal mol-1"
+        )
+
+        # Save so that restarts can reload without re-generating.
+        _sr.stream.save(restraints, restraint_file)
+
+        return restraints
+
     def _check_space(self):
         """
         Check if the system has a periodic space.
@@ -1496,6 +1619,8 @@ def _compare_configs(config1, config2):
             "log_file",
             "overwrite",
             "timeout",
+            "restraint_search_time",
+            "restraint_search_frequency",
         ]
         for key in config1.keys():
             if key not in allowed_diffs:
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 7da3ba3..56a1b28 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -740,6 +740,20 @@ def __init__(self, system, config):
         else:
             self._num_gpus = min(self._config.max_gpus, len(gpu_devices))
 
+        # Auto-generate a Boresch restraint for ABFE runs with no user-supplied
+        # restraint. This must happen 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
+
         # 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 fb3980f..56dc31e 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -217,6 +217,18 @@ 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
+
         import concurrent.futures as _futures
         import multiprocessing as _mp
 

From 2c572de07595fe4b7a62ff91e84cfdaf69955d6d Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Thu, 2 Jul 2026 13:17:19 +0100
Subject: [PATCH 07/67] Add support for split restraint levers in ABFE
 schedules.

---
 src/somd2/_utils/_schedules.py |  97 +++++++++++++++++++--
 tests/schedules/test_abfe.py   | 153 +++++++++++++++++++++++++++++++++
 2 files changed, 244 insertions(+), 6 deletions(-)
 create mode 100644 tests/schedules/test_abfe.py

diff --git a/src/somd2/_utils/_schedules.py b/src/somd2/_utils/_schedules.py
index 4c49a0e..f40dd52 100644
--- a/src/somd2/_utils/_schedules.py
+++ b/src/somd2/_utils/_schedules.py
@@ -27,7 +27,44 @@
 ]
 
 
-def annihilate(fix_epsilon=True):
+def _set_boresch_lever_equations(s, stage_dihedral, stage_distance_angle):
+    """
+    Set the equations for a "split" restraint_lever Boresch restraint (see
+    sire.restraints.boresch's restraint_lever parameter), reproducing the
+    RXRX protocol's staged restraint turn-on (Table S1 of the RXRX paper's
+    SI): within each of the two named stages, the corresponding restraint
+    group ramps from ~0 to 1 following a geometric progression, while the
+    other group is held fixed. 'stage_dihedral' is the stage over which
+    the dihedral restraint group ramps on (with the distance/angle group
+    held at 0); 'stage_distance_angle' is the stage over which the
+    distance/angle group ramps on (with the dihedral group held at 1,
+    already fully on).
+
+    Note: this aligns the restraint turn-on with SOMD2's own decharge/
+    annihilate(or decouple) stage boundaries, rather than reproducing the
+    RXRX paper's exact global 50/50 window split (which falls partway
+    through the annihilate/decouple stage, since the paper's own decharge
+    stage is only 21 of 64 total bound-leg windows) - the relative sizes of
+    SOMD2's stages are controlled by lambda_values weighting, not fixed.
+    """
+    from sire.legacy.CAS import Exp as _Exp
+    import math as _math
+
+    # Geometric progression from ~0.01 (fully off) to 1.0 (fully on), matching
+    # the ratio observed in the RXRX paper's published lambda schedule.
+    ramp_on = _Exp((1 - s.lam()) * _math.log(0.01))
+
+    s.set_equation(stage=stage_dihedral, lever="restraint_dihedral", equation=ramp_on)
+    s.set_equation(stage=stage_dihedral, lever="restraint_distance_angle", equation=0)
+    s.set_equation(stage=stage_distance_angle, lever="restraint_dihedral", equation=1)
+    s.set_equation(
+        stage=stage_distance_angle,
+        lever="restraint_distance_angle",
+        equation=ramp_on,
+    )
+
+
+def annihilate(fix_epsilon=True, restraint_lever="split"):
     """
     Build the ABFE lambda schedule using decharge → annihilate.
 
@@ -44,12 +81,28 @@ def annihilate(fix_epsilon=True):
         If False, epsilon is scaled normally from initial to final and the LRC
         follows naturally.
 
+    restraint_lever : str, optional
+        How the Boresch restraint is controlled by this schedule, matching
+        sire.restraints.boresch's restraint_lever parameter. Either "split"
+        (default), where the dihedral restraint terms are turned on during
+        decharge and the distance/angle terms are turned on during
+        annihilate, reproducing the RXRX protocol's staged restraint
+        turn-on, or "combined", where the whole restraint is turned on
+        together during the decharge stage. The Boresch restraint object
+        passed to the simulation must have a matching restraint_lever value.
+
     Returns
     -------
 
     schedule : sire.legacy.CAS.LambdaSchedule
         The lambda schedule.
     """
+    if restraint_lever not in ("combined", "split"):
+        raise ValueError(
+            "'restraint_lever' must be either 'combined' or 'split', "
+            f"got {restraint_lever!r}"
+        )
+
     from sire.cas import LambdaSchedule as _LambdaSchedule
 
     # Start with the standard decouple schedule and modify the stages and
@@ -65,14 +118,22 @@ def annihilate(fix_epsilon=True):
         lever="charge",
         equation=s.lam() * s.final() + s.initial() * (1 - s.lam()),
     )
-    s.set_equation(stage="decharge", force="restraint", equation=s.lam() * s.final())
 
     s.add_stage(
         "annihilate",
         equation=(-s.lam() + 1) * s.initial() + s.lam() * s.final(),
     )
     s.set_equation(stage="annihilate", lever="charge", equation=s.final())
-    s.set_equation(stage="annihilate", force="restraint", equation=s.final())
+
+    if restraint_lever == "split":
+        _set_boresch_lever_equations(
+            s, stage_dihedral="decharge", stage_distance_angle="annihilate"
+        )
+    else:
+        s.set_equation(
+            stage="decharge", lever="restraint", equation=s.lam() * s.final()
+        )
+        s.set_equation(stage="annihilate", lever="restraint", equation=s.final())
 
     if fix_epsilon:
         s.set_equation(stage="annihilate", lever="epsilon", equation=s.initial())
@@ -86,7 +147,7 @@ def annihilate(fix_epsilon=True):
     return s
 
 
-def decouple(fix_epsilon=True):
+def decouple(fix_epsilon=True, restraint_lever="split"):
     """
     Build the ABFE lambda schedule using decharge → decouple.
 
@@ -101,12 +162,28 @@ def decouple(fix_epsilon=True):
         ghost-LRC force is then explicitly scaled to zero over the stage.
         If False, epsilon is scaled normally and the LRC follows naturally.
 
+    restraint_lever : str, optional
+        How the Boresch restraint is controlled by this schedule, matching
+        sire.restraints.boresch's restraint_lever parameter. Either "split"
+        (default), where the dihedral restraint terms are turned on during
+        decharge and the distance/angle terms are turned on during decouple,
+        reproducing the RXRX protocol's staged restraint turn-on, or
+        "combined", where the whole restraint is turned on together during
+        the decharge stage. The Boresch restraint object passed to the
+        simulation must have a matching restraint_lever value.
+
     Returns
     -------
 
     schedule : sire.legacy.CAS.LambdaSchedule
         The lambda schedule.
     """
+    if restraint_lever not in ("combined", "split"):
+        raise ValueError(
+            "'restraint_lever' must be either 'combined' or 'split', "
+            f"got {restraint_lever!r}"
+        )
+
     from sire.cas import LambdaSchedule as _LambdaSchedule
 
     # Start with the standard decouple schedule and modify the stages and
@@ -114,7 +191,6 @@ def decouple(fix_epsilon=True):
     # we will use this approach for prototyping.
     s = _LambdaSchedule.standard_decouple()
 
-    s.set_equation(stage="decouple", lever="restraint", equation=s.final())
     s.set_equation(stage="decouple", lever="kappa", force="ghost/ghost", equation=0)
     s.set_equation(stage="decouple", lever="kappa", force="ghost-14", equation=0)
     s.set_equation(stage="decouple", lever="charge", equation=s.final())
@@ -142,7 +218,16 @@ def decouple(fix_epsilon=True):
     s.set_equation(
         stage="decharge", lever="kappa", force="ghost-14", equation=-s.lam() + 1
     )
-    s.set_equation(stage="decharge", lever="restraint", equation=s.initial() * s.lam())
+
+    if restraint_lever == "split":
+        _set_boresch_lever_equations(
+            s, stage_dihedral="decharge", stage_distance_angle="decouple"
+        )
+    else:
+        s.set_equation(stage="decouple", lever="restraint", equation=s.final())
+        s.set_equation(
+            stage="decharge", lever="restraint", equation=s.initial() * s.lam()
+        )
 
     return s
 
diff --git a/tests/schedules/test_abfe.py b/tests/schedules/test_abfe.py
new file mode 100644
index 0000000..e44a25a
--- /dev/null
+++ b/tests/schedules/test_abfe.py
@@ -0,0 +1,153 @@
+import pytest
+
+from somd2._utils._schedules import 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
+# identical expected lever values.
+BUILDERS = [annihilate, decouple]
+
+_LAMBDA_VALUES = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
+
+# Expected "restraint" lever values (restraint_lever="combined"): ramps
+# linearly 0 -> 1 across the whole first stage (decharge), then held at 1.
+_COMBINED_RESTRAINT = [0.0, 0.4, 0.8, 1.0, 1.0, 1.0]
+
+# Expected "restraint_dihedral"/"restraint_distance_angle" lever values
+# (restraint_lever="split"): a geometric progression from ~0.01 to 1.0 across
+# each stage in turn, reproducing the ratio in Table S1 of the RXRX paper's
+# SI. dihedral ramps during the first stage then holds at 1; distance_angle
+# holds at 0 during the first stage then ramps during the second.
+_SPLIT_DIHEDRAL = [0.01, 0.06309573444801934, 0.39810717055349737, 1.0, 1.0, 1.0]
+_SPLIT_DISTANCE_ANGLE = [0.0, 0.0, 0.0, 0.025118864315095805, 0.15848931924611143, 1.0]
+
+
+def _morph(schedule, lever, lambda_value):
+    """
+    Query a lever's value exactly as SireOpenMM's LambdaLever does at
+    runtime: lambda_schedule.morph("*", restraint_name, 1.0, 1.0, lambda_value).
+    """
+    return schedule.morph("*", lever, 1.0, 1.0, lambda_value)
+
+
+@pytest.mark.parametrize("builder", BUILDERS)
+def test_restraint_lever_defaults_to_split(builder):
+    """
+    annihilate()/decouple() default to restraint_lever="split", matching
+    boresch_search()'s own default of restraint_lever="split" for its
+    default protocol="rxrx"
+    """
+    schedule = builder()
+    assert "restraint_dihedral" in schedule.get_levers()
+    assert "restraint_distance_angle" in schedule.get_levers()
+    assert "restraint" not in schedule.get_levers()
+
+
+@pytest.mark.parametrize("builder", BUILDERS)
+def test_restraint_lever_invalid_raises(builder):
+    with pytest.raises(ValueError, match="restraint_lever"):
+        builder(restraint_lever="not_a_real_lever")
+
+
+@pytest.mark.parametrize("builder", BUILDERS)
+def test_restraint_lever_combined(builder):
+    """
+    restraint_lever="combined" sets a single "restraint" lever that ramps
+    0 -> 1 across the first stage (decharge), then holds at 1.
+    """
+    schedule = builder(restraint_lever="combined")
+    assert "restraint" in schedule.get_levers()
+    assert "restraint_dihedral" not in schedule.get_levers()
+    assert "restraint_distance_angle" not in schedule.get_levers()
+
+    values = [_morph(schedule, "restraint", lv) for lv in _LAMBDA_VALUES]
+    assert values == pytest.approx(_COMBINED_RESTRAINT, abs=1e-6)
+
+
+@pytest.mark.parametrize("builder", BUILDERS)
+def test_restraint_lever_split(builder):
+    """
+    restraint_lever="split" sets two independent levers ("restraint_dihedral"
+    and "restraint_distance_angle"), each following a geometric progression
+    across its own stage while the other is held fixed, reproducing the RXRX
+    protocol's staged restraint turn-on.
+    """
+    schedule = builder(restraint_lever="split")
+
+    dihedral_values = [
+        _morph(schedule, "restraint_dihedral", lv) for lv in _LAMBDA_VALUES
+    ]
+    distance_angle_values = [
+        _morph(schedule, "restraint_distance_angle", lv) for lv in _LAMBDA_VALUES
+    ]
+
+    assert dihedral_values == pytest.approx(_SPLIT_DIHEDRAL, abs=1e-6)
+    assert distance_angle_values == pytest.approx(_SPLIT_DISTANCE_ANGLE, abs=1e-6)
+
+
+@pytest.mark.skipif(
+    "openmm" not in __import__("sire").convert.supported_formats(),
+    reason="openmm support is not available",
+)
+def test_restraint_lever_split_openmm_system():
+    """
+    End-to-end regression test: build a real dynamics object with a "split"
+    Boresch restraint and confirm the two independently-lambda-addressable
+    OpenMM Forces (distance/angle and dihedral) exist, with 'rho' values
+    that follow the expected geometric progression as lambda changes,
+    matching what test_restraint_lever_split checks at the schedule level.
+    """
+    import xml.etree.ElementTree as ET
+
+    import sire as sr
+
+    mols = sr.load_test_files("boresch_restraints.prm7", "boresch_restraints.dcd")
+    mols.update(sr.morph.decouple(mols.molecule(1), as_new_molecule=False))
+
+    restraints = sr.restraints.boresch(
+        mols,
+        receptor=[692, 702, 704],
+        ligand=[1496, 1498, 1499],
+        kr="1 kcal mol-1 A-2",
+        ktheta=["80 kcal mol-1 rad-2"] * 2,
+        kphi=["80 kcal mol-1 rad-2"] * 3,
+        r0="4.56908 A",
+        theta0=["82.5581 degrees", "94.9595 degrees"],
+        phi0=["27.9429 degrees", "125.68 degrees", "-107.008 degrees"],
+        angle_potential="restricted_bending",
+        restraint_lever="split",
+    )
+
+    schedule = decouple(restraint_lever="split")
+
+    d = mols.dynamics(
+        timestep="2fs",
+        temperature="298 K",
+        schedule=schedule,
+        lambda_value=0.0,
+        map={"restraints": restraints},
+    )
+
+    expected_rho = {
+        0.0: {"distance_angle": 0.0, "dihedral": 0.01},
+        0.5: {"distance_angle": 0.01, "dihedral": 1.0},
+        1.0: {"distance_angle": 1.0, "dihedral": 1.0},
+    }
+
+    for lam, expected in expected_rho.items():
+        d.set_lambda(lam)
+        root = ET.fromstring(d.to_xml())
+
+        rhos = {}
+        for force in root.iter("Force"):
+            if force.get("name") == "BoreschRestraintForce":
+                kind = (
+                    "distance_angle" if "e_bond" in force.get("energy") else "dihedral"
+                )
+                rhos[kind] = float(force.find("Bonds")[0].get("param1"))
+
+        assert set(rhos.keys()) == {"distance_angle", "dihedral"}
+        assert rhos["distance_angle"] == pytest.approx(
+            expected["distance_angle"], abs=1e-6
+        )
+        assert rhos["dihedral"] == pytest.approx(expected["dihedral"], abs=1e-6)

From a496234564444a25c4d2b3de5b6cd5380f1eda01 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 3 Jul 2026 09:04:22 +0100
Subject: [PATCH 08/67] Store auto-generated standard state correction in
 parquet metadata.

---
 src/somd2/runner/_base.py | 32 +++++++++++++++++++++++++++++---
 1 file changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index d5101f2..fa3932f 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -232,6 +232,11 @@ def __init__(self, system, config):
         except:
             self._has_protein = False
 
+        # Set by _generate_boresch_restraint() when a Boresch restraint is
+        # auto-generated, then written into the energy trajectory parquet
+        # metadata (see _checkpoint).
+        self._standard_state_correction = None
+
         # Warn if dispersion correction is requested but can't be applied.
         if self._config.use_dispersion_correction and not self._has_water:
             msg = "Cannot use dispersion correction for vacuum simulations. Disabling!"
@@ -1008,7 +1013,18 @@ def _generate_boresch_restraint(self, device=None):
         # On restart, load the restraint saved from the previous run.
         if _Path(restraint_file).exists():
             _logger.info(f"Loading existing Boresch restraint from {restraint_file}")
-            return _sr.stream.load(restraint_file)
+            restraints = _sr.stream.load(restraint_file)
+
+            from sire.restraints import get_standard_state_correction
+
+            correction = get_standard_state_correction(
+                restraints[0], temperature=self._config.temperature
+            )
+            self._standard_state_correction = float(
+                correction.to(_sr.units.kcal_per_mol)
+            )
+
+            return restraints
 
         _logger.info(
             "No restraint supplied for ABFE. Running Boresch restraint search."
@@ -1066,10 +1082,13 @@ def _generate_boresch_restraint(self, device=None):
 
         restraints, correction = boresch_search(search_system, **search_kwargs)
 
-        correction_kcal_mol = float(correction.to(_sr.units.kcal_per_mol))
+        # Cache so it can be written into the energy trajectory parquet
+        # metadata (see _checkpoint), letting analysis code automatically
+        # apply the correction without needing to scan the logs.
+        self._standard_state_correction = float(correction.to(_sr.units.kcal_per_mol))
         _logger.info(
             f"Boresch restraint generated. Standard state correction: "
-            f"{correction_kcal_mol:.4f} kcal mol-1"
+            f"{self._standard_state_correction:.4f} kcal mol-1"
         )
 
         # Save so that restarts can reload without re-generating.
@@ -2070,6 +2089,13 @@ def _checkpoint(
                 if lambda_grad is not None:
                     metadata["lambda_grad"] = [f"{v:.5f}" for v in lambda_grad]
 
+                # Add the standard state correction, if a Boresch restraint
+                # was auto-generated for this ABFE run.
+                if self._standard_state_correction is not None:
+                    metadata["standard_state_correction"] = (
+                        f"{self._standard_state_correction:.6f}"
+                    )
+
             if is_final_block:
                 # Save the end-state GCMC topologies for trajectory analysis and visualisation.
                 # This topology contains additional water molecules that are used for GCMC

From d276232fbe9a26f8e805d1d3f2fa20d032fdd42d Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 3 Jul 2026 09:16:20 +0100
Subject: [PATCH 09/67] Harmonise dependency version formatting.

---
 src/somd2/__init__.py     | 14 ++++++++++++--
 src/somd2/runner/_base.py |  3 +--
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/somd2/__init__.py b/src/somd2/__init__.py
index b3070df..7f0a24a 100644
--- a/src/somd2/__init__.py
+++ b/src/somd2/__init__.py
@@ -34,10 +34,20 @@
 # Store the somd2 version.
 from ._version import __version__
 
-# Store the sire version.
+# Store the sire version. Unlike somd2/BioSimSpace/ghostly/loch (which use
+# versioningit and only append a "+g<revisionid>" local version segment
+# for non-release (".dev") builds, omitting it entirely for a clean tagged
+# release), sire exposes its version and revision id as separate attributes,
+# with __revisionid__ always set regardless of release status. Build a
+# composite string using the same "+g<revisionid>" convention as the other
+# packages, only appending it for non-release (".dev") builds, so that all
+# five version strings are formatted consistently.
 from sire import __version__ as _sire_version
 from sire import __revisionid__ as _sire_revisionid
 
+if ".dev" in _sire_version:
+    _sire_version = f"{_sire_version}+g{_sire_revisionid}"
+
 # Store the BioSimSpace version.
 from BioSimSpace import __version__ as _biosimspace_version
 
@@ -61,7 +71,7 @@ def get_versions():
     """
     return {
         "somd2": __version__,
-        "sire": f"{_sire_version}+{_sire_revisionid}",
+        "sire": _sire_version,
         "biosimspace": _biosimspace_version,
         "ghostly": _ghostly_version,
         "loch": _loch_version,
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index fa3932f..85f6e82 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -2078,8 +2078,7 @@ def _checkpoint(
             if not is_post_equilibration:
                 metadata = {
                     "attrs": df.attrs,
-                    "somd2 version": versions["somd2"],
-                    "sire version": versions["sire"],
+                    "versions": versions,
                     "lambda": f"{lam:.5f}",
                     "speed": speed,
                     "temperature": str(self._config.temperature.value()),

From 3ba282adb0d8abc4781c1b673b21ca3a41f89eb8 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 3 Jul 2026 16:58:07 +0100
Subject: [PATCH 10/67] Give alchemical ions their own plain morph lambda
 schedule.

---
 CHANGELOG.md                         |  1 +
 src/somd2/runner/_base.py            | 67 +++++++++++++++++++---------
 tests/runner/test_alchemical_ions.py | 67 +++++++++++++++++++++++++++-
 3 files changed, 113 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0190663..63c9ad1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,7 @@ Changelog
 
 * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
 * Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166).
+* Give alchemical ions their own plain morph lambda schedule so they interpolate correctly under non-standard lambda schedules [#169](https://github.com/OpenBioSim/somd2/pull/169).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 85f6e82..7a116fd 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -374,11 +374,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.
@@ -401,24 +404,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
 
@@ -428,6 +417,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
@@ -1203,6 +1206,15 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
 
         system: :class: `System <sire.system.System>`
             The perturbed system with alchemical ions added.
+
+        restraints: :class: `Restraints <sire.restraints.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 <sire.cas.LambdaSchedule>`.
         """
 
         from sire.legacy.IO import createChlorineIon as _createChlorineIon
@@ -1268,6 +1280,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
@@ -1383,6 +1398,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())
 
@@ -1398,7 +1416,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 6024d1a..f4a4bfe 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()

From 5eb4265cb98a0cc95bcb3b33801cafd1ef82e92b Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 3 Jul 2026 22:07:40 +0100
Subject: [PATCH 11/67] Persist alchemical ion identity across restarts,
 independent of GCMC.

---
 CHANGELOG.md                         |   1 +
 src/somd2/runner/_base.py            | 104 ++++++++++++++++++++++++---
 tests/runner/test_alchemical_ions.py |  96 ++++++++++++++++++++++++-
 3 files changed, 190 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 63c9ad1..41970ae 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ Changelog
 * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
 * Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166).
 * Give alchemical ions their own plain morph lambda schedule so they interpolate correctly under non-standard lambda schedules [#169](https://github.com/OpenBioSim/somd2/pull/169).
+* Persist alchemical ion identity across restarts so the same molecule is reused regardless of GCMC state [#172](https://github.com/OpenBioSim/somd2/pull/172).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 7a116fd..ee08598 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -376,14 +376,27 @@ def __init__(self, system, config):
         # Create alchemical ions.
         ion_indices = []
         if charge_diff != 0:
-            self._system, coalchemical_restraints, ion_indices = (
+            # On restart, reuse the exact molecule(s) chosen as alchemical ions
+            # in the original run, rather than re-running the "furthest waters"
+            # search. This makes ion selection independent of GCMC state (or
+            # anything else that might change between runs), since the search
+            # is otherwise only reproducible by assumption, not by construction.
+            mol_indices = None
+            if self._config.restart:
+                mol_indices = self._load_alchemical_ion_indices()
+
+            self._system, coalchemical_restraints, ion_indices, ion_mol_indices = (
                 self._create_alchemical_ions(
                     self._system,
                     charge_diff,
                     restraint_distance=self._config.coalchemical_restraint_dist,
+                    mol_indices=mol_indices,
                 )
             )
 
+            # Keep the stored indices in sync for any future restart.
+            self._save_alchemical_ion_indices(ion_mol_indices)
+
             # Add the coalchemical restraints to the extra args.
             if coalchemical_restraints is not None:
                 self._config._extra_args["coalchemical_restraints"] = (
@@ -1187,8 +1200,49 @@ def _get_charge_difference(system):
 
         return perturbed - reference
 
+    def _save_alchemical_ion_indices(self, mol_indices):
+        """
+        Persist the absolute molecule index of each alchemical ion to a small,
+        dedicated file in the output directory, independent of the per-window
+        (regular runner) or shared (repex) checkpoint formats. This allows a
+        restart to reuse the exact same ion(s) chosen in the original run.
+
+        Parameters
+        ----------
+
+        mol_indices: [int]
+            The absolute molecule index of each alchemical ion.
+        """
+        import numpy as _np
+
+        path = self._config.output_directory / "alchemical_ions.npz"
+        _np.savez(path, mol_indices=_np.array(mol_indices, dtype=int))
+
+    def _load_alchemical_ion_indices(self):
+        """
+        Load the absolute molecule index of each alchemical ion previously
+        stored by `_save_alchemical_ion_indices`, if present.
+
+        Returns
+        -------
+
+        mol_indices: [int], None
+            The absolute molecule index of each alchemical ion, or None if no
+            stored indices are available (e.g. a fresh run, or a restart from
+            an output directory that predates this feature).
+        """
+        import numpy as _np
+
+        path = self._config.output_directory / "alchemical_ions.npz"
+        try:
+            return _np.load(path)["mol_indices"].tolist()
+        except Exception:
+            return None
+
     @staticmethod
-    def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
+    def _create_alchemical_ions(
+        system, charge_diff, restraint_distance=None, mol_indices=None
+    ):
         """
         Internal function to create alchemical ions to maintain a constant charge.
 
@@ -1201,6 +1255,14 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
         charge_diff: int
             The charge difference between perturbed and reference states.
 
+        mol_indices: [int]
+            The absolute molecule index (position in `system.molecules()`) of
+            each water to convert into an alchemical ion. If provided, these
+            molecules are converted directly, bypassing the "furthest waters"
+            search. Used on restart to reproduce the exact same ion(s) chosen
+            in the original run, independent of any GCMC state or changes to
+            the search heuristic. Must have the same length as `abs(charge_diff)`.
+
         Returns
         -------
 
@@ -1215,6 +1277,11 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
             The perturbable-molecule index of each alchemical ion that was
             added, suitable for use with
             `LambdaSchedule.set_molecule_schedule <sire.cas.LambdaSchedule>`.
+
+        ion_mol_indices: [int]
+            The absolute molecule index (position in `system.molecules()`,
+            prior to any conversion) of each alchemical ion that was added.
+            Suitable for passing back in as `mol_indices` on a restart.
         """
 
         from sire.legacy.IO import createChlorineIon as _createChlorineIon
@@ -1258,12 +1325,28 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
                 f"{len(system['water'].molecules())} available."
             )
 
-        # Reference coordinates.
-        coords = system.molecules("property is_perturbable").coordinates()
-        coord_string = f"{coords[0].value()}, {coords[1].value()}, {coords[2].value()}"
+        if mol_indices is not None:
+            if len(mol_indices) != num_waters:
+                raise ValueError(
+                    f"Number of stored alchemical-ion molecule indices "
+                    f"({len(mol_indices)}) does not match the current charge "
+                    f"difference ({num_waters} waters required)."
+                )
+
+            # Reuse the exact molecules chosen in the original run.
+            all_mols = system.molecules()
+            waters = [all_mols[idx] for idx in mol_indices]
+        else:
+            # Reference coordinates.
+            coords = system.molecules("property is_perturbable").coordinates()
+            coord_string = (
+                f"{coords[0].value()}, {coords[1].value()}, {coords[2].value()}"
+            )
 
-        # Find the furthest N waters from the perturbable molecule.
-        waters = system[f"furthest {num_waters} waters from {coord_string}"].molecules()
+            # Find the furthest N waters from the perturbable molecule.
+            waters = system[
+                f"furthest {num_waters} waters from {coord_string}"
+            ].molecules()
 
         # Determine the water model.
         if waters[0].num_atoms() == 3:
@@ -1283,6 +1366,10 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
         # Store the molecule numbers of the alchemical ions.
         ion_numbers = []
 
+        # Store the absolute molecule index of each alchemical ion (prior to
+        # conversion), for persisting across restarts.
+        ion_mol_indices = []
+
         # Create the ions.
         for water in waters:
             # Flag to indicate whether we need to reverse the alchemical ion
@@ -1403,6 +1490,7 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
 
             # Get the index of the perturbed water.
             index = numbers.index(water.number())
+            ion_mol_indices.append(index)
 
             # Log that we are adding an alchemical ion.
             if is_reverse:
@@ -1425,7 +1513,7 @@ def _create_alchemical_ions(system, charge_diff, restraint_distance=None):
         perturbable_mols = system.molecules()["perturbable"].molecules()
         ion_indices = [perturbable_mols.find(system[number]) for number in ion_numbers]
 
-        return system, restraints, ion_indices
+        return system, restraints, ion_indices, ion_mol_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 f4a4bfe..2c5d7fe 100644
--- a/tests/runner/test_alchemical_ions.py
+++ b/tests/runner/test_alchemical_ions.py
@@ -1,7 +1,8 @@
 import math
+import pytest
 import tempfile
 
-import pytest
+from pathlib import Path
 
 from somd2.config import Config
 from somd2.runner import Runner
@@ -15,20 +16,109 @@ def test_alchemical_ions(mols, request):
     mols = request.getfixturevalue(mols).clone()
 
     # Add 10 Cl- ions.
-    new_mols, _, ion_indices = Runner._create_alchemical_ions(mols, 10)
+    new_mols, _, ion_indices, ion_mol_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
+    assert len(ion_mol_indices) == 10
 
     # Add 10 Na+ ions.
-    new_mols, _, ion_indices = Runner._create_alchemical_ions(mols, -10)
+    new_mols, _, ion_indices, ion_mol_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
+    assert len(ion_mol_indices) == 10
+
+
+@pytest.mark.parametrize("mols", ["ethane_methanol", "ethane_methanol_ions"])
+def test_alchemical_ion_mol_indices_reproducible(mols, request):
+    """
+    Ensure that passing the molecule indices returned by a previous call to
+    `_create_alchemical_ions` reproduces the exact same ion(s), bypassing the
+    "furthest waters" search entirely. This is what a restart relies on.
+    """
+    mols = request.getfixturevalue(mols).clone()
+
+    # Pick ions via the heuristic search, recording which molecule(s) were
+    # converted.
+    heuristic_mols, _, _, ion_mol_indices = Runner._create_alchemical_ions(mols, 3)
+    heuristic_ion_numbers = {
+        mol.number()
+        for mol in heuristic_mols.molecules()["perturbable"].molecules()
+        if mol.has_property("is_alchemical_ion")
+    }
+
+    # Reuse the stored indices directly - should convert the exact same
+    # molecules, without running the search.
+    replayed_mols, _, _, replayed_mol_indices = Runner._create_alchemical_ions(
+        mols, 3, mol_indices=ion_mol_indices
+    )
+    replayed_ion_numbers = {
+        mol.number()
+        for mol in replayed_mols.molecules()["perturbable"].molecules()
+        if mol.has_property("is_alchemical_ion")
+    }
+
+    assert replayed_ion_numbers == heuristic_ion_numbers
+    assert replayed_mol_indices == ion_mol_indices
+    assert math.isclose(
+        Runner._get_charge_difference(replayed_mols), -3.0, rel_tol=1e-6
+    )
+
+
+def test_alchemical_ion_mol_indices_mismatch_raises(ethane_methanol):
+    """A stored index count that doesn't match the charge difference should
+    raise a clear error, rather than silently converting the wrong number of
+    waters."""
+    mols = ethane_methanol.clone()
+
+    with pytest.raises(ValueError, match="does not match the current charge"):
+        Runner._create_alchemical_ions(mols, 3, mol_indices=[0, 1])
+
+
+def test_alchemical_ion_restart_reuses_same_ion(ethane_methanol_ions):
+    """
+    Ensure that restarting a run picks the exact same alchemical ion as the
+    original run, via the persisted `alchemical_ions.npz` file, rather than
+    re-running the "furthest waters" search from scratch.
+    """
+    mols = ethane_methanol_ions.clone()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_config = dict(
+            output_directory=tmpdir,
+            platform="cpu",
+            charge_difference=1,
+        )
+
+        # Fresh run: picks an ion via the heuristic search and persists its
+        # molecule index to alchemical_ions.npz.
+        runner1 = Runner(mols.clone(), Config(restart=False, **base_config))
+        ion_number_1 = next(
+            mol.number()
+            for mol in runner1._system.molecules()["perturbable"].molecules()
+            if mol.has_property("is_alchemical_ion")
+        )
+
+        assert (Path(tmpdir) / "alchemical_ions.npz").exists()
+
+        # "Restart": construct a new Runner against the same input and output
+        # directory. It should reuse the stored ion index rather than
+        # re-running the search.
+        runner2 = Runner(mols.clone(), Config(restart=True, **base_config))
+        ion_number_2 = next(
+            mol.number()
+            for mol in runner2._system.molecules()["perturbable"].molecules()
+            if mol.has_property("is_alchemical_ion")
+        )
+
+        assert ion_number_1 == ion_number_2
 
 
 @pytest.mark.parametrize("schedule_name", ["decouple", "annihilate"])

From 196237b1b46bcb6be157c87c228c511e2493ad7a Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sat, 4 Jul 2026 16:27:10 +0100
Subject: [PATCH 12/67] Seed simulation from least strained restraint search
 structure.

---
 src/somd2/runner/_base.py | 53 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 49 insertions(+), 4 deletions(-)

diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index ee08598..def58bb 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -1021,13 +1021,45 @@ def _generate_boresch_restraint(self, device=None):
         -------
 
         restraints : sire.mm.BoreschRestraints
+
+        Notes
+        -----
+
+        As a side effect, when the restraint is generated (rather than reused
+        from a checkpointed restart), ``self._system`` is re-seeded from the
+        least-strained trajectory frame returned by ``boresch_search``. The
+        restraint equilibrium values are trajectory averages, so the input
+        structure is generally not consistent with them; starting production
+        from it instead would leave the restraint badly strained at t=0 and can
+        blow the simulation up as the restraint is switched on.
+
+        The cached restraint is only reused on a genuine (checkpointed) restart,
+        where it must be kept identical to the one the accumulated free energy
+        was computed with. When there is no checkpoint (a fresh run, or a crash
+        before any progress was checkpointed) the restraint is regenerated. This
+        avoids getting pinned to a restraint that is itself the cause of the
+        crash: reusing it and re-seeding the same frame would just reproduce the
+        crash on every restart, whereas a fresh search may pick a different frame
+        or anchor, and re-seeds ``self._system`` naturally.
         """
         from sire.restraints import boresch_search
 
         restraint_file = str(self._config.output_directory / "abfe_restraint.s3")
 
-        # On restart, load the restraint saved from the previous run.
-        if _Path(restraint_file).exists():
+        # Only reuse a saved restraint when continuing an actual checkpointed
+        # restart: the restraint must match the one the accumulated free energy
+        # was computed with, and the coordinates come from the checkpoint (so no
+        # re-seeding is needed). Deliberately do NOT reuse it when there is no
+        # checkpoint - see the Notes above.
+        if self._is_restart:
+            if not _Path(restraint_file).exists():
+                raise FileNotFoundError(
+                    "Restarting an ABFE bound-leg simulation, but no saved "
+                    f"restraint was found at {restraint_file}. The restraint "
+                    "cannot be regenerated mid-simulation without invalidating "
+                    "the accumulated free energy."
+                )
+
             _logger.info(f"Loading existing Boresch restraint from {restraint_file}")
             restraints = _sr.stream.load(restraint_file)
 
@@ -1096,7 +1128,9 @@ def _generate_boresch_restraint(self, device=None):
                 self._config.restraint_search_receptor_selection
             )
 
-        restraints, correction = boresch_search(search_system, **search_kwargs)
+        restraints, correction, starting_structure = boresch_search(
+            search_system, **search_kwargs
+        )
 
         # Cache so it can be written into the energy trajectory parquet
         # metadata (see _checkpoint), letting analysis code automatically
@@ -1107,7 +1141,18 @@ def _generate_boresch_restraint(self, device=None):
             f"{self._standard_state_correction:.4f} kcal mol-1"
         )
 
-        # Save so that restarts can reload without re-generating.
+        # Re-seed production from the least-strained frame so the restraint is
+        # essentially relaxed at t=0 (see the docstring Notes). The frame comes
+        # from the perturbable search system, so link its properties back to the
+        # reference (lambda=0) end state, matching how the seed systems are
+        # handled elsewhere (see _perturbed_system), and drop the search
+        # trajectory frames so only the single starting snapshot is retained.
+        starting_structure = _sr.morph.link_to_reference(starting_structure)
+        starting_structure.delete_all_frames()
+        self._system = starting_structure
+
+        # Save so that a genuine (checkpointed) restart can reuse the exact same
+        # restraint without re-running the search.
         _sr.stream.save(restraints, restraint_file)
 
         return restraints

From bcc17419079e0e69b0ddf5026d8ffbd975b1c667 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 21 Jul 2026 09:41:34 +0100
Subject: [PATCH 13/67] Use persistent thread pools for the production block.

---
 CHANGELOG.md               |   1 +
 src/somd2/runner/_repex.py | 163 +++++++++++++++++++------------------
 2 files changed, 87 insertions(+), 77 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 41970ae..3005f3a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ Changelog
 * Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166).
 * Give alchemical ions their own plain morph lambda schedule so they interpolate correctly under non-standard lambda schedules [#169](https://github.com/OpenBioSim/somd2/pull/169).
 * Persist alchemical ion identity across restarts so the same molecule is reused regardless of GCMC state [#172](https://github.com/OpenBioSim/somd2/pull/172).
+* Use perisistent `ThreadPoolExector` objects within the main replica exchange dynamics block [#175](https://github.com/OpenBioSim/somd2/pull/175).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 56a1b28..ea73ea3 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1084,6 +1084,10 @@ def run(self):
         num_batches = ceil(self._config.num_lambda / num_workers)
         num_checkpoint_batches = ceil(self._config.num_lambda / num_checkpoint_workers)
 
+        # Persistent thread pools, reused across every batch and cycle.
+        dynamics_executor = ThreadPoolExecutor(max_workers=num_workers)
+        checkpoint_executor = ThreadPoolExecutor(max_workers=num_checkpoint_workers)
+
         # Create the replica list.
         replica_list = list(range(self._config.num_lambda))
 
@@ -1134,29 +1138,27 @@ def run(self):
                     replicas = replica_list[
                         j * num_checkpoint_workers : (j + 1) * num_checkpoint_workers
                     ]
-                    with ThreadPoolExecutor(
-                        max_workers=num_checkpoint_workers
-                    ) as executor:
-                        try:
-                            for index, error in executor.map(
-                                self._checkpoint,
-                                replicas,
-                                repeat(self._lambda_values),
-                                repeat(-1),
-                                repeat(cycles),
-                            ):
-                                if error is not None:
-                                    msg = (
-                                        f"Post-equilibration checkpoint failed for {_lam_sym} = "
-                                        f"{self._lambda_values[index]:.5f}:\n{error}"
-                                    )
-                                    _logger.error(msg)
-                                    raise error
-                        except KeyboardInterrupt:
-                            _logger.error(
-                                "Post-equilibration checkpoint cancelled. Exiting."
-                            )
-                            _sys.exit(1)
+                    try:
+                        for index, error in checkpoint_executor.map(
+                            self._checkpoint,
+                            replicas,
+                            repeat(self._lambda_values),
+                            repeat(-1),
+                            repeat(cycles),
+                        ):
+                            if error is not None:
+                                msg = (
+                                    f"Post-equilibration checkpoint failed for {_lam_sym} = "
+                                    f"{self._lambda_values[index]:.5f}:\n{error}"
+                                )
+                                _logger.error(msg)
+                                raise error
+                    except KeyboardInterrupt:
+                        checkpoint_executor.shutdown(wait=False, cancel_futures=True)
+                        _logger.error(
+                            "Post-equilibration checkpoint cancelled. Exiting."
+                        )
+                        _sys.exit(1)
 
         # Current block number.
         block = self._start_block
@@ -1244,26 +1246,26 @@ def run(self):
             # oversubscribed by a factor of self._config.oversubscription_factor.
             for j in range(num_batches):
                 replicas = replica_list[j * num_workers : (j + 1) * num_workers]
-                with ThreadPoolExecutor(max_workers=num_workers) as executor:
-                    try:
-                        for result, index, energies in executor.map(
-                            self._run_block,
-                            replicas,
-                            repeat(self._lambda_values),
-                            repeat(is_gcmc),
-                            repeat(write_gcmc_ghosts),
-                            repeat(is_terminal_flip),
-                            repeat(time_ns),
-                        ):
-                            if not result:
-                                _logger.error(
-                                    f"Dynamics failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {energies}"
-                                )
-                                raise energies
-                            results.append((index, energies))
-                    except KeyboardInterrupt:
-                        _logger.error("Dynamics cancelled. Exiting.")
-                        _sys.exit(1)
+                try:
+                    for result, index, energies in dynamics_executor.map(
+                        self._run_block,
+                        replicas,
+                        repeat(self._lambda_values),
+                        repeat(is_gcmc),
+                        repeat(write_gcmc_ghosts),
+                        repeat(is_terminal_flip),
+                        repeat(time_ns),
+                    ):
+                        if not result:
+                            _logger.error(
+                                f"Dynamics failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {energies}"
+                            )
+                            raise energies
+                        results.append((index, energies))
+                except KeyboardInterrupt:
+                    dynamics_executor.shutdown(wait=False, cancel_futures=True)
+                    _logger.error("Dynamics cancelled. Exiting.")
+                    _sys.exit(1)
 
             # Checkpoint.
             if is_checkpoint or i == cycles - 1:
@@ -1280,21 +1282,23 @@ def run(self):
                             j * num_checkpoint_workers : (j + 1)
                             * num_checkpoint_workers
                         ]
-                        with ThreadPoolExecutor(max_workers=num_workers) as executor:
-                            try:
-                                for index, error in executor.map(
-                                    self._backup_checkpoint,
-                                    replicas,
-                                ):
-                                    if not result:
-                                        _logger.error(
-                                            f"Backup failed for {_lam_sym} = "
-                                            f"{self._lambda_values[index]:.5f}: {error}"
-                                        )
-                                        raise error
-                            except KeyboardInterrupt:
-                                _logger.error("Backup cancelled. Exiting.")
-                                _sys.exit(1)
+                        try:
+                            for index, error in checkpoint_executor.map(
+                                self._backup_checkpoint,
+                                replicas,
+                            ):
+                                if not result:
+                                    _logger.error(
+                                        f"Backup failed for {_lam_sym} = "
+                                        f"{self._lambda_values[index]:.5f}: {error}"
+                                    )
+                                    raise error
+                        except KeyboardInterrupt:
+                            checkpoint_executor.shutdown(
+                                wait=False, cancel_futures=True
+                            )
+                            _logger.error("Backup cancelled. Exiting.")
+                            _sys.exit(1)
 
                     # Now write the new checkpoint files.
                     for j in range(num_checkpoint_batches):
@@ -1303,25 +1307,27 @@ def run(self):
                             j * num_checkpoint_workers : (j + 1)
                             * num_checkpoint_workers
                         ]
-                        with ThreadPoolExecutor(max_workers=num_workers) as executor:
-                            try:
-                                for index, error in executor.map(
-                                    self._checkpoint,
-                                    replicas,
-                                    repeat(self._lambda_values),
-                                    repeat(block),
-                                    repeat(num_blocks + int(rem > 0)),
-                                    repeat(i == cycles - 1),
-                                ):
-                                    if error:
-                                        _logger.error(
-                                            f"Checkpoint failed for {_lam_sym} = "
-                                            f"{self._lambda_values[index]:.5f}: {error}"
-                                        )
-                                        raise error
-                            except KeyboardInterrupt:
-                                _logger.error("Checkpoint cancelled. Exiting.")
-                                _sys.exit(1)
+                        try:
+                            for index, error in checkpoint_executor.map(
+                                self._checkpoint,
+                                replicas,
+                                repeat(self._lambda_values),
+                                repeat(block),
+                                repeat(num_blocks + int(rem > 0)),
+                                repeat(i == cycles - 1),
+                            ):
+                                if error:
+                                    _logger.error(
+                                        f"Checkpoint failed for {_lam_sym} = "
+                                        f"{self._lambda_values[index]:.5f}: {error}"
+                                    )
+                                    raise error
+                        except KeyboardInterrupt:
+                            checkpoint_executor.shutdown(
+                                wait=False, cancel_futures=True
+                            )
+                            _logger.error("Checkpoint cancelled. Exiting.")
+                            _sys.exit(1)
 
             # Assemble an energy matrix from the results.
             _logger.info("Assembling energy matrix")
@@ -1377,6 +1383,9 @@ def run(self):
                     with open(self._repex_state, "wb") as f:
                         _pickle.dump(self._dynamics_cache, f)
 
+        dynamics_executor.shutdown(wait=True)
+        checkpoint_executor.shutdown(wait=True)
+
         # Record the end time for the production block.
         prod_end = time()
 

From c0b7d5bed681903f7f350edd04ee485ce2d365fc Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 21 Jul 2026 15:15:25 +0100
Subject: [PATCH 14/67] Add oversubscription_factor to allowed_diffs.

---
 CHANGELOG.md              | 1 +
 src/somd2/runner/_base.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3005f3a..a9681d0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ Changelog
 * Give alchemical ions their own plain morph lambda schedule so they interpolate correctly under non-standard lambda schedules [#169](https://github.com/OpenBioSim/somd2/pull/169).
 * Persist alchemical ion identity across restarts so the same molecule is reused regardless of GCMC state [#172](https://github.com/OpenBioSim/somd2/pull/172).
 * Use perisistent `ThreadPoolExector` objects within the main replica exchange dynamics block [#175](https://github.com/OpenBioSim/somd2/pull/175).
+* Allow `oversubscription_factor` to change on restart [#177](https://github.com/OpenBioSim/somd2/pull/177).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index def58bb..26487b8 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -1798,6 +1798,7 @@ def _compare_configs(config1, config2):
             "log_file",
             "overwrite",
             "timeout",
+            "oversubscription_factor",
             "restraint_search_time",
             "restraint_search_frequency",
         ]

From 59ae1ec7800cc07f048345211e91850b5d51d64d Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 22 Jul 2026 11:30:51 +0100
Subject: [PATCH 15/67] Restrict energy decomposition to force groups used for
 integration.

---
 CHANGELOG.md              |  1 +
 src/somd2/runner/_base.py | 12 ++++++++++++
 2 files changed, 13 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a9681d0..3fcfe2e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,7 @@ Changelog
 * Persist alchemical ion identity across restarts so the same molecule is reused regardless of GCMC state [#172](https://github.com/OpenBioSim/somd2/pull/172).
 * Use perisistent `ThreadPoolExector` objects within the main replica exchange dynamics block [#175](https://github.com/OpenBioSim/somd2/pull/175).
 * Allow `oversubscription_factor` to change on restart [#177](https://github.com/OpenBioSim/somd2/pull/177).
+* Restrict energy component decomposition to force groups that are used for integration [#180](https://github.com/OpenBioSim/somd2/pull/180).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 26487b8..1c915e6 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -615,6 +615,9 @@ def __init__(self, system, config):
         # used to skip duplicate rows on restart.
         self._last_ec_time = {}
 
+        # Per-window cache of the integrator's integration force groups bitmask.
+        self._integration_groups = {}
+
         # Store the current system as a reference.
         self._reference_system = self._system.clone()
 
@@ -2514,10 +2517,19 @@ def _save_energy_components(self, index, context, time_ns):
         if time_ns <= self._last_ec_time[index]:
             return
 
+        if index not in self._integration_groups:
+            self._integration_groups[index] = (
+                context.getIntegrator().getIntegrationForceGroups()
+            )
+        integration_groups = self._integration_groups[index]
+
         # Use the named force groups already assigned by sire_to_openmm_system,
         # sorted alphabetically for a consistent column order across runs.
+        # Skip any group not actually used for integration.
         energies = {}
         for name, grp in sorted(context._force_group_map.items()):
+            if not integration_groups & (1 << grp):
+                continue
             state = context.getState(getEnergy=True, groups=(1 << grp))
             energies[name] = state.getPotentialEnergy().value_in_unit(
                 openmm.unit.kilocalories_per_mole

From 761c691be4b73ecb0993f62e97046debca4322d0 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 22 Jul 2026 11:44:28 +0100
Subject: [PATCH 16/67] Parallelise replica mixing.

---
 CHANGELOG.md               |  1 +
 src/somd2/runner/_repex.py | 82 ++++++++++++++++++++++++--------------
 2 files changed, 54 insertions(+), 29 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3fcfe2e..44bb716 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ Changelog
 * Use perisistent `ThreadPoolExector` objects within the main replica exchange dynamics block [#175](https://github.com/OpenBioSim/somd2/pull/175).
 * Allow `oversubscription_factor` to change on restart [#177](https://github.com/OpenBioSim/somd2/pull/177).
 * Restrict energy component decomposition to force groups that are used for integration [#180](https://github.com/OpenBioSim/somd2/pull/180).
+* Parallelise replica mixing [#181](https://github.com/OpenBioSim/somd2/pull/181).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index ea73ea3..4afe6e2 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -538,7 +538,37 @@ def set_states(self, states):
         """
         self._states = states
 
-    def mix_states(self, old_states):
+    def _seed_replica(self, i):
+        """
+        Apply the (possibly new, post-mix) state to replica i's context,
+        including any GCMC water-state swap. Only touches replica i's own
+        context/sampler, so this is safe to run concurrently with other
+        replicas' calls from a thread pool.
+        """
+        state = self._states[i]
+
+        _logger.debug(f"Replica {i} seeded from state {state}")
+        self._apply_openmm_state(
+            self._dynamics[i].context(), self._openmm_states[state]
+        )
+
+        # Swap the water state in the GCMCSamplers.
+        if self._gcmc_samplers[i] is not None:
+            # Find the indices of the water states that differ.
+            water_idxs = _np.where(self._gcmc_states[i] != self._gcmc_states[state])[0]
+
+            # Update the water state in the GCMCSampler.
+            self._gcmc_samplers[i].push()
+            try:
+                self._gcmc_samplers[i]._set_water_state(
+                    self._dynamics[i].context(),
+                    indices=water_idxs,
+                    states=self._gcmc_states[state][water_idxs],
+                )
+            finally:
+                self._gcmc_samplers[i].pop()
+
+    def mix_states(self, old_states, executor=None):
         """
         Mix the states of the dynamics objects.
 
@@ -546,35 +576,29 @@ def mix_states(self, old_states):
         ----------
         old_states : numpy.ndarray
             The state indices from before the last replica mix.
-        """
-        # Mix the states.
-        for i, state in enumerate(self._states):
-            # The state has changed.
-            if i != state:
-                _logger.debug(f"Replica {i} seeded from state {state}")
-                self._apply_openmm_state(
-                    self._dynamics[i].context(), self._openmm_states[state]
-                )
 
-                # Swap the water state in the GCMCSamplers.
-                if self._gcmc_samplers[i] is not None:
-                    # Find the indices of the water states that differ.
-                    water_idxs = _np.where(
-                        self._gcmc_states[i] != self._gcmc_states[state]
-                    )[0]
-
-                    # Update the water state in the GCMCSampler.
-                    self._gcmc_samplers[i].push()
-                    try:
-                        self._gcmc_samplers[i]._set_water_state(
-                            self._dynamics[i].context(),
-                            indices=water_idxs,
-                            states=self._gcmc_states[state][water_idxs],
-                        )
-                    finally:
-                        self._gcmc_samplers[i].pop()
+        executor : concurrent.futures.ThreadPoolExecutor, optional
+            Executor used to apply the per-replica state changes (an
+            OpenMM setPositions/setVelocities/setPeriodicBoxVectors call
+            per changed replica, each against a different context) in
+            parallel. Each replica's context is independent of every
+            other, so this is safe. Falls back to a serial loop if not
+            provided.
+        """
+        # Replicas whose state actually changed.
+        changed = [i for i, state in enumerate(self._states) if i != state]
+
+        if executor is not None and len(changed) > 1:
+            # Consume the map so we block until every replica is seeded
+            # and any exception raised in a worker thread propagates here.
+            list(executor.map(self._seed_replica, changed))
+        else:
+            for i in changed:
+                self._seed_replica(i)
 
-            # Update the swap matrix.
+        # Update the swap matrix. Cheap, CPU-only bookkeeping - kept out of
+        # the parallel section above to avoid any shared-array races.
+        for i, state in enumerate(self._states):
             self._num_swaps[old_states[i], state] += 1
 
     def get_proposed(self):
@@ -1344,7 +1368,7 @@ def run(self):
                     self._dynamics_cache.get_accepted(),
                 )
             )
-            self._dynamics_cache.mix_states(old_states)
+            self._dynamics_cache.mix_states(old_states, executor=dynamics_executor)
 
             # Snapshot the pre-run state for crash recovery.
             if self._config.auto_fix_minimise:

From 00d637e6bf246bc3c32d1a201a36a7756ed7fd89 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 27 Jul 2026 14:16:54 +0100
Subject: [PATCH 17/67] Query the correct GPU when checking device memory.
 [closes #182]

---
 CHANGELOG.md               |  1 +
 src/somd2/runner/_repex.py | 90 +++++++++++++++++++++++++++++++++-----
 tests/runner/test_repex.py | 87 ++++++++++++++++++++++++++++++++++++
 3 files changed, 166 insertions(+), 12 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 44bb716..4420b78 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@ Changelog
 * Allow `oversubscription_factor` to change on restart [#177](https://github.com/OpenBioSim/somd2/pull/177).
 * Restrict energy component decomposition to force groups that are used for integration [#180](https://github.com/OpenBioSim/somd2/pull/180).
 * Parallelise replica mixing [#181](https://github.com/OpenBioSim/somd2/pull/181).
+* Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#182](https://github.com/OpenBioSim/somd2/issues/182).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 4afe6e2..2cb7369 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -54,6 +54,7 @@ def __init__(
         output_directory=None,
         perturbed_system=None,
         xml_filenames=None,
+        gpu_devices=None,
     ):
         """
         Constructor.
@@ -89,6 +90,11 @@ def __init__(
         xml_filenames: list of str
             A list of file paths for the OpenMM XML output, one per replica.
             If None, XML files are not written.
+
+        gpu_devices: list
+            The physical devices backing each OpenMM device index, i.e. the
+            entries of CUDA_VISIBLE_DEVICES. Used to query the memory of the
+            right device. If None, the OpenMM index is used directly.
         """
 
         # Warn if the number of replicas is not a multiple of the number of GPUs.
@@ -101,6 +107,7 @@ def __init__(
         # Initialise attributes.
         self._lambdas = lambdas
         self._rest2_scale_factors = rest2_scale_factors
+        self._gpu_devices = gpu_devices
         self._states = _np.array(range(len(lambdas)))
         self._time = None
         self._openmm_states = [None] * len(lambdas)
@@ -250,7 +257,9 @@ def _create_dynamics(
 
             # Record baseline memory before the first replica on this device.
             if device not in device_mem:
-                used_before, _, total_mem = self._check_device_memory(device)
+                used_before, _, total_mem = self._check_device_memory(
+                    self._physical_device(device)
+                )
                 device_mem[device] = {
                     "before": used_before,
                     "total": total_mem,
@@ -335,7 +344,9 @@ def _create_dynamics(
 
             # Estimate memory after the first or second replica.
             if info["count"] == 1:
-                used_mem, _, _ = self._check_device_memory(device)
+                used_mem, _, _ = self._check_device_memory(
+                    self._physical_device(device)
+                )
                 info["after_first"] = used_mem
 
                 if num_contexts == 1:
@@ -346,7 +357,9 @@ def _create_dynamics(
                     est_total = None
 
             elif info["count"] == 2:
-                used_mem, _, _ = self._check_device_memory(device)
+                used_mem, _, _ = self._check_device_memory(
+                    self._physical_device(device)
+                )
                 # The first replica includes one-time context overhead.
                 # The marginal cost of subsequent replicas is the difference
                 # between the second and first.
@@ -619,24 +632,65 @@ def get_swaps(self):
         """
         return self._num_swaps
 
+    def _physical_device(self, device):
+        """
+        Map an OpenMM device index to the physical device backing it.
+
+        Parameters
+        ----------
+
+        device: int
+            The OpenMM device index, which is relative to the visible set.
+
+        Returns
+        -------
+
+        int, str
+            The physical device, i.e. the matching entry from
+            CUDA_VISIBLE_DEVICES. Falls back to the OpenMM index if the visible
+            set is unknown, which is correct whenever it starts at zero and is
+            contiguous.
+        """
+        gpu_devices = getattr(self, "_gpu_devices", None)
+
+        if gpu_devices is None or device >= len(gpu_devices):
+            return device
+
+        return gpu_devices[device]
+
     @staticmethod
-    def _check_device_memory(device_index=0):
+    def _check_device_memory(device=0):
         """
         Check the memory usage of the specified GPU device.
 
         Parameters
         ----------
 
-        index: int
-            The index of the GPU device.
+        device: int, str
+            The device to query. This is the physical device, i.e. an entry
+            from CUDA_VISIBLE_DEVICES (or the equivalent for other platforms),
+            not the index used by OpenMM. OpenMM numbers devices relative to
+            the visible set, whereas pynvml and pyopencl enumerate all devices
+            on the machine, so the two only agree when the visible set starts
+            at zero and is contiguous. CUDA_VISIBLE_DEVICES entries may be
+            either an index or a UUID.
         """
 
+        device = str(device).strip()
+
+        # A UUID cannot be used to index into the OpenCL device list.
+        is_uuid = device.startswith("GPU-") or device.startswith("MIG-")
+        device_index = None if is_uuid else int(device)
+
         # Try to use pyopencl to detect the GPU vendor.
         vendor = None
         ocl_device = None
         try:
             import pyopencl as cl
 
+            if device_index is None:
+                raise ValueError("Cannot index OpenCL devices by UUID")
+
             platforms = cl.get_platforms()
             all_devices = []
             for platform in platforms:
@@ -666,15 +720,20 @@ def _check_device_memory(device_index=0):
                 import pynvml
 
                 pynvml.nvmlInit()
-                handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
+                if is_uuid:
+                    handle = pynvml.nvmlDeviceGetHandleByUUID(device.encode())
+                else:
+                    handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
                 memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
                 pynvml.nvmlShutdown()
                 return (memory.used, memory.free, memory.total)
             except Exception as e:
                 if vendor is None:
-                    msg = f"Could not get GPU memory info for device {device_index} via OpenCL or pynvml: {e}"
+                    msg = f"Could not get GPU memory info for device {device} via OpenCL or pynvml: {e}"
                 else:
-                    msg = f"Could not get NVIDIA GPU memory info for device {device_index}: {e}"
+                    msg = (
+                        f"Could not get NVIDIA GPU memory info for device {device}: {e}"
+                    )
                 _logger.error(msg)
                 raise RuntimeError(msg) from e
 
@@ -692,9 +751,7 @@ def _check_device_memory(device_index=0):
                 used = total - free
                 return (used, free, total)
             except Exception as e:
-                msg = (
-                    f"Could not get AMD GPU memory info for device {device_index}: {e}"
-                )
+                msg = f"Could not get AMD GPU memory info for device {device}: {e}"
                 _logger.error(msg)
                 raise RuntimeError(msg) from e
 
@@ -764,6 +821,10 @@ def __init__(self, system, config):
         else:
             self._num_gpus = min(self._config.max_gpus, len(gpu_devices))
 
+        # The physical devices backing each OpenMM device index. OpenMM numbers
+        # devices relative to the visible set, so index i is gpu_devices[i].
+        self._gpu_devices = list(gpu_devices)[: self._num_gpus]
+
         # Auto-generate a Boresch restraint for ABFE runs with no user-supplied
         # restraint. This must happen before the dynamics cache is built below,
         # since the per-replica OpenMM contexts it creates are fixed at
@@ -851,6 +912,7 @@ def __init__(self, system, config):
                 perturbed_system=self._perturbed_system,
                 output_directory=self._config.output_directory,
                 xml_filenames=xml_filenames,
+                gpu_devices=self._gpu_devices,
             )
 
         else:
@@ -921,6 +983,10 @@ def __init__(self, system, config):
             if not isinstance(self._system, list):
                 self._system.set_time(time)
 
+            # The physical device list is not pickled, since the run may be
+            # restarted against a different set of GPUs.
+            self._dynamics_cache._gpu_devices = self._gpu_devices
+
             # Create the dynamics objects.
             self._dynamics_cache._create_dynamics(
                 self._system,
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 9336053..e7b970d 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -148,3 +148,90 @@ def test_rest2_selection(ethane_methanol, rest2_selection, is_valid):
         else:
             with pytest.raises(ValueError):
                 runner = RunnerBase(ethane_methanol, Config(**config))
+
+
+@pytest.mark.parametrize(
+    "gpu_devices, expected",
+    [
+        # Visible set starts at zero and is contiguous, so the OpenMM index and
+        # the physical device agree.
+        (["0", "1", "2"], ["0", "1", "2"]),
+        # Offset visible set: OpenMM index 0 is physical device 1.
+        (["1", "2"], ["1", "2"]),
+        # A single device that is not device zero.
+        (["3"], ["3"]),
+        # CUDA_VISIBLE_DEVICES may hold UUIDs rather than indices.
+        (["GPU-abc123", "GPU-def456"], ["GPU-abc123", "GPU-def456"]),
+        # Unknown visible set falls back to the OpenMM index.
+        (None, [0, 1]),
+    ],
+)
+def test_physical_device_mapping(gpu_devices, expected):
+    """
+    Validate that an OpenMM device index is mapped to the physical device
+    backing it.
+
+    OpenMM numbers devices relative to the visible set, whereas pynvml and
+    pyopencl enumerate every device on the machine. Querying the memory of a
+    device by its OpenMM index therefore reports the wrong GPU whenever the
+    visible set does not start at zero.
+    """
+    from somd2.runner._repex import DynamicsCache
+
+    cache = object.__new__(DynamicsCache)
+    cache._gpu_devices = gpu_devices
+
+    assert [cache._physical_device(i) for i in range(len(expected))] == expected
+
+
+@pytest.mark.parametrize(
+    "device, key, value",
+    [
+        ("2", "index", 2),
+        (2, "index", 2),
+        ("GPU-abc123", "uuid", b"GPU-abc123"),
+    ],
+)
+def test_check_device_memory_queries_requested_device(monkeypatch, device, key, value):
+    """
+    Validate that the memory query is made against the device it was asked
+    for, by index or by UUID.
+    """
+    import sys
+    import types
+
+    from somd2.runner._repex import DynamicsCache
+
+    pynvml = pytest.importorskip("pynvml")
+
+    # Force the OpenCL branch to fail so that the pynvml path is always taken,
+    # regardless of what the machine running the tests has installed.
+    broken = types.SimpleNamespace()
+
+    def get_platforms():
+        raise RuntimeError("no OpenCL")
+
+    broken.get_platforms = get_platforms
+    monkeypatch.setitem(sys.modules, "pyopencl", broken)
+
+    requested = {}
+
+    class Memory:
+        used, free, total = 1, 2, 3
+
+    def by_index(index):
+        requested["index"] = index
+        return "handle"
+
+    def by_uuid(uuid):
+        requested["uuid"] = uuid
+        return "handle"
+
+    monkeypatch.setattr(pynvml, "nvmlInit", lambda: None)
+    monkeypatch.setattr(pynvml, "nvmlShutdown", lambda: None)
+    monkeypatch.setattr(pynvml, "nvmlDeviceGetHandleByIndex", by_index)
+    monkeypatch.setattr(pynvml, "nvmlDeviceGetHandleByUUID", by_uuid)
+    monkeypatch.setattr(pynvml, "nvmlDeviceGetMemoryInfo", lambda handle: Memory)
+
+    assert DynamicsCache._check_device_memory(device) == (1, 2, 3)
+    assert requested == {key: value}

From 688ac1fe5a7301a3c1d59b4bbc92cd6da727d6d5 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 28 Jul 2026 10:15:47 +0100
Subject: [PATCH 18/67] Store GCMC statistics per lambda value.

---
 CHANGELOG.md                    |  3 +-
 src/somd2/runner/_base.py       | 47 ++++++++++++++++
 src/somd2/runner/_repex.py      | 51 ++++++++++++++---
 src/somd2/runner/_runner.py     |  7 ++-
 tests/runner/test_gcmc_stats.py | 97 +++++++++++++++++++++++++++++++++
 5 files changed, 195 insertions(+), 10 deletions(-)
 create mode 100644 tests/runner/test_gcmc_stats.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4420b78..f1578fa 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,7 +12,8 @@ Changelog
 * Allow `oversubscription_factor` to change on restart [#177](https://github.com/OpenBioSim/somd2/pull/177).
 * Restrict energy component decomposition to force groups that are used for integration [#180](https://github.com/OpenBioSim/somd2/pull/180).
 * Parallelise replica mixing [#181](https://github.com/OpenBioSim/somd2/pull/181).
-* Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#182](https://github.com/OpenBioSim/somd2/issues/182).
+* Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#183](https://github.com/OpenBioSim/somd2/issues/183).
+* Store GCMC sampling statistics per lambda value, converting those from earlier checkpoints on restart [#184](https://github.com/OpenBioSim/somd2/pull/184).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 1c915e6..69da02d 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -2415,6 +2415,53 @@ def _write_checkpoint_system(self, system, index, context=None, gcmc_sampler=Non
         system.delete_all_frames()
         _sr.stream.save(system, self._filenames[index]["checkpoint"])
 
+    @staticmethod
+    def _is_legacy_gcmc_stats(stats):
+        """
+        Whether GCMC statistics are in the format used before a sampler could
+        be re-used across lambda values.
+
+        Those were a flat dictionary of counters for a single lambda value,
+        rather than a dictionary of counters keyed by lambda value.
+
+        Parameters
+        ----------
+
+        stats: dict
+            The GCMC sampling statistics.
+
+        Returns
+        -------
+
+        bool
+            Whether the statistics are in the old format.
+        """
+        return isinstance(stats, dict) and "num_moves" in stats
+
+    @staticmethod
+    def _convert_legacy_gcmc_stats(stats, lambda_value):
+        """
+        Convert GCMC statistics from the old format to the current one.
+
+        Parameters
+        ----------
+
+        stats: dict
+            A flat dictionary of counters, for a single lambda value.
+
+        lambda_value: float
+            The lambda value that the statistics belong to.
+
+        Returns
+        -------
+
+        dict
+            The statistics, keyed by lambda value.
+        """
+        from loch import GCMCSampler as _GCMCSampler
+
+        return {_GCMCSampler.stats_key(lambda_value): dict(stats)}
+
     def _backup_checkpoint(self, index):
         """
         Create a backup of the previous checkpoint files.
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 2cb7369..57f8a32 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -113,7 +113,8 @@ def __init__(
         self._openmm_states = [None] * len(lambdas)
         self._gcmc_samplers = [None] * len(lambdas)
         self._gcmc_states = [None] * len(lambdas)
-        self._gcmc_stats = [None] * len(lambdas)
+        # GCMC statistics for the whole simulation, keyed by lambda value.
+        self._gcmc_stats = None
         self._terminal_flip_stats = [[0, 0]] * len(lambdas)
         self._num_proposed = _np.matrix(_np.zeros((len(lambdas), len(lambdas))))
         self._num_accepted = _np.matrix(_np.zeros((len(lambdas), len(lambdas))))
@@ -143,7 +144,7 @@ def __setstate__(self, state):
         # so that old checkpoint files can still be loaded.
         n = len(self._lambdas)
         if not hasattr(self, "_gcmc_stats"):
-            self._gcmc_stats = [None] * n
+            self._gcmc_stats = None
         if not hasattr(self, "_gcmc_states"):
             self._gcmc_states = [None] * n
         if not hasattr(self, "_terminal_flip_stats"):
@@ -151,6 +152,16 @@ def __setstate__(self, state):
         if not hasattr(self, "_time"):
             self._time = None
 
+        # Checkpoints written before a sampler could be re-used across lambda
+        # values stored the GCMC statistics as a list of counters per replica.
+        # Convert these to a single dictionary keyed by lambda value.
+        if isinstance(self._gcmc_stats, list):
+            converted = {}
+            for lam, stats in zip(self._lambdas, self._gcmc_stats):
+                if _RunnerBase._is_legacy_gcmc_stats(stats):
+                    converted.update(_RunnerBase._convert_legacy_gcmc_stats(stats, lam))
+            self._gcmc_stats = converted if converted else None
+
     def __getstate__(self):
         """
         Get the state of the object.
@@ -1020,8 +1031,11 @@ def __init__(self, system, config):
                         )
                     finally:
                         gcmc_sampler.pop()
-                    if self._dynamics_cache._gcmc_stats[i] is not None:
-                        gcmc_sampler.restore_stats(self._dynamics_cache._gcmc_stats[i])
+
+                    # Samplers keep only the lambda values they visit, so it's
+                    # safe to hand each of them the whole simulation's stats.
+                    if self._dynamics_cache._gcmc_stats is not None:
+                        gcmc_sampler.restore_stats(self._dynamics_cache._gcmc_stats)
 
         # Log the GCMC sphere centre for each replica using the actual context
         # positions (accurate for both fresh runs and restarts).
@@ -2232,15 +2246,36 @@ def _mix_replicas(num_replicas, energy_matrix, proposed, accepted):
 
         return states
 
-    def _save_sampler_stats(self):
+    def _merge_gcmc_stats(self):
         """
-        Save GCMC and terminal flip sampler statistics to the dynamics cache
-        prior to pickling.
+        Merge the GCMC sampling statistics from every sampler.
+
+        A sampler accumulates statistics for each lambda value it visits, so
+        the results are gathered into a single dictionary keyed by lambda
+        value. Samplers only report the lambda values they visit, so the keys
+        are disjoint and the merge order doesn't matter.
+
+        Returns
+        -------
+
+        dict
+            The statistics for each lambda value, or None if not using GCMC.
         """
+        stats = {}
+
         for i in range(len(self._lambda_values)):
             _, gcmc_sampler = self._dynamics_cache.get(i)
             if gcmc_sampler is not None:
-                self._dynamics_cache._gcmc_stats[i] = gcmc_sampler.get_stats()
+                stats.update(gcmc_sampler.get_stats())
+
+        return stats if stats else None
+
+    def _save_sampler_stats(self):
+        """
+        Save GCMC and terminal flip sampler statistics to the dynamics cache
+        prior to pickling.
+        """
+        self._dynamics_cache._gcmc_stats = self._merge_gcmc_stats()
 
         if self._terminal_flip_samplers is not None:
             self._dynamics_cache._terminal_flip_stats = [
diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index 56dc31e..fb0dac7 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -823,7 +823,12 @@ def generate_lam_vals(lambda_base, increment=0.001):
             stats = self._load_sampler_stats(index)
             if stats is not None:
                 if gcmc_sampler is not None and "gcmc" in stats:
-                    gcmc_sampler.restore_stats(stats["gcmc"])
+                    gcmc_stats = stats["gcmc"]
+                    if self._is_legacy_gcmc_stats(gcmc_stats):
+                        gcmc_stats = self._convert_legacy_gcmc_stats(
+                            gcmc_stats, self._lambda_values[index]
+                        )
+                    gcmc_sampler.restore_stats(gcmc_stats)
                 if terminal_flip_sampler is not None and "terminal_flip" in stats:
                     attempted, accepted = stats["terminal_flip"]
                     terminal_flip_sampler.reset(attempted, accepted)
diff --git a/tests/runner/test_gcmc_stats.py b/tests/runner/test_gcmc_stats.py
new file mode 100644
index 0000000..c6ff648
--- /dev/null
+++ b/tests/runner/test_gcmc_stats.py
@@ -0,0 +1,97 @@
+import numpy as np
+import pytest
+
+from somd2.runner._base import RunnerBase
+from somd2.runner._repex import DynamicsCache
+
+
+def counters(num_moves=0):
+    """A flat dictionary of counters, as written before the format changed."""
+    return {
+        "num_moves": num_moves,
+        "num_accepted": num_moves,
+        "num_insertions": 0,
+        "num_deletions": 0,
+        "num_accepted_attempts": 0,
+    }
+
+
+class TestLegacyDetection:
+    """Tests for detecting the format of GCMC statistics."""
+
+    def test_flat_counters_are_legacy(self):
+        assert RunnerBase._is_legacy_gcmc_stats(counters(5))
+
+    def test_keyed_by_lambda_is_current(self):
+        assert not RunnerBase._is_legacy_gcmc_stats({"0.00000": counters(5)})
+
+    @pytest.mark.parametrize("stats", [None, {}, []])
+    def test_other_values_are_not_legacy(self, stats):
+        assert not RunnerBase._is_legacy_gcmc_stats(stats)
+
+    def test_conversion_keys_by_lambda(self):
+        converted = RunnerBase._convert_legacy_gcmc_stats(counters(5), 0.33333)
+        assert converted == {"0.33333": counters(5)}
+
+    def test_conversion_is_a_copy(self):
+        """The original must not be aliased into the converted result."""
+        original = counters(5)
+        converted = RunnerBase._convert_legacy_gcmc_stats(original, 0.0)
+        original["num_moves"] = 99
+        assert converted["0.00000"]["num_moves"] == 5
+
+
+class TestLegacyRepexCheckpoint:
+    """Tests for restoring a replica exchange checkpoint."""
+
+    @staticmethod
+    def make_state(gcmc_stats, lambdas=(0.0, 0.5, 1.0)):
+        n = len(lambdas)
+        return {
+            "_lambdas": list(lambdas),
+            "_rest2_scale_factors": [1.0] * n,
+            "_states": np.arange(n),
+            "_time": None,
+            "_openmm_states": [None] * n,
+            "_gcmc_samplers": [None] * n,
+            "_gcmc_states": [None] * n,
+            "_gcmc_stats": gcmc_stats,
+            "_terminal_flip_stats": [[0, 0]] * n,
+            "_num_proposed": np.zeros((n, n)),
+            "_num_accepted": np.zeros((n, n)),
+            "_num_swaps": np.zeros((n, n)),
+        }
+
+    def restore(self, gcmc_stats, **kwargs):
+        cache = object.__new__(DynamicsCache)
+        cache.__setstate__(self.make_state(gcmc_stats, **kwargs))
+        return cache
+
+    def test_per_replica_list_is_converted(self):
+        """A list of counters per replica becomes a map keyed by lambda."""
+        cache = self.restore([counters(i) for i in range(3)])
+
+        assert cache._gcmc_stats == {
+            "0.00000": counters(0),
+            "0.50000": counters(1),
+            "1.00000": counters(2),
+        }
+
+    def test_no_gcmc_gives_none(self):
+        """A checkpoint from a run without GCMC has no statistics."""
+        assert self.restore([None, None, None])._gcmc_stats is None
+
+    def test_current_format_is_untouched(self):
+        """A checkpoint already in the current format is left alone."""
+        stats = {"0.00000": counters(4), "1.00000": counters(9)}
+        assert self.restore(dict(stats))._gcmc_stats == stats
+
+    def test_missing_attribute_defaults_to_none(self):
+        """A checkpoint predating GCMC statistics has none."""
+        state = self.make_state(None)
+        del state["_gcmc_stats"]
+
+        cache = object.__new__(DynamicsCache)
+        cache.__setstate__(state)
+
+        assert cache._gcmc_stats is None

From 6b21af9d780e53cbcff51e578b9ff8993391e102 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 28 Jul 2026 12:13:06 +0100
Subject: [PATCH 19/67] Use pip_check: false to avoid NumPy pin errors with
 AmberTools.

---
 recipes/somd2/recipe.yaml | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/recipes/somd2/recipe.yaml b/recipes/somd2/recipe.yaml
index d8916d6..ba83354 100644
--- a/recipes/somd2/recipe.yaml
+++ b/recipes/somd2/recipe.yaml
@@ -37,6 +37,10 @@ tests:
       - python:
           imports:
             - somd2
+          # AmberTools can be pulled in transitively, and installs tools into
+          # site-packages whose metadata still pins numpy <2, so 'pip check'
+          # fails even though nothing here uses them.
+          pip_check: false
       - script:
           - PYTHONPATH=. pytest -vvv --color=yes --import-mode=importlib ./tests
         files:

From 696344f2fe347dd879289b7d39f02f8bd97c3e62 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 4 Aug 2026 18:37:20 +0100
Subject: [PATCH 20/67] Link restart systems to the reference end state.

---
 CHANGELOG.md              | 1 +
 src/somd2/runner/_base.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f1578fa..d699ec4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@ Changelog
 * Parallelise replica mixing [#181](https://github.com/OpenBioSim/somd2/pull/181).
 * Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#183](https://github.com/OpenBioSim/somd2/issues/183).
 * Store GCMC sampling statistics per lambda value, converting those from earlier checkpoints on restart [#184](https://github.com/OpenBioSim/somd2/pull/184).
+* Link restart systems to the reference end state rather than the perturbed one, since that is the coordinate set that dynamics maintains. Perturbable molecules were otherwise resumed from the coordinates they were built with [#189](https://github.com/OpenBioSim/somd2/pull/189).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 69da02d..158f176 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -1723,7 +1723,7 @@ def _check_restart(self):
                         raise ValueError(msg)
 
                 # Store the system to the list.
-                systems[i] = _sr.morph.link_to_perturbed(system)
+                systems[i] = _sr.morph.link_to_reference(system)
 
         # If this is a GCMC simulation, then remove all ghost waters from each of the systems.
         if self._config.gcmc:

From d4e98bdc8333116156b942135d502f263f312175 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 27 Jul 2026 12:58:41 +0100
Subject: [PATCH 21/67] Cap the number of OpenMM contexts used for replica
 exchange.

---
 CHANGELOG.md                |    1 +
 src/somd2/config/_config.py |   57 ++
 src/somd2/runner/_base.py   |    1 +
 src/somd2/runner/_repex.py  | 1069 +++++++++++++++++++++++++++--------
 tests/runner/test_repex.py  |  490 +++++++++++++++-
 5 files changed, 1378 insertions(+), 240 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d699ec4..4082038 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,7 @@ Changelog
 * Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#183](https://github.com/OpenBioSim/somd2/issues/183).
 * Store GCMC sampling statistics per lambda value, converting those from earlier checkpoints on restart [#184](https://github.com/OpenBioSim/somd2/pull/184).
 * Link restart systems to the reference end state rather than the perturbed one, since that is the coordinate set that dynamics maintains. Perturbable molecules were otherwise resumed from the coordinates they were built with [#189](https://github.com/OpenBioSim/somd2/pull/189).
+* Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas.
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index 9efd79c..4096cac 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -138,6 +138,8 @@ def __init__(
         opencl_platform_index=0,
         oversubscription_factor=1,
         replica_exchange=False,
+        max_contexts=None,
+        update_constraints=True,
         randomise_velocities=False,
         perturbed_system=None,
         terminal_flip_frequency=None,
@@ -377,6 +379,27 @@ def __init__(
             Whether to run replica exchange simulation. Currently this can only be used when
             GPU resources are available.
 
+        max_contexts: int
+            The maximum number of OpenMM contexts to create for a replica exchange
+            simulation. If None, then one context is created per replica, which is
+            fastest, but limits the number of replicas to those that fit in GPU memory.
+            If fewer contexts than replicas are requested, then each context is re-used
+            to propagate several replicas per cycle, changing its lambda value as it
+            goes. This lifts the memory limit at the cost of some performance. When
+            re-using contexts, 'frame_frequency' must equal 'checkpoint_frequency'.
+
+        update_constraints: bool
+            Whether the constraints are updated when the lambda value of a context is
+            changed, i.e. whether constrained bond lengths are allowed to perturb with
+            lambda. This is only used when contexts are re-used across lambda values,
+            i.e. when 'max_contexts' is less than the number of replicas. Updating the
+            constraints is correct, but requires the OpenMM context to be reinitialised
+            whenever a constrained bond length actually changes, which is slow. Set this
+            to False if that overhead is significant; the constrained bond lengths are
+            then frozen at those of the lambda value the context was created at. Note
+            that this is distinct from 'dynamic_constraints', which controls where the
+            constraint lengths are taken from rather than whether they track lambda.
+
         randomise_velocities: bool
             Whether to randomise velocities at the start of each replica exchange cycle
             or following a terminal flip Monte Carlo move.
@@ -630,6 +653,8 @@ def __init__(
         self.opencl_platform_index = opencl_platform_index
         self.oversubscription_factor = oversubscription_factor
         self.replica_exchange = replica_exchange
+        self.max_contexts = max_contexts
+        self.update_constraints = update_constraints
         self.randomise_velocities = randomise_velocities
         self.perturbed_system = perturbed_system
         self.terminal_flip_frequency = terminal_flip_frequency
@@ -1778,6 +1803,38 @@ def replica_exchange(self, replica_exchange):
             raise ValueError("'replica_exchange' must be of type 'bool'")
         self._replica_exchange = replica_exchange
 
+    @property
+    def max_contexts(self):
+        return self._max_contexts
+
+    @max_contexts.setter
+    def max_contexts(self, max_contexts):
+        if max_contexts is None or (
+            isinstance(max_contexts, str)
+            and max_contexts.lower().replace(" ", "") == "none"
+        ):
+            self._max_contexts = None
+            return
+
+        if not isinstance(max_contexts, int):
+            try:
+                max_contexts = int(max_contexts)
+            except Exception:
+                raise ValueError("'max_contexts' must be of type 'int'")
+        if max_contexts < 1:
+            raise ValueError("'max_contexts' must be greater than 0")
+        self._max_contexts = max_contexts
+
+    @property
+    def update_constraints(self):
+        return self._update_constraints
+
+    @update_constraints.setter
+    def update_constraints(self, update_constraints):
+        if not isinstance(update_constraints, bool):
+            raise ValueError("'update_constraints' must be of type 'bool'")
+        self._update_constraints = update_constraints
+
     @property
     def randomise_velocities(self):
         return self._randomise_velocities
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 158f176..c42adde 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -1802,6 +1802,7 @@ def _compare_configs(config1, config2):
             "overwrite",
             "timeout",
             "oversubscription_factor",
+            "max_contexts",
             "restraint_search_time",
             "restraint_search_frequency",
         ]
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 57f8a32..d006da9 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -54,6 +54,8 @@ def __init__(
         output_directory=None,
         perturbed_system=None,
         xml_filenames=None,
+        num_slots=None,
+        update_constraints=True,
         gpu_devices=None,
     ):
         """
@@ -91,34 +93,62 @@ def __init__(
             A list of file paths for the OpenMM XML output, one per replica.
             If None, XML files are not written.
 
+        num_slots: int
+            The number of dynamics objects (slots) to create. If None, then one
+            is created per replica. If fewer, then each slot is re-used to
+            propagate several replicas per cycle, changing its lambda value as
+            it goes.
+
+        update_constraints: bool
+            Whether to update the constraints when changing the lambda value of
+            a slot.
+
         gpu_devices: list
             The physical devices backing each OpenMM device index, i.e. the
             entries of CUDA_VISIBLE_DEVICES. Used to query the memory of the
             right device. If None, the OpenMM index is used directly.
         """
 
-        # Warn if the number of replicas is not a multiple of the number of GPUs.
-        if len(lambdas) > num_gpus and len(lambdas) % num_gpus != 0:
+        num_replicas = len(lambdas)
+
+        if num_slots is None:
+            num_slots = num_replicas
+
+        # Warn if the number of slots is not a multiple of the number of GPUs.
+        if num_slots > num_gpus and num_slots % num_gpus != 0:
             _logger.warning(
-                "The number of replicas is not a multiple of the number of GPUs. "
+                "The number of contexts is not a multiple of the number of GPUs. "
                 "This may result in suboptimal performance."
             )
 
         # Initialise attributes.
         self._lambdas = lambdas
         self._rest2_scale_factors = rest2_scale_factors
+        self._num_replicas = num_replicas
+        self._num_slots = num_slots
+        self._update_constraints = update_constraints
         self._gpu_devices = gpu_devices
-        self._states = _np.array(range(len(lambdas)))
+        self._states = _np.array(range(num_replicas))
         self._time = None
-        self._openmm_states = [None] * len(lambdas)
-        self._gcmc_samplers = [None] * len(lambdas)
-        self._gcmc_states = [None] * len(lambdas)
+        self._openmm_states = [None] * num_replicas
+        self._gcmc_states = [None] * num_replicas
         # GCMC statistics for the whole simulation, keyed by lambda value.
         self._gcmc_stats = None
-        self._terminal_flip_stats = [[0, 0]] * len(lambdas)
-        self._num_proposed = _np.matrix(_np.zeros((len(lambdas), len(lambdas))))
-        self._num_accepted = _np.matrix(_np.zeros((len(lambdas), len(lambdas))))
-        self._num_swaps = _np.matrix(_np.zeros((len(lambdas), len(lambdas))))
+        self._energy_trajectories = [None] * num_replicas
+        self._ghost_files = [None] * num_replicas
+        # Waters in the GCMC volume, recorded per replica while it is resident
+        # in its slot. Derived, so not stored in a checkpoint.
+        self._gcmc_num_waters = [None] * num_replicas
+        # Whether the last mix moved a replica's state, and so whether it must
+        # be pushed into the context before the next block.
+        self._state_moved = [False] * num_replicas
+        self._terminal_flip_stats = [[0, 0]] * num_replicas
+        self._num_proposed = _np.matrix(_np.zeros((num_replicas, num_replicas)))
+        self._num_accepted = _np.matrix(_np.zeros((num_replicas, num_replicas)))
+        self._num_swaps = _np.matrix(_np.zeros((num_replicas, num_replicas)))
+
+        # Build the slot layout and the per-slot attributes.
+        self._build_slot_layout()
 
         # Create the dynamics objects.
         self._create_dynamics(
@@ -133,10 +163,44 @@ def __init__(
             xml_filenames=xml_filenames,
         )
 
+    def _build_slot_layout(self):
+        """
+        Assign replicas to slots.
+
+        Each slot is given a contiguous group of replicas, so that it only ever
+        moves between neighbouring lambda values. That keeps the change in the
+        force field parameters, and hence the chance of a constraint update
+        forcing the context to be reinitialised, as small as possible.
+
+        The layout is derived from the number of replicas and slots, so it is
+        rebuilt rather than stored in a checkpoint.
+        """
+        self._gcmc_samplers = [None] * self._num_slots
+        self._slot_replica = [None] * self._num_slots
+
+        self._groups = [
+            [int(r) for r in group]
+            for group in _np.array_split(
+                _np.arange(self._num_replicas), self._num_slots
+            )
+        ]
+
+        # The slot that hosts each replica.
+        self._replica_slot = [None] * self._num_replicas
+        for slot, group in enumerate(self._groups):
+            for replica in group:
+                self._replica_slot[replica] = slot
+
     def __setstate__(self, state):
         """
         Set the state of the object.
         """
+
+        # Checkpoints written before slots were introduced stored the states
+        # unpermuted, with self._states holding the mapping to apply on
+        # restart. They are detected by the absence of "_num_slots".
+        is_legacy = "_num_slots" not in state
+
         for key, value in state.items():
             setattr(self, key, value)
 
@@ -151,6 +215,35 @@ def __setstate__(self, state):
             self._terminal_flip_stats = [[0, 0]] * n
         if not hasattr(self, "_time"):
             self._time = None
+        if not hasattr(self, "_num_replicas"):
+            self._num_replicas = n
+        if not hasattr(self, "_energy_trajectories"):
+            self._energy_trajectories = [None] * n
+        if not hasattr(self, "_ghost_files"):
+            self._ghost_files = [None] * n
+        self._gcmc_num_waters = [None] * n
+
+        # The slot layout is not pickled, since it is rebuilt by
+        # _create_dynamics() when the run is restarted. Older checkpoints
+        # predate slots entirely, in which case there was one per replica.
+        if not hasattr(self, "_num_slots"):
+            self._num_slots = n
+        if not hasattr(self, "_update_constraints"):
+            self._update_constraints = True
+
+        # Convert a legacy checkpoint to the current convention, in which the
+        # stored state of a replica is its own, with the last mix already
+        # applied.
+        if is_legacy:
+            self._openmm_states = [self._openmm_states[s] for s in self._states]
+            self._gcmc_states = [self._gcmc_states[s] for s in self._states]
+
+        # Every replica is seeded from its stored state on a restart, since the
+        # contexts are created from the input system rather than the checkpoint.
+        self._state_moved = [True] * n
+
+        # Rebuild the slot layout, which is derived rather than stored.
+        self._build_slot_layout()
 
         # Checkpoints written before a sampler could be re-used across lambda
         # values stored the GCMC statistics as a list of counters per replica.
@@ -171,6 +264,9 @@ def __getstate__(self):
         d = {
             "_lambdas": self._lambdas,
             "_rest2_scale_factors": self._rest2_scale_factors,
+            "_num_replicas": self._num_replicas,
+            "_num_slots": self._num_slots,
+            "_update_constraints": self._update_constraints,
             "_states": self._states,
             "_time": self._time,
             "_openmm_states": self._openmm_states,
@@ -252,21 +348,35 @@ def _create_dynamics(
         # Per-device memory tracking for estimation.
         device_mem = {}
 
-        # Work out how many replicas are assigned to each device.
-        # Replicas are assigned round-robin, so the first (num_replicas % num_gpus)
-        # devices get one extra replica.
-        base = floor(num_replicas / num_gpus)
-        remainder = num_replicas % num_gpus
+        # Work out how many slots are assigned to each device.
+        # Slots are assigned round-robin, so the first (num_slots % num_gpus)
+        # devices get one extra slot.
+        base = floor(self._num_slots / num_gpus)
+        remainder = self._num_slots % num_gpus
         contexts_per_device = [
             base + (1 if i < remainder else 0) for i in range(num_gpus)
         ]
 
-        # Create the dynamics objects in serial.
-        for i, (lam, scale) in enumerate(zip(lambdas, rest2_scale_factors)):
+        # Record the ghost file for each replica. A slot writes to the file of
+        # whichever replica it currently hosts.
+        if gcmc_kwargs is not None:
+            self._ghost_files = [
+                str(output_directory / f"gcmc_ghosts_{lam:.5f}.txt") for lam in lambdas
+            ]
+
+        # Create the dynamics objects in serial. Each slot is created at the
+        # lambda value of the first replica that it hosts.
+        for i in range(self._num_slots):
+            # The replica that seeds this slot.
+            seed = self._groups[i][0]
+
+            lam = lambdas[seed]
+            scale = rest2_scale_factors[seed]
+
             # Work out the device index.
             device = i % num_gpus
 
-            # Record baseline memory before the first replica on this device.
+            # Record baseline memory before the first slot on this device.
             if device not in device_mem:
                 used_before, _, total_mem = self._check_device_memory(
                     self._physical_device(device)
@@ -277,9 +387,9 @@ def _create_dynamics(
                     "count": 0,
                 }
 
-            # This is a restart, get the system for this replica.
+            # This is a restart, get the system for the seeding replica.
             if isinstance(system, list):
-                mols = system[i]
+                mols = system[seed]
             # This is a new simulation. For lambda > 0.5, use the perturbed
             # system to seed the starting coordinates and periodic space.
             elif perturbed_system is not None and lam > 0.5:
@@ -302,15 +412,17 @@ def _create_dynamics(
                     msg = "loch is not installed. GCMC sampling cannot be performed."
                     _logger.error(msg)
 
-                ghost_file = str(output_directory / f"gcmc_ghosts_{lam:.5f}.txt")
-
-                # Create the GCMC sampler.
+                # Create the GCMC sampler, telling it every lambda value that
+                # this slot will host so that switching between them doesn't
+                # need to build an OpenMM context.
                 gcmc_sampler = GCMCSampler(
                     mols,
                     device=device,
                     lambda_value=lam,
                     rest2_scale=scale,
-                    ghost_file=ghost_file,
+                    lambda_values=[lambdas[r] for r in self._groups[i]],
+                    rest2_scales=[rest2_scale_factors[r] for r in self._groups[i]],
+                    ghost_file=self._ghost_files[seed],
                     **gcmc_kwargs,
                 )
 
@@ -341,12 +453,13 @@ def _create_dynamics(
             # Append the dynamics object.
             self._dynamics.append(dynamics)
 
-            # Write the OpenMM XML file to the output directory.
+            # Write the OpenMM XML file to the output directory. This is
+            # indexed by replica, so use the replica that seeded the slot.
             if xml_filenames is not None:
                 _logger.info(
                     f"Writing OpenMM XML for lambda {lam:.5f} on device {device}"
                 )
-                dynamics.to_xml(xml_filenames[i])
+                dynamics.to_xml(xml_filenames[seed])
 
             # Track memory footprint for this device.
             info = device_mem[device]
@@ -424,65 +537,120 @@ def _create_dynamics(
                 f"Created dynamics object for lambda {lam:.5f} on device {device}"
             )
 
-    def get(self, index):
+            # Leave the slot marked as holding no replica, so that the first
+            # call to load_replica() does the full setup (lambda value, GCMC
+            # parameters, ghost file and sampling statistics) rather than
+            # assuming the seeding replica is already fully installed.
+            self._slot_replica[i] = None
+
+        # Give each replica its own energy trajectory. These are seeded from a
+        # slot's own so that the "ensemble" property is carried over. A slot
+        # accumulates into the trajectory of whichever replica it hosts.
+        for replica in range(self._num_replicas):
+            slot = self._replica_slot[replica]
+            self._energy_trajectories[replica] = self._dynamics[
+                slot
+            ]._d.energy_trajectory()
+
+        # Seed the starting state for every replica from the context of the
+        # slot that hosts it. The GCMC water state must be seeded too, since
+        # load_replica() diffs against it: a replica with no stored state would
+        # be skipped, leaving the sampler holding the water configuration of
+        # whichever replica used the slot last.
+        for replica in range(self._num_replicas):
+            slot = self._replica_slot[replica]
+
+            if self._openmm_states[replica] is None:
+                self.save_openmm_state(slot, replica)
+
+            if self._gcmc_samplers[slot] is not None and (
+                self._gcmc_states[replica] is None
+            ):
+                self.save_gcmc_state(slot, replica)
+
+    def slot_for(self, replica):
         """
-        Get the dynamics object (and GCMC sampler) for a given index.
+        Return the index of the slot that hosts a given replica.
 
         Parameters
         ----------
 
-        index: int
+        replica: int
             The index of the replica.
 
         Returns
         -------
 
+        int
+            The index of the slot.
+        """
+        return self._replica_slot[replica]
+
+    def get(self, slot):
+        """
+        Get the dynamics object (and GCMC sampler) for a given slot.
+
+        When there is one slot per replica the slot and replica indices are
+        the same. Otherwise use slot_for() to map a replica to its slot.
+
+        Parameters
+        ----------
+
+        slot: int
+            The index of the slot.
+
+        Returns
+        -------
+
         tuple
-            The dynamics object for the replica and its GCMC sampler.
+            The dynamics object for the slot and its GCMC sampler.
         """
-        return self._dynamics[index], self._gcmc_samplers[index]
+        return self._dynamics[slot], self._gcmc_samplers[slot]
 
-    def set(self, index, dynamics):
+    def set(self, slot, dynamics):
         """
-        Set the dynamics object for a given index.
+        Set the dynamics object for a given slot.
 
         Parameters
         ----------
 
-        index: int
-            The index of the replica.
+        slot: int
+            The index of the slot.
 
         dynamics: sire.legacy.Convert.SOMMContext
             The dynamics object.
         """
-        self._dynamics[index] = dynamics
+        self._dynamics[slot] = dynamics
 
-    def delete(self, index):
+    def delete(self, slot):
         """
-        Delete the dynamics object for a given index.
+        Delete the dynamics object for a given slot.
 
         Parameters
         ----------
 
-        index: int
-            The index of the replica.
+        slot: int
+            The index of the slot.
         """
-        self._dynamics[index] = None
+        self._dynamics[slot] = None
 
-    def save_openmm_state(self, index):
+    def save_openmm_state(self, slot, replica):
         """
-        Save the state of the dynamics object.
+        Save the state of a slot's dynamics object as the state of a replica.
 
         Parameters
         ----------
 
-        index: int
-            The index of the replica.
+        slot: int
+            The index of the slot.
+
+        replica: int
+            The index of the replica whose state this is.
         """
 
         # Get the current OpenMM state.
         state = (
-            self._dynamics[index]
+            self._dynamics[slot]
             .context()
             .getState(getPositions=True, getVelocities=True)
         )
@@ -490,7 +658,7 @@ def save_openmm_state(self, index):
         # Store positions, velocities, and box vectors as compact numpy arrays
         # rather than the OpenMM State object, which serialises to XML when
         # pickled and is orders of magnitude larger.
-        self._openmm_states[index] = {
+        self._openmm_states[replica] = {
             "positions": state.getPositions(asNumpy=True),
             "velocities": state.getVelocities(asNumpy=True),
             "box": state.getPeriodicBoxVectors(asNumpy=True),
@@ -522,21 +690,39 @@ def _apply_openmm_state(context, state):
             # format change.
             context.setState(state)
 
-    def save_gcmc_state(self, index):
+    def save_gcmc_state(self, slot, replica):
         """
-        Save the current GCMC water state for the replica.
+        Save the current GCMC water state of a slot as that of a replica.
 
         Parameters
         ----------
 
-        index: int
-            The index of the replica.
+        slot: int
+            The index of the slot.
+
+        replica: int
+            The index of the replica whose state this is.
         """
         # Get the GCMC sampler.
-        gcmc_sampler = self._gcmc_samplers[index]
+        gcmc_sampler = self._gcmc_samplers[slot]
 
         # Store the state.
-        self._gcmc_states[index] = gcmc_sampler.water_state()
+        self._gcmc_states[replica] = gcmc_sampler.water_state()
+
+    def get_clock(self):
+        """
+        Get the simulation clock.
+
+        Every replica advances by the same amount each cycle, so the clock is
+        common to all of them and is read from the first slot.
+
+        Returns
+        -------
+
+        dict
+            The clock, as returned by Dynamics._get_clock().
+        """
+        return self._dynamics[0]._get_clock()
 
     def get_states(self):
         """
@@ -562,66 +748,161 @@ def set_states(self, states):
         """
         self._states = states
 
-    def _seed_replica(self, i):
+    def load_replica(self, replica, clock=None):
         """
-        Apply the (possibly new, post-mix) state to replica i's context,
-        including any GCMC water-state swap. Only touches replica i's own
-        context/sampler, so this is safe to run concurrently with other
-        replicas' calls from a thread pool.
+        Make a replica resident in its slot, ready for a dynamics block.
+
+        This sets the slot's lambda value, pushes the replica's state into the
+        OpenMM context, and points the slot at the replica's clock, energy
+        trajectory, GCMC water state and ghost file.
+
+        Work that isn't needed is skipped. When there is one slot per replica
+        and no swap has taken place, the slot already holds everything the
+        replica needs and this reduces to a handful of comparisons.
+
+        Only touches the replica's own slot, so this is safe to call
+        concurrently for replicas in different slots.
+
+        Parameters
+        ----------
+
+        replica: int
+            The index of the replica.
+
+        clock: dict
+            The simulation clock to restore, as returned by
+            Dynamics._get_clock(). If None, the slot's clock is left alone.
         """
-        state = self._states[i]
+        slot = self._replica_slot[replica]
+        dynamics = self._dynamics[slot]
+        gcmc_sampler = self._gcmc_samplers[slot]
 
-        _logger.debug(f"Replica {i} seeded from state {state}")
-        self._apply_openmm_state(
-            self._dynamics[i].context(), self._openmm_states[state]
-        )
+        # The replica that the slot currently holds.
+        resident = self._slot_replica[slot]
+
+        # The slot needs new positions and velocities if it is being handed a
+        # different replica, or if the last mix moved this replica's state.
+        if resident != replica or self._state_moved[replica]:
+            self._apply_openmm_state(dynamics.context(), self._openmm_states[replica])
+
+            # Positions have changed underneath the context, so any cached
+            # energies are stale.
+            dynamics.clear_energy_cache()
 
-        # Swap the water state in the GCMCSamplers.
-        if self._gcmc_samplers[i] is not None:
-            # Find the indices of the water states that differ.
-            water_idxs = _np.where(self._gcmc_states[i] != self._gcmc_states[state])[0]
+            self._state_moved[replica] = False
+
+        # Set the lambda value and REST2 scaling factor. This is a no-op if
+        # the slot is already at this lambda value.
+        if resident != replica:
+            dynamics.set_lambda(
+                self._lambdas[replica],
+                rest2_scale=self._rest2_scale_factors[replica],
+                update_constraints=self._update_constraints,
+            )
 
-            # Update the water state in the GCMCSampler.
-            self._gcmc_samplers[i].push()
+        # Restore the clock and point the slot at this replica's energy
+        # trajectory, so that energies are accumulated against the replica
+        # rather than the slot.
+        if clock is not None:
+            dynamics._set_clock(clock)
+        dynamics.set_energy_trajectory(self._energy_trajectories[replica])
+
+        if gcmc_sampler is not None:
+            gcmc_sampler.push()
             try:
-                self._gcmc_samplers[i]._set_water_state(
-                    self._dynamics[i].context(),
-                    indices=water_idxs,
-                    states=self._gcmc_states[state][water_idxs],
+                # Swap the water state into the sampler. Diff against what the
+                # sampler currently holds, which is the state of whichever
+                # replica was last resident.
+                target = self._gcmc_states[replica]
+                if target is not None:
+                    current = gcmc_sampler.water_state()
+                    water_idxs = _np.where(current != target)[0]
+
+                    if len(water_idxs) > 0:
+                        gcmc_sampler._set_water_state(
+                            dynamics.context(),
+                            indices=water_idxs,
+                            states=target[water_idxs],
+                        )
+
+                if resident != replica:
+                    # Update the lambda dependent non-bonded parameters used to
+                    # evaluate insertion and deletion energies, and append ghost
+                    # residues to this replica's file. The sampler keeps its
+                    # statistics per lambda value, so switching also switches
+                    # to this replica's.
+                    gcmc_sampler.set_lambda(
+                        self._lambdas[replica], self._rest2_scale_factors[replica]
+                    )
+                    gcmc_sampler.set_ghost_file(self._ghost_files[replica])
+            finally:
+                gcmc_sampler.pop()
+
+        self._slot_replica[slot] = replica
+
+    def store_replica(self, replica):
+        """
+        Save the state of a replica back out of its slot, so that the slot can
+        be handed to another replica.
+
+        Parameters
+        ----------
+
+        replica: int
+            The index of the replica.
+        """
+        slot = self._replica_slot[replica]
+
+        self.save_openmm_state(slot, replica)
+
+        if self._gcmc_samplers[slot] is not None:
+            self.save_gcmc_state(slot, replica)
+
+            # Count the waters against the slot's context rather than whichever
+            # one the sampler happens to be bound to, which is unset after the
+            # sampler has been reset.
+            gcmc_sampler = self._gcmc_samplers[slot]
+            gcmc_sampler.push()
+            try:
+                self._gcmc_num_waters[replica] = gcmc_sampler.num_waters(
+                    context=self._dynamics[slot].context()
                 )
             finally:
-                self._gcmc_samplers[i].pop()
+                gcmc_sampler.pop()
 
-    def mix_states(self, old_states, executor=None):
+    def mix_states(self, old_states):
         """
-        Mix the states of the dynamics objects.
+        Apply the result of a replica mix.
+
+        The states are permuted here, but not pushed into the OpenMM contexts.
+        They are applied lazily by load_replica(), which is the only point at
+        which a slot is known to be free. This does the same amount of work as
+        applying them eagerly, since load_replica() pushes a state exactly when
+        the mix moved it.
+
+        The permutation must happen here rather than being resolved lazily
+        through self._states. A slot is re-used within a cycle, so a replica
+        may be loaded after another replica has already stored its post-run
+        state; reading through the indirection at that point would pick up the
+        new state rather than the pre-mix one.
 
         Parameters
         ----------
+
         old_states : numpy.ndarray
             The state indices from before the last replica mix.
+        """
+        # Permute the travelling state. This is a reference shuffle, so it is
+        # cheap even for large systems. Statistics and output files stay with
+        # the lambda window, so are not permuted.
+        self._openmm_states = [self._openmm_states[state] for state in self._states]
+        self._gcmc_states = [self._gcmc_states[state] for state in self._states]
 
-        executor : concurrent.futures.ThreadPoolExecutor, optional
-            Executor used to apply the per-replica state changes (an
-            OpenMM setPositions/setVelocities/setPeriodicBoxVectors call
-            per changed replica, each against a different context) in
-            parallel. Each replica's context is independent of every
-            other, so this is safe. Falls back to a serial loop if not
-            provided.
-        """
-        # Replicas whose state actually changed.
-        changed = [i for i, state in enumerate(self._states) if i != state]
-
-        if executor is not None and len(changed) > 1:
-            # Consume the map so we block until every replica is seeded
-            # and any exception raised in a worker thread propagates here.
-            list(executor.map(self._seed_replica, changed))
-        else:
-            for i in changed:
-                self._seed_replica(i)
+        # Flag the replicas whose state moved, so that load_replica() knows it
+        # has to push new positions and velocities into the context.
+        self._state_moved = [bool(state != i) for i, state in enumerate(self._states)]
 
-        # Update the swap matrix. Cheap, CPU-only bookkeeping - kept out of
-        # the parallel section above to avoid any shared-array races.
+        # Update the swap matrix.
         for i, state in enumerate(self._states):
             self._num_swaps[old_states[i], state] += 1
 
@@ -836,6 +1117,12 @@ def __init__(self, system, config):
         # devices relative to the visible set, so index i is gpu_devices[i].
         self._gpu_devices = list(gpu_devices)[: self._num_gpus]
 
+        # Work out how many OpenMM contexts (slots) to create. When there are
+        # fewer slots than replicas, each slot is re-used to propagate several
+        # replicas per cycle, changing its lambda value as it goes.
+        self._num_replicas = len(self._lambda_values)
+        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,
         # since the per-replica OpenMM contexts it creates are fixed at
@@ -923,6 +1210,8 @@ def __init__(self, system, config):
                 perturbed_system=self._perturbed_system,
                 output_directory=self._config.output_directory,
                 xml_filenames=xml_filenames,
+                num_slots=self._num_slots,
+                update_constraints=self._config.update_constraints,
                 gpu_devices=self._gpu_devices,
             )
 
@@ -998,6 +1287,13 @@ def __init__(self, system, config):
             # restarted against a different set of GPUs.
             self._dynamics_cache._gpu_devices = self._gpu_devices
 
+            # Rebuild the slot layout from the current config, so that
+            # 'max_contexts' can change on restart. Everything that is restored
+            # is per-replica, so it doesn't depend on the grouping.
+            self._dynamics_cache._num_slots = self._num_slots
+            self._dynamics_cache._update_constraints = self._config.update_constraints
+            self._dynamics_cache._build_slot_layout()
+
             # Create the dynamics objects.
             self._dynamics_cache._create_dynamics(
                 self._system,
@@ -1009,46 +1305,30 @@ def __init__(self, system, config):
                 output_directory=self._config.output_directory,
             )
 
-            # Reset the state of the OpenMM contexts and GCMC samplers.
-            for i in range(len(self._lambda_values)):
-                dynamics, gcmc_sampler = self._dynamics_cache.get(i)
-
-                # Reset the OpenMM state, applying the last replica exchange
-                # mixing so the correct post-mix state is restored.
-                state = self._dynamics_cache._states[i]
-                DynamicsCache._apply_openmm_state(
-                    dynamics.context(), self._dynamics_cache._openmm_states[state]
-                )
-
-                # Reset the GCMC water state and restore statistics.
-                if gcmc_sampler is not None:
-                    gcmc_sampler.push()
-                    try:
-                        gcmc_sampler._set_water_state(
-                            dynamics.context(),
-                            states=self._dynamics_cache._gcmc_states[state],
-                            force=True,
-                        )
-                    finally:
-                        gcmc_sampler.pop()
-
-                    # Samplers keep only the lambda values they visit, so it's
-                    # safe to hand each of them the whole simulation's stats.
-                    if self._dynamics_cache._gcmc_stats is not None:
-                        gcmc_sampler.restore_stats(self._dynamics_cache._gcmc_stats)
-
-        # Log the GCMC sphere centre for each replica using the actual context
-        # positions (accurate for both fresh runs and restarts).
+            # The OpenMM contexts are not reset here. Each replica's state is
+            # pushed into its slot by load_replica() at the start of its first
+            # block, which is the only point at which the slot is known to be
+            # free.
+
+            # Restore the sampling statistics. A sampler keeps only the lambda
+            # values it visits, so each can be handed the whole simulation's.
+            if self._dynamics_cache._gcmc_stats is not None:
+                for slot in range(self._dynamics_cache._num_slots):
+                    _, gcmc_sampler = self._dynamics_cache.get(slot)
+                    gcmc_sampler.restore_stats(self._dynamics_cache._gcmc_stats)
+
+        # Log the GCMC sphere centre for each replica. This uses the stored
+        # state rather than the context, since a slot only holds the positions
+        # of the replica it last hosted.
         import openmm.unit as _omm_unit
 
         for i, lam in enumerate(self._lambda_values):
-            dynamics, gcmc_sampler = self._dynamics_cache.get(i)
+            _, gcmc_sampler = self._dynamics_cache.get(self._dynamics_cache.slot_for(i))
             if gcmc_sampler is not None and gcmc_sampler._reference is not None:
-                state = dynamics.context().getState(getPositions=True)
-                positions = state.getPositions(asNumpy=True).value_in_unit(
-                    _omm_unit.angstrom
+                positions = self._dynamics_cache._openmm_states[i]["positions"]
+                target = gcmc_sampler._get_target_position(
+                    positions.value_in_unit(_omm_unit.angstrom)
                 )
-                target = gcmc_sampler._get_target_position(positions)
                 _logger.info(
                     f"Initial GCMC sphere centre for lambda {lam:.5f}: "
                     f"[{target[0]:.3f}, {target[1]:.3f}, {target[2]:.3f}] A"
@@ -1109,6 +1389,158 @@ def __init__(self, system, config):
         # Create a lock to guard the dynamics cache.
         self._lock = Lock()
 
+        # Systems committed while a replica was resident in its slot, awaiting
+        # the end of cycle checkpoint. Only used when contexts are shared,
+        # keyed by replica index and emptied by _checkpoint().
+        self._committed = {}
+
+    def _replica_passes(self, cycle):
+        """
+        Work out which replicas to propagate in each pass of a cycle.
+
+        A slot can only host one replica at a time, so each pass takes at most
+        one replica from each slot. Groups are traversed in alternating
+        directions on successive cycles, so that a slot always moves to a
+        neighbouring lambda window, including across the cycle boundary. That
+        keeps the change in force field parameters as small as possible.
+
+        When there is one slot per replica there is a single pass containing
+        every replica, which is the same as propagating them all together.
+
+        Parameters
+        ----------
+
+        cycle: int
+            The index of the current cycle.
+
+        Returns
+        -------
+
+        list of list of int
+            The replicas to propagate in each pass.
+        """
+        groups = self._dynamics_cache._groups
+
+        # Traverse the groups backwards on odd cycles.
+        is_reversed = cycle % 2 == 1
+
+        passes = []
+        for i in range(max(len(group) for group in groups)):
+            batch = []
+            for group in groups:
+                if i < len(group):
+                    batch.append(group[len(group) - 1 - i] if is_reversed else group[i])
+            passes.append(batch)
+
+        return passes
+
+    def _safe_batches(self, num_workers, cycle=0):
+        """
+        Yield batches of replicas that can be processed concurrently.
+
+        Two replicas that share a slot must never be processed at the same
+        time, since they would be using the same dynamics object and OpenMM
+        context. Batches are therefore taken from within a single pass, which
+        holds at most one replica per slot, and then split by the number of
+        workers.
+
+        With one slot per replica there is a single pass containing every
+        replica, so this is just a split by the number of workers.
+
+        Parameters
+        ----------
+
+        num_workers: int
+            The maximum number of replicas in a batch.
+
+        cycle: int
+            The index of the current cycle, which sets the traversal order.
+
+        Yields
+        ------
+
+        list of int
+            A batch of replicas that is safe to process concurrently.
+        """
+        from math import ceil
+
+        for batch in self._replica_passes(cycle):
+            for i in range(ceil(len(batch) / num_workers)):
+                yield batch[i * num_workers : (i + 1) * num_workers]
+
+    def _set_num_slots(self):
+        """
+        Work out the number of OpenMM contexts (slots) to create, validating
+        the configuration options that only apply when contexts are re-used
+        across lambda values.
+
+        Sets self._num_slots and self._is_cached.
+        """
+
+        num_replicas = self._num_replicas
+
+        if self._config.max_contexts is None:
+            self._num_slots = num_replicas
+        else:
+            self._num_slots = min(self._config.max_contexts, num_replicas)
+
+        # There is a context per replica, so nothing is re-used and all of the
+        # constraints below are irrelevant.
+        self._is_cached = self._num_slots < num_replicas
+
+        if not self._is_cached:
+            if self._config.max_contexts is not None:
+                _logger.info(
+                    f"Creating one OpenMM context per replica ({num_replicas})"
+                )
+            return
+
+        # Frames can only be saved on checkpoint cycles when contexts are
+        # re-used. Within a cycle a context propagates several replicas in
+        # turn, so frames from different replicas would otherwise accumulate
+        # in the same internal trajectory. Tying frames to checkpoints means
+        # each one is written out and cleared before the context is handed to
+        # the next replica.
+        if (
+            self._save_frames
+            and self._config.frame_frequency != self._config.checkpoint_frequency
+        ):
+            msg = (
+                "'frame_frequency' must equal 'checkpoint_frequency' when "
+                "'max_contexts' is less than the number of replicas."
+            )
+            _logger.error(msg)
+            raise ValueError(msg)
+
+        num_workers = self._num_gpus * self._config.oversubscription_factor
+
+        if self._num_slots < num_workers:
+            _logger.warning(
+                f"'max_contexts' ({self._num_slots}) is less than the number of "
+                f"workers ({num_workers}). Some GPUs will be left idle."
+            )
+        elif self._num_slots % self._num_gpus != 0:
+            _logger.warning(
+                f"'max_contexts' ({self._num_slots}) is not a multiple of the "
+                f"number of GPUs ({self._num_gpus}). This may result in "
+                "suboptimal performance."
+            )
+
+        if not self._config.update_constraints:
+            _logger.warning(
+                "'update_constraints' is False. Constrained bond lengths will not "
+                "perturb with lambda, and are frozen at those of the lambda value "
+                "each context was created at."
+            )
+
+        from math import ceil
+
+        _logger.info(
+            f"Re-using {self._num_slots} OpenMM context(s) across "
+            f"{num_replicas} replicas: {ceil(num_replicas / self._num_slots)} "
+            "pass(es) per cycle"
+        )
+
     def __str__(self):
         """Return a string representation of the object."""
         return f"RepexRunner(system={self._system}, config={self._config})"
@@ -1185,7 +1617,6 @@ def run(self):
             )
 
         # Work out the required number of batches.
-        num_batches = ceil(self._config.num_lambda / num_workers)
         num_checkpoint_batches = ceil(self._config.num_lambda / num_checkpoint_workers)
 
         # Persistent thread pools, reused across every batch and cycle.
@@ -1197,12 +1628,12 @@ def run(self):
 
         # Minimise at each lambda value.
         if self._config.minimise:
-            for i in range(num_batches):
+            for batch in self._safe_batches(num_workers):
                 with ThreadPoolExecutor(max_workers=num_workers) as executor:
                     try:
                         for success, index, e in executor.map(
                             self._minimise,
-                            replica_list[i * num_workers : (i + 1) * num_workers],
+                            batch,
                         ):
                             if not success:
                                 msg = f"Minimisation failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {e}"
@@ -1217,12 +1648,12 @@ def run(self):
 
         # Equilibrate the system.
         if self._is_equilibration and not self._is_restart:
-            for i in range(num_batches):
+            for batch in self._safe_batches(num_workers):
                 with ThreadPoolExecutor(max_workers=num_workers) as executor:
                     try:
                         for success, index, e in executor.map(
                             self._equilibrate,
-                            replica_list[i * num_workers : (i + 1) * num_workers],
+                            batch,
                         ):
                             if not success:
                                 _logger.error(
@@ -1236,6 +1667,28 @@ def run(self):
         # Write a checkpoint immediately after equilibration so that a restart
         # after an early production crash doesn't need to re-equilibrate.
         if self._is_equilibration and not self._is_restart:
+            # When contexts are shared, commit each replica while it is
+            # resident so that every checkpoint file can still be written
+            # under a single lock below.
+            if self._is_cached:
+                for batch in self._safe_batches(num_checkpoint_workers):
+                    try:
+                        for index, error in checkpoint_executor.map(
+                            self._load_and_commit,
+                            batch,
+                        ):
+                            if error is not None:
+                                msg = (
+                                    f"Post-equilibration commit failed for {_lam_sym} = "
+                                    f"{self._lambda_values[index]:.5f}:\n{error}"
+                                )
+                                _logger.error(msg)
+                                raise error
+                    except KeyboardInterrupt:
+                        checkpoint_executor.shutdown(wait=False, cancel_futures=True)
+                        _logger.error("Post-equilibration commit cancelled. Exiting.")
+                        _sys.exit(1)
+
             lock = _FileLock(self._lock_file)
             with lock.acquire(timeout=self._config.timeout.to("seconds")):
                 for j in range(num_checkpoint_batches):
@@ -1346,33 +1799,80 @@ def run(self):
                 else None
             )
 
-            # Run a dynamics block for each replica, making sure only each GPU is only
-            # oversubscribed by a factor of self._config.oversubscription_factor.
-            for j in range(num_batches):
-                replicas = replica_list[j * num_workers : (j + 1) * num_workers]
-                try:
-                    for result, index, energies in dynamics_executor.map(
-                        self._run_block,
-                        replicas,
-                        repeat(self._lambda_values),
-                        repeat(is_gcmc),
-                        repeat(write_gcmc_ghosts),
-                        repeat(is_terminal_flip),
-                        repeat(time_ns),
-                    ):
-                        if not result:
-                            _logger.error(
-                                f"Dynamics failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {energies}"
+            # Whether the checkpoint files are written at the end of this cycle.
+            do_checkpoint = is_checkpoint or i == cycles - 1
+
+            # Capture the simulation clock at the start of the cycle. Every
+            # replica advances by the same amount each cycle, so a slot has to
+            # be rewound to this point before it propagates the next replica.
+            clock = self._dynamics_cache.get_clock()
+
+            # Propagate the replicas, one pass at a time. Each pass takes at
+            # most one replica from each slot, and there is a single pass when
+            # there is a slot per replica.
+            for batch in self._replica_passes(i):
+                # Run a dynamics block for each replica in the pass, making sure
+                # each GPU is only oversubscribed by a factor of
+                # self._config.oversubscription_factor.
+                for j in range(ceil(len(batch) / num_workers)):
+                    replicas = batch[j * num_workers : (j + 1) * num_workers]
+                    try:
+                        for result, index, energies in dynamics_executor.map(
+                            self._run_block,
+                            replicas,
+                            repeat(self._lambda_values),
+                            repeat(is_gcmc),
+                            repeat(write_gcmc_ghosts),
+                            repeat(is_terminal_flip),
+                            repeat(time_ns),
+                            repeat(clock),
+                        ):
+                            if not result:
+                                _logger.error(
+                                    f"Dynamics failed for {_lam_sym} = {self._lambda_values[index]:.5f}: {energies}"
+                                )
+                                raise energies
+                            results.append((index, energies))
+                    except KeyboardInterrupt:
+                        dynamics_executor.shutdown(wait=False, cancel_futures=True)
+                        _logger.error("Dynamics cancelled. Exiting.")
+                        _sys.exit(1)
+
+                # When contexts are shared, commit the replicas of this pass
+                # while they are still resident in their slots. The committed
+                # systems are held until the end of the cycle, so that every
+                # checkpoint file is still written under a single lock. With a
+                # context per replica everything is still resident once the
+                # cycle finishes, so committing is left to _checkpoint().
+                if do_checkpoint and self._is_cached:
+                    for j in range(ceil(len(batch) / num_checkpoint_workers)):
+                        replicas = batch[
+                            j * num_checkpoint_workers : (j + 1)
+                            * num_checkpoint_workers
+                        ]
+                        try:
+                            for index, error in checkpoint_executor.map(
+                                self._commit_replica,
+                                replicas,
+                            ):
+                                if error:
+                                    _logger.error(
+                                        f"Commit failed for {_lam_sym} = "
+                                        f"{self._lambda_values[index]:.5f}: {error}"
+                                    )
+                                    raise error
+                        except KeyboardInterrupt:
+                            checkpoint_executor.shutdown(
+                                wait=False, cancel_futures=True
                             )
-                            raise energies
-                        results.append((index, energies))
-                except KeyboardInterrupt:
-                    dynamics_executor.shutdown(wait=False, cancel_futures=True)
-                    _logger.error("Dynamics cancelled. Exiting.")
-                    _sys.exit(1)
-
-            # Checkpoint.
-            if is_checkpoint or i == cycles - 1:
+                            _logger.error("Commit cancelled. Exiting.")
+                            _sys.exit(1)
+
+            # Checkpoint. This happens once the whole cycle is complete, with
+            # every checkpoint file written under a single lock, so that an
+            # external process reading the output directory always sees a
+            # coherent set rather than a mixture of new and old files.
+            if do_checkpoint:
                 # Create the lock.
                 lock = _FileLock(self._lock_file)
 
@@ -1391,7 +1891,7 @@ def run(self):
                                 self._backup_checkpoint,
                                 replicas,
                             ):
-                                if not result:
+                                if error:
                                     _logger.error(
                                         f"Backup failed for {_lam_sym} = "
                                         f"{self._lambda_values[index]:.5f}: {error}"
@@ -1448,16 +1948,11 @@ def run(self):
                     self._dynamics_cache.get_accepted(),
                 )
             )
-            self._dynamics_cache.mix_states(old_states, executor=dynamics_executor)
-
-            # Snapshot the pre-run state for crash recovery.
-            if self._config.auto_fix_minimise:
-                for i, state in enumerate(self._dynamics_cache.get_states()):
-                    self._dynamics_cache._dynamics[i]._d._pre_run_state = (
-                        self._dynamics_cache._dynamics[i]
-                        .context()
-                        .getState(getPositions=True, getVelocities=True)
-                    )
+
+            # This only permutes the stored states. They are pushed into the
+            # contexts by load_replica() at the start of the next block, which
+            # is also where the pre-run state for crash recovery is captured.
+            self._dynamics_cache.mix_states(old_states)
 
             # This is a checkpoint cycle.
             if is_checkpoint:
@@ -1548,10 +2043,16 @@ def _run_block(
         write_gcmc_ghosts=False,
         is_terminal_flip=False,
         time_ns=None,
+        clock=None,
     ):
         """
         Run a dynamics block for a given replica.
 
+        The replica is made resident in its slot, propagated, then stored back
+        out again so that the slot can be handed to the next replica. When
+        there is one slot per replica, loading and storing reduce to the state
+        bookkeeping that replica exchange does anyway.
+
         Parameters
         ----------
 
@@ -1579,6 +2080,12 @@ def _run_block(
             The current simulation time in nanoseconds, used when saving energy
             components. If None, energy components are not saved.
 
+        clock: dict
+            The simulation clock at the start of the cycle, as returned by
+            Dynamics._get_clock(). Every replica advances by the same amount
+            each cycle, so a slot must be rewound to the start of the cycle
+            before propagating the next replica.
+
         Returns
         -------
 
@@ -1597,13 +2104,14 @@ def _run_block(
         lam = lambdas[index]
 
         try:
+            # Make the replica resident in its slot.
+            self._dynamics_cache.load_replica(index, clock=clock)
+
             # Get the dynamics object (and GCMC sampler).
-            dynamics, gcmc_sampler = self._dynamics_cache.get(index)
+            dynamics, gcmc_sampler = self._dynamics_cache.get(
+                self._dynamics_cache.slot_for(index)
+            )
 
-            # Track whether any MC move changed the context positions so we
-            # can update _pre_run_state once at the end. Only needed when
-            # crash recovery is enabled.
-            needs_pre_run_snapshot = False
             auto_fix_minimise = self._config.auto_fix_minimise
 
             # Perform the GCMC move before dynamics so that the energies
@@ -1617,9 +2125,6 @@ def _run_block(
                 finally:
                     gcmc_sampler.pop()
 
-                if auto_fix_minimise:
-                    needs_pre_run_snapshot = True
-
                 # Write ghost residues immediately after the GCMC move so the
                 # ghost state and frame (saved during dynamics) are consistent.
                 if write_gcmc_ghosts:
@@ -1628,13 +2133,12 @@ def _run_block(
             # Perform a terminal flip move before dynamics if requested.
             if self._terminal_flip_samplers is not None and is_terminal_flip:
                 _logger.info(f"Performing terminal flip move at {_lam_sym} = {lam:.5f}")
-                if self._terminal_flip_samplers[index].move(dynamics.context()):
-                    if auto_fix_minimise:
-                        needs_pre_run_snapshot = True
+                self._terminal_flip_samplers[index].move(dynamics.context())
 
-            # Snapshot the context state for crash recovery if any MC move
-            # changed positions.
-            if needs_pre_run_snapshot:
+            # Snapshot the context state for crash recovery. The slot was just
+            # seeded with this replica's state, and any MC move above may have
+            # changed it again, so this is always required.
+            if auto_fix_minimise:
                 dynamics._d._pre_run_state = dynamics.context().getState(
                     getPositions=True, getVelocities=True
                 )
@@ -1668,12 +2172,9 @@ def _run_block(
                 ),
             )
 
-            if gcmc_sampler is not None:
-                # Save the GCMC state.
-                self._dynamics_cache.save_gcmc_state(index)
-
-            # Save the OpenMM state.
-            self._dynamics_cache.save_openmm_state(index)
+            # Save the replica's state back out of the slot, so that the slot
+            # can be handed to the next replica.
+            self._dynamics_cache.store_replica(index)
 
             # Save the energy contribution for each force.
             if self._config.save_energy_components and time_ns is not None:
@@ -1717,8 +2218,12 @@ def _minimise(self, index):
         _logger.info(f"Minimising at {_lam_sym} = {self._lambda_values[index]:.5f}")
 
         try:
+            # Make the replica resident in its slot.
+            self._dynamics_cache.load_replica(index)
+
             # Get the dynamics object (and GCMC sampler).
-            dynamics, gcmc_sampler = self._dynamics_cache.get(index)
+            slot = self._dynamics_cache.slot_for(index)
+            dynamics, gcmc_sampler = self._dynamics_cache.get(slot)
 
             if gcmc_sampler is not None and not self._is_restart:
                 gcmc_sampler.push()
@@ -1749,10 +2254,10 @@ def _minimise(self, index):
                     system = dynamics.commit()
 
                     # Delete the dynamics object.
-                    self._dynamics_cache.delete(index)
+                    self._dynamics_cache.delete(slot)
 
                     # Work out the device index.
-                    device = index % self._num_gpus
+                    device = slot % self._num_gpus
 
                     # Copy the dynamics keyword arguments.
                     dynamics_kwargs = self._dynamics_kwargs.copy()
@@ -1771,13 +2276,20 @@ def _minimise(self, index):
                     if gcmc_sampler is not None:
                         self._reset_gcmc_sampler(gcmc_sampler, dynamics)
 
-                    # Set the new dynamics object.
-                    self._dynamics_cache.set(index, dynamics)
+                    # Set the new dynamics object. The rebuilt object has a
+                    # fresh, empty energy trajectory, but no energies are
+                    # recorded during minimisation or equilibration, and
+                    # load_replica() re-attaches the replica's trajectory
+                    # before the first production block.
+                    self._dynamics_cache.set(slot, dynamics)
 
                     _logger.info(
                         f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}"
                     )
 
+            # Save the minimised state back out of the slot.
+            self._dynamics_cache.store_replica(index)
+
         except Exception as e:
             return False, index, e
 
@@ -1808,8 +2320,12 @@ def _equilibrate(self, index):
         _logger.info(f"Equilibrating at {_lam_sym} = {self._lambda_values[index]:.5f}")
 
         try:
+            # Make the replica resident in its slot.
+            self._dynamics_cache.load_replica(index)
+
             # Get the dynamics object (and GCMC sampler).
-            dynamics, gcmc_sampler = self._dynamics_cache.get(index)
+            slot = self._dynamics_cache.slot_for(index)
+            dynamics, gcmc_sampler = self._dynamics_cache.get(slot)
 
             if gcmc_sampler is not None:
                 gcmc_sampler.push()
@@ -1848,10 +2364,10 @@ def _equilibrate(self, index):
                     system = dynamics.commit()
 
                     # Delete the current dynamics object.
-                    self._dynamics_cache.delete(index)
+                    self._dynamics_cache.delete(slot)
 
                     # Work out the device index.
-                    device = index % self._num_gpus
+                    device = slot % self._num_gpus
 
                     # Copy the dynamics keyword arguments.
                     dynamics_kwargs = self._dynamics_kwargs.copy()
@@ -1872,7 +2388,7 @@ def _equilibrate(self, index):
                         self._reset_gcmc_sampler(gcmc_sampler, dynamics)
 
                     # Update the dynamics object in the cache.
-                    self._dynamics_cache.set(index, dynamics)
+                    self._dynamics_cache.set(slot, dynamics)
 
             # Equilibrate.
             dynamics.run(
@@ -1894,10 +2410,10 @@ def _equilibrate(self, index):
                 system.set_time(_sr.u("0ps"))
 
             # Delete the dynamics object.
-            self._dynamics_cache.delete(index)
+            self._dynamics_cache.delete(slot)
 
             # Work out the device index.
-            device = index % self._num_gpus
+            device = slot % self._num_gpus
 
             # Copy the dynamics keyword arguments.
             dynamics_kwargs = self._dynamics_kwargs.copy()
@@ -1925,12 +2441,15 @@ def _equilibrate(self, index):
                     gcmc_sampler.pop()
 
             # Set the new dynamics object.
-            self._dynamics_cache.set(index, dynamics)
+            self._dynamics_cache.set(slot, dynamics)
 
             _logger.info(
                 f"Created dynamics object for {_lam_sym} = {self._lambda_values[index]:.5f}"
             )
 
+            # Save the equilibrated state back out of the slot.
+            self._dynamics_cache.store_replica(index)
+
         except Exception as e:
             return False, index, e
 
@@ -1966,7 +2485,7 @@ def _compute_energies(self, index):
         )
 
         # Get the dynamics object.
-        dynamics, _ = self._dynamics_cache.get(index)
+        dynamics, _ = self._dynamics_cache.get(self._dynamics_cache.slot_for(index))
 
         # Create an array to hold the energies.
         energies = _np.zeros(self._config.num_lambda)
@@ -2044,6 +2563,80 @@ def _write_checkpoint_system(self, system, index, context=None, gcmc_sampler=Non
         """
         self._dynamics_cache._time = system.time()
 
+    def _load_and_commit(self, index):
+        """
+        Make a replica resident in its slot and commit it, for the
+        post-equilibration checkpoint.
+
+        Equilibration stores every replica's state back out of its slot, so a
+        replica has to be loaded again before it can be committed.
+
+        Parameters
+        ----------
+
+        index: int
+            The index of the replica.
+
+        Returns
+        -------
+
+        index: int
+            The index of the replica.
+
+        exception: Exception
+            The exception if the commit failed.
+        """
+        try:
+            self._dynamics_cache.load_replica(index)
+        except Exception as e:
+            return index, e
+
+        return self._commit_replica(index)
+
+    def _commit_replica(self, index):
+        """
+        Commit a replica while it is still resident in its slot, holding the
+        result until the checkpoint files are written at the end of the cycle.
+
+        This is only used when contexts are shared between replicas. A slot
+        holds the positions and trajectory frames of the replica it last
+        hosted, so a replica has to be committed before its slot is handed on.
+        Deferring only the file writes keeps every checkpoint file within a
+        single lock, so that an external reader never sees a partial set.
+
+        Parameters
+        ----------
+
+        index: int
+            The index of the replica.
+
+        Returns
+        -------
+
+        index: int
+            The index of the replica.
+
+        exception: Exception
+            The exception if the commit failed.
+        """
+        try:
+            slot = self._dynamics_cache.slot_for(index)
+            dynamics, _ = self._dynamics_cache.get(slot)
+
+            # commit() returns a clone, so the frames can be cleared straight
+            # away, ready for the next replica to use the slot.
+            system = dynamics.commit()
+            speed = dynamics.time_speed()
+            dynamics._d._sire_mols.delete_all_frames()
+
+            with self._lock:
+                self._committed[index] = (system, speed)
+
+        except Exception as e:
+            return index, e
+
+        return index, None
+
     def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False):
         """
         Checkpoint the simulation.
@@ -2080,13 +2673,23 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False):
             lam = lambdas[index]
 
             # Get the dynamics object (and GCMC sampler).
-            dynamics, gcmc_sampler = self._dynamics_cache.get(index)
+            slot = self._dynamics_cache.slot_for(index)
+            dynamics, gcmc_sampler = self._dynamics_cache.get(slot)
 
-            # Commit the current system.
-            system = dynamics.commit()
+            # Use the system committed by _commit_replica() if there is one.
+            # When contexts are shared the replica is no longer resident in its
+            # slot by the time the cycle finishes, so it was committed earlier,
+            # while it still was.
+            committed = self._committed.pop(index, None)
 
-            # Get the simulation speed.
-            speed = dynamics.time_speed()
+            if committed is None:
+                # Commit the current system.
+                system = dynamics.commit()
+
+                # Get the simulation speed.
+                speed = dynamics.time_speed()
+            else:
+                system, speed = committed
 
             # Call the base class checkpoint method to save the system state.
             with self._lock:
@@ -2098,7 +2701,8 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False):
                     return index, error
 
             # Delete all trajectory frames from the Sire system within the
-            # dynamics object.
+            # dynamics object. This is a no-op when the replica was committed
+            # earlier, since the frames were cleared then to free the slot.
             dynamics._d._sire_mols.delete_all_frames()
 
             if block == -1:
@@ -2111,23 +2715,24 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False):
                     f"for {_lam_sym} = {lam:.5f}"
                 )
 
-            # Log the number of waters within the GCMC sampling volume.
-            if gcmc_sampler is not None:
-                gcmc_sampler.push()
-                try:
-                    n_moves = gcmc_sampler._num_moves
-                    acc_str = (
-                        f", acceptance rate = {gcmc_sampler.move_acceptance_probability():.3f}"
-                        f" (ins = {gcmc_sampler.num_insertions()}, del = {gcmc_sampler.num_deletions()})"
-                        if n_moves > 0
-                        else ""
-                    )
-                    _logger.info(
-                        f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} "
-                        f"is {gcmc_sampler.num_waters()}{acc_str}"
-                    )
-                finally:
-                    gcmc_sampler.pop()
+            # Log the number of waters within the GCMC sampling volume. Both
+            # the water count and the statistics are read from what was
+            # recorded while this replica was resident in its slot, since by
+            # the time the cycle finishes the slot may hold another one.
+            num_waters = self._dynamics_cache._gcmc_num_waters[index]
+            if gcmc_sampler is not None and num_waters is not None:
+                stats = gcmc_sampler.get_stats().get(gcmc_sampler.stats_key(lam))
+                n_moves = stats["num_moves"] if stats is not None else 0
+                acc_str = (
+                    f", acceptance rate = {stats['num_accepted'] / n_moves:.3f}"
+                    f" (ins = {stats['num_insertions']}, del = {stats['num_deletions']})"
+                    if n_moves > 0
+                    else ""
+                )
+                _logger.info(
+                    f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} "
+                    f"is {num_waters}{acc_str}"
+                )
 
             # Log terminal flip acceptance rate for this replica.
             if self._terminal_flip_samplers is not None:
@@ -2263,8 +2868,8 @@ def _merge_gcmc_stats(self):
         """
         stats = {}
 
-        for i in range(len(self._lambda_values)):
-            _, gcmc_sampler = self._dynamics_cache.get(i)
+        for slot in range(self._dynamics_cache._num_slots):
+            _, gcmc_sampler = self._dynamics_cache.get(slot)
             if gcmc_sampler is not None:
                 stats.update(gcmc_sampler.get_stats())
 
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index e7b970d..24b3ec7 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -150,6 +150,328 @@ def test_rest2_selection(ethane_methanol, rest2_selection, is_valid):
                 runner = RunnerBase(ethane_methanol, Config(**config))
 
 
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+@pytest.mark.parametrize("max_contexts", [1, 2, 3, 4])
+def test_repex_bounded_contexts(ethane_methanol, max_contexts):
+    """
+    Validate that a replica exchange simulation runs when there are fewer
+    OpenMM contexts than replicas, so that each context is re-used to
+    propagate several replicas per cycle.
+    """
+    num_lambda = 4
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "12fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": max_contexts,
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+
+        # Only the requested number of contexts should have been created.
+        assert len(runner._dynamics_cache._dynamics) == max_contexts
+
+        # Every replica must be assigned to exactly one slot.
+        groups = runner._dynamics_cache._groups
+        assert sorted(r for group in groups for r in group) == list(range(num_lambda))
+
+        runner.run()
+
+        # Output is per replica, regardless of how many contexts were used.
+        assert (Path(tmpdir) / "repex_matrix.txt").exists()
+        for i in range(num_lambda):
+            assert Path(runner._filenames[i]["energy_traj"]).exists()
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_frame_frequency_constraint(ethane_methanol):
+    """
+    Validate that frames can only be saved on checkpoint cycles when contexts
+    are re-used across lambda values.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "12fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "8fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": 4,
+            "replica_exchange": True,
+            "max_contexts": 2,
+        }
+
+        with pytest.raises(ValueError, match="frame_frequency"):
+            RepexRunner(ethane_methanol, Config(**config))
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+@pytest.mark.parametrize("update_constraints", [True, False])
+def test_repex_update_constraints(ethane_methanol, update_constraints):
+    """
+    Validate both constraint modes. Ethane to methanol does perturb a
+    constrained bond length, so update_constraints=True forces the context to
+    be reinitialised on every lambda change.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "12fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": 4,
+            "replica_exchange": True,
+            "max_contexts": 2,
+            "update_constraints": update_constraints,
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+        runner.run()
+
+        assert (Path(tmpdir) / "repex_matrix.txt").exists()
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_bounded_contexts_output_equivalence(ethane_methanol):
+    """
+    Validate that re-using contexts produces the same output structure as one
+    context per replica. Energies are not compared: a shared context consumes
+    the integrator's random number stream in a different order, so the
+    trajectories legitimately differ.
+    """
+    import pandas as pd
+
+    num_lambda = 4
+
+    def run(max_contexts, tmpdir):
+        config = {
+            "runtime": "16fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": max_contexts,
+        }
+        runner = RepexRunner(ethane_methanol, Config(**config))
+        runner.run()
+        return [
+            pd.read_parquet(runner._filenames[i]["energy_traj"])
+            for i in range(num_lambda)
+        ]
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        full = run(num_lambda, tmpdir)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        cached = run(1, tmpdir)
+
+    for i in range(num_lambda):
+        assert len(cached[i]) == len(full[i])
+        assert list(cached[i].columns) == list(full[i].columns)
+        assert cached[i].index.equals(full[i].index)
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_bounded_contexts_restart(ethane_methanol):
+    """
+    Validate that a replica exchange simulation using fewer contexts than
+    replicas can be restarted, and that the energy trajectory is extended
+    rather than restarted.
+    """
+    import pandas as pd
+
+    num_lambda = 4
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "8fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": 2,
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+        runner.run()
+
+        num_rows = [
+            len(pd.read_parquet(runner._filenames[i]["energy_traj"]))
+            for i in range(num_lambda)
+        ]
+
+        # Restart, extending the runtime.
+        config["runtime"] = "16fs"
+        config["restart"] = True
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+        runner.run()
+
+        for i in range(num_lambda):
+            extended = pd.read_parquet(runner._filenames[i]["energy_traj"])
+            assert len(extended) > num_rows[i]
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+@pytest.mark.parametrize("max_contexts", [1, 4])
+def test_repex_checkpoint_single_lock(ethane_methanol, max_contexts):
+    """
+    Validate that every checkpoint file is written within a single acquisition
+    of the file lock, so that a process streaming the output off the machine
+    always sees a coherent set rather than a mixture of new and old files.
+    """
+    import somd2.runner._repex as repex_module
+
+    num_lambda = 4
+    acquisitions = []
+
+    real_filelock = repex_module._FileLock
+
+    class CountingFileLock(real_filelock):
+        def acquire(self, *args, **kwargs):
+            acquisitions.append(1)
+            return super().acquire(*args, **kwargs)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "8fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": max_contexts,
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+
+        repex_module._FileLock = CountingFileLock
+        try:
+            runner.run()
+        finally:
+            repex_module._FileLock = real_filelock
+
+    # Two cycles, each taking the lock once for the checkpoint files and once
+    # for the repex state, plus a final acquisition. This must not scale with
+    # the number of passes.
+    assert len(acquisitions) == 5
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+@pytest.mark.parametrize("max_contexts", [1, 4])
+def test_repex_gcmc_bounded_contexts(ethane_methanol, max_contexts):
+    """
+    Validate that GCMC sampling works when contexts are re-used across lambda
+    values, so a slot's single sampler is re-parameterised and re-pointed at
+    the ghost file of whichever replica it hosts.
+    """
+    pytest.importorskip("loch")
+
+    num_lambda = 4
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "8fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": max_contexts,
+            "gcmc": True,
+            "gcmc_selection": "resname LIG",
+            "gcmc_frequency": "4fs",
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+        runner.run()
+
+        assert (Path(tmpdir) / "repex_matrix.txt").exists()
+
+        # One ghost file per lambda, each with a line per saved frame. A slot
+        # writing to the wrong file would leave these unbalanced.
+        counts = []
+        for lam in runner._lambda_values:
+            ghost_file = Path(tmpdir) / f"gcmc_ghosts_{lam:.5f}.txt"
+            assert ghost_file.exists()
+            counts.append(len(ghost_file.read_text().strip().splitlines()))
+
+        assert len(set(counts)) == 1, f"unbalanced ghost files: {counts}"
+        assert counts[0] > 0
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_concurrent_slots(ethane_methanol):
+    """
+    Validate that replicas sharing a slot are never propagated concurrently.
+    Oversubscribing exercises this on a single GPU, since the worker count is
+    the number of GPUs times the oversubscription factor.
+    """
+    num_lambda = 4
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "12fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": 2,
+            "oversubscription_factor": 2,
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+
+        # Every batch must contain at most one replica per slot.
+        num_workers = runner._num_gpus * config["oversubscription_factor"]
+        for batch in runner._safe_batches(num_workers):
+            slots = [runner._dynamics_cache.slot_for(r) for r in batch]
+            assert len(slots) == len(set(slots)), f"batch {batch} shares a slot"
+
+        runner.run()
+
+        assert (Path(tmpdir) / "repex_matrix.txt").exists()
+
+
 @pytest.mark.parametrize(
     "gpu_devices, expected",
     [
@@ -169,12 +491,8 @@ def test_rest2_selection(ethane_methanol, rest2_selection, is_valid):
 def test_physical_device_mapping(gpu_devices, expected):
     """
     Validate that an OpenMM device index is mapped to the physical device
-    backing it.
-
-    OpenMM numbers devices relative to the visible set, whereas pynvml and
-    pyopencl enumerate every device on the machine. Querying the memory of a
-    device by its OpenMM index therefore reports the wrong GPU whenever the
-    visible set does not start at zero.
+    backing it, since OpenMM numbers devices relative to the visible set
+    whereas pynvml and pyopencl enumerate all of them.
     """
     from somd2.runner._repex import DynamicsCache
 
@@ -204,8 +522,7 @@ def test_check_device_memory_queries_requested_device(monkeypatch, device, key,
 
     pynvml = pytest.importorskip("pynvml")
 
-    # Force the OpenCL branch to fail so that the pynvml path is always taken,
-    # regardless of what the machine running the tests has installed.
+    # Force the OpenCL branch to fail so the pynvml path is always taken.
     broken = types.SimpleNamespace()
 
     def get_platforms():
@@ -235,3 +552,160 @@ def by_uuid(uuid):
 
     assert DynamicsCache._check_device_memory(device) == (1, 2, 3)
     assert requested == {key: value}
+
+
+def test_legacy_checkpoint_restore():
+    """
+    Validate that a checkpoint written before slots existed can still be
+    loaded. These stored the replica states unpermuted, with the states array
+    holding the mapping to apply on restart.
+    """
+    from somd2.runner._repex import DynamicsCache
+
+    n = 4
+    legacy = {
+        "_lambdas": [0.0, 0.33, 0.67, 1.0],
+        "_rest2_scale_factors": [1.0] * n,
+        "_states": np.array([2, 0, 1, 3]),
+        "_time": None,
+        "_openmm_states": [f"state{i}" for i in range(n)],
+        "_gcmc_samplers": [None] * n,
+        "_gcmc_states": [f"water{i}" for i in range(n)],
+        "_gcmc_stats": [None] * n,
+        "_terminal_flip_stats": [[0, 0]] * n,
+        "_num_proposed": np.zeros((n, n)),
+        "_num_accepted": np.zeros((n, n)),
+        "_num_swaps": np.zeros((n, n)),
+    }
+
+    cache = object.__new__(DynamicsCache)
+    cache.__setstate__(dict(legacy))
+
+    # Converted to the current convention: each replica's own state, with the
+    # last mix applied.
+    assert cache._openmm_states == ["state2", "state0", "state1", "state3"]
+    assert cache._gcmc_states == ["water2", "water0", "water1", "water3"]
+
+    # Every replica is seeded from its stored state on a restart.
+    assert cache._state_moved == [True] * n
+
+    # Attributes postdating the checkpoint are defaulted, one slot per replica.
+    assert cache._num_slots == n
+    assert cache._groups == [[0], [1], [2], [3]]
+    assert cache._energy_trajectories == [None] * n
+    assert cache._ghost_files == [None] * n
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_max_contexts_change_on_restart(ethane_methanol):
+    """
+    Validate that the number of contexts can change on restart. The slot
+    layout is rebuilt from the configuration rather than restored.
+    """
+    num_lambda = 4
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "8fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": 2,
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+        assert len(runner._dynamics_cache._dynamics) == 2
+        runner.run()
+
+        # Restart with a different number of contexts.
+        config["runtime"] = "16fs"
+        config["restart"] = True
+        config["max_contexts"] = 4
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+        assert len(runner._dynamics_cache._dynamics) == 4
+        assert runner._dynamics_cache._groups == [[0], [1], [2], [3]]
+        runner.run()
+
+        assert (Path(tmpdir) / "repex_matrix.txt").exists()
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_gcmc_lambda_cache_warm(ethane_methanol, monkeypatch):
+    """
+    Validate that a GCMC sampler builds exactly one OpenMM context, scanning
+    it over every lambda it will host, and none once running. A mismatch
+    between the cached lambdas and those passed to set_lambda would show up
+    here as an extra build.
+    """
+    loch = pytest.importorskip("loch")
+
+    num_lambda = 4
+    calls = []
+
+    real_precompute = loch.GCMCSampler._precompute_lambdas
+
+    def counting_precompute(self, lambda_values, rest2_scales):
+        # Only record calls with work to do. Deduplicated, since the caller
+        # may name the same lambda twice but the scan extracts it once.
+        missing = sorted(
+            {
+                (float(lam), float(scale))
+                for lam, scale in zip(lambda_values, rest2_scales)
+                if (float(lam), float(scale)) not in self._lambda_params
+            }
+        )
+        if missing:
+            calls.append(missing)
+        return real_precompute(self, lambda_values, rest2_scales)
+
+    monkeypatch.setattr(loch.GCMCSampler, "_precompute_lambdas", counting_precompute)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "8fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": 1,
+            "gcmc": True,
+            "gcmc_selection": "resname LIG",
+            "gcmc_frequency": "4fs",
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+
+        # A single context build, covering every lambda value in one scan.
+        assert len(calls) == 1
+        assert sorted(lam for lam, _ in calls[0]) == sorted(runner._lambda_values)
+
+        runner.run()
+
+        # No further context may be built once the simulation is running.
+        assert len(calls) == 1
+
+        # An uncached lambda still works, building and caching on demand.
+        _, sampler = runner._dynamics_cache.get(0)
+        uncached = 0.123456
+        assert uncached not in runner._lambda_values
+
+        sampler.push()
+        try:
+            sampler.set_lambda(uncached)
+        finally:
+            sampler.pop()
+
+        assert len(calls) == 2
+        assert (uncached, sampler._rest2_scale) in sampler._lambda_params

From 602b50288f1cd6c8e3f11b735781c2bc45110d69 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 28 Jul 2026 16:21:45 +0100
Subject: [PATCH 22/67] Fix constrained bond lengths at a common lambda value.

---
 src/somd2/config/_config.py | 37 +++++++++++++--
 src/somd2/runner/_base.py   |  7 +++
 src/somd2/runner/_repex.py  | 90 +++++++++++++++++++++++++++++++++----
 tests/runner/test_repex.py  | 14 ++++++
 4 files changed, 136 insertions(+), 12 deletions(-)

diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index 4096cac..82994f9 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -140,6 +140,7 @@ def __init__(
         replica_exchange=False,
         max_contexts=None,
         update_constraints=True,
+        constraint_lambda_index=0,
         randomise_velocities=False,
         perturbed_system=None,
         terminal_flip_frequency=None,
@@ -396,9 +397,21 @@ def __init__(
             constraints is correct, but requires the OpenMM context to be reinitialised
             whenever a constrained bond length actually changes, which is slow. Set this
             to False if that overhead is significant; the constrained bond lengths are
-            then frozen at those of the lambda value the context was created at. Note
-            that this is distinct from 'dynamic_constraints', which controls where the
-            constraint lengths are taken from rather than whether they track lambda.
+            then frozen at those of the lambda value given by
+            'constraint_lambda_index'. Note that this is distinct from
+            'dynamic_constraints', which controls where the constraint lengths are
+            taken from rather than whether they track lambda.
+
+        constraint_lambda_index: int
+            The index of the lambda value at which to fix the constrained bond lengths
+            when 'update_constraints' is False. Every context is created at this lambda
+            value, so that the constraints are the same for all replicas rather than
+            depending on which context a replica is assigned to. The default of zero is
+            arbitrary but consistent; a lambda schedule that perturbs bonds away from
+            the end states may warrant a different choice. This is only used for
+            replica exchange simulations, and only when 'max_contexts' is less than the
+            number of replicas, 'update_constraints' is False, and a constrained bond
+            length actually perturbs with lambda.
 
         randomise_velocities: bool
             Whether to randomise velocities at the start of each replica exchange cycle
@@ -655,6 +668,7 @@ def __init__(
         self.replica_exchange = replica_exchange
         self.max_contexts = max_contexts
         self.update_constraints = update_constraints
+        self.constraint_lambda_index = constraint_lambda_index
         self.randomise_velocities = randomise_velocities
         self.perturbed_system = perturbed_system
         self.terminal_flip_frequency = terminal_flip_frequency
@@ -1835,6 +1849,23 @@ def update_constraints(self, update_constraints):
             raise ValueError("'update_constraints' must be of type 'bool'")
         self._update_constraints = update_constraints
 
+    @property
+    def constraint_lambda_index(self):
+        return self._constraint_lambda_index
+
+    @constraint_lambda_index.setter
+    def constraint_lambda_index(self, constraint_lambda_index):
+        if not isinstance(constraint_lambda_index, int):
+            try:
+                constraint_lambda_index = int(constraint_lambda_index)
+            except Exception:
+                raise ValueError("'constraint_lambda_index' must be of type 'int'")
+        if constraint_lambda_index < 0:
+            raise ValueError(
+                "'constraint_lambda_index' must be greater than or equal to 0"
+            )
+        self._constraint_lambda_index = constraint_lambda_index
+
     @property
     def randomise_velocities(self):
         return self._randomise_velocities
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index c42adde..b7ef5ab 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -1188,8 +1188,13 @@ def _check_end_state_constraints(self):
         """
         Internal function to check whether the constraints are the same at the two
         end states.
+
+        Sets self._end_state_constraints_differ, which records whether any
+        constrained bond length changes with lambda.
         """
 
+        self._end_state_constraints_differ = False
+
         # Find all perturbable molecules in the system..
         pert_mols = self._system.molecules("property is_perturbable")
 
@@ -1214,12 +1219,14 @@ def _check_end_state_constraints(self):
 
             # Check for equivalence.
             if len(constraints0) != len(constraints1):
+                self._end_state_constraints_differ = True
                 _logger.info(
                     f"Constraints are at not the same at {_lam_sym} = 0 and {_lam_sym} = 1."
                 )
             else:
                 for c0, c1 in zip(constraints0, constraints1):
                     if c0 != c1:
+                        self._end_state_constraints_differ = True
                         _logger.info(
                             f"Constraints are at not the same at {_lam_sym} = 0 and {_lam_sym} = 1."
                         )
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index d006da9..5dee8f9 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -56,6 +56,7 @@ def __init__(
         xml_filenames=None,
         num_slots=None,
         update_constraints=True,
+        constraint_lambda_index=None,
         gpu_devices=None,
     ):
         """
@@ -103,6 +104,15 @@ def __init__(
             Whether to update the constraints when changing the lambda value of
             a slot.
 
+        constraint_lambda_index: int
+            The index of the lambda value to create every context at, so that
+            constrained bond lengths are the same for every replica. If None,
+            each context is created at the lambda value of the first replica it
+            hosts, which is only consistent between replicas when there is a
+            context each. Only meaningful when 'update_constraints' is False,
+            since the constraints are otherwise updated whenever lambda
+            changes.
+
         gpu_devices: list
             The physical devices backing each OpenMM device index, i.e. the
             entries of CUDA_VISIBLE_DEVICES. Used to query the memory of the
@@ -127,6 +137,7 @@ def __init__(
         self._num_replicas = num_replicas
         self._num_slots = num_slots
         self._update_constraints = update_constraints
+        self._constraint_lambda_index = constraint_lambda_index
         self._gpu_devices = gpu_devices
         self._states = _np.array(range(num_replicas))
         self._time = None
@@ -373,6 +384,14 @@ def _create_dynamics(
             lam = lambdas[seed]
             scale = rest2_scale_factors[seed]
 
+            # Create the context at a common lambda value, so that constrained
+            # bond lengths are the same for every replica. Only set when the
+            # constraints aren't updated as the slot changes lambda.
+            if self._constraint_lambda_index is None:
+                build_lam = lam
+            else:
+                build_lam = lambdas[self._constraint_lambda_index]
+
             # Work out the device index.
             device = i % num_gpus
 
@@ -402,7 +421,7 @@ def _create_dynamics(
 
             # Overload the device and lambda value.
             dynamics_kwargs["device"] = device
-            dynamics_kwargs["lambda_value"] = lam
+            dynamics_kwargs["lambda_value"] = build_lam
             dynamics_kwargs["rest2_scale"] = scale
 
             if gcmc_kwargs is not None:
@@ -1212,6 +1231,7 @@ def __init__(self, system, config):
                 xml_filenames=xml_filenames,
                 num_slots=self._num_slots,
                 update_constraints=self._config.update_constraints,
+                constraint_lambda_index=self._constraint_lambda_index,
                 gpu_devices=self._gpu_devices,
             )
 
@@ -1292,6 +1312,9 @@ def __init__(self, system, config):
             # is per-replica, so it doesn't depend on the grouping.
             self._dynamics_cache._num_slots = self._num_slots
             self._dynamics_cache._update_constraints = self._config.update_constraints
+            self._dynamics_cache._constraint_lambda_index = (
+                self._constraint_lambda_index
+            )
             self._dynamics_cache._build_slot_layout()
 
             # Create the dynamics objects.
@@ -1394,6 +1417,32 @@ def __init__(self, system, config):
         # keyed by replica index and emptied by _checkpoint().
         self._committed = {}
 
+    def _build_lambda(self, replica):
+        """
+        Return the lambda value to create a replica's context at.
+
+        This is the replica's own lambda value, unless the constraints are
+        fixed at a common one, in which case a rebuilt context has to use that
+        too or it would pick up the constrained bond lengths of its own lambda
+        value instead.
+
+        Parameters
+        ----------
+
+        replica: int
+            The index of the replica.
+
+        Returns
+        -------
+
+        float
+            The lambda value to create the context at.
+        """
+        if self._constraint_lambda_index is None:
+            return self._lambda_values[replica]
+
+        return self._lambda_values[self._constraint_lambda_index]
+
     def _replica_passes(self, cycle):
         """
         Work out which replicas to propagate in each pass of a cycle.
@@ -1485,8 +1534,11 @@ def _set_num_slots(self):
             self._num_slots = min(self._config.max_contexts, num_replicas)
 
         # There is a context per replica, so nothing is re-used and all of the
-        # constraints below are irrelevant.
+        # constraints below are irrelevant. A context then keeps the lambda
+        # value it was created at, so there is no need to fix the constraints
+        # at a common one.
         self._is_cached = self._num_slots < num_replicas
+        self._constraint_lambda_index = None
 
         if not self._is_cached:
             if self._config.max_contexts is not None:
@@ -1526,11 +1578,31 @@ def _set_num_slots(self):
                 "suboptimal performance."
             )
 
-        if not self._config.update_constraints:
+        # When the constraints aren't updated as a slot changes lambda, they stay
+        # as they were when its context was created. Create every context at the
+        # same lambda value, so that the constrained bond lengths are uniform
+        # across replicas rather than depending on which slot a replica happens
+        # to be assigned to. Only needed if they actually perturb.
+        if not self._config.update_constraints and self._end_state_constraints_differ:
+            # Which lambda value is used matters less than every replica using
+            # the same one, since where the bonds actually perturb depends on
+            # the lambda schedule.
+            if self._config.constraint_lambda_index >= num_replicas:
+                msg = (
+                    f"'constraint_lambda_index' "
+                    f"({self._config.constraint_lambda_index}) is out of range "
+                    f"for {num_replicas} {_lam_sym} values."
+                )
+                _logger.error(msg)
+                raise ValueError(msg)
+
+            self._constraint_lambda_index = self._config.constraint_lambda_index
             _logger.warning(
-                "'update_constraints' is False. Constrained bond lengths will not "
-                "perturb with lambda, and are frozen at those of the lambda value "
-                "each context was created at."
+                f"'update_constraints' is False. Constrained bond lengths will not "
+                f"perturb with lambda, and are fixed at those of "
+                f"{_lam_sym} = "
+                f"{self._lambda_values[self._constraint_lambda_index]:.5f} "
+                f"for every replica."
             )
 
         from math import ceil
@@ -2264,7 +2336,7 @@ def _minimise(self, index):
 
                     # Overload the device and lambda value.
                     dynamics_kwargs["device"] = device
-                    dynamics_kwargs["lambda_value"] = self._lambda_values[index]
+                    dynamics_kwargs["lambda_value"] = self._build_lambda(index)
                     dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index]
 
                     # Create the production dynamics object.
@@ -2374,7 +2446,7 @@ def _equilibrate(self, index):
 
                     # Overload the device and lambda value.
                     dynamics_kwargs["device"] = device
-                    dynamics_kwargs["lambda_value"] = self._lambda_values[index]
+                    dynamics_kwargs["lambda_value"] = self._build_lambda(index)
                     dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index]
                     dynamics_kwargs["timestep"] = self._config._equilibration_timestep
                     dynamics_kwargs["constraint"] = constraint
@@ -2420,7 +2492,7 @@ def _equilibrate(self, index):
 
             # Overload the device and lambda value.
             dynamics_kwargs["device"] = device
-            dynamics_kwargs["lambda_value"] = self._lambda_values[index]
+            dynamics_kwargs["lambda_value"] = self._build_lambda(index)
             dynamics_kwargs["rest2_scale"] = self._rest2_scale_factors[index]
 
             # Create the production dynamics object.
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 24b3ec7..8eb157c 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -440,6 +440,10 @@ def test_repex_concurrent_slots(ethane_methanol):
     Validate that replicas sharing a slot are never propagated concurrently.
     Oversubscribing exercises this on a single GPU, since the worker count is
     the number of GPUs times the oversubscription factor.
+
+    This is also the only test that equilibrates, so it covers moving replicas
+    in and out of their slots during equilibration, the context rebuild when
+    the constraints change, and the post-equilibration checkpoint.
     """
     num_lambda = 4
 
@@ -451,6 +455,7 @@ def test_repex_concurrent_slots(ethane_methanol):
             "energy_frequency": "4fs",
             "checkpoint_frequency": "4fs",
             "frame_frequency": "4fs",
+            "equilibration_time": "4fs",
             "platform": "cuda",
             "max_threads": 1,
             "num_lambda": num_lambda,
@@ -461,6 +466,15 @@ def test_repex_concurrent_slots(ethane_methanol):
 
         runner = RepexRunner(ethane_methanol, Config(**config))
 
+        # Guard against the equilibration coverage being lost silently.
+        assert runner._is_equilibration
+
+        # Minimising without constraints and equilibrating with them means the
+        # contexts are rebuilt part way through, which is the path being
+        # covered here.
+        assert not runner._config.minimisation_constraints
+        assert runner._config.equilibration_constraints
+
         # Every batch must contain at most one replica per slot.
         num_workers = runner._num_gpus * config["oversubscription_factor"]
         for batch in runner._safe_batches(num_workers):

From 7aa1d55e040f8f1b555879c9ebe78af675ea2f7f Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 09:48:24 +0100
Subject: [PATCH 23/67] Seed the end state from the middle of a group.

---
 src/somd2/runner/_repex.py | 10 ++++-
 tests/runner/test_repex.py | 82 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 5dee8f9..7299b80 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -381,6 +381,14 @@ def _create_dynamics(
             # The replica that seeds this slot.
             seed = self._groups[i][0]
 
+            # The replica in the middle of the group, used to choose which end
+            # state the starting coordinates come from. A slot's context is
+            # created from a single system and every replica it hosts starts
+            # from that context, so taking the middle rather than the first
+            # keeps any mismatch to at most half a group, next to the lambda
+            # value at which the end state switches.
+            middle = self._groups[i][len(self._groups[i]) // 2]
+
             lam = lambdas[seed]
             scale = rest2_scale_factors[seed]
 
@@ -411,7 +419,7 @@ def _create_dynamics(
                 mols = system[seed]
             # This is a new simulation. For lambda > 0.5, use the perturbed
             # system to seed the starting coordinates and periodic space.
-            elif perturbed_system is not None and lam > 0.5:
+            elif perturbed_system is not None and lambdas[middle] > 0.5:
                 mols = perturbed_system
             else:
                 mols = system
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 8eb157c..07ecc86 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -723,3 +723,85 @@ def counting_precompute(self, lambda_values, rest2_scales):
 
         assert len(calls) == 2
         assert (uncached, sampler._rest2_scale) in sampler._lambda_params
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_perturbed_system_seeding(ethane_methanol):
+    """
+    Validate that the end states are seeded from the right coordinates when
+    contexts are shared.
+
+    A context is created from a single system and every replica it hosts starts
+    from that context, so the end state a replica starts from is chosen from
+    the middle of the group rather than its first replica. That keeps any
+    mismatch next to the lambda value at which the end state switches, instead
+    of it depending on where the groups happen to fall.
+    """
+    import sire as sr
+
+    # A perturbed end state, displaced so that its coordinates are distinct.
+    perturbed = ethane_methanol.clone()
+    perturbed.set_property("space", ethane_methanol.property("space"))
+    coords = sr.io.get_coords_array(ethane_methanol)
+    from sire.legacy.IO import setCoordinates
+
+    perturbed = sr.system.System(
+        setCoordinates(perturbed._system, (coords + 1.0).tolist())
+    )
+
+    # Ten replicas across three contexts, a layout in which the switch falls
+    # inside a group.
+    num_lambda = 10
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "4fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": 3,
+            "perturbed_system": perturbed,
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+
+        reference = sr.io.get_coords_array(runner._system)
+        target = sr.io.get_coords_array(runner._perturbed_system)
+
+        import openmm.unit as omm_unit
+
+        seeded = []
+        for i in range(num_lambda):
+            positions = runner._dynamics_cache._openmm_states[i][
+                "positions"
+            ].value_in_unit(omm_unit.angstrom)
+            from_reference = np.allclose(positions, reference, atol=1e-3)
+            from_target = np.allclose(positions, target, atol=1e-3)
+            assert from_reference != from_target, f"replica {i} matches neither"
+            seeded.append("perturbed" if from_target else "reference")
+
+        # The end states themselves must always be right.
+        assert seeded[0] == "reference"
+        assert seeded[-1] == "perturbed"
+
+        # Both systems must be used, otherwise the option does nothing.
+        assert set(seeded) == {"reference", "perturbed"}
+
+        # Only the group containing the switch can be seeded from the wrong end
+        # state, and then for no more than half of it. Choosing the end state
+        # from the first replica of a group rather than its middle breaks this.
+        lambdas = runner._lambda_values
+        for group in runner._dynamics_cache._groups:
+            wrong = [
+                i for i in group if (seeded[i] == "perturbed") != (lambdas[i] > 0.5)
+            ]
+            assert len(wrong) <= len(group) // 2, (
+                f"group {group} has {len(wrong)} replicas seeded from the "
+                f"wrong end state: {wrong}"
+            )

From dbf257cf5e4cf24f2d330f171ee9d2491c32eed8 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 09:55:39 +0100
Subject: [PATCH 24/67] Check that replicas resume from where they stopped.

---
 tests/runner/test_repex.py | 60 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 58 insertions(+), 2 deletions(-)

diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 07ecc86..0f51514 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -296,8 +296,9 @@ def run(max_contexts, tmpdir):
 def test_repex_bounded_contexts_restart(ethane_methanol):
     """
     Validate that a replica exchange simulation using fewer contexts than
-    replicas can be restarted, and that the energy trajectory is extended
-    rather than restarted.
+    replicas can be restarted, that each replica resumes from the state it
+    stopped at, and that the energy trajectory is extended rather than
+    restarted.
     """
     import pandas as pd
 
@@ -326,11 +327,66 @@ def test_repex_bounded_contexts_restart(ethane_methanol):
             for i in range(num_lambda)
         ]
 
+        # The state each replica finished at, which the checkpoint holds.
+        import openmm.unit as omm_unit
+
+        stopped = [
+            {
+                "positions": state["positions"].value_in_unit(omm_unit.nanometer),
+                "velocities": state["velocities"].value_in_unit(
+                    omm_unit.nanometer / omm_unit.picosecond
+                ),
+                "box": state["box"].value_in_unit(omm_unit.nanometer),
+            }
+            for state in runner._dynamics_cache._openmm_states
+        ]
+
         # Restart, extending the runtime.
         config["runtime"] = "16fs"
         config["restart"] = True
 
         runner = RepexRunner(ethane_methanol, Config(**config))
+
+        # Every replica must resume from where it stopped. The contexts are
+        # created from the input system, so the only thing carrying the
+        # simulated state across a restart is the checkpoint.
+        for i in range(num_lambda):
+            state = runner._dynamics_cache._openmm_states[i]
+            for key, unit in (
+                ("positions", omm_unit.nanometer),
+                ("velocities", omm_unit.nanometer / omm_unit.picosecond),
+                ("box", omm_unit.nanometer),
+            ):
+                assert np.allclose(
+                    state[key].value_in_unit(unit), stopped[i][key], atol=1e-6
+                ), f"replica {i} {key} not restored"
+
+        # The input coordinates must not be what was restored, otherwise the
+        # checks above would pass even if the checkpoint were ignored.
+        import sire as sr
+
+        inputs = sr.io.get_coords_array(runner._system)
+        restored = runner._dynamics_cache._openmm_states[0]["positions"].value_in_unit(
+            omm_unit.angstrom
+        )
+        assert not np.allclose(restored, inputs, atol=1e-3)
+
+        # Restoring the checkpoint into the cache is not enough: the contexts
+        # are created from the input system, so the state has to reach them
+        # too. Loading a replica is what pushes it.
+        for i in range(num_lambda):
+            runner._dynamics_cache.load_replica(i)
+            dynamics, _ = runner._dynamics_cache.get(runner._dynamics_cache.slot_for(i))
+            positions = (
+                dynamics.context()
+                .getState(getPositions=True)
+                .getPositions(asNumpy=True)
+                .value_in_unit(omm_unit.nanometer)
+            )
+            assert np.allclose(positions, stopped[i]["positions"], atol=1e-5), (
+                f"replica {i} positions not pushed into its context"
+            )
+
         runner.run()
 
         for i in range(num_lambda):

From da97dafdc52566c3fbcc04b213b38c2113706150 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 09:58:26 +0100
Subject: [PATCH 25/67] Read positions from a legacy checkpoint when logging
 the GCMC sphere.

---
 src/somd2/runner/_repex.py | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 7299b80..ce93f90 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -691,6 +691,30 @@ def save_openmm_state(self, slot, replica):
             "box": state.getPeriodicBoxVectors(asNumpy=True),
         }
 
+    @staticmethod
+    def _get_positions(state):
+        """
+        Return the positions from a saved OpenMM state.
+
+        Parameters
+        ----------
+
+        state: dict or openmm.State
+            The state to read. Dicts (new format) hold the positions directly.
+            A bare openmm.State is accepted for backwards compatibility with
+            old checkpoint files, as it is by _apply_openmm_state().
+
+        Returns
+        -------
+
+        openmm.unit.Quantity
+            The positions.
+        """
+        if isinstance(state, dict):
+            return state["positions"]
+
+        return state.getPositions(asNumpy=True)
+
     @staticmethod
     def _apply_openmm_state(context, state):
         """
@@ -1356,7 +1380,9 @@ def __init__(self, system, config):
         for i, lam in enumerate(self._lambda_values):
             _, gcmc_sampler = self._dynamics_cache.get(self._dynamics_cache.slot_for(i))
             if gcmc_sampler is not None and gcmc_sampler._reference is not None:
-                positions = self._dynamics_cache._openmm_states[i]["positions"]
+                positions = DynamicsCache._get_positions(
+                    self._dynamics_cache._openmm_states[i]
+                )
                 target = gcmc_sampler._get_target_position(
                     positions.value_in_unit(_omm_unit.angstrom)
                 )

From dd98085f0cd73bdd08eea196b5cbfa0e17a18f5a Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 10:12:53 +0100
Subject: [PATCH 26/67] Set the box vectors before the positions.

---
 src/somd2/runner/_repex.py  |  9 +++++++--
 src/somd2/runner/_runner.py | 16 ++++++++++------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index ce93f90..2c7b6fa 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -730,12 +730,17 @@ def _apply_openmm_state(context, state):
             The state to apply. Dicts (new format) contain "positions",
             "velocities", and "box" numpy arrays. A bare openmm.State is
             accepted for backwards compatibility with old checkpoint files.
+
+        Note that the step count and simulation time carried by an openmm.State
+        are deliberately not restored for the dict format. They are held
+        separately and applied by Dynamics._set_clock().
         """
         if isinstance(state, dict):
+            # Set the box before the positions, since a barostat may have
+            # changed it between the state being saved and restored.
+            context.setPeriodicBoxVectors(*state["box"])
             context.setPositions(state["positions"])
             context.setVelocities(state["velocities"])
-            if state["box"] is not None:
-                context.setPeriodicBoxVectors(*state["box"])
         else:
             # Legacy openmm.State from checkpoint files written before this
             # format change.
diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index fb0dac7..5571fb6 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -730,12 +730,9 @@ def generate_lam_vals(lambda_base, increment=0.001):
             import openmm.unit as _omm_unit
 
             _npz_state = _np.load(self._filenames[index]["checkpoint_state"])
-            dynamics.context().setPositions(
-                _npz_state["positions"] * _omm_unit.nanometer
-            )
-            dynamics.context().setVelocities(
-                _npz_state["velocities"] * _omm_unit.nanometer / _omm_unit.picosecond
-            )
+
+            # Set the box before the positions, since a barostat may have
+            # changed it between the state being saved and restored.
             if "box" in _npz_state:
                 from openmm import Vec3 as _Vec3
 
@@ -746,6 +743,13 @@ def generate_lam_vals(lambda_base, increment=0.001):
                     _Vec3(*_box[2]) * _omm_unit.nanometer,
                 )
 
+            dynamics.context().setPositions(
+                _npz_state["positions"] * _omm_unit.nanometer
+            )
+            dynamics.context().setVelocities(
+                _npz_state["velocities"] * _omm_unit.nanometer / _omm_unit.picosecond
+            )
+
         # Reset the GCMC sampler. This resets the sampling statistics and clears
         # the associated OpenMM forces.
         if gcmc_sampler is not None:

From 78c8d684b575d2263f59ee2cdbe1ddb4fd23305b Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 10:15:37 +0100
Subject: [PATCH 27/67] Accumulate the flags for replicas whose state has
 moved.

---
 src/somd2/runner/_repex.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 2c7b6fa..c1985e4 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -955,8 +955,14 @@ def mix_states(self, old_states):
         self._gcmc_states = [self._gcmc_states[state] for state in self._states]
 
         # Flag the replicas whose state moved, so that load_replica() knows it
-        # has to push new positions and velocities into the context.
-        self._state_moved = [bool(state != i) for i, state in enumerate(self._states)]
+        # has to push new positions and velocities into the context. The flags
+        # accumulate, and are only cleared once the state has been pushed, so
+        # that a replica mixed twice without being loaded in between isn't left
+        # starting from whatever its context happens to hold.
+        self._state_moved = [
+            moved or bool(state != i)
+            for i, (state, moved) in enumerate(zip(self._states, self._state_moved))
+        ]
 
         # Update the swap matrix.
         for i, state in enumerate(self._states):

From 4b19dfd8b78ec73b7d0e7d3b5e0ea0aef1a1d341 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 10:28:08 +0100
Subject: [PATCH 28/67] Don't minimise on restart.

---
 CHANGELOG.md                |  1 +
 src/somd2/runner/_repex.py  | 10 +++++++---
 src/somd2/runner/_runner.py |  4 ++--
 3 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4082038..be35684 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ Changelog
 * Store GCMC sampling statistics per lambda value, converting those from earlier checkpoints on restart [#184](https://github.com/OpenBioSim/somd2/pull/184).
 * Link restart systems to the reference end state rather than the perturbed one, since that is the coordinate set that dynamics maintains. Perturbable molecules were otherwise resumed from the coordinates they were built with [#189](https://github.com/OpenBioSim/somd2/pull/189).
 * Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas.
+* Skip minimisation on restart.
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index c1985e4..53845f8 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1219,7 +1219,11 @@ def __init__(self, system, config):
         perturbable_constraint = self._config.perturbable_constraint
 
         # Don't use constraints during minimisation.
-        if self._config.minimise and not self._config.minimisation_constraints:
+        if (
+            self._config.minimise
+            and not self._is_restart
+            and not self._config.minimisation_constraints
+        ):
             constraint = "none"
             perturbable_constraint = "none"
 
@@ -1743,8 +1747,8 @@ def run(self):
         # Create the replica list.
         replica_list = list(range(self._config.num_lambda))
 
-        # Minimise at each lambda value.
-        if self._config.minimise:
+        # Minimise at each lambda value. Don't minimise on restart.
+        if self._config.minimise and not self._is_restart:
             for batch in self._safe_batches(num_workers):
                 with ThreadPoolExecutor(max_workers=num_workers) as executor:
                     try:
diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index 5571fb6..3497eba 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -560,8 +560,8 @@ def generate_lam_vals(lambda_base, increment=0.001):
         else:
             terminal_flip_sampler = None
 
-        # Minimisation.
-        if self._config.minimise:
+        # Minimisation. Don't minimise on restart.
+        if self._config.minimise and not is_restart:
             constraint = self._config.constraint
             perturbable_constraint = self._config.perturbable_constraint
 

From b05c162c414b881e0c73c8721c5d696aa0260c03 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 14:22:11 +0100
Subject: [PATCH 29/67] Consolidate the duplicated context count and GPU load
 warnings.

---
 src/somd2/runner/_repex.py | 37 ++++++++++++++++---------------------
 1 file changed, 16 insertions(+), 21 deletions(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 53845f8..1c02a8c 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -124,13 +124,6 @@ def __init__(
         if num_slots is None:
             num_slots = num_replicas
 
-        # Warn if the number of slots is not a multiple of the number of GPUs.
-        if num_slots > num_gpus and num_slots % num_gpus != 0:
-            _logger.warning(
-                "The number of contexts is not a multiple of the number of GPUs. "
-                "This may result in suboptimal performance."
-            )
-
         # Initialise attributes.
         self._lambdas = lambdas
         self._rest2_scale_factors = rest2_scale_factors
@@ -1589,6 +1582,22 @@ def _set_num_slots(self):
         self._is_cached = self._num_slots < num_replicas
         self._constraint_lambda_index = None
 
+        # Contexts run concurrently, so it's the number of them that determines
+        # how evenly the GPUs are loaded, not the number of replicas.
+        num_workers = self._num_gpus * self._config.oversubscription_factor
+
+        if self._num_slots < num_workers:
+            _logger.warning(
+                f"The number of contexts ({self._num_slots}) is less than the "
+                f"number of workers ({num_workers}). Some GPUs will be left idle."
+            )
+        elif self._num_slots % self._num_gpus != 0:
+            _logger.warning(
+                f"The number of contexts ({self._num_slots}) is not a multiple "
+                f"of the number of GPUs ({self._num_gpus}). This may result in "
+                "suboptimal performance."
+            )
+
         if not self._is_cached:
             if self._config.max_contexts is not None:
                 _logger.info(
@@ -1613,20 +1622,6 @@ def _set_num_slots(self):
             _logger.error(msg)
             raise ValueError(msg)
 
-        num_workers = self._num_gpus * self._config.oversubscription_factor
-
-        if self._num_slots < num_workers:
-            _logger.warning(
-                f"'max_contexts' ({self._num_slots}) is less than the number of "
-                f"workers ({num_workers}). Some GPUs will be left idle."
-            )
-        elif self._num_slots % self._num_gpus != 0:
-            _logger.warning(
-                f"'max_contexts' ({self._num_slots}) is not a multiple of the "
-                f"number of GPUs ({self._num_gpus}). This may result in "
-                "suboptimal performance."
-            )
-
         # When the constraints aren't updated as a slot changes lambda, they stay
         # as they were when its context was created. Create every context at the
         # same lambda value, so that the constrained bond lengths are uniform

From 93591100d92a31cabb2defc832bb548cb44d4ba3 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 14:22:28 +0100
Subject: [PATCH 30/67] Document max_contexts in the replica exchange section
 of the README.

---
 README.md | 39 ++++++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index cad485e..20bcb0a 100644
--- a/README.md
+++ b/README.md
@@ -161,16 +161,37 @@ GPUs 0 and 1.
 ## Replica exchange
 
 `SOMD2` supports Hamiltonian replica exchange (HREX) simulations, which can be
-enabled using the `--replica-exchange` option. Note that dynamics contexts will
-be created up-front for all replicas, so this can be memory intensive. As such,
+enabled using the `--replica-exchange` option. By default, dynamics contexts are
+created up-front for all replicas, so this can be memory intensive. As such,
 replica exchange is intended for use on multi-GPU nodes with a large amount of
-memory. For optimal performance, it is recommended that the number of replicas
-be a multiple of the number of GPUs. It is also possible to oversubscribe the
-GPUs, i.e. have more than one replica running on a GPU at a time. This can be
-controlled via the `--oversubscription-factor` option, e.g. a value of 2 would
-allow 2 replicas to run on each GPU at a time.
-
-The swap frequency for replica exchange is controlled by the `energy-frequency`
+memory. It is also possible to oversubscribe the GPUs, i.e. have more than one
+replica running on a GPU at a time. This can be controlled via the
+`--oversubscription-factor` option, e.g. a value of 2 would allow 2 replicas to
+run on each GPU at a time.
+
+If the number of replicas you want doesn't fit in GPU memory, use the
+`--max-contexts` option to cap the number of contexts that are created. Each
+context is then re-used to propagate several replicas per cycle, changing its
+lambda value as it goes, so the number of replicas is no longer limited by
+memory. For example, `--num-lambda 24 --max-contexts 4` runs 24 replicas using
+the memory of 4. This costs some performance, since the replicas sharing a
+context run one after another rather than at the same time, so only use it when
+one context per replica won't fit. When contexts are re-used, `--frame-frequency`
+must equal `--checkpoint-frequency`.
+
+For optimal performance, it is recommended that the number of contexts, i.e. the
+number of replicas, or `--max-contexts` if it is set, be a multiple of the number
+of GPUs, and no smaller than the number of GPUs multiplied by the
+oversubscription factor. `SOMD2` will warn you if this isn't the case.
+
+Changing the lambda value of a context requires it to be reinitialised whenever a
+constrained bond length actually perturbs with lambda, which is slow. If this
+overhead is significant, set `--update-constraints=False` to freeze the
+constrained bond lengths at those of a single lambda value, chosen with
+`--constraint-lambda-index`. Both options are ignored unless contexts are being
+re-used.
+
+The swap frequency for replica exchange is controlled by the `--energy-frequency`
 option, i.e. we compute the energies for all replicas at this frequency, then
 attempt to mix the replicas. A larger value will improve performance, but may
 reduce the efficiency of the exchange.

From e297c245aaf20ffe3a2a48223df8f8b486b036c9 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 14:31:19 +0100
Subject: [PATCH 31/67] Fix typos and formatting errors in README.

---
 README.md | 37 +++++++++++++++++++------------------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 20bcb0a..f9b8e63 100644
--- a/README.md
+++ b/README.md
@@ -74,7 +74,7 @@ Then install `somd2` into the environment:
 pip install -e .
 ```
 
-> [!Note]
+> [!NOTE]
 > Pixi does not run conda post-link scripts, so the `ocl-icd-system`
 > symlink needed for OpenCL won't be created automatically. After
 > creating the environment (or after a pixi update), run the following
@@ -117,7 +117,7 @@ In order to run an alchemical free-energy simulation you will need to
 first create a stream file containing the _perturbable_ system of interest.
 This can be created using [BioSimSpace](https://github.com/OpenBioSim/biosimspace).
 For example, following the tutorial
-[here](https://biosimspace.openbiosim.org/versions/2023.4.0/tutorials/hydration_freenrg.html).
+[here](https://biosimspace.openbiosim.org/tutorials/hydration_freenrg.html).
 Once the system is created, it can be streamed to file using, e.g.:
 
 ```python
@@ -129,23 +129,23 @@ BSS.Stream.save(system, "perturbable_system")
 You can then run a simulation with:
 
 ```
-somd2 perturtbable_system.bss
+somd2 perturbable_system.bss
 ```
 
 The help message provides information on all of the supported options, along
 with their default values. Options can be specified on the command line, or
 using a YAML configuration file, passed with the `--config` option. Any options
-explicity set on the command line will override those set via the config file.
+explicitly set on the command line will override those set via the config file.
 
 An example perturbable system for a methane to ethanol perturbation in solvent
 can be found [here](https://sire.openbiosim.org/m/merged_molecule.s3.bz2).
 This is a `bzip2` compressed file that will need to be extracted before use.
 
-#### Running SOMD2 using one or more GPUs
+### Running SOMD2 using one or more GPUs
 
 In order to run using GPUs you will first need to set the relevant environment
-variable. For example, to run using 4 CUDA enabled GPUS set `CUDA_VISIBLE_DEVICES=0,1,2,3`
-(for openCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively).
+variable. For example, to run using 4 CUDA enabled GPUs set `CUDA_VISIBLE_DEVICES=0,1,2,3`
+(for OpenCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively).
 
 By default `SOMD2` will run using the CPU platform, however if the relevant
 environment variable has been set (as above) the new platform will be detected
@@ -154,7 +154,7 @@ available, the `--platform` option can be set (for example `--platform cuda`).
 
 By default, `SOMD2` will automatically manage the distribution of lambda windows
 across all listed devices. In order to restrict the number of devices used
-the `--max_gpus` option can be set, for example setting `max_gpus=2` while
+the `--max-gpus` option can be set, for example setting `--max-gpus 2` while
 `CUDA_VISIBLE_DEVICES` are set as above would restrict `SOMD2` to using only
 GPUs 0 and 1.
 
@@ -167,7 +167,8 @@ replica exchange is intended for use on multi-GPU nodes with a large amount of
 memory. It is also possible to oversubscribe the GPUs, i.e. have more than one
 replica running on a GPU at a time. This can be controlled via the
 `--oversubscription-factor` option, e.g. a value of 2 would allow 2 replicas to
-run on each GPU at a time.
+run on each GPU at a time. This requires the NVIDIA multi-process service (MPS)
+to be enabled, see [GPU oversubscription](#gpu-oversubscription) below.
 
 If the number of replicas you want doesn't fit in GPU memory, use the
 `--max-contexts` option to cap the number of contexts that are created. Each
@@ -186,7 +187,7 @@ oversubscription factor. `SOMD2` will warn you if this isn't the case.
 
 Changing the lambda value of a context requires it to be reinitialised whenever a
 constrained bond length actually perturbs with lambda, which is slow. If this
-overhead is significant, set `--update-constraints=False` to freeze the
+overhead is significant, pass `--no-update-constraints` to freeze the
 constrained bond lengths at those of a single lambda value, chosen with
 `--constraint-lambda-index`. Both options are ignored unless contexts are being
 re-used.
@@ -200,7 +201,7 @@ reduce the efficiency of the exchange.
 
 We also support Replica Exchange with Solute Scaling
 ([REST2](https://pubs.acs.org/doi/10.1021/jp204407d)) simulations to facilitate sampling for perturbations
-involving conformational changes, e.g.  ring flips. This can be enabled
+involving conformational changes, e.g. ring flips. This can be enabled
 using the `--rest2-scale` option, which specifies the "temperature" of the
 REST2 region relative to the rest of the system. By default, the REST2 region
 comprises _all_ atoms in perturbable molecules, but can be controlled via the
@@ -244,7 +245,7 @@ from `conda-forge`.
 
 SOMD2 supports terminal ring flip Monte Carlo (MC) moves to improve sampling
 of terminal aromatic rings in perturbable ligands, as described in
-[this paper](https://chemrxiv.org/doi/full/10.26434/chemrxiv-2025-2zkx5).
+[this paper](https://doi.org/10.26434/chemrxiv-2025-2zkx5).
 Each move attempts a discrete rotation of a terminal ring around the bond
 connecting it to the rest of the molecule, accepted or rejected via the
 Metropolis criterion. Terminal ring groups are detected automatically from
@@ -351,7 +352,7 @@ free_nrg = BSS.FreeEnergy.Relative.difference(pmf1, pmf2)
 
 When running HREX with a large number of replicas it can become computationally
 expensive to compute energies. (We need the energies of each replica at each
-lamdba value.) As a shortcut, it's possible to truncate the neighbourhood of
+lambda value.) As a shortcut, it's possible to truncate the neighbourhood of
 windows for which we compute energies, then use a large null energy for the
 remaining windows. This can be controlled via the `--num-energy-neighbours` option.
 For example, setting this to 2 would compute energies for the current window and
@@ -366,7 +367,7 @@ We support modification of ghost atom bonded terms to avoid spurious coupling
 to the physical system using the approach described in
 [this](https://pubs.acs.org/doi/10.1021/acs.jctc.0c01328) paper.
 These are enabled by default, but can be disabled using the ``--no-ghost-modifications``
-option. Modifications are implemented using the [ghostly](https://gitbub.com/OpenBioSim/ghostly)
+option. Modifications are implemented using the [ghostly](https://github.com/OpenBioSim/ghostly)
 package.
 
 ## Note for SOMD1 users
@@ -381,7 +382,7 @@ that the perturbation used is consistent with the approach from `somd1`, i.e.
 it uses the same modifications for bonded-terms involving dummy atoms as `somd1`.
 
 Finally, it is also possible to run `somd2` using an existing `somd1` perturbation
-file. To do so, you will also need to create a stream file representating the
+file. To do so, you will also need to create a stream file representing the
 λ = 0 state. For existing input generated by `prepareFEP.py`, this can be done as
 follows. (This assumes that the output has a prefix `somd1`.)
 
@@ -412,7 +413,7 @@ then simply omit the `--somd1-compatibility` option.
 ## GPU oversubscription
 
 If you have an NVIDIA GPU that supports the multi-process service (MPS), you can
-oversubscibe the GPU to run multiple OpenMM contexts on the same GPU at once,
+oversubscribe the GPU to run multiple OpenMM contexts on the same GPU at once,
 increasing the throughput of your simulation. To do this, you will need to first
 enable MPS by running the following command:
 
@@ -424,7 +425,7 @@ The number of contexts that can be run in parallel is then controlled by the
 `--oversubscription-factor` option, which defaults to 1.
 
 More details on MPS, including tuning options, can be found in the following
-[techical blog](https://developer.nvidia.com/blog/maximizing-openmm-molecular-dynamics-throughput-with-nvidia-multi-process-service/).
+[technical blog](https://developer.nvidia.com/blog/maximizing-openmm-molecular-dynamics-throughput-with-nvidia-multi-process-service/).
 
 ## Python API
 
@@ -446,5 +447,5 @@ frequently. (Frames are written to disk and cleared from memory at each
 checkpoint.)
 
 PyMBAR uses JAX by default for GPU acceleration, which can cause issues in
-some environments. If you encounter issues when analysing simlation output,
+some environments. If you encounter issues when analysing simulation output,
 try setting the `PYMBAR_DISABLE_JAX` environment variable to `1`.

From 9e97ff107436a3a74dddbc6fb5e97197bae4e1e7 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 20:31:09 +0100
Subject: [PATCH 32/67] Pin the perturbed system coordinate path in the seeding
 test.

---
 tests/runner/test_repex.py | 25 +++++++++++++++++++++++--
 1 file changed, 23 insertions(+), 2 deletions(-)

diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 0f51514..058237c 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -792,17 +792,38 @@ def test_repex_perturbed_system_seeding(ethane_methanol):
     the middle of the group rather than its first replica. That keeps any
     mismatch next to the lambda value at which the end state switches, instead
     of it depending on where the groups happen to fall.
+
+    Only the perturbable molecule is displaced, so the assertions also cover
+    the property path that its coordinates travel along: read from
+    'coordinates1' via link_to_perturbed, written to 'coordinates0', then read
+    back via link_to_reference.
     """
     import sire as sr
 
-    # A perturbed end state, displaced so that its coordinates are distinct.
+    # A perturbed end state, with the perturbable molecule displaced so that
+    # its coordinates are distinct.
     perturbed = ethane_methanol.clone()
     perturbed.set_property("space", ethane_methanol.property("space"))
     coords = sr.io.get_coords_array(ethane_methanol)
     from sire.legacy.IO import setCoordinates
 
+    # Flag the atoms of the perturbable molecules. get_coords_array returns the
+    # atoms in molecule order, so the offset tracks the array index.
+    is_perturbable = np.zeros(len(coords), dtype=bool)
+    offset = 0
+    for mol in ethane_methanol.molecules():
+        num_atoms = mol.num_atoms()
+        if mol.has_property("is_perturbable"):
+            is_perturbable[offset : offset + num_atoms] = True
+        offset += num_atoms
+    assert is_perturbable.any(), "no perturbable molecules in the test system"
+
+    coords[is_perturbable] += 1.0
+
+    # Write to 'coordinates1', which is the property that the runner reads via
+    # link_to_perturbed.
     perturbed = sr.system.System(
-        setCoordinates(perturbed._system, (coords + 1.0).tolist())
+        setCoordinates(perturbed._system, coords.tolist(), True)
     )
 
     # Ten replicas across three contexts, a layout in which the switch falls

From 06534ffe9dc30ab9029bba2a408b26f64d95abf3 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 20:44:38 +0100
Subject: [PATCH 33/67] Respect swap_end_states when seeding from the perturbed
 system.

---
 src/somd2/runner/_base.py  |  5 +++--
 src/somd2/runner/_repex.py | 13 ++++++++++---
 tests/runner/test_repex.py | 22 +++++++++++++++++-----
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index b7ef5ab..4974db5 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -624,8 +624,9 @@ def __init__(self, system, config):
         # Create a clone of the fully-prepared reference system with the
         # perturbed end-state coordinates and periodic space. This is done
         # after all system preparation so that the clone inherits the same
-        # topology and properties. It is used to seed starting coordinates
-        # for lambda > 0.5 replicas.
+        # topology and properties. It is used to seed starting coordinates for
+        # the replicas closest to the perturbed end state, i.e. lambda > 0.5,
+        # or lambda < 0.5 when the end states are swapped.
         if self._config.replica_exchange and self._config.perturbed_system is not None:
             from sire.legacy.IO import setCoordinates as _setCoordinates
 
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 1c02a8c..4b6e13c 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -407,12 +407,19 @@ def _create_dynamics(
                     "count": 0,
                 }
 
+            # Whether to seed from the perturbed end state. Swapping the end
+            # states reverses the lambda schedule, so the perturbed end state
+            # is then at lambda = 0 rather than lambda = 1.
+            seed_perturbed = (lambdas[middle] > 0.5) != dynamics_kwargs[
+                "swap_end_states"
+            ]
+
             # This is a restart, get the system for the seeding replica.
             if isinstance(system, list):
                 mols = system[seed]
-            # This is a new simulation. For lambda > 0.5, use the perturbed
-            # system to seed the starting coordinates and periodic space.
-            elif perturbed_system is not None and lambdas[middle] > 0.5:
+            # This is a new simulation. Use the perturbed system to seed the
+            # starting coordinates and periodic space.
+            elif perturbed_system is not None and seed_perturbed:
                 mols = perturbed_system
             else:
                 mols = system
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 058237c..20ddcc4 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -782,7 +782,8 @@ def counting_precompute(self, lambda_values, rest2_scales):
 
 
 @pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
-def test_repex_perturbed_system_seeding(ethane_methanol):
+@pytest.mark.parametrize("swap_end_states", [False, True])
+def test_repex_perturbed_system_seeding(ethane_methanol, swap_end_states):
     """
     Validate that the end states are seeded from the right coordinates when
     contexts are shared.
@@ -793,6 +794,9 @@ def test_repex_perturbed_system_seeding(ethane_methanol):
     mismatch next to the lambda value at which the end state switches, instead
     of it depending on where the groups happen to fall.
 
+    Swapping the end states reverses the lambda schedule, so the perturbed end
+    state moves to lambda = 0 and the seeding must mirror with it.
+
     Only the perturbable molecule is displaced, so the assertions also cover
     the property path that its coordinates travel along: read from
     'coordinates1' via link_to_perturbed, written to 'coordinates0', then read
@@ -844,6 +848,7 @@ def test_repex_perturbed_system_seeding(ethane_methanol):
             "replica_exchange": True,
             "max_contexts": 3,
             "perturbed_system": perturbed,
+            "swap_end_states": swap_end_states,
         }
 
         runner = RepexRunner(ethane_methanol, Config(**config))
@@ -863,9 +868,14 @@ def test_repex_perturbed_system_seeding(ethane_methanol):
             assert from_reference != from_target, f"replica {i} matches neither"
             seeded.append("perturbed" if from_target else "reference")
 
-        # The end states themselves must always be right.
-        assert seeded[0] == "reference"
-        assert seeded[-1] == "perturbed"
+        # The end states themselves must always be right. Swapping the end
+        # states puts the perturbed one at lambda = 0.
+        if swap_end_states:
+            assert seeded[0] == "perturbed"
+            assert seeded[-1] == "reference"
+        else:
+            assert seeded[0] == "reference"
+            assert seeded[-1] == "perturbed"
 
         # Both systems must be used, otherwise the option does nothing.
         assert set(seeded) == {"reference", "perturbed"}
@@ -876,7 +886,9 @@ def test_repex_perturbed_system_seeding(ethane_methanol):
         lambdas = runner._lambda_values
         for group in runner._dynamics_cache._groups:
             wrong = [
-                i for i in group if (seeded[i] == "perturbed") != (lambdas[i] > 0.5)
+                i
+                for i in group
+                if (seeded[i] == "perturbed") != ((lambdas[i] > 0.5) != swap_end_states)
             ]
             assert len(wrong) <= len(group) // 2, (
                 f"group {group} has {len(wrong)} replicas seeded from the "

From 6596422736b7bbf5e887d53b77c9805115dc5067 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 29 Jul 2026 20:47:05 +0100
Subject: [PATCH 34/67] Clarify what the perturbed_system input must contain.

---
 src/somd2/config/_config.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index 82994f9..5d653c3 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -419,9 +419,12 @@ def __init__(
 
         perturbed_system: str
             The path to a stream file containing a Sire system for the equilibrated perturbed
-            end state (lambda = 1). This will be used as the starting conformation all lambda
-            windows > 0.5 when performing a replica exchange simulation. (Note that this assumes
-            that the "coordinates1" property specifies the coordinates for perturbable molecules.)
+            end state (lambda = 1). This is the same system as the input, but with the
+            "coordinates1" property of any perturbable molecules holding the equilibrated
+            coordinates for the lambda = 1 state. It is used as the starting conformation for
+            the lambda windows closest to the perturbed end state when performing a replica
+            exchange simulation, i.e. those with lambda > 0.5, or lambda < 0.5 when
+            'swap_end_states' is True.
 
         terminal_flip_frequency: str
             Frequency at which to attempt terminal ring flip Monte Carlo moves. If None

From 3a95b4beb177a1ec5027acd5e105226ae2261f98 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sat, 1 Aug 2026 18:07:03 +0100
Subject: [PATCH 35/67] Suggest a max_contexts value when the replicas don't
 fit in memory.

---
 src/somd2/runner/_repex.py | 54 ++++++++++++++++++++++++++++++++++++--
 tests/runner/test_repex.py | 18 +++++++++++++
 2 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 4b6e13c..365d9ed 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -467,7 +467,19 @@ def _create_dynamics(
             try:
                 dynamics = mols.dynamics(**dynamics_kwargs)
             except Exception as e:
-                msg = f"Could not create dynamics object for lambda {lam:.5f} on device {device}: {e}"
+                # Contexts are appended as they are created, so this is the
+                # number that fit before the failure.
+                num_created = len(self._dynamics)
+                msg = (
+                    f"Could not create dynamics object for lambda {lam:.5f} on "
+                    f"device {device} after creating {num_created} of "
+                    f"{self._num_slots} contexts: {e}"
+                )
+                # Failing partway through points at exhausted resources.
+                # Failing on the first context does not, so don't advise on
+                # the number of contexts in that case.
+                if num_created > 0:
+                    msg += f". {self._max_contexts_advice(num_created)}"
                 _logger.error(msg)
                 raise RuntimeError(msg) from e
 
@@ -534,12 +546,25 @@ def _create_dynamics(
                 if est_total > total_mem:
                     baseline = info["before"]
                     replica_cost = first_cost + marginal_cost * (num_contexts - 1)
+
+                    # How many contexts fit on this device, given the measured
+                    # cost of the first and the marginal cost of the rest. The
+                    # slots are spread evenly over the devices, so the total is
+                    # capped by what the busiest device can hold.
+                    if marginal_cost > 0:
+                        per_device = 1 + floor(
+                            (total_mem - baseline - first_cost) / marginal_cost
+                        )
+                    else:
+                        per_device = info["count"]
+
                     msg = (
                         f"Not enough memory on device {device} for all assigned replicas. "
                         f"Baseline usage before simulation: {baseline / (1024**3):.2f} GB "
                         f"Estimated replica memory: {replica_cost / (1024**3):.2f} GB, "
                         f"Total estimated: {est_total / (1024**3):.2f} GB, "
-                        f"Available memory: {total_mem / (1024**3):.2f} GB."
+                        f"Available memory: {total_mem / (1024**3):.2f} GB. "
+                        f"{self._max_contexts_advice(per_device * num_gpus)}"
                     )
                     _logger.error(msg)
                     raise MemoryError(msg)
@@ -1012,6 +1037,31 @@ def _physical_device(self, device):
 
         return gpu_devices[device]
 
+    @staticmethod
+    def _max_contexts_advice(num_contexts):
+        """
+        Advice for a user whose replicas don't fit in GPU memory.
+
+        Parameters
+        ----------
+
+        num_contexts: int
+            The number of OpenMM contexts that are known to fit.
+        """
+
+        if num_contexts < 1:
+            return (
+                "A single OpenMM context does not fit on this device. Reduce "
+                "the size of the system, or use a device with more memory."
+            )
+
+        return (
+            f"Set 'max_contexts' to {num_contexts} or fewer to re-use contexts "
+            "across lambda values, which removes the memory limit on the "
+            "number of replicas. This also requires 'frame_frequency' to equal "
+            "'checkpoint_frequency'."
+        )
+
     @staticmethod
     def _check_device_memory(device=0):
         """
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 20ddcc4..c3158af 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -572,6 +572,24 @@ def test_physical_device_mapping(gpu_devices, expected):
     assert [cache._physical_device(i) for i in range(len(expected))] == expected
 
 
+def test_max_contexts_advice():
+    """
+    Validate the advice given when the replicas don't fit in GPU memory. It
+    must name a number that the user can pass to 'max_contexts', and mention
+    the frame frequency constraint that comes with it.
+    """
+    from somd2.runner._repex import DynamicsCache
+
+    advice = DynamicsCache._max_contexts_advice(6)
+    assert "'max_contexts' to 6 or fewer" in advice
+    assert "frame_frequency" in advice and "checkpoint_frequency" in advice
+
+    # Nothing fits, so there is no number to suggest.
+    advice = DynamicsCache._max_contexts_advice(0)
+    assert "max_contexts" not in advice
+    assert "does not fit" in advice
+
+
 @pytest.mark.parametrize(
     "device, key, value",
     [

From d28a6dfb615caf363200f1f2b246d91d1e593036 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 14:26:42 +0100
Subject: [PATCH 36/67] Remove the redundant post-equilibration GCMC water
 count.

---
 src/somd2/runner/_repex.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 365d9ed..609f62b 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -2605,14 +2605,6 @@ def _equilibrate(self, index):
             if gcmc_sampler is not None:
                 self._reset_gcmc_sampler(gcmc_sampler, dynamics)
 
-                # Compute the current number of waters in the GCMC sampling
-                # volume after equilibration.
-                gcmc_sampler.push()
-                try:
-                    gcmc_sampler.num_waters(context=dynamics.context())
-                finally:
-                    gcmc_sampler.pop()
-
             # Set the new dynamics object.
             self._dynamics_cache.set(slot, dynamics)
 

From 0c2e3024d629a887d44c7dee1378cc61f357cb48 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 14:28:37 +0100
Subject: [PATCH 37/67] Assert that GCMC water occupancy follows the replica
 through a mix.

---
 tests/runner/test_repex.py | 46 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index c3158af..7a489eb 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -590,6 +590,52 @@ def test_max_contexts_advice():
     assert "does not fit" in advice
 
 
+def test_gcmc_state_follows_replica():
+    """
+    Validate that the GCMC water occupancy travels with the configuration it
+    belongs to when replicas are mixed.
+
+    A slot holds one GCMC sampler but hosts several replicas, and load_replica()
+    installs a replica's water state by diffing it against whatever the last
+    resident left in the sampler. If the occupancy did not follow the positions
+    through a mix, a replica would run with another replica's waters, which
+    gives plausible numbers rather than an obviously wrong output file.
+
+    The ghost files and the sampling statistics belong to the lambda window
+    rather than the configuration, so they must not be permuted.
+    """
+    from somd2.runner._repex import DynamicsCache
+
+    num_replicas = 4
+
+    cache = object.__new__(DynamicsCache)
+
+    # Label both states with the replica they came from, so that a replica
+    # holding mismatched positions and waters is detectable.
+    cache._openmm_states = list(range(num_replicas))
+    cache._gcmc_states = list(range(num_replicas))
+    cache._ghost_files = [f"ghosts_{i}.txt" for i in range(num_replicas)]
+    cache._state_moved = [False] * num_replicas
+    cache._num_swaps = np.zeros((num_replicas, num_replicas))
+
+    # Mix twice, since a slot is re-used within a cycle.
+    for states in ([2, 0, 3, 1], [1, 3, 0, 2]):
+        old_states = list(range(num_replicas))
+        expected = [cache._gcmc_states[state] for state in states]
+
+        cache._states = states
+        cache.mix_states(old_states)
+
+        # The water occupancy follows the same permutation as the positions.
+        assert cache._gcmc_states == expected
+
+        # Every replica holds the positions and waters of the same origin.
+        assert cache._openmm_states == cache._gcmc_states
+
+    # The ghost files stay with the lambda window.
+    assert cache._ghost_files == [f"ghosts_{i}.txt" for i in range(num_replicas)]
+
+
 @pytest.mark.parametrize(
     "device, key, value",
     [

From 23d29fa2d6f3b2b75b7491b3caab8fb7b3907d77 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 15:22:06 +0100
Subject: [PATCH 38/67] Count the GCMC waters against the context in the
 regular runner.

---
 src/somd2/runner/_runner.py |  8 +++++--
 tests/runner/test_repex.py  | 44 +++++++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index 3497eba..2bebfa7 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -818,7 +818,6 @@ def generate_lam_vals(lambda_base, increment=0.001):
                         dynamics.context(),
                         force=True,
                     )
-                    gcmc_sampler.num_waters(context=dynamics.context())
                 finally:
                     gcmc_sampler.pop()
 
@@ -1156,9 +1155,14 @@ def generate_lam_vals(lambda_base, increment=0.001):
                                 if n_moves > 0
                                 else ""
                             )
+                            # Count against the context, since dynamics have
+                            # run since the last move.
+                            num_waters = gcmc_sampler.num_waters(
+                                context=dynamics.context()
+                            )
                             _logger.info(
                                 f"Current number of waters in GCMC volume at {_lam_sym} = {lambda_value:.5f} "
-                                f"is {gcmc_sampler.num_waters()}{acc_str}"
+                                f"is {num_waters}{acc_str}"
                             )
                         finally:
                             gcmc_sampler.pop()
diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 7a489eb..3a1c388 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -490,6 +490,50 @@ def test_repex_gcmc_bounded_contexts(ethane_methanol, max_contexts):
         assert counts[0] > 0
 
 
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_repex_gcmc_without_a_selection(ethane_methanol):
+    """
+    Validate GCMC sampling with no 'gcmc_selection', where moves are attempted
+    within the entire simulation box rather than a region around a selection.
+
+    The sampler then has no reference, so it cannot count the waters within a
+    region, and every move samples the whole box. Counting the waters on each
+    replica handover has to account for that.
+    """
+    pytest.importorskip("loch")
+
+    num_lambda = 4
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": "8fs",
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": "4fs",
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": num_lambda,
+            "replica_exchange": True,
+            "max_contexts": 2,
+            "gcmc": True,
+            "gcmc_frequency": "4fs",
+        }
+
+        runner = RepexRunner(ethane_methanol, Config(**config))
+
+        # The bulk-only path is the point of this test, so make sure it can't
+        # stop being exercised without the test failing.
+        assert runner._dynamics_cache._gcmc_samplers[0]._reference is None
+
+        runner.run()
+
+        assert (Path(tmpdir) / "repex_matrix.txt").exists()
+        for lam in runner._lambda_values:
+            assert (Path(tmpdir) / f"gcmc_ghosts_{lam:.5f}.txt").exists()
+
+
 @pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
 def test_repex_concurrent_slots(ethane_methanol):
     """

From b1d0e6f0a49075fca8440c9b1928e4f969622ea7 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 16:06:19 +0100
Subject: [PATCH 39/67] Cover GCMC with no sampling region in the regular
 runner.

---
 tests/runner/test_gcmc.py | 63 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 63 insertions(+)
 create mode 100644 tests/runner/test_gcmc.py

diff --git a/tests/runner/test_gcmc.py b/tests/runner/test_gcmc.py
new file mode 100644
index 0000000..c952966
--- /dev/null
+++ b/tests/runner/test_gcmc.py
@@ -0,0 +1,63 @@
+"""
+Tests for GCMC sampling via the regular (non replica exchange) runner.
+"""
+
+import pytest
+import re
+import tempfile
+
+from pathlib import Path
+
+from somd2.config import Config
+from somd2.runner import Runner
+
+from tests.conftest import has_cuda
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+def test_runner_gcmc_without_a_selection(ethane_methanol):
+    """
+    Validate GCMC sampling with no 'gcmc_selection', where moves are attempted
+    within the entire simulation box rather than a region around a selection.
+
+    The sampler then has no reference, so it cannot count the waters within a
+    region. Reporting the water count has to account for that, which is what
+    this exercises: counting it raised before there was a path for the
+    reference-free case.
+    """
+    pytest.importorskip("loch")
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = Config(
+            runtime="8fs",
+            output_directory=tmpdir,
+            energy_frequency="4fs",
+            checkpoint_frequency="4fs",
+            frame_frequency="4fs",
+            platform="cuda",
+            max_threads=1,
+            num_lambda=2,
+            gcmc=True,
+            gcmc_frequency="4fs",
+        )
+
+        # The bulk-only path is the point of this test, so make sure it can't
+        # stop being exercised without the test failing.
+        assert config.gcmc_selection is None
+
+        runner = Runner(ethane_methanol, config)
+        runner.run()
+
+        # GCMC ran, so the ghost residues were written.
+        for lam in runner._lambda_values:
+            assert (Path(tmpdir) / f"gcmc_ghosts_{lam:.5f}.txt").exists()
+
+        # With no region the count is the number of non-ghost waters in the
+        # box, which is never zero for a solvated system.
+        log = (Path(tmpdir) / config.log_file).read_text()
+        counts = [
+            int(x)
+            for x in re.findall(r"number of waters in GCMC volume.*? is (\d+)", log)
+        ]
+        assert counts, "no water count was logged"
+        assert all(count > 0 for count in counts), f"zero water count logged: {counts}"

From 682e5c92a31a60dde3da1e072afe6e1062687449 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 16:18:24 +0100
Subject: [PATCH 40/67] Pre-equilibrate the water before minimising in the
 regular runner.

---
 CHANGELOG.md                |  1 +
 src/somd2/runner/_runner.py | 41 ++++++++++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index be35684..d51646d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ Changelog
 * Link restart systems to the reference end state rather than the perturbed one, since that is the coordinate set that dynamics maintains. Perturbable molecules were otherwise resumed from the coordinates they were built with [#189](https://github.com/OpenBioSim/somd2/pull/189).
 * Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas.
 * Skip minimisation on restart.
+* Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket.
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index 2bebfa7..0017b44 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -584,6 +584,7 @@ def generate_lam_vals(lambda_base, increment=0.001):
                     device=device,
                     constraint=constraint,
                     perturbable_constraint=perturbable_constraint,
+                    gcmc_sampler=gcmc_sampler,
                 )
             except Exception as e:
                 msg = f"Minimisation failed for {_lam_sym} = {lambda_value:.5f}: {e}"
@@ -638,6 +639,15 @@ def generate_lam_vals(lambda_base, increment=0.001):
                     # Bind the GCMC sampler to the dynamics object.
                     gcmc_sampler.bind_dynamics(dynamics)
 
+                    # This context is built from the system as it was committed
+                    # by minimisation, so it doesn't carry the water state that
+                    # the moves performed there left behind.
+                    gcmc_sampler.push()
+                    try:
+                        gcmc_sampler._set_water_state(dynamics.context(), force=True)
+                    finally:
+                        gcmc_sampler.pop()
+
                     _logger.info(
                         f"Equilibrating with GCMC moves at {_lam_sym} = {lambda_value:.5f}"
                     )
@@ -808,9 +818,11 @@ def generate_lam_vals(lambda_base, increment=0.001):
                     finally:
                         gcmc_sampler.pop()
 
-            # Otherwise, if we've performed equilibration, then we need to reset
-            # the water state in the new context to match the equilibrated system.
-            elif is_equilibrated:
+            # Otherwise, reset the water state in the new context to match the
+            # system that the preparation stages left behind. Both minimisation
+            # and equilibration perform GCMC moves, so the state held by the
+            # sampler is ahead of the one the context was built with.
+            else:
                 # Reset the water state.
                 gcmc_sampler.push()
                 try:
@@ -1537,6 +1549,7 @@ def _minimisation(
         device=None,
         constraint="none",
         perturbable_constraint="none",
+        gcmc_sampler=None,
     ):
         """
         Minimise a system.
@@ -1562,6 +1575,10 @@ def _minimisation(
         perturbable_constraint: str
             The constraint for perturbable molecules.
 
+        gcmc_sampler: :class: `GCMCSampler <loch.GCMCSampler>`
+            A GCMC sampler to pre-equilibrate the water with before minimising.
+            If None, then no GCMC moves are performed.
+
         Returns
         -------
 
@@ -1589,6 +1606,24 @@ def _minimisation(
             # Create a dynamics object.
             dynamics = system.dynamics(**dynamics_kwargs)
 
+            # Pre-equilibrate the water before minimising, so that a dry pocket
+            # is filled before the geometry relaxes into it. The context is
+            # created from the sampler's own system, so its water state already
+            # matches and only needs binding.
+            if gcmc_sampler is not None:
+                gcmc_sampler.bind_dynamics(dynamics)
+
+                _logger.info(
+                    f"Pre-equilibrating with GCMC moves at {_lam_sym} = {lambda_value:.5f}"
+                )
+
+                gcmc_sampler.push()
+                try:
+                    for i in range(100):
+                        gcmc_sampler.move(dynamics.context())
+                finally:
+                    gcmc_sampler.pop()
+
             # Run the minimisation.
             dynamics.minimise(timeout=self._config.timeout)
 

From 6a80168602385ecd14e6a9feeee1e2ddd377b4e6 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 3 Aug 2026 09:27:51 +0100
Subject: [PATCH 41/67] Give each replica its own terminal flip statistics
 list.

---
 src/somd2/runner/_repex.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 609f62b..fe6c4a6 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -146,7 +146,7 @@ def __init__(
         # Whether the last mix moved a replica's state, and so whether it must
         # be pushed into the context before the next block.
         self._state_moved = [False] * num_replicas
-        self._terminal_flip_stats = [[0, 0]] * num_replicas
+        self._terminal_flip_stats = [[0, 0] for _ in range(num_replicas)]
         self._num_proposed = _np.matrix(_np.zeros((num_replicas, num_replicas)))
         self._num_accepted = _np.matrix(_np.zeros((num_replicas, num_replicas)))
         self._num_swaps = _np.matrix(_np.zeros((num_replicas, num_replicas)))
@@ -216,7 +216,7 @@ def __setstate__(self, state):
         if not hasattr(self, "_gcmc_states"):
             self._gcmc_states = [None] * n
         if not hasattr(self, "_terminal_flip_stats"):
-            self._terminal_flip_stats = [[0, 0]] * n
+            self._terminal_flip_stats = [[0, 0] for _ in range(n)]
         if not hasattr(self, "_time"):
             self._time = None
         if not hasattr(self, "_num_replicas"):

From c51911d2fa1576a2ca001ca97c02cef1135a5033 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 4 Aug 2026 19:40:56 +0100
Subject: [PATCH 42/67] Link CHANGELOG entries to PR.

---
 CHANGELOG.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d51646d..c2a0277 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,9 +15,9 @@ Changelog
 * Fixed the replica exchange GPU memory check querying the wrong device when `CUDA_VISIBLE_DEVICES` does not start at zero, since OpenMM numbers devices relative to the visible set whereas `pynvml` enumerates all of them [#183](https://github.com/OpenBioSim/somd2/issues/183).
 * Store GCMC sampling statistics per lambda value, converting those from earlier checkpoints on restart [#184](https://github.com/OpenBioSim/somd2/pull/184).
 * Link restart systems to the reference end state rather than the perturbed one, since that is the coordinate set that dynamics maintains. Perturbable molecules were otherwise resumed from the coordinates they were built with [#189](https://github.com/OpenBioSim/somd2/pull/189).
-* Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas.
-* Skip minimisation on restart.
-* Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket.
+* Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas [#191](https://github.com/OpenBioSim/somd2/pull/191).
+* Skip minimisation on restart [#191](https://github.com/OpenBioSim/somd2/pull/191).
+* Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket [#191](https://github.com/OpenBioSim/somd2/pull/191).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------

From ac9b1e6d2c78c0cebfcb70fef7d059b46813c056 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 4 Aug 2026 20:03:52 +0100
Subject: [PATCH 43/67] Add a precision option for GPU platforms.

---
 CHANGELOG.md                |  1 +
 src/somd2/config/_config.py | 24 ++++++++++++++++++++++++
 src/somd2/runner/_base.py   |  1 +
 3 files changed, 26 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c2a0277..d9dbe38 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,7 @@ Changelog
 * Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Skip minimisation on restart [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket [#191](https://github.com/OpenBioSim/somd2/pull/191).
+* Add a `precision` option for GPU platforms, defaulting to `mixed` [#191](https://github.com/OpenBioSim/somd2/pull/191).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index 5d653c3..adaa0e6 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -75,6 +75,7 @@ class Config:
         ],
         "log_level": [level.lower() for level in _logger._core.levels],
         "softcore_form": ["zacharias", "taylor", "beutler"],
+        "precision": ["single", "mixed", "double"],
     }
 
     # A dictionary of nargs for the various options.
@@ -132,6 +133,7 @@ def __init__(
         num_energy_neighbours=None,
         null_energy="1e6 kcal/mol",
         platform="auto",
+        precision="mixed",
         max_threads=None,
         max_gpus=None,
         max_sire_threads=None,
@@ -356,6 +358,12 @@ def __init__(
         platform: str
             Platform to run simulation on.
 
+        precision: str
+            The floating point precision to use on GPU platforms. 'single' is fastest,
+            'double' is slowest, and 'mixed' computes forces in single precision but
+            accumulates and integrates in double. Ignored by platforms that do not
+            support it, such as CPU.
+
         max_threads: int
             Maximum number of CPU threads to use for simulation. (Default None, uses all available)
             Does nothing if platform is set to CUDA.
@@ -663,6 +671,7 @@ def __init__(
         self.checkpoint_frequency = checkpoint_frequency
         self.num_checkpoint_workers = num_checkpoint_workers
         self.platform = platform
+        self.precision = precision
         self.max_threads = max_threads
         self.max_gpus = max_gpus
         self.max_sire_threads = max_sire_threads
@@ -1709,6 +1718,21 @@ def platform(self, platform):
             else:
                 self._platform = "cpu"
 
+    @property
+    def precision(self):
+        return self._precision
+
+    @precision.setter
+    def precision(self, precision):
+        if not isinstance(precision, str):
+            raise TypeError("'precision' must be of type 'str'")
+        precision = precision.lower().replace(" ", "")
+        if precision not in self._choices["precision"]:
+            raise ValueError(
+                f"'precision' not recognised. Valid options are: {', '.join(self._choices['precision'])}"
+            )
+        self._precision = precision
+
     @property
     def max_threads(self):
         return self._max_threads
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 4974db5..044e3a9 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -893,6 +893,7 @@ def __init__(self, system, config):
             "cutoff": self._config.cutoff,
             "cutoff_type": self._config.cutoff_type,
             "platform": self._config.platform,
+            "precision": self._config.precision,
             "rest2_selection": self._config.rest2_selection,
             "shift_coulomb": self._config.shift_coulomb,
             "shift_delta": self._config.shift_delta,

From b5b7cb5388b9d3e6e670e05f98bd76c66d558ead Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 5 Aug 2026 10:53:46 +0100
Subject: [PATCH 44/67] Pin the recorded times in the energy trajectory tests.

---
 tests/runner/test_repex.py | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py
index 3a1c388..d0a7bea 100644
--- a/tests/runner/test_repex.py
+++ b/tests/runner/test_repex.py
@@ -278,19 +278,31 @@ def run(max_contexts, tmpdir):
         return [
             pd.read_parquet(runner._filenames[i]["energy_traj"])
             for i in range(num_lambda)
-        ]
+        ], list(runner._lambda_values)
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        full = run(num_lambda, tmpdir)
+        full, lambdas = run(num_lambda, tmpdir)
 
     with tempfile.TemporaryDirectory() as tmpdir:
-        cached = run(1, tmpdir)
+        cached, _ = run(1, tmpdir)
+
+    # One record per energy_frequency interval, starting at the first. Pinning
+    # the values rather than only comparing the two runs catches a clock that
+    # is wrong the same way in both.
+    expected_times = [0.004, 0.008, 0.012, 0.016]
 
     for i in range(num_lambda):
         assert len(cached[i]) == len(full[i])
         assert list(cached[i].columns) == list(full[i].columns)
         assert cached[i].index.equals(full[i].index)
 
+        for records in (full[i], cached[i]):
+            times = [round(t, 6) for t in records.index.get_level_values(0)]
+            assert times == expected_times, f"replica {i} recorded times {times}"
+
+            # Each replica records energies at every lambda value.
+            assert list(records.columns) == lambdas
+
 
 @pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
 def test_repex_bounded_contexts_restart(ethane_methanol):
@@ -389,10 +401,18 @@ def test_repex_bounded_contexts_restart(ethane_methanol):
 
         runner.run()
 
+        # 8 fs at 4 fs intervals, extended to 16 fs. The clock has to continue
+        # from where it stopped rather than restarting at zero, so the records
+        # must run to the new runtime with no repeats or gaps.
+        expected_times = [0.004, 0.008, 0.012, 0.016]
+
         for i in range(num_lambda):
             extended = pd.read_parquet(runner._filenames[i]["energy_traj"])
             assert len(extended) > num_rows[i]
 
+            times = [round(t, 6) for t in extended.index.get_level_values(0)]
+            assert times == expected_times, f"replica {i} recorded times {times}"
+
 
 @pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
 @pytest.mark.parametrize("max_contexts", [1, 4])

From fdec2e726ef8a707d2bae695a54bd9656d3d9cdf Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 24 Aug 2026 10:15:44 +0100
Subject: [PATCH 45/67] Add section on keeping up-to-date during a development
 cycle. [ci skip]

---
 README.md | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/README.md b/README.md
index f9b8e63..dd22fec 100644
--- a/README.md
+++ b/README.md
@@ -93,6 +93,31 @@ You should now have a `somd2` executable in your path. To test, run:
 somd2 --help
 ```
 
+### Keeping up to date
+
+During a development cycle the OpenBioSim packages are pinned only to a
+`X.X.X.dev` version, not to a specific build. `somd2` and its dependencies
+therefore need to be kept in sync, so always update the whole stack together
+rather than `somd2` alone.
+
+For a conda install, update everything in one go:
+
+```
+conda update -c conda-forge -c openbiosim/label/dev sire biosimspace ghostly loch somd2
+```
+
+For a standalone pixi install, pull the latest `somd2` and refresh the
+pre-built dependencies:
+
+```
+git pull
+pixi update
+```
+
+For a full source install, `git pull` in *every* repository you have installed
+(`sire`, `biosimspace`, `ghostly`, `loch` and `somd2`), not just `somd2`. Since
+`sire` is compiled, you will also need to rebuild it.
+
 ## Development
 
 Pre-commit hooks are used to ensure consistent code formatting and linting.

From 771c8cd918c77ce4b8b05c8c23a82b07b30440b3 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 24 Aug 2026 12:11:11 +0100
Subject: [PATCH 46/67] Add support for auto-generating Morse restraints for
 ring-breaking.

---
 CHANGELOG.md                               |   1 +
 src/somd2/config/_config.py                |  99 +++++++++
 src/somd2/runner/_base.py                  |  88 ++++++++
 tests/runner/test_config.py                |  40 ++++
 tests/runner/test_ring_break_restraints.py | 232 +++++++++++++++++++++
 5 files changed, 460 insertions(+)
 create mode 100644 tests/runner/test_ring_break_restraints.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d9dbe38..ddfcb67 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ Changelog
 * Skip minimisation on restart [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Add a `precision` option for GPU platforms, defaulting to `mixed` [#191](https://github.com/OpenBioSim/somd2/pull/191).
+* Add support for generating Morse restraints for ring-breaking perturbations.
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index adaa0e6..28f2470 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -180,6 +180,9 @@ def __init__(
         restraint_search_time="1 ns",
         restraint_search_frequency="10 ps",
         restraint_search_receptor_selection=None,
+        morse_hard_well_depth="150 kcal mol-1",
+        morse_soft_well_depth="50 kcal mol-1",
+        morse_soft_force_constant="125 kcal mol-1 A-2",
     ):
         """
         Constructor.
@@ -622,6 +625,20 @@ def __init__(
             Sire selection string for receptor anchor atom candidates used
             during automatic Boresch restraint generation. If None, the default
             backbone selection is used (CA, C, N atoms in non-water molecules).
+
+        morse_hard_well_depth: str
+            The well depth of the "hard" Morse potential that replaces the
+            broken bond when auto-generating restraints for a ring-breaking
+            simulation.
+
+        morse_soft_well_depth: str
+            The well depth of the "soft" Morse restraint that holds the broken
+            fragment in place when auto-generating restraints for a
+            ring-breaking simulation.
+
+        morse_soft_force_constant: str
+            The force constant of the "soft" Morse restraint used when
+            auto-generating restraints for a ring-breaking simulation.
         """
 
         # Setup logger before doing anything else
@@ -717,6 +734,9 @@ def __init__(
         self.restraint_search_time = restraint_search_time
         self.restraint_search_frequency = restraint_search_frequency
         self.restraint_search_receptor_selection = restraint_search_receptor_selection
+        self.morse_hard_well_depth = morse_hard_well_depth
+        self.morse_soft_well_depth = morse_soft_well_depth
+        self.morse_soft_force_constant = morse_soft_force_constant
         self.write_config = write_config
         self.overwrite = overwrite
 
@@ -2793,6 +2813,85 @@ def restraint_search_receptor_selection(self, restraint_search_receptor_selectio
                 )
         self._restraint_search_receptor_selection = restraint_search_receptor_selection
 
+    @property
+    def morse_hard_well_depth(self):
+        return self._morse_hard_well_depth
+
+    @morse_hard_well_depth.setter
+    def morse_hard_well_depth(self, morse_hard_well_depth):
+        self._morse_hard_well_depth = self._parse_well_depth(
+            morse_hard_well_depth, "morse_hard_well_depth"
+        )
+
+    @property
+    def morse_soft_well_depth(self):
+        return self._morse_soft_well_depth
+
+    @morse_soft_well_depth.setter
+    def morse_soft_well_depth(self, morse_soft_well_depth):
+        self._morse_soft_well_depth = self._parse_well_depth(
+            morse_soft_well_depth, "morse_soft_well_depth"
+        )
+
+    @property
+    def morse_soft_force_constant(self):
+        return self._morse_soft_force_constant
+
+    @morse_soft_force_constant.setter
+    def morse_soft_force_constant(self, morse_soft_force_constant):
+        if not isinstance(morse_soft_force_constant, str):
+            raise TypeError("'morse_soft_force_constant' must be of type 'str'")
+
+        from sire.units import angstrom, kcal_per_mol
+
+        try:
+            k = _sr.u(morse_soft_force_constant)
+        except:
+            raise ValueError(
+                "Unable to parse 'morse_soft_force_constant' as a Sire "
+                f"GeneralUnit: {morse_soft_force_constant}"
+            )
+
+        if not k.has_same_units(kcal_per_mol / (angstrom * angstrom)):
+            raise ValueError("'morse_soft_force_constant' units are invalid.")
+
+        self._morse_soft_force_constant = k
+
+    @staticmethod
+    def _parse_well_depth(value, name):
+        """
+        Internal helper to validate a Morse potential well depth.
+
+        Parameters
+        ----------
+
+        value: str
+            The well depth as a string, e.g. "150 kcal mol-1".
+
+        name: str
+            The name of the option, used in error messages.
+
+        Returns
+        -------
+
+        well_depth: sire.units.GeneralUnit
+            The parsed well depth.
+        """
+        if not isinstance(value, str):
+            raise TypeError(f"'{name}' must be of type 'str'")
+
+        from sire.units import kcal_per_mol
+
+        try:
+            de = _sr.u(value)
+        except:
+            raise ValueError(f"Unable to parse '{name}' as a Sire GeneralUnit: {value}")
+
+        if not de.has_same_units(kcal_per_mol):
+            raise ValueError(f"'{name}' units are invalid.")
+
+        return de
+
     def _reset_logger(self, logger):
         """
         Internal method to reset the logger.
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 044e3a9..2f4eea8 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -267,6 +267,24 @@ def __init__(self, system, config):
             self._config._extra_args["use_gcmc_lrc"] = True
             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:
+            try:
+                self._config.restraints = self._generate_morse_restraints()
+            except Exception as e:
+                msg = (
+                    "Unable to generate Morse restraints for ring-breaking "
+                    f"perturbation: {e}. If the Morse potential has already been "
+                    "applied to the input system, then pass the corresponding "
+                    "restraints using the 'restraints' option."
+                )
+                _logger.error(msg)
+                raise RuntimeError(msg)
+
         # We're running in SOMD1 compatibility mode.
         if self._config.somd1_compatibility:
             from .._utils._somd1 import make_compatible
@@ -1002,6 +1020,76 @@ def _is_abfe_bound(self):
             and self._has_water
         )
 
+    @property
+    def _is_ring_break(self):
+        """
+        Whether this is a ring-breaking (or ring-making) simulation, i.e. one
+        using the 'ring_break_morph' lambda schedule, or its reverse.
+        """
+        return self._config._lambda_schedule_name in (
+            "ring_break_morph",
+            "reverse_ring_break_morph",
+        )
+
+    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.
+
+        The "hard" restraint directly replaces the harmonic bond that is broken
+        (or formed) by the perturbation, inheriting its force constant and
+        equilibrium length. The "soft" restraint acts on the same pair of atoms
+        and holds the broken fragment in place while the hard restraint is
+        switched off.
+
+        Returns
+        -------
+
+        restraints: [sire.mm.MorsePotentialRestraints]
+            The hard and soft Morse restraints, in the order expected by the
+            schedule's 'morse_hard' and 'morse_soft' levers.
+
+        Notes
+        -----
+
+        As a side effect, ``self._system`` is updated with the replacement of
+        the broken bond by the hard Morse potential.
+        """
+        from sire.restraints import morse_potential as _morse_potential
+
+        _logger.info(
+            "No restraints supplied for ring-breaking perturbation. "
+            "Generating default Morse restraints."
+        )
+
+        hard_restraints, self._system = _morse_potential(
+            self._system,
+            de=self._config.morse_hard_well_depth,
+            auto_parametrise=True,
+            direct_morse_replacement=True,
+            name="morse_hard",
+        )
+
+        # Restrain the same pair of atoms as the hard restraint, at the same
+        # equilibrium distance.
+        soft_restraints, _ = _morse_potential(
+            self._system,
+            atoms0=hard_restraints[0].atom0(),
+            atoms1=hard_restraints[0].atom1(),
+            r0=hard_restraints[0].r0(),
+            k=self._config.morse_soft_force_constant,
+            de=self._config.morse_soft_well_depth,
+            auto_parametrise=False,
+            direct_morse_replacement=False,
+            name="morse_soft",
+        )
+
+        _logger.info(f"Hard Morse restraint: {hard_restraints[0]}")
+        _logger.info(f"Soft Morse restraint: {soft_restraints[0]}")
+
+        return [hard_restraints, soft_restraints]
+
     def _generate_boresch_restraint(self, device=None):
         """
         Return a Boresch restraint for the ABFE simulation, either by loading
diff --git a/tests/runner/test_config.py b/tests/runner/test_config.py
index 39bd948..27d8e4a 100644
--- a/tests/runner/test_config.py
+++ b/tests/runner/test_config.py
@@ -83,3 +83,43 @@ def test_logfile_creation():
         assert Path.exists(runner._config.output_directory / runner._config.log_file)
 
         somd2._logger.remove()
+
+
+def test_morse_restraint_options():
+    """Validate that the Morse restraint options are parsed correctly."""
+    import math
+
+    import pytest
+
+    # The defaults are parsed as Sire units.
+    config = Config()
+    assert config.morse_hard_well_depth == sr.u("150 kcal mol-1")
+    assert config.morse_soft_well_depth == sr.u("50 kcal mol-1")
+    assert config.morse_soft_force_constant == sr.u("125 kcal mol-1 A-2")
+
+    # Equivalent units are accepted, and converted.
+    config = Config(morse_hard_well_depth="418.4 kJ mol-1")
+    assert math.isclose(
+        config.morse_hard_well_depth.to(sr.units.kcal_per_mol), 100.0, rel_tol=1e-6
+    )
+
+    # Well depths must be energies.
+    for option in ("morse_hard_well_depth", "morse_soft_well_depth"):
+        with pytest.raises(TypeError):
+            Config(**{option: 150})
+
+        with pytest.raises(ValueError, match="Unable to parse"):
+            Config(**{option: "not a unit"})
+
+        with pytest.raises(ValueError, match="units are invalid"):
+            Config(**{option: "150 kcal mol-1 A-2"})
+
+    # The force constant must be an energy per unit area.
+    with pytest.raises(TypeError):
+        Config(morse_soft_force_constant=125)
+
+    with pytest.raises(ValueError, match="Unable to parse"):
+        Config(morse_soft_force_constant="not a unit")
+
+    with pytest.raises(ValueError, match="units are invalid"):
+        Config(morse_soft_force_constant="125 kcal mol-1")
diff --git a/tests/runner/test_ring_break_restraints.py b/tests/runner/test_ring_break_restraints.py
new file mode 100644
index 0000000..255931a
--- /dev/null
+++ b/tests/runner/test_ring_break_restraints.py
@@ -0,0 +1,232 @@
+import tempfile
+
+import pytest
+
+from somd2.config import Config
+from somd2.runner import Runner
+
+
+def _config(tmpdir, **kwargs):
+    """Return a minimal ring-breaking config rooted at 'tmpdir'."""
+    options = {
+        "output_directory": tmpdir,
+        "lambda_schedule": "ring_break_morph",
+        "num_lambda": 3,
+        "runtime": "12fs",
+        "energy_frequency": "4fs",
+        "frame_frequency": "4fs",
+        "checkpoint_frequency": "4fs",
+        "equilibration_time": "0fs",
+        "minimise": False,
+        "platform": "CPU",
+        "max_threads": 1,
+    }
+    options.update(kwargs)
+    return Config(**options)
+
+
+def _restraint(restraints, name):
+    """Return the single restraint from the set called 'name'."""
+    for restraint_set in restraints:
+        if str(restraint_set.name()) == name:
+            assert len(restraint_set) == 1
+            return restraint_set[0]
+    raise AssertionError(f"No restraint set named {name!r} in {restraints}")
+
+
+def test_restraints_are_generated(syk_ring_break_mols):
+    """
+    Ensure that a pair of Morse restraints is automatically generated for a
+    ring-breaking perturbation when no restraint is supplied, and that they act
+    on the same pair of atoms, at the same equilibrium distance.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        runner = Runner(syk_ring_break_mols.clone(), _config(tmpdir))
+
+        restraints = runner._config.restraints
+        assert restraints is not None
+        assert len(restraints) == 2
+
+        hard = _restraint(restraints, "morse_hard")
+        soft = _restraint(restraints, "morse_soft")
+
+        # Both restraints act on the bond that is broken.
+        assert hard.atom0() == soft.atom0()
+        assert hard.atom1() == soft.atom1()
+        assert hard.r0() == soft.r0()
+
+        # The restraints are passed through to the dynamics.
+        assert runner._dynamics_kwargs["restraints"] is restraints
+
+
+def test_restraints_match_config(syk_ring_break_mols):
+    """
+    Ensure that the generated restraints use the well depths and force constant
+    from the config, and that the hard restraint inherits the force constant of
+    the bond that it replaces.
+    """
+    import sire as sr
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = _config(
+            tmpdir,
+            morse_hard_well_depth="123 kcal mol-1",
+            morse_soft_well_depth="45 kcal mol-1",
+            morse_soft_force_constant="67 kcal mol-1 A-2",
+        )
+        runner = Runner(syk_ring_break_mols.clone(), config)
+
+        hard = _restraint(runner._config.restraints, "morse_hard")
+        soft = _restraint(runner._config.restraints, "morse_soft")
+
+        assert hard.de() == sr.u("123 kcal mol-1")
+        assert soft.de() == sr.u("45 kcal mol-1")
+        assert soft.k() == sr.u("67 kcal mol-1 A-2")
+
+        # The hard restraint is auto-parametrised from the broken bond, so its
+        # force constant comes from the bond, not the config.
+        assert hard.k() != soft.k()
+        assert hard.k().value() > 0
+
+
+def test_broken_bond_is_replaced(syk_ring_break_mols):
+    """
+    Ensure that the hard restraint replaces the harmonic bond that is broken by
+    the perturbation, i.e. that the bond is removed from the runner's system.
+    Leaving both in place would double count the interaction.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        mols = syk_ring_break_mols.clone()
+        runner = Runner(mols, _config(tmpdir))
+
+        hard = _restraint(runner._config.restraints, "morse_hard")
+
+        def num_bonds(system, idx0, idx1):
+            """Count the bond potentials between a pair of atom indices."""
+            atoms = system.atoms()
+            atom0 = atoms[idx0]
+            atom1 = atoms[idx1]
+            mol = system[atom0.molecule().number()]
+            info = mol.info()
+            expected = {atom0.index().value(), atom1.index().value()}
+
+            count = 0
+            for bond_prop in ("bond0", "bond1"):
+                for potential in mol.property(bond_prop).potentials():
+                    idxs = {
+                        info.atom_idx(potential.atom0()).value(),
+                        info.atom_idx(potential.atom1()).value(),
+                    }
+                    if idxs == expected:
+                        count += 1
+            return count
+
+        # The unmodified input still has the bond, in the reference end state
+        # only, since it is broken by the perturbation.
+        assert num_bonds(syk_ring_break_mols, hard.atom0(), hard.atom1()) == 1
+
+        # The runner's system has it removed, replaced by the Morse restraint.
+        assert num_bonds(runner._system, hard.atom0(), hard.atom1()) == 0
+
+
+def test_restraints_are_deterministic(syk_ring_break_mols):
+    """
+    Ensure that generating the restraints twice from the same input gives
+    identical restraints. A restart regenerates them from the input system
+    rather than reloading them, so they must not drift between runs, otherwise
+    the accumulated free energy would be invalidated.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir0:
+        runner0 = Runner(syk_ring_break_mols.clone(), _config(tmpdir0))
+
+    with tempfile.TemporaryDirectory() as tmpdir1:
+        runner1 = Runner(syk_ring_break_mols.clone(), _config(tmpdir1))
+
+    for name in ("morse_hard", "morse_soft"):
+        assert _restraint(runner0._config.restraints, name) == _restraint(
+            runner1._config.restraints, name
+        )
+
+
+def test_reverse_schedule_generates_restraints(syk_ring_break_mols):
+    """
+    Ensure that restraints are also generated for the ring-making direction,
+    which uses the reversed schedule.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = _config(tmpdir, lambda_schedule="reverse_ring_break_morph")
+        runner = Runner(syk_ring_break_mols.clone(), config)
+
+        assert runner._config.restraints is not None
+        assert len(runner._config.restraints) == 2
+
+
+def test_user_restraints_are_not_overridden(syk_ring_break_mols):
+    """
+    Ensure that a user-supplied restraint is left alone, and that the system is
+    not modified behind their back.
+    """
+    import sire as sr
+
+    mols = syk_ring_break_mols.clone()
+
+    restraints = sr.restraints.distance(
+        mols,
+        atoms0=0,
+        atoms1=1,
+        k="10 kcal mol-1 A-2",
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        runner = Runner(mols, _config(tmpdir, restraints=restraints))
+
+        assert len(runner._config.restraints) == 1
+        assert runner._config.restraints[0] == restraints
+
+
+@pytest.mark.parametrize("schedule", ["standard_morph", "charge_scaled_morph"])
+def test_no_restraints_for_other_schedules(schedule, ethane_methanol):
+    """
+    Ensure that Morse restraints are only generated for ring-breaking
+    schedules.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = _config(tmpdir, lambda_schedule=schedule)
+        runner = Runner(ethane_methanol.clone(), config)
+
+        assert runner._config.restraints is None
+
+
+def test_no_broken_bond_raises(ethane_methanol):
+    """
+    Ensure that a clear error is raised when a ring-breaking schedule is used
+    for a perturbation that doesn't break (or form) a bond.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with pytest.raises(RuntimeError, match="Unable to generate Morse restraints"):
+            Runner(ethane_methanol.clone(), _config(tmpdir))
+
+
+def test_already_applied_raises(syk_ring_break_mols):
+    """
+    Ensure that a helpful error is raised if the Morse potential has already
+    been applied to the input system, but the corresponding restraints were not
+    passed via the config. The replacement must not be applied twice.
+    """
+    import sire as sr
+
+    mols = syk_ring_break_mols.clone()
+
+    # Apply the Morse replacement, as a user following the existing workflow
+    # would, but don't pass the restraints to the config.
+    _, mols = sr.restraints.morse_potential(
+        mols,
+        de="150 kcal mol-1",
+        auto_parametrise=True,
+        direct_morse_replacement=True,
+        name="morse_hard",
+    )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        with pytest.raises(RuntimeError, match="already been applied"):
+            Runner(mols, _config(tmpdir))

From 57f3693ef5e8be61a9c89ccd265f0928dc4b3356 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 24 Aug 2026 12:34:23 +0100
Subject: [PATCH 47/67] Document schedules, ABFE, and ring-breaking. [ci skip]

---
 README.md | 81 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 81 insertions(+)

diff --git a/README.md b/README.md
index dd22fec..26966d0 100644
--- a/README.md
+++ b/README.md
@@ -289,6 +289,87 @@ geometry. To override this for all groups:
 somd2 perturbable_system.bss --terminal-flip-frequency "1 ps" --terminal-flip-angle "180 degrees"
 ```
 
+## Lambda schedules
+
+The way that the perturbation is applied across the lambda coordinate is
+controlled by the `--lambda-schedule` option, which defaults to
+`standard_morph`. The available schedules are:
+
+| Schedule | Description |
+| --- | --- |
+| `standard_morph` | Linear interpolation between the two end states. |
+| `charge_scaled_morph` | As above, but with charges scaled at intermediate lambda values. |
+| `annihilate` | Absolute binding free energies, removing all non-bonded interactions. |
+| `decouple` | Absolute binding free energies, removing only intermolecular interactions. |
+| `ring_break_morph` | Ring-breaking perturbations. |
+| `reverse_ring_break_morph` | Ring-making perturbations, i.e. the reverse of the above. |
+
+For the `annihilate`, `decouple`, and ring-breaking schedules, appropriate
+restraints can be generated automatically. See the sections below.
+
+## Absolute binding free energies
+
+Absolute binding free energy (ABFE) calculations are supported using the
+`annihilate` and `decouple` lambda schedules. Both first discharge the ligand,
+then remove its Lennard-Jones interactions: `annihilate` removes all non-bonded
+interactions, including those within the ligand, whereas `decouple` retains the
+intramolecular terms.
+
+```
+somd2 perturbable_system.bss --lambda-schedule decouple
+```
+
+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 both a protein and water. This is done by minimising
+the system, running a short trajectory at lambda = 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
+default the receptor anchor atoms are chosen from the protein backbone; use
+`--restraint-search-receptor-selection` to pass a `Sire` selection string
+instead.
+
+The restraint is written to `abfe_restraint.s3` in the output directory and is
+reloaded on restart, since the accumulated free energy corresponds to that
+particular restraint. The standard state correction is logged and written to
+the metadata of the energy trajectory, so analysis code can apply it without
+needing to scan the log.
+
+> [!NOTE]
+> The Beutler soft-core form, enabled with `--softcore-form beutler`, is only
+> supported with the ABFE schedules, or a custom schedule.
+
+## Ring-breaking perturbations
+
+Perturbations that break (or form) a ring are supported using the
+`ring_break_morph` schedule, or `reverse_ring_break_morph` for the ring-making
+direction.
+
+```
+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`,
+`--morse-soft-well-depth`, and `--morse-soft-force-constant` options.
+
+Unlike the ABFE restraints, these are regenerated on each run rather than being
+cached, since they are derived from the bond parameters alone and are therefore
+identical every time.
+
+> [!NOTE]
+> The defaults are a reasonable starting point, but ring-breaking
+> perturbations are demanding. A non-uniform spacing of lambda values, set with
+> `--lambda-values`, is typically needed to obtain good overlap around the point
+> at which the bond is broken. The
+> [alchemate](https://github.com/akalpokas/alchemate) package provides
+> workflows for iteratively optimising the lambda schedule.
+
 ## Debugging with energy components
 
 To help diagnose simulation instabilities, `SOMD2` can record the potential

From 666898e1ef19dc2909adc6534f213d4b7fd8401e Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 24 Aug 2026 13:52:40 +0100
Subject: [PATCH 48/67] Document intended use of standard_morph schedule. [ci
 skip]

---
 README.md | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 26966d0..5bd3d16 100644
--- a/README.md
+++ b/README.md
@@ -139,7 +139,7 @@ pre-commit run --all-files
 ## Usage
 
 In order to run an alchemical free-energy simulation you will need to
-first create a stream file containing the _perturbable_ system of interest.
+first create a stream file containing the *perturbable* system of interest.
 This can be created using [BioSimSpace](https://github.com/OpenBioSim/biosimspace).
 For example, following the tutorial
 [here](https://biosimspace.openbiosim.org/tutorials/hydration_freenrg.html).
@@ -229,7 +229,7 @@ We also support Replica Exchange with Solute Scaling
 involving conformational changes, e.g. ring flips. This can be enabled
 using the `--rest2-scale` option, which specifies the "temperature" of the
 REST2 region relative to the rest of the system. By default, the REST2 region
-comprises _all_ atoms in perturbable molecules, but can be controlled via the
+comprises *all* atoms in perturbable molecules, but can be controlled via the
 `--rest2-selection` option. This should be a `Sire` selection string that specifies
 additional atoms of interest, i.e. those in regular, non-perturbable molecules.
 If the selection does contain atoms within perturbable molecules, then only
@@ -293,7 +293,8 @@ somd2 perturbable_system.bss --terminal-flip-frequency "1 ps" --terminal-flip-an
 
 The way that the perturbation is applied across the lambda coordinate is
 controlled by the `--lambda-schedule` option, which defaults to
-`standard_morph`. The available schedules are:
+`standard_morph`, which is intended for use with relative binding free
+energy (RBFE) simulations. The available schedules are:
 
 | Schedule | Description |
 | --- | --- |
@@ -482,7 +483,7 @@ For existing users of `somd1`, it's possible to generate input for `somd2` by pa
 `--somd2 True` to the `prepareFEP.py` setup script. This will write a `somd2` compatible
 stream file.
 
-Additionally, `somd2` can be run in `somd1` _compatibility_ mode by passing the
+Additionally, `somd2` can be run in `somd1` *compatibility* mode by passing the
 ``--somd1-compatibility`` command-line option to the `somd2` executable. This ensures
 that the perturbation used is consistent with the approach from `somd1`, i.e.
 it uses the same modifications for bonded-terms involving dummy atoms as `somd1`.

From c681d7fbf0cbcf58ba859d42e2e3f1a2348bc8e8 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 24 Aug 2026 15:57:18 +0100
Subject: [PATCH 49/67] Remove redundant prepareFEP option. [ci skip]

---
 README.md | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 5bd3d16..1cf6e21 100644
--- a/README.md
+++ b/README.md
@@ -479,11 +479,7 @@ package.
 
 ## Note for SOMD1 users
 
-For existing users of `somd1`, it's possible to generate input for `somd2` by passing
-`--somd2 True` to the `prepareFEP.py` setup script. This will write a `somd2` compatible
-stream file.
-
-Additionally, `somd2` can be run in `somd1` *compatibility* mode by passing the
+`somd2` can be run in `somd1` *compatibility* mode by passing the
 ``--somd1-compatibility`` command-line option to the `somd2` executable. This ensures
 that the perturbation used is consistent with the approach from `somd1`, i.e.
 it uses the same modifications for bonded-terms involving dummy atoms as `somd1`.

From 0ebf66f3301773874433cee01f571b85a912c2cf Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 24 Aug 2026 16:00:12 +0100
Subject: [PATCH 50/67] Make package name formatting consistent. [ci skip]

---
 README.md | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/README.md b/README.md
index 1cf6e21..35a69af 100644
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 <p align="center">
     <picture align="center">
-        <img alt="SOMD" src="./.img/somd2.png" width="50%"/>
+        <img alt="SOMD2" src="./.img/somd2.png" width="50%"/>
     </picture>
 </p>
 
@@ -172,20 +172,20 @@ In order to run using GPUs you will first need to set the relevant environment
 variable. For example, to run using 4 CUDA enabled GPUs set `CUDA_VISIBLE_DEVICES=0,1,2,3`
 (for OpenCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively).
 
-By default `SOMD2` will run using the CPU platform, however if the relevant
+By default SOMD2 will run using the CPU platform, however if the relevant
 environment variable has been set (as above) the new platform will be detected
 and set. In the case that this detection fails, or if there are multiple platforms
 available, the `--platform` option can be set (for example `--platform cuda`).
 
-By default, `SOMD2` will automatically manage the distribution of lambda windows
+By default, SOMD2 will automatically manage the distribution of lambda windows
 across all listed devices. In order to restrict the number of devices used
 the `--max-gpus` option can be set, for example setting `--max-gpus 2` while
-`CUDA_VISIBLE_DEVICES` are set as above would restrict `SOMD2` to using only
+`CUDA_VISIBLE_DEVICES` are set as above would restrict SOMD2 to using only
 GPUs 0 and 1.
 
 ## Replica exchange
 
-`SOMD2` supports Hamiltonian replica exchange (HREX) simulations, which can be
+SOMD2 supports Hamiltonian replica exchange (HREX) simulations, which can be
 enabled using the `--replica-exchange` option. By default, dynamics contexts are
 created up-front for all replicas, so this can be memory intensive. As such,
 replica exchange is intended for use on multi-GPU nodes with a large amount of
@@ -208,7 +208,7 @@ must equal `--checkpoint-frequency`.
 For optimal performance, it is recommended that the number of contexts, i.e. the
 number of replicas, or `--max-contexts` if it is set, be a multiple of the number
 of GPUs, and no smaller than the number of GPUs multiplied by the
-oversubscription factor. `SOMD2` will warn you if this isn't the case.
+oversubscription factor. SOMD2 will warn you if this isn't the case.
 
 Changing the lambda value of a context requires it to be reinitialised whenever a
 constrained bond length actually perturbs with lambda, which is slow. If this
@@ -373,7 +373,7 @@ identical every time.
 
 ## Debugging with energy components
 
-To help diagnose simulation instabilities, `SOMD2` can record the potential
+To help diagnose simulation instabilities, SOMD2 can record the potential
 energy contribution from each OpenMM force group. This is enabled with the
 `--save-energy-components` flag:
 
@@ -400,7 +400,7 @@ The recording interval depends on the runner and active samplers:
 
 ## Copying output files during a simulation
 
-When `SOMD2` writes checkpoint files it acquires an exclusive
+When SOMD2 writes checkpoint files it acquires an exclusive
 [file lock](https://py-filelock.readthedocs.io) on `somd2.lock` inside the output
 directory. This guarantees that checkpoint files are always in a consistent
 state on disk.
@@ -426,7 +426,7 @@ with FileLock("/path/to/output/somd2.lock"):
 ```
 
 > [!NOTE]
-> The `--timeout` option (default: `300 s`) controls how long `SOMD2` will
+> The `--timeout` option (default: `300 s`) controls how long SOMD2 will
 > wait to re-acquire the lock after your copy completes. If you hold the lock
 > for longer than this, the simulation will raise a `Timeout` error.
 
@@ -473,18 +473,18 @@ be chosen as a trade off between accuracy and computational cost. A value of aro
 We support modification of ghost atom bonded terms to avoid spurious coupling
 to the physical system using the approach described in
 [this](https://pubs.acs.org/doi/10.1021/acs.jctc.0c01328) paper.
-These are enabled by default, but can be disabled using the ``--no-ghost-modifications``
+These are enabled by default, but can be disabled using the `--no-ghost-modifications`
 option. Modifications are implemented using the [ghostly](https://github.com/OpenBioSim/ghostly)
 package.
 
 ## Note for SOMD1 users
 
-`somd2` can be run in `somd1` *compatibility* mode by passing the
-``--somd1-compatibility`` command-line option to the `somd2` executable. This ensures
-that the perturbation used is consistent with the approach from `somd1`, i.e.
-it uses the same modifications for bonded-terms involving dummy atoms as `somd1`.
+SOMD2 can be run in SOMD1 *compatibility* mode by passing the
+`--somd1-compatibility` command-line option to the `somd2` executable. This ensures
+that the perturbation used is consistent with the approach from SOMD1, i.e.
+it uses the same modifications for bonded-terms involving dummy atoms as SOMD1.
 
-Finally, it is also possible to run `somd2` using an existing `somd1` perturbation
+Finally, it is also possible to run SOMD2 using an existing SOMD1 perturbation
 file. To do so, you will also need to create a stream file representing the
 λ = 0 state. For existing input generated by `prepareFEP.py`, this can be done as
 follows. (This assumes that the output has a prefix `somd1`.)
@@ -501,7 +501,7 @@ BSS.Stream.save(system, "somd1")
 
 (This will write a stream file called `somd1.bss`.)
 
-This can then be run with `somd2` using the following:
+This can then be run with SOMD2 using the following:
 
 ```
 somd2 somd1.bss --pert-file somd1.pert --somd1-compatibility
@@ -510,7 +510,7 @@ somd2 somd1.bss --pert-file somd1.pert --somd1-compatibility
 (This only shows the limited options required. Others will take default values and can be set accordingly.)
 
 If you want to load an existing system from a perturbation file and use the
-new `somd2` [ghost atom bonded-term modifications](https://github.com/OpenBioSim/ghostly),
+new SOMD2 [ghost atom bonded-term modifications](https://github.com/OpenBioSim/ghostly),
 then simply omit the `--somd1-compatibility` option.
 
 ## GPU oversubscription
@@ -532,7 +532,7 @@ More details on MPS, including tuning options, can be found in the following
 
 ## Python API
 
-`SOMD2` can also be used as a Python API, allowing it to be embedded
+SOMD2 can also be used as a Python API, allowing it to be embedded
 within other Python scripts.
 
 ## Known issues

From 72894403c5ef2a30b67ae57e5d824edcd75959ab Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 24 Aug 2026 17:58:51 +0100
Subject: [PATCH 51/67] Correct version scheme. [ci skip]

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 35a69af..166cb0a 100644
--- a/README.md
+++ b/README.md
@@ -96,7 +96,7 @@ somd2 --help
 ### Keeping up to date
 
 During a development cycle the OpenBioSim packages are pinned only to a
-`X.X.X.dev` version, not to a specific build. `somd2` and its dependencies
+`YYYY.N.0.dev` version, not to a specific build. `somd2` and its dependencies
 therefore need to be kept in sync, so always update the whole stack together
 rather than `somd2` alone.
 

From beb490cf0cbe724345f6d521e80217864e818b5d Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 26 Aug 2026 11:25:48 +0100
Subject: [PATCH 52/67] Remove redundant kappa lever equations from
 ring-breaking/making schedules.

---
 CHANGELOG.md                       |   1 +
 src/somd2/_utils/_schedules.py     |  42 +++----
 tests/schedules/test_ring_break.py | 169 ++++++++++++++---------------
 3 files changed, 99 insertions(+), 113 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ddfcb67..112246d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,7 @@ Changelog
 * Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Add a `precision` option for GPU platforms, defaulting to `mixed` [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Add support for generating Morse restraints for ring-breaking perturbations.
+* Remove the unused `kappa` lever equations from the ring-breaking/making lambda schedules.
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/_utils/_schedules.py b/src/somd2/_utils/_schedules.py
index f40dd52..1f1f0e2 100644
--- a/src/somd2/_utils/_schedules.py
+++ b/src/somd2/_utils/_schedules.py
@@ -239,13 +239,19 @@ def ring_break_morph():
     Three stages: potential_swap → restraints_off → morph.
 
     During restraints_off the Morse restraint ramps off (morse_soft: 1→0) while
-    the ring-break softcore simultaneously ramps on (alpha: 1→0, kappa: 0→1),
-    equations mirror ring-break so that ``ring_break_morph().reverse()`` is the
-    providing a smooth handover with no gap between the two forces. The ring-make
-    correct schedule for the ring-making direction (used by
-    :func:`reverse_ring_break_morph`). Because ring_break_morph is only used for
-    ring-breaking perturbations (no ring-make force present), the ring-make
-    equations have no effect on forward simulations.
+    the ring-break softcore LJ simultaneously ramps on (alpha: 1→0), providing a
+    smooth handover with no gap between the two forces.
+
+    Coulomb is decoupled from the LJ and driven by its own coul_kappa lever,
+    which is held at zero through both bonded stages and ramps 0→1 during morph
+    only, once the softcore LJ has already separated the pair.
+
+    The ring-make equations mirror ring-break so that
+    ``ring_break_morph().reverse()`` is the correct schedule for the ring-making
+    direction (used by :func:`reverse_ring_break_morph`). Because
+    ring_break_morph is only used for ring-breaking perturbations (no ring-make
+    force present), the ring-make equations have no effect on forward
+    simulations.
 
     Returns
     -------
@@ -257,8 +263,8 @@ def ring_break_morph():
 
     s = _LambdaSchedule.standard_morph()
 
-    # restraints_off [1/3, 2/3): Morse ramps off while ring-break softcore ramps
-    # on simultaneously (alpha: 1→0, kappa: 0→1). Bonded terms (angles, torsions)
+    # restraints_off [1/3, 2/3): Morse ramps off while the ring-break softcore LJ
+    # ramps on simultaneously (alpha: 1→0). Bonded terms (angles, torsions)
     # interpolate initial→final over the same stage. ring-make mirrors ring-break
     # so that after .reverse(), the ring-make softcore ramps off as morse_soft ramps
     # on in the reversed restraints_off stage, correct for ring-making perturbations.
@@ -290,15 +296,9 @@ def ring_break_morph():
     s.set_equation(
         stage="restraints_off", force="ring-break", lever="alpha", equation=1 - s.lam()
     )
-    s.set_equation(
-        stage="restraints_off", force="ring-break", lever="kappa", equation=s.lam()
-    )
     s.set_equation(
         stage="restraints_off", force="ring-make", lever="alpha", equation=1 - s.lam()
     )
-    s.set_equation(
-        stage="restraints_off", force="ring-make", lever="kappa", equation=s.lam()
-    )
 
     s.prepend_stage("potential_swap", s.initial())
     s.set_equation(stage="potential_swap", lever="morse_hard", equation=1 - s.lam())
@@ -322,15 +322,11 @@ def ring_break_morph():
     s.set_equation(
         stage="potential_swap", force="ring-break", lever="alpha", equation=1
     )
-    s.set_equation(
-        stage="potential_swap", force="ring-break", lever="kappa", equation=0
-    )
     s.set_equation(stage="potential_swap", force="ring-make", lever="alpha", equation=1)
-    s.set_equation(stage="potential_swap", force="ring-make", lever="kappa", equation=0)
 
     # morph [2/3, 1]: standard nonbonded morphing with ring-break/ring-make fixed
-    # at fully open (kappa=1, alpha=0). ring-make mirrors ring-break so .reverse()
-    # gives kappa=1 at lam=0 of the reversed morph stage (ring-making start).
+    # at fully open (alpha=0). ring-make mirrors ring-break so .reverse() gives
+    # alpha=0 at lam=0 of the reversed morph stage (ring-making start).
     s.set_equation(stage="morph", lever="morse_hard", equation=0)
     s.set_equation(stage="morph", lever="morse_soft", equation=0)
     s.set_equation(stage="morph", lever="bond_k", equation=s.final())
@@ -340,9 +336,7 @@ def ring_break_morph():
     s.set_equation(stage="morph", lever="torsion_k", equation=s.final())
     s.set_equation(stage="morph", lever="torsion_phase", equation=s.final())
     s.set_equation(stage="morph", force="ring-break", lever="alpha", equation=0)
-    s.set_equation(stage="morph", force="ring-break", lever="kappa", equation=1)
     s.set_equation(stage="morph", force="ring-make", lever="alpha", equation=0)
-    s.set_equation(stage="morph", force="ring-make", lever="kappa", equation=1)
 
     # coul_kappa: zero through both bonded stages so the CLJ exception carries no
     # charge while atoms are at covalent distances; ramps 0→1 in morph only once
@@ -381,7 +375,7 @@ def reverse_ring_break_morph():
     This schedule is correct for two equivalent use-cases:
 
     1. A ring-making perturbation run with ``swap_end_states=False``: the
-       ring-make softcore force (kappa=1 at λ=0, ramping to 0) is controlled
+       ring-make softcore force (alpha=0 at λ=0, ramping to 1) is controlled
        directly by the ring-make lever equations.
     2. A ring-breaking perturbation run with ``swap_end_states=True`` (the
        runner reverses the schedule automatically, yielding the same effective
diff --git a/tests/schedules/test_ring_break.py b/tests/schedules/test_ring_break.py
index bfc0146..730c3c0 100644
--- a/tests/schedules/test_ring_break.py
+++ b/tests/schedules/test_ring_break.py
@@ -6,7 +6,7 @@
     reason="openmm support is not available",
 )
 
-# Energy threshold (kcal/mol) for the "active" state: kappa=1 should give
+# Energy threshold (kcal/mol) for the "active" state: alpha=0 should give
 # clearly non-zero CustomBondForce energy.
 _ACTIVE_THRESHOLD = 0.1
 
@@ -122,59 +122,60 @@ def test_reverse_has_ring_make_not_ring_break(reverse_dynamics):
     )
 
 
-# ── schedule kappa/alpha tests ────────────────────────────────────────────────
+# ── schedule alpha/coul_kappa tests ───────────────────────────────────────────
 #
-# These tests verify kappa and alpha values by calling schedule.morph() directly,
-# using the same initial/final values that lambdalever passes in production.
-# They are completely independent of Sire's energy formula and will continue to
-# work correctly regardless of changes to the softcore implementation.
-
-# ring_break_morph kappa/alpha points (3 equal stages: [0,1/3), [1/3,2/3), [2/3,1]):
-#   λ=0.00  potential_swap start     kappa=0, alpha=1
-#   λ=0.15  potential_swap mid       kappa=0, alpha=1
-#   λ=1/3   restraints_off start     kappa=0, alpha=1  (within-stage lam=0)
-#   λ=0.45  restraints_off mid       kappa=0.35, alpha=0.65  (within-stage lam=0.35)
-#   λ=0.50  restraints_off mid       kappa=0.5, alpha=0.5   (within-stage lam=0.5)
-#   λ=0.55  restraints_off mid       kappa=0.65, alpha=0.35  (within-stage lam=0.65)
-#   λ=0.60  restraints_off near end  kappa=0.8, alpha=0.2   (within-stage lam=0.8)
-#   λ=2/3   morph start              kappa=1, alpha=0
-#   λ=0.85  morph mid                kappa=1, alpha=0
-#   λ=1.00  morph end                kappa=1, alpha=0
-_FWD_KAPPA_ALPHA = [
-    (0.00, 0.0, 1.0),
-    (0.15, 0.0, 1.0),
-    (1 / 3, 0.0, 1.0),
-    (0.45, 0.35, 0.65),
-    (0.50, 0.5, 0.5),
-    (0.55, 0.65, 0.35),
-    (0.60, 0.8, 0.2),
-    (2 / 3, 1.0, 0.0),
-    (0.85, 1.0, 0.0),
-    (1.00, 1.0, 0.0),
+# These tests verify alpha and coul_kappa values by calling schedule.morph()
+# directly, using the same initial/final values that lambdalever passes in
+# production. They are completely independent of Sire's energy formula and will
+# continue to work correctly regardless of changes to the softcore
+# implementation.
+
+# ring_break_morph alpha points (3 equal stages: [0,1/3), [1/3,2/3), [2/3,1]):
+#   λ=0.00  potential_swap start     alpha=1
+#   λ=0.15  potential_swap mid       alpha=1
+#   λ=1/3   restraints_off start     alpha=1     (within-stage lam=0)
+#   λ=0.45  restraints_off mid       alpha=0.65  (within-stage lam=0.35)
+#   λ=0.50  restraints_off mid       alpha=0.5   (within-stage lam=0.5)
+#   λ=0.55  restraints_off mid       alpha=0.35  (within-stage lam=0.65)
+#   λ=0.60  restraints_off near end  alpha=0.2   (within-stage lam=0.8)
+#   λ=2/3   morph start              alpha=0
+#   λ=0.85  morph mid                alpha=0
+#   λ=1.00  morph end                alpha=0
+_FWD_ALPHA = [
+    (0.00, 1.0),
+    (0.15, 1.0),
+    (1 / 3, 1.0),
+    (0.45, 0.65),
+    (0.50, 0.5),
+    (0.55, 0.35),
+    (0.60, 0.2),
+    (2 / 3, 0.0),
+    (0.85, 0.0),
+    (1.00, 0.0),
 ]
 
-# reverse_ring_break_morph ring-make kappa/alpha points (mirror of forward):
-#   λ=0.00  reversed morph start      kappa=1, alpha=0
-#   λ=0.15  reversed morph mid        kappa=1, alpha=0
-#   λ=1/3   reversed restraints_off start  kappa=1, alpha=0  (within-stage lam=0)
-#   λ=0.45  reversed restraints_off mid    kappa=0.65, alpha=0.35
-#   λ=0.50  reversed restraints_off mid    kappa=0.5, alpha=0.5
-#   λ=0.55  reversed restraints_off mid    kappa=0.35, alpha=0.65
-#   λ=0.60  reversed restraints_off near end  kappa=0.2, alpha=0.8
-#   λ=2/3   reversed potential_swap start  kappa=0, alpha=1
-#   λ=0.85  reversed potential_swap mid    kappa=0, alpha=1
-#   λ=1.00  reversed potential_swap end    kappa=0, alpha=1
-_REV_KAPPA_ALPHA = [
-    (0.00, 1.0, 0.0),
-    (0.15, 1.0, 0.0),
-    (1 / 3, 1.0, 0.0),
-    (0.45, 0.65, 0.35),
-    (0.50, 0.5, 0.5),
-    (0.55, 0.35, 0.65),
-    (0.60, 0.2, 0.8),
-    (2 / 3, 0.0, 1.0),
-    (0.85, 0.0, 1.0),
-    (1.00, 0.0, 1.0),
+# reverse_ring_break_morph ring-make alpha points (mirror of forward):
+#   λ=0.00  reversed morph start      alpha=0
+#   λ=0.15  reversed morph mid        alpha=0
+#   λ=1/3   reversed restraints_off start  alpha=0  (within-stage lam=0)
+#   λ=0.45  reversed restraints_off mid    alpha=0.35
+#   λ=0.50  reversed restraints_off mid    alpha=0.5
+#   λ=0.55  reversed restraints_off mid    alpha=0.65
+#   λ=0.60  reversed restraints_off near end  alpha=0.8
+#   λ=2/3   reversed potential_swap start  alpha=1
+#   λ=0.85  reversed potential_swap mid    alpha=1
+#   λ=1.00  reversed potential_swap end    alpha=1
+_REV_ALPHA = [
+    (0.00, 0.0),
+    (0.15, 0.0),
+    (1 / 3, 0.0),
+    (0.45, 0.35),
+    (0.50, 0.5),
+    (0.55, 0.65),
+    (0.60, 0.8),
+    (2 / 3, 1.0),
+    (0.85, 1.0),
+    (1.00, 1.0),
 ]
 
 # ring_break_morph coul_kappa points (initial=0, final=1):
@@ -213,43 +214,35 @@ def test_reverse_has_ring_make_not_ring_break(reverse_dynamics):
 ]
 
 
-@pytest.mark.parametrize("lam,expected_kappa,expected_alpha", _FWD_KAPPA_ALPHA)
-def test_ring_break_morph_schedule(lam, expected_kappa, expected_alpha):
+@pytest.mark.parametrize("lam,expected_alpha", _FWD_ALPHA)
+def test_ring_break_morph_schedule(lam, expected_alpha):
     """
-    ring_break_morph() produces the correct ring-break kappa and alpha at each λ.
+    ring_break_morph() produces the correct ring-break alpha at each λ.
 
-    Uses lambdalever's initial/final values (kappa: 0→1, alpha: 1→0) to ensure
-    the test matches production behaviour exactly.
+    Uses lambdalever's initial/final values (alpha: 1→0) to ensure the test
+    matches production behaviour exactly.
     """
     from somd2._utils._schedules import ring_break_morph
 
     s = ring_break_morph()
-    kappa = s.morph("ring-break", "kappa", 0.0, 1.0, lam)
     alpha = s.morph("ring-break", "alpha", 1.0, 0.0, lam)
-    assert abs(kappa - expected_kappa) < 1e-10, (
-        f"ring-break kappa={kappa:.8f} at λ={lam:.4f}, expected {expected_kappa}"
-    )
     assert abs(alpha - expected_alpha) < 1e-10, (
         f"ring-break alpha={alpha:.8f} at λ={lam:.4f}, expected {expected_alpha}"
     )
 
 
-@pytest.mark.parametrize("lam,expected_kappa,expected_alpha", _REV_KAPPA_ALPHA)
-def test_reverse_ring_break_morph_schedule(lam, expected_kappa, expected_alpha):
+@pytest.mark.parametrize("lam,expected_alpha", _REV_ALPHA)
+def test_reverse_ring_break_morph_schedule(lam, expected_alpha):
     """
-    reverse_ring_break_morph() produces the correct ring-make kappa and alpha at each λ.
+    reverse_ring_break_morph() produces the correct ring-make alpha at each λ.
 
-    Uses lambdalever's initial/final values (kappa: 1→0, alpha: 0→1) to ensure
-    the test matches production behaviour exactly.
+    Uses lambdalever's initial/final values (alpha: 0→1) to ensure the test
+    matches production behaviour exactly.
     """
     from somd2._utils._schedules import reverse_ring_break_morph
 
     s = reverse_ring_break_morph()
-    kappa = s.morph("ring-make", "kappa", 1.0, 0.0, lam)
     alpha = s.morph("ring-make", "alpha", 0.0, 1.0, lam)
-    assert abs(kappa - expected_kappa) < 1e-10, (
-        f"ring-make kappa={kappa:.8f} at λ={lam:.4f}, expected {expected_kappa}"
-    )
     assert abs(alpha - expected_alpha) < 1e-10, (
         f"ring-make alpha={alpha:.8f} at λ={lam:.4f}, expected {expected_alpha}"
     )
@@ -297,30 +290,30 @@ def test_reverse_ring_break_morph_coul_kappa(lam, expected_coul_kappa):
 @pytest.mark.parametrize("lam", [2 / 3, 1.0])
 def test_ring_break_active_in_morph(forward_dynamics, lam):
     """
-    Ring-break energy is clearly non-zero (kappa=1) throughout the morph stage.
+    Ring-break energy is clearly non-zero (alpha=0) throughout the morph stage.
     """
     e = _force_energy_kcal(forward_dynamics, lam, "ring-break")
     assert abs(e) > _ACTIVE_THRESHOLD, (
         f"ring-break energy {e:.4f} kcal/mol at λ={lam:.4f} is below active "
-        f"threshold {_ACTIVE_THRESHOLD} kcal/mol (kappa should be 1)"
+        f"threshold {_ACTIVE_THRESHOLD} kcal/mol (alpha should be 0)"
     )
 
 
 def test_ring_make_active_at_lambda_zero(reverse_dynamics):
     """
-    Ring-make energy is non-zero at λ=0: the morph stage fixes kappa=1
+    Ring-make energy is non-zero at λ=0: the morph stage fixes alpha=0
     so the ring-make interaction is fully on from the start.
     """
     e = _force_energy_kcal(reverse_dynamics, 0.0, "ring-make")
     assert abs(e) > _ACTIVE_THRESHOLD, (
         f"ring-make energy {e:.4f} kcal/mol at λ=0 is below active threshold "
-        f"{_ACTIVE_THRESHOLD} kcal/mol (kappa should be 1 in morph stage)"
+        f"{_ACTIVE_THRESHOLD} kcal/mol (alpha should be 0 in morph stage)"
     )
 
 
 def test_ring_make_inactive_at_lambda_one(reverse_dynamics):
     """
-    Ring-make energy is near-zero at λ=1 (potential_swap end, kappa=0).
+    Ring-make energy is near-zero at λ=1 (potential_swap end, alpha=1).
 
     At λ=1 the system is at the ring-open end state; the hard-hard correction
     term in the CustomBondForce is small because the pair is at nonbonded
@@ -329,25 +322,25 @@ def test_ring_make_inactive_at_lambda_one(reverse_dynamics):
     e = _force_energy_kcal(reverse_dynamics, 1.0, "ring-make")
     assert abs(e) < _ACTIVE_THRESHOLD, (
         f"ring-make energy {e:.4f} kcal/mol at λ=1 exceeds threshold "
-        f"{_ACTIVE_THRESHOLD} kcal/mol (kappa should be 0)"
+        f"{_ACTIVE_THRESHOLD} kcal/mol (alpha should be 1)"
     )
 
 
 # ── energy symmetry tests ─────────────────────────────────────────────────────
 #
 # The invariant ring_break_morph().reverse() == reverse_ring_break_morph() means
-# that the softcore kappa/alpha values at (forward, λ) and (reverse, 1-λ) are
-# equal.  Both forces act on the same bond (the original ring_breaking_bond,
-# which swap_end_states=True maps to ring_making_pairs), so the energies must
-# also match.  The hard-hard correction appears identically on both sides and
+# that the softcore alpha values at (forward, λ) and (reverse, 1-λ) are equal.
+# Both forces act on the same bond (the original ring_breaking_bond, which
+# swap_end_states=True maps to ring_making_pairs), so the energies must also
+# match.  The hard-hard correction appears identically on both sides and
 # cancels in the comparison, making this test robust to formula changes.
 #
 # Test points span zero and non-zero energy regions:
-#   λ=0.0  → forward kappa=0, reverse at 1-λ=1.0 kappa=0 (both ≈0)
-#   λ=0.55 → forward restraints_off (kappa=0.65), reverse restraints_off at 0.45 (kappa=0.65)
-#   λ=2/3  → forward morph start (kappa=1), reverse restraints_off start at 1/3 (kappa=1)
-#   λ=0.85 → forward morph (kappa=1), reverse reversed-morph at 0.15 (kappa=1)
-#   λ=1.0  → forward morph end (kappa=1), reverse at 0.0 reversed-morph (kappa=1)
+#   λ=0.0  → forward alpha=1, reverse at 1-λ=1.0 alpha=1 (both ≈0)
+#   λ=0.55 → forward restraints_off (alpha=0.35), reverse restraints_off at 0.45 (alpha=0.35)
+#   λ=2/3  → forward morph start (alpha=0), reverse restraints_off start at 1/3 (alpha=0)
+#   λ=0.85 → forward morph (alpha=0), reverse reversed-morph at 0.15 (alpha=0)
+#   λ=1.0  → forward morph end (alpha=0), reverse at 0.0 reversed-morph (alpha=0)
 
 
 @pytest.mark.parametrize("lam", [0.0, 0.55, 2 / 3, 0.85, 1.0])
@@ -356,7 +349,7 @@ def test_energy_symmetry_forward_reverse(forward_dynamics, reverse_dynamics, lam
     Single-point energy symmetry: E_ring_break_forward(λ) == E_ring_make_reverse(1-λ).
 
     Verifies that reverse_ring_break_morph() == ring_break_morph().reverse() and
-    that the mirrored kappa/alpha produce identical corrections on the same bond.
+    that the mirrored alpha values produce identical corrections on the same bond.
     """
     e_fwd = _force_energy_kcal(forward_dynamics, lam, "ring-break")
     e_rev = _force_energy_kcal(reverse_dynamics, 1.0 - lam, "ring-make")
@@ -373,8 +366,8 @@ def test_schedule_symmetry():
     reverse_ring_break_morph() must equal ring_break_morph().reverse().
 
     Checks that the simplified implementation produces identical schedules by
-    comparing kappa values at a dense grid of lambda points using the default
-    initial/final values that lambdalever passes for ring-break kappa.
+    comparing lever values at a dense grid of lambda points using the default
+    initial/final values that lambdalever passes for each lever.
     """
     from somd2._utils._schedules import ring_break_morph, reverse_ring_break_morph
 
@@ -385,10 +378,8 @@ def test_schedule_symmetry():
     test_lambdas = [i / 20 for i in range(21)]
     for lam in test_lambdas:
         for force, lever, init, fin in [
-            ("ring-break", "kappa", 0.0, 1.0),
             ("ring-break", "alpha", 1.0, 0.0),
             ("ring-break", "coul_kappa", 0.0, 1.0),
-            ("ring-make", "kappa", 1.0, 0.0),
             ("ring-make", "alpha", 0.0, 1.0),
             ("ring-make", "coul_kappa", 1.0, 0.0),
         ]:

From e8c1b13fbb727163518707fa56fea455d6e46a76 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 28 Aug 2026 09:07:08 +0100
Subject: [PATCH 53/67] Don't silently swallow errors when adding coalchemical
 restraints. [ci skip]

---
 src/somd2/runner/_base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 2f4eea8..4f2517c 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -1622,10 +1622,10 @@ def _create_alchemical_ions(
                     restraint_distance,
                 )
 
-                try:
-                    restraints.add(restraint)
-                except:
+                if restraints is None:
                     restraints = restraint
+                else:
+                    restraints.add(restraint)
 
             # Update the system.
             system.update(merged)

From 95eb838c22c2b7b7dc0e4972260c0ee72363244a Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 28 Aug 2026 09:10:51 +0100
Subject: [PATCH 54/67] Add section on charge-change perturbations. [ci skip]

---
 README.md | 40 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 40 insertions(+)

diff --git a/README.md b/README.md
index 166cb0a..c79d5c9 100644
--- a/README.md
+++ b/README.md
@@ -371,6 +371,46 @@ identical every time.
 > [alchemate](https://github.com/akalpokas/alchemate) package provides
 > workflows for iteratively optimising the lambda schedule.
 
+## Charge-change perturbations
+
+Perturbations that change the net charge of the system are handled
+automatically using the co-alchemical ion method. The charge difference between
+the two end states is computed when the system is loaded, and, if it is
+non-zero, a number of water molecules equal to the absolute charge difference
+are perturbed into counter-ions alongside the main perturbation, keeping the
+total charge constant at every lambda value. The waters furthest from the
+perturbable molecule are chosen, and the ion type is picked to offset the
+charge change, re-using the parameters of a free ion already present in the
+system where possible.
+
+No options are needed to enable this. The automatically detected value can be
+overridden with `--charge-difference`, which takes the perturbed charge minus
+the reference charge:
+
+```
+somd2 perturbable_system.bss --charge-difference -1
+```
+
+The molecules chosen as alchemical ions are written to `alchemical_ions.npz` in
+the output directory and reused on restart, so that ion selection does not
+depend on anything that might have changed between runs.
+
+Since a co-alchemical ion is only meaningful in the bulk, SOMD2 can restrain it
+away from the perturbable region. Passing a distance to
+`--coalchemical-restraint-dist` adds an inverse-distance restraint between each
+ion and the atom closest to the centre of geometry of the perturbable molecule,
+preventing the ion from drifting into the binding site and interacting with the
+protein or ligand:
+
+```
+somd2 perturbable_system.bss --coalchemical-restraint-dist "10 A"
+```
+
+> [!NOTE]
+> These restraints are *added* to any others in use. Restraints passed via the
+> Python API, and those generated automatically for the ABFE and ring-breaking
+> schedules described above, are all retained.
+
 ## Debugging with energy components
 
 To help diagnose simulation instabilities, SOMD2 can record the potential

From 97092ac545f103880d7df7339f6ab6b0eaaa3346 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 28 Aug 2026 09:20:37 +0100
Subject: [PATCH 55/67] Add features section to the top of the README. [ci
 skip]

---
 README.md | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/README.md b/README.md
index c79d5c9..0590863 100644
--- a/README.md
+++ b/README.md
@@ -13,6 +13,24 @@
 Open-source GPU accelerated molecular dynamics engine for alchemical free-energy
 simulations. Built on top of [Sire](https://github.com/OpenBioSim/sire) and [OpenMM](https://github.com/openmm/openmm).
 
+## Features
+
+- **Perturbations**: relative binding free energies,
+  [absolute binding free energies](#absolute-binding-free-energies),
+  [ring-breaking](#ring-breaking-perturbations),
+  [charge-change](#charge-change-perturbations), and protein mutations.
+- **[GCMC](#gcmc)**: grand canonical Monte Carlo water sampling.
+- **[Replica exchange](#replica-exchange)**: Hamiltonian replica exchange
+  between lambda windows.
+- **[REST2](#rest2)**: replica exchange with solute scaling.
+- **[Terminal ring flips](#terminal-ring-flip-monte-carlo)**: Monte Carlo moves
+  to improve sampling of terminal aromatic rings.
+- **[Ghost atom modifications](#ghost-atom-modifications)**: modification of
+  ghost atom bonded terms to avoid spurious coupling to the physical system.
+- **[Multiple GPUs](#running-somd2-using-one-or-more-gpus)**: lambda windows are
+  distributed across the available devices, with optional
+  [oversubscription](#gpu-oversubscription).
+
 ## Installation
 
 ### Conda package

From d1bc0d60a24408daf6835ce1fcb48513b1512bd1 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 28 Aug 2026 09:40:03 +0100
Subject: [PATCH 56/67] Add link to somd2_examples repository. [ci skip]

---
 README.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/README.md b/README.md
index 0590863..c9f1e65 100644
--- a/README.md
+++ b/README.md
@@ -184,6 +184,10 @@ An example perturbable system for a methane to ethanol perturbation in solvent
 can be found [here](https://sire.openbiosim.org/m/merged_molecule.s3.bz2).
 This is a `bzip2` compressed file that will need to be extracted before use.
 
+A larger collection of input files and end-to-end tutorials, covering everything
+from a simple charge-change validation system to full case studies, can be found
+in the [somd2_examples](https://github.com/OpenBioSim/somd2_examples) repository.
+
 ### Running SOMD2 using one or more GPUs
 
 In order to run using GPUs you will first need to set the relevant environment

From 3384ca75512b83dfc2f326bdb8076eb24f0f2b4e Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 28 Aug 2026 10:52:33 +0100
Subject: [PATCH 57/67] Accept stream file paths for restraints and lambda
 schedules.

---
 README.md                   |  60 ++++++++++++
 src/somd2/config/_config.py | 180 ++++++++++++++++++++++++++----------
 tests/runner/test_config.py | 103 +++++++++++++++++++++
 3 files changed, 296 insertions(+), 47 deletions(-)

diff --git a/README.md b/README.md
index c9f1e65..828a9f0 100644
--- a/README.md
+++ b/README.md
@@ -205,6 +205,48 @@ the `--max-gpus` option can be set, for example setting `--max-gpus 2` while
 `CUDA_VISIBLE_DEVICES` are set as above would restrict SOMD2 to using only
 GPUs 0 and 1.
 
+## Restarting
+
+A simulation can be continued from the files in its output directory using the
+`--restart` option:
+
+```
+somd2 perturbable_system.bss --restart --output-directory output
+```
+
+Each λ window (or replica) resumes from its most recent checkpoint. The
+configuration used for the original run is written to `config.yaml` in the
+output directory, controlled by `--write-config`, which is enabled by default.
+This file is required in order to restart, since the current configuration is
+validated against it.
+
+Only a limited set of options may be changed on restart. Broadly, anything that
+would change the perturbation or the Hamiltonian is fixed, whereas options
+controlling how long to run for, what to write out, and which hardware to use
+can be varied. The most useful of these is `--runtime`, which allows a completed
+simulation to be extended. SOMD2 will tell you which option is at fault if you
+change one that isn't allowed.
+
+> [!NOTE]
+> If the most recent checkpoint files are incomplete or corrupt, for example
+> when recovering from a crash, pass `--use-backup` to restart from the last
+> but one checkpoint instead.
+
+## Hydrogen mass repartitioning
+
+By default SOMD2 applies hydrogen mass repartitioning (HMR), scaling hydrogen
+masses by the factor given by `--h-mass-factor` (default 1.5). This is what
+allows the default `--timestep` of 4 fs.
+
+If the masses of your input system have already been repartitioned, or you want
+to use a different repartitioning scheme, pass `--no-hmr` so that the masses of
+the input system are used as they are.
+
+> [!NOTE]
+> A 4 fs timestep is not stable without repartitioning, so if you disable HMR
+> you will need to reduce `--timestep` accordingly, or supply a system that has
+> already been repartitioned.
+
 ## Replica exchange
 
 SOMD2 supports Hamiltonian replica exchange (HREX) simulations, which can be
@@ -597,6 +639,24 @@ More details on MPS, including tuning options, can be found in the following
 SOMD2 can also be used as a Python API, allowing it to be embedded
 within other Python scripts.
 
+A few options take objects rather than values, so cannot be set directly on the
+command line. A custom lambda schedule can be passed to `lambda_schedule` as a
+`sire.cas.LambdaSchedule`, rather than one of the named schedules, and
+user-defined restraints can be passed to `restraints`.
+
+Both options can also be set via a YAML configuration file, where they are
+stored as a hex string of the serialised object. This is the form written to
+`config.yaml`, so the simplest way to obtain one is to configure the option in
+Python, run a simulation, and re-use the value from the resulting file.
+
+Alternatively, both accept a path to a [Sire](https://github.com/OpenBioSim/sire)
+stream file containing the serialised object, which can be written with
+`sire.stream.save`:
+
+```
+somd2 perturbable_system.bss --lambda-schedule my_schedule.s3 --restraints my_restraints.s3
+```
+
 ## Known issues
 
 If using the regular `Runner` class via the Python API, then you will need to
diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index 28f2470..cfa0a99 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -78,11 +78,17 @@ class Config:
         "precision": ["single", "mixed", "double"],
     }
 
+    # Options that advertise a set of choices, but which also accept other
+    # forms, e.g. the path to a stream file. These are validated by the setter,
+    # rather than by argparse.
+    _open_choices = ["lambda_schedule"]
+
     # A dictionary of nargs for the various options.
     _nargs = {
         "lambda_values": "+",
         "lambda_energy": "+",
         "rest2_scale": "+",
+        "restraints": "+",
     }
 
     def __init__(
@@ -239,8 +245,10 @@ def __init__(
             then this will be set to the same as 'lambda_values', or the values
             defined by 'num_lambda' if 'lambda_values' is not set.
 
-        lambda_schedule: str
-            Lambda schedule to use for alchemical free energy simulations.
+        lambda_schedule: str, sire.cas.LambdaSchedule
+            Lambda schedule to use for alchemical free energy simulations. This
+            can be the name of one of the standard schedules, or the path to a
+            Sire stream file containing a custom LambdaSchedule.
 
         charge_scale_factor: float
             Factor by which to scale charges for charge scaled morph.
@@ -256,9 +264,10 @@ def __init__(
             The soft-core shift-delta parameter. This is used to soften the
             Lennard-Jones interaction.
 
-        restraints: sire.mm._MM.Restraints
-            A single set of restraints, or a list of sets of restraints that
-            will be applied to the atoms during the simulation.
+        restraints: str, sire.mm._MM.Restraints
+            One or more paths to Sire stream files containing the sets of
+            restraints that will be applied to the atoms during the simulation.
+            A stream file may hold a single set, or a list of sets.
 
         constraint: str
             Constraint type to use for non-perturbable molecules.
@@ -1180,44 +1189,51 @@ def lambda_schedule(self, lambda_schedule):
                     "'lambda_schedule' must be of type 'str' or 'LambdaSchedule' object"
                 )
             if isinstance(lambda_schedule, str):
-                # Strip whitespace and convert to lower case.
-                lambda_schedule = lambda_schedule.strip().lower()
-                if lambda_schedule == "standard_morph":
+                # Strip whitespace. The keyword comparison is made against a
+                # lower case copy, since the string may also be a path, which
+                # is case sensitive.
+                lambda_schedule = lambda_schedule.strip()
+                keyword = lambda_schedule.lower()
+                if keyword == "standard_morph":
                     self._lambda_schedule = _LambdaSchedule.standard_morph()
                     self._lambda_schedule_name = "standard_morph"
-                elif lambda_schedule == "charge_scaled_morph":
+                elif keyword == "charge_scaled_morph":
                     self._lambda_schedule = _LambdaSchedule.charge_scaled_morph(0.2)
                     self._lambda_schedule_name = "charge_scaled_morph"
-                elif lambda_schedule == "ring_break_morph":
+                elif keyword == "ring_break_morph":
                     from .._utils._schedules import (
                         ring_break_morph as _ring_break_morph,
                     )
 
                     self._lambda_schedule = _ring_break_morph()
                     self._lambda_schedule_name = "ring_break_morph"
-                elif lambda_schedule == "reverse_ring_break_morph":
+                elif keyword == "reverse_ring_break_morph":
                     from .._utils._schedules import (
                         reverse_ring_break_morph as _reverse_ring_break_morph,
                     )
 
                     self._lambda_schedule = _reverse_ring_break_morph()
                     self._lambda_schedule_name = "reverse_ring_break_morph"
-                elif lambda_schedule == "annihilate":
+                elif keyword == "annihilate":
                     self._lambda_schedule = None
                     self._lambda_schedule_name = "annihilate"
-                elif lambda_schedule == "decouple":
+                elif keyword == "decouple":
                     self._lambda_schedule = None
                     self._lambda_schedule_name = "decouple"
                 else:
-                    try:
-                        self._lambda_schedule = self._from_hex(lambda_schedule)
-                        self._lambda_schedule_name = None
-                    except Exception:
+                    schedule = self._from_string(
+                        lambda_schedule,
+                        "lambda_schedule",
+                        hint=", or one of the following strings: "
+                        f"{', '.join(self._choices['lambda_schedule'])}",
+                    )
+                    if not isinstance(schedule, _LambdaSchedule):
                         raise ValueError(
-                            "Unable to deserialise 'lambda_schedule'. Ensure that this is a "
-                            "hex string representation of a valid LambdaSchedule object, or "
-                            f"one of the following strings: {', '.join(self._choices['lambda_schedule'])}"
+                            f"'lambda_schedule' deserialised to a "
+                            f"'{type(schedule).__name__}', not a 'LambdaSchedule'."
                         )
+                    self._lambda_schedule = schedule
+                    self._lambda_schedule_name = None
             else:
                 self._lambda_schedule = lambda_schedule
                 self._lambda_schedule_name = None
@@ -1303,32 +1319,34 @@ def restraints(self):
 
     @restraints.setter
     def restraints(self, restraints):
-        # If not supplied as a list, convert to a list.
+        # If not supplied as a list, convert to a list. Note that a string is
+        # itself iterable, so must be wrapped explicitly.
         if restraints is not None:
-            if not isinstance(restraints, _Iterable):
+            if isinstance(restraints, str) or not isinstance(restraints, _Iterable):
                 restraints = [restraints]
 
-            # Check that all restraints are of the correct type.
-            deserialised_restraints = []
+            # Resolve each entry, keeping objects and deserialised strings in
+            # the order they were given.
+            resolved_restraints = []
             for restraint in restraints:
-                if isinstance(restraint, _sr.mm._MM.Restraints):
-                    continue
-                elif isinstance(restraint, str):
-                    try:
-                        restraint = self._from_hex(restraint)
-                    except Exception:
-                        raise ValueError(
-                            "Unable to deserialise restraint. Ensure that this "
-                            "is a hex string representation of a valid sire.mm._MM.Restraints object."
-                        )
-                    deserialised_restraints.append(restraint)
+                if isinstance(restraint, str):
+                    restraint = self._from_string(restraint.strip(), "restraints")
+
+                # A stream file may hold a list of sets of restraints, e.g. the
+                # pair used for a ring-breaking perturbation.
+                if isinstance(restraint, _Iterable):
+                    resolved_restraints.extend(restraint)
                 else:
+                    resolved_restraints.append(restraint)
+
+            # Check that all restraints are of the correct type.
+            for restraint in resolved_restraints:
+                if not isinstance(restraint, _sr.mm._MM.Restraints):
                     raise ValueError(
                         "'restraints' must be a sire.mm._MM.Restraints object, or a list of these objects."
                     )
 
-            if len(deserialised_restraints) > 0:
-                restraints = deserialised_restraints
+            restraints = resolved_restraints
 
         self._restraints = restraints
 
@@ -2622,6 +2640,61 @@ def _from_hex(hex):
 
         return obj
 
+    @classmethod
+    def _from_string(cls, string, name, hint=""):
+        """
+        Internal method to deserialise a Sire object from a string, which can
+        either be the path to a stream file, or the hex string representation
+        of the serialised object.
+
+        Parameters
+        ----------
+
+        string: str
+            The path to a stream file, or a hex string representation of the
+            Sire object.
+
+        name: str
+            The name of the option being set, used for error messages.
+
+        hint: str
+            An additional clause appended to the error message, e.g. listing
+            the keywords that the option also accepts.
+
+        Returns
+        -------
+
+        obj:
+            The deserialised Sire object.
+        """
+        from pathlib import Path as _Path
+
+        # Work out whether this is a path to an existing file. A hex string can
+        # exceed the maximum filename length, which raises rather than simply
+        # returning False on some platforms.
+        try:
+            is_file = _Path(string).is_file()
+        except Exception:
+            is_file = False
+
+        if is_file:
+            from sire.stream import load
+
+            try:
+                return load(string)
+            except Exception as e:
+                raise ValueError(
+                    f"Unable to load '{name}' from stream file '{string}': {e}"
+                )
+        else:
+            try:
+                return cls._from_hex(string)
+            except Exception:
+                raise ValueError(
+                    f"Unable to interpret '{name}'. Expected the path to a Sire "
+                    f"stream file, or a hex string of a serialised object{hint}."
+                )
+
     def __getstate__(self):
         """
         Hex-encode the same fields that to_yaml()/from_yaml() already
@@ -2658,6 +2731,7 @@ def _create_parser(cls):
 
         import argparse
         import inspect
+        import re
 
         # Inspect the signature to get the parameters.
         sig = inspect.signature(Config.__init__)
@@ -2665,7 +2739,7 @@ def _create_parser(cls):
         params = {
             key: value
             for key, value in params.items()
-            if key not in ["self", "args", "kwargs", "restraints"]
+            if key not in ["self", "args", "kwargs"]
         }
 
         # Get the docstring.
@@ -2682,7 +2756,7 @@ def _create_parser(cls):
             # Loop over all lines in the docstring until we find the parameter.
             for line in doc:
                 line = line.strip()
-                if line.startswith(param):
+                if re.match(rf"{re.escape(param)}\s*:", line):
                     found_param = True
                 elif found_param:
                     if line == "":
@@ -2724,14 +2798,26 @@ def _create_parser(cls):
 
             # This parameter has choices.
             if param in cls._choices:
-                parser.add_argument(
-                    f"--{cli_param}",
-                    type=typ,
-                    default=params[param].default,
-                    choices=cls._choices[param],
-                    help=help[param],
-                    required=False,
-                )
+                # Other forms are also accepted, so advertise the choices in the
+                # help text, but leave the validation to the setter.
+                if param in cls._open_choices:
+                    parser.add_argument(
+                        f"--{cli_param}",
+                        type=typ,
+                        default=params[param].default,
+                        metavar="{" + ",".join(cls._choices[param]) + "}",
+                        help=help[param],
+                        required=False,
+                    )
+                else:
+                    parser.add_argument(
+                        f"--{cli_param}",
+                        type=typ,
+                        default=params[param].default,
+                        choices=cls._choices[param],
+                        help=help[param],
+                        required=False,
+                    )
             # This is a standard parameter.
             else:
                 if typ == bool:
diff --git a/tests/runner/test_config.py b/tests/runner/test_config.py
index 27d8e4a..ce0168e 100644
--- a/tests/runner/test_config.py
+++ b/tests/runner/test_config.py
@@ -123,3 +123,106 @@ def test_morse_restraint_options():
 
     with pytest.raises(ValueError, match="units are invalid"):
         Config(morse_soft_force_constant="125 kcal mol-1")
+
+
+def test_lambda_schedule_input_forms():
+    """Validate that all supported lambda schedule input forms are accepted."""
+    import os
+
+    import pytest
+
+    schedule = sr.cas.LambdaSchedule.standard_morph()
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = os.path.join(tmpdir, "schedule.s3")
+        sr.stream.save(schedule, path)
+
+        # A named schedule, which is case insensitive.
+        config = Config(lambda_schedule="DECOUPLE")
+        assert config._lambda_schedule_name == "decouple"
+
+        # The path to a stream file.
+        config = Config(lambda_schedule=path)
+        assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule)
+        assert config._lambda_schedule_name is None
+
+        # A hex string of the serialised object.
+        config = Config(lambda_schedule=Config._to_hex(schedule))
+        assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule)
+
+        # The object itself.
+        config = Config(lambda_schedule=schedule)
+        assert isinstance(config.lambda_schedule, sr.cas.LambdaSchedule)
+
+        # Anything else is rejected.
+        with pytest.raises(ValueError, match="Unable to interpret"):
+            Config(lambda_schedule="not_a_schedule")
+
+        # A stream file holding the wrong type of object.
+        wrong_path = os.path.join(tmpdir, "wrong.s3")
+        sr.stream.save(sr.cas.Symbol("x"), wrong_path)
+        with pytest.raises(ValueError, match="not a 'LambdaSchedule'"):
+            Config(lambda_schedule=wrong_path)
+
+
+def test_restraints_input_forms():
+    """Validate that all supported restraint input forms are accepted."""
+    import os
+
+    import pytest
+
+    mols = sr.load_test_files("ala.top", "ala.crd")
+    restraint0 = sr.restraints.positional(mols, atoms="atomidx 0")
+    restraint1 = sr.restraints.positional(mols, atoms="atomidx 1")
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path0 = os.path.join(tmpdir, "restraint0.s3")
+        both_path = os.path.join(tmpdir, "both.s3")
+        sr.stream.save(restraint0, path0)
+        sr.stream.save([restraint0, restraint1], both_path)
+
+        # A single object, or a list of objects.
+        assert len(Config(restraints=restraint0).restraints) == 1
+        assert len(Config(restraints=[restraint0, restraint1]).restraints) == 2
+
+        # The path to a stream file, or a list of paths.
+        assert len(Config(restraints=path0).restraints) == 1
+        assert len(Config(restraints=[path0, path0]).restraints) == 2
+
+        # A stream file holding a list of sets of restraints.
+        assert len(Config(restraints=both_path).restraints) == 2
+
+        # A hex string of the serialised object.
+        assert len(Config(restraints=Config._to_hex(restraint0)).restraints) == 1
+
+        # Objects and paths can be mixed, and all are retained.
+        config = Config(restraints=[restraint0, path0])
+        assert len(config.restraints) == 2
+        assert all(
+            isinstance(restraint, sr.mm._MM.Restraints)
+            for restraint in config.restraints
+        )
+
+        # Anything else is rejected.
+        with pytest.raises(ValueError, match="Unable to interpret"):
+            Config(restraints="not_a_restraint")
+
+        # A stream file holding the wrong type of object.
+        wrong_path = os.path.join(tmpdir, "wrong.s3")
+        sr.stream.save(sr.cas.LambdaSchedule.standard_morph(), wrong_path)
+        with pytest.raises(ValueError, match="must be a sire.mm._MM.Restraints"):
+            Config(restraints=wrong_path)
+
+
+def test_help_text_scraping():
+    """Validate that help text isn't truncated by the parameter name."""
+    parser = Config._create_parser()
+
+    for action in parser._actions:
+        if action.dest == "restraints":
+            break
+
+    # The description wraps onto a line starting with the parameter name, which
+    # must not be mistaken for the start of the next parameter.
+    assert "applied to the atoms" in action.help
+    assert "a list of sets" in action.help

From 9217a47cb8baf9a3a18f75cb86b62aaebb86f44e Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 28 Aug 2026 13:27:56 +0100
Subject: [PATCH 58/67] Fix log message.

---
 src/somd2/runner/_base.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 4f2517c..2deb608 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -207,7 +207,8 @@ def __init__(self, system, config):
                     if c0 != c1:
                         msg = (
                             "End-state connectivities are different. If this is a ring-breaking "
-                            "perturbation, please set 'lambda_schedule_name' to 'ring_breaking'."
+                            "perturbation, please set 'lambda_schedule_name' to 'ring_break_morph' "
+                            "(or 'reverse_ring_break_morph' for the reverse perturbation)."
                         )
                         _logger.warning(msg)
                         break

From b4628436c67efa47c8224c496923fff6dd9e4337 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 1 Sep 2026 14:38:04 +0100
Subject: [PATCH 59/67] Fix platform section. [ci skip]

---
 README.md | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 828a9f0..0f9cf96 100644
--- a/README.md
+++ b/README.md
@@ -194,10 +194,13 @@ In order to run using GPUs you will first need to set the relevant environment
 variable. For example, to run using 4 CUDA enabled GPUs set `CUDA_VISIBLE_DEVICES=0,1,2,3`
 (for OpenCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively).
 
-By default SOMD2 will run using the CPU platform, however if the relevant
-environment variable has been set (as above) the new platform will be detected
-and set. In the case that this detection fails, or if there are multiple platforms
-available, the `--platform` option can be set (for example `--platform cuda`).
+This is required in all cases, since SOMD2 enumerates the devices to run on from
+the environment variable itself.
+
+By default SOMD2 uses `--platform auto`, which selects the first platform
+registered by OpenMM in order of preference: CUDA, OpenCL, HIP, Metal, Reference,
+then CPU. If detection fails, or if you want a specific platform, use the
+`--platform` option (for example `--platform cuda`).
 
 By default, SOMD2 will automatically manage the distribution of lambda windows
 across all listed devices. In order to restrict the number of devices used

From b46b16bb99c7a6e689af4709c2c866f22c5f84ce Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Fri, 4 Sep 2026 12:26:30 +0100
Subject: [PATCH 60/67] Account for v-site charges when checking for charge
 change.

---
 CHANGELOG.md                         |  1 +
 src/somd2/runner/_base.py            | 11 +++++++++++
 tests/runner/test_alchemical_ions.py | 22 ++++++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 112246d..648817d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,7 @@ Changelog
 * Add a `precision` option for GPU platforms, defaulting to `mixed` [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Add support for generating Morse restraints for ring-breaking perturbations.
 * Remove the unused `kappa` lever equations from the ring-breaking/making lambda schedules.
+* Account for off-site charges (virtual sites) when computing the charge difference between the end states. They are held as a molecule property rather than on the atoms, so a charge-preserving perturbation could appear to change charge and be given spurious alchemical ions.
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 2deb608..809e189 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -1344,6 +1344,17 @@ def _get_charge_difference(system):
         reference = _sr.morph.link_to_reference(system).charge().value()
         perturbed = _sr.morph.link_to_perturbed(system).charge().value()
 
+        # Off-site charges are held as a molecule property, not on the atoms,
+        # so they are invisible to charge() above.
+        try:
+            vsite_mols = system.molecules("property vs_charges0")
+        except KeyError:
+            vsite_mols = []
+
+        for mol in vsite_mols:
+            reference += sum(float(x) for x in mol.property("vs_charges0"))
+            perturbed += sum(float(x) for x in mol.property("vs_charges1"))
+
         return perturbed - reference
 
     def _save_alchemical_ion_indices(self, mol_indices):
diff --git a/tests/runner/test_alchemical_ions.py b/tests/runner/test_alchemical_ions.py
index 2c5d7fe..a255372 100644
--- a/tests/runner/test_alchemical_ions.py
+++ b/tests/runner/test_alchemical_ions.py
@@ -175,3 +175,25 @@ def test_alchemical_ion_abfe_schedule(schedule_name, ethane_methanol_ions):
         ion_schedule = schedule.get_molecule_schedule(ion_idx)
         assert ion_schedule.get_stages() == ["morph"]
         assert ion_schedule.to_string() == LambdaSchedule.standard_morph().to_string()
+
+
+def test_charge_difference_with_virtual_sites(ethane_methanol):
+    """
+    Off-site charges are held as a molecule property rather than on the atoms,
+    so they are invisible to charge() and must be added separately.
+    """
+
+    mols = ethane_methanol.clone()
+
+    # No virtual sites, so the end states have the same charge.
+    assert math.isclose(Runner._get_charge_difference(mols), 0.0, abs_tol=1e-6)
+
+    # Give the perturbable molecule off-site charges that differ by one unit
+    # between the end states.
+    mol = mols.molecules("property is_perturbable")[0]
+    cursor = mol.cursor()
+    cursor.set("vs_charges0", [0.5, 0.5])
+    cursor.set("vs_charges1", [1.0, 1.0])
+    mols.update(cursor.commit())
+
+    assert math.isclose(Runner._get_charge_difference(mols), 1.0, abs_tol=1e-6)

From 4734d57dc5f0a92aedf4a339bcc4d6c05d0cdbb5 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 8 Sep 2026 11:01:50 +0100
Subject: [PATCH 61/67] Handle num_lambda=1.

---
 src/somd2/runner/_base.py  |  2 ++
 src/somd2/runner/_repex.py | 12 +++++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 809e189..1d28609 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -466,6 +466,8 @@ def __init__(self, system, config):
         # Set the lambda values.
         if self._config.lambda_values:
             self._lambda_values = self._config.lambda_values
+        elif self._config.num_lambda == 1:
+            self._lambda_values = [0.0]
         else:
             self._lambda_values = [
                 round(i / (self._config.num_lambda - 1), 5)
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index fe6c4a6..f6aa553 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1233,6 +1233,16 @@ def __init__(self, system, config):
         # fewer slots than replicas, each slot is re-used to propagate several
         # replicas per cycle, changing its lambda value as it goes.
         self._num_replicas = len(self._lambda_values)
+
+        # There is nothing to exchange with a single replica.
+        if self._num_replicas < 2:
+            msg = (
+                "Replica exchange requires at least two "
+                f"{_lam_sym} values. Please use the Runner class instead."
+            )
+            _logger.error(msg)
+            raise ValueError(msg)
+
         self._set_num_slots()
 
         # Auto-generate a Boresch restraint for ABFE runs with no user-supplied
@@ -3068,7 +3078,7 @@ def _save_transition_matrix(self):
                         swaps[i_state, j_state] + swaps[j_state, i_state]
                     ) / denom
             else:
-                t[i_state, i_state] = 1.0
+                t_ij[i_state, i_state] = 1.0
 
         # Backup the existing transition matrix, if it exists.
         if self._repex_matrix.exists():

From 05e447bfd6f98948b4cb592b9806c20f715f8f0e Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 8 Sep 2026 11:02:09 +0100
Subject: [PATCH 62/67] Update CHANGELOG and add missing entries and links.

---
 CHANGELOG.md | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 648817d..cbb8405 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,9 +19,11 @@ Changelog
 * Skip minimisation on restart [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Add a `precision` option for GPU platforms, defaulting to `mixed` [#191](https://github.com/OpenBioSim/somd2/pull/191).
-* Add support for generating Morse restraints for ring-breaking perturbations.
-* Remove the unused `kappa` lever equations from the ring-breaking/making lambda schedules.
-* Account for off-site charges (virtual sites) when computing the charge difference between the end states. They are held as a molecule property rather than on the atoms, so a charge-preserving perturbation could appear to change charge and be given spurious alchemical ions.
+* Add support for generating Morse restraints for ring-breaking perturbations [#194](https://github.com/OpenBioSim/somd2/pull/194).
+* Remove the unused `kappa` lever equations from the ring-breaking/making lambda schedules [#195](https://github.com/OpenBioSim/somd2/pull/195).
+* Accept stream file paths for the `restraints` and `lambda_schedule` configuration options, so they can be set from the command line [#198](https://github.com/OpenBioSim/somd2/pull/198).
+* Account for off-site charges (virtual sites) when computing the charge difference between the end states. They are held as a molecule property rather than on the atoms, so a charge-preserving perturbation could appear to change charge and be given spurious alchemical ions [#200](https://github.com/OpenBioSim/somd2/pull/200).
+* Handle `num_lambda=1`, which previously raised a `ZeroDivisionError` when generating the lambda values. The `RepexRunner` now rejects a single lambda window, since there is nothing to exchange with and the regular `Runner` is faster [#203](https://github.com/OpenBioSim/somd2/pull/203).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------

From 08f7b8d9c39d0e9b414850ddfde980cac9f086d0 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 14 Sep 2026 16:09:57 +0100
Subject: [PATCH 63/67] Default to single precision until benchmarks are done.

---
 CHANGELOG.md                | 2 +-
 src/somd2/config/_config.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index cbb8405..225993e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,7 +18,7 @@ Changelog
 * Add `max_contexts` to cap the number of OpenMM contexts used for replica exchange, re-using each across lambda values so that GPU memory no longer limits the number of replicas [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Skip minimisation on restart [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Pre-equilibrate the water with GCMC moves before minimising in the regular `Runner`, making it consistent with the `RepexRunner`, which already did so to stop the geometry relaxing into a dry pocket [#191](https://github.com/OpenBioSim/somd2/pull/191).
-* Add a `precision` option for GPU platforms, defaulting to `mixed` [#191](https://github.com/OpenBioSim/somd2/pull/191).
+* Add a `precision` option for GPU platforms, defaulting to `single` [#191](https://github.com/OpenBioSim/somd2/pull/191).
 * Add support for generating Morse restraints for ring-breaking perturbations [#194](https://github.com/OpenBioSim/somd2/pull/194).
 * Remove the unused `kappa` lever equations from the ring-breaking/making lambda schedules [#195](https://github.com/OpenBioSim/somd2/pull/195).
 * Accept stream file paths for the `restraints` and `lambda_schedule` configuration options, so they can be set from the command line [#198](https://github.com/OpenBioSim/somd2/pull/198).
diff --git a/src/somd2/config/_config.py b/src/somd2/config/_config.py
index cfa0a99..b2b6598 100644
--- a/src/somd2/config/_config.py
+++ b/src/somd2/config/_config.py
@@ -139,7 +139,7 @@ def __init__(
         num_energy_neighbours=None,
         null_energy="1e6 kcal/mol",
         platform="auto",
-        precision="mixed",
+        precision="single",
         max_threads=None,
         max_gpus=None,
         max_sire_threads=None,

From 2e4d5191be5035607fc93f82ae509f66cb659ca1 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 15 Sep 2026 09:03:44 +0100
Subject: [PATCH 64/67] Use the configured platform when detecting available
 GPUs.

---
 CHANGELOG.md                |  1 +
 src/somd2/runner/_base.py   | 15 +++++++++++----
 src/somd2/runner/_repex.py  | 16 +++++-----------
 src/somd2/runner/_runner.py |  4 +---
 4 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 225993e..6df6613 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -24,6 +24,7 @@ Changelog
 * Accept stream file paths for the `restraints` and `lambda_schedule` configuration options, so they can be set from the command line [#198](https://github.com/OpenBioSim/somd2/pull/198).
 * Account for off-site charges (virtual sites) when computing the charge difference between the end states. They are held as a molecule property rather than on the atoms, so a charge-preserving perturbation could appear to change charge and be given spurious alchemical ions [#200](https://github.com/OpenBioSim/somd2/pull/200).
 * Handle `num_lambda=1`, which previously raised a `ZeroDivisionError` when generating the lambda values. The `RepexRunner` now rejects a single lambda window, since there is nothing to exchange with and the regular `Runner` is faster [#203](https://github.com/OpenBioSim/somd2/pull/203).
+* Detect the available GPUs once in the base runner and re-use the list, rather than the `RepexRunner` querying `CUDA_VISIBLE_DEVICES` regardless of the chosen platform. Replica exchange is now also permitted on the HIP platform [#206](https://github.com/OpenBioSim/somd2/pull/206).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py
index 1d28609..8e8890b 100644
--- a/src/somd2/runner/_base.py
+++ b/src/somd2/runner/_base.py
@@ -973,19 +973,26 @@ def __init__(self, system, config):
                 "schedule"
             ].reverse()
 
+        # The physical GPU devices available to this run, as listed by the
+        # visible-devices environment variable for the chosen platform.
+        self._gpu_devices = []
+
         # Limit the number of CPU threads available to Sire when running in parallel.
         if self._is_gpu:
             # First get the total number of threads that are available to Sire.
             total_threads = _sr.legacy.Base.get_max_num_threads()
 
-            # Get the number of GPU devices.
-            devices = self._get_gpu_devices(
+            # Get the available GPU devices. Subclasses re-use this list rather
+            # than querying the environment again.
+            self._gpu_devices = self._get_gpu_devices(
                 self._config.platform,
-                log=False,
+                self._config.oversubscription_factor,
             )
 
             # Work out the number of GPU workers.
-            num_gpu_workers = len(devices) * self._config.oversubscription_factor
+            num_gpu_workers = (
+                len(self._gpu_devices) * self._config.oversubscription_factor
+            )
 
             # Adjust based on the maximum number of GPUs.
             if self._config.max_gpus is not None:
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index f6aa553..7a8c7da 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1197,23 +1197,17 @@ def __init__(self, system, config):
         # Call the base class constructor.
         super().__init__(system, config)
 
-        # Make sure we're using the CUDA or OpenCL platform.
-        if self._config.platform not in ["cuda", "opencl"]:
+        # Make sure we're using a GPU platform.
+        if not self._is_gpu:
             msg = (
                 "Currently replica exchange simulations can only be "
-                "run on the CUDA and OpenCL platforms."
+                "run on the CUDA, OpenCL, and HIP platforms."
             )
             _logger.error(msg)
             raise ValueError(msg)
 
-        # Get the number of available GPUs.
-        try:
-            gpu_devices = self._get_gpu_devices(
-                "cuda", self._config.oversubscription_factor
-            )
-        except Exception as e:
-            _logger.error(f"Could not determine available GPU devices: {e}")
-            raise e
+        # The available devices were detected by the base class constructor.
+        gpu_devices = self._gpu_devices
 
         # We can only use replica exchange if we have a GPU.
         if len(gpu_devices) == 0:
diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py
index 0017b44..5b6283e 100644
--- a/src/somd2/runner/_runner.py
+++ b/src/somd2/runner/_runner.py
@@ -112,9 +112,7 @@ def _create_shared_resources(self):
         Also intialises the list with all available GPUs.
         """
         if self._is_gpu:
-            devices = self._get_gpu_devices(
-                self._config.platform, self._config.oversubscription_factor
-            )
+            devices = self._gpu_devices
             if self._config.max_gpus is not None:
                 if self._config.max_gpus > len(devices):
                     _logger.warning(

From e25aa48b08e27f7c949ef90414c271b0f584885e Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 15 Sep 2026 09:04:03 +0100
Subject: [PATCH 65/67] Fix the OpenCL constant used to query AMD GPU free
 memory.

---
 CHANGELOG.md               |  1 +
 src/somd2/runner/_repex.py | 20 ++++++++++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6df6613..a3a3fce 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ Changelog
 * Account for off-site charges (virtual sites) when computing the charge difference between the end states. They are held as a molecule property rather than on the atoms, so a charge-preserving perturbation could appear to change charge and be given spurious alchemical ions [#200](https://github.com/OpenBioSim/somd2/pull/200).
 * Handle `num_lambda=1`, which previously raised a `ZeroDivisionError` when generating the lambda values. The `RepexRunner` now rejects a single lambda window, since there is nothing to exchange with and the regular `Runner` is faster [#203](https://github.com/OpenBioSim/somd2/pull/203).
 * Detect the available GPUs once in the base runner and re-use the list, rather than the `RepexRunner` querying `CUDA_VISIBLE_DEVICES` regardless of the chosen platform. Replica exchange is now also permitted on the HIP platform [#206](https://github.com/OpenBioSim/somd2/pull/206).
+* Query the free memory of AMD GPUs with `CL_DEVICE_GLOBAL_FREE_MEMORY_AMD` rather than `CL_DEVICE_BOARD_NAME_AMD`, which returns the device name [#206](https://github.com/OpenBioSim/somd2/pull/206).
 
 [2026.1.0](https://github.com/openbiosim/somd2/compare/2025.1.0...2026.1.0) - Jun 2026
 --------------------------------------------------------------------------------------
diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py
index 7a8c7da..5dc4d66 100644
--- a/src/somd2/runner/_repex.py
+++ b/src/somd2/runner/_repex.py
@@ -1144,11 +1144,18 @@ def _check_device_memory(device=0):
         # AMD: Use OpenCL extension.
         elif "AMD" in vendor or "Advanced Micro Devices" in vendor:
             try:
+                import pyopencl as cl
+
                 total = ocl_device.global_mem_size
-                free_memory_info = ocl_device.get_info(0x4038)
+
+                # cl_amd_device_attribute_query reports the free memory in
+                # KBytes, as a list with the whole heap first.
+                free_memory_info = ocl_device.get_info(
+                    cl.device_info.GLOBAL_FREE_MEMORY_AMD
+                )
                 free_kb = (
                     free_memory_info[0]
-                    if isinstance(free_memory_info, list)
+                    if isinstance(free_memory_info, (list, tuple))
                     else free_memory_info
                 )
                 free = free_kb * 1024
@@ -1159,6 +1166,15 @@ def _check_device_memory(device=0):
                 _logger.error(msg)
                 raise RuntimeError(msg) from e
 
+        # Any other vendor: we have no way of querying the free memory.
+        else:
+            msg = (
+                f"Unable to query the memory of device {device}: "
+                f"unsupported GPU vendor '{vendor}'."
+            )
+            _logger.error(msg)
+            raise RuntimeError(msg)
+
 
 class RepexRunner(_RunnerBase):
     """

From ca848fc7e3744b11d2d4eaa7de42b752b4158bbe Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 15 Sep 2026 09:53:40 +0100
Subject: [PATCH 66/67] Update BioSimSpace pin.

---
 pixi.toml                 | 4 ++--
 recipes/somd2/recipe.yaml | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pixi.toml b/pixi.toml
index 8a43389..612b388 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -6,9 +6,9 @@ platforms = ["linux-64", "osx-arm64"]
 [dependencies]
 python = ">=3.10"
 # main
-#biosimspace = ">=2026.1.0,<2026.2.0"
+biosimspace = ">=2026.2.0,<2026.3.0"
 # devel
-biosimspace = "==2026.2.0.dev"
+#biosimspace = "==2026.3.0.dev"
 filelock = "*"
 ghostly = "*"
 loch = "*"
diff --git a/recipes/somd2/recipe.yaml b/recipes/somd2/recipe.yaml
index ba83354..74073ae 100644
--- a/recipes/somd2/recipe.yaml
+++ b/recipes/somd2/recipe.yaml
@@ -20,9 +20,9 @@ requirements:
     - versioningit
   run:
     # main
-    #- biosimspace >=2026.1.0,<2026.2.0
+    - biosimspace >=2026.2.0,<2026.3.0
     # devel
-    - biosimspace ==2026.2.0.dev
+    #- biosimspace ==2026.3.0.dev
     - filelock
     - ghostly
     - loch

From b60f2ab8a77e125726a07c2635e584fda59d3a38 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 15 Sep 2026 09:53:55 +0100
Subject: [PATCH 67/67] Update CHANGELOG.

---
 CHANGELOG.md | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a3a3fce..dab67f8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,10 +1,9 @@
 Changelog
 =========
 
-[2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - ********
+[2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026
 --------------------------------------------------------------------------------------
 
-* Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
 * Add support for generating Boresch restraints for absolute binding free energy calculations [#166](https://github.com/OpenBioSim/somd2/pull/166).
 * Give alchemical ions their own plain morph lambda schedule so they interpolate correctly under non-standard lambda schedules [#169](https://github.com/OpenBioSim/somd2/pull/169).
 * Persist alchemical ion identity across restarts so the same molecule is reused regardless of GCMC state [#172](https://github.com/OpenBioSim/somd2/pull/172).