nvalchemi-dynamics-implementation

nvalchemi Dynamics Implementation

Overview

To implement a dynamics class (integrator) in nvalchemi, subclass BaseDynamics and override two methods: pre_update() and post_update(). The base class handles the model forward pass, hook dispatch, convergence checking, and the step/run loop.

from nvalchemi.dynamics.base import BaseDynamics, ConvergenceHook
from nvalchemi.data import Batch

Step execution flow

The first step(batch) after admission dispatches ON_ADMISSION, while subsequent steps skip it until admission is explicitly reset. The per-step sequence is:

0. ON_ADMISSION hooks (once after reset, before the compiled step)
1. BEFORE_STEP hooks
2. BEFORE_PRE_UPDATE hooks  →  pre_update(batch)  →  AFTER_PRE_UPDATE hooks
3. BEFORE_COMPUTE hooks     →  compute(batch)      →  AFTER_COMPUTE hooks
4. BEFORE_POST_UPDATE hooks →  post_update(batch)  →  AFTER_POST_UPDATE hooks
5. AFTER_STEP hooks
6. Check convergence → ON_CONVERGE hooks if converged
7. ON_GRADUATE hooks (if any registered) with ctx.graduated_mask = graphs that reached exit_status
8. Increment step_count

• The base step() calls pre_update() and post_update() with autograd enabled — it does not wrap them in torch.no_grad(). Your implementation must wrap its own state updates in torch.no_grad() itself (as the example below and DemoDynamics do)
• compute() calls the model forward pass and writes forces/energy to the batch in-place
• You implement pre_update() and post_update(); everything else is inherited

Implementation guide

1. Define the class

Set __needs_keys__ (model outputs your integrator requires) and __provides_keys__ (state your integrator produces).

class MyDynamics(BaseDynamics):
    __needs_keys__: set[str] = {"forces"}
    __provides_keys__: set[str] = {"velocities", "positions"}

2. Implement __init__

Store integrator parameters. Always call super().__init__() and forward **kwargs (needed for cooperative multiple inheritance with the communication mixin).

def __init__(
    self,
    model: BaseModelMixin,
    n_steps: int,
    dt: float = 1.0,
    hooks: list[Hook] | None = None,
    convergence_hook: ConvergenceHook | dict | None = None,
    **kwargs: Any,
) -> None:
    super().__init__(
        model=model,
        hooks=hooks,
        convergence_hook=convergence_hook,
        n_steps=n_steps,
        **kwargs,
    )
    self.dt = dt

BaseDynamics constructor parameters:

3. Implement pre_update(batch)

Update positions based on current velocities and forces. Modify the batch in-place.

def pre_update(self, batch: Batch) -> None:
    positions = batch.positions       # [V, 3]
    velocities = batch.velocities     # [V, 3]
    forces = batch.forces             # [V, 3] or None
    masses = batch.atomic_masses.unsqueeze(-1)  # [V] -> [V, 1]

    with torch.no_grad():
        if forces is not None and not torch.all(forces == 0):
            accelerations = forces / masses
            # x(t+dt) = x(t) + v(t)*dt + 0.5*a(t)*dt^2
            positions.add_(velocities * self.dt + 0.5 * accelerations * self.dt**2)
        else:
            # First step fallback (no forces yet)
            positions.add_(velocities * self.dt)

4. Implement post_update(batch)

Update velocities based on new forces (computed between pre_update and post_update by the inherited compute() method). Modify the batch in-place.

def post_update(self, batch: Batch) -> None:
    velocities = batch.velocities     # [V, 3]
    forces = batch.forces             # [V, 3]
    masses = batch.atomic_masses.unsqueeze(-1)

    with torch.no_grad():
        new_accelerations = forces / masses
        # v(t+dt) = v(t) + a(t+dt)*dt
        velocities.add_(new_accelerations * self.dt)

Inherited methods (do NOT override)

Inherited attributes

Usage

from nvalchemi.models.demo import DemoModelWrapper
from nvalchemi.data import AtomicData, Batch
import torch

