MACE v0.3.16 Release Notes

Excited to announce MACE v0.3.16, featuring MACE-POLAR-1, a new family of polarisable electrostatic foundation models, a major performance push (torch.compile with reduce-overhead, OpenEquivariance and hybrid kernels, and a torch-sim interface), and a range of training and stability fixes.

🏗️ Foundation Models

MACE-POLAR-1

Introduced MACE-POLAR-1, a new family of foundation models that extends the MACE architecture with explicit long-range electrostatics. MACE-POLAR-1 augments local atomic energies with a non-self-consistent polarisable field formalism, learning per-atom charge and spin densities — represented as multipole expansions in a Gaussian-type orbital basis — directly from energy and force labels alone. Global charge and spin constraints are enforced through learnable Fukui equilibration functions, so the model handles arbitrary charge and spin states, responds to external electric fields, and exposes physically interpretable spin-resolved charge densities.

MACE-POLAR-1 is trained on OMol25 (100 M structures at the ωB97M-V hybrid DFT level of theory) covering 83 elements. Two variants are available:

• polar-1-m (medium) — 2 interaction layers, 12 Å receptive field
• polar-1-l (large) — 3 interaction layers, 18 Å receptive field

Hessian computation is now also supported for PolarMACE models.

Example usage:

from mace.calculators import mace_polar

calc = mace_polar(
    model="polar-1-m",        # or "polar-1-l"
    device="cuda",
    default_dtype="float64",  # use float32 for faster MD
)

atoms.info["charge"] = 0
atoms.info["spin"] = 1 # Number of unpaired electrons + 1
atoms.info["external_field"] = [0.0, 0.0, 0.0]
atoms.calc = calc

energy = atoms.get_potential_energy()
forces = atoms.get_forces()

# extra polar outputs
mu       = calc.results["dipole"]                # total molecular dipole
rho      = calc.results["density_coefficients"]  # per-atom multipole coefficients
rho_spin = calc.results["spin_charge_density"]   # spin-resolved density
charges  = calc.results["charges"]               # per-atom monopoles

Model weights and documentation: mace_polar_1 release on mace-foundations (ASL license), polar_mace docs.

See code and test.

References: Batatia, Baldwin, Kuryla, Hart, Kasoar, Elena, Moore, Gawkowski, Shi, Kapil, Kourtis, Magdau, Csányi, "MACE-POLAR-1: A Polarisable Electrostatic Foundation Model for Molecular Chemistry", arXiv:2602.19411

William J. Baldwin, Ilyes Batatia, Martin Vondrák, Johannes T. Margraf, Gábor Csányi, "Design Space of Self-Consistent Electrostatic Machine Learning Interatomic Potentials", arXiv:2603.14700

🔬 Models

use_edge_irreps_first Flag

Exposed use_edge_irreps_first as a model construction argument, allowing the first interaction layer to use edge-only irreps for a smaller, faster initial layer. Models that do not inherit from the base MACE class accept the flag for compatibility but ignore it.

Example usage:

python mace_run_train.py \
  --train_file=data.xyz \
  --use_edge_irreps_first=True

See code.

GatedEquivariantBlock for torch.compile

Replaced the e3nn Gate layer with a layout-aware GatedEquivariantBlock so the entire forward graph compiles cleanly under torch.compile. This is the building block that unlocks the reduce-overhead and max-autotune paths described below.

See code and test.

⚡ Performance Improvements

torch.compile with reduce-overhead / CUDA graphs

Stabilised the torch.compile path: padded data-dicts so the compiled graph is shape-stable across batches, added isolated-system padding tools, and enabled reduce-overhead and max-autotune modes for end-to-end CUDA-graph capture. Combined with the new GatedEquivariantBlock, this unlocks fully-compiled inference and training.

Example usage:

from mace.calculators import mace_mp

calc = mace_mp(model="mh-1", device="cuda", compile_mode="reduce-overhead")

See code and test.

OpenEquivariance and Hybrid Kernels

Added an OpenEquivariance backend (enable_oeq=True) and a hybrid mode (enable_cueq=True, enable_oeq=True) that uses cuEquivariance for symmetric contractions and linear layers and OpenEquivariance for the convolution tensor product. The calculator automatically falls back to hybrid when cuEq's conv_fusion cannot handle non-uniform edge irreps.

Example usage:

from mace.calculators import mace_mp

calc_oeq    = mace_mp(model="mh-1", device="cuda", enable_oeq=True)
calc_hybrid = mace_mp(model="mh-1", device="cuda",
                      enable_cueq=True, enable_oeq=True)

See code and test.

TorchSim Interface

Added MaceTorchSimModel, a torch-sim ModelInterface wrapper around any MACE or PolarMACE model. Supports cuEq + OEQ + hybrid acceleration, padded torch.compile, multi-system batching, head selection, and PolarMACE-specific outputs (charges, dipole, density_coefficients, …).

Example usage:

import torch
from mace.calculators.mace_torchsim import MaceTorchSimModel

model = MaceTorchSimModel(
    model="mh-1.model",
    device=torch.device("cuda"),
    dtype=torch.float64,
    head="omat_pbe",
    enable_cueq=True,
)
results = model(state)   # state: ts.SimState

See code and test.

🔧 Training and Infrastructure Improvements

Validation Log Head Identification

The validation log now includes the active head name for stress-virials and energy-only loss tables, making multi-head training runs easier to read.

See code.

Preserved Error Table Type

The --error_table choice given by the user is no longer silently rewritten when stress / virial computation is enabled; the original RMSE/MAE form is kept, and stress/virials are appended to the appropriate variant.

See code.

Multi-GPU Fine-tuning Robustness

Fixed a multi-GPU fine-tuning crash caused by zero-size bias parameters under DistributedDataParallel.

See code.

🐛 Bug Fixes and Improvements

• Atomic foundation-model downloads — checkpoint downloads now stream to a .part sibling and os.replace to the final path on success; interrupted downloads no longer leave a truncated file that breaks torch.load on the next call (#1363).
• torch_sim.SimState compatibility — SimState dropped charge/spin from default fields in 0.6+; MaceTorchSimModel now reads them via total_charge/total_spin or extras and skips them when absent.
• mace-off URL handling — fixed broken raw URL resolution.
• dtype pollution when constructing the calculator with a non-default dtype.
• Positions gradient restored in AtomicData.
• PolarMACE finetuning compatibility after the upstream develop merge.
• Foundation finetuning robustness: corrected source range, external-field 3D vector format, removed center-of-mass for the external field, fixed node_feats handling.

🔨 Infrastructure

• Python 3.8 support dropped (lmdb requirement) — minimum supported version is now Python 3.9.
• Python 3.11 compatibility for the torch-sim install path.
• Pre-commit / Black / pylint cleanup on develop.
• New tests: test_polar_models.py, test_polar_cueq.py, test_torchsim.py, test_padding.py, test_gate.py.
• CI improvements: deferred foundation-model construction at test time, balanced pytest splits, foundation-test ID preservation, dtype-leak guard in test_foundations.

🙏 Acknowledgments

We thank all contributors to this release!

Full Changelog: https://github.com/ACEsuit/mace/compare/v0.3.15...v0.3.16

For detailed documentation and examples, visit our GitHub repository and documentation.