nvalchemi-dynamics-api

OverviewHistoryStatsSecurity

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

Documentation

SKILL.md

npx skills add ...

Documentation

SKILL.md

How to configure and run dynamics simulations, compose multi-stage pipelines (FusedStage, DistributedPipeline), use inflight batching, and manage data sinks. Use when writing any simulation script — molecular dynamics (NVE/NVT), structure relaxation or geometry optimization (e.g. FIRE), equation-of-state or adsorption scans — or orchestrating many structures through a batched GPU pipeline.

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

Documentation

SKILL.mdView source

nvalchemi Dynamics API

Overview

The dynamics API provides tools to discover available dynamics classes, configure them, and scale simulations up (single GPU) and out (multi-rank pipelines).

from nvalchemi.dynamics import (
    BaseDynamics,
    DemoDynamics,
    FusedStage,
    DistributedPipeline,
    ConvergenceHook,
    Hook,
    DynamicsStage,
    SizeAwareSampler,
    DataSink, GPUBuffer, HostMemory, ZarrData,
    hooks,
)

Available dynamics classes

Class	Description
`BaseDynamics`	Abstract base — subclass to create integrators
`DemoDynamics`	Velocity Verlet reference implementation (testing only)

To find all dynamics classes in a codebase, search for subclasses of BaseDynamics.

Configuring a dynamics run

Basic setup

from nvalchemi.dynamics import DemoDynamics, ConvergenceHook
from nvalchemi.models.demo import DemoModelWrapper
from nvalchemi.data import AtomicData, Batch
import torch
 
model = DemoModelWrapper()
dynamics = DemoDynamics(
    model=model,
    n_steps=1000,       # total steps for run()
    dt=0.5,             # timestep (fs)
)
 
# Create batch with required state
data = AtomicData(
    atomic_numbers=torch.tensor([6, 8, 1], 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)
 
result = dynamics.run(batch)

With convergence detection

dynamics = DemoDynamics(
    model=model,
    n_steps=10000,
    dt=0.5,
    convergence_hook=ConvergenceHook(
        criteria=[
            {"key": "fmax", "threshold": 0.05, "reduce_op": "max"},
            {"key": "energy_change", "threshold": 1e-6},
        ],
        frequency=1,
    ),
)

Shorthand for force-based convergence:

dynamics = DemoDynamics(
    model=model,
    n_steps=10000,
    dt=0.5,
    convergence_hook=ConvergenceHook.from_fmax(threshold=0.05),
)

With hooks

from nvalchemi.dynamics.hooks import MaxForceClampHook, LoggingHook
 
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),
    ],
)

Scaling up: FusedStage (single GPU, multiple stages)

FusedStage composes multiple dynamics stages on one GPU with a single shared model forward pass per step. Samples migrate between stages via convergence.

Composition with `+` operator

# Stage 0: geometry optimization
opt = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    convergence_hook=ConvergenceHook.from_fmax(0.05),
)
 
# Stage 1: production MD
md = DemoDynamics(model=model, n_steps=500, dt=0.5)
 
# Compose — auto-registers convergence hook to migrate status 0 → 1
fused = opt + md
 
result = fused.run(batch)  # runs until all samples reach exit_status

Constructor (explicit)

fused = FusedStage(
    sub_stages=[(0, opt), (1, md)],
    entry_status=0,          # initial sample status
    exit_status=2,           # auto-set to len(sub_stages) if -1
    compile_step=False,      # enable torch.compile
    compile_kwargs=None,     # kwargs for torch.compile
    reprime_on_entry=None,   # status codes needing a fresh force compute on entry
)

If a sub-stage needs freshly computed forces before it starts integrating (e.g. its AFTER_COMPUTE hooks differ from the stage the sample came from), pass its status code in reprime_on_entry. Newly entering samples skip one pre_update/post_update cycle while the shared compute and that stage's AFTER_COMPUTE hooks refresh their forces, then integrate normally the following iteration. This is separate from the initial batch-wide force prime and does not affect samples already in that status.

With torch.compile

