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