nvalchemi.dynamics.BaseDynamics

class nvalchemi.dynamics.BaseDynamics(model, hooks=None, convergence_hook=None, n_steps=None, exit_status=1, **kwargs)

Base class for all dynamics simulations.

This class coordinates a BaseModelMixin model with a numerical integrator to evolve a Batch of atomic systems over time. It manages the step loop, hook execution at stage boundaries, and model evaluation.

BaseDynamics inherits from HookRegistryMixin for hook storage and from _CommunicationMixin for inter-rank communication and buffer management for pipeline execution. All dynamics subclasses automatically have communication capabilities.

The public interface centers on three methods. run(batch) is the top-level entry point: it repeatedly calls step() for n_steps iterations and is the only method most users need. n_steps can be set at construction time or passed to run(). step(batch) executes a single simulation step, orchestrating the full hook-wrapped sequence pre_update → compute → post_update, with hooks fired at each stage boundary, followed by convergence checking. Subclasses should generally NOT override step. compute(batch) performs the model forward pass: it calls model(batch) which must return a fully adapted ModelOutputs dict, validates outputs against __needs_keys__, and writes results (forces, energies, stresses) back to the batch in-place. Subclasses should generally NOT override compute.

Parameters:

• model (BaseModelMixin)

• hooks (list[Hook])

• convergence_hook (Any)

• n_steps (int | None)

• exit_status (int)

• kwargs (Any)

model

The neural network potential model.

Type:

BaseModelMixin

step_count

The current step number, starting from 0.

Type:

int

hooks

Flat list of registered hooks.

Type:

list[Hook]

model_is_conservative

Indicates that the model uses automatic differentiation to obtain forces.

Type:

bool

convergence_hook

Hook that evaluates composable convergence criteria. Defaults to a single forces-based criterion with threshold 0.05.

Type:

ConvergenceHook

n_steps

Total number of simulation steps for run(). None means the step count must be supplied when calling run().

Type:

int | None

exit_status

Status code threshold for graduated samples. Samples with status >= exit_status are treated as no-ops during step() — their positions and velocities are preserved through the integrator. Default is 1.

Type:

int

__needs_keys__

Set of output keys that this dynamics requires from the model. Empty by default on BaseDynamics. Subclasses declare their own requirements (e.g., typically forces for optimization and MD). Checked in _validate_model_outputs() after each forward pass.

Type:

set[str]

__provides_keys__

Set of keys that this dynamics produces or updates on the batch beyond model outputs. Empty by default. Subclasses declare what additional state they provide (e.g., {"velocities", "positions"} for velocity verlet). Used for validation and buffer preallocation.

Type:

set[str]

Notes

Developers implementing a new integrator should override pre_update(batch) and post_update(batch) to implement the integration scheme. These are called around compute() — pre_update before, post_update after. For example, Velocity Verlet updates positions in pre_update and velocities in post_update. The class-level sets __needs_keys__ and __provides_keys__ declare what outputs the dynamics requires from the model and what additional state it produces; requirements are checked in _validate_model_outputs() after each forward pass. masked_update(batch, mask) is used by FusedStage to apply pre_update/post_update only to a subset of samples in a batched setting. Models must be BaseModelMixin instances — plain nn.Module is not accepted.

Examples

>>> model = MyPotentialModel()
>>> dynamics = BaseDynamics(model, n_steps=1000)
>>> dynamics.run(batch)

compute(batch)

Perform the model forward pass to compute forces and energies.

This method:

1. Runs the model forward pass, which should enable gradients

2. Adapts outputs to the standard format

3. Validates outputs against dynamics requirements

4. Writes known keys back to the batch in-place via _OUTPUT_KEY_TO_BATCH_ATTR

5. Detaches all output tensors from the computation graph and exposes them as _last_outputs for custom dynamics subclasses that need charges, embeddings, or other non-standard outputs.

6. Clears requires_grad on batch tensors that the model enabled for autograd (e.g. positions), so downstream hooks can safely perform in-place operations.

The detach in step 5 is deliberate: model wrappers may return tensors that are still attached to the autograd graph (e.g. MACE returns energies on the graph even after computing forces internally). Since compute() is a terminal consumer — values have already been copied into the batch — holding the graph would cause memory to grow without bound across dynamics steps. Callers that need the live graph (e.g. training loops computing a loss) should call model(batch) directly instead of compute().

