xnn.gnn.models.aimnet2_foundation#

Load the published AIMNet2 models into the xnn AIMNet2.

The AIMNet2 models of Anstine, Zubatyuk and Isayev (Chem. Sci. 16, 10228, 2025; MIT license) are published by isayevlab/aimnetcentral as six families of four ensemble members each:

  • aimnet2-wb97m-d3: the general organic model (wB97M-D3/def2-TZVPP; H, B, C, N, O, F, Si, P, S, Cl, As, Se, Br, I);

  • aimnet2-b973c-d3: the same chemistry at the B97-3c level;

  • aimnet2-b973c-2025-d3: the 2025 B97-3c retraining with improved intermolecular interactions;

  • aimnet2-nse: the open-shell variant (two charge channels, spin multiplicity input);

  • aimnet2-pd: palladium organometallics (B97-3c, no Coulomb term);

  • aimnet2-rxn: reactive CHNO chemistry (wB97M, neutral systems).

Each member is a self-describing artifact: a torch.save file loadable with weights_only=True holding the state dict, the architecture as YAML and the metadata (cutoff, elements, Coulomb mode, D3 parameters). The same content is distributed as config.json + ensemble_<k>.safetensors directories; both are read here. Since the xnn AIMNet2 follows the reference architecture block by block, the weights transplant directly:

  • the per-element embedding rows, the Gaussian centers, the shell-combination weights, the MLPs of every pass and the energy readout copy over;

  • the per-element energy shifts go to atom_ref (float64, as stored);

  • the short-range Coulomb cutoff and the Gaussian exponent are read off the artifact into the model config, so a model rebuilt from its config reproduces the conversion exactly;

  • the D3(BJ) dispersion the models are served with (the reference calculator adds it post hoc, the labels were fitted without it) becomes the config’s subtracted_dispersion record, which the model hub adds back through the shared D3Dispersion (two-body, 15 Angstrom cutoff with a 3 Angstrom switching window, as the reference does). dispersion=False at load time serves the bare network.

The models are entries of the model hub registry in the "aimnet2" format, registered by this module: xnn.common.models.from_pretrained() downloads an artifact, converts it once and caches the converted model as a portable xnn directory. The conversion needs nothing beyond torch and PyYAML.

Not covered: the legacy TorchScript (.jpt) checkpoints of the first AIMNet2 release (the reference package converts them to the artifact format read here), and architectures the reference code can express but never published (shared-shell features, a second Gaussian basis for the vector part); the converter raises NotImplementedError naming the piece.

Module Attributes

FOUNDATION_MODELS

name -> (download URL, license), a view of the "aimnet2" entries of the model hub registry (models.json), which is the single list.

ALIASES

alias -> registered name (member 0 of each family, as the reference package resolves them).

Functions

convert_artifact(artifact[, dtype])

Convert a reference artifact; return the model and the config that rebuilds it.

detect_artifact(path)

Whether path is an AIMNet2 artifact file or directory.

foundation_to_xnn(source[, dtype, ...])

Resolve, load and convert a published AIMNet2 model (the body of AIMNet2.from_foundation()).

from_aimnet_artifact(artifact[, dtype])

Convert a reference AIMNet2 artifact to an xnn AIMNet2.

load_artifact(source[, member])

Read an AIMNet2 artifact into one mapping.

read_safetensors(path)

Read a .safetensors file into a dict of CPU tensors.
Show Source