fused = (opt + md)
fused.compile(fullgraph=True, mode="reduce-overhead")
 
with fused:  # lazy compilation on context entry
    result = fused.run(batch)

CUDA streams: `with` vs bare `run()`

Prefer to use context manager:

with dynamics:
    dynamics.run(batch)

As best pattern, even without explicit CUDA stream usage, as it in turn supports compile(mode="reduce-overhead")

Calling dynamics.run(batch) without context manager will reuse the current stream.

How it works

Each sample has a status field (integer)
Each sub-stage processes only samples matching its status code
When a sub-stage's ConvergenceHook fires, converged samples' status increments
Samples at exit_status are graduated (no longer updated)
run() loops until all samples reach exit_status or sampler is exhausted

Chaining more stages

fused_3 = opt + md + analysis   # 3 sub-stages: status 0, 1, 2
fused_3 = fused + extra_stage   # append to existing FusedStage

Scaling out: DistributedPipeline (multi-rank)

DistributedPipeline chains dynamics stages across multiple ranks using torch.distributed. Each rank runs one stage; converged samples are sent to the next rank. This is pipeline parallelism for dynamics, distinct from data-parallel multi-GPU training (DDP).

Composition with `|` operator

opt_stage = DemoDynamics(model=model, n_steps=100, dt=1.0)   # rank 0
md_stage = DemoDynamics(model=model, n_steps=500, dt=0.5)    # rank 1
 
pipeline = opt_stage | md_stage
 
with pipeline:   # initializes torch.distributed + setup
    pipeline.run()

Constructor (explicit)

pipeline = DistributedPipeline(
    stages={0: opt_stage, 1: md_stage},
    synchronized=False,   # True for debugging (adds barriers)
)

Communication modes

Control how inter-rank buffers synchronize:

stage = DemoDynamics(
    model=model, n_steps=100, dt=0.5,
    comm_mode="async_recv",   # default: deferred blocking
    # comm_mode="sync",       # immediate blocking (debugging)
    # comm_mode="fully_async", # maximum overlap
)

The default comm_mode is "async_recv". The three modes differ in when blocking occurs:

"sync": irecv completes inline in _prestep_sync_buffers; simplest and good for debugging.
"async_recv": irecv is posted in _prestep_sync_buffers, but wait() is deferred to _complete_pending_recv for communication overlap.
"fully_async": send and receive are both deferred for maximum overlap. Pending sends from the prior step are drained at the start of the next _prestep_sync_buffers.

Pre-allocated buffers

For high-throughput pipelines, pre-allocate send/recv buffers:

from nvalchemi.dynamics.base import BufferConfig
 
buffer_cfg = BufferConfig(
    num_systems=100,   # max graphs
    num_nodes=5000,    # total node capacity
    num_edges=20000,   # total edge capacity
)
 
stage = DemoDynamics(
    model=model, n_steps=100, dt=0.5,
    buffer_config=buffer_cfg,
)

Buffers are lazily initialized on the first step using the first concrete batch as a template for attribute keys, dtypes, and shapes. This means the first step has slightly more overhead.

Adjacent stages must use identical BufferConfig values. This is validated in DistributedPipeline.setup().

Buffer semantics and communication

Three buffer layers

The dynamics framework manages data flow through three layers:

Layer	Location	Purpose
Active batch	`_CommunicationMixin.active_batch`	Working set being integrated
Communication buffers	`send_buffer` / `recv_buffer`	Pre-allocated `Batch.empty()` for zero-copy inter-rank transfer
Overflow sinks	`DataSink` list (priority-ordered)	Staging when active batch is full

Communication protocol (DistributedPipeline)

Each pipeline step follows a four-phase protocol:

_prestep_sync_buffers() zeros the send buffer and posts irecv from the prior rank.
_complete_pending_recv() waits on deferred receive, routes into the active batch, and drains overflow sinks.
step() runs dynamics integration.
_poststep_sync_buffers(converged_indices) extracts converged samples into the send buffer and sends them to the next rank.

