xnn.common.models.base.InteratomicPotential#

class xnn.common.models.base.InteratomicPotential(*args, **kwargs)[source]#

Bases: Module

Abstract base class every interatomic-potential model must subclass.

A model is an torch.nn.Module that maps an AtomicGraph to a dict of tensors with at least the keys "node_energy" (shape (N,), per-atom energies) and "energy" (shape (B,), per-structure energies). Models do not compute forces or stress themselves; those are added uniformly by xnn.common.models.outputs.ForceStressOutput via autograd, so every model gets correct conservative forces for free.

Notes

The radial cutoff (in Angstrom, used to build neighbor lists) is set as the instance attribute cutoff in each subclass __init__. It is intentionally not declared as a bare class annotation here, so subclasses remain compatible with torch.jit.script().

head_modules names the attributes that form one readout head (the readout layers, the per-element reference energies and any energy scale/shift), so that MultiHead can give the model several heads over one shared trunk. A model that leaves it empty (the classical force fields) cannot be given several heads.

abstractmethod forward(data)[source]#

Compute per-atom and per-structure energies for a graph.

Parameters:

data (AtomicGraph) – The batched atomic graph to evaluate.

Returns:

A mapping containing at least "node_energy" of shape (N,) (per-atom energies) and "energy" of shape (B,) (per-structure energies), where N is the number of atoms and B the number of structures in the batch.

Return type:

dict of str to torch.Tensor

Raises:

NotImplementedError – Always, unless overridden by a concrete subclass.

abstractmethod classmethod from_config(cfg)[source]#

Construct a model instance from a configuration dataclass.

Parameters:

cfg (ModelConfig) – The model configuration dataclass describing hyperparameters.

Returns:

A newly constructed model instance.

Return type:

InteratomicPotential

Raises:

NotImplementedError – Always, unless overridden by a concrete subclass.

classmethod from_pretrained(source, **kwargs)[source]#

Load a pre-trained model of this class.

A thin front of xnn.common.models.hub.from_pretrained() that returns the bare potential (no force wrapper) and checks its class, e.g. MACE.from_pretrained("mace-off23-small"). Dispersion and long-range wrappers built from the config are kept around it.

Parameters:

  • source (str or pathlib.Path) – Registry name, local checkpoint or model directory, URL or Zenodo DOI.

  • **kwargs – Options of load_pretrained().

Returns:

The model, in eval mode.

Return type:

InteratomicPotential

Raises:

TypeError – If the loaded model is not an instance of this class.

aggregate_energy(node_energy, data)[source]#

Sum per-atom energies into per-structure energies (locality).

Parameters:

  • node_energy (torch.Tensor) – Per-atom energies of shape (N,).

  • data (AtomicGraph) – The batched graph, providing num_graphs and the batch vector that maps each atom to its structure index.

Returns:

Per-structure total energies of shape (B,), where B == data.num_graphs, in float64 (see structure_sum()).

Return type:

torch.Tensor