Optimise checkpoints

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

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"]
 
@@ -2745,26 +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:
-                        # 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, append=self._config.restart
+                        )
 
                 # Write the checkpoint system to file.
                 self._write_checkpoint_system(
@@ -2831,6 +2814,53 @@ def _checkpoint(
 
         return index, None
 
+    def _join_trajectory_chunks(self, chunks, filename, append=False):
+        """
+        Join the trajectory chunks, in order, into the final trajectory, then
+        remove them. The first frame is saved after one 'frame_frequency', and
+        the rest follow at the same interval.
+
+        Parameters
+        ----------
+
+        chunks: list of str
+            The DCD files to join, in order.
+
+        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"
+        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:
+            _Path(chunk).unlink()
+
     def _write_checkpoint_system(self, system, index, context=None, gcmc_sampler=None):
         """
         Write the system state to the checkpoint file.

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

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,21 +2553,23 @@ 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
 
+            # 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()
+
             # 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 +2602,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.
@@ -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,23 +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
 
-            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, append=True)
 
     def _merge_gcmc_stats(self):
         """

src/somd2/runner/_runner.py

@@ -1181,17 +1181,18 @@ 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}"
                 )
 
                 self._log_samplers(window, system)
 
+                # 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()
+
                 if is_final_block:
                     _logger.success(
                         f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1"

tests/runner/test_restart.py

@@ -96,6 +96,12 @@ 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, 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 + 1))
+
         config_difftimestep = config_new.copy()
         config_difftimestep["runtime"] = "36fs"
         config_difftimestep["timestep"] = "2fs"