Models: Wrapping ML Interatomic Potentials

The ALCHEMI Toolkit uses a standardized interface — BaseModelMixin — that sits between your PyTorch model and the rest of the framework (dynamics, data loading, active learning). Any machine-learning interatomic potential (MLIP) can be used with the toolkit as long as it is wrapped with this interface.

Tip

AI coding assistant? Load the nvalchemi-model-wrapping agent skill for concise instructions on wrapping an arbitrary MLIP with the BaseModelMixin interface.

This guide covers:

1. What models are currently supported out of the box.

2. The two building blocks: ModelConfig and BaseModelMixin.

3. How to wrap your own model, using DemoModelWrapper as a worked example.

4. How to compose multiple models using the + operator or the explicit PipelineModelWrapper API.

Supported models

The nvalchemi.models package ships wrappers for the following potentials:

Wrapper class	Underlying model	Notes
DemoModelWrapper	DemoModel	Non-invariant demo; useful for testing and tutorials
AIMNet2Wrapper	AIMNet2Calculator	Requires the aimnet2 optional dependency
MACEWrapper	Any MACE variant	Requires the mace optional dependency with a CUDA extra, such as cu13 or cu12
UMAWrapper	fairchem-core UMA (MLIPPredictUnit)	Requires the uma optional dependency; conflicts with mace (incompatible e3nn pins)

AIMNet2Wrapper, MACEWrapper, and UMAWrapper are lazily imported — they only load when accessed, so missing dependencies will not break other imports.

Note

UMA resolves to a different torch than the cu12 / cu13 GPU stack. fairchem-core caps torch below 2.9 (fairchem-core>=2.8 requires torch>=2.8,<2.9), while the cu12 / cu13 extras pull nvalchemi-toolkit-ops[torch-cuXX], which floors torch at >=2.11. The uma extra is therefore declared mutually exclusive with cu12, cu13, and mace, and uv sync --extra uma forks a standalone resolution that installs a PyPI CUDA torch wheel (~2.8) instead of the NVIDIA-indexed cuXX build. Keep UMA in its own environment, e.g.:

uv venv .venv-uma && uv sync --extra uma           # UMA (fairchem's torch)
uv venv .venv-mace && uv sync --extra cu13 --extra mace   # MACE on the cu13 GPU stack

The core nvalchemi-toolkit-ops package (and its Warp kernels) is still installed in the UMA environment — only the cuXX GPU-acceleration extras (cuEquivariance, cuML, the NVIDIA-indexed torch build) are dropped, none of which UMA uses, since it builds its neighbor graph inside fairchem. The toolkit-ops torch-cuXX extras pin torch>=2.11, but that tracks the cuXX wheel builds rather than a toolkit-ops API requirement — the base package already declares torch>=2.8, so the Warp path is expected to work against the ~2.8 torch in the UMA environment.

MACE checkpoints in training

When starting from an existing MACE checkpoint, prefer from_checkpoint() over manually loading the underlying MACE module and wrapping it. The wrapper records a factory-based reconstruction spec that strategy checkpoints can use later. This matters for optimized variants such as cuEquivariance, where the live transformed module is not reliably reconstructible from its Python constructor.

import torch

from nvalchemi.models.mace import MACEWrapper
from nvalchemi.training import EMAHook, TrainingStrategy

model = MACEWrapper.from_checkpoint(
    "small-0b",
    device=torch.device("cuda"),
    dtype=torch.float32,
    enable_cueq=True,
)

ema = EMAHook(model_key="main", decay=0.999)
strategy = TrainingStrategy(
    models=model,
    ...,
    hooks=[ema],
)
strategy.save_checkpoint(checkpoint_dir)

restored_ema = EMAHook(model_key="main", decay=0.999)
restored = TrainingStrategy.load_checkpoint(
    checkpoint_dir,
    map_location=torch.device("cuda"),
    hooks=[restored_ema],
    training_fn=training_fn,
)

Avoid saving only ema.state_dict() for MACE training restarts. Strategy checkpoints preserve the model reconstruction recipe, model weights, optimizer state, runtime counters, and checkpointable hook state together.

Repairing methods on EMA copies

EMAHook uses torch.optim.swa_utils.AveragedModel, which deep-copies the live model when constructing its averaged copy. Most wrappers require no special handling. If a third-party model attaches runtime methods that do not survive deepcopy, its wrapper can optionally implement modify_ema_methods():

class MyPotentialWrapper(nn.Module, BaseModelMixin):
    def modify_ema_methods(self) -> None:
        """Restore runtime methods on this EMA model copy."""
        my_method_to_patch_runtime_methods()

