nvalchemi-model-wrapping

name	nvalchemi-model-wrapping
description	How to wrap an arbitrary MLIP (Machine Learning Interatomic Potential) using the BaseModelMixin interface to standardize inputs, outputs, and embeddings. Use when integrating a model such as MACE or AIMNet2 (e.g. MACEWrapper, loading pretrained checkpoints) so dynamics, training, or fine-tuning stages can call it, or when exposing energies, forces, or embeddings from a custom PyTorch model.

nvalchemi Model Wrapping

Overview

To use an arbitrary MLIP (Machine Learning Interatomic Potential) within nvalchemi, pair it with the BaseModelMixin interface. This standardizes how models receive AtomicData/Batch inputs and produce ModelOutputs.

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

---

Architecture

A wrapper subclasses nn.Module and BaseModelMixin, and holds the underlying model by composition (self.model = ...). This is the pattern used by every built-in wrapper (DemoModelWrapper, MACEWrapper, AIMNet2Wrapper, LennardJonesModelWrapper).

┌──────────────────────┐    ┌──────────────────┐
│  YourModel(nn.Module)│    │  BaseModelMixin   │
│  - forward()         │    │  - model_config   │
│  - your layers       │    │  - adapt_input()  │
└──────────────────────┘    │  - adapt_output() │
        held via            └────────┬─────────┘
      composition                    │
             ┌──────────▼───────────────────────┐
             │  YourModelWrapper                 │
             │  (nn.Module, BaseModelMixin)      │
             │  self.model = YourModel(...)      │
             │  self.model_config = ModelConfig(…)│
             └───────────────────────────────────┘

nn.Module must come first in the bases so PyTorch initializes correctly.

---

Step-by-step guide

1. Set model_config in __init__ (capabilities & runtime control)

ModelConfig unifies two kinds of fields:

• Capability fields (frozen frozenset/bool at construction) describe what the checkpoint can do: outputs, autograd_outputs, autograd_inputs, required_inputs, optional_inputs, supports_pbc, needs_pbc, neighbor_config.
• Runtime fields (mutable) control what to compute each pass: active_outputs (defaults to outputs) and gradient_keys.

BaseModelMixin enforces that every wrapper sets self.model_config in __init__ (a missing one raises TypeError at construction).

def __init__(self, model: nn.Module) -> None:
    super().__init__()
    self.model = model
    self.model_config = ModelConfig(
        outputs=frozenset({"energy", "forces"}),    # everything the model CAN produce
        autograd_outputs=frozenset({"forces"}),     # subset computed via autograd
        autograd_inputs=frozenset({"positions"}),   # inputs needing requires_grad
        required_inputs=frozenset(),                # extra required beyond positions/atomic_numbers
        optional_inputs=frozenset(),                # used if present, skipped if absent
        supports_pbc=False,
        needs_pbc=False,
        neighbor_config=None,                       # NeighborConfig(...) if needed
    )

Well-known output keys: energy, forces, stress, hessians, dipoles, charges, embeddings. outputs/required_inputs are free-form strings, so new properties can be added without changing ModelConfig.

2. Define embedding_shapes

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

3. Implement adapt_input

Converts AtomicData/Batch to a dict of keyword arguments for the underlying model's forward().

Always call super().adapt_input() first — it enables requires_grad on autograd_inputs (when an autograd output is active) plus any gradient_keys, and collects the keys declared by input_data().

def adapt_input(self, data: AtomicData | Batch, **kwargs: Any) -> 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

    # Gate behavior on the active outputs, not a compute_* flag
    model_inputs["compute_forces"] = "forces" in self.model_config.active_outputs
    return model_inputs

4. Implement adapt_output

Converts the model's raw output to ModelOutputs (an OrderedDict[str, Tensor | None]).

Always call super().adapt_output() first — it returns an OrderedDict pre-filled with the output_data() keys (set to None) and auto-maps matching key names (unsqueezing a 1-D energy to [B, 1]).

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

    # Map model outputs to standardized keys
    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"]

    return output

Standard output keys and shapes:

Key	Shape	Notes
energy	[B, 1]	Per-graph energy (eV)
forces	[V, 3]	Per-node forces
stress	[B, 3, 3]	Per-graph stress tensor
hessian	[V, 3, 3]	Energy Hessian
dipole	[B, 3]	Dipole moment
charges	[V]	Partial charges

5. Implement compute_embeddings

