From 62353639e3659d04cb8c3324605d2895bb8fe8d8 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 21 Sep 2026 14:33:26 +0100 Subject: [PATCH] Save the replica exchange state once at the end of a run. --- CHANGELOG.md | 1 + src/somd2/runner/_repex.py | 78 ++++++++++++++++++-------------------- tests/runner/test_repex.py | 48 +++++++++++++++++++++-- 3 files changed, 83 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 33a768b..34661f4 100644 --- a/CHANGELOG.md +++ b/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 -------------------------------------------------------------------------------------- diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 5dc4d66..f87a80c 100644 --- a/src/somd2/runner/_repex.py +++ b/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,25 +2158,7 @@ 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) @@ -2181,28 +2166,10 @@ def run(self): # 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 diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index d0a7bea..71660b8 100644 --- a/tests/runner/test_repex.py +++ b/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.")