From c125092e4fa2e9a20bb7eb22c69325578c92ffbe Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 12:38:57 +0100 Subject: [PATCH 1/7] Free committed repex systems before deleting their frames. --- src/somd2/runner/_repex.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 0de187e..9e316a9 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2570,6 +2570,12 @@ def _checkpoint( if error is not None: return index, error + # Free the committed system with the GIL released. This comes + # before the frames are deleted below, which would otherwise copy + # every molecule that is still shared with it. + system_time = system.time() + system._free() + # Delete the trajectory frames from the Sire system within the # dynamics object. A replica committed earlier, with shared # contexts, had them cleared then, so its slot isn't touched here. @@ -2602,7 +2608,7 @@ def _checkpoint( ) _logger.info( f"Current number of waters in GCMC volume at {_lam_sym} = {lam:.5f} " - f"(time = {system.time().to('ns'):.4f} ns) is {num_waters}{acc_str}" + f"(time = {system_time.to('ns'):.4f} ns) is {num_waters}{acc_str}" ) # Log terminal flip acceptance rate for this replica. From bdbc533e50d5f75063d8a68d630b7e15fe9a802f Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 12:39:09 +0100 Subject: [PATCH 2/7] Free the committed system before deleting frames in the regular runner. --- src/somd2/runner/_runner.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 594a6cc..3ab6f89 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -1181,10 +1181,6 @@ def _run_checkpoint_blocks(self, window, time): self._write_checkpoint(window, system, block, speed, is_final_block) - # Delete all trajectory frames from the Sire system within the - # dynamics object. - dynamics._d._sire_mols.delete_all_frames() - _logger.info( f"Finished block {block + 1} of {total_blocks} " f"for {_lam_sym} = {lambda_value:.5f}" @@ -1192,6 +1188,12 @@ def _run_checkpoint_blocks(self, window, time): self._log_samplers(window, system) + # Free the committed system before deleting the trajectory + # frames from the dynamics object, which would otherwise copy + # every molecule that is still shared with it. + del system + dynamics._d._sire_mols.delete_all_frames() + if is_final_block: _logger.success( f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1" From 2974fb28e9af5eb5ddf957e2fa06a8fd1d3d3eb2 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 12:47:08 +0100 Subject: [PATCH 3/7] Shorten checkpoint comments. --- src/somd2/runner/_repex.py | 5 ++--- src/somd2/runner/_runner.py | 5 ++--- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 9e316a9..cc17853 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2570,9 +2570,8 @@ def _checkpoint( if error is not None: return index, error - # Free the committed system with the GIL released. This comes - # before the frames are deleted below, which would otherwise copy - # every molecule that is still shared with it. + # Free the committed system with the GIL released, before deleting + # frames, which would otherwise copy every molecule it shares. system_time = system.time() system._free() diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index 3ab6f89..7b34b3e 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -1188,9 +1188,8 @@ def _run_checkpoint_blocks(self, window, time): self._log_samplers(window, system) - # Free the committed system before deleting the trajectory - # frames from the dynamics object, which would otherwise copy - # every molecule that is still shared with it. + # Free the committed system before deleting the frames, which + # would otherwise copy every molecule it shares. del system dynamics._d._sire_mols.delete_all_frames() From c786edddf83da75f14c99fb7116415c7c8fe88d6 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 13:59:17 +0100 Subject: [PATCH 4/7] Join trajectory chunks with DCD.concatenate instead of reloading the topology. --- src/somd2/runner/_base.py | 40 +++++++++++++++++++++++++----------- src/somd2/runner/_repex.py | 30 +++++++++------------------ tests/runner/test_restart.py | 5 +++++ 3 files changed, 43 insertions(+), 32 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index ee39ba0..eca5308 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2734,9 +2734,6 @@ def _checkpoint( format=["DCD"], ) - # Create the final topology file name. - topology0 = self._filenames["topology0"] - # Create the final trajectory file name. traj_filename = self._filenames[index]["trajectory"] @@ -2756,15 +2753,7 @@ def _checkpoint( # Make sure there are trajectory chunks to process. if len(traj_chunks) > 0: - # Load the topology and chunked trajectory files. - mols = _sr.load([topology0] + traj_chunks) - - # Save the final trajectory to a single file. - _sr.save(mols.trajectory(), traj_filename, format=["DCD"]) - - # Now remove the chunked trajectory files. - for chunk in traj_chunks: - _Path(chunk).unlink() + self._join_trajectory_chunks(traj_chunks, traj_filename) # Write the checkpoint system to file. self._write_checkpoint_system( @@ -2831,6 +2820,33 @@ def _checkpoint( return index, None + def _join_trajectory_chunks(self, chunks, filename): + """ + Join the trajectory chunks, in order, into the final trajectory, then + remove them. Frames are 'frame_frequency' apart. + + Parameters + ---------- + + chunks: list of str + The DCD files to join, in order. + + filename: str + The final trajectory file. + """ + from os import replace as _replace + + from sire.legacy.IO import DCD as _DCD + + # Write to a temporary file, so a failure can't leave a partial + # trajectory for a restart to build on. + temp_filename = f"{filename}.tmp" + _DCD.concatenate(chunks, temp_filename, self._config.frame_frequency) + _replace(temp_filename, filename) + + for chunk in chunks: + _Path(chunk).unlink() + def _write_checkpoint_system(self, system, index, context=None, gcmc_sampler=None): """ Write the system state to the checkpoint file. diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index cc17853..3fb1ac2 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -21,7 +21,6 @@ __all__ = ["RepexRunner"] -from contextlib import nullcontext as _nullcontext from filelock import FileLock as _FileLock from numba import njit as _njit from shutil import copyfile as _copyfile @@ -482,8 +481,7 @@ def __init__(self, system, config): from threading import Lock - # Create a lock to guard the committed systems, which also makes the - # replicas take turns at the final checkpoint. + # Create a lock to guard the committed systems. self._lock = Lock() # Systems committed while a replica was resident in its slot, awaiting @@ -2555,17 +2553,14 @@ def _checkpoint( system, speed = committed # Call the base class checkpoint method to save the system state. - # The final block reassembles each trajectory, which loads the - # topology, so replicas take turns to bound the memory used. - with self._lock if is_final_block else _nullcontext(): - index, error = super()._checkpoint( - system, - index, - block, - speed, - is_final_block=is_final_block, - num_energy_rows=num_energy_rows, - ) + index, error = super()._checkpoint( + system, + index, + block, + speed, + is_final_block=is_final_block, + num_energy_rows=num_energy_rows, + ) if error is not None: return index, error @@ -2659,12 +2654,7 @@ def _consolidate_trajectories(self): if not traj_chunks: continue - topology0 = self._filenames["topology0"] - mols = _sr.load([topology0] + traj_chunks) - _sr.save(mols.trajectory(), traj_filename, format=["DCD"]) - - for chunk in traj_chunks: - _Path_local(chunk).unlink() + self._join_trajectory_chunks(traj_chunks, traj_filename) def _merge_gcmc_stats(self): """ diff --git a/tests/runner/test_restart.py b/tests/runner/test_restart.py index fda3fcf..9ad02c7 100644 --- a/tests/runner/test_restart.py +++ b/tests/runner/test_restart.py @@ -96,6 +96,11 @@ def test_restart(mols, request): # Check that the trajectory is twice as long as the first. assert traj_2.num_frames() == 2 * traj_1.num_frames() + # Frames are a frame_frequency apart, across the restart. + for i in range(traj_2.num_frames()): + time = traj_2.trajectory()[i].current().time() + assert time.to(sr.units.femtosecond) == pytest.approx(4 * i) + config_difftimestep = config_new.copy() config_difftimestep["runtime"] = "36fs" config_difftimestep["timestep"] = "2fs" From 472128b95729c450b6158cd4880f2e7ef6a3f5d5 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 14:11:55 +0100 Subject: [PATCH 5/7] Append trajectory chunks in place and time frames from the first interval. --- src/somd2/runner/_base.py | 42 ++++++++++++++++++++++++------------ src/somd2/runner/_repex.py | 12 +---------- tests/runner/test_restart.py | 5 +++-- 3 files changed, 32 insertions(+), 27 deletions(-) diff --git a/src/somd2/runner/_base.py b/src/somd2/runner/_base.py index eca5308..98f985a 100644 --- a/src/somd2/runner/_base.py +++ b/src/somd2/runner/_base.py @@ -2742,18 +2742,12 @@ def _checkpoint( _glob(f"{self._filenames[index]['trajectory_chunk']}*") ) - # If this is a restart, then we need to check for an existing - # trajectory file with the same name. If it exists and is non-empty, - # then copy it to a backup file and prepend it to the list of chunks. - if self._config.restart: - path = _Path(traj_filename) - if path.exists() and path.stat().st_size > 0: - _copyfile(traj_filename, f"{traj_filename}.prev") - traj_chunks = [f"{traj_filename}.prev"] + traj_chunks - - # Make sure there are trajectory chunks to process. + # On a restart, the chunks are added to the end of any + # existing final trajectory. if len(traj_chunks) > 0: - self._join_trajectory_chunks(traj_chunks, traj_filename) + self._join_trajectory_chunks( + traj_chunks, traj_filename, append=self._config.restart + ) # Write the checkpoint system to file. self._write_checkpoint_system( @@ -2820,10 +2814,11 @@ def _checkpoint( return index, None - def _join_trajectory_chunks(self, chunks, filename): + def _join_trajectory_chunks(self, chunks, filename, append=False): """ Join the trajectory chunks, in order, into the final trajectory, then - remove them. Frames are 'frame_frequency' apart. + remove them. The first frame is saved after one 'frame_frequency', and + the rest follow at the same interval. Parameters ---------- @@ -2833,15 +2828,34 @@ def _join_trajectory_chunks(self, chunks, filename): filename: str The final trajectory file. + + append: bool + Whether to add the chunks to the end of an existing final + trajectory, e.g. on a restart. """ from os import replace as _replace from sire.legacy.IO import DCD as _DCD + path = _Path(filename) + if append and path.exists() and path.stat().st_size > 0: + inputs = [filename] + chunks + else: + inputs = chunks + # Write to a temporary file, so a failure can't leave a partial # trajectory for a restart to build on. temp_filename = f"{filename}.tmp" - _DCD.concatenate(chunks, temp_filename, self._config.frame_frequency) + try: + _DCD.concatenate( + inputs, + temp_filename, + self._config.frame_frequency, + self._config.frame_frequency, + ) + except Exception: + _Path(temp_filename).unlink(missing_ok=True) + raise _replace(temp_filename, filename) for chunk in chunks: diff --git a/src/somd2/runner/_repex.py b/src/somd2/runner/_repex.py index 3fb1ac2..36d9ee1 100644 --- a/src/somd2/runner/_repex.py +++ b/src/somd2/runner/_repex.py @@ -2632,8 +2632,6 @@ def _consolidate_trajectories(self): skipped automatically. """ from glob import glob as _glob_local - from pathlib import Path as _Path_local - from shutil import copyfile as _copyfile_local if not self._config.save_trajectories: return @@ -2643,18 +2641,10 @@ def _consolidate_trajectories(self): chunk_pattern = f"{self._filenames[i]['trajectory_chunk']}*" traj_chunks = sorted(_glob_local(chunk_pattern)) - # On a restart, prepend an existing final DCD as .prev so frames - # from a previous (possibly partial) consolidation are preserved. - path = _Path_local(traj_filename) - if path.exists() and path.stat().st_size > 0: - prev = f"{traj_filename}.prev" - _copyfile_local(traj_filename, prev) - traj_chunks = [prev] + traj_chunks - if not traj_chunks: continue - self._join_trajectory_chunks(traj_chunks, traj_filename) + self._join_trajectory_chunks(traj_chunks, traj_filename, append=True) def _merge_gcmc_stats(self): """ diff --git a/tests/runner/test_restart.py b/tests/runner/test_restart.py index 9ad02c7..853114b 100644 --- a/tests/runner/test_restart.py +++ b/tests/runner/test_restart.py @@ -96,10 +96,11 @@ def test_restart(mols, request): # Check that the trajectory is twice as long as the first. assert traj_2.num_frames() == 2 * traj_1.num_frames() - # Frames are a frame_frequency apart, across the restart. + # Frames are a frame_frequency apart, from the first frame saved after + # one frame_frequency, across the restart. for i in range(traj_2.num_frames()): time = traj_2.trajectory()[i].current().time() - assert time.to(sr.units.femtosecond) == pytest.approx(4 * i) + assert time.to(sr.units.femtosecond) == pytest.approx(4 * (i + 1)) config_difftimestep = config_new.copy() config_difftimestep["runtime"] = "36fs" From 49e74f9d0f1bff1a03672cec867366dab167b2df Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 16:25:23 +0100 Subject: [PATCH 6/7] Use plain arrays rather than numpy matrices for swap counts. --- src/somd2/runner/_dynamics_cache.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/somd2/runner/_dynamics_cache.py b/src/somd2/runner/_dynamics_cache.py index 787e9e3..0b956f5 100644 --- a/src/somd2/runner/_dynamics_cache.py +++ b/src/somd2/runner/_dynamics_cache.py @@ -156,9 +156,9 @@ def __init__( # be pushed into the context before the next block. self._state_moved = [False] * num_replicas self._terminal_flip_stats = [[0, 0] for _ in range(num_replicas)] - self._num_proposed = _np.matrix(_np.zeros((num_replicas, num_replicas))) - self._num_accepted = _np.matrix(_np.zeros((num_replicas, num_replicas))) - self._num_swaps = _np.matrix(_np.zeros((num_replicas, num_replicas))) + self._num_proposed = _np.zeros((num_replicas, num_replicas)) + self._num_accepted = _np.zeros((num_replicas, num_replicas)) + self._num_swaps = _np.zeros((num_replicas, num_replicas)) # Build the slot layout and the per-slot attributes. self._build_slot_layout() From a7c54db0695036c0ae6de6daddb8e2a8f70f01af Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 20 Oct 2026 16:29:45 +0100 Subject: [PATCH 7/7] Update CHANGELOG. --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 83783cf..bfb43c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,8 @@ Changelog * Require the `runtime`, `checkpoint_frequency`, `energy_frequency`, `frame_frequency`, `gcmc_frequency` and `terminal_flip_frequency` to be multiples of the `timestep` for the regular runner, which runs in whole integration steps. GCMC and terminal flip moves now keep their frequency across checkpoints [#248](https://github.com/OpenBioSim/somd2/pull/248). * Speed up checkpointing, which could take minutes for large systems. Sire now uses every CPU thread while checkpoint files are written, unless `max_sire_threads` is set, and replica exchange checkpoints replicas in parallel, writing the shared GCMC topologies once rather than once per replica. The regular runner's worker processes now also apply the Sire thread limit, which previously only took effect in the parent process [#251](https://github.com/OpenBioSim/somd2/pull/251). * Reduce the setup time of replica exchange simulations that re-use OpenMM contexts across lambda values. Minimisation and equilibration now create their dynamics objects once per context, rather than once per replica [#255](https://github.com/OpenBioSim/somd2/pull/255). +* Further speed up checkpointing, and join the final trajectory without reloading the topology [#257](https://github.com/OpenBioSim/somd2/pull/257). +* Record frame times in final trajectories, rather than zero for every frame [#257](https://github.com/OpenBioSim/somd2/pull/257). [2026.2.0](https://github.com/openbiosim/somd2/compare/2026.1.0...2026.2.0) - Sep 2026 --------------------------------------------------------------------------------------