Add option to always save the energy when .run() exits.

src/sire/mol/_dynamics.py

@@ -891,6 +891,7 @@ def run(
         rest2_scale_factors=None,
         save_velocities: bool = None,
         save_frame_on_exit: bool = False,
+        save_energy_on_exit: bool = False,
         auto_fix_minimise: bool = True,
     ):
         if self.is_null():
@@ -905,6 +906,7 @@ def run(
             "rest2_scale_factors": rest2_scale_factors,
             "save_velocities": save_velocities,
             "save_frame_on_exit": save_frame_on_exit,
+            "save_energy_on_exit": save_energy_on_exit,
             "auto_fix_minimise": auto_fix_minimise,
         }
 
@@ -1093,6 +1095,14 @@ class NeedsMinimiseError(Exception):
                         ):
                             save_frame = True
 
+                        # save the last energy if we're about to exit and the user
+                        # has requested it
+                        if (
+                            save_energy_on_exit
+                            and completed + block_size >= steps_to_run
+                        ):
+                            save_energy = True
+
                         self._enter_dynamics_block()
 
                         # process the last block in the foreground
@@ -1447,6 +1457,7 @@ def run(
         rest2_scale_factors=None,
         save_velocities: bool = None,
         save_frame_on_exit: bool = False,
+        save_energy_on_exit: bool = False,
         auto_fix_minimise: bool = True,
     ):
         """
@@ -1517,6 +1528,10 @@ def run(
             Whether to save a trajectory frame on exit, regardless of
             whether the frame frequency has been reached.
 
+        save_energy_on_exit: bool
+            Whether to save the energy on exit, regardless of whether
+            the energy frequency has been reached.
+
         auto_fix_minimise: bool
             Whether or not to automatically run minimisation if the
             trajectory exits with an error in the first few steps.
@@ -1540,6 +1555,7 @@ def run(
                 rest2_scale_factors=rest2_scale_factors,
                 save_velocities=save_velocities,
                 save_frame_on_exit=save_frame_on_exit,
+                save_energy_on_exit=save_energy_on_exit,
                 auto_fix_minimise=auto_fix_minimise,
             )
 

tests/mol/test_dynamics.py

@@ -123,9 +123,22 @@ def test_sample_frequency(ala_mols):
     # Recreate the dynamics object.
     d = mols.dynamics(platform="Reference", timestep="1 fs")
 
-    # Run 10 cycles of dynamics saving frames every 5 fs.
+    # Run 10 cycles of dynamics, saving energies frames on exit.
     for i in range(10):
-        d.run("1 fs", frame_frequency="5 fs", save_frame_on_exit=True)
+        d.run(
+            "1 fs",
+            energy_frequency="2 fs",
+            frame_frequency="5 fs",
+            lambda_windows=lambdas,
+            save_frame_on_exit=True,
+            save_energy_on_exit=True,
+        )
+
+    # Get the energy trajectory.
+    nrg_traj = d.energy_trajectory()
+
+    # Make sure the trajectory has 10 frames.
+    assert len(nrg_traj) == 10
 
     # Get the updated system.
     new_mols = d.commit()