nvalchemi-dynamics-hooks

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):

class Hook(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), except ON_ADMISSION, which fires once per admission regardless of frequency.

HookContext — base snapshot shared by hook-enabled workflows:

@dataclass(kw_only=True)
class HookContext:
    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:

@dataclass(kw_only=True)
class DynamicsContext(HookContext):
    step_count: int = 0
    converged_mask: torch.Tensor | None = None
    active_graph_mask: torch.Tensor | None = None
    graduated_mask: torch.Tensor | None = None   # ON_GRADUATE only

ctx.active_graph_mask marks the active graphs, those whose status is below exit_status, as they were at the start of the step. A hook that must see a status change made earlier in the same step (for example one registered after a ConvergenceHook) reads the current status with BaseDynamics.active_graph_mask(ctx.batch, ctx.workflow.exit_status).

Access batch data via ctx.batch and dynamics step info via ctx.step_count.

Execution stages

Dynamics — DynamicsStage

Dynamics exposes 11 lifecycle stages. ON_ADMISSION fires once when a batch is admitted, while the remaining 10 stages fire within each step():

ON_ADMISSION (-1)  ← once before force priming and the first step
BEFORE_STEP (0)
  BEFORE_PRE_UPDATE (1)  →  pre_update()  →  AFTER_PRE_UPDATE (2)
  BEFORE_COMPUTE (3)     →  compute()      →  AFTER_COMPUTE (4)
  BEFORE_POST_UPDATE (5) →  post_update()  →  AFTER_POST_UPDATE (6)
AFTER_STEP (7)
ON_CONVERGE (8)   ← BaseDynamics: if detected; fused sub-stage: frequency-eligible steps
ON_GRADUATE (9)   ← whenever a hook listens; ctx.graduated_mask = graphs that reached exit_status

Stage selection guidelines (dynamics):

ON_ADMISSION is reset for every new run() and for managed membership changes such as refill or pipeline communication. In FusedStage, it runs outside compiled _step_impl, making it suitable for shape-dependent allocation and Python setup that per-step hooks cannot safely perform under fullgraph=True. It ignores the step-based frequency gate; a multi-stage hook's frequency still applies at its other stages.

In FusedStage, fused-level hooks wrap sub-stage hooks at every shared boundary: fused BEFORE_* hooks run before the corresponding sub-stage loop, and fused AFTER_* hooks run after it. Every hook receives ctx.active_graph_mask for the graphs participating at that boundary. Fused-level masks span all participating sub-stages; sub-stage masks are further restricted to graphs owned by that sub-stage. The shared model forward can evaluate the whole batch, but publication back to the batch uses the fused-level active mask. Graph-, atom-, and edge-level rows for inactive or graduated graphs retain their prior values. Mutating sub-stage hooks must then restrict their own writes to the sub-stage mask so they do not overwrite another stage's rows.

During a force-reprime iteration, a graph participates in the step and shared compute but skips both integrator updates. Therefore:

• BEFORE_STEP, BEFORE_COMPUTE, AFTER_COMPUTE, and AFTER_STEP include reprime-pending graphs.
• BEFORE_PRE_UPDATE, AFTER_PRE_UPDATE, BEFORE_POST_UPDATE, and AFTER_POST_UPDATE exclude them, at both the fused and sub-stage level.

Update masks are intentionally fixed at step start, so clearing reprime_pending after compute enables integrator updates on the next iteration without enabling post-update in the current one.

ON_CONVERGE remains sub-stage-only because convergence is evaluated independently per sub-stage. Fused sub-stages evaluate convergence every step; registered ON_CONVERGE hooks run when allowed by hook.frequency and must inspect ctx.converged_mask. BaseDynamics.step() calls them only when convergence is detected.

ON_GRADUATE reports graduation: a graph graduates on the step its status reaches exit_status, whether a criterion, a step budget, or another hook changed it. The stage fires after ON_CONVERGE (in FusedStage, also after the step-budget migration), and ctx.graduated_mask marks the graphs that graduated during the step. In a FusedStage it fires at both levels: first on each sub-stage, restricted to the graphs that sub-stage owned, then on the fused stage. It ignores hook.frequency and is dispatched on every step on which a hook is registered for it and the batch carries a status column, so the mask may be all False; read it rather than assume a graph graduated. DomainParallel does not dispatch it.

