nvalchemi.dynamics.hooks.SnapshotHook

class nvalchemi.dynamics.hooks.SnapshotHook(sink, frequency=1)

Save a snapshot of the active batch to a DataSink at a given frequency.

This hook writes the full batch state — positions, velocities, forces, energies, and any other tensors present on the Batch — to the configured sink every frequency steps. It is the primary mechanism for recording trajectories and creating restart checkpoints during dynamics runs.

The hook delegates serialization entirely to the DataSink interface, meaning the same SnapshotHook instance works with any backend:

• GPUBuffer — pre-allocated device memory for high-speed, in-simulation buffering.

• HostMemory — CPU-resident list-of-AtomicData storage, useful for staging before disk I/O.

• ZarrData — persistent, Zarr-backed storage with CSR-style layout for variable-length graph data; supports local, in-memory, and remote (S3/GCS) stores.

Because SnapshotHook inherits from _ObserverHook, it fires at AFTER_STEP — after all integrator updates, force clamping, and convergence checks have completed — guaranteeing that the snapshot reflects the fully resolved state for each recorded step.

Parameters:

• sink (DataSink) – The storage backend to write snapshots to.

• frequency (int, optional) – Write a snapshot every frequency steps. Default 1 (every step).

sink

The storage backend.

Type:

DataSink

frequency

Snapshot frequency in steps.

Type:

int

stage

Fixed to AFTER_STEP.

Type:

HookStageEnum

Examples

>>> from nvalchemi.dynamics.hooks import SnapshotHook
>>> from nvalchemi.dynamics.sinks import HostMemory
>>> sink = HostMemory(capacity=10_000)
>>> hook = SnapshotHook(sink=sink, frequency=10)
>>> dynamics = DemoDynamics(model=model, n_steps=1000, dt=0.5, hooks=[hook])
>>> dynamics.run(batch)  # 100 snapshots written
>>> trajectory = sink.read()

Notes

• The hook does not clone the batch before writing. Whether data is copied depends on the sink implementation (e.g. HostMemory moves to CPU; GPUBuffer copies into pre-allocated slots).

• For long simulations, prefer ZarrData to avoid accumulating the full trajectory in memory.

• When used inside a FusedStage, the snapshot includes samples at all status codes in a single write.

__init__(sink, frequency=1)

Parameters:

• sink (DataSink)

• frequency (int)

Return type:

None

Methods

__init__(sink[, frequency])