Deadlock prevention: when no samples converge, an empty send buffer is still sent so the downstream irecv completes.

Back-pressure

When send_buffer has limited capacity (via BufferConfig):

Only min(converged_count, remaining_capacity) samples are extracted
Excess converged samples remain in the active batch as no-ops. Their positions and velocities are saved before the integrator and restored after it runs.
Without BufferConfig, all converged samples are sent without constraints (backward compatible).

Buffer lifecycle: put/defrag/zero

# Pre-allocated buffer created via Batch.empty()
buffer = Batch.empty(num_systems=100, num_nodes=5000, num_edges=20000, template=batch)
 
# Copy selected graphs into buffer (Warp GPU kernels, float32 only)
mask = converged_mask  # bool tensor, True = copy this graph
buffer.put(src_batch, mask)
 
# Remove copied graphs from source in-place
src_batch.defrag()
 
# Reset buffer for reuse (preserves allocated memory)
buffer.zero()

Important: Batch.put() uses Warp GPU kernels that only handle float32 attributes. Adjacent pipeline stages must have identical BufferConfig values.

Data routing methods

Method	Purpose
`_recv_to_batch(incoming)`	Route received data through recv buffer into active batch
`_buffer_to_batch(incoming)`	Append to active batch, overflow to sinks if full
`_batch_to_buffer(mask)`	Copy graduated samples into send buffer, defrag active batch
`_overflow_to_sinks(batch)`	Write to first non-full sink in priority order
`_drain_sinks_to_batch()`	Pull from sinks back into active batch when room available

Inflight batching with SizeAwareSampler

For streaming workflows, SizeAwareSampler manages dataset access with bin-packing for size-matched batching. As samples converge and leave the batch, new samples are pulled from the dataset.

from nvalchemi.dynamics import SizeAwareSampler
 
sampler = SizeAwareSampler(
    dataset=dataset,                # must have __len__, __getitem__, get_metadata
    max_atoms=1000,                 # max total atoms per batch (None = auto from GPU)
    max_edges=5000,                 # max total edges per batch
    max_batch_size=32,              # max graphs per batch
    bin_width=10,                   # group samples by atom count bins
    shuffle=False,
    max_gpu_memory_fraction=0.8,    # for auto max_atoms estimation
)
 
# Build initial batch via greedy bin-packing
batch = sampler.build_initial_batch()
 
# Request a replacement sample that fits constraints
replacement = sampler.request_replacement(num_atoms=50, num_edges=200)
 
# Check if all samples consumed
sampler.exhausted  # bool

How inflight replacement works (`_refill_check`)

When refill_frequency triggers (every N steps), _refill_check():

Identifies graduated graphs (status >= exit_status)
Writes graduated graphs to sinks
Extracts remaining graphs via Batch.index_select
Requests replacements from sampler (one per graduated slot, matching atom/edge budget)
Appends replacements via Batch.append
Rebuilds status (replacements get 0) and fmax (replacements get inf) tensors

This produces a new Batch object, not an in-place mutation. It returns None when the sampler is exhausted and no active samples remain.

With FusedStage

opt = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    sampler=sampler,
    refill_frequency=1,    # check for replacements every N steps
    convergence_hook=ConvergenceHook.from_fmax(0.05),
)
md = DemoDynamics(model=model, n_steps=500, dt=0.5)
 
fused = opt + md
with fused:
    result = fused.run()  # no batch arg — built from sampler

With DistributedPipeline (first stage only)

first_stage = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    sampler=sampler,
    refill_frequency=1,
)
pipeline = first_stage | second_stage
with pipeline:
    pipeline.run()  # first stage auto-refills from sampler

Data sinks

Sinks store graduated/snapshot data. Used by SnapshotHook and the communication layer.

from nvalchemi.dynamics import GPUBuffer, HostMemory, ZarrData

DataSink interface

class DataSink(ABC):
    def write(self, batch: Batch, mask: Tensor | None = None) -> None: ...
    def read(self) -> Batch: ...
    def zero(self) -> None: ...           # clear contents
    def drain(self) -> Batch: ...         # read() + zero()
    def is_full(self) -> bool: ...        # len(self) >= capacity
    def __len__(self) -> int: ...
    @property
    def capacity(self) -> int: ...

