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.

See also

• User guide: Hooks — conceptual overview, writing custom hooks, and composing hook pipelines.

• Core framework: Hooks — Core Framework — the Hook protocol, HookContext, and HookRegistryMixin.

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 organized into four categories. General-purpose hooks (NeighborListHook, BiasedPotentialHook, WrapPeriodicHook) are documented in Hooks — Core Framework.

Observer hooks (read-only, fire at AFTER_STEP)

These hooks do not modify the batch — they record, log, or monitor simulation state.

Hook	Purpose
LoggingHook	Log scalar observables (energy, fmax, temperature) to loguru, CSV, TensorBoard, or a custom backend.
SnapshotHook	Write the full batch state to a DataSink (GPUBuffer, HostMemory, or ZarrData).
ConvergedSnapshotHook	Write only newly converged samples to a DataSink. Fires at ON_CONVERGE; ideal for persisting optimized structures from FusedStage pipelines.
EnergyDriftMonitorHook	Track cumulative energy drift in NVE runs; warn or halt on excessive drift.
ProfilerHook	Instrument steps with NVTX ranges and wall-clock timing for Nsight Systems profiling. Fires at multiple stages via _runs_on_stage and uses plum dispatch to support dynamics and custom workflows.

Post-compute hooks (modify batch, fire at AFTER_COMPUTE)

These hooks modify the batch after the model forward pass and before the velocity update.

Hook	Purpose
NaNDetectorHook	Detect NaN/Inf in forces and energy; raise with diagnostic info (affected graph indices, step count).
MaxForceClampHook	Clamp per-atom force magnitudes to a safe maximum, preserving force direction. Prevents numerical explosions.

Constraint hooks (modify batch, fire at BEFORE_PRE_UPDATE)

These hooks enforce geometric constraints across integration steps.

Hook	Purpose
FreezeAtomsHook	Freeze atoms by category (e.g. substrate, boundary). Snapshots positions at BEFORE_PRE_UPDATE and restores them (with zeroed velocities) at AFTER_POST_UPDATE.

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 + force clamping

from nvalchemi.dynamics.hooks import MaxForceClampHook, NaNDetectorHook

dynamics = DemoDynamics(
    model=model,
    dt=0.5,
    hooks=[
        # Clamp first, then check — both fire at AFTER_COMPUTE
        # in registration order.
        MaxForceClampHook(max_force=50.0, log_clamps=True),
        NaNDetectorHook(extra_keys=["stress"]),
    ],
)

Enhanced sampling with a bias potential

from nvalchemi.hooks import BiasedPotentialHook
from nvalchemi.dynamics.base import DynamicsStage

def harmonic_restraint(batch):
    """Restrain center of mass to the origin."""
    k = 10.0  # eV/Ų
    com = batch.positions.mean(dim=0, keepdim=True)
    bias_energy = 0.5 * k * (com ** 2).sum().unsqueeze(0).unsqueeze(0)
    bias_forces = -k * com.expand_as(batch.positions) / batch.num_nodes
    return bias_energy, bias_forces

hook = BiasedPotentialHook(bias_fn=harmonic_restraint, stage=DynamicsStage.AFTER_COMPUTE)
dynamics = DemoDynamics(model=model, dt=0.5, hooks=[hook])

Profiling with Nsight Systems

from nvalchemi.dynamics.hooks import ProfilerHook

hook = ProfilerHook(enable_nvtx=True, enable_timer=True, frequency=10)
dynamics = DemoDynamics(model=model, n_steps=1_000, dt=0.5, hooks=[hook])

# Run under: nsys profile python my_script.py
dynamics.run(batch)

NVE energy drift monitoring

from nvalchemi.dynamics.hooks import EnergyDriftMonitorHook

hook = EnergyDriftMonitorHook(
    threshold=1e-5,
    metric="per_atom_per_step",
    action="raise",    # or "warn" for production
    frequency=100,
)
dynamics = DemoDynamics(model=model, dt=0.5, hooks=[hook])

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().