Architecture Overview

The nvalchemi.dynamics module provides a composable framework for running batched molecular-dynamics (MD) and optimization workflows on GPUs. It is built around four core ideas:

1. One base class, two extension points — subclass BaseDynamics and override only pre_update / post_update.

2. A pluggable hook system — observe or modify every stage of the integration step via the Hook protocol.

3. Composable convergence criteria — ConvergenceHook lets you combine multiple criteria with AND semantics and automatically migrate samples between stages.

4. Operator-based composition — fuse stages on a single GPU with + or distribute across GPUs with |.

A single BaseDynamics.step() and its hook stages.

Inheritance hierarchy

Inheritance hierarchy of the dynamics module.

Every BaseDynamics subclass inherits communication capabilities automatically — no extra mixins required.

Quick-start

from nvalchemi.dynamics import DemoDynamics
from nvalchemi.dynamics.hooks import LoggingHook, NaNDetectorHook

# 1. Wrap your model
dynamics = DemoDynamics(
    model=my_model,
    n_steps=10_000,
    dt=0.5,
    hooks=[LoggingHook(backend="csv", log_path="log.csv", frequency=100), NaNDetectorHook()],
)

# 2. Run
dynamics.run(batch)

Single-GPU multi-stage (relax → MD):

fused = optimizer + md_dynamics          # uses the  +  operator
fused.run(batch)                         # one forward pass per step

Multi-GPU pipeline (2 ranks):

pipeline = optimizer | md_dynamics       # uses the  |  operator
with pipeline:
    pipeline.run()                       # each rank runs its stage

For a detailed walkthrough of how data flows through buffers, communication channels, and sinks — including the back-pressure mechanism and sample lifecycle — see Buffers & Data Flow.

Key concepts

Concept	Description
BaseDynamics	Coordinates a BaseModelMixin with a numerical integrator. Manages the step loop, hook execution, and model forward pass.
Hook	A runtime_checkable Protocol with frequency, stage, and __call__.
DynamicsStage	Nine insertion points covering every phase of a dynamics step.
ConvergenceHook	Composable convergence detector with AND semantics and optional status migration for FusedStage.
FusedStage	Composes N dynamics on one GPU with a shared forward pass. Each sub-stage processes only samples matching its status code.
DistributedPipeline	Maps one BaseDynamics per GPU rank and coordinates isend / irecv for sample graduation between stages.
DataSink	Pluggable storage for graduated / snapshotted samples: GPUBuffer, HostMemory, ZarrData.
SizeAwareSampler	Bin-packing sampler for inflight batching: graduated samples are replaced on the fly without rebuilding the batch.
BufferConfig	Required pre-allocation capacities for inter-rank communication buffers. See Buffers & Data Flow.