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

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

diff --git a/pixi.toml b/pixi.toml
index 16c6cff..56780fd 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -6,9 +6,9 @@ platforms = ["linux-64", "osx-arm64", "win-64"]
 [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"
 loguru = "*"
 pyopencl = "*"
 
diff --git a/recipes/loch/recipe.yaml b/recipes/loch/recipe.yaml
index 1e854de..1fb1f79 100644
--- a/recipes/loch/recipe.yaml
+++ b/recipes/loch/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
     - loguru
     - pyopencl
     - python

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

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 50de0be..ca9dbb1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,11 @@
 Changelog
 =========
 
+[2026.2.0](https://github.com/openbiosim/loch/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/loch/compare/2025.2.0...2026.1.0) - Jun 2026
 -------------------------------------------------------------------------------------
 

From a56b935655e31efe90f216e69acc8ebbe6a6b480 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 22 Jul 2026 11:53:41 +0100
Subject: [PATCH 03/16] Restrict PME calculation to specific force groups.

---
 CHANGELOG.md         |  1 +
 src/loch/_sampler.py | 29 +++++++++++++++++++++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ca9dbb1..87f18d6 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.
+* Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29).
 
 [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026
 -------------------------------------------------------------------------------------
diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index 0c17461..c1ffc54 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -731,6 +731,8 @@ def __init__(
         # Null the nonbonded forces.
         self._nonbonded_force = None
         self._custom_nonbonded_force = None
+        self._integration_groups = None
+        self._pme_groups = None
 
         # Flag for whether the last move was a bulk sampling move.
         self._is_bulk = False
@@ -1353,6 +1355,8 @@ def reset(self) -> None:
         # Clear the forces.
         self._nonbonded_force = None
         self._custom_nonbonded_force = None
+        self._integration_groups = None
+        self._pme_groups = None
 
         # Clear the OpenMM context.
         self._openmm_context = None
@@ -1482,7 +1486,11 @@ def move(self, context: _openmm.Context) -> list[int]:
                         self._init_gcmc_lrc(context)
 
                     # Get the OpenMM state.
-                    state = context.getState(getPositions=True, getEnergy=self._is_pme)
+                    state = context.getState(
+                        getPositions=True,
+                        getEnergy=self._is_pme,
+                        groups=self._pme_groups,
+                    )
 
                     # Get the current positions in OpenMM format and in Angstrom.
                     positions_openmm = state.getPositions(asNumpy=True)
@@ -1768,7 +1776,7 @@ def move(self, context: _openmm.Context) -> list[int]:
 
                         # Get the new energy.
                         final_energy = context.getState(
-                            getEnergy=True
+                            getEnergy=True, groups=self._pme_groups
                         ).getPotentialEnergy()
 
                         # Add the analytic LRC delta so the PME correction sees only
@@ -1870,7 +1878,7 @@ def move(self, context: _openmm.Context) -> list[int]:
 
                         # Get the new energy.
                         final_energy = context.getState(
-                            getEnergy=True
+                            getEnergy=True, groups=self._pme_groups
                         ).getPotentialEnergy()
 
                         # Add the analytic LRC delta.
@@ -3007,12 +3015,19 @@ def _set_nonbonded_forces(self, context):
         context: openmm.Context
             The OpenMM context to use.
         """
+        if self._integration_groups is None:
+            self._integration_groups = (
+                context.getIntegrator().getIntegrationForceGroups()
+            )
+
         if self._nonbonded_force is None or (
             self._is_fep and self._custom_nonbonded_force is None
         ):
             for force in context.getSystem().getForces():
                 if isinstance(force, _openmm.NonbondedForce):
-                    self._nonbonded_force = force
+                    # Only accept a force actually used for integration.
+                    if self._integration_groups & (1 << force.getForceGroup()):
+                        self._nonbonded_force = force
                 elif self._is_fep and force.getName() == "GhostNonGhostNonbondedForce":
                     self._custom_nonbonded_force = force
                 elif self._pressure is None and "Barostat" in force.getName():
@@ -3034,6 +3049,12 @@ def _set_nonbonded_forces(self, context):
             _logger.error(msg)
             raise ValueError(msg)
 
+        if self._pme_groups is None:
+            groups = 1 << self._nonbonded_force.getForceGroup()
+            if self._is_fep:
+                groups |= 1 << self._custom_nonbonded_force.getForceGroup()
+            self._pme_groups = groups
+
     def _get_target_position(self, positions):
         """
         Get the current centre of the GCMC sphere.

From 40a660a69801008571a0dcefe3ddcd93582003c1 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 27 Jul 2026 12:58:27 +0100
Subject: [PATCH 04/16] Allow a sampler to be re-used across lambda values.

---
 CHANGELOG.md         |   1 +
 src/loch/_sampler.py | 335 ++++++++++++++++++++++++++++++++++---------
 2 files changed, 269 insertions(+), 67 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 87f18d6..acb79af 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.
 * Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29).
+* Add `set_lambda` and `precompute_lambda` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context.
 
 [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026
 -------------------------------------------------------------------------------------
diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index c1ffc54..7fe20e3 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -89,6 +89,8 @@ def __init__(
         lambda_schedule: _Optional[_Any] = None,
         lambda_value: float = 0.0,
         rest2_scale: float = 1.0,
+        lambda_values: _Optional[list[float]] = None,
+        rest2_scales: _Optional[list[float]] = None,
         rest2_selection: _Optional[str] = None,
         shift_coulomb: str = "1 A",
         shift_delta: str = "1.5 A",
@@ -205,6 +207,16 @@ def __init__(
             (REST2) for alchemical systems. This should specify the temperature
             of the REST2 system relative to the rest of the system.
 
+        lambda_values: [float]
+            The lambda values that the sampler will be switched between via
+            set_lambda(). The non-bonded parameters for each are cached at
+            construction, so that switching later doesn't need to build an
+            OpenMM context. If None, only 'lambda_value' is cached.
+
+        rest2_scales: [float]
+            The REST2 scaling factor for each entry of 'lambda_values'. If
+            None, 'rest2_scale' is used for all of them.
+
         rest2_selection: str
             A selection string for atoms to include in the REST2 region in
             addition to any perturbable molecules. For example, "molidx 0 and
@@ -506,6 +518,37 @@ def __init__(
             raise ValueError("'rest2_scale' must be greater than or equal to 1.0")
         self._rest2_scale = rest2_scale
 
+        # Cache of the lambda dependent non-bonded parameters, keyed by
+        # (lambda_value, rest2_scale). Populated below and consumed by
+        # set_lambda().
+        self._lambda_params = {}
+
+        if lambda_values is None:
+            self._lambda_values = None
+        else:
+            try:
+                self._lambda_values = [float(x) for x in lambda_values]
+            except:
+                raise ValueError("'lambda_values' must be a list of 'float'")
+            if not all(0.0 <= x <= 1.0 for x in self._lambda_values):
+                raise ValueError("'lambda_values' must be between 0 and 1")
+
+            if rest2_scales is None:
+                self._rest2_scales = [rest2_scale] * len(self._lambda_values)
+            else:
+                try:
+                    self._rest2_scales = [float(x) for x in rest2_scales]
+                except:
+                    raise ValueError("'rest2_scales' must be a list of 'float'")
+                if len(self._rest2_scales) != len(self._lambda_values):
+                    raise ValueError(
+                        "'rest2_scales' must be the same length as 'lambda_values'"
+                    )
+                if any(x < 1.0 for x in self._rest2_scales):
+                    raise ValueError(
+                        "'rest2_scales' must be greater than or equal to 1.0"
+                    )
+
         if rest2_selection is not None:
             if not isinstance(rest2_selection, str):
                 raise ValueError("'rest2_selection' must be of type 'str'")
@@ -694,6 +737,17 @@ def __init__(
         self._water_blocks = self._num_waters // self._num_threads + 1
 
         # Initialise the GPU memory.
+        # Cache the non-bonded parameters for every lambda value that the
+        # sampler will be switched to, so that set_lambda() never has to build
+        # an OpenMM context mid-simulation. This is done before the GPU memory
+        # is initialised, so that the current lambda value is part of the same
+        # scan rather than needing a context of its own.
+        if self._lambda_values is not None:
+            self._precompute_lambdas(
+                [self._lambda_value] + self._lambda_values,
+                [self._rest2_scale] + self._rest2_scales,
+            )
+
         self._initialise_gpu_memory()
 
         # Set the box information.
@@ -2296,6 +2350,210 @@ def _prepare_system(system, water_template, rng, num_ghost_waters):
             _np.array(water_residues),
         )
 
+    def _precompute_lambdas(
+        self, lambda_values: list[float], rest2_scales: list[float]
+    ) -> None:
+        """
+        Cache the lambda dependent non-bonded parameters for a set of lambda
+        values, so that set_lambda() can switch between them without touching
+        OpenMM.
+
+        The system only holds the end-state properties, so the parameters are
+        obtained from an OpenMM context built using the specified lambda
+        schedule. A single context is created and scanned over the requested
+        lambda values, re-reading the GhostNonGhostNonbondedForce at each one,
+        since building the context is by far the expensive part.
+
+        This is a no-op for non-alchemical systems, and for lambda values that
+        have already been cached.
+
+        Parameters
+        ----------
+
+        lambda_values: [float]
+            The lambda values to cache.
+
+        rest2_scales: [float]
+            The REST2 scaling factor for each lambda value.
+        """
+
+        if not self._is_fep:
+            return
+
+        # Only build a context if something is actually missing.
+        wanted = []
+        for lambda_value, rest2_scale in zip(lambda_values, rest2_scales):
+            key = (float(lambda_value), float(rest2_scale))
+            if key not in self._lambda_params and key not in wanted:
+                wanted.append(key)
+
+        if not wanted:
+            return
+
+        # Link to the reference state.
+        mols = _sr.morph.link_to_reference(self._system)
+
+        # Build map of extra options for the dynamics object.
+        _map = {}
+        if self._softcore_form == _SoftcoreForm.TAYLOR:
+            _map["use_taylor_softening"] = True
+            _map["taylor_power"] = self._taylor_power
+        elif self._softcore_form == _SoftcoreForm.BEUTLER:
+            _map["use_beutler_softening"] = True
+            _map["beutler_alpha"] = self._beutler_alpha
+
+        # Create a single dynamics object, which is then scanned over the
+        # requested lambda values.
+        d = mols.dynamics(
+            cutoff_type=self._cutoff,
+            cutoff=self._cutoff,
+            lambda_value=wanted[0][0],
+            schedule=self._lambda_schedule,
+            pressure=None,
+            timestep="2fs",
+            constraint="h_bonds",
+            perturbable_constraint="h_bonds_not_heavy_perturbed",
+            rest2_scale=wanted[0][1],
+            rest2_selection=self._rest2_selection,
+            swap_end_states=self._swap_end_states,
+            platform="cpu",
+            map=_map,
+        )
+
+        # Find the required force. set_lambda() updates it in place, so it
+        # only needs to be located once.
+        gng_force = None
+        for force in d.context().getSystem().getForces():
+            if force.getName() == "GhostNonGhostNonbondedForce":
+                gng_force = force
+                break
+
+        if gng_force is None:
+            raise ValueError(
+                "Could not find the GhostNonGhostNonbondedForce in the system"
+            )
+
+        num_particles = gng_force.getNumParticles()
+
+        for lambda_value, rest2_scale in wanted:
+            d.set_lambda(lambda_value, rest2_scale=rest2_scale)
+
+            # Get the parameters for the GhostNonGhostNonbondedForce.
+            charges = _np.zeros(self._num_atoms, dtype=_np.float32)
+            sigmas = _np.zeros(self._num_atoms, dtype=_np.float32)
+            epsilons = _np.zeros(self._num_atoms, dtype=_np.float32)
+            alphas = _np.zeros(self._num_atoms, dtype=_np.float32)
+            for i in range(num_particles):
+                # Custom force parameters are returned as floats.
+                q, half_sigma, two_sqrt_epsilon, alpha, _ = (
+                    gng_force.getParticleParameters(i)
+                )
+                # Charge in |e|, sigma in nm, epsilon in kJ/mol.
+                charges[i] = q
+                # Rescale and convert units.
+                sigmas[i] = _sr.u(f"{2.0 * half_sigma} nm").to("angstrom")
+                epsilons[i] = _sr.u(f"{(0.5 * two_sqrt_epsilon) ** 2} kJ/mol").to(
+                    "kcal/mol"
+                )
+                # Store the softening parameter.
+                alphas[i] = alpha
+
+            self._lambda_params[(lambda_value, rest2_scale)] = (
+                charges,
+                sigmas,
+                epsilons,
+                alphas,
+            )
+
+    def set_lambda(
+        self, lambda_value: float, rest2_scale: _Optional[float] = None
+    ) -> None:
+        """
+        Set the lambda value for the sampler, updating the non-bonded
+        parameters used to evaluate insertion and deletion energies.
+
+        Parameters for lambda values passed to the constructor are already
+        cached. Any other value is computed here, which requires building an
+        OpenMM context and is slow, though the result is then cached too.
+
+        This must be kept consistent with the lambda value of the OpenMM
+        context that the sampler is used with.
+
+        Parameters
+        ----------
+
+        lambda_value: float
+            The lambda value.
+
+        rest2_scale: float
+            The REST2 scaling factor. If None, the current value is retained.
+        """
+
+        try:
+            lambda_value = float(lambda_value)
+        except:
+            raise ValueError("'lambda_value' must be of type 'float'")
+        if not 0.0 <= lambda_value <= 1.0:
+            raise ValueError("'lambda_value' must be between 0 and 1")
+
+        if rest2_scale is None:
+            rest2_scale = self._rest2_scale
+        else:
+            try:
+                rest2_scale = float(rest2_scale)
+            except:
+                raise ValueError("'rest2_scale' must be of type 'float'")
+            if rest2_scale < 1.0:
+                raise ValueError("'rest2_scale' must be greater than or equal to 1.0")
+
+        # Nothing to do.
+        if lambda_value == self._lambda_value and rest2_scale == self._rest2_scale:
+            return
+
+        # There are no lambda dependent parameters for a non-alchemical system.
+        if not self._is_fep:
+            self._lambda_value = lambda_value
+            self._rest2_scale = rest2_scale
+            return
+
+        # Make sure the parameters are cached. This builds an OpenMM context,
+        # so is slow; pass 'lambda_values' to the constructor to avoid it.
+        self._precompute_lambdas([lambda_value], [rest2_scale])
+
+        charges, sigmas, epsilons, alphas = self._lambda_params[
+            (lambda_value, rest2_scale)
+        ]
+
+        # Upload the new parameters to the GPU.
+        self._gpu_charge = self._backend.to_gpu(charges)
+        self._gpu_sigma = self._backend.to_gpu(sigmas)
+        self._gpu_epsilon = self._backend.to_gpu(epsilons)
+        self._gpu_alpha = self._backend.to_gpu(alphas)
+
+        self._lambda_value = lambda_value
+        self._rest2_scale = rest2_scale
+
+    def set_ghost_file(self, ghost_file: _Optional[str]) -> None:
+        """
+        Set the file that write_ghost_residues() appends to.
+
+        Unlike the constructor, this does not create or truncate the file,
+        since it is intended for switching between files that are already
+        being written to.
+
+        Parameters
+        ----------
+
+        ghost_file: str
+            The path to the ghost residue file. If None, ghost residues
+            cannot be written.
+        """
+
+        if ghost_file is not None and not isinstance(ghost_file, str):
+            raise TypeError("'ghost_file' must be of type 'str'")
+
+        self._ghost_file = ghost_file
+
     def _initialise_gpu_memory(self):
         """
         Initialise the GPU memory.
@@ -2360,76 +2618,19 @@ def _initialise_gpu_memory(self):
         # schedule and value, then extract the required properties from the forces
         # within the context. (The system just contains the end-state properties.)
         else:
-            # Link to the reference state.
-            mols = _sr.morph.link_to_reference(self._system)
-
-            # Build map of extra options for the dynamics object.
-            _map = {}
-            if self._softcore_form == _SoftcoreForm.TAYLOR:
-                _map["use_taylor_softening"] = True
-                _map["taylor_power"] = self._taylor_power
-            elif self._softcore_form == _SoftcoreForm.BEUTLER:
-                _map["use_beutler_softening"] = True
-                _map["beutler_alpha"] = self._beutler_alpha
-
-            # Create a dynamics object.
-            d = mols.dynamics(
-                cutoff_type=self._cutoff,
-                cutoff=self._cutoff,
-                lambda_value=self._lambda_value,
-                schedule=self._lambda_schedule,
-                pressure=None,
-                timestep="2fs",
-                constraint="h_bonds",
-                perturbable_constraint="h_bonds_not_heavy_perturbed",
-                rest2_scale=self._rest2_scale,
-                rest2_selection=self._rest2_selection,
-                swap_end_states=self._swap_end_states,
-                platform="cpu",
-                map=_map,
-            )
-
-            # Flag for the required force.
-            has_gng = False
-
-            # Find the required forces.
-            for force in d.context().getSystem().getForces():
-                if force.getName() == "GhostNonGhostNonbondedForce":
-                    gng_force = force
-                    has_gng = True
-                    break
+            # Cache the host arrays so that set_lambda() can switch back to
+            # this lambda value without rebuilding an OpenMM context.
+            self._precompute_lambdas([self._lambda_value], [self._rest2_scale])
 
-            # Make sure the force was found.
-            if not has_gng:
-                raise ValueError(
-                    "Could not find the GhostNonGhostNonbondedForce in the system"
-                )
-
-            # Get the parameters for the GhostNonGhostNonbondedForce.
-            charges = _np.zeros(self._num_atoms, dtype=_np.float32)
-            sigmas = _np.zeros(self._num_atoms, dtype=_np.float32)
-            epsilons = _np.zeros(self._num_atoms, dtype=_np.float32)
-            alphas = _np.zeros(self._num_atoms, dtype=_np.float32)
-            for i in range(gng_force.getNumParticles()):
-                # Custom force parameters are returned as floats.
-                q, half_sigma, two_sqrt_epsilon, alpha, _ = (
-                    gng_force.getParticleParameters(i)
-                )
-                # Charge in |e|, sigma in nm, epsilon in kJ/mol.
-                charges[i] = q
-                # Rescale and convert units.
-                sigmas[i] = _sr.u(f"{2.0 * half_sigma} nm").to("angstrom")
-                epsilons[i] = _sr.u(f"{(0.5 * two_sqrt_epsilon) ** 2} kJ/mol").to(
-                    "kcal/mol"
-                )
-                # Store the softening parameter.
-                alphas[i] = alpha
+            charges, sigmas, epsilons, alphas = self._lambda_params[
+                (self._lambda_value, self._rest2_scale)
+            ]
 
             # Convert to GPU arrays.
-            charges = self._backend.to_gpu(charges.astype(_np.float32))
-            sigmas = self._backend.to_gpu(sigmas.astype(_np.float32))
-            epsilons = self._backend.to_gpu(epsilons.astype(_np.float32))
-            alphas = self._backend.to_gpu(alphas.astype(_np.float32))
+            charges = self._backend.to_gpu(charges)
+            sigmas = self._backend.to_gpu(sigmas)
+            epsilons = self._backend.to_gpu(epsilons)
+            alphas = self._backend.to_gpu(alphas)
 
             # Create the ghost atom array.
             is_ghost_fep = _np.zeros(self._num_atoms, dtype=_np.int32)

From 671331ac650ebf49ec3b72882ce170a5cabeb1d0 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Mon, 27 Jul 2026 16:07:27 +0100
Subject: [PATCH 05/16] Mark method as private. [ci skip]

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index acb79af..4298cd3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,7 +6,7 @@ Changelog
 
 * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
 * Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29).
-* Add `set_lambda` and `precompute_lambda` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context.
+* Add `set_lambda` and `_precompute_lambdas` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context.
 
 [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026
 -------------------------------------------------------------------------------------

From 473c89e9c01fa528f7db2723184d1b8a12a46614 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 28 Jul 2026 09:59:08 +0100
Subject: [PATCH 06/16] Handle per-lambda sampler statistics.

---
 src/loch/_sampler.py | 150 +++++++++++++++++++++++++++++++++++--------
 1 file changed, 124 insertions(+), 26 deletions(-)

diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index 7fe20e3..6b1d24d 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -775,13 +775,18 @@ def __init__(
         # Zero the number of waters in the sampling volume.
         self._N = 0
 
-        # Zero the statistics.
+        # Zero the statistics for the current lambda value.
         self._num_moves = 0
         self._num_accepted = 0
         self._num_accepted_attempts = 0
         self._num_insertions = 0
         self._num_deletions = 0
 
+        # Statistics for lambda values other than the current one, keyed by
+        # formatted lambda value. The current one is held in the counters above
+        # and archived here by set_lambda().
+        self._stats = {}
+
         # Null the nonbonded forces.
         self._nonbonded_force = None
         self._custom_nonbonded_force = None
@@ -1397,14 +1402,9 @@ def reset(self) -> None:
         """
         Reset the sampler.
         """
-        # Zero the number of accepted moves.
-        self._num_accepted = 0
-        self._num_insertions = 0
-        self._num_deletions = 0
-        self._num_moves = 0
-        self._num_accepted_attempts = 0
-        self._num_accepted_insertions = 0
-        self._num_accepted_deletions = 0
+        # Zero the number of accepted moves, for every lambda value.
+        self._zero_stats()
+        self._stats = {}
 
         # Clear the forces.
         self._nonbonded_force = None
@@ -1415,43 +1415,137 @@ def reset(self) -> None:
         # Clear the OpenMM context.
         self._openmm_context = None
 
-    def restore_stats(self, stats: dict) -> None:
+    @staticmethod
+    def stats_key(lambda_value: float) -> str:
         """
-        Restore sampler statistics from a dictionary.
+        Return the key used to store statistics for a lambda value.
 
         Parameters
         ----------
 
-        stats : dict
-            Dictionary of sampler statistics as returned by ``get_stats()``.
+        lambda_value: float
+            The lambda value.
+
+        Returns
+        -------
+
+        str
+            The key.
+        """
+        return f"{float(lambda_value):.5f}"
+
+    def _zero_stats(self) -> None:
+        """
+        Zero the statistics for the current lambda value.
+        """
+        self._num_moves = 0
+        self._num_accepted = 0
+        self._num_insertions = 0
+        self._num_deletions = 0
+        self._num_accepted_attempts = 0
+
+    def _stats_keys(self) -> set:
+        """
+        Return the keys of the lambda values that this sampler visits.
+        """
+        keys = {self.stats_key(self._lambda_value)}
+
+        if self._lambda_values is not None:
+            keys.update(self.stats_key(x) for x in self._lambda_values)
+
+        return keys
+
+    def _switch_stats(self, lambda_value: float) -> None:
+        """
+        Archive the statistics for the current lambda value and load those for
+        a new one, zeroing them if it hasn't been visited before.
+
+        Parameters
+        ----------
+
+        lambda_value: float
+            The lambda value being switched to.
+        """
+        old_key = self.stats_key(self._lambda_value)
+        new_key = self.stats_key(lambda_value)
+
+        if old_key == new_key:
+            return
+
+        self._stats[old_key] = self._get_current_stats()
+
+        if new_key in self._stats:
+            self._set_current_stats(self._stats.pop(new_key))
+        else:
+            self._zero_stats()
+
+    def _get_current_stats(self) -> dict:
+        """
+        Return the statistics for the current lambda value.
+        """
+        return {
+            "num_moves": self._num_moves,
+            "num_accepted": self._num_accepted,
+            "num_insertions": self._num_insertions,
+            "num_deletions": self._num_deletions,
+            "num_accepted_attempts": self._num_accepted_attempts,
+        }
+
+    def _set_current_stats(self, stats: dict) -> None:
+        """
+        Set the statistics for the current lambda value. Unrecognised entries
+        are ignored, so that statistics written by an older version can still
+        be restored.
         """
         self._num_moves = stats["num_moves"]
         self._num_accepted = stats["num_accepted"]
         self._num_insertions = stats["num_insertions"]
         self._num_deletions = stats["num_deletions"]
         self._num_accepted_attempts = stats["num_accepted_attempts"]
-        self._num_accepted_insertions = stats["num_accepted_insertions"]
-        self._num_accepted_deletions = stats["num_accepted_deletions"]
+
+    def restore_stats(self, stats: dict) -> None:
+        """
+        Restore sampler statistics.
+
+        Parameters
+        ----------
+
+        stats : dict
+            Statistics as returned by ``get_stats()``, i.e. keyed by lambda
+            value. Entries for lambda values that this sampler doesn't visit
+            are ignored, so it is safe to pass the statistics for a whole
+            simulation to each of several samplers. Were they retained, a
+            sampler would report stale statistics for another's lambda values,
+            which could then overwrite the live ones when merged.
+        """
+        keys = self._stats_keys()
+        self._stats = {key: dict(value) for key, value in stats.items() if key in keys}
+
+        # Load the statistics for the current lambda value, if present.
+        key = self.stats_key(self._lambda_value)
+        if key in self._stats:
+            self._set_current_stats(self._stats.pop(key))
+        else:
+            self._zero_stats()
 
     def get_stats(self) -> dict:
         """
-        Return the current sampler statistics as a dictionary.
+        Return the sampler statistics, keyed by lambda value.
+
+        A sampler may be switched between lambda values via ``set_lambda()``,
+        and accumulates statistics for each of them separately. Non-alchemical
+        systems have a single key, for the lambda value the sampler was created
+        with.
 
         Returns
         -------
 
         dict
-            Dictionary of sampler statistics.
+            Dictionary of sampler statistics for each lambda value.
         """
-        return {
-            "num_moves": self._num_moves,
-            "num_accepted": self._num_accepted,
-            "num_insertions": self._num_insertions,
-            "num_deletions": self._num_deletions,
-            "num_accepted_attempts": self._num_accepted_attempts,
-            "num_accepted_insertions": self._num_accepted_insertions,
-            "num_accepted_deletions": self._num_accepted_deletions,
-        }
+        stats = {key: dict(value) for key, value in self._stats.items()}
+        stats[self.stats_key(self._lambda_value)] = self._get_current_stats()
+        return stats
 
     def ghost_residues(self) -> _np.ndarray:
         """
@@ -2510,6 +2604,10 @@ def set_lambda(
         if lambda_value == self._lambda_value and rest2_scale == self._rest2_scale:
             return
 
+        # Statistics are accumulated per lambda value, so archive those for the
+        # current one and load those for the new one.
+        self._switch_stats(lambda_value)
+
         # There are no lambda dependent parameters for a non-alchemical system.
         if not self._is_fep:
             self._lambda_value = lambda_value

From 1309869bfbec6af12a5cfcb24152e1b6ea115f30 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 28 Jul 2026 12:57:02 +0100
Subject: [PATCH 07/16] Avoid per-atom unit conversion when extracting lambda
 parameters.

---
 CHANGELOG.md         |  3 ++-
 src/loch/_sampler.py | 11 +++++++----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4298cd3..faf3291 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,7 +6,8 @@ Changelog
 
 * Please add an item to this CHANGELOG for any new features or bug fixes when creating a PR.
 * Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29).
-* Add `set_lambda` and `_precompute_lambdas` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context.
+* Add `set_lambda` and `_precompute_lambdas` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context [#30](https://github.com/OpenBioSim/loch/pull/30).
+* Resolve unit conversions once when extracting the lambda dependent non-bonded parameters, rather than per atom, which dominated sampler setup [#33](https://github.com/OpenBioSim/loch/pull/33).
 
 [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026
 -------------------------------------------------------------------------------------
diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index 6b1d24d..e60881a 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -2529,6 +2529,11 @@ def _precompute_lambdas(
 
         num_particles = gng_force.getNumParticles()
 
+        # Resolve the unit conversions once. Doing this per atom, by formatting
+        # and re-parsing a string, dominates the setup time for large systems.
+        nm_to_angstrom = _sr.u("1 nm").to("angstrom")
+        kj_per_mol_to_kcal_per_mol = _sr.u("1 kJ/mol").to("kcal/mol")
+
         for lambda_value, rest2_scale in wanted:
             d.set_lambda(lambda_value, rest2_scale=rest2_scale)
 
@@ -2545,10 +2550,8 @@ def _precompute_lambdas(
                 # Charge in |e|, sigma in nm, epsilon in kJ/mol.
                 charges[i] = q
                 # Rescale and convert units.
-                sigmas[i] = _sr.u(f"{2.0 * half_sigma} nm").to("angstrom")
-                epsilons[i] = _sr.u(f"{(0.5 * two_sqrt_epsilon) ** 2} kJ/mol").to(
-                    "kcal/mol"
-                )
+                sigmas[i] = 2.0 * half_sigma * nm_to_angstrom
+                epsilons[i] = (0.5 * two_sqrt_epsilon) ** 2 * kj_per_mol_to_kcal_per_mol
                 # Store the softening parameter.
                 alphas[i] = alpha
 

From b90d41a18c219f1b46a83204b87ad04f21fcbb89 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 15:08:52 +0100
Subject: [PATCH 08/16] Add missing test file.

---
 tests/test_stats.py | 224 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 224 insertions(+)
 create mode 100644 tests/test_stats.py

diff --git a/tests/test_stats.py b/tests/test_stats.py
new file mode 100644
index 0000000..41baeb6
--- /dev/null
+++ b/tests/test_stats.py
@@ -0,0 +1,224 @@
+import pytest
+
+from loch import GCMCSampler
+
+
+def make_sampler(lambda_value=0.0, lambda_values=None, is_fep=True):
+    """
+    Create a sampler with only the attributes the statistics use, so that the
+    bookkeeping can be tested without a system or a GPU.
+    """
+    sampler = object.__new__(GCMCSampler)
+    sampler._lambda_value = lambda_value
+    sampler._rest2_scale = 1.0
+    sampler._is_fep = is_fep
+    sampler._lambda_values = lambda_values
+    sampler._stats = {}
+    sampler._zero_stats()
+    return sampler
+
+
+def do_moves(sampler, num_moves):
+    """Pretend that a number of moves were performed and all were accepted."""
+    sampler._num_moves += num_moves
+    sampler._num_accepted += num_moves
+
+
+def switch(sampler, lambda_value):
+    """Switch lambda, as set_lambda() does."""
+    sampler._switch_stats(lambda_value)
+    sampler._lambda_value = lambda_value
+
+
+class TestStatsKey:
+    """Tests for the key used to store statistics."""
+
+    @pytest.mark.parametrize(
+        "lambda_value, expected",
+        [
+            (0.0, "0.00000"),
+            (1, "1.00000"),
+            (0.33333, "0.33333"),
+            (1.0 / 3.0, "0.33333"),
+        ],
+    )
+    def test_key_format(self, lambda_value, expected):
+        """Keys are formatted to five decimal places, as SOMD2 does."""
+        assert GCMCSampler.stats_key(lambda_value) == expected
+
+    def test_key_is_stable_across_representations(self):
+        """Values that agree to five decimal places share a key."""
+        assert GCMCSampler.stats_key(0.1 + 0.2) == GCMCSampler.stats_key(0.3)
+
+
+class TestPerLambdaStats:
+    """Tests for statistics accumulated per lambda value."""
+
+    def test_isolated_between_lambdas(self):
+        """Moves at one lambda must not be counted at another."""
+        sampler = make_sampler(lambda_values=[0.0, 0.5])
+
+        do_moves(sampler, 3)
+        switch(sampler, 0.5)
+
+        # The new lambda starts from zero.
+        assert sampler._num_moves == 0
+
+        do_moves(sampler, 7)
+        switch(sampler, 0.0)
+
+        # Returning restores the original count, not the total.
+        assert sampler._num_moves == 3
+
+        stats = sampler.get_stats()
+        assert stats["0.00000"]["num_moves"] == 3
+        assert stats["0.50000"]["num_moves"] == 7
+
+    def test_accumulates_across_visits(self):
+        """Revisiting a lambda continues from where it left off."""
+        sampler = make_sampler(lambda_values=[0.0, 0.5])
+
+        do_moves(sampler, 3)
+        switch(sampler, 0.5)
+        do_moves(sampler, 7)
+        switch(sampler, 0.0)
+        do_moves(sampler, 2)
+
+        stats = sampler.get_stats()
+        assert stats["0.00000"]["num_moves"] == 5
+        assert stats["0.50000"]["num_moves"] == 7
+
+    def test_current_lambda_is_reported(self):
+        """The lambda in use is included alongside the archived ones."""
+        sampler = make_sampler(lambda_values=[0.0, 0.5])
+        do_moves(sampler, 4)
+
+        assert sampler.get_stats() == {
+            "0.00000": {
+                "num_moves": 4,
+                "num_accepted": 4,
+                "num_insertions": 0,
+                "num_deletions": 0,
+                "num_accepted_attempts": 0,
+            }
+        }
+
+    def test_non_alchemical_has_a_single_key(self):
+        """A non-alchemical system reports the same shape, with one key."""
+        sampler = make_sampler(is_fep=False)
+        do_moves(sampler, 6)
+
+        stats = sampler.get_stats()
+        assert list(stats) == ["0.00000"]
+        assert stats["0.00000"]["num_moves"] == 6
+
+    def test_reset_clears_every_lambda(self):
+        """reset() zeroes the current lambda and discards the others."""
+        sampler = make_sampler(lambda_values=[0.0, 0.5])
+        do_moves(sampler, 3)
+        switch(sampler, 0.5)
+        do_moves(sampler, 7)
+
+        sampler.reset()
+
+        assert sampler.get_stats() == {
+            "0.50000": {
+                "num_moves": 0,
+                "num_accepted": 0,
+                "num_insertions": 0,
+                "num_deletions": 0,
+                "num_accepted_attempts": 0,
+            }
+        }
+
+
+class TestRestoreStats:
+    """Tests for restoring statistics, e.g. from a checkpoint."""
+
+    def test_round_trip(self):
+        """Statistics survive a save and restore."""
+        sampler = make_sampler(lambda_values=[0.0, 0.5])
+        do_moves(sampler, 3)
+        switch(sampler, 0.5)
+        do_moves(sampler, 7)
+        stats = sampler.get_stats()
+
+        restored = make_sampler(lambda_value=0.5, lambda_values=[0.0, 0.5])
+        restored.restore_stats(stats)
+
+        assert restored.get_stats() == stats
+        assert restored._num_moves == 7
+
+    def test_unvisited_lambdas_are_ignored(self):
+        """
+        A sampler keeps only its own lambda values.
+
+        Each sampler can be handed the statistics for a whole simulation. If it
+        kept the others, it would report stale values for lambdas it never
+        samples, which could overwrite the live ones when merged.
+        """
+        stats = {
+            "0.00000": {
+                "num_moves": 5,
+                "num_accepted": 5,
+                "num_insertions": 0,
+                "num_deletions": 0,
+                "num_accepted_attempts": 0,
+            },
+            "1.00000": {
+                "num_moves": 9,
+                "num_accepted": 9,
+                "num_insertions": 0,
+                "num_deletions": 0,
+                "num_accepted_attempts": 0,
+            },
+        }
+
+        sampler = make_sampler(lambda_values=[0.0])
+        sampler.restore_stats(stats)
+
+        assert list(sampler.get_stats()) == ["0.00000"]
+
+    def test_merge_order_cannot_clobber(self):
+        """Merging several samplers is safe regardless of order."""
+        first = make_sampler(lambda_value=0.0, lambda_values=[0.0])
+        second = make_sampler(lambda_value=1.0, lambda_values=[1.0])
+
+        do_moves(first, 5)
+        do_moves(second, 9)
+
+        merged = {}
+        merged.update(first.get_stats())
+        merged.update(second.get_stats())
+
+        # Restart both from the merged statistics, then advance one of them.
+        first = make_sampler(lambda_value=0.0, lambda_values=[0.0])
+        second = make_sampler(lambda_value=1.0, lambda_values=[1.0])
+        first.restore_stats(merged)
+        second.restore_stats(merged)
+        do_moves(first, 100)
+
+        for order in ([first, second], [second, first]):
+            remerged = {}
+            for sampler in order:
+                remerged.update(sampler.get_stats())
+            assert remerged["0.00000"]["num_moves"] == 105
+            assert remerged["1.00000"]["num_moves"] == 9
+
+    def test_missing_lambda_is_zeroed(self):
+        """A lambda absent from the statistics starts from zero."""
+        sampler = make_sampler(lambda_value=0.5, lambda_values=[0.0, 0.5])
+        sampler.restore_stats(
+            {
+                "0.00000": {
+                    "num_moves": 5,
+                    "num_accepted": 5,
+                    "num_insertions": 0,
+                    "num_deletions": 0,
+                    "num_accepted_attempts": 0,
+                }
+            }
+        )
+
+        assert sampler._num_moves == 0
+        assert sampler.get_stats()["0.00000"]["num_moves"] == 5

From b06185f71190de43ec6e07ea82b042d597669889 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 15:09:59 +0100
Subject: [PATCH 09/16] Recount the GCMC region rather than trusting a cached
 count.

---
 CHANGELOG.md             |   1 +
 src/loch/_sampler.py     | 129 ++++++++++++++++++++++++---------------
 tests/test_num_waters.py |  51 ++++++++++++++++
 3 files changed, 131 insertions(+), 50 deletions(-)
 create mode 100644 tests/test_num_waters.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index faf3291..c88e13e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ Changelog
 * Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29).
 * Add `set_lambda` and `_precompute_lambdas` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context [#30](https://github.com/OpenBioSim/loch/pull/30).
 * Resolve unit conversions once when extracting the lambda dependent non-bonded parameters, rather than per atom, which dominated sampler setup [#33](https://github.com/OpenBioSim/loch/pull/33).
+* Recount the waters in the GCMC region rather than returning a cached count that was not invalidated when the positions changed, and return the whole box count directly when no region is defined, which previously raised [#36](https://github.com/OpenBioSim/loch/pull/36).
 
 [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026
 -------------------------------------------------------------------------------------
diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index e60881a..f2353ad 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -796,6 +796,11 @@ def __init__(
         # Flag for whether the last move was a bulk sampling move.
         self._is_bulk = False
 
+        # The number of waters in the GCMC region, as of the last count. This
+        # is what num_waters() reports, and is separate from self._N, which is
+        # the count for the volume that move() samples. None when unknown.
+        self._N_region = None
+
         import sys
 
         # Create a logger that writes to stderr and the log file.
@@ -1248,74 +1253,88 @@ def num_waters(self, context=None) -> int:
         """
         Return the number of waters in the GCMC region.
 
+        Parameters
+        ----------
+
+        context: openmm.Context, optional
+            The OpenMM context to count the waters from. If None, then the
+            internal context is used if one is available, otherwise the count
+            from the last move is returned.
+
         Returns
         -------
 
         num_waters: int
             The number of waters.
-
-        context: openmm.Context, optional
-            The OpenMM context to use for counting the waters. If None, then the
-            internal context will be used if available.
         """
 
-        # Whether we need to recalculate the number of waters in the GCMC sphere.
-        recalculate = context is not None or (
-            self._reference is not None and self._is_bulk
-        )
+        # Without a region every move samples the whole box, so the count that
+        # move() maintains is already the answer. There is also no reference to
+        # take a sphere centre from.
+        if self._reference is None:
+            return self._N
 
-        # We need to recalculate the number of waters.
-        if recalculate:
-            if context is None:
-                if not self._openmm_context:
-                    msg = "OpenMM context is not set!"
-                    _logger.error(msg)
-                    raise RuntimeError(msg)
-                else:
-                    context = self._openmm_context
+        # Fall back to the internal context, which is stored by a bulk move.
+        if context is None:
+            context = self._openmm_context
 
-            # Get the OpenMM state.
-            state = context.getState(getPositions=True)
+        # There is nothing to count from, so return the count from the last
+        # move. A bulk move clears this, since it counts the whole box rather
+        # than the region, and cannot answer for the region.
+        if context is None:
+            if self._N_region is None:
+                msg = "OpenMM context is not set!"
+                _logger.error(msg)
+                raise RuntimeError(msg)
 
-            # Get the current positions in Angstrom.
-            positions = state.getPositions(asNumpy=True) / _openmm.unit.angstrom
+            return self._N_region
 
-            # Get the position of the GCMC sphere centre.
-            target = self._backend.to_gpu(
-                self._get_target_position(positions).astype(_np.float32)
-            )
+        # Recount. The positions change outside of the sampler's control, via
+        # dynamics between moves, or a context being handed to another replica,
+        # so a stored count cannot be re-used when there is a context to count
+        # from.
 
-            # Upload atom positions to GPU.
-            self._gpu_position = self._backend.to_gpu(_as_float32(positions).flatten())
+        # Get the OpenMM state.
+        state = context.getState(getPositions=True)
 
-            # Find the non-ghost waters within the GCMC region.
-            self._kernels["deletion"](
-                _np.int32(self._num_waters),
-                self._deletion_candidates,
-                self._backend.to_gpu(target.astype(_np.float32)),
-                _np.float32(self._radius.value()),
-                self._gpu_position,
-                self._gpu_water_idx,
-                self._gpu_water_state,
-                self._gpu_cell_matrix_inverse,
-                self._gpu_M,
-                block=(self._num_threads, 1, 1),
-                grid=(self._water_blocks, 1, 1),
-            )
+        # Get the current positions in Angstrom.
+        positions = state.getPositions(asNumpy=True) / _openmm.unit.angstrom
+
+        # Get the position of the GCMC sphere centre.
+        target = self._backend.to_gpu(
+            self._get_target_position(positions).astype(_np.float32)
+        )
 
-            # Get the candidates.
-            candidates = self._backend.from_gpu(self._deletion_candidates).flatten()
+        # Upload atom positions to GPU. This is re-uploaded by the next move,
+        # so overwriting it here is safe.
+        self._gpu_position = self._backend.to_gpu(_as_float32(positions).flatten())
 
-            # Find the waters within the GCMC sphere.
-            candidates = _np.where(candidates == 1)[0]
+        # Find the non-ghost waters within the GCMC region.
+        self._kernels["deletion"](
+            _np.int32(self._num_waters),
+            self._deletion_candidates,
+            self._backend.to_gpu(target.astype(_np.float32)),
+            _np.float32(self._radius.value()),
+            self._gpu_position,
+            self._gpu_water_idx,
+            self._gpu_water_state,
+            self._gpu_cell_matrix_inverse,
+            self._gpu_M,
+            block=(self._num_threads, 1, 1),
+            grid=(self._water_blocks, 1, 1),
+        )
 
-            # Set the number of waters.
-            self._N = len(candidates)
+        # Get the candidates.
+        candidates = self._backend.from_gpu(self._deletion_candidates).flatten()
 
-            # Reset the bulk sampling flag.
-            self._is_bulk = False
+        # Find the waters within the GCMC sphere.
+        candidates = _np.where(candidates == 1)[0]
+
+        # Store the number of waters in the region. self._N is left alone, as
+        # it belongs to move(), where it must match the volume being sampled.
+        self._N_region = len(candidates)
 
-        return self._N
+        return self._N_region
 
     def num_accepted_moves(self) -> int:
         """
@@ -1415,6 +1434,9 @@ def reset(self) -> None:
         # Clear the OpenMM context.
         self._openmm_context = None
 
+        # The stored region count refers to the cleared context.
+        self._N_region = None
+
     @staticmethod
     def stats_key(lambda_value: float) -> str:
         """
@@ -1721,6 +1743,13 @@ def move(self, context: _openmm.Context) -> list[int]:
                 # Set the number of waters.
                 self._N = len(deletion_candidates)
 
+                # A bulk move counts the whole box, so it cannot report the
+                # region. Anything else counts the region directly.
+                if self._is_bulk:
+                    self._N_region = None
+                else:
+                    self._N_region = self._N
+
             # Reset the batch acceptance flag.
             is_accepted = False
 
diff --git a/tests/test_num_waters.py b/tests/test_num_waters.py
new file mode 100644
index 0000000..97dfa58
--- /dev/null
+++ b/tests/test_num_waters.py
@@ -0,0 +1,51 @@
+import pytest
+
+from loch import GCMCSampler
+
+
+def make_sampler(reference="resname LIG", N=0, N_region=None, openmm_context=None):
+    """
+    Create a sampler with only the attributes num_waters() uses, so that the
+    counting logic can be tested without a system or a GPU.
+    """
+    sampler = object.__new__(GCMCSampler)
+    sampler._reference = reference
+    sampler._N = N
+    sampler._N_region = N_region
+    sampler._openmm_context = openmm_context
+    sampler._is_bulk = False
+    return sampler
+
+
+def test_num_waters_without_a_region():
+    """
+    Without a GCMC region every move samples the whole box, so the count that
+    move() maintains is already the answer. Counting a region would need a
+    reference to take a sphere centre from, which does not exist in this case.
+    """
+    sampler = make_sampler(reference=None, N=7)
+
+    assert sampler.num_waters() == 7
+
+    # Passing a context must not send it down the recount path either, which
+    # would dereference the reference indices that were never set.
+    assert sampler.num_waters(context=object()) == 7
+
+
+def test_num_waters_reports_the_stored_region_count():
+    """With a region and nothing to count from, the stored count is returned."""
+    sampler = make_sampler(N=99, N_region=4)
+
+    assert sampler.num_waters() == 4
+
+
+def test_num_waters_refuses_a_whole_box_count():
+    """
+    A bulk move leaves self._N counting the whole box, so it cannot answer for
+    the region. With no context to recount from, that must raise rather than
+    report the box count as though it were the region count.
+    """
+    sampler = make_sampler(N=99, N_region=None)
+
+    with pytest.raises(RuntimeError, match="OpenMM context is not set"):
+        sampler.num_waters()

From a89b0aced5d8cdf589881232f0fb23fe9338f7d2 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 15:45:01 +0100
Subject: [PATCH 10/16] Seed the water count so that it is right before the
 first move.

---
 src/loch/_sampler.py     | 18 +++++++++++++++---
 tests/test_num_waters.py | 24 ++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index f2353ad..ba0c852 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -1253,6 +1253,13 @@ def num_waters(self, context=None) -> int:
         """
         Return the number of waters in the GCMC region.
 
+        With no region there is nothing to recount. self._N then comes from
+        self._water_state, an occupancy array that only insertion and deletion
+        modify, rather than from a geometric test against a sphere centre as
+        the region count is. A water cannot leave the box, so dynamics cannot
+        change it, and consulting the context would cost a state fetch and a
+        kernel launch to arrive at the same number.
+
         Parameters
         ----------
 
@@ -1268,9 +1275,9 @@ def num_waters(self, context=None) -> int:
             The number of waters.
         """
 
-        # Without a region every move samples the whole box, so the count that
-        # move() maintains is already the answer. There is also no reference to
-        # take a sphere centre from.
+        # Without a region the count that move() maintains is already the
+        # answer, and there is no reference to take a sphere centre from. See
+        # the docstring for why the context is not consulted.
         if self._reference is None:
             return self._N
 
@@ -2849,6 +2856,11 @@ def _initialise_gpu_memory(self):
         self._non_ghost_waters_cache = None
         self._invalidate_water_caches()
 
+        # Seed the water count. Every move sets this for the volume that it
+        # samples, but with no region num_waters() reports it directly, and
+        # would otherwise report zero until the first move has run.
+        self._N = len(self._non_ghost_waters_cache)
+
         # Pre-allocate zero target array for bulk sampling.
         self._zero_target_gpu = self._backend.to_gpu(_np.zeros(3, dtype=_np.float32))
 
diff --git a/tests/test_num_waters.py b/tests/test_num_waters.py
index 97dfa58..641d821 100644
--- a/tests/test_num_waters.py
+++ b/tests/test_num_waters.py
@@ -32,6 +32,30 @@ def test_num_waters_without_a_region():
     assert sampler.num_waters(context=object()) == 7
 
 
+def test_num_waters_is_seeded_before_the_first_move():
+    """
+    self._N is set as each move runs, so without a region it would report zero
+    until the first one had. A cycle can complete without any GCMC moves when
+    'gcmc_frequency' is coarser than the checkpoint interval, so seeding it
+    during setup is what stops a checkpoint logging zero waters.
+    """
+    import numpy as np
+
+    from loch import GCMCSampler
+
+    sampler = object.__new__(GCMCSampler)
+
+    # Four waters, of which the last two are ghosts.
+    sampler._water_state = np.array([1, 1, 0, 0], dtype=np.int32)
+    sampler._invalidate_water_caches()
+
+    # The seeding that _initialise_gpu_memory() performs.
+    sampler._N = len(sampler._non_ghost_waters_cache)
+
+    sampler._reference = None
+    assert sampler.num_waters() == 2
+
+
 def test_num_waters_reports_the_stored_region_count():
     """With a region and nothing to count from, the stored count is returned."""
     sampler = make_sampler(N=99, N_region=4)

From 82e64a8ca9b99eaf2f64cd834fd334b49691b998 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Sun, 2 Aug 2026 16:27:59 +0100
Subject: [PATCH 11/16] Document when num_waters() may be called without a
 context.

---
 src/loch/_sampler.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index ba0c852..17aff4f 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -1266,7 +1266,10 @@ def num_waters(self, context=None) -> int:
         context: openmm.Context, optional
             The OpenMM context to count the waters from. If None, then the
             internal context is used if one is available, otherwise the count
-            from the last move is returned.
+            from the last move is returned. Only omit this immediately after a
+            move. Pass a context if dynamics have been run since, otherwise
+            waters will have crossed the region boundary and the stored count
+            will be out of date.
 
         Returns
         -------

From 6c6b580ec2df69780e05ad6c9f78b20ee6018da5 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 5 Aug 2026 11:27:31 +0100
Subject: [PATCH 12/16] Test that set_lambda uploads the parameters for the new
 lambda value.

---
 tests/test_energy.py | 63 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/tests/test_energy.py b/tests/test_energy.py
index 703a8a2..8a1d2f2 100644
--- a/tests/test_energy.py
+++ b/tests/test_energy.py
@@ -616,3 +616,66 @@ def test_finalise_system(platform, water_box):
     assert actual_waters == expected_waters, (
         f"Water count mismatch: got {actual_waters}, expected {expected_waters}"
     )
+
+
+@pytest.mark.skipif(
+    "CUDA_VISIBLE_DEVICES" not in os.environ,
+    reason="Requires CUDA enabled GPU.",
+)
+@pytest.mark.parametrize("platform", ["cuda", "opencl"])
+def test_set_lambda_uploads_parameters(sd12, platform):
+    """
+    Test that set_lambda() uploads the parameters for the new lambda value.
+
+    A sampler may be re-used across lambda values, re-uploading the lambda
+    dependent non-bonded parameters rather than rebuilding a context. If the
+    upload were skipped, or the wrong entry taken from the cache, the sampler
+    would evaluate insertion and deletion energies against the parameters of
+    the lambda value it was built at, while reporting the new one.
+    """
+
+    mols, reference = sd12
+
+    # The end states, so that the parameters differ as much as they can.
+    build_lambda = 0.0
+    run_lambda = 1.0
+
+    sampler = GCMCSampler(
+        mols,
+        cutoff_type="rf",
+        cutoff="10 A",
+        reference=reference,
+        lambda_schedule=sr.cas.LambdaSchedule.standard_morph(),
+        lambda_value=build_lambda,
+        lambda_values=[run_lambda],
+        log_level="error",
+        ghost_file=None,
+        log_file=None,
+        test=True,
+        platform=platform,
+    )
+
+    assert sampler._is_fep, "the test system must be perturbable"
+
+    names = ("_gpu_charge", "_gpu_sigma", "_gpu_epsilon", "_gpu_alpha")
+    built = sampler._lambda_params[(build_lambda, sampler._rest2_scale)]
+    wanted = sampler._lambda_params[(run_lambda, sampler._rest2_scale)]
+
+    # The two end points must differ, otherwise the test cannot tell whether
+    # the upload happened.
+    assert any(not np.allclose(a, b) for a, b in zip(built, wanted)), (
+        "the parameters are the same at both lambda values"
+    )
+
+    sampler.set_lambda(run_lambda)
+    assert sampler._lambda_value == run_lambda
+
+    sampler.push()
+    try:
+        for name, expected in zip(names, wanted):
+            on_gpu = np.asarray(sampler._backend.from_gpu(getattr(sampler, name)))
+            assert np.allclose(on_gpu, expected), (
+                f"{name} does not hold the parameters for lambda {run_lambda}"
+            )
+    finally:
+        sampler.pop()

From 61acd3fba5d0b295be6b50f271f46eebd6984ea3 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 5 Aug 2026 11:50:36 +0100
Subject: [PATCH 13/16] Fix the GCMC region methods on OpenCL and without a
 region.

---
 CHANGELOG.md         |  1 +
 src/loch/_sampler.py | 10 ++++++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c88e13e..db12ea6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ Changelog
 * Add `set_lambda` and `_precompute_lambdas` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context [#30](https://github.com/OpenBioSim/loch/pull/30).
 * Resolve unit conversions once when extracting the lambda dependent non-bonded parameters, rather than per atom, which dominated sampler setup [#33](https://github.com/OpenBioSim/loch/pull/33).
 * Recount the waters in the GCMC region rather than returning a cached count that was not invalidated when the positions changed, and return the whole box count directly when no region is defined, which previously raised [#36](https://github.com/OpenBioSim/loch/pull/36).
+* Stop uploading the GCMC region centre to the GPU twice in `delete_waters` and `num_waters`, which raised on the OpenCL platform, and make `delete_waters` a no-op when no region is defined, which previously raised [#39](https://github.com/OpenBioSim/loch/pull/39).
 
 [2026.1.0](https://github.com/openbiosim/loch/compare/2025.2.0...2026.1.0) - Jun 2026
 -------------------------------------------------------------------------------------
diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py
index 17aff4f..3c11e24 100644
--- a/src/loch/_sampler.py
+++ b/src/loch/_sampler.py
@@ -1195,6 +1195,8 @@ def delete_waters(self, context: _openmm.Context) -> None:
         """
         Delete any waters within the GCMC sphere. (Convert to ghosts.)
 
+        This does nothing when there is no sphere.
+
         Parameters
         ----------
 
@@ -1202,6 +1204,10 @@ def delete_waters(self, context: _openmm.Context) -> None:
             The OpenMM context to use.
         """
 
+        # There is no sphere to empty.
+        if self._reference is None:
+            return
+
         # Set the NonBondedForce(s).
         self._set_nonbonded_forces(context)
 
@@ -1223,7 +1229,7 @@ def delete_waters(self, context: _openmm.Context) -> None:
         self._kernels["deletion"](
             _np.int32(self._num_waters),
             self._deletion_candidates,
-            self._backend.to_gpu(target.astype(_np.float32)),
+            target,
             _np.float32(self._radius.value()),
             self._gpu_position,
             self._gpu_water_idx,
@@ -1323,7 +1329,7 @@ def num_waters(self, context=None) -> int:
         self._kernels["deletion"](
             _np.int32(self._num_waters),
             self._deletion_candidates,
-            self._backend.to_gpu(target.astype(_np.float32)),
+            target,
             _np.float32(self._radius.value()),
             self._gpu_position,
             self._gpu_water_idx,

From 0da76ce15e5852a101cd0d388527a4074d6440f0 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Wed, 5 Aug 2026 11:51:03 +0100
Subject: [PATCH 14/16] Test the energy after switching a sampler to another
 lambda value.

---
 tests/test_energy.py | 109 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 109 insertions(+)

diff --git a/tests/test_energy.py b/tests/test_energy.py
index 8a1d2f2..4a5ea06 100644
--- a/tests/test_energy.py
+++ b/tests/test_energy.py
@@ -48,6 +48,9 @@ def test_energy(fixture, softcore_form, platform, request):
         lambda_schedule=schedule,
         lambda_value=lambda_value,
         softcore_form=softcore_form,
+        # Sample within the region when there is one. Without a reference
+        # every move is a bulk move regardless.
+        bulk_sampling_probability=0.0 if reference is not None else 0.1,
         log_level="debug",
         ghost_file=None,
         log_file=None,
@@ -79,6 +82,10 @@ def test_energy(fixture, softcore_form, platform, request):
         map=dyn_map,
     )
 
+    # Empty the region, since at equilibrium it is full and insertions into it
+    # are rejected. A no-op when there is no region.
+    sampler.delete_waters(d.context())
+
     # Loop until we accept an insertion move.
     is_accepted = False
     while not is_accepted:
@@ -679,3 +686,105 @@ def test_set_lambda_uploads_parameters(sd12, platform):
             )
     finally:
         sampler.pop()
+
+
+@pytest.mark.skipif(
+    "CUDA_VISIBLE_DEVICES" not in os.environ,
+    reason="Requires CUDA enabled GPU.",
+)
+@pytest.mark.parametrize("platform", ["cuda", "opencl"])
+def test_energy_after_set_lambda(sd12, platform):
+    """
+    Test that the RF energy difference agrees with OpenMM after the sampler has
+    been switched to a different lambda value.
+
+    This checks the uploaded parameters through the physics rather than by
+    inspecting them, so it also covers the kernel using them correctly.
+
+    The move has to happen where the perturbation is. Bulk sampling would
+    place the water anywhere in the box, typically tens of Angstrom from the
+    perturbable molecule, where the lambda dependent parameters contribute
+    nothing and the comparison holds however wrong they are. The sphere is
+    emptied first, since at equilibrium it is full and insertions into it are
+    rejected.
+    """
+
+    mols, reference = sd12
+
+    schedule = sr.cas.LambdaSchedule.standard_morph()
+
+    # The end states, so that the parameters differ as much as they can.
+    build_lambda = 0.0
+    run_lambda = 1.0
+
+    sampler = GCMCSampler(
+        mols,
+        cutoff_type="rf",
+        cutoff="10 A",
+        reference=reference,
+        lambda_schedule=schedule,
+        lambda_value=build_lambda,
+        lambda_values=[run_lambda],
+        bulk_sampling_probability=0.0,
+        log_level="debug",
+        ghost_file=None,
+        log_file=None,
+        test=True,
+        platform=platform,
+    )
+
+    sampler.set_lambda(run_lambda)
+
+    d = sampler.system().dynamics(
+        cutoff_type="rf",
+        cutoff="10 A",
+        temperature="298 K",
+        pressure=None,
+        constraint="h_bonds",
+        timestep="2 fs",
+        schedule=schedule,
+        lambda_value=run_lambda,
+        shift_coulomb=str(sampler._shift_coulomb),
+        shift_delta=str(sampler._shift_delta),
+        platform=platform,
+    )
+
+    def potential_energy():
+        return (
+            d.context()
+            .getState(getEnergy=True)
+            .getPotentialEnergy()
+            .value_in_unit(openmm.unit.kilocalories_per_mole)
+        )
+
+    # Empty the sphere so that the next accepted move is an insertion into it.
+    sampler.delete_waters(d.context())
+
+    for _ in range(50):
+        initial_energy = potential_energy()
+        moves = sampler.move(d.context())
+        if moves:
+            break
+    else:
+        pytest.fail("no GCMC move was accepted")
+
+    energy_difference = potential_energy() - initial_energy
+    sampler_energy = sampler._debug["energy_coul"] + sampler._debug["energy_lj"]
+
+    # The move must be near the atoms whose parameters perturb, otherwise the
+    # comparison cannot see them.
+    charges0 = np.asarray(sampler._lambda_params[(build_lambda, 1.0)][0])
+    charges1 = np.asarray(sampler._lambda_params[(run_lambda, 1.0)][0])
+    changed = np.where(np.abs(charges0 - charges1) > 1e-9)[0]
+    positions = (
+        d.context().getState(getPositions=True).getPositions(asNumpy=True)
+        / omm_unit.angstrom
+    )
+    oxygen = positions[sampler._water_indices[sampler._debug["idx"]]]
+    distances = np.linalg.norm(positions[changed] - oxygen, axis=1)
+    assert (distances <= 10).any(), (
+        "no perturbing atom within the cutoff of the move, so the lambda "
+        "dependent parameters are not being tested"
+    )
+
+    assert math.isclose(energy_difference, sampler_energy, abs_tol=1e-2)

From 42cd05b12ecab1d5ba5ad9ca8b7e5543edcb6ec5 Mon Sep 17 00:00:00 2001
From: Lester Hedges <lester.hedges@gmail.com>
Date: Tue, 15 Sep 2026 09:51:31 +0100
Subject: [PATCH 15/16] Update BioSimSpace pin.

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

diff --git a/pixi.toml b/pixi.toml
index 56780fd..35498d1 100644
--- a/pixi.toml
+++ b/pixi.toml
@@ -6,9 +6,9 @@ platforms = ["linux-64", "osx-arm64", "win-64"]
 [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"
 loguru = "*"
 pyopencl = "*"
 
diff --git a/recipes/loch/recipe.yaml b/recipes/loch/recipe.yaml
index 1fb1f79..f93af36 100644
--- a/recipes/loch/recipe.yaml
+++ b/recipes/loch/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
     - loguru
     - pyopencl
     - python

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

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index db12ea6..a0b1e4d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,10 +1,9 @@
 Changelog
 =========
 
-[2026.2.0](https://github.com/openbiosim/loch/compare/2026.1.0...2026.2.0) - ********
+[2026.2.0](https://github.com/openbiosim/loch/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.
 * Restrict PME energy calculation to required force groups [#29](https://github.com/OpenBioSim/loch/pull/29).
 * Add `set_lambda` and `_precompute_lambdas` so that a sampler can be re-used across lambda values without rebuilding an OpenMM context [#30](https://github.com/OpenBioSim/loch/pull/30).
 * Resolve unit conversions once when extracting the lambda dependent non-bonded parameters, rather than per atom, which dominated sampler setup [#33](https://github.com/OpenBioSim/loch/pull/33).