# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import abc
import warnings
from collections import OrderedDict
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Any
import torch
from pydantic import BaseModel, ConfigDict, Field, model_validator
from nvalchemi._typing import AtomsLike, ModelOutputs
from nvalchemi.data import AtomicData, Batch
if TYPE_CHECKING:
from nvalchemi.distributed.config import StrategyKind
from nvalchemi.distributed.spec import MLIPSpec
warnings.simplefilter("once", UserWarning)
class NeighborListFormat(str, Enum):
"""Storage format for neighbor data written to the batch.
Attributes
----------
COO : str
Coordinate (sparse) format. Internally ``edge_index`` is stored as
``[E, 2]`` (each row is a ``[source, target]`` pair). Model boundary
adapters (e.g. ``MACEWrapper.adapt_input``) transpose to the
conventional ``[2, E]`` layout expected by most GNN-based MLIPs.
MATRIX : str
Dense neighbor-matrix format. Neighbors are stored as a
``neighbor_matrix`` tensor of shape ``[N, max_neighbors]`` (global
atom indices) together with a ``num_neighbors`` tensor of shape
``[N]``. Used by Warp interaction kernels (e.g. Lennard-Jones) that
benefit from fixed-width rows.
"""
COO = "coo" # internal (E, 2); model boundary adapters transpose to (2, E)
MATRIX = "matrix"
[docs]
class NeighborConfig(BaseModel):
"""Configuration for on-the-fly neighbor list construction.
An instance of this class attached to a :class:`ModelConfig` (via its
``neighbor_config`` field) signals that the model requires a neighbor list
and describes the format and parameters it expects. At runtime a
:class:`~nvalchemi.hooks.NeighborListHook` reads this config to compute and
cache the appropriate neighbor data on each :class:`~nvalchemi.data.Batch`
before the model's forward pass, rebuilding only when atoms have moved
beyond the Verlet ``skin``.
The single required argument is ``cutoff`` (in the same length units as
``positions``). ``format`` selects the storage layout: ``COO`` produces a
sparse ``[E, 2]`` edge index (the default, used by most GNN-based MLIPs),
while ``MATRIX`` produces a dense ``[N, max_neighbors]`` neighbor matrix
(used by Warp interaction kernels) — see :class:`NeighborListFormat`.
``half_list`` and ``skin`` tune pair deduplication and rebuild frequency,
respectively.
Examples
--------
A minimal sparse (COO) neighbor list with a 5 A cutoff:
>>> from nvalchemi.models.base import NeighborConfig
>>> nc = NeighborConfig(cutoff=5.0)
>>> nc.format
<NeighborListFormat.COO: 'coo'>
A dense half neighbor list with a Verlet skin, e.g. for a Warp kernel:
>>> from nvalchemi.models.base import NeighborListFormat
>>> nc = NeighborConfig(
... cutoff=5.0,
... format=NeighborListFormat.MATRIX,
... half_list=True,
... skin=1.0,
... )
Attach it to a :class:`ModelConfig` to declare a neighbor-list requirement:
>>> from nvalchemi.models.base import ModelConfig
>>> cfg = ModelConfig(neighbor_config=NeighborConfig(cutoff=5.0))
>>> cfg.needs_neighborlist
True
Notes
-----
- ``skin=0.0`` (the default) rebuilds the neighbor list every step; a
positive skin defers rebuilds until any atom moves more than ``skin / 2``,
trading memory staleness for fewer rebuilds.
- ``half_list=True`` stores each ``(i, j)`` pair once; the interaction kernel
applies Newton's third law to recover forces on both atoms, so it is only
appropriate for kernels that expect a half list.
"""
cutoff: Annotated[
float,
Field(
description=(
"Interaction cutoff radius in the same length units as positions."
)
),
]
format: Annotated[
NeighborListFormat,
Field(
description=(
"Whether to build a dense neighbor matrix (``MATRIX``) or a sparse "
"edge-index list (``COO``). Defaults to ``COO``."
)
),
] = NeighborListFormat.COO
half_list: Annotated[
bool,
Field(
description=(
"If ``True``, each pair ``(i, j)`` with ``i < j`` appears only once. "
"Newton's third law is applied inside the interaction kernel to recover "
"forces on both atoms. Defaults to ``False``."
)
),
] = False
skin: Annotated[
float,
Field(
description=(
"Verlet skin distance. The neighbor list is only rebuilt when any atom "
"has moved more than ``skin / 2`` since the last build. Set to ``0.0`` "
"(default) to rebuild every step."
)
),
] = 0.0
[docs]
class ModelConfig(BaseModel):
"""Unified model configuration combining capability declaration and
runtime control.
``ModelConfig`` is the contract between a model wrapper and the rest of
nvalchemi: dynamics engines, composition pipelines, loss functions, and the
:class:`BaseModelMixin` adapters all read it to decide which inputs to
prepare, which gradients to enable, and which outputs to compute. Every
:class:`BaseModelMixin` subclass must set a ``self.model_config`` instance
in its ``__init__`` (there is deliberately no class-level default, so each
wrapper owns its own config object).
A ``ModelConfig`` has two kinds of fields:
- **Capability fields** (frozen at construction) describe what the
model checkpoint can do. These use ``frozenset`` to signal
immutability. They are set once by the wrapper's ``__init__`` and
should not be changed at runtime. Examples: ``outputs``,
``autograd_outputs``, ``autograd_inputs``, ``required_inputs``,
``optional_inputs``, ``supports_pbc``, ``needs_pbc``,
``neighbor_config``.
- **Runtime fields** (mutable) control what the model should compute
on each forward pass. These can be changed freely by the user:
``active_outputs`` selects the subset of ``outputs`` to compute this
run, and ``gradient_keys`` enables gradients on extra input tensors.
``outputs`` and ``required_inputs`` use free-form strings so new
properties can be added without modifying this class. Well-known
output keys: ``energy``, ``forces``, ``stress``, ``hessian``,
``dipole``, ``charges``, ``embeddings``. Declare a neighbor-list
requirement by attaching a :class:`NeighborConfig` (see
:attr:`needs_neighborlist`).
Examples
--------
An energy-and-forces model whose forces come from autograd on positions:
>>> from nvalchemi.models.base import ModelConfig
>>> cfg = ModelConfig(
... outputs=frozenset({"energy", "forces"}),
... autograd_outputs=frozenset({"forces"}),
... autograd_inputs=frozenset({"positions"}),
... )
>>> cfg.active_outputs == {"energy", "forces"}
True
Restrict a run to a single output without changing the model's capabilities:
>>> cfg.active_outputs = {"energy"}
Declare a PBC-aware model that needs a neighbor list:
>>> from nvalchemi.models.base import NeighborConfig
>>> cfg = ModelConfig(
... outputs=frozenset({"energy", "forces", "stress"}),
... autograd_outputs=frozenset({"forces", "stress"}),
... supports_pbc=True,
... needs_pbc=True,
... neighbor_config=NeighborConfig(cutoff=5.0),
... )
>>> cfg.needs_neighborlist
True
Notes
-----
- ``extra="forbid"``: unknown constructor keywords raise a
:class:`pydantic.ValidationError`, guarding against typo'd field names.
- ``active_outputs`` defaults to a mutable copy of ``outputs`` when left as
``None``; set it to narrow the per-run output set, and reset it to ``None``
to fall back to all ``outputs``.
- Capability fields are ``frozenset`` values: rebind the whole field to
change them (in-place mutation is impossible), which keeps the declared
capabilities effectively immutable after construction.
"""
# ── Capability fields (frozen at construction) ──────────────────────
outputs: Annotated[
frozenset[str],
Field(
default_factory=lambda: frozenset({"energy"}),
description="All properties the model can produce.",
),
]
autograd_outputs: Annotated[
frozenset[str],
Field(
default_factory=frozenset,
description="Subset of outputs computed via autograd.",
),
]
autograd_inputs: Annotated[
frozenset[str],
Field(
default_factory=lambda: frozenset({"positions"}),
description="Input keys needing requires_grad for autograd outputs.",
),
]
required_inputs: Annotated[
frozenset[str],
Field(
default_factory=frozenset,
description="Extra required inputs beyond {positions, atomic_numbers}.",
),
]
optional_inputs: Annotated[
frozenset[str],
Field(
default_factory=frozenset,
description="Extra inputs used if present, silently skipped if absent.",
),
]
supports_pbc: Annotated[
bool,
Field(
default=False,
description="Whether the model supports periodic boundary conditions.",
),
]
needs_pbc: Annotated[
bool,
Field(
default=False,
description="Whether the model requires PBC inputs.",
),
]
neighbor_config: Annotated[
NeighborConfig | None,
Field(
default=None,
description="Neighbor list requirements. None means no neighbor list.",
),
]
# ── Runtime fields (mutable) ────────────────────────────────────────
active_outputs: Annotated[
set[str] | None,
Field(
default=None,
description=(
"Properties to compute this run. "
"None means use all outputs (the default)."
),
),
]
gradient_keys: Annotated[
set[str],
Field(
default_factory=set,
description="Extra input keys to enable gradients for.",
),
]
model_config = ConfigDict(extra="forbid")
@model_validator(mode="after")
def _default_active_outputs(self) -> "ModelConfig":
"""Default active_outputs to outputs if not explicitly set."""
if self.active_outputs is None:
# Use object.__setattr__ because we're inside validation
object.__setattr__(self, "active_outputs", set(self.outputs))
return self
@property
def needs_neighborlist(self) -> bool:
"""Convenience accessor: ``True`` when the model requires a neighbor list."""
return self.neighbor_config is not None
[docs]
class BaseModelMixin(abc.ABC):
"""Abstract mixin providing a standardized interface for model wrappers.
All external MLIP wrappers should inherit from this mixin (alongside
``nn.Module``) to ensure a consistent interface for dynamics engines,
composition pipelines, and downstream tooling.
Concrete implementations must provide:
- ``model_config`` attribute — a :class:`ModelConfig` instance set in
``__init__``.
- ``embedding_shapes`` property — expected shapes of computed
embeddings.
- ``compute_embeddings()`` — compute and attach embeddings to the
input data structure.
The mixin provides default implementations of:
- ``input_data()`` — set of required input keys derived from the
model config.
- ``output_data()`` — set of active outputs intersected with
supported outputs (warns on unsupported requests).
- ``adapt_input()`` — enable gradients on required tensors and
collect input dict.
- ``adapt_output()`` — map raw model output to :class:`ModelOutputs`
ordered dict.
"""
# model_config must be set as an instance attribute in each subclass __init__:
# self.model_config = ModelConfig(outputs=..., ...)
# There is intentionally NO class-level default to prevent all instances from
# sharing a single ModelConfig object (which would cause mutations in one wrapper
# to silently affect all others). __init_subclass__ wraps __init__ to enforce
# this at construction time — a missing model_config raises TypeError.
# Per-scope distributed runtime context. ``None`` outside a
# ``DistributedModel`` scope; set to the
# :class:`DistributedContext` the scope owns by
# :meth:`distributed_setup`. Read via ``self._dist_ctx`` from
# :meth:`adapt_input` to get live ``halo_meta`` / ``gather_meta``.
_dist_ctx: Any = None
def __init_subclass__(cls, **kwargs: Any) -> None:
"""Hook applied to every concrete subclass at class-creation time.
Performs two injections:
1. **extra_repr** — ``nn.Module.__repr__`` calls
``self.extra_repr()`` but its default returns ``""``. Since
wrappers inherit ``nn.Module`` before ``BaseModelMixin``
(required for PyTorch), ``Module.extra_repr`` wins in the MRO.
This hook injects our version directly onto each concrete
wrapper class so it takes precedence.
2. **model_config post-init check** — wraps the subclass
``__init__`` so that after construction,
``self.model_config`` is verified to exist. This catches
the common mistake of forgetting to set ``model_config`` in
``__init__`` with a clear error instead of a late
``AttributeError`` deep in a forward pass.
"""
super().__init_subclass__(**kwargs)
if "extra_repr" not in cls.__dict__:
cls.extra_repr = BaseModelMixin._config_extra_repr
# Wrap __init__ to verify model_config is set after construction.
if "__init__" in cls.__dict__:
import functools
original_init = cls.__init__
@functools.wraps(original_init)
def _checked_init(self: Any, *args: Any, **kw: Any) -> None:
original_init(self, *args, **kw)
if not hasattr(self, "model_config"):
raise TypeError(
f"{type(self).__name__}.__init__() must set "
f"self.model_config = ModelConfig(...). "
f"See BaseModelMixin docstring for details."
)
cls.__init__ = _checked_init # type: ignore[attr-defined]
@staticmethod
def _config_extra_repr(self: Any) -> str:
"""Format the model config for ``nn.Module.__repr__``."""
cfg = getattr(self, "model_config", None)
if cfg is None:
return "model_config=<not set>"
parts = []
outputs = sorted(cfg.outputs)
active = sorted(cfg.active_outputs)
parts.append(f"outputs={{{', '.join(outputs)}}}")
if set(active) != set(outputs):
parts.append(f"active_outputs={{{', '.join(active)}}}")
if cfg.autograd_outputs:
parts.append(
f"autograd_outputs={{{', '.join(sorted(cfg.autograd_outputs))}}}"
)
if cfg.required_inputs:
parts.append(
f"required_inputs={{{', '.join(sorted(cfg.required_inputs))}}}"
)
if cfg.optional_inputs:
parts.append(
f"optional_inputs={{{', '.join(sorted(cfg.optional_inputs))}}}"
)
if cfg.supports_pbc or cfg.needs_pbc:
pbc_parts = []
if cfg.supports_pbc:
pbc_parts.append("supports_pbc")
if cfg.needs_pbc:
pbc_parts.append("needs_pbc")
parts.append(f"pbc=[{', '.join(pbc_parts)}]")
if cfg.neighbor_config is not None:
nc = cfg.neighbor_config
nc_str = f"cutoff={nc.cutoff}, format={nc.format.value}"
if nc.half_list:
nc_str += ", half_list"
parts.append(f"neighbors=({nc_str})")
return "\n".join(parts)
@property
@abc.abstractmethod
def embedding_shapes(self) -> dict[str, tuple[int, ...]]:
"""Retrieves the expected shapes of the node, edge, and graph embeddings."""
...
[docs]
@abc.abstractmethod
def compute_embeddings(
self, data: AtomicData | Batch, **kwargs: Any
) -> AtomicData | Batch:
"""
Compute embeddings at different levels of a batch of atomic graphs.
Parameters
----------
data : AtomicData | Batch
Input atomic data containing positions, atomic numbers, etc.
Returns
-------
AtomicData | Batch
Data structure with embeddings attached in-place.
Raises
------
NotImplementedError
If the model does not support embeddings computation
"""
...
[docs]
def direct_derivative_keys(self) -> set[str]:
"""Return output keys this model computes analytically in ``forward()``.
When this model participates in a pipeline autograd group, the
pipeline strips ``"forces"`` and ``"stress"`` from sub-model
``active_outputs`` so it can compute them via autograd on the
summed energy. Keys returned by this method are **kept** in
``active_outputs`` — the pipeline collects them from the model
output and sums them with the autograd-derived derivatives.
Override this in models that produce analytical forces or stress
alongside an energy that carries autograd information (e.g.
Ewald/PME with ``hybrid_forces=True``).
Returns
-------
set[str]
Keys (e.g. ``{"forces", "stress"}``) that the model produces
analytically and should be summed with autograd derivatives.
Default: empty set (all derivatives come from autograd).
"""
return set()
[docs]
def set_config(self, key: str, value: Any) -> None:
"""Set a mutable field on :attr:`model_config`.
Convenience method equivalent to
``self.model_config.<key> = value`` with validation that the
field exists and is mutable.
Parameters
----------
key : str
Name of a mutable ``ModelConfig`` field (e.g.
``"active_outputs"``, ``"gradient_keys"``).
value
New value for the field.
Raises
------
AttributeError
If *key* is not a field on :class:`ModelConfig`.
"""
if not hasattr(self.model_config, key):
raise AttributeError(
f"ModelConfig has no field '{key}'. "
f"Available fields: {list(self.model_config.model_fields)}"
)
setattr(self.model_config, key, value)
[docs]
def adapt_output(self, model_output: Any, data: AtomicData | Batch) -> ModelOutputs:
"""Adapt external model output to :class:`ModelOutputs` format.
Returns an OrderedDict keyed by :meth:`output_data` entries,
populated from *model_output* where keys match.
.. note::
Returned tensors may still be attached to the autograd
computation graph (e.g. energies from autograd-force models
like MACE). This is intentional — the model does not know
whether the caller needs the graph (e.g. pipeline
shared-autograd groups). **Callers that do not need the
graph are responsible for detaching.**
Parameters
----------
model_output : Any
Raw output from the external model.
data : AtomicData | Batch
Original input data (may be needed for context/metadata).
Returns
-------
ModelOutputs
OrderedDict with expected output keys and their values
(or ``None`` if not present). Tensors may be graph-attached.
"""
# ``output_data()`` returns a set, whose iteration order is randomized
# per process (str hashing). Sort so every rank seeds the output dict in
# an identical key order.
output = OrderedDict((key, None) for key in sorted(self.output_data()))
if isinstance(model_output, dict):
for key in output:
value = model_output.get(key)
if value is not None:
if key == "energy" and value.ndim == 1:
value = value.unsqueeze(-1)
output[key] = value
return output
[docs]
def add_output_head(self, prefix: str) -> None:
"""
Add an output head to the model.
Parameters
----------
prefix : str
Prefix for the output head
"""
raise NotImplementedError
# ------------------------------------------------------------------
# Distributed hooks
# ------------------------------------------------------------------
[docs]
def distribution_spec(
self, strategy: "StrategyKind | None" = None
) -> "MLIPSpec | None":
r"""Return the :class:`MLIPSpec` describing the primitives this model
needs under domain parallelism *for the given parallelization strategy*.
Parameters
----------
strategy
The :class:`~nvalchemi.distributed.config.StrategyKind` the scope runs
under (``HALO`` / ``GRAPH_PARTITION``); ``None``
is treated as ``HALO``. The framework passes the config-selected
strategy — models must not sniff the environment. The spec content is a
joint :math:`(\mathrm{model} \times \mathrm{strategy})` product, so a model that supports graph
parallel returns a different ``(policy, adapters, shard_fields,
consolidation)`` bundle per strategy.
Returns
-------
MLIPSpec | None
Default ``None`` = model doesn't declare distributed support.
``DomainParallel`` raises if asked to shard a model whose spec is
``None``. Per-model wrappers override to return a preset
(``SPEC_MPNN_HALO`` / ``SPEC_UMA_HALO``) or a custom spec.
"""
return None
[docs]
def distributed_setup(self, ctx: Any) -> None:
"""Called once by :class:`DistributedModel` when entering a
distributed scope, after the spec's adapters have been
installed. ``ctx`` is the :class:`DistributedContext` that the
framework will mutate per-step (``ctx.halo_meta`` /
``ctx.gather_meta``); the wrapper should stash a reference to
``ctx`` (commonly as ``self._dist_ctx``) so its
:meth:`adapt_input` can read the live values at forward time.
Override to also build per-rank closures (e.g. distributed
helper replacements that close over ``ctx.gather_meta``) at
scope-entry time.
Default: no-op. Simple halo-mode models (MACE, LJ) don't
override beyond stashing the ctx reference, which the default
implementation handles for them — see :attr:`_dist_ctx`.
"""
self._dist_ctx = ctx
[docs]
def distributed_teardown(self) -> None:
"""Restore anything :meth:`distributed_setup` mutated. Default
clears the ctx reference; override to also restore monkey-patches
the framework's adapter registry doesn't manage."""
self._dist_ctx = None
[docs]
def output_data(self) -> set[str]:
"""Return the set of keys the model will compute this run.
Intersects ``active_outputs`` with ``outputs``.
Warns if any active keys are not supported by the model.
Returns
-------
set[str]
Set of output keys that are both active and supported.
"""
active = self.model_config.active_outputs
supported = self.model_config.outputs
unsupported = active - supported
if unsupported:
warnings.warn(
f"Requested {unsupported} but model only supports {supported}.",
UserWarning,
stacklevel=2,
)
return active & supported
[docs]
def export_model(self, path: Path, as_state_dict: bool = False) -> None:
"""
Export the current model without the ``BaseModelMixin`` interface.
"""
raise NotImplementedError
def __add__(self, other: "BaseModelMixin") -> "PipelineModelWrapper":
"""Compose two models additively via the ``+`` operator.
Returns a :class:`~nvalchemi.models.pipeline.PipelineModelWrapper`
where each model occupies its own group with
``use_autograd=False``, so energy, forces, and stress from
both models are summed element-wise.
This is the simplest composition pattern — suitable when each model
computes its own forces independently (analytically or via its own
internal autograd). For dependent pipelines where one model's
output feeds into another's input, or for shared-autograd groups
that differentiate the summed energy of multiple models, use the
explicit :class:`~nvalchemi.models.pipeline.PipelineModelWrapper`
constructor with :class:`~nvalchemi.models.pipeline.PipelineGroup`
and :class:`~nvalchemi.models.pipeline.PipelineStep`.
Parameters
----------
other : BaseModelMixin
Another model to compose with.
Returns
-------
PipelineModelWrapper
A pipeline that sums the outputs of both models.
Examples
--------
>>> combined = lj_model + ewald_model
>>> combined = mace_model + dftd3_model
>>> combined = model_a + model_b + model_c # chains naturally
"""
from nvalchemi.models.pipeline import ( # noqa: PLC0415
PipelineGroup,
PipelineModelWrapper,
)
# If the left-hand side is already a pipeline of direct groups
# (produced by a previous +), flatten into it instead of nesting.
if isinstance(self, PipelineModelWrapper):
new_groups = list(self.groups) + [PipelineGroup(steps=[other])]
return PipelineModelWrapper(groups=new_groups)
return PipelineModelWrapper(
groups=[
PipelineGroup(steps=[self]),
PipelineGroup(steps=[other]),
]
)
[docs]
def make_neighbor_hooks(
self,
max_neighbors: int | None = None,
neighbor_list_method: str | None = None,
) -> list:
"""Return a list of :class:`~nvalchemi.hooks.NeighborListHook` instances
for this model's neighbor configuration.
Returns an empty list if the model does not require a neighbor list.
Defers the import to avoid circular imports.
Parameters
----------
max_neighbors : int | None, optional
Maximum neighbors per atom for MATRIX format. When ``None``
(default), auto-estimated from the cutoff at first use.
neighbor_list_method : str | None, optional
Explicit ``nvalchemiops`` neighbor-list method to use. When
``None`` (default), the hook selects an appropriate method.
"""
from nvalchemi.dynamics.base import DynamicsStage # noqa: PLC0415
from nvalchemi.hooks import NeighborListHook # noqa: PLC0415
nc = self.model_config.neighbor_config
if nc is None:
return []
return [
NeighborListHook(
nc,
skin=nc.skin,
max_neighbors=max_neighbors,
method=neighbor_list_method,
stage=DynamicsStage.BEFORE_COMPUTE,
)
]
