Allow a dynamics object to propagate several independent trajectories

doc/source/changelog.rst

@@ -42,6 +42,11 @@ organisation on `GitHub <https://github.com/openbiosim/sire>`__.
   terms into two independently lambda-addressable OpenMM Forces, allowing them to be
   turned on according to different lambda schedule equations.
 
+* Added ``Dynamics.set_energy_trajectory()``, along with internal
+  ``Dynamics._get_clock()``/``_set_clock()``, which allow a single dynamics object to
+  propagate several independent trajectories by swapping the energy trajectory and
+  simulation clock between blocks.
+
 `2026.1.0 <https://github.com/openbiosim/sire/compare/2025.4.0...2026.1.0>`__ - June 2026
 -----------------------------------------------------------------------------------------
 

src/sire/mol/_dynamics.py

@@ -8,6 +8,26 @@ class DynamicsData:
     of molecule(s).
     """
 
+    # The attributes that make up the simulation clock. These track the
+    # progress of the simulation and schedule the saving of frames and
+    # energies. They are captured and restored together by _get_clock()
+    # and _set_clock(), which allows a single dynamics object to be
+    # re-used to propagate multiple independent trajectories.
+    _CLOCK_ATTRS = (
+        "_current_step",
+        "_current_time",
+        "_elapsed_time",
+        "_prev_step",
+        "_prev_current_time",
+        "_prev_elapsed_time",
+        "_next_save_frame",
+        "_next_save_energy",
+        "_prev_frame_frequency_steps",
+        "_prev_energy_frequency_steps",
+        "_prev_no_frame",
+        "_prev_no_energy",
+    )
+
     def __init__(self, mols=None, map=None, **kwargs):
         from ..base import create_map
 
@@ -182,6 +202,16 @@ def __init__(self, mols=None, map=None, **kwargs):
             self._is_running = False
             self._schedule_changed = False
 
+            # Save frequency counters. These are set on the first call to
+            # run(), but are initialised here so that the full clock state
+            # can be captured and restored before any dynamics has been run.
+            self._next_save_frame = None
+            self._next_save_energy = None
+            self._prev_frame_frequency_steps = None
+            self._prev_energy_frequency_steps = None
+            self._prev_no_frame = None
+            self._prev_no_energy = None
+
             # Initialise the GCMC sampler. This will be updated externally.
             # if the dynamics object is coupled to a sampler.
             self._gcmc_sampler = None
@@ -808,6 +838,29 @@ def platform(self):
         else:
             return self._omm_mols.getPlatform().getName()
 
+    def _get_clock(self):
+        if self.is_null():
+            return None
+        else:
+            return {attr: getattr(self, attr) for attr in self._CLOCK_ATTRS}
+
+    def _set_clock(self, clock):
+        if self.is_null():
+            return
+
+        from openmm.unit import picosecond
+
+        for attr, value in clock.items():
+            if attr not in self._CLOCK_ATTRS:
+                raise KeyError(f"'{attr}' is not a valid clock attribute")
+            setattr(self, attr, value)
+
+        # The OpenMM context has its own clock, which must be kept in sync
+        # with the elapsed time. _exit_dynamics_block() computes the time
+        # delta for a block as the difference between the two, so restoring
+        # one without the other would corrupt the recorded times.
+        self._omm_mols.setTime(self._elapsed_time.to("picosecond") * picosecond)
+
     def current_step(self):
         if self.is_null():
             return 0
@@ -883,6 +936,17 @@ def current_kinetic_energy(self):
     def energy_trajectory(self):
         return self._energy_trajectory.clone()
 
+    def set_energy_trajectory(self, energy_trajectory):
+        if self.is_null():
+            return
+
+        from ..legacy.Maths import EnergyTrajectory
+
+        if not isinstance(energy_trajectory, EnergyTrajectory):
+            raise TypeError("'energy_trajectory' must be of type 'EnergyTrajectory'")
+
+        self._energy_trajectory = energy_trajectory
+
     def _current_energy_array(self):
         try:
             import numpy as np
@@ -1290,10 +1354,13 @@ class NeedsMinimiseError(Exception):
 
         nsteps_before_run = self._current_step
 
-        # if this is the first call, then set the save frequencies
-        if nsteps_before_run == 0:
-            self._next_save_frame = frame_frequency_steps
-            self._next_save_energy = energy_frequency_steps
+        # if this is the first call, then set the save frequencies. The
+        # counters are also unset if a clock was restored from a dynamics
+        # object that had yet to be run, in which case schedule the first
+        # save relative to the restored step count.
+        if nsteps_before_run == 0 or self._next_save_frame is None:
+            self._next_save_frame = nsteps_before_run + frame_frequency_steps
+            self._next_save_energy = nsteps_before_run + energy_frequency_steps
             self._prev_frame_frequency_steps = frame_frequency_steps
             self._prev_energy_frequency_steps = energy_frequency_steps
             self._prev_no_frame = no_save_frame
@@ -2084,6 +2151,32 @@ def timestep(self):
         """
         return self._d.timestep()
 
