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