Dynamics Hooks — Stages & Usage#

This page covers hook behaviour specific to dynamics simulations. For the general hook protocol, context, and registry see Hooks — Core Framework.

DynamicsStage#

DynamicsStage enumerates the nine hook-firing points within a single dynamics step:

DynamicsStage hook firing points within a single step.#

Dynamics stages reference#
Stage	Value	When it fires
`BEFORE_STEP`	0	Very start of each step, before any operations.
`BEFORE_PRE_UPDATE`	1	Before the first integrator half-step (positions).
`AFTER_PRE_UPDATE`	2	After positions are updated, before the forward pass.
`BEFORE_COMPUTE`	3	Before the model forward pass.
`AFTER_COMPUTE`	4	After forces/energy are written to the batch.
`BEFORE_POST_UPDATE`	5	Before the second integrator half-step (velocities).
`AFTER_POST_UPDATE`	6	After velocities are updated.
`AFTER_STEP`	7	Very end of the step, after all operations.
`ON_CONVERGE`	8	Only when the convergence hook detects converged samples.

Built-in dynamics hooks#

The nvalchemi.dynamics.hooks package ships production-ready hooks in three categories. NeighborListHook, BiasedPotentialHook, and WrapPeriodicHook are general-purpose hooks documented in Hooks — Core Framework.

Observer hooks#

Observer hooks fire at AFTER_STEP and do not modify the batch.

LoggingHook#

LoggingHook writes per-step scalar observables to a backend. The default scalars are energy (per atom), fmax (maximum force component across all atoms), temperature (derived from kinetic energy when velocities are present), and converged_fraction (fraction of samples that have met the convergence criterion).

backend is a required argument that selects the output destination. It must be one of "csv", "tensorboard", or "custom":

"csv" — writes one row per step to log_path. Use when you need per-step data for post-run analysis in Python or a spreadsheet.
"tensorboard" — writes scalar events to log_path as a TensorBoard event file. Use when comparing scalar trends across experiments.
"custom" — routes each snapshot to a custom writer callable passed via the separate writer_fn parameter (signature fn(step_count, rows) -> None), such as a W&B or MLflow sink.

frequency throttles writes to every N steps. For long runs, frequency=10 or higher keeps output manageable without losing trends.

SnapshotHook#

SnapshotHook writes the full batch state — positions, velocities, forces, energy, cell, and atom types — to a DataSink at a specified frequency.

sink accepts one of three DataSink types:

GPUBuffer — stores batches in GPU memory. Fastest write path; capacity bounded by GPU memory.
HostMemory — stores in pinned CPU memory. Slightly slower; larger capacity and works without GPU.
ZarrData — streams to disk in Zarr format. Unbounded capacity; suitable for long trajectories and persistent storage.

After the run, call sink.read() to retrieve the accumulated trajectory as a Batch. Use this hook when you need full atomic-detail trajectories for analysis, visualization, or continuation from a specific frame.

ConvergedSnapshotHook#

ConvergedSnapshotHook writes only newly-converged samples at ON_CONVERGE — once per sample, exactly when convergence is detected — rather than periodically. The same DataSink types apply as for SnapshotHook.

This hook is designed for FusedStage pipelines where samples converge at different steps. A periodic snapshot would produce ragged data or miss samples; this hook captures each sample exactly once. Call sink.read() after the run to collect all converged structures.

EnergyDriftMonitorHook#

EnergyDriftMonitorHook tracks cumulative energy drift in NVE (constant-energy) simulations and takes a configurable action when drift exceeds a threshold.

Key arguments:

threshold — allowed drift, in the model’s energy output units.
metric — how drift is measured. "per_atom_per_step" normalises by system size and simulation length, making the threshold transferable across systems and time steps.
action — "raise" (default) halts the run; "warn" logs and continues. Use "warn" in production, "raise" during model validation.
frequency — check every N steps. Checking every step is accurate but adds overhead for large batches; frequency=100 is typical.

StageTimingHook and TorchProfilerHook are described in Hooks — Core Framework.

Post-compute hooks#

Post-compute hooks fire at AFTER_COMPUTE, after forces and energy are written to the batch but before the velocity update. They may modify the batch.

NaNDetectorHook#

NaNDetectorHook checks energy and forces for NaN or Inf values after the model forward pass. On detection it raises a RuntimeError that includes the affected graph indices and the current step count so the offending sample can be identified.

