Merge pull request #403 from OpenBioSim/fix_somd_124

Fix SOMD #124

Author: lohedges

doc/source/changelog.rst

@@ -21,6 +21,8 @@ organisation on `GitHub <https://github.com/openbiosim/sire>`__.
   incorrectly for systems containing more than one molecule with CMAP terms (e.g.
   multi-chain glycoproteins).
 
+* Add convenience function to ``sire.mol.dynamics`` to get current energy trajectory records.
+
 `2025.4.0 <https://github.com/openbiosim/sire/compare/2025.3.0...2025.4.0>`__ - February 2026
 ---------------------------------------------------------------------------------------------
 

src/sire/mol/_dynamics.py

@@ -434,6 +434,7 @@ def _exit_dynamics_block(
         if save_energy:
             # should save energy here
             nrgs = {}
+            nrgs_array = []
 
             nrgs["kinetic"] = (
                 self._omm_state.getKineticEnergy().value_in_unit(
@@ -475,7 +476,8 @@ def _exit_dynamics_block(
                     nrg += self._pressure * volume
                 if excess_chemical_potential is not None:
                     nrg += excess_chemical_potential * num_waters
-                nrgs[str(sim_lambda_value)] = nrg * kcal_per_mol
+                # Store the potential energy for the current lambda value.
+                nrg_sim_lambda_value = nrg
 
                 if lambda_windows is not None:
                     # get the index of the simulation lambda value in the
@@ -494,6 +496,7 @@ def _exit_dynamics_block(
                                 not has_lambda_index
                                 or abs(lambda_index - i) <= num_energy_neighbours
                             ):
+                                key = f"{lambda_value:.5f}"
                                 self._omm_mols.set_lambda(
                                     lambda_value,
                                     rest2_scale=rest2_scale,
@@ -506,9 +509,16 @@ def _exit_dynamics_block(
                                     nrg += self._pressure * volume
                                 if excess_chemical_potential is not None:
                                     nrg += excess_chemical_potential * num_waters
-                                nrgs[str(lambda_value)] = nrg * kcal_per_mol
+                                nrgs[key] = nrg * kcal_per_mol
+                                nrgs_array.append(nrg)
                             else:
-                                nrgs[str(lambda_value)] = null_energy * kcal_per_mol
+                                nrgs_array.append(null_energy)
+                                nrgs[key] = null_energy * kcal_per_mol
+                        else:
+                            nrgs[f"{sim_lambda_value:.5f}"] = (
+                                nrg_sim_lambda_value * kcal_per_mol
+                            )
+                            nrgs_array.append(nrg_sim_lambda_value)
 
                 self._omm_mols.set_lambda(
                     sim_lambda_value,
@@ -525,6 +535,10 @@ def _exit_dynamics_block(
                     self._current_time, nrgs, {"lambda": str(sim_lambda_value)}
                 )
 
+            # Store the current energies.
+            self._nrgs = nrgs
+            self._nrgs_array = nrgs_array
+
             # update the interpolation lambda value
             if self._is_interpolate:
                 if delta_lambda:
@@ -861,6 +875,14 @@ def current_kinetic_energy(self):
     def energy_trajectory(self):
         return self._energy_trajectory.clone()
 
+    def _current_energy_array(self):
+        try:
+            import numpy as np
+
+            return np.array(self._nrgs_array)
+        except Exception:
+            return None
+
     def step(self, num_steps: int = 1):
         """
         Just perform 'num_steps' steps of dynamics, without saving
@@ -2195,6 +2217,13 @@ def energy_trajectory(self, to_pandas: bool = False, to_alchemlyb: bool = False)
         else:
             return t
 
+    def _current_energy_array(self):
+        """
+        Return the current energies as a numpy array, in the same order
+        as the energy trajectory columns.
+        """
+        return self._d._current_energy_array()
+
     def to_xml(self, f=None):
         """
         Save the current state of the dynamics to XML.