From a9a0f94e05c4b79339b5376644421b03fe459bfe Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Mon, 27 Jul 2026 12:58:07 +0100 Subject: [PATCH] Allow a dynamics object to propagate several independent trajectories. --- doc/source/changelog.rst | 5 ++ src/sire/mol/_dynamics.py | 118 +++++++++++++++++++++++++++++++++++-- tests/mol/test_dynamics.py | 105 +++++++++++++++++++++++++++++++++ 3 files changed, 224 insertions(+), 4 deletions(-) diff --git a/doc/source/changelog.rst b/doc/source/changelog.rst index 30587b2e8..7365b4623 100644 --- a/doc/source/changelog.rst +++ b/doc/source/changelog.rst @@ -42,6 +42,11 @@ organisation on `GitHub `__. 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 `__ - June 2026 ----------------------------------------------------------------------------------------- diff --git a/src/sire/mol/_dynamics.py b/src/sire/mol/_dynamics.py index 65feefb4a..13ea7e2fd 100644 --- a/src/sire/mol/_dynamics.py +++ b/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 ` + 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 diff --git a/tests/mol/test_dynamics.py b/tests/mol/test_dynamics.py index 9d6c43b48..acda5d55d 100644 --- a/tests/mol/test_dynamics.py +++ b/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]