[BUG] Frames and energies are always saved after dynamics.run()

[lohedges]

I am using Sire dynamics for replica exchange. In doing so, I need to run short cycles at the energy_frequency then swap replicas. However, if I set a frame_frequency to larger than my simulation run time I find that a frame is still saved after dynamics.run() completes. In fact, it seems that energies and frames are always saved when it exits. This is massively impacting performance, since I am effectively saving frames at the (very short) energy_frequency.

As a solution, the code should be updated so that calls to .run() are not independent of each other, i.e. it remembers how many steps were performed before and calculates the time till saving the energy or a frame accordingly.

[lohedges]

Given the way that things are currently implemented, this is really tricky to fix. For example, the total number of steps is stored as a class attribute, but this is only updated if the state is saved when the dynamics block exits, which only happens when coordinates are saved, or when .commit() is called. This isn't what I want. I need to be able to run multiple blocks in sequence, only saving frames when a multiple of the frame frequency is passed. I can hack things, but I think it will cause issues with the way that the OpenMM trajectory state is synchronised with the Sire system, e.g. for crash recovery, etc.

[lohedges]

This diff seems to fix things, although I'll need to make sure that there are no unintended consequences.

diff --git a/src/sire/mol/_dynamics.py b/src/sire/mol/_dynamics.py
index 40c8b125..f174a521 100644
--- a/src/sire/mol/_dynamics.py
+++ b/src/sire/mol/_dynamics.py
@@ -209,6 +209,8 @@ class DynamicsData:
         if self.is_null():
             return
 
+        self._current_step = nsteps_completed
+
         if not state_has_cv[0]:
             # there is no information to update
             return
@@ -247,8 +249,6 @@ class DynamicsData:
                 map=self._map,
             )
 
-        self._current_step = nsteps_completed
-
         self._sire_mols.update(mols_to_update.molecules())
 
         if self._ffinfo.space().is_periodic():
@@ -924,50 +924,6 @@ class DynamicsData:
                     }
                 )
 
-        def get_steps_till_save(completed: int, total: int):
-            """Internal function to calculate the number of steps
-            to run before the next save. This returns a tuple
-            of number of steps, and then if a frame should be
-            saved and if the energy should be saved
-            """
-            if completed < 0:
-                completed = 0
-
-            if completed >= total:
-                return (
-                    0,
-                    frame_frequency_steps > 0,
-                    energy_frequency_steps > 0,
-                )
-
-            elif frame_frequency_steps <= 0 and energy_frequency_steps <= 0:
-                return (total, False, False)
-
-            n_to_end = total - completed
-
-            if frame_frequency_steps > 0:
-                n_to_frame = min(
-                    frame_frequency_steps - (completed % frame_frequency_steps),
-                    n_to_end,
-                )
-            else:
-                n_to_frame = total - completed
-
-            if energy_frequency_steps > 0:
-                n_to_energy = min(
-                    energy_frequency_steps - (completed % energy_frequency_steps),
-                    n_to_end,
-                )
-            else:
-                n_to_energy = total - completed
-
-            if n_to_frame == n_to_energy:
-                return (n_to_frame, True, True)
-            elif n_to_frame < n_to_energy:
-                return (n_to_frame, True, False)
-            else:
-                return (n_to_energy, False, True)
-
         block_size = 50
 
         state = None
@@ -978,6 +934,9 @@ class DynamicsData:
             pass
 
         nsteps_before_run = self._current_step
+        if nsteps_before_run == 0:
+            self._next_save_frame = frame_frequency_steps
+            self._next_save_energy = energy_frequency_steps
 
         from ..base import ProgressBar
         from ..units import second
@@ -992,13 +951,27 @@ class DynamicsData:
 
                 with ThreadPoolExecutor() as pool:
                     while completed < steps_to_run:
-                        (
-                            nrun_till_save,
-                            save_frame,
-                            save_energy,
-                        ) = get_steps_till_save(completed, steps_to_run)
+                        steps_till_frame = self._next_save_frame - (
+                            completed + nsteps_before_run
+                        )
+                        if steps_till_frame <= 0 or steps_till_frame == steps_to_run:
+                            save_frame = True
+                            self._next_save_frame += frame_frequency_steps
+                        else:
+                            save_frame = False
+
+                        steps_till_energy = self._next_save_energy - (
+                            completed + nsteps_before_run
+                        )
+                        if steps_till_energy <= 0 or steps_till_energy == steps_to_run:
+                            save_energy = True
+                            self._next_save_energy += energy_frequency_steps
+                        else:
+                            save_energy = False
+
+                        nrun_till_save = min(steps_till_frame, steps_till_energy)
 
-                        assert nrun_till_save > 0
+                        assert nrun_till_save >= 0
 
                         self._enter_dynamics_block()