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

Each call to step(batch) executes:

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