The hook calls this method once on the copied wrapper immediately after EMA construction and before applying any pending EMA checkpoint weights. The method must modify only self, must not change parameters or buffers, and should be safe when called on a freshly copied model.

Tip

MACEWrapper implements this interface to restore cuEquivariance’s fused convolution method. Take a look at the MACEWrapper.modify_ema_methods to see how this is done.

Using UMA (fairchem-core)

UMA (Universal Models for Atoms) is a multi-task foundation model: one checkpoint ships task heads for molecules (omol), bulk crystals (omat), catalysis (oc20), direct air capture (odac), and molecular crystals (omc). UMAWrapper pins a single task at construction; active_outputs is {energy, forces} for molecular tasks and {energy, forces, stress} for periodic ones.

1. Install the optional dependency (in its own environment, per the note above):

uv venv .venv-uma && uv sync --extra uma
# or, with pip:  pip install 'nvalchemi-toolkit[uma]'

2. Get HuggingFace access. UMA checkpoints live in the gated facebook/UMA repository, so a (free) HuggingFace account and a one-time access approval are required:

1. Sign in and click “Agree and access repository” on the model page.

2. Create a read token at https://huggingface.co/settings/tokens.

3. Make the token available to your shell, either by logging in once (it is cached under ~/.cache/huggingface):

   huggingface-cli login
   

   or by exporting it for the session:

   export HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxx
   

3. Load a checkpoint. The first call downloads and caches the weights (under ~/.cache/fairchem); later calls reuse the cache and need no network:

from nvalchemi.models.uma import UMAWrapper

# Molecular potential (OMol head)
mol = UMAWrapper.from_checkpoint("uma-s-1p1", task_name="omol", device="cuda")

# Bulk-crystal potential (OMat head, same checkpoint family)
mat = UMAWrapper.from_checkpoint("uma-m-1p1", task_name="omat", device="cuda")

Registered checkpoint names (see fairchem.core.calculate.pretrained_mlip.available_models for the full list):

Checkpoint	Size	Notes
uma-s-1p1	small	Default in the examples / tests
uma-s-1p2	small	Updated small release
uma-m-1p1	medium	Higher accuracy, larger / slower

torch.compile / turbo. UMAWrapper does not add a compile_model flag (unlike the MACE / AIMNet2 wrappers) because fairchem owns compilation internally as a field on its InferenceSettings. Reach it through from_checkpoint’s inference_settings argument — pass "turbo" for fairchem’s compiled preset (torch.compile + TF32 + MoLE merge, for runs with fixed atomic composition), or an InferenceSettings instance for finer control:

fast = UMAWrapper.from_checkpoint(
    "uma-s-1p1", task_name="omat", device="cuda", inference_settings="turbo"
)

See the UMA NVE/NVT example for a runnable end-to-end molecular-dynamics walkthrough.

Architecture overview

A wrapped model uses multiple inheritance: your existing Module subclass provides the forward pass, while BaseModelMixin adds the standardized interface.

Multiple-inheritance pattern for model wrapping.

The wrapper’s forward method follows a three-step pipeline:

1. adapt_input — convert AtomicData / Batch into the keyword arguments your model expects.

2. super().forward — call the underlying model unchanged.

3. adapt_output — map raw model outputs to the framework’s ModelOutputs ordered dictionary.

ModelConfig: capability declaration and runtime control

ModelConfig is a single Pydantic model that serves two purposes:

1. Capability fields (frozen at construction) describe what the model checkpoint can do. These use frozenset to signal immutability.

2. Runtime fields (mutable) control what the model computes on each forward pass. These can be changed freely at any time.

Every wrapper sets self.model_config in its __init__. The config uses free-form strings for outputs and inputs, so new properties (e.g. "magnetic_moment", "charges") can be added without modifying the schema.

Capability fields (frozen)

Field	Default	Meaning
outputs	frozenset({"energy"})	All property names the model can produce. Well-known keys: energy, forces, stress, hessian, dipole, charges.
autograd_outputs	frozenset()	Subset of outputs computed via autograd (e.g. {"forces"} for conservative MLIP forces). Empty for analytical-force models.
autograd_inputs	frozenset({"positions"})	Input keys that need requires_grad_(True) when any autograd output is requested.
required_inputs	frozenset()	Extra inputs beyond {positions, atomic_numbers} that the model requires (error if missing). Neighbor-list keys are auto-derived from neighbor_config.
optional_inputs	frozenset()	Extra inputs the model can optionally use if present, silently skipped if absent.
supports_pbc	False	Model handles periodic boundary conditions.
needs_pbc	False	Model requires pbc and cell in its input.
neighbor_config	None	NeighborConfig describing neighbor list requirements, or None if the model does not use a neighbor list.