Implementations

GPUBuffer — GPU-resident, pre-allocated.

gpu_sink = GPUBuffer(
    capacity=100,       # max graphs
    max_atoms=5000,     # total node capacity
    max_edges=20000,    # total edge capacity
    device="cuda",
)

HostMemory — CPU-resident list.

cpu_sink = HostMemory(capacity=1000)

ZarrData — Disk-backed persistent storage.

zarr_sink = ZarrData(
    store="trajectory.zarr",   # path, S3 URI, or dict
    capacity=1_000_000,
)

Using sinks with dynamics

from nvalchemi.dynamics.hooks import SnapshotHook
 
dynamics = DemoDynamics(
    model=model, n_steps=1000, dt=0.5,
    sinks=[gpu_sink, zarr_sink],   # for pipeline communication
    hooks=[
        SnapshotHook(sink=zarr_sink, frequency=10),
    ],
)

ConvergenceHook reference

ConvergenceHook(
    criteria=[                              # AND semantics (all must pass)
        {
            "key": "fmax",                  # batch attribute to check
            "threshold": 0.05,              # convergence threshold
            "reduce_op": "max",             # min, max, norm, mean, sum
            "reduce_dims": -1,              # dimensions to reduce
            "custom_op": None,              # custom callable
        },
    ],
    source_status=0,                        # check samples with this status
    target_status=1,                        # migrate converged samples
    frequency=1,                            # check every N steps
)
 
# Shorthand
ConvergenceHook.from_fmax(
    threshold=0.05,
    source_status=None,
    target_status=None,
    frequency=1,
)

Full workflow example

import torch
from nvalchemi.data import AtomicData, Batch
from nvalchemi.data.datapipes import AtomicDataZarrReader, Dataset
from nvalchemi.models.demo import DemoModelWrapper
from nvalchemi.dynamics import (
    DemoDynamics, ConvergenceHook, SizeAwareSampler, ZarrData,
)
from nvalchemi.dynamics.hooks import MaxForceClampHook, LoggingHook, SnapshotHook
 
# Model
model = DemoModelWrapper()
 
# Dataset + sampler for inflight batching
reader = AtomicDataZarrReader("structures.zarr")
dataset = Dataset(reader, device="cuda")
sampler = SizeAwareSampler(
    dataset=dataset, max_atoms=1000, max_edges=5000, max_batch_size=32,
)
 
# Output sink
output = ZarrData("results.zarr", capacity=100000)
 
# Stage 0: optimize
opt = DemoDynamics(
    model=model, n_steps=500, dt=1.0,
    sampler=sampler, refill_frequency=1,
    convergence_hook=ConvergenceHook.from_fmax(0.05),
    hooks=[MaxForceClampHook(max_force=10.0)],
)
 
# Stage 1: MD
md = DemoDynamics(
    model=model, n_steps=1000, dt=0.5,
    convergence_hook=ConvergenceHook(
        criteria=[{"key": "fmax", "threshold": 0.01}],
    ),
    hooks=[
        LoggingHook(backend="csv", log_path="md_log.csv", frequency=100),
        SnapshotHook(sink=output, frequency=50),
    ],
)
 
# Compose and run
fused = opt + md
with fused:
    fused.run()  # batch built from sampler, runs until exhausted

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

Documentation

SKILL.md

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

Documentation

SKILL.mdView source

nvalchemi Dynamics API

Overview

The dynamics API provides tools to discover available dynamics classes, configure them, and scale simulations up (single GPU) and out (multi-rank pipelines).

from nvalchemi.dynamics import (
    BaseDynamics,
    DemoDynamics,
    FusedStage,
    DistributedPipeline,
    ConvergenceHook,
    Hook,
    DynamicsStage,
    SizeAwareSampler,
    DataSink, GPUBuffer, HostMemory, ZarrData,
    hooks,
)

Available dynamics classes

