xnn.gnn.models.mace_foundation#

Load pretrained MACE foundation models into the xnn MACE.

The MACE foundation models – the MACE-MP materials series (MP-0, MP-0b/0b2/ 0b3, MPA-0, OMAT-0, MATPES, the multi-head MH series; Batatia et al., arXiv:2401.00096) and the MACE-OFF23 organic series (Kovacs et al., arXiv:2312.15211) – are distributed as pickled mace-torch modules of the upstream ScaleShiftMACE class. Since the xnn MACE reproduces upstream block by block, those checkpoints convert weight-for-weight into MACE:

  • the ScaleShiftMACE energy expression maps onto the model’s scale_shift block,

  • the Agnesi distance transform, ZBL pair repulsion and the density-normalized interaction blocks map onto the same-named xnn options,

  • multi-head checkpoints are sliced to one head: the per-head rows of the atomic energies, readout weights and scale/shift are extracted, so the converted model is an ordinary single-head potential (the readout hidden channels of head h occupy one contiguous block, and only the final readout linear needs a sqrt(1/n_heads) fan-in renormalization).

Loading a checkpoint requires the mace-torch package (the pickle references its classes); the conversion itself and the converted model do not. Checkpoints are cached under the xnn dataset cache (<cache>/foundations/); a file already in mace-torch’s own cache (~/.cache/mace) is reused instead of re-downloading.

Not covered: checkpoints trained with apply_cutoff=False radial embeddings or interaction blocks outside the xnn registry (currently only mace-mh-1, whose RealAgnosticResidualNonLinearInteractionBlock is a different architecture generation) – the converter raises NotImplementedError naming the offending piece.

Functions

foundation_to_xnn(source[, head, dtype])

Resolve, load and convert a foundation checkpoint (or model).

from_mace_torch(upstream[, head, dtype])

Convert a mace-torch model (plain or ScaleShift) to an xnn MACE.

load_foundation(source)

Load an upstream foundation checkpoint as a mace-torch module.
Show Source