Parameters:

batch (Batch) – The current batch of atomic data. Will have forces and energies updated in-place.

Returns:

OrderedDict containing the model outputs (energies, forces, and any other computed properties). All tensors are detached from the computation graph.

Return type:

ModelOutputs

Raises:

RuntimeError – If the model outputs do not satisfy the dynamics requirements specified by __needs_keys__.

masked_update(batch, mask)

Apply pre_update and post_update only to selected samples in the batch.

This method allows selective updates where only some graphs in the batch are modified. Unmasked samples retain their original positions and velocities.

The mask is a boolean tensor of shape (B,) where B is the number of graphs. True values indicate samples that should be updated.

Parameters:

• batch (Batch) – The current batch of atomic data, modified in-place.

• mask (Bool[Tensor, "B"]) – Boolean mask selecting which graphs to update. Shape (B,) where B is the number of graphs in the batch.

Return type:

None

Notes

This method expands the graph-level mask to node-level using batch.batch_idx to correctly index per-node tensors like positions and velocities.

property model_is_conservative: bool

Returns whether or not the model uses conservative forces.

post_update(batch)

Perform the second half of the integration step.

This method is a no-op in the base class and should be overridden by integrator subclasses (e.g., Velocity Verlet would update velocities here).

Parameters:

batch (Batch) – The current batch of atomic data, modified in-place.

Return type:

None

pre_update(batch)

Perform the first half of the integration step.

This method is a no-op in the base class and should be overridden by integrator subclasses (e.g., Velocity Verlet would update positions here).

Parameters:

batch (Batch) – The current batch of atomic data, modified in-place.

Return type:

None

refill_check(batch, exit_status)

Replace graduated samples via index-select and append.

Graduated graphs (status >= exit_status) are written to sinks, then removed via Batch.index_select() on the remaining indices. Replacement samples from the sampler are appended via Batch.append(). Dynamics-specific bookkeeping fields are written into the result batch via the _bookkeeping_keys registry.

Parameters:

• batch (Batch) – The current batch with a status field.

• exit_status (int) – Status code indicating graduation.

Returns:

A new batch with graduated graphs replaced by fresh samples, or None if no active samples remain (sampler exhausted and all graduated) — in which case self.done is set to True.

Return type:

Batch | None

Raises:

RuntimeError – If self.sampler is None.

classmethod register_bookkeeping_key(key, init_fn)

Register a graph-level bookkeeping field to survive refill_check.

Parameters:

• key (str) – Field name on Batch.

• init_fn (Callable[[int, torch.device], torch.Tensor]) – Factory that creates a zero-initialized tensor of shape (n, 1) for n systems on the given device.

Return type:

None

run(batch, n_steps=None)

Run the dynamics simulation for a specified number of steps.

This is a convenience method that repeatedly calls step(). The step count can be set at construction time via the n_steps parameter, or passed directly to this method. A value passed here takes precedence over the instance attribute.

Parameters:

• batch (Batch) – The initial batch of atomic data.

• n_steps (int | None, optional) – The number of steps to run. If None, falls back to self.n_steps. If both are None, raises ValueError.

Returns:

The batch after all steps have been executed.

Return type:

Batch

Raises:

ValueError – If no step count is available (both the argument and self.n_steps are None).

step(batch)

Execute a single dynamics step with the full hook-wrapped sequence.

The step proceeds as follows: 1. BEFORE_STEP hooks 2. BEFORE_PRE_UPDATE hooks -> pre_update() -> AFTER_PRE_UPDATE hooks 3. BEFORE_COMPUTE hooks -> compute() -> AFTER_COMPUTE hooks 4. BEFORE_POST_UPDATE hooks -> post_update() -> AFTER_POST_UPDATE hooks 5. AFTER_STEP hooks 6. Check convergence and fire ON_CONVERGE hooks if any samples converged 7. Increment step_count

Samples with status >= exit_status are treated as no-ops for the integrator (pre_update/post_update). Their positions and velocities are preserved through the step. This enables back-pressure handling in pipeline mode where converged samples may remain in the active batch when the send buffer is full.

Parameters:

batch (Batch) – The current batch of atomic data.

Returns:

The updated batch after the step, and a 1-D integer tensor of converged sample indices (or None if nothing converged).

Return type:

tuple[Batch, torch.Tensor | None]