Class	Description
`BaseDynamics`	Abstract base — subclass to create integrators
`DemoDynamics`	Velocity Verlet reference implementation (testing only)

To find all dynamics classes in a codebase, search for subclasses of BaseDynamics.

Configuring a dynamics run

Basic setup

from nvalchemi.dynamics import DemoDynamics, ConvergenceHook
from nvalchemi.models.demo import DemoModelWrapper
from nvalchemi.data import AtomicData, Batch
import torch
 
model = DemoModelWrapper()
dynamics = DemoDynamics(
    model=model,
    n_steps=1000,       # total steps for run()
    dt=0.5,             # timestep (fs)
)
 
# Create batch with required state
data = AtomicData(
    atomic_numbers=torch.tensor([6, 8, 1], 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)
 
result = dynamics.run(batch)

With convergence detection

dynamics = DemoDynamics(
    model=model,
    n_steps=10000,
    dt=0.5,
    convergence_hook=ConvergenceHook(
        criteria=[
            {"key": "fmax", "threshold": 0.05, "reduce_op": "max"},
            {"key": "energy_change", "threshold": 1e-6},
        ],
        frequency=1,
    ),
)

Shorthand for force-based convergence:

dynamics = DemoDynamics(
    model=model,
    n_steps=10000,
    dt=0.5,
    convergence_hook=ConvergenceHook.from_fmax(threshold=0.05),
)

With hooks

from nvalchemi.dynamics.hooks import MaxForceClampHook, LoggingHook
 
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),
    ],
)

Scaling up: FusedStage (single GPU, multiple stages)

FusedStage composes multiple dynamics stages on one GPU with a single shared model forward pass per step. Samples migrate between stages via convergence.

Composition with `+` operator

# Stage 0: geometry optimization
opt = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    convergence_hook=ConvergenceHook.from_fmax(0.05),
)
 
# Stage 1: production MD
md = DemoDynamics(model=model, n_steps=500, dt=0.5)
 
# Compose — auto-registers convergence hook to migrate status 0 → 1
fused = opt + md
 
result = fused.run(batch)  # runs until all samples reach exit_status

Constructor (explicit)

fused = FusedStage(
    sub_stages=[(0, opt), (1, md)],
    entry_status=0,          # initial sample status
    exit_status=2,           # auto-set to len(sub_stages) if -1
    compile_step=False,      # enable torch.compile
    compile_kwargs=None,     # kwargs for torch.compile
    reprime_on_entry=None,   # status codes needing a fresh force compute on entry
)

With torch.compile

fused = (opt + md)
fused.compile(fullgraph=True, mode="reduce-overhead")
 
with fused:  # lazy compilation on context entry
    result = fused.run(batch)

CUDA streams: `with` vs bare `run()`

Prefer to use context manager:

with dynamics:
    dynamics.run(batch)

As best pattern, even without explicit CUDA stream usage, as it in turn supports compile(mode="reduce-overhead")

Calling dynamics.run(batch) without context manager will reuse the current stream.

How it works

Each sample has a status field (integer)
Each sub-stage processes only samples matching its status code
When a sub-stage's ConvergenceHook fires, converged samples' status increments
Samples at exit_status are graduated (no longer updated)
run() loops until all samples reach exit_status or sampler is exhausted

Chaining more stages

fused_3 = opt + md + analysis   # 3 sub-stages: status 0, 1, 2
fused_3 = fused + extra_stage   # append to existing FusedStage

Scaling out: DistributedPipeline (multi-rank)

Composition with `|` operator

opt_stage = DemoDynamics(model=model, n_steps=100, dt=1.0)   # rank 0
md_stage = DemoDynamics(model=model, n_steps=500, dt=0.5)    # rank 1
 
pipeline = opt_stage | md_stage
 
with pipeline:   # initializes torch.distributed + setup
    pipeline.run()

Constructor (explicit)

pipeline = DistributedPipeline(
    stages={0: opt_stage, 1: md_stage},
    synchronized=False,   # True for debugging (adds barriers)
)

