From 8d069d19c6ee80749d56b38f49832ab2020c82df Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Sun, 2 Aug 2026 15:26:57 +0100 Subject: [PATCH] Backport fix from PR #36. [ci skip] --- src/loch/_sampler.py | 129 ++++++++++++++++++++++++--------------- tests/test_num_waters.py | 51 ++++++++++++++++ 2 files changed, 130 insertions(+), 50 deletions(-) create mode 100644 tests/test_num_waters.py diff --git a/src/loch/_sampler.py b/src/loch/_sampler.py index df65570..54df198 100644 --- a/src/loch/_sampler.py +++ b/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 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()