Registering hooks

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.

Deprecated. Use nvalchemi.enhanced_sampling (ConservativeBias) for new biases: it derives forces and stress from one energy definition by autograd. A bias_fn bias contributes no stress, so it is invisible to an NPT/NPH barostat — safe under NVE/NVT, silently wrong under a barostat. Run new biases through EnhancedSampling, which covers everything this hook does. The hook stays functional and no removal date is set; constructing it emits a DeprecationWarning.

def my_bias(batch: Batch) -> tuple[torch.Tensor, torch.Tensor]:
    """Return (bias_energy [B, 1], bias_forces [V, 3])."""
    bias_e = torch.zeros(batch.num_graphs, 1, device=batch.device)
    bias_f = torch.zeros_like(batch.positions)
    # ... compute bias ...
    return bias_e, bias_f

BiasedPotentialHook(
    bias_fn=my_bias,
    stage=DynamicsStage.AFTER_COMPUTE,
    frequency=1,
)

Observer hooks (stage: AFTER_STEP)

LoggingHook — log scalar observables.

LoggingHook(
    backend="csv",                  # "csv", "tensorboard", or "custom"
    frequency=100,
    log_path="md_log.csv",          # for file-based backends
    custom_scalars={                # additional scalars to log
        "max_velocity": lambda ctx: ctx.batch.velocities.norm(dim=-1).max(),
    },
    writer_fn=None,                 # custom writer for "custom" backend
)

SnapshotHook — save full batch state to a DataSink.

from nvalchemi.dynamics.sinks import GPUBuffer, HostMemory, ZarrData

SnapshotHook(
    sink=ZarrData("trajectory.zarr", capacity=10000),
    frequency=10,
)

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.

WrapPeriodicHook(frequency=10, stage=DynamicsStage.AFTER_POST_UPDATE)

Profiling hooks (multi-stage)

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

class TemperatureLogger:
    stage = DynamicsStage.AFTER_STEP
    frequency = 50

    def __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/K
        print(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:

from enum import Enum
from nvalchemi.dynamics.base import DynamicsStage
from nvalchemi.hooks import DynamicsContext

class StepTimerHook:
    stage = DynamicsStage.BEFORE_STEP  # primary stage (protocol compliance)
    frequency = 1

    def __init__(self):
        self._stages = {DynamicsStage.BEFORE_STEP, DynamicsStage.AFTER_STEP}
        self._t0 = None

    def _runs_on_stage(self, stage: Enum) -> bool:
        return stage in self._stages

    def __call__(self, ctx: DynamicsContext, stage: DynamicsStage) -> None:
        import time
        if stage == DynamicsStage.BEFORE_STEP:
            self._t0 = time.perf_counter()
        elif stage == DynamicsStage.AFTER_STEP and self._t0 is not None:
            dt = time.perf_counter() - self._t0
            print(f"Step {ctx.step_count}: {dt*1000:.1f} ms")

Option 3: Cross-category hook with plum dispatch

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 pipeline
class MyPipelineStage(Enum):
    BEFORE_PROCESS = 0
    AFTER_PROCESS = 1


@dataclass(kw_only=True)
class PipelineContext(HookContext):
    step_count: int = 0


class UniversalLoggerHook:
    stage = DynamicsStage.AFTER_STEP
    frequency = 10

    def __init__(self):
        self._stages = {DynamicsStage.AFTER_STEP, MyPipelineStage.AFTER_PROCESS}

    def _runs_on_stage(self, stage: Enum) -> bool:
        return stage in self._stages

    @dispatch
    def __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}")

    @dispatch
    def __call__(self, ctx: PipelineContext, stage: MyPipelineStage) -> None:
        print(f"[pipeline] step {ctx.step_count}: processed")

    @dispatch
    def __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 hook
class StepPrinter:
    stage = DynamicsStage.AFTER_STEP
    frequency = 10

    def __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)