Runtime fields (mutable)

Field	Default	Meaning
active_outputs	None (defaults to outputs)	Set of property names to compute this run. Change this to narrow or expand what the model computes.
gradient_keys	set()	Additional tensor keys that need requires_grad_(True) beyond those implied by autograd_inputs.

The method output_data() intersects active_outputs with outputs and warns if any requested keys are unsupported.

from nvalchemi.models.base import ModelConfig, NeighborConfig

# An autograd-forces MLIP with PBC support
cfg = ModelConfig(
    outputs={"energy", "forces", "stress"},
    autograd_outputs={"forces", "stress"},
    supports_pbc=True,
    needs_pbc=False,
    neighbor_config=NeighborConfig(cutoff=5.0, format="coo"),
)

# An analytical-forces model (e.g. Lennard-Jones)
cfg = ModelConfig(
    outputs={"energy", "forces", "stress"},
    autograd_outputs=set(),   # forces computed by kernel, not autograd
    supports_pbc=True,
    needs_pbc=False,
    neighbor_config=NeighborConfig(cutoff=8.5, format="matrix"),
)

# A model that requires charges as input (e.g. Ewald)
cfg = ModelConfig(
    outputs={"energy", "forces", "stress"},
    required_inputs={"charges"},
    needs_pbc=True,
    supports_pbc=True,
    neighbor_config=NeighborConfig(cutoff=10.0, format="matrix"),
)

# A model with optional inputs (e.g. AIMNet2 — works with or without PBC)
cfg = ModelConfig(
    outputs={"energy", "forces", "charges"},
    autograd_outputs={"forces"},
    required_inputs={"charge"},       # system charge is required
    optional_inputs={"cell", "mult"}, # PBC cell and multiplicity are optional
)

Changing active outputs at runtime

The active_outputs field is the primary lever for controlling what a model computes on each forward pass. It defaults to outputs (i.e. compute everything the model supports), but you can narrow or expand it at any time:

# Start with full computation
model = MyWrapper()
out = model(batch)  # computes energies + forces (the defaults)

# Switch to energy-only evaluation (faster — skips force computation)
model.model_config.active_outputs = {"energy"}
out = model(batch)  # only energies

# Enable stress computation for NPT dynamics
model.model_config.active_outputs = {"energy", "forces", "stress"}
out = model(batch)  # energies + forces + stresses

# Restore defaults (compute everything the model supports)
model.model_config.active_outputs = set(model.model_config.outputs)

This is particularly useful in multi-stage workflows: use energy-only evaluation during screening, then switch to forces + stresses for production dynamics.

Wrapping your own model: step by step

This section walks through every method you need to implement, using DemoModelWrapper as the running example.

Required interface checklist

Your wrapper class must provide the following. Methods marked abstract will raise TypeError at instantiation if missing:

Method / Property	Abstract?	Classical potential stub
model_config attribute	— (enforced by post-init check)	Set self.model_config = ModelConfig(...) in __init__
embedding_shapes (property)	Yes	return {}
compute_embeddings()	Yes	raise NotImplementedError
adapt_input()	No (has default)	Override to collect model-specific inputs
adapt_output()	No (has default)	Override to map raw outputs
forward()	No (inherit from nn.Module)	Implement the three-step pipeline
export_model()	No (base default raises NotImplementedError)	Override to enable export

For classical potentials with no learned embeddings, stub both embedding methods:

@property
def embedding_shapes(self) -> dict[str, tuple[int, ...]]:
    return {}

def compute_embeddings(self, data, **kwargs):
    raise NotImplementedError("No embeddings for this model.")

Step 1 — Create the wrapper class

Subclass Module and mix in BaseModelMixin, then hold the underlying model as self.model:

from torch import nn
from nvalchemi.models.base import BaseModelMixin, ModelConfig

class DemoModelWrapper(nn.Module, BaseModelMixin):
    def __init__(self, model: DemoModel) -> None:
        super().__init__()
        self.model: DemoModel = model
        ...

Step 2 — Set model_config in __init__

Create a ModelConfig describing your model’s capabilities and set it as self.model_config in __init__:

def __init__(self, model: DemoModel) -> None:
    super().__init__()
    self.model = model
    self.model_config = ModelConfig(
        outputs={"energy", "forces"},
        autograd_outputs={"forces"},
        needs_pbc=False,
    )

