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 _CommunicationMixin, which provides 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] | None)

• 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

Dictionary mapping each stage to a list of registered hooks.

Type:

dict[HookStageEnum, 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 fmax 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)

__init__(model, hooks=None, convergence_hook=None, n_steps=None, exit_status=1, **kwargs)

Initialize the dynamics engine.

Parameters:

• model (BaseModelMixin) – The neural network potential model.

• hooks (list[Hook] | None, optional) – Initial list of hooks to register. Each hook will be organized by its stage attribute.

• convergence_hook (ConvergenceHook | dict | None, optional) – Hook that evaluates composable convergence criteria.

• n_steps (int | None, optional) – Total number of simulation steps. If provided, run() will use this value when called without an explicit n_steps argument. Default is None. If a dict is provided, it is unpacked as ConvergenceHook(**convergence_hook). If None, no convergence will be assessed.

• exit_status (int, optional) – 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. Subclasses like FusedStage may compute this dynamically.

• **kwargs (Any) – Additional keyword arguments forwarded to the next class in the MRO (for cooperative multiple inheritance).

Return type:

None

Methods

__init__(model[, hooks, convergence_hook, ...])	Initialize the dynamics engine.
compute(batch)	Perform the model forward pass to compute forces and energies.
masked_update(batch, mask)	Apply pre_update and post_update only to selected samples in the batch.
post_update(batch)	Perform the second half of the integration step.
pre_update(batch)	Perform the first half of the integration step.
refill_check(batch, exit_status)	Replace graduated samples via index-select and append.
register_bookkeeping_key(key, init_fn)	Register a graph-level bookkeeping field to survive refill_check.
register_hook(hook)	Register a hook to be executed at its designated stage(s).
run(batch[, n_steps])	Run the dynamics simulation for a specified number of steps.
step(batch)	Execute a single dynamics step with the full hook-wrapped sequence.

Attributes

active_batch_has_room	Return whether the active batch can accept more samples.
active_batch_size	Return the number of samples currently in the active batch.
device	Compute the torch device for this rank.
global_rank	Get the global rank for this process.
has_neighbor	Convenient property to see if rank is isolated
inflight_mode	Return whether inflight batching is enabled.
is_final_stage	Return whether this is the last stage in the pipeline.
is_first_stage	Return whether this is the first stage in the pipeline.
local_rank	Get the node-local rank for this process.
model_is_conservative	Returns whether or not the model uses conservative forces
room_in_active_batch	Return the number of additional samples the active batch can hold.
stream	Return the active CUDA stream, if any.