embedding_shapes and compute_embeddings are abstract on BaseModelMixin, so every wrapper must define them (raise NotImplementedError if the model has no embeddings). compute_embeddings writes embeddings to the data structure in-place and returns it.

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

    # Run model layers to get intermediate representations
    atom_z = self.model.embedding(model_inputs["atomic_numbers"])
    coord_z = self.model.coord_embedding(model_inputs["positions"])
    embedding = self.model.joint_mlp(torch.cat([atom_z, coord_z], dim=-1))

    # Aggregate to graph level
    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_embedding = torch.zeros(
        (num_graphs, *self.embedding_shapes["graph_embedding"]),
        device=embedding.device, dtype=embedding.dtype,
    )
    graph_embedding.scatter_add_(0, batch_indices.unsqueeze(-1), embedding)

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

6. Implement forward

The main entry point. Adapts input, calls the underlying model, adapts output.

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

7. (Optional) Override export_model / add_output_head

BaseModelMixin.export_model and add_output_head default to raising NotImplementedError. Override them if your model needs to be exported without the mixin (e.g. for ASE calculators) or supports extra output heads.

def export_model(self, path: Path, as_state_dict: bool = False) -> None:
    if as_state_dict:
        torch.save(self.model.state_dict(), path)
    else:
        torch.save(self.model, path)

---

Runtime control: active_outputs

active_outputs selects what to compute on each forward pass. Change it with set_config, which validates that the field exists and is mutable:

model = MyModelWrapper(MyPotential())

# Enable stress (e.g. for NPT/NPH) — must already be in outputs
model.set_config("active_outputs", {"energy", "forces", "stress"})

# Add extra gradient inputs beyond those implied by autograd_inputs
model.set_config("gradient_keys", {"positions"})

set_config(key, value) is equivalent to model.model_config.<key> = value. output_data() returns active_outputs & outputs and warns if you request a key the model does not support.

---

Helper methods

Method	Returns	Description
input_data()	set[str]	Required input keys from model_config (positions, atomic_numbers, neighbor-list keys, pbc, required_inputs)
output_data()	set[str]	active_outputs & outputs (warns on unsupported requests)
set_config(key, value)	None	Set a mutable ModelConfig field with validation
direct_derivative_keys()	set[str]	Outputs computed analytically alongside an autograd energy (pipeline autograd); default empty
add_output_head(prefix)	None	Override to add an MLP output head; default raises NotImplementedError
export_model(path, as_state_dict=False)	None	Override to export the raw model; default raises NotImplementedError

---

Complete example

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

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

    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, atomic_numbers=None, batch_indices=None):
        h = self.encoder(positions)
        node_energy = self.energy_head(h)
        if batch_indices is not None:
            num_graphs = int(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)
        :
            energy = node_energy.(dim=, keepdim=)
         {: energy}


 (nn.Module, BaseModelMixin):
    

     ():
        ().__init__()
        .model = MyPotential(hidden_dim)
        .model_config = ModelConfig(
            outputs=({, }),
            autograd_outputs=({}),
            autograd_inputs=({}),
            supports_pbc=,
            needs_pbc=,
            neighbor_config=,
        )


     () -> [, [, ...]]:
         {: (.model.hidden_dim,)}

     () -> AtomicData | Batch:
        model_inputs = .adapt_input(data, **kwargs)
        data.node_embeddings = .model.encoder(model_inputs[])
         data

     () -> [, ]:
        model_inputs = ().adapt_input(data, **kwargs)
        model_inputs[] = data.positions
         (data, Batch):
            model_inputs[] = data.batch_idx
        :
            model_inputs[] = 
         model_inputs

     () -> ModelOutputs:
        output = ().adapt_output(model_output, data)
        output[] = model_output[]
           .model_config.active_outputs:
            output[] = -torch.autograd.grad(
                model_output[],
                data.positions,
                grad_outputs=torch.ones_like(model_output[]),
                create_graph=.training,
            )[]
         output

     () -> ModelOutputs:
        model_inputs = .adapt_input(data, **kwargs)
        model_outputs = .model(**model_inputs)
         .adapt_output(model_outputs, data)



model = MyPotentialWrapper(hidden_dim=)
model.set_config(, {, })

data = AtomicData(
    positions=torch.randn(, ),
    atomic_numbers=torch.tensor([, , , , ], dtype=torch.long),
)
batch = Batch.from_data_list([data])
outputs = model(batch)