Important

Always set model_config as an instance attribute in __init__. There is intentionally no class-level default — a shared class attribute would cause mutations in one wrapper to silently affect all others.

Step 3 — Implement embedding_shapes

Return a dictionary mapping embedding names to their trailing shapes. This is used by downstream consumers (e.g. active learning) to know what representations the model can provide:

@property
def embedding_shapes(self) -> dict[str, tuple[int, ...]]:
    return {
        "node_embeddings": (self.hidden_dim,),
        "graph_embedding": (self.hidden_dim,),
    }

Step 4 — Implement adapt_input

Convert framework data to the keyword arguments your underlying model’s forward() expects. Always call super().adapt_input() first — the base implementation enables gradients on the required tensors (using autograd_inputs and autograd_outputs from the model card) and validates that all required input keys are present:

def adapt_input(self, data: AtomicData | Batch, **kwargs) -> dict[str, Any]:
    model_inputs = super().adapt_input(data, **kwargs)

    # Extract tensors in the format your model expects
    model_inputs["atomic_numbers"] = data.atomic_numbers
    model_inputs["positions"] = data.positions.to(self.dtype)

    # Handle batched vs. single input
    if isinstance(data, Batch):
        model_inputs["batch_indices"] = data.batch_idx
    else:
        model_inputs["batch_indices"] = None

    # Pass config flags to control model behavior
    model_inputs["compute_forces"] = "forces" in self.model_config.active_outputs
    return model_inputs

Step 5 — Implement adapt_output

Map the model’s raw output dictionary to ModelOutputs, an OrderedDict[str, Tensor | None] with standardized keys. Always call super().adapt_output() first — it creates the OrderedDict pre-filled with expected keys (derived from the intersection of model_config.active_outputs and model_config.outputs) and auto-maps any keys whose names already match:

def adapt_output(self, model_output, data: AtomicData | Batch) -> ModelOutputs:
    output = super().adapt_output(model_output, data)

    energy = model_output["energy"]
    if isinstance(data, AtomicData) and energy.ndim == 1:
        energy = energy.unsqueeze(-1)  # must be [B, 1]
    output["energy"] = energy

    if "forces" in self.model_config.active_outputs:
        output["forces"] = model_output["forces"]

    # Validate: no expected key should be None
    for key, value in output.items():
        if value is None:
            raise KeyError(
                f"Key '{key}' not found in model output "
                "but is supported and requested."
            )
    return output

The standard output shapes are:

Key	Shape	Description
energy	[B, 1]	Per-graph total energy
forces	[V, 3]	Per-atom forces
stress	[B, 3, 3]	Per-graph stress tensor
hessian	[V, 3, 3]	Per-atom Hessian
dipole	[B, 3]	Per-graph dipole moment
charges	[V]	Per-atom partial charges

Step 6 — Implement compute_embeddings

This method is abstract — you must implement it even if your model has no learned embeddings. For classical potentials, a one-line stub suffices:

def compute_embeddings(self, data, **kwargs):
    raise NotImplementedError("No embeddings for this model.")

For learned models, extract intermediate representations and write them to the data structure in-place. This is used by active learning and other downstream consumers:

def compute_embeddings(self, data: AtomicData | Batch, **kwargs) -> AtomicData | Batch:
    model_inputs = self.adapt_input(data, **kwargs)

    # Run the model's internal layers
    atom_z = self.embedding(model_inputs["atomic_numbers"])
    coord_z = self.coord_embedding(model_inputs["positions"])
    embedding = self.joint_mlp(torch.cat([atom_z, coord_z], dim=-1))
    embedding = embedding + atom_z + coord_z

    # Aggregate to graph level via scatter
    if isinstance(data, Batch):
        batch_indices = data.batch_idx
        num_graphs = data.batch_size
    else:
        batch_indices = torch.zeros_like(model_inputs["atomic_numbers"])
        num_graphs = 1

    graph_shape = self.embedding_shapes["graph_embedding"]
    graph_embedding = torch.zeros(
        (num_graphs, *graph_shape),
        device=embedding.device,
        dtype=embedding.dtype,
    )
    graph_embedding.scatter_add_(0, batch_indices.unsqueeze(-1), embedding)

    # Write in-place
    data.node_embeddings = embedding
    data.graph_embeddings = graph_embedding
    return data

Step 7 — Implement forward

Wire the three-step pipeline together:

def forward(self, data: AtomicData | Batch, **kwargs) -> ModelOutputs:
    model_inputs = self.adapt_input(data, **kwargs)
    model_outputs = self.model(**model_inputs)
    return self.adapt_output(model_outputs, data)

