nvalchemi.dynamics.hooks.EnergyDriftMonitorHook

class nvalchemi.dynamics.hooks.EnergyDriftMonitorHook(threshold, metric='per_atom_per_step', action='warn', frequency=1, include_kinetic=True)

Track energy drift and warn or stop if it exceeds a threshold.

In a well-behaved NVE (microcanonical) simulation with a symplectic integrator, the total energy should be conserved to within numerical precision. Significant energy drift indicates problems with:

• The integration timestep (too large for the force magnitudes).

• The ML potential (non-smooth or discontinuous energy surface).

• Numerical precision (single vs. double precision accumulation).

• Force clamping or other hook-induced modifications breaking energy conservation.

This hook monitors the total energy (potential + kinetic) over the simulation and computes drift metrics. It supports two modes:

Absolute drift mode (metric="absolute")

Tracks |E(t) - E(0)|, the absolute deviation from the initial total energy. Suitable for NVE validation runs.

Per-atom-per-step drift mode (metric="per_atom_per_step")

Tracks |E(t) - E(0)| / (N_atoms * step_count), a normalized metric that allows comparison across systems of different size and simulation length. This is the standard metric reported in ML potential benchmarks.

When the drift exceeds threshold, the hook either emits a warning or raises a RuntimeError, controlled by the action parameter.

The hook records the reference energy on the first firing and computes drift on all subsequent firings. For NVT or NPT simulations, energy drift is expected (the thermostat/barostat injects or removes energy), so use this hook primarily for NVE validation.

Parameters:

• threshold (float) – Maximum acceptable drift before triggering the action. Units depend on metric: eV for "absolute", eV/atom/step for "per_atom_per_step".

• metric ({"absolute", "per_atom_per_step"}, optional) – Drift metric to use. Default "per_atom_per_step".

• action ({"warn", "raise"}, optional) – What to do when the threshold is exceeded. "warn" emits a loguru warning; "raise" raises a RuntimeError. Default "warn".

• frequency (int, optional) – Evaluate drift every frequency steps. Default 1.

• include_kinetic (bool, optional) – Whether to include kinetic energy in the total energy calculation. Set to False if only monitoring potential energy drift (e.g. for optimizers). Default True.

threshold

Drift threshold.

Type:

float

metric

Drift metric mode.

Type:

str

action

Threshold violation behavior.

Type:

str

include_kinetic

Whether kinetic energy is included.

Type:

bool

frequency

Evaluation frequency in steps.

Type:

int

stage

Fixed to AFTER_STEP.

Type:

HookStageEnum

Examples

NVE validation with strict drift tolerance:

>>> from nvalchemi.dynamics.hooks import EnergyDriftMonitorHook
>>> hook = EnergyDriftMonitorHook(
...     threshold=1e-5,
...     metric="per_atom_per_step",
...     action="raise",
...     frequency=100,
... )
>>> dynamics = DemoDynamics(model=model, n_steps=10_000, dt=0.5, hooks=[hook])
>>> dynamics.run(batch)

Soft monitoring during production:

>>> hook = EnergyDriftMonitorHook(
...     threshold=1e-3,
...     action="warn",
...     frequency=1000,
... )

Notes

• The reference energy is captured on the first hook firing (step 0 by default), not at construction time. This allows the hook to be registered before the batch is available.

• For batched simulations, drift is computed per graph and the maximum drift across all graphs is compared to the threshold.

__init__(threshold, metric='per_atom_per_step', action='warn', frequency=1, include_kinetic=True)

Parameters:

• threshold (float)

• metric (Literal['absolute', 'per_atom_per_step'])

• action (Literal['warn', 'raise'])

• frequency (int)

• include_kinetic (bool)

Return type:

None

Methods

__init__(threshold[, metric, action, ...])