Optimization and Integrators#

This page covers the concrete simulation types provided by the dynamics module. All of them follow the execution loop described in the dynamics overview — they generally differ only in what pre_update and post_update do.

Geometry optimization#

Geometry optimization finds the nearest local energy minimum by iteratively moving atoms downhill on the potential energy surface. The toolkit provides the FIRE (Fast Inertial Relaxation Engine) algorithm in two variants.

Fixed-cell optimization#

FIRE optimizes atomic positions while keeping the simulation cell fixed:

from nvalchemi.dynamics import FIRE, ConvergenceHook

with FIRE(
    model=model,
    dt=0.1,           # initial timestep (femtoseconds)
    n_steps=500,
    hooks=[ConvergenceHook(fmax=0.05)],
) as opt:
    relaxed = opt.run(batch)

FIRE uses an adaptive timestep and velocity mixing: when the system is moving downhill (forces aligned with velocities), the timestep grows and velocities are biased toward the force direction. When the system overshoots, the timestep shrinks and velocities are zeroed. This makes it robust across a wide range of systems without manual tuning.

Variable-cell optimization#

FIREVariableCell extends FIRE to simultaneously optimize both atomic positions and the simulation cell. This is useful for finding equilibrium crystal structures where the lattice parameters are not known a priori:

from nvalchemi.dynamics.optimizers.fire import FIREVariableCell
from nvalchemi.dynamics import ConvergenceHook

with FIREVariableCell(
    model=model,
    dt=0.1,
    n_steps=500,
    hooks=[ConvergenceHook(fmax=0.05)],
) as opt:
    relaxed = opt.run(batch)

The cell degrees of freedom are propagated using an NPH-like scheme at zero target pressure. The model must return stresses (or virials) in addition to forces.

Choosing between fixed and variable cell#

Use fixed-cell FIRE when the cell is known (e.g. a bulk crystal at experimental lattice parameters, or a molecule in vacuum where the cell is just a bounding box). Use variable-cell FIRE when the equilibrium cell shape or volume is unknown, such as when screening candidate crystal structures or computing equations of state.

Molecular dynamics#

Molecular dynamics (MD) propagates the equations of motion forward in time, sampling the trajectory of a system at finite temperature. The toolkit provides integrators for three standard ensembles.

NVE: energy conservation#

NVE uses the Velocity Verlet algorithm — a symplectic integrator that conserves total energy in the microcanonical ensemble:

from nvalchemi.dynamics import NVE

with NVE(model=model, dt=1.0, n_steps=1000) as md:
    trajectory = md.run(batch)

NVE is the natural choice for verifying that a model’s energy surface is smooth enough for stable dynamics: if the total energy drifts significantly, the force field is likely too noisy for the chosen timestep.

NVT: constant temperature#

NVTLangevin implements the BAOAB Langevin splitting scheme, which samples the canonical (NVT) ensemble exactly — the thermostat does not introduce systematic bias:

from nvalchemi.dynamics import NVTLangevin

with NVTLangevin(
    model=model,
    dt=1.0,              # femtoseconds
    temperature=300.0,    # Kelvin
    friction=0.01,        # collision frequency (1/fs)
    n_steps=10000,
) as md:
    trajectory = md.run(batch)

The friction parameter controls how strongly the thermostat couples to the system. A low value gives longer correlation times (closer to NVE); a high value thermalises quickly but damps real dynamics.

NPT: constant pressure#

NPT uses the Martyna–Tobias–Klein (MTK) barostat with Nose–Hoover chains to sample the isothermal-isobaric ensemble. Both the atomic positions and the simulation cell evolve:

from nvalchemi.dynamics import NPT

with NPT(
    model=model,
    dt=1.0,
    temperature=300.0,
    pressure=1.0,            # target pressure (bar)
    n_steps=10000,
) as md:
    trajectory = md.run(batch)

The model must return stresses for NPT to propagate the cell degrees of freedom.