self.model(**model_inputs) calls the underlying DemoModel.forward with the unpacked keyword arguments — your original model is never modified. For additional flair, the @beartype.beartype decorator can be applied to the forward method, which will provide runtime type checking on the inputs and outputs, as well as shape checking.

Step 8 (optional) — Implement export_model

Export the model without the BaseModelMixin interface, for use with external tools (e.g. ASE calculators):

def export_model(self, path: Path, as_state_dict: bool = False) -> None:
    base_cls = self.__class__.__mro__[1]  # the original nn.Module
    base_model = base_cls()
    for name, module in self.named_children():
        setattr(base_model, name, module)
    if as_state_dict:
        torch.save(base_model.state_dict(), path)
    else:
        torch.save(base_model, path)

Putting it all together

A complete minimal wrapper for a custom potential:

import torch
from torch import nn
from typing import Any
from pathlib import Path

from nvalchemi.data import AtomicData, Batch
from nvalchemi.models.base import BaseModelMixin, ModelConfig
from nvalchemi._typing import ModelOutputs


class MyPotential(nn.Module):
    """Your existing PyTorch MLIP."""

    def __init__(self, hidden_dim: int = 128):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.encoder = nn.Linear(3, hidden_dim)
        self.energy_head = nn.Linear(hidden_dim, 1)

    def forward(self, positions, batch_indices=None, **kwargs):
        h = self.encoder(positions)
        node_energy = self.energy_head(h)
        if batch_indices is not None:
            num_graphs = batch_indices.max() + 1
            energy = torch.zeros(num_graphs, 1, device=h.device, dtype=h.dtype)
            energy.scatter_add_(0, batch_indices.unsqueeze(-1), node_energy)
        else:
            energy = node_energy.sum(dim=0, keepdim=True)
        return {"energy": energy}


class MyPotentialWrapper(MyPotential, BaseModelMixin):
    """Wrapped version for use in nvalchemi."""

    def __init__(self, hidden_dim: int = 128):
        super().__init__(hidden_dim=hidden_dim)
        self.model_config = ModelConfig(
            outputs={"energy", "forces"},
            autograd_outputs={"forces"},
            needs_pbc=False,
        )

    @property
    def embedding_shapes(self) -> dict[str, tuple[int, ...]]:
        return {"node_embeddings": (self.hidden_dim,)}

    def adapt_input(self, data: AtomicData | Batch, **kwargs: Any) -> dict[str, Any]:
        model_inputs = super().adapt_input(data, **kwargs)
        model_inputs["positions"] = data.positions
        model_inputs["batch_indices"] = data.batch_idx if isinstance(data, Batch) else None
        return model_inputs

    def adapt_output(self, model_output: Any, data: AtomicData | Batch) -> ModelOutputs:
        output = super().adapt_output(model_output, data)
        output["energy"] = model_output["energy"]
        if "forces" in self.model_config.active_outputs:
            output["forces"] = -torch.autograd.grad(
                model_output["energy"],
                data.positions,
                grad_outputs=torch.ones_like(model_output["energy"]),
                create_graph=self.training,
            )[0]
        return output

    def compute_embeddings(self, data: AtomicData | Batch, **kwargs) -> AtomicData | Batch:
        model_inputs = self.adapt_input(data, **kwargs)
        data.node_embeddings = self.encoder(model_inputs["positions"])
        return data

    def forward(self, data: AtomicData | Batch, **kwargs: Any) -> ModelOutputs:
        model_inputs = self.adapt_input(data, **kwargs)
        model_outputs = super().forward(**model_inputs)
        return self.adapt_output(model_outputs, data)

Usage:

model = MyPotentialWrapper(hidden_dim=128)

data = AtomicData(
    positions=torch.randn(5, 3),
    atomic_numbers=torch.tensor([6, 6, 8, 1, 1], dtype=torch.long),
)
batch = Batch.from_data_list([data])
outputs = model(batch)
# outputs["energy"] shape: [1, 1]
# outputs["forces"] shape: [5, 3]

Composing multiple models

nvalchemi provides three tiers of model composition, from simplest to most powerful. Choose the simplest tier that fits your use case.

Tier 1: The + operator (independent additive sum)

The + operator is the simplest way to combine models whose outputs should be summed element-wise. Each model computes its own forces independently (analytically or via its own internal autograd) and the pipeline sums energies, forces, and stresses across all models:

from nvalchemi.models.lj import LennardJonesModelWrapper
from nvalchemi.models.ewald import EwaldModelWrapper