+    def _get_clock(self):
+        """
+        Return the current state of the simulation clock as a dictionary.
+
+        This captures the completed time and step count, along with the
+        counters that schedule the saving of frames and energies. Passing
+        the result to _set_clock() rewinds (or advances) the simulation to
+        that point, which allows a single dynamics object to be re-used to
+        propagate several independent trajectories.
+        """
+        return self._d._get_clock()
+
+    def _set_clock(self, clock):
+        """
+        Restore the state of the simulation clock from a dictionary
+        returned by _get_clock(). This also updates the time held by the
+        OpenMM context, so that the two remain in sync.
+
+        Parameters
+        ----------
+
+        clock: dict
+            The clock state, as returned by _get_clock().
+        """
+        self._d._set_clock(clock)
+
     def current_step(self):
         """
         Return the current number of completed steps of dynamics
@@ -2241,6 +2334,23 @@ def energy_trajectory(self, to_pandas: bool = False, to_alchemlyb: bool = False)
         else:
             return t
 
+    def set_energy_trajectory(self, energy_trajectory):
+        """
+        Replace the energy trajectory that is accumulated during dynamics.
+
+        Subsequent energy saves are appended to 'energy_trajectory', and it
+        is the trajectory that is attached to the system by commit(). This
+        allows a single dynamics object to be re-used to propagate several
+        independent trajectories, each accumulating its own energies.
+
+        Parameters
+        ----------
+
+        energy_trajectory: :class: `EnergyTrajectory <sire.legacy.Maths.EnergyTrajectory>`
+            The energy trajectory to accumulate into.
+        """
+        self._d.set_energy_trajectory(energy_trajectory)
+
     def _current_energy_array(self):
         """
         Return the current energies as a numpy array, in the same order

tests/mol/test_dynamics.py

@@ -196,3 +196,108 @@ def test_crash_report(merged_ethane_methanol, openmm_platform):
     finally:
         # Change back to the old directory.
         os.chdir(old_dir)
+
+
+@pytest.mark.skipif(
+    "openmm" not in sr.convert.supported_formats(),
+    reason="openmm support is not available",
+)
+def test_clock_and_energy_trajectory_swap(ala_mols):
+    """
+    Test that a single dynamics object can propagate several independent
+    trajectories by swapping the clock and energy trajectory between blocks.
+
+    This underpins replica exchange runs that re-use a bounded number of
+    OpenMM contexts across a larger number of replicas.
+
+    The NVE ensemble is used on the Reference platform so that the integrator
+    is deterministic and has no random number stream. A thermostat would make
+    the comparison meaningless, since the cached run interleaves the replicas
+    through a single integrator and so consumes its RNG stream in a different
+    order to separate dynamics objects.
+    """
+
+    from sire.base import ProgressBar
+
+    ProgressBar.set_silent()
+
+    mols = ala_mols.clone()
+    mols.delete_all_frames()
+
+    num_cycles = 5
+    num_replicas = 2
+
+    # No temperature, so this is NVE and the integrator has no RNG.
+    kwargs = dict(platform="Reference", timestep="1 fs")
+    run_kwargs = dict(
+        energy_frequency="2 fs",
+        frame_frequency="2 fs",
+        lambda_windows=[0.0, 0.5, 1.0],
+    )
+
+    def potentials(traj):
+        return [round(float(v), 8) for v in traj.to_pandas()["potential"]]
+
+    # Build distinct starting states, so that the replicas follow genuinely
+    # different trajectories and the test is not vacuous.
+    seed = mols.dynamics(**kwargs)
+    start_states = []
+    for r in range(num_replicas):
+        if r > 0:
+            seed.run("10 fs", energy_frequency=0, frame_frequency=0)
+        start_states.append(
+            seed.context().getState(getPositions=True, getVelocities=True)
+        )
+
+    # Reference: one dynamics object per replica.
+    ref = []
+    for r in range(num_replicas):
+        d = mols.dynamics(**kwargs)
+        d.context().setState(start_states[r])
+        d._d._clear_state()
+        ref.append(d)
+
+    for i in range(num_cycles):
+        for d in ref:
+            d.run("2 fs", **run_kwargs)
+
+    ref_nrgs = [potentials(d.energy_trajectory()) for d in ref]
+    ref_steps = [d.current_step() for d in ref]
+
+    # Cached: a single dynamics object, with a clock and energy trajectory
+    # per replica.
+    slot = mols.dynamics(**kwargs)
+
+    # Seed the per-replica trajectories from the slot's own, so that the
+    # "ensemble" property is carried over.
+    trajs = [slot._d.energy_trajectory() for _ in range(num_replicas)]
+    assert all(len(t) == 0 for t in trajs)
+
+    clocks = [slot._get_clock() for _ in range(num_replicas)]
+    states = list(start_states)
+
+    for i in range(num_cycles):
+        for r in range(num_replicas):
+            slot.context().setState(states[r])
+            slot._set_clock(clocks[r])
+            slot.set_energy_trajectory(trajs[r])
+            slot._d._clear_state()
+
+            slot.run("2 fs", **run_kwargs)
+
+            clocks[r] = slot._get_clock()
+            states[r] = slot.context().getState(getPositions=True, getVelocities=True)
+            slot._d._sire_mols.delete_all_frames()
+
+    cache_nrgs = [potentials(t) for t in trajs]
+    cache_steps = [c["_current_step"] for c in clocks]
+
+    # The replicas must be distinct, otherwise nothing is being tested.
+    assert ref_nrgs[0] != ref_nrgs[1]
+
+    # Each replica must have accumulated its own energies, and the clock must
+    # have advanced as if it had a dynamics object to itself.
+    assert cache_steps == ref_steps
+    for r in range(num_replicas):
+        assert len(cache_nrgs[r]) == num_cycles
+        assert cache_nrgs[r] == ref_nrgs[r]