Guard energy rescale against unphysical steps (silent NaN)

docs/energy.md

@@ -42,3 +42,48 @@ atoms.calc = calculator
 In general, the energy models are slower and have a larger memory footprint compared to
 the FlashMD models. As summarized above, you should use `do_gradients_with_energy=False`
 to save computation and memory when you do not need forces.
+
+Monitoring the rescaling factor
+--------------------------------
+
+Every time ``rescale_energy=True`` triggers a rescale, the momenta are multiplied by a
+factor ``alpha = sqrt(1 - (E_new - E_old) / E_kin)``. If a step increases the total energy
+by more than the post-step kinetic energy can absorb, both the ASE and i-PI
+integrators raise a ``RuntimeError``. This is a sign that the step was
+unphysical (e.g. atomic overlap or a model extrapolation error) -- consider using a
+smaller time step.
+
+You can also monitor ``alpha`` directly, to catch large corrections before they become
+outright failures.
+
+**ASE**: the last computed value is available as ``dyn.alpha`` (``None`` until the first
+rescaled step). Attach an observer to log it during a run:
+
+```
+dyn.attach(lambda: print(dyn.alpha), interval=1)
+```
+
+**i-PI**: the value is written each step to ``motion.flashmd_alpha`` (``nan`` on any step
+where rescaling did not run), which you can expose as a genuine column in the ``.out``
+file, by registering it as a custom property when you
+build the ``InteractiveSimulation``:
+
+```
+sim = InteractiveSimulation(
+    input_xml,
+    custom_properties={
+        "flashmd_alpha": {
+            "func": lambda self: getattr(self.motion, "flashmd_alpha", float("nan")),
+            "dimension": "undefined",
+            "help": "FlashMD momentum rescale factor (energy conservation).",
+        }
+    },
+)
+```
+
+and adding ``flashmd_alpha`` to the ``<properties>`` list in your xml file. The custom
+property must be registered this way, in the Python script that builds the
+``InteractiveSimulation``, before you can reference it in the xml: i-PI validates the
+``<properties>`` list against its property registry as soon as the simulation object is
+built, so adding ``flashmd_alpha`` to the xml alone, without this registration, raises
+``KeyError: flashmd_alpha is not a recognized property``.

src/flashmd/ase/velocity_verlet.py

@@ -60,6 +60,11 @@ def __init__(
             self.stepper = flashmd_stepper
         self.rescale_energy = rescale_energy
         self.random_rotation = random_rotation
+        # last momentum rescaling factor applied to conserve energy; None until
+        # the first step, and only set when rescale_energy=True. Attach an
+        # observer (e.g. dyn.attach(lambda: print(dyn.alpha), interval=1)) to
+        # monitor it during a run.
+        self.alpha = None
 
     def step(self):
         if self.rescale_energy:
@@ -108,7 +113,17 @@ def step(self):
         if self.rescale_energy:
             new_energy = self.atoms.get_total_energy()
             old_kinetic_energy = self.atoms.get_kinetic_energy()
-            alpha = np.sqrt(1.0 - (new_energy - old_energy) / old_kinetic_energy)
+            discriminant = 1.0 - (new_energy - old_energy) / old_kinetic_energy
+            if discriminant < 0.0:
+                raise RuntimeError(
+                    "Energy rescale failed: the step increased the total "
+                    f"energy by {new_energy - old_energy:.6g}, which exceeds "
+                    f"the post-step kinetic energy ({old_kinetic_energy:.6g}). "
+                    "Try a smaller timestep, or disable rescale_energy to "
+                    "inspect the trajectory."
+                )
+            alpha = np.sqrt(discriminant)
+            self.alpha = alpha
             self.atoms.set_momenta(alpha * self.atoms.get_momenta())
 
     def irun(self, steps=50):

src/flashmd/ipi.py

@@ -30,6 +30,8 @@ def get_standard_vv_step(
     """
 
     def vv_step(motion):
+        motion.flashmd_alpha = float("nan")
+
         if random_rotation:
             raise NotImplementedError(
                 "Random rotation is not implemented in the standard VV stepper."
@@ -50,7 +52,17 @@ def vv_step(motion):
             info("@flashmd: Energy rescale", verbosity.debug)
             new_energy = sim.properties("conserved")
             kinetic_energy = sim.properties("kinetic_md")
-            alpha = np.sqrt(1.0 - (new_energy - old_energy) / kinetic_energy)
+            discriminant = 1.0 - (new_energy - old_energy) / kinetic_energy
+            if discriminant < 0.0:
+                raise RuntimeError(
+                    "Energy rescale failed: the step increased the total "
+                    f"energy by {new_energy - old_energy:.6g}, which exceeds "
+                    f"the post-step kinetic energy ({kinetic_energy:.6g}). "
+                    "Try a smaller timestep, or disable rescale_energy to "
+                    "inspect the trajectory."
+                )
+            alpha = np.sqrt(discriminant)
+            motion.flashmd_alpha = alpha
             motion.beads.p[:] = alpha * dstrip(motion.beads.p)
 
     return vv_step
@@ -90,6 +102,8 @@ def get_flashmd_vv_step(
         stepper = flashmd_stepper
 
     def flashmd_vv(motion):
+        motion.flashmd_alpha = float("nan")
+
         info("@flashmd: Starting VV", verbosity.debug)
         if rescale_energy:
             info("@flashmd: Old energy", verbosity.debug)
@@ -128,7 +142,17 @@ def flashmd_vv(motion):
             info("@flashmd: Energy rescale", verbosity.debug)
             new_energy = sim.properties("conserved")
             kinetic_energy = sim.properties("kinetic_md")
-            alpha = np.sqrt(1.0 - (new_energy - old_energy) / kinetic_energy)
+            discriminant = 1.0 - (new_energy - old_energy) / kinetic_energy
+            if discriminant < 0.0:
+                raise RuntimeError(
+                    "Energy rescale failed: the step increased the total "
+                    f"energy by {new_energy - old_energy:.6g}, which exceeds "
+                    f"the post-step kinetic energy ({kinetic_energy:.6g}). "
+                    "Try a smaller timestep, or disable rescale_energy to "
+                    "inspect the trajectory."
+                )
+            alpha = np.sqrt(discriminant)
+            motion.flashmd_alpha = alpha
             motion.beads.p[:] = alpha * dstrip(motion.beads.p)
             motion.integrator.pconstraints()  # just to be sure