lj = LennardJonesModelWrapper(epsilon=0.05, sigma=2.5, cutoff=8.0)
ewald = EwaldModelWrapper(cutoff=8.0)

combined = lj + ewald            # sums energies, forces, stresses
# With more models:
# combined = model_a + model_b + model_c  # chains naturally (3 groups)

The result is a PipelineModelWrapper where each model occupies its own group with use_autograd=False. Use this when:

• Each model computes its outputs independently (no data flows between them).

• Each model handles its own force computation (analytical kernels or self-contained autograd).

• You just want to sum energies, forces, and stresses.

The + operator does not support:

• Wiring one model’s output into another’s input (e.g. charges -> electrostatics).

• Shared autograd groups (differentiating the summed energy of multiple models).

For those cases, use the explicit pipeline API (Tier 2).

Tier 2: Explicit PipelineModelWrapper (dependent pipelines & shared autograd)

PipelineModelWrapper gives full control over derivative strategy, inter-model data wiring, and autograd scope. Models are organized into groups, where each group is a mini-pipeline with its own derivative computation strategy:

• use_autograd=False (default) — each model computes its own outputs. The group sums them.

• use_autograd=True — the group sums all model energies, then computes derivatives (forces, stresses, etc.) from the summed energy. This is required when one model’s output feeds into another’s energy computation and forces must backpropagate through the full chain.

from nvalchemi.models.pipeline import (
    PipelineModelWrapper, PipelineGroup, PipelineStep,
)

# AIMNet2 predicts charges + energy; Ewald uses those charges.
# Both use the key "charges" — auto-wired, no explicit mapping needed.
# Forces must backpropagate through both → shared autograd.
pipe = PipelineModelWrapper(groups=[
    PipelineGroup(
        steps=[aimnet2, ewald],
        use_autograd=True,
    ),
    PipelineGroup(steps=[dftd3]),
])

Key concepts:

• PipelineStep(model, wire={...}) — wraps a model with an output rename mapping. Only needed when a model’s output key doesn’t match the downstream input key. For models that don’t need renaming (like AIMNet2 + Ewald above where both use "charges"), pass the bare model directly.

• PipelineGroup(steps=[...], use_autograd=True|False) — a group of steps with a shared derivative strategy.

• Auto-wiring — if an upstream model’s output key matches a downstream model’s input key, the pipeline connects them automatically.

• Cross-group data flow — a model in group 2 can read outputs produced by group 1 (the forward context accumulates across groups).

• active_outputs drives derivatives — the pipeline’s default model_config.active_outputs is synthesized as the union of all sub-model active output sets. Add "stress" to request stress computation.

Tier 3: Fully custom composition (utility functions)

For total control, write a custom nn.Module, BaseModelMixin subclass and use the utility functions in nvalchemi.models._utils:

from nvalchemi.models._utils import (
    autograd_forces,
    autograd_forces_and_stresses,
    autograd_stresses,
    prepare_strain,
    sum_outputs,
)

• autograd_forces(energy, positions) — compute forces as -dE/dr.

• autograd_forces_and_stresses(energy, positions, displacement, cell, num_graphs) — compute forces and stresses from one autograd call.

• autograd_stresses(energy, displacement, cell, num_graphs) — compute tensile-positive Cauchy stresses as 1/V * dE/d(strain).

• prepare_strain(positions, cell, batch_idx) — set up the affine strain trick for autograd stress computation (see below).

• sum_outputs(*outputs) — element-wise sum on additive keys (energies, forces, stresses), last-write-wins for everything else.

Autograd derivatives: forces, stresses, and beyond

This section explains how autograd-based derivatives work in nvalchemi — both for single-model wrapper authors and for pipeline composition.

Single-model wrappers: you own the derivatives

When writing a model wrapper, you decide how to compute derivatives. The framework imposes no constraints. If your model computes forces analytically (like LJ or Ewald via Warp kernels), declare autograd_outputs=set() in your ModelConfig and compute forces directly in your forward() method. If your model uses autograd (like MACE), declare autograd_outputs={"forces"} and call torch.autograd.grad in your forward().

The key expectation is that your forward() returns a ModelOutputs dict with whatever keys your model_config.active_outputs requests, however you choose to compute them.

Example: autograd forces in a wrapper

def forward(self, data, **kwargs):
    model_inputs = self.adapt_input(data, **kwargs)
    raw = self.model(**model_inputs)  # returns {"energy": tensor}

    energy = raw["energy"]
    result = {"energy": energy.unsqueeze(-1)}

    if "forces" in self.model_config.active_outputs:
        result["forces"] = -torch.autograd.grad(
            energy, data.positions,
            grad_outputs=torch.ones_like(energy),
            create_graph=False,  # set True for training
        )[0]

    return self.adapt_output(result, data)

