Backport fix from PR #36

src/loch/_sampler.py

@@ -735,6 +735,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.
@@ -1187,74 +1192,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:
         """
@@ -1357,6 +1376,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
+
     def restore_stats(self, stats: dict) -> None:
         """
         Restore sampler statistics from a dictionary.
@@ -1565,6 +1587,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
 

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