Handle per-lambda sampler statistics

src/loch/_sampler.py

@@ -775,13 +775,18 @@ def __init__(
         # Zero the number of waters in the sampling volume.
         self._N = 0
 
-        # Zero the statistics.
+        # Zero the statistics for the current lambda value.
         self._num_moves = 0
         self._num_accepted = 0
         self._num_accepted_attempts = 0
         self._num_insertions = 0
         self._num_deletions = 0
 
+        # Statistics for lambda values other than the current one, keyed by
+        # formatted lambda value. The current one is held in the counters above
+        # and archived here by set_lambda().
+        self._stats = {}
+
         # Null the nonbonded forces.
         self._nonbonded_force = None
         self._custom_nonbonded_force = None
@@ -1397,14 +1402,9 @@ def reset(self) -> None:
         """
         Reset the sampler.
         """
-        # Zero the number of accepted moves.
-        self._num_accepted = 0
-        self._num_insertions = 0
-        self._num_deletions = 0
-        self._num_moves = 0
-        self._num_accepted_attempts = 0
-        self._num_accepted_insertions = 0
-        self._num_accepted_deletions = 0
+        # Zero the number of accepted moves, for every lambda value.
+        self._zero_stats()
+        self._stats = {}
 
         # Clear the forces.
         self._nonbonded_force = None
@@ -1415,43 +1415,137 @@ def reset(self) -> None:
         # Clear the OpenMM context.
         self._openmm_context = None
 
-    def restore_stats(self, stats: dict) -> None:
+    @staticmethod
+    def stats_key(lambda_value: float) -> str:
         """
-        Restore sampler statistics from a dictionary.
+        Return the key used to store statistics for a lambda value.
 
         Parameters
         ----------
 
-        stats : dict
-            Dictionary of sampler statistics as returned by ``get_stats()``.
+        lambda_value: float
+            The lambda value.
+
+        Returns
+        -------
+
+        str
+            The key.
+        """
+        return f"{float(lambda_value):.5f}"
+
+    def _zero_stats(self) -> None:
+        """
+        Zero the statistics for the current lambda value.
+        """
+        self._num_moves = 0
+        self._num_accepted = 0
+        self._num_insertions = 0
+        self._num_deletions = 0
+        self._num_accepted_attempts = 0
+
+    def _stats_keys(self) -> set:
+        """
+        Return the keys of the lambda values that this sampler visits.
+        """
+        keys = {self.stats_key(self._lambda_value)}
+
+        if self._lambda_values is not None:
+            keys.update(self.stats_key(x) for x in self._lambda_values)
+
+        return keys
+
+    def _switch_stats(self, lambda_value: float) -> None:
+        """
+        Archive the statistics for the current lambda value and load those for
+        a new one, zeroing them if it hasn't been visited before.
+
+        Parameters
+        ----------
+
+        lambda_value: float
+            The lambda value being switched to.
+        """
+        old_key = self.stats_key(self._lambda_value)
+        new_key = self.stats_key(lambda_value)
+
+        if old_key == new_key:
+            return
+
+        self._stats[old_key] = self._get_current_stats()
+
+        if new_key in self._stats:
+            self._set_current_stats(self._stats.pop(new_key))
+        else:
+            self._zero_stats()
+
+    def _get_current_stats(self) -> dict:
+        """
+        Return the statistics for the current lambda value.
+        """
+        return {
+            "num_moves": self._num_moves,
+            "num_accepted": self._num_accepted,
+            "num_insertions": self._num_insertions,
+            "num_deletions": self._num_deletions,
+            "num_accepted_attempts": self._num_accepted_attempts,
+        }
+
+    def _set_current_stats(self, stats: dict) -> None:
+        """
+        Set the statistics for the current lambda value. Unrecognised entries
+        are ignored, so that statistics written by an older version can still
+        be restored.
         """
         self._num_moves = stats["num_moves"]
         self._num_accepted = stats["num_accepted"]
         self._num_insertions = stats["num_insertions"]
         self._num_deletions = stats["num_deletions"]
         self._num_accepted_attempts = stats["num_accepted_attempts"]
-        self._num_accepted_insertions = stats["num_accepted_insertions"]
-        self._num_accepted_deletions = stats["num_accepted_deletions"]
+
+    def restore_stats(self, stats: dict) -> None:
+        """
+        Restore sampler statistics.
+
+        Parameters
+        ----------
+
+        stats : dict
+            Statistics as returned by ``get_stats()``, i.e. keyed by lambda
+            value. Entries for lambda values that this sampler doesn't visit
+            are ignored, so it is safe to pass the statistics for a whole
+            simulation to each of several samplers. Were they retained, a
+            sampler would report stale statistics for another's lambda values,
+            which could then overwrite the live ones when merged.
+        """
+        keys = self._stats_keys()
+        self._stats = {key: dict(value) for key, value in stats.items() if key in keys}
+
+        # Load the statistics for the current lambda value, if present.
+        key = self.stats_key(self._lambda_value)
+        if key in self._stats:
+            self._set_current_stats(self._stats.pop(key))
+        else:
+            self._zero_stats()
 
     def get_stats(self) -> dict:
         """
-        Return the current sampler statistics as a dictionary.
+        Return the sampler statistics, keyed by lambda value.
+
+        A sampler may be switched between lambda values via ``set_lambda()``,
+        and accumulates statistics for each of them separately. Non-alchemical
+        systems have a single key, for the lambda value the sampler was created
+        with.
 
         Returns
         -------
 
         dict
-            Dictionary of sampler statistics.
+            Dictionary of sampler statistics for each lambda value.
         """
-        return {
-            "num_moves": self._num_moves,
-            "num_accepted": self._num_accepted,
-            "num_insertions": self._num_insertions,
-            "num_deletions": self._num_deletions,
-            "num_accepted_attempts": self._num_accepted_attempts,
-            "num_accepted_insertions": self._num_accepted_insertions,
-            "num_accepted_deletions": self._num_accepted_deletions,
-        }
+        stats = {key: dict(value) for key, value in self._stats.items()}
+        stats[self.stats_key(self._lambda_value)] = self._get_current_stats()
+        return stats
 
     def ghost_residues(self) -> _np.ndarray:
         """
@@ -2510,6 +2604,10 @@ def set_lambda(
         if lambda_value == self._lambda_value and rest2_scale == self._rest2_scale:
             return
 
+        # Statistics are accumulated per lambda value, so archive those for the
+        # current one and load those for the new one.
+        self._switch_stats(lambda_value)
+
         # There are no lambda dependent parameters for a non-alchemical system.
         if not self._is_fep:
             self._lambda_value = lambda_value