Example: autograd stresses using prepare_strain

Computing stresses via autograd requires the “affine strain trick” — a non-trivial setup step that scales positions and cell through a displacement tensor. The prepare_strain() helper handles this:

from nvalchemi.models._utils import autograd_forces_and_stresses, prepare_strain

def forward(self, data, **kwargs):
    compute_stresses = "stress" in self.model_config.active_outputs

    if compute_stresses:
        scaled_pos, scaled_cell, displacement = prepare_strain(
            data.positions, data.cell, data.batch_idx
        )
        # Run model on scaled tensors
        energy = self.model(scaled_pos, scaled_cell, ...)
    else:
        energy = self.model(data.positions, data.cell, ...)

    result = {"energy": energy.unsqueeze(-1)}

    if "forces" in self.model_config.active_outputs and compute_stresses:
        result["forces"], result["stress"] = autograd_forces_and_stresses(
            energy, scaled_pos, displacement, data.cell, data.num_graphs
        )
    elif "forces" in self.model_config.active_outputs:
        result["forces"] = -torch.autograd.grad(
            energy, data.positions,
            grad_outputs=torch.ones_like(energy),
        )[0]

    if compute_stresses and "stress" not in result:
        grad = torch.autograd.grad(
            energy, displacement,
            grad_outputs=torch.ones_like(energy),
        )[0]
        volume = torch.det(data.cell).abs().view(-1, 1, 1)
        result["stress"] = grad.view(data.num_graphs, 3, 3) / volume

    return self.adapt_output(result, data)

You don’t have to use prepare_strain — it’s a convenience. MACE uses its own internal displacement trick via compute_displacement=True. The only requirement is that your forward() returns the requested outputs.

See Conventions for the project-wide virial, stress, and pressure sign conventions.

Example: Hessians and Jacobians

These are standard torch.autograd operations — nvalchemi does not wrap them:

# Hessian (second derivative of energy w.r.t. positions)
# Models expect a Batch, not raw positions — define a closure.
def energy_fn(pos):
    data.positions = pos
    return model(data)["energy"].sum()

hessian = torch.autograd.functional.hessian(energy_fn, data.positions)

# Born effective charges (Jacobian of dipoles w.r.t. positions)
dipoles = model(data)["dipole"]  # [B, 3]
Z_star = torch.autograd.functional.jacobian(
    lambda pos: model_dipoles(pos), data.positions
)

Pipeline autograd groups: default and custom derivatives

When models are composed in a PipelineModelWrapper with use_autograd=True, the pipeline sums sub-model energies and computes derivatives from the total. What gets computed is driven by model_config.active_outputs:

pipe = PipelineModelWrapper(groups=[
    PipelineGroup(steps=[aimnet2, ewald], use_autograd=True),
])

# Default: pipeline inherits sub-model active output sets
# (typically {"energy", "forces"}).  Forces computed via autograd.
out = pipe(batch)

# Request stresses: pipeline uses affine strain trick automatically.
pipe.model_config.active_outputs = {"energy", "forces", "stress"}
out = pipe(batch)  # now includes stresses

The pipeline’s default model_config.active_outputs is the union of all sub-model active output sets at construction time. If sub-models default to {"energy", "forces"}, the pipeline does too. You can expand it (add "stress") or narrow it (remove "forces").

Default behavior: The pipeline’s built-in derivative function computes forces as -dE/dr and stresses via the affine strain trick. This covers the vast majority of inference use cases.

Custom derivative_fn: For anything beyond forces and stresses, provide a custom function that receives the summed energy, the batch, and the set of requested keys. You write whatever torch.autograd.grad calls you want — the same power as a single-model wrapper’s forward():

def my_derivatives(energy, data, requested):
    """Custom derivative function for a pipeline autograd group.

    Parameters
    ----------
    energy : torch.Tensor
        Summed energy from all models in the group.  On the autograd
        graph --- ready for torch.autograd.grad.
    data : Batch
        The batch.  data.positions has requires_grad=True.
    requested : set[str]
        Output keys still needed (e.g. {"forces", "hessian"}).

    Returns
    -------
    dict[str, torch.Tensor]
        Computed derivatives.
    """
    result = {}
    if "forces" in requested:
        result["forces"] = -torch.autograd.grad(
            energy, data.positions,
            grad_outputs=torch.ones_like(energy),
            retain_graph="hessian" in requested,
        )[0]
    if "hessian" in requested:
        # Your custom Hessian implementation
        result["hessian"] = compute_chunked_hessian(energy, data.positions)
    return result

