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])

• 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.

Type:

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

post_update(batch)

Perform the velocity update (second half of Velocity Verlet).

Updates velocities according to: v(t+dt) = v(t) + 0.5*(a(t) + a(t+dt))*dt

where a(t+dt) = F(t+dt) / m are the forces computed after the position update.

If previous accelerations are unavailable (first step), falls back to Euler velocity update: v(t+dt) = v(t) + a(t+dt)*dt.

The update is performed inside a torch.no_grad() context to avoid conflicts with autograd when forces_via_autograd=True.

Parameters:

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

Return type:

None

pre_update(batch)

Perform the position update (first half of Velocity Verlet).

Updates positions according to: x(t+dt) = x(t) + v(t)*dt + 0.5*a(t)*dt^2

where a(t) = F(t) / m.

If forces are not yet computed (first call before any compute), this method falls back to a simple Euler position update: x(t+dt) = x(t) + v(t)*dt.

The update is performed inside a torch.no_grad() context to avoid conflicts with autograd when forces_via_autograd=True (which causes compute() to set requires_grad_(True) on positions).

Parameters:

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

Return type:

None