Communication modes

Control how inter-rank buffers synchronize:

stage = DemoDynamics(
    model=model, n_steps=100, dt=0.5,
    comm_mode="async_recv",   # default: deferred blocking
    # comm_mode="sync",       # immediate blocking (debugging)
    # comm_mode="fully_async", # maximum overlap
)

The default comm_mode is "async_recv". The three modes differ in when blocking occurs:

"sync": irecv completes inline in _prestep_sync_buffers; simplest and good for debugging.
"async_recv": irecv is posted in _prestep_sync_buffers, but wait() is deferred to _complete_pending_recv for communication overlap.
"fully_async": send and receive are both deferred for maximum overlap. Pending sends from the prior step are drained at the start of the next _prestep_sync_buffers.

Pre-allocated buffers

For high-throughput pipelines, pre-allocate send/recv buffers:

from nvalchemi.dynamics.base import BufferConfig
 
buffer_cfg = BufferConfig(
    num_systems=100,   # max graphs
    num_nodes=5000,    # total node capacity
    num_edges=20000,   # total edge capacity
)
 
stage = DemoDynamics(
    model=model, n_steps=100, dt=0.5,
    buffer_config=buffer_cfg,
)

Buffers are lazily initialized on the first step using the first concrete batch as a template for attribute keys, dtypes, and shapes. This means the first step has slightly more overhead.

Adjacent stages must use identical BufferConfig values. This is validated in DistributedPipeline.setup().

Buffer semantics and communication

Three buffer layers

The dynamics framework manages data flow through three layers:

Layer	Location	Purpose
Active batch	`_CommunicationMixin.active_batch`	Working set being integrated
Communication buffers	`send_buffer` / `recv_buffer`	Pre-allocated `Batch.empty()` for zero-copy inter-rank transfer
Overflow sinks	`DataSink` list (priority-ordered)	Staging when active batch is full

Communication protocol (DistributedPipeline)

Each pipeline step follows a four-phase protocol:

_prestep_sync_buffers() zeros the send buffer and posts irecv from the prior rank.
_complete_pending_recv() waits on deferred receive, routes into the active batch, and drains overflow sinks.
step() runs dynamics integration.
_poststep_sync_buffers(converged_indices) extracts converged samples into the send buffer and sends them to the next rank.

Deadlock prevention: when no samples converge, an empty send buffer is still sent so the downstream irecv completes.

Back-pressure

When send_buffer has limited capacity (via BufferConfig):

Only min(converged_count, remaining_capacity) samples are extracted
Excess converged samples remain in the active batch as no-ops. Their positions and velocities are saved before the integrator and restored after it runs.
Without BufferConfig, all converged samples are sent without constraints (backward compatible).

Buffer lifecycle: put/defrag/zero

# Pre-allocated buffer created via Batch.empty()
buffer = Batch.empty(num_systems=100, num_nodes=5000, num_edges=20000, template=batch)
 
# Copy selected graphs into buffer (Warp GPU kernels, float32 only)
mask = converged_mask  # bool tensor, True = copy this graph
buffer.put(src_batch, mask)
 
# Remove copied graphs from source in-place
src_batch.defrag()
 
# Reset buffer for reuse (preserves allocated memory)
buffer.zero()

Important: Batch.put() uses Warp GPU kernels that only handle float32 attributes. Adjacent pipeline stages must have identical BufferConfig values.

Data routing methods

Method	Purpose
`_recv_to_batch(incoming)`	Route received data through recv buffer into active batch
`_buffer_to_batch(incoming)`	Append to active batch, overflow to sinks if full
`_batch_to_buffer(mask)`	Copy graduated samples into send buffer, defrag active batch
`_overflow_to_sinks(batch)`	Write to first non-full sink in priority order
`_drain_sinks_to_batch()`	Pull from sinks back into active batch when room available

Inflight batching with SizeAwareSampler

For streaming workflows, SizeAwareSampler manages dataset access with bin-packing for size-matched batching. As samples converge and leave the batch, new samples are pulled from the dataset.

