nvalchemi-model-wrapping

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