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

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

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

Execution stages

Dynamics — DynamicsStage

Each step() call fires hooks at 9 stages in this order:

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)   ← only if convergence detected

Stage selection guidelines (dynamics):

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.

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)