from nvalchemi.dynamics import SizeAwareSampler
 
sampler = SizeAwareSampler(
    dataset=dataset,                # must have __len__, __getitem__, get_metadata
    max_atoms=1000,                 # max total atoms per batch (None = auto from GPU)
    max_edges=5000,                 # max total edges per batch
    max_batch_size=32,              # max graphs per batch
    bin_width=10,                   # group samples by atom count bins
    shuffle=False,
    max_gpu_memory_fraction=0.8,    # for auto max_atoms estimation
)
 
# Build initial batch via greedy bin-packing
batch = sampler.build_initial_batch()
 
# Request a replacement sample that fits constraints
replacement = sampler.request_replacement(num_atoms=50, num_edges=200)
 
# Check if all samples consumed
sampler.exhausted  # bool

How inflight replacement works (`_refill_check`)

When refill_frequency triggers (every N steps), _refill_check():

Identifies graduated graphs (status >= exit_status)
Writes graduated graphs to sinks
Extracts remaining graphs via Batch.index_select
Requests replacements from sampler (one per graduated slot, matching atom/edge budget)
Appends replacements via Batch.append
Rebuilds status (replacements get 0) and fmax (replacements get inf) tensors

This produces a new Batch object, not an in-place mutation. It returns None when the sampler is exhausted and no active samples remain.

With FusedStage

opt = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    sampler=sampler,
    refill_frequency=1,    # check for replacements every N steps
    convergence_hook=ConvergenceHook.from_fmax(0.05),
)
md = DemoDynamics(model=model, n_steps=500, dt=0.5)
 
fused = opt + md
with fused:
    result = fused.run()  # no batch arg — built from sampler

With DistributedPipeline (first stage only)

first_stage = DemoDynamics(
    model=model, n_steps=100, dt=1.0,
    sampler=sampler,
    refill_frequency=1,
)
pipeline = first_stage | second_stage
with pipeline:
    pipeline.run()  # first stage auto-refills from sampler

Data sinks

Sinks store graduated/snapshot data. Used by SnapshotHook and the communication layer.

from nvalchemi.dynamics import GPUBuffer, HostMemory, ZarrData

DataSink interface

class DataSink(ABC):
    def write(self, batch: Batch, mask: Tensor | None = None) -> None: ...
    def read(self) -> Batch: ...
    def zero(self) -> None: ...           # clear contents
    def drain(self) -> Batch: ...         # read() + zero()
    def is_full(self) -> bool: ...        # len(self) >= capacity
    def __len__(self) -> int: ...
    @property
    def capacity(self) -> int: ...

Implementations

GPUBuffer — GPU-resident, pre-allocated.

gpu_sink = GPUBuffer(
    capacity=100,       # max graphs
    max_atoms=5000,     # total node capacity
    max_edges=20000,    # total edge capacity
    device="cuda",
)

HostMemory — CPU-resident list.

cpu_sink = HostMemory(capacity=1000)

ZarrData — Disk-backed persistent storage.

zarr_sink = ZarrData(
    store="trajectory.zarr",   # path, S3 URI, or dict
    capacity=1_000_000,
)

Using sinks with dynamics

from nvalchemi.dynamics.hooks import SnapshotHook
 
dynamics = DemoDynamics(
    model=model, n_steps=1000, dt=0.5,
    sinks=[gpu_sink, zarr_sink],   # for pipeline communication
    hooks=[
        SnapshotHook(sink=zarr_sink, frequency=10),
    ],
)

ConvergenceHook reference

ConvergenceHook(
    criteria=[                              # AND semantics (all must pass)
        {
            "key": "fmax",                  # batch attribute to check
            "threshold": 0.05,              # convergence threshold
            "reduce_op": "max",             # min, max, norm, mean, sum
            "reduce_dims": -1,              # dimensions to reduce
            "custom_op": None,              # custom callable
        },
    ],
    source_status=0,                        # check samples with this status
    target_status=1,                        # migrate converged samples
    frequency=1,                            # check every N steps
)
 
