nvalchemi-dynamics-hooks

OverviewHistoryStatsSecurity

npx skills add nvidia/nvalchemi-toolkit --skill nvalchemi-dynamics-hooks

Documentation

SKILL.md

npx skills add ...

Documentation

SKILL.md

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.

npx skills add nvidia/nvalchemi-toolkit --skill nvalchemi-dynamics-hooks

Documentation

SKILL.mdView source

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

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

Execution stages

Dynamics — `DynamicsStage`

Dynamics exposes 10 lifecycle stages. ON_ADMISSION fires once when a batch is admitted, while the remaining 9 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

Stage selection guidelines (dynamics):

Goal	Stage
Validate or allocate for a newly admitted batch	`DynamicsStage.ON_ADMISSION`
Modify forces/energy after model	`DynamicsStage.AFTER_COMPUTE`
Observe final state (logging, snapshots)	`DynamicsStage.AFTER_STEP`
Wrap positions after velocity update	`DynamicsStage.AFTER_POST_UPDATE`
Instrument timing / profiling	`DynamicsStage.BEFORE_STEP`
React to convergence	`DynamicsStage.ON_CONVERGE`

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.

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.

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

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)

npx skills add nvidia/nvalchemi-toolkit --skill nvalchemi-dynamics-hooks

Documentation

SKILL.md

npx skills add nvidia/nvalchemi-toolkit --skill nvalchemi-dynamics-hooks

Documentation

SKILL.mdView source

nvalchemi Hooks

Overview

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

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

Execution stages

Dynamics — `DynamicsStage`

Dynamics exposes 10 lifecycle stages. ON_ADMISSION fires once when a batch is admitted, while the remaining 9 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

Stage selection guidelines (dynamics):

Goal	Stage
Validate or allocate for a newly admitted batch	`DynamicsStage.ON_ADMISSION`
Modify forces/energy after model	`DynamicsStage.AFTER_COMPUTE`
Observe final state (logging, snapshots)	`DynamicsStage.AFTER_STEP`
Wrap positions after velocity update	`DynamicsStage.AFTER_POST_UPDATE`
Instrument timing / profiling	`DynamicsStage.BEFORE_STEP`
React to convergence	`DynamicsStage.ON_CONVERGE`

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.

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

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}")

Hook ordering recommendations

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)

nvalchemi-dynamics-hooks

Documentation

Documentation

nvalchemi-dynamics-hooks

Documentation

nvalchemi Hooks

Overview

Hook protocol

Execution stages

Dynamics — DynamicsStage

Registering hooks

Built-in hooks

Safety hooks (stage: AFTER_COMPUTE)

Bias hook (stage: AFTER_COMPUTE)

Observer hooks (stage: AFTER_STEP)

Periodic boundary hook (stage: AFTER_POST_UPDATE)

Profiling hooks (multi-stage)

Writing a custom hook

Option 1: Simple single-stage hook (dynamics)

Option 2: Multi-stage hook with _runs_on_stage

Option 3: Cross-category hook with plum dispatch

Hook ordering recommendations

Complete example

Documentation

Documentation

nvalchemi Hooks

Overview

Hook protocol

Execution stages

Dynamics — DynamicsStage

Registering hooks

Built-in hooks

Safety hooks (stage: AFTER_COMPUTE)

Bias hook (stage: AFTER_COMPUTE)

Observer hooks (stage: AFTER_STEP)

Periodic boundary hook (stage: AFTER_POST_UPDATE)

Profiling hooks (multi-stage)

Writing a custom hook

Option 1: Simple single-stage hook (dynamics)

Option 2: Multi-stage hook with _runs_on_stage

Option 3: Cross-category hook with plum dispatch

Hook ordering recommendations

Complete example

Dynamics — `DynamicsStage`

Option 2: Multi-stage hook with `_runs_on_stage`

Option 3: Cross-category hook with `plum` dispatch

Dynamics — `DynamicsStage`

Option 2: Multi-stage hook with `_runs_on_stage`

Option 3: Cross-category hook with `plum` dispatch