extra_keys extends the check to additional batch fields beyond energy and forces. For models that output stress tensors, pass extra_keys=["stress"].

When used with MaxForceClampHook, register the clamping hook first so the detector sees the clamped values and only catches what clamping did not prevent.

MaxForceClampHook#

MaxForceClampHook rescales per-atom forces whose magnitude exceeds max_force back to the threshold, preserving direction. Energy is not modified.

max_force is in the same units as the model’s force output (typically eV/Å). Clamping is applied in-place to any per-atom force whose magnitude exceeds the threshold. Frequent clamping during model development is a signal to identify problem configurations.

Clamping prevents numerical blow-up from large forces in high-energy or poorly-sampled configurations. It is a safety net, not a model fix: if clamping fires frequently, the model has accuracy problems for those structures.

Constraint hooks#

Constraint hooks enforce geometric constraints across integration steps. They fire at both BEFORE_PRE_UPDATE (to snapshot positions) and AFTER_POST_UPDATE (to restore them).

FreezeAtomsHook#

FreezeAtomsHook keeps selected atoms fixed: it snapshots their positions at BEFORE_PRE_UPDATE and restores them — with zeroed velocities — at AFTER_POST_UPDATE. The integrator runs normally and the positions are overwritten afterward, so no integrator modification is required.

categories is a string or list of strings matching atom type categories in the batch (for example, "substrate" or ["substrate", "boundary"]). Only atoms in the listed categories are frozen; all others evolve freely.

Use this hook for partial-system relaxations (freeze the substrate, relax the adsorbate), slab calculations (freeze bottom layers), or any configuration where part of the system must remain rigid.

Usage examples#

Logging to CSV every 100 steps#

from nvalchemi.dynamics.hooks import LoggingHook

hook = LoggingHook(frequency=100, backend="csv", log_path="md_log.csv")
dynamics = DemoDynamics(model=model, n_steps=10_000, dt=0.5, hooks=[hook])
dynamics.run(batch)

Recording trajectories to a data sink#

from nvalchemi.dynamics.hooks import SnapshotHook
from nvalchemi.dynamics import HostMemory

sink = HostMemory(capacity=10_000)
hook = SnapshotHook(sink=sink, frequency=10)
dynamics = DemoDynamics(model=model, n_steps=1_000, dt=0.5, hooks=[hook])
dynamics.run(batch)   # 100 snapshots
trajectory = sink.read()

Safety: NaN detection and force clamping#

Registration order determines execution order at the same stage. Clamp before checking so the detector sees the corrected forces:

from nvalchemi.dynamics.hooks import MaxForceClampHook, NaNDetectorHook

dynamics = DemoDynamics(
    model=model,
    dt=0.5,
    hooks=[
        MaxForceClampHook(max_force=50.0),
        NaNDetectorHook(extra_keys=["stress"]),
    ],
)

Hooks inside `FusedStage`#

When hooks are registered on sub-stage dynamics inside a FusedStage, their firing semantics differ slightly from standalone execution:

Fired on each sub-stage:

BEFORE_STEP, AFTER_COMPUTE, BEFORE_PRE_UPDATE, AFTER_POST_UPDATE, AFTER_STEP, ON_CONVERGE

Not fired on sub-stages (because the forward pass is shared):

BEFORE_COMPUTE, AFTER_PRE_UPDATE, BEFORE_POST_UPDATE

This means safety hooks (NaNDetectorHook, MaxForceClampHook) and observer hooks (LoggingHook, SnapshotHook) work as expected inside fused stages, since they fire at AFTER_COMPUTE or AFTER_STEP.

Hook ordering inside a fused step:

Hook ordering inside a single FusedStage.step().#

API reference#

`LoggingHook`	Log per-sample scalar observables from the simulation.
`SnapshotHook`	Save a snapshot of the active batch to a `DataSink` at a given frequency.
`ConvergedSnapshotHook`	Write only newly converged samples to a `DataSink`.
`EnergyDriftMonitorHook`	Track energy drift and warn or stop if it exceeds a threshold.
`NaNDetectorHook`	Detect NaN or Inf values in model outputs and raise immediately.
`MaxForceClampHook`	Clamp per-atom force vectors to a maximum magnitude.
`FreezeAtomsHook`	Freeze selected atoms during molecular dynamics simulation.

The general-purpose profiling hooks StageTimingHook and TorchProfilerHook also work with dynamics and are documented in Hooks — Core Framework.