Store GCMC statistics per lambda value

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

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.

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,14 +144,24 @@ 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"):
             self._terminal_flip_stats = [[0, 0]] * n
         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 = [

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)

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