# Create model and dynamics
model = DemoModelWrapper()
dynamics = MyDynamics(model=model, n_steps=100, dt=0.5)

# Create batch
data = AtomicData(
    atomic_numbers=torch.tensor([6, 6, 8], dtype=torch.long),
    positions=torch.randn(3, 3),
)
batch = Batch.from_data_list([data])

# Initialize required fields (forces/energy must exist for copy_())
batch.forces = torch.zeros(3, 3)
batch.energy = torch.zeros(1, 1)

# Run
result = dynamics.run(batch)
# Or step-by-step
dynamics.step(batch)

Complete example: Velocity Verlet

This mirrors DemoDynamics, the reference implementation.

from __future__ import annotations
from typing import Any, TYPE_CHECKING
import torch
from nvalchemi.data import Batch
from nvalchemi.dynamics.base import BaseDynamics, ConvergenceHook

if TYPE_CHECKING:
    from nvalchemi.dynamics.base import Hook
    from nvalchemi.models.base import BaseModelMixin


class VelocityVerlet(BaseDynamics):
    """Velocity Verlet integrator."""

    __needs_keys__: set[str] = {"forces"}
    __provides_keys__: set[str] = {"velocities", "positions"}

    def __init__(
        self,
        model: BaseModelMixin,
        n_steps: int,
        dt: float = 1.0,
        hooks: list[Hook] | None = None,
        convergence_hook: ConvergenceHook | dict | None = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            model=model, hooks=hooks, convergence_hook=convergence_hook,
            n_steps=n_steps, **kwargs,
        )
        self.dt = dt
        self._prev_accelerations: torch.Tensor | None = None

    def pre_update(self, batch: Batch) -> None:
        """x(t+dt) = x(t) + v(t)*dt + 0.5*a(t)*dt^2"""
        positions = batch.positions
        velocities = batch.velocities
        forces = batch.forces
        masses = batch.atomic_masses.unsqueeze(-1)

        with torch.no_grad():
            if forces is not None and not torch.all(forces == 0):
                accelerations = forces / masses
                self._prev_accelerations = accelerations.clone()
                positions.add_(velocities * self.dt + 0.5 * accelerations * self.dt**2)
            else:
                positions.add_(velocities * self.dt)

    def post_update(self, batch: Batch) -> None:
        """v(t+dt) = v(t) + 0.5*(a(t) + a(t+dt))*dt"""
        velocities = batch.velocities
        forces = batch.forces
        masses = batch.atomic_masses.unsqueeze(-1)

        with torch.no_grad():
            new_accelerations = forces / masses
            if self._prev_accelerations is not None:
                velocities.add_(
                    0.5 * (self._prev_accelerations + new_accelerations) * self.dt
                )
            else:
                velocities.add_(new_accelerations * self.dt)

Convergence

Use ConvergenceHook to stop early or migrate samples in a pipeline:

from nvalchemi.dynamics.base import ConvergenceHook

hook = ConvergenceHook(
    criteria=[
        {"key": "fmax", "threshold": 0.05},
        {"key": "energy_change", "threshold": 1e-6},
    ],
    source_status=0,   # check samples with this status
    target_status=1,   # migrate converged samples to this status
    frequency=1,       # check every N steps
)

dynamics = MyDynamics(model=model, n_steps=1000, convergence_hook=hook)

Composition with FusedStage

Chain multiple dynamics stages that share a single model forward pass:

relax = MyDynamics(model, n_steps=100, dt=0.5)
md = MyDynamics(model, n_steps=500, dt=0.1)

# Compose with + operator
fused = relax + md

# Samples start in relax, converge, then move to md
fused.run(batch)

Distributed pipeline

Chain stages across ranks with the | operator:

opt_stage = MyDynamics(model, n_steps=100, dt=0.5)   # rank 0
md_stage = MyDynamics(model, n_steps=500, dt=0.1)    # rank 1

pipeline = opt_stage | md_stage

with pipeline:
    pipeline.run()