Fix issue #217

CHANGELOG.md

@@ -8,6 +8,7 @@ Changelog
 * Increase the default `cutoff` from 7.5 Å to 9 Å, matching common practice and giving faster PME on current GPUs, since the shorter cutoff shifts too much work onto the reciprocal space grid [#209](https://github.com/OpenBioSim/somd2/pull/209).
 * Buffer energy components and write them at checkpoint time, rather than rewriting the parquet file on every energy save, which cost a few milliseconds per replica per cycle and grew with the length of the run [#212](https://github.com/OpenBioSim/somd2/pull/212).
 * Silence Sire's progress bars when a runner is constructed rather than when `somd2` is imported, so that importing `somd2` as a library no longer changes how Sire reports progress [#215](https://github.com/OpenBioSim/somd2/pull/215).
+* Save the replica exchange state once at the end of a run rather than twice when the last cycle is a checkpoint cycle, and include the GCMC statistics in the final save [#218](https://github.com/OpenBioSim/somd2/pull/218).
 
 [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026
 --------------------------------------------------------------------------------------

src/somd2/runner/_repex.py

@@ -1953,6 +1953,9 @@ def run(self):
         # to handle non-integer ratios between the checkpoint and energy frequencies.
         next_checkpoint = cycles_per_checkpoint
 
+        # Whether the most recent cycle saved the replica exchange state.
+        is_checkpoint = False
+
         # Perform the replica exchange simulation.
         for i in range(cycles):
             _logger.info(f"Running dynamics for cycle {i + 1} of {cycles}")
@@ -2155,54 +2158,18 @@ def run(self):
                 # Advance the checkpoint threshold.
                 next_checkpoint += cycles_per_checkpoint
 
-                # Guard the repex state and transition matrix saving with a file lock.
-                lock = _FileLock(self._lock_file)
-                with lock.acquire(timeout=self._config.timeout.to("seconds")):
-                    # Save the transition matrix.
-                    _logger.info("Saving replica exchange transition matrix")
-                    self._save_transition_matrix()
-
-                    # Backup the dynamics cache pickle file, if it exists.
-                    if self._repex_state.exists():
-                        _copyfile(
-                            self._repex_state,
-                            self._repex_state.with_suffix(".pkl.bak"),
-                        )
-
-                    # Pickle the dynamics cache.
-                    _logger.info("Saving replica exchange state")
-                    self._save_sampler_stats()
-                    with open(self._repex_state, "wb") as f:
-                        _pickle.dump(self._dynamics_cache, f)
+                self._save_repex_state()
 
         dynamics_executor.shutdown(wait=True)
         checkpoint_executor.shutdown(wait=True)
 
         # Record the end time for the production block.
         prod_end = time()
 
-        lock = _FileLock(self._lock_file)
-        with lock.acquire(timeout=self._config.timeout.to("seconds")):
-            # Save the final transition matrix.
-            _logger.info("Saving final replica exchange transition matrix")
-            self._save_transition_matrix()
-
-            # Backup the dynamics cache pickle file, if it exists.
-            if self._repex_state.exists():
-                _copyfile(
-                    self._repex_state,
-                    self._repex_state.with_suffix(".pkl.bak"),
-                )
-
-            # Pickle final state of the dynamics cache.
-            _logger.info("Saving final replica exchange state")
-            if self._terminal_flip_samplers is not None:
-                self._dynamics_cache._terminal_flip_stats = [
-                    [s.num_attempted, s.num_accepted]
-                    for s in self._terminal_flip_samplers
-                ]
-            with open(self._repex_state, "wb") as f:
-                _pickle.dump(self._dynamics_cache, f)
+        # Save the final state, unless the last cycle was a checkpoint cycle
+        # and has just done so.
+        if not is_checkpoint:
+            self._save_repex_state(final=True)
 
         # Record the end time.
         end = time()
@@ -3060,6 +3027,35 @@ def _merge_gcmc_stats(self):
 
         return stats if stats else None
 
+    def _save_repex_state(self, final=False):
+        """
+        Save the transition matrix and pickle the dynamics cache, backing up
+        the previous pickle, under the file lock.
+
+        Parameters
+        ----------
+
+        final: bool
+            Whether this is the final save of the run, for logging.
+        """
+        label = "final replica exchange" if final else "replica exchange"
+
+        lock = _FileLock(self._lock_file)
+        with lock.acquire(timeout=self._config.timeout.to("seconds")):
+            _logger.info(f"Saving {label} transition matrix")
+            self._save_transition_matrix()
+
+            if self._repex_state.exists():
+                _copyfile(
+                    self._repex_state,
+                    self._repex_state.with_suffix(".pkl.bak"),
+                )
+
+            _logger.info(f"Saving {label} state")
+            self._save_sampler_stats()
+            with open(self._repex_state, "wb") as f:
+                _pickle.dump(self._dynamics_cache, f)
+
     def _save_sampler_stats(self):
         """
         Save GCMC and terminal flip sampler statistics to the dynamics cache

tests/runner/test_repex.py

@@ -458,9 +458,51 @@ def acquire(self, *args, **kwargs):
             repex_module._FileLock = real_filelock
 
     # Two cycles, each taking the lock once for the checkpoint files and once
-    # for the repex state, plus a final acquisition. This must not scale with
-    # the number of passes.
-    assert len(acquisitions) == 5
+    # for the repex state. The last cycle is a checkpoint cycle, so there is no
+    # separate final save. This must not scale with the number of passes.
+    assert len(acquisitions) == 4
+
+
+@pytest.mark.skipif(not has_cuda, reason="CUDA not available.")
+@pytest.mark.parametrize(
+    "runtime, checkpoint_frequency, expected",
+    [("8fs", "4fs", [False, False]), ("12fs", "8fs", [False, True])],
+)
+def test_repex_state_saved_once(
+    ethane_methanol, runtime, checkpoint_frequency, expected
+):
+    """
+    Validate that the replica exchange state is saved once per checkpoint
+    cycle, with a separate final save only when the last cycle is not a
+    checkpoint cycle.
+    """
+    with tempfile.TemporaryDirectory() as tmpdir:
+        config = {
+            "runtime": runtime,
+            "restart": False,
+            "output_directory": tmpdir,
+            "energy_frequency": "4fs",
+            "checkpoint_frequency": checkpoint_frequency,
+            "frame_frequency": "4fs",
+            "platform": "cuda",
+            "max_threads": 1,
+            "num_lambda": 2,
+            "replica_exchange": True,
+        }
+        runner = RepexRunner(ethane_methanol, Config(**config))
+
+        saves = []
+        save = runner._save_repex_state
+
+        def counting_save(final=False):
+            saves.append(final)
+            return save(final=final)
+
+        runner._save_repex_state = counting_save
+        runner.run()
+
+        assert saves == expected
+        assert (Path(tmpdir) / "repex_state.pkl").exists()
 
 
 @pytest.mark.skipif(not has_cuda, reason="CUDA not available.")