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