# Shorthand
ConvergenceHook.from_fmax(
    threshold=0.05,
    source_status=None,
    target_status=None,
    frequency=1,
)

Full workflow example

import torch
from nvalchemi.data import AtomicData, Batch
from nvalchemi.data.datapipes import AtomicDataZarrReader, Dataset
from nvalchemi.models.demo import DemoModelWrapper
from nvalchemi.dynamics import (
    DemoDynamics, ConvergenceHook, SizeAwareSampler, ZarrData,
)
from nvalchemi.dynamics.hooks import MaxForceClampHook, LoggingHook, SnapshotHook
 
# Model
model = DemoModelWrapper()
 
# Dataset + sampler for inflight batching
reader = AtomicDataZarrReader("structures.zarr")
dataset = Dataset(reader, device="cuda")
sampler = SizeAwareSampler(
    dataset=dataset, max_atoms=1000, max_edges=5000, max_batch_size=32,
)
 
# Output sink
output = ZarrData("results.zarr", capacity=100000)
 
# Stage 0: optimize
opt = DemoDynamics(
    model=model, n_steps=500, dt=1.0,
    sampler=sampler, refill_frequency=1,
    convergence_hook=ConvergenceHook.from_fmax(0.05),
    hooks=[MaxForceClampHook(max_force=10.0)],
)
 
# Stage 1: MD
md = DemoDynamics(
    model=model, n_steps=1000, dt=0.5,
    convergence_hook=ConvergenceHook(
        criteria=[{"key": "fmax", "threshold": 0.01}],
    ),
    hooks=[
        LoggingHook(backend="csv", log_path="md_log.csv", frequency=100),
        SnapshotHook(sink=output, frequency=50),
    ],
)
 
# Compose and run
fused = opt + md
with fused:
    fused.run()  # batch built from sampler, runs until exhausted

nvalchemi-dynamics-api

Documentation

Documentation

nvalchemi-dynamics-api

Documentation

nvalchemi Dynamics API

Overview

Available dynamics classes

Configuring a dynamics run

Basic setup

With convergence detection

With hooks

Scaling up: FusedStage (single GPU, multiple stages)

Composition with + operator

Constructor (explicit)

With torch.compile

CUDA streams: with vs bare run()

How it works

Chaining more stages

Scaling out: DistributedPipeline (multi-rank)

Composition with | operator

Constructor (explicit)

Communication modes

Pre-allocated buffers

Buffer semantics and communication

Three buffer layers

Communication protocol (DistributedPipeline)

Back-pressure

Buffer lifecycle: put/defrag/zero

Data routing methods

Inflight batching with SizeAwareSampler

How inflight replacement works (_refill_check)

With FusedStage

With DistributedPipeline (first stage only)

Data sinks

DataSink interface

Implementations

Using sinks with dynamics

ConvergenceHook reference

Full workflow example

Documentation

Documentation

nvalchemi Dynamics API

Overview

Available dynamics classes

Configuring a dynamics run

Basic setup

With convergence detection

With hooks

Scaling up: FusedStage (single GPU, multiple stages)

Composition with + operator

Constructor (explicit)

With torch.compile

CUDA streams: with vs bare run()

How it works

Chaining more stages

Scaling out: DistributedPipeline (multi-rank)

Composition with | operator

Constructor (explicit)

Communication modes

Pre-allocated buffers

Buffer semantics and communication

Three buffer layers

Communication protocol (DistributedPipeline)

Back-pressure

Buffer lifecycle: put/defrag/zero

Data routing methods

Inflight batching with SizeAwareSampler

How inflight replacement works (_refill_check)

With FusedStage

With DistributedPipeline (first stage only)

Data sinks

DataSink interface

Implementations

Using sinks with dynamics

ConvergenceHook reference

Full workflow example

Composition with `+` operator

CUDA streams: `with` vs bare `run()`

Composition with `|` operator

How inflight replacement works (`_refill_check`)

Composition with `+` operator

CUDA streams: `with` vs bare `run()`

Composition with `|` operator

How inflight replacement works (`_refill_check`)