Fix SOMD #124

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

@@ -475,7 +475,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 +495,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 +508,13 @@ 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
                             else:
-                                nrgs[str(lambda_value)] = null_energy * kcal_per_mol
+                                nrgs[key] = null_energy * kcal_per_mol
+                        else:
+                            nrgs[f"{sim_lambda_value:.5f}"] = (
+                                nrg_sim_lambda_value * kcal_per_mol
+                            )
 
                 self._omm_mols.set_lambda(
                     sim_lambda_value,
@@ -525,6 +531,9 @@ def _exit_dynamics_block(
                     self._current_time, nrgs, {"lambda": str(sim_lambda_value)}
                 )
 
+            # Store the current energies.
+            self._nrgs = nrgs
+
             # update the interpolation lambda value
             if self._is_interpolate:
                 if delta_lambda:
@@ -861,6 +870,18 @@ def current_kinetic_energy(self):
     def energy_trajectory(self):
         return self._energy_trajectory.clone()
 
+    def current_energies(self, sort: bool = False):
+        try:
+            if sort:
+                nrgs = self._nrgs.copy()
+                sorted_items = sorted(list(nrgs.items())[2:], key=lambda x: x[0])
+                nrgs = dict(list(nrgs.items())[:2] + sorted_items)
+                return nrgs
+            else:
+                return self._nrgs
+        except Exception:
+            return {}
+
     def step(self, num_steps: int = 1):
         """
         Just perform 'num_steps' steps of dynamics, without saving
@@ -2195,6 +2216,13 @@ def energy_trajectory(self, to_pandas: bool = False, to_alchemlyb: bool = False)
         else:
             return t
 
+    def current_energies(self, sort: bool = False):
+        """
+        Return a dictionary of the most recent energy trajectory entry.
+        If 'sort' is True, then the dictionary will be sorted by key.
+        """
+        return self._d.current_energies(sort=sort)
+
     def to_xml(self, f=None):
         """
         Save the current state of the dynamics to XML.