pipe = PipelineModelWrapper(groups=[
    PipelineGroup(
        steps=[aimnet2, ewald],
        use_autograd=True,
        derivative_fn=my_derivatives,
    ),
])
pipe.model_config.active_outputs = {"energy", "forces", "hessian"}
out = pipe(batch)  # forces + hessian via your function

When derivative_fn is provided, the pipeline does not apply the strain trick or compute forces automatically — your function has full control. If you want stresses, use prepare_strain() inside your function.

Neighbor list handling and make_neighbor_hooks()

All composition tiers handle neighbor lists centrally:

1. The pipeline (or + result) builds a private neighbor-list plan from sub-model NeighborConfig values.

2. By default (neighbor_adaptation="auto"), a source neighbor list can serve a smaller target cutoff only when source_cutoff <= target_cutoff * max_cutoff_ratio. The default ratio is 1.5.

3. If the cutoff gap is larger, the pipeline creates a separate source list for that cutoff group.

4. make_neighbor_hooks() returns all hooks needed by the plan. This is the recommended registration API for composed models.

5. Before each sub-model call, the pipeline temporarily shadows the selected neighbor tensors onto the batch and filters/converts them only when the step plan requires it.

neighbor_adaptation accepts:

• "auto": default. Adapt only when the source cutoff is at most max_cutoff_ratio times the target cutoff; otherwise build another source list.

• "always": build one max-cutoff source and adapt every tighter model from it.

• "never": do not perform cutoff filtering. The pipeline builds exact cutoff source groups. Runtime format conversion is still allowed.

pipe = PipelineModelWrapper(
    groups=[PipelineGroup(steps=[short_range, long_range])],
    neighbor_adaptation="auto",
    max_cutoff_ratio=1.5,
)

for hook in pipe.make_neighbor_hooks():
    dynamics.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE)

Compile behavior: exact source lists use the same NeighborListHook preallocation model as single-model dynamics. First use, shape changes, and K resizing may allocate; steady-state rebuilds reuse persistent buffers. Runtime cutoff filtering and MATRIX/COO conversion still use the existing adapter utilities and may allocate intermediate tensors.

The pipeline synthesizes one or more source lists according to neighbor_adaptation. model_config.neighbor_config exposes the largest (default) source for compatibility.

Choosing a registration pattern:

make_neighbor_hooks() works for both single models and composed models. For a single model it is equivalent to constructing a NeighborListHook manually from the model’s NeighborConfig:

# These two are equivalent for a single model:

# (a) make_neighbor_hooks — recommended
for hook in model.make_neighbor_hooks():
    dynamics.register_hook(hook, stage=DynamicsStage.BEFORE_COMPUTE)

# (b) manual construction — use when you need extra control (e.g. skin distance)
from nvalchemi.hooks import NeighborListHook
dynamics.register_hook(
    NeighborListHook(model.model_config.neighbor_config, skin=0.5),
    stage=DynamicsStage.BEFORE_COMPUTE,
)

For composed models (pipeline), composed.make_neighbor_hooks() returns every hook required by the pipeline’s neighbor-list plan. Manual construction from composed.model_config.neighbor_config only covers the largest/default source and is not sufficient when the plan builds multiple source lists.

Hooks returned by make_neighbor_hooks() are configured for DynamicsStage.BEFORE_COMPUTE, matching the usual dynamics registration stage.

How models integrate with dynamics

Once wrapped, a model plugs directly into the dynamics framework. The dynamics integrator calls the wrapper’s forward method internally via BaseDynamics.compute(), and the resulting forces and energy are written back to the batch:

from nvalchemi.dynamics import DemoDynamics

model = MyPotentialWrapper(hidden_dim=128)
dynamics = DemoDynamics(model=model, n_steps=1000, dt=0.5)
# DemoDynamics expects forces to exist on the batch.
batch.forces = torch.zeros_like(batch.positions)
dynamics.run(batch)

The __needs_keys__ set on the dynamics class (e.g. {"forces"}) is validated against the model’s output after every compute() call, so mismatches between the model’s declared capabilities and the integrator’s requirements are caught immediately at runtime.

See also

• Examples: The gallery includes dynamics examples that demonstrate model usage in context.

• API: nvalchemi.models for the full reference of BaseModelMixin and ModelConfig.

• Dynamics guide: dynamics for how models are used inside optimization and MD workflows.