Writing your own dynamics#

All integrators and optimizers inherit from BaseDynamics. To implement a custom one, you subclass it and override pre_update and post_update — the two methods that define how the batch state evolves within a single step.

The minimal contract#

Your subclass must provide:

__needs_keys__ — a set of strings naming the model outputs your dynamics reads (e.g. {"forces"}, or {"forces", "stresses"} for cell-aware schemes).
__provides_keys__ — a set of strings naming the batch keys your dynamics writes (e.g. {"positions", "velocities"}).
pre_update(batch) — called before the model forward pass. Typically updates positions using current velocities and/or forces.
post_update(batch) — called after the model forward pass. Typically completes the velocity update with the newly computed forces.

Both methods receive the Batch and modify it in-place. Return value is None.

Example: a Velocity Verlet integrator#

The DemoDynamics class in nvalchemi.dynamics.demo is a complete, minimal Velocity Verlet implementation that is useful as a template:

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

class MyVerlet(BaseDynamics):
    __needs_keys__ = {"forces"}
    __provides_keys__ = {"positions", "velocities"}

    def __init__(self, model, n_steps, dt=1.0, hooks=None, convergence_hook=None, **kwargs):
        super().__init__(
            model=model, hooks=hooks, convergence_hook=convergence_hook,
            n_steps=n_steps, **kwargs,
        )
        self.dt = dt
        self._prev_accelerations = None

    def pre_update(self, batch: Batch) -> None:
        """Position half-step: x(t+dt) = x(t) + v*dt + 0.5*a*dt^2."""
        import torch
        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):
                acc = forces / masses
                self._prev_accelerations = acc.clone()
                positions.add_(velocities * self.dt + 0.5 * acc * self.dt**2)
            else:
                positions.add_(velocities * self.dt)

    def post_update(self, batch: Batch) -> None:
        """Velocity half-step: v(t+dt) = v(t) + 0.5*(a_old + a_new)*dt."""
        import torch
        velocities = batch.velocities
        forces = batch.forces
        masses = batch.atomic_masses.unsqueeze(-1)

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

Important

The demo VelocityVerlet class is intended for debugging and pedagogy only. Do not use this class for production runs, and instead, see the NVE class instead.

Data flow through a step#

Understanding what the batch contains at each point is key to writing correct updates:

Point in step	What just happened	What the batch contains
`pre_update` entry	Hooks ran	Positions and velocities from the previous step; forces may be from the previous `compute` (or absent on step 0)
`pre_update` exit	You updated positions	New positions; velocities partially updated (or unchanged)
After `compute`	Model ran	Fresh `forces` (and `energies`, `stresses`, etc.) for the new positions
`post_update` entry	Forces are fresh	Complete the velocity update with new forces
`post_update` exit	Step is done	Consistent positions, velocities, and forces for the current timestep

Gotchas and tips#

Use torch.no_grad(): Wrap in-place updates in torch.no_grad() to avoid conflicts with autograd. When forces_via_autograd=True, compute() sets requires_grad_(True) on positions to compute forces via backprop.
In-place operations: Modify batch tensors in-place (positions.add_(...)) rather than reassigning. The batch’s storage model expects tensors to be updated in place.
First-step fallback: On the first call to pre_update, forces may be None or zero (no model evaluation has happened yet). Guard against this and fall back to an Euler step.
Per-system state: If your integrator needs auxiliary state (e.g. thermostat variables, previous accelerations), store it as instance attributes. The _prev_accelerations pattern above is typical.
__needs_keys__ matters: BaseDynamics uses this set to verify the model produces the required outputs before the simulation starts. If your dynamics needs stresses, declare {"forces", "stresses"}.
FusedStage compatibility: When your dynamics runs inside a FusedStage, a save-and-restore mask is applied around pre_update and post_update so that only systems belonging to your stage are modified. You do not need to handle masking yourself.