nvalchemi.dynamics.DemoDynamics

class nvalchemi.dynamics.DemoDynamics(model, n_steps, dt=1.0, hooks=None, convergence_hook=None, **kwargs)

Velocity Verlet integrator for molecular dynamics simulations.

Implements the standard Velocity Verlet algorithm, a symplectic, time-reversible integration scheme commonly used in molecular dynamics. The algorithm splits the integration into two half-steps:

1. pre_update: Update positions using current velocities and accelerations x(t+dt) = x(t) + v(t)*dt + 0.5*a(t)*dt^2

2. post_update: Update velocities using averaged accelerations v(t+dt) = v(t) + 0.5*(a(t) + a(t+dt))*dt

On the first step when previous accelerations are unavailable, falls back to Euler integration for velocities: v(t+dt) = v(t) + a(t+dt)*dt.

Inter-rank communication capabilities for use in pipeline workflows are inherited from BaseDynamics. Communication attributes include prior_rank, next_rank, sinks, active_batch, max_batch_size, and done.

This class is intended entirely for testing and debugging, demonstrating how to implement an integrator by overriding pre_update and post_update. Do NOT use this class for production.

Parameters:

• model (BaseModelMixin)

• n_steps (int)

• dt (float)

• hooks (list[Hook] | None)

• convergence_hook (ConvergenceHook | dict | None)

• kwargs (Any)

__needs_keys__

Set of output keys required from the model. Set to {"forces"}.

Type:

set[str]

__provides_keys__

Set of keys this dynamics produces beyond model outputs. Set to {"velocities", "positions"}.

Type:

set[str]

model

The neural network potential model.

Type:

BaseModelMixin

dt

The integration timestep.

Type:

float

step_count

The current step number.

Type:

int

hooks

Registered hooks organized by stage.

Type:

dict[HookStageEnum, list[Hook]]

_prev_accelerations

Cached accelerations from the previous step for the velocity half-step. None on the first step.

Type:

torch.Tensor | None

prior_rank

Rank of the previous pipeline stage (inherited from BaseDynamics).

Type:

int | None

next_rank

Rank of the next pipeline stage (inherited from BaseDynamics).

Type:

int | None

Examples

>>> model = DemoModelWrapper()
>>> dynamics = DemoDynamics(model, dt=0.5, n_steps=100)
>>> dynamics.run(batch)
>>> # DistributedPipeline composition:
>>> pipeline = dynamics | other_dynamics

__init__(model, n_steps, dt=1.0, hooks=None, convergence_hook=None, **kwargs)

Initialize the Velocity Verlet integrator.

Parameters:

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

• n_steps (int) – Total number of simulation steps for run().

• dt (float, optional) – The integration timestep in femtoseconds. Default is 1.0.

• 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) – Convergence hook configuration. Accepts a ConvergenceHook instance, a dict (unpacked as ConvergenceHook(**dict)), or None for no convergence checking. Default is None.

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

Return type:

None

Methods

__init__(model, n_steps[, dt, hooks, ...])	Initialize the Velocity Verlet integrator.
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 velocity update (second half of Velocity Verlet).
pre_update(batch)	Perform the position update (first half of Velocity Verlet).
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.