How to use and write dynamics hooks — callbacks that observe or modify batch state at specific points during each simulation step. Use when a simulation needs neighbor-list rebuilds, convergence checks or early stopping, temperature control, per-step logging or trajectory capture, or any custom per-step behavior attached to a dynamics run.
Install with Codex or Claude Copy this prompt, paste it into Codex, Claude, or another assistant, and let it review the skill page and install it for you.
A direct command skips the review prompt. Inspect the source before running it.
How to use and write dynamics hooks — callbacks that observe or modify batch state at specific points during each simulation step. Use when a simulation needs neighbor-list rebuilds, convergence checks or early stopping, temperature control, per-step logging or trajectory capture, or any custom per-step behavior attached to a dynamics run.
nvalchemi Hooks
Overview
Hooks are callbacks that fire at specific points during each workflow step.
They observe or modify batch state without changing the engine itself.
The hook system is framework-wide: the same Hook protocol works for
dynamics and custom pipelines. Dynamics engines pass DynamicsContext;
custom engines can pass HookContext or their own context subclass.
from nvalchemi.hooks import (
BiasedPotentialHook,
DynamicsContext,
Hook,
HookContext,
HookRegistryMixin,
NeighborListHook,
WrapPeriodicHook,
)
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.dynamics.hooks import (
EnergyDriftMonitorHook,
LoggingHook,
MaxForceClampHook,
NaNDetectorHook,
SnapshotHook,
StageTimingHook,
TorchProfilerHook,
)
Hook protocol
Any object with these attributes satisfies the Hook protocol (runtime-checkable):
classHook(Protocol):
frequency: int# execute every N steps (1 = every step)
stage: Enum | None
() -> :
...
# stage enum value (None for stage-agnostic hooks)
def
__call__
self, ctx: HookContext, stage: Enum
None
"""Called with a context snapshot and the current stage."""
A hook fires when step_count % hook.frequency == 0 (so all hooks fire at
step 0).
HookContext — base snapshot shared by hook-enabled workflows:
@dataclass(kw_only=True)classHookContext:
batch: Batch # current batch (all engines)
model: BaseModelMixin | None = None
global_rank: int = 0# distributed rank
workflow: Any = None# back-reference to the engine
DynamicsContext — context passed by dynamics engines:
from nvalchemi.dynamics.demo import DemoDynamics
# At construction
dynamics = DemoDynamics(
model=model,
n_steps=1000,
dt=0.5,
hooks=[
MaxForceClampHook(max_force=10.0),
LoggingHook(backend="csv", log_path="md_log.csv", frequency=100),
],
)
# After construction
dynamics.register_hook(NaNDetectorHook(frequency=10))
Multiple hooks at the same stage fire in registration order.
Stage type enforcement: each engine declares _stage_type to restrict
which enum types are accepted. For example, BaseDynamics sets
_stage_type = DynamicsStage.
Built-in hooks
Safety hooks (stage: AFTER_COMPUTE)
NaNDetectorHook — detect NaN/Inf in forces and energy.
NaNDetectorHook(
frequency=1, # check every N steps
extra_keys=["stress"], # additional batch keys to check (optional)
)
MaxForceClampHook — clamp per-atom force vectors to a maximum L2 norm.
MaxForceClampHook(
max_force=10.0, # max force norm (eV/A)
frequency=1,
)
Bias hook (stage: AFTER_COMPUTE)
BiasedPotentialHook — add an external bias potential for enhanced sampling.
EnergyDriftMonitorHook — track total energy drift.
EnergyDriftMonitorHook(
threshold=1e-4, # drift threshold
metric="per_atom_per_step", # or "absolute"
action="warn", # or "raise"
frequency=1,
include_kinetic=True, # include kinetic energy
)
Periodic boundary hook (stage: AFTER_POST_UPDATE)
WrapPeriodicHook — wrap positions back into the unit cell.
StageTimingHook — per-stage NVTX ranges and wall-clock timing. Registers
itself at every profiled stage via _runs_on_stage, records timestamps, and
computes per-transition deltas (optionally written to CSV or console).
StageTimingHook(
profiled_stages="all", # "all", "step", "detailed", or a set[Enum]
frequency=1,
enable_nvtx=True, # NVTX push/pop ranges for Nsight Systems
timer_backend="auto", # "auto", "cuda_event", or "perf_counter"
log_path="timing.csv", # optional CSV of per-transition timings
show_console=False, # print a timing table via loguru
)
Call profiler.summary() after the run for aggregated per-stage timings. For
full kernel-level PyTorch profiler traces, use TorchProfilerHook, which
captures traces through PhysicsNeMo's profiler wrapper.
Writing a custom hook
Option 1: Simple single-stage hook (dynamics)
Implement the protocol directly — no inheritance needed.
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.hooks import DynamicsContext
classTemperatureLogger:
stage = DynamicsStage.AFTER_STEP
frequency = 50def__call__(self, ctx: DynamicsContext, stage: DynamicsStage) -> None:
ke = ctx.batch.kinetic_energies.sum()
n_atoms = ctx.batch.num_nodes
temp = 2.0 * ke / (3.0 * n_atoms * 8.617e-5) # kB in eV/Kprint(f"Step {ctx.step_count}: T = {temp:.1f} K")
Option 2: Multi-stage hook with _runs_on_stage
Fire at multiple stages by defining _runs_on_stage(stage) -> bool:
For hooks that work with multiple stage enum types (e.g. DynamicsStage and
a custom enum), use plum.dispatch to overload __call__ with different
stage types:
from dataclasses import dataclass
from enum import Enum
from plum import dispatch
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.hooks import DynamicsContext, HookContext
# Example custom stage enum for a hypothetical pipelineclassMyPipelineStage(Enum):
BEFORE_PROCESS = 0
AFTER_PROCESS = 1@dataclass(kw_only=True)classPipelineContext(HookContext):
step_count: int = 0classUniversalLoggerHook:
stage = DynamicsStage.AFTER_STEP
frequency = 10def__init__(self):
self._stages = {DynamicsStage.AFTER_STEP, MyPipelineStage.AFTER_PROCESS}
def_runs_on_stage(self, stage: Enum) -> bool:
return stage inself._stages
@dispatchdef__call__(self, ctx: DynamicsContext, stage: DynamicsStage) -> None:
fmax = ctx.batch.forces.norm(dim=-1).max().item()
print(f"[dynamics] step {ctx.step_count}: fmax={fmax:.4f}")
@dispatchdef__call__(self, ctx: PipelineContext, stage: MyPipelineStage) -> None:
print(f"[pipeline] step {ctx.step_count}: processed")
@dispatchdef__call__(self, ctx: HookContext, stage: Enum) -> None:
print(f"[custom] stage={stage.name}, graphs={ctx.batch.num_graphs}")
Use this plum.dispatch pattern when one hook must handle several
context/stage types. Built-in multi-stage hooks like StageTimingHook
instead use the simpler _runs_on_stage approach from Option 2.
Hook ordering recommendations
Register hooks in this order for correct behavior:
hooks = [
# 1. Bias (modifies forces/energy)
BiasedPotentialHook(bias_fn=my_bias, stage=DynamicsStage.AFTER_COMPUTE),
# 2. Safety (clamp after all force modifications)
MaxForceClampHook(max_force=10.0),
# 3. NaN detection (check final forces)
NaNDetectorHook(),
# 4. Periodic wrapping
WrapPeriodicHook(frequency=10, stage=DynamicsStage.AFTER_POST_UPDATE),
# 5. Observers (read final state)
LoggingHook(backend="csv", log_path="md_log.csv", frequency=100),
SnapshotHook(sink=my_sink, frequency=50),
EnergyDriftMonitorHook(threshold=1e-4),
# 6. Profiling
StageTimingHook(),
]
dynamics = DemoDynamics(model=model, n_steps=10000, dt=0.5, hooks=hooks)
Complete example
import torch
from nvalchemi.data import AtomicData, Batch
from nvalchemi.models.demo import DemoModel, DemoModelWrapper
from nvalchemi.dynamics.demo import DemoDynamics
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.hooks import DynamicsContext
from nvalchemi.dynamics.hooks import MaxForceClampHook, NaNDetectorHook
# Custom hookclassStepPrinter:
stage = DynamicsStage.AFTER_STEP
frequency = 10def__call__(self, ctx: DynamicsContext, stage: DynamicsStage) -> None:
fmax = ctx.batch.forces.norm(dim=-1).max().item()
print(f"Step {ctx.step_count}: fmax={fmax:.4f}")
# Setup
model = DemoModelWrapper(DemoModel())
dynamics = DemoDynamics(
model=model,
n_steps=100,
dt=0.5,
hooks=[
MaxForceClampHook(max_force=10.0),
NaNDetectorHook(),
StepPrinter(),
],
)
data = AtomicData(
atomic_numbers=torch.tensor([6, 6, 8], dtype=torch.long),
positions=torch.randn(3, 3),
)
batch = Batch.from_data_list([data])
batch.forces = torch.zeros(3, 3)
batch.energy = torch.zeros(1, 1)
dynamics.run(batch)
