MACE

MACE: the MACE foundation models and custom MACE potentials, through mace-torch.

The codec loads MACE's raw PyTorch models (the loaders' return_raw_model=True) and calls them directly in the calling process, as MACE's ASE calculator does internally; install MACE next to qccompute with pip install mace-torch. mace-torch and fairchem-core pin incompatible versions of e3nn, so they need separate environments. Models download on first use, and each is loaded once per process and reused.

Inputs map to MACE as follows:

• model.method names a foundation model as <family> or <family>-<model>, where <model> is the argument MACE's loader takes:

• mace-off: MACE-OFF23 organic force field (mace-off-small, mace-off-medium, mace-off-large; medium by default);

• mace-mp: Materials Project models (mace-mp-medium, mace-mp-medium-mpa-0, ...; see mace.calculators.mace_mp);
• mace-omol: MACE-OMOL-0, trained on OMol25 (mace-omol-extra_large, the default), the one family that reads charge and spin;

or a path to a .model file. Models have no basis set. The codec uses a model's omol head for MACE-OMOL, its only head, or its default head. - mace-omol gets structure.charge and structure.multiplicity through atoms.info["charge"] and atoms.info["spin"]. The other models describe neutral closed-shell molecules only, so the codec raises for anything else. - keywords: device ("cpu", "cuda"; CUDA when available by default) and default_dtype ("float64" or "float32"; MACE's default for the model otherwise: float64 for MACE-OFF and MACE-OMOL, float32 for MACE-MP). Nothing else is accepted. A model on a GPU declares one for qccompute, which gives each calculation of a list its own; an explicit index such as "cuda:1" opts out of that. - cmdline_args and files are not accepted.

Calculation types: energy, gradient, and hessian (by automatic differentiation, for MACE models). MACE predicts energies in eV; the codec converts them, and the forces and Hessian, to Hartree and bohr.

MACE does not move the molecule, so results are in the input's frame. No input field carries atom indices.

Energies and gradients of several structures can be computed as one batch: qccompute.compute does this for a list of inputs that share a model, keywords, and calculation type, evaluating them in one pass instead of one call each. Hessians are computed one structure at a time.

Examples

Each script runs as written. Copy it, delete the options you don't need, and run it with python script.py. They live in the repository's examples/ folder.

Single point: energy, gradient, Hessian

examples/mace/single_point.py

"""MACE single point: energy, gradient, or Hessian from a MACE potential.

Requires: ``pip install qccompute mace-torch``, in an environment without
``fairchem-core`` (the two need different versions of e3nn). MACE is a Python library,
so nothing has to be on PATH: the model runs inside this Python process, on the GPU
when PyTorch sees one. Models download on first use; later calculations in the same
process reuse the loaded model.

How the codec talks to MACE (https://github.com/ACEsuit/mace):

- ``model.method`` is a foundation-model family, optionally with the model argument
  MACE's loader takes: "mace-off" or "mace-off-small"/"-medium"/"-large" (MACE-OFF23
  organic force field), "mace-mp" or "mace-mp-medium-mpa-0" etc. (materials), or
  "mace-omol" (MACE-OMOL-0, trained on OMol25). A path to a ``.model`` file runs your
  own model. There is no basis set.
- Only MACE-OMOL reads charge and multiplicity; MACE-OFF and MACE-MP describe neutral
  closed-shell molecules, and the codec refuses anything else for them.
- ``default_dtype`` is "float64" or "float32" (MACE's default otherwise: float64 for
  MACE-OFF and MACE-OMOL, float32 for MACE-MP).
- Hessians come from automatic differentiation.
- Energies come back in Hartree and gradients in Hartree/bohr, converted from MACE's
  eV and angstrom.
"""

from qccompute import compute
from qcdata import CalcType, Model, ProgramInput, Structure

calctype = CalcType.gradient  # or CalcType.energy, CalcType.hessian

# Coordinates are in bohr (Structure.from_xyz / Structure.open read angstrom files).
structure = Structure(
    symbols=["O", "H", "H"],
    geometry=[[0.0, 0.0, 0.0], [0.524, 1.687, 0.48], [1.146, -0.45, -1.354]],
    charge=0,  # other charges and multiplicities need "mace-omol"
    multiplicity=1,
)

prog_input = ProgramInput(
    program="mace",
    calctype=calctype,
    structure=structure,
    model=Model(method="mace-off-small"),  # or "mace-off", "mace-omol", "mace-mp", ...
    keywords={
        "device": "cpu",  # or "cuda"; leave out to use a GPU when present
        # "default_dtype": "float32",  # faster, less precise
    },
)

output = compute(prog_input)

results = output.results
print("Energy (Hartree):", results.energy)
if results.gradient is not None:
    print("Gradient (Hartree/bohr):\n", results.gradient)
if results.hessian is not None:
    print("Hessian (Hartree/bohr^2):", results.hessian.shape)
print("mace-torch version:", results.provenance.program_version)

Many structures at once: batches

examples/mace/many_calculations.py

"""MACE for many structures at once: a bond scan evaluated in batches.

Requires: what examples/mace/single_point.py requires, and qccompute 0.21.0 or newer.

How batching works:

- Pass ``compute`` a list of inputs. Energy and gradient inputs with the same program,
  model, and keywords are evaluated in batches, one model call for up to
  ``get_codec("mace").max_batch`` (32) structures, in this Python process. On a GPU
  that is many times faster than one call per structure.
- Each structure's results are what computing it alone gives, to the model's precision.
  Structures in one batch may differ in size, composition, and charge.
- Inputs with different models or keywords go in separate batches. Hessians are
  computed one structure at a time.
- A batch that runs out of GPU memory is split in half until it fits. A batch that
  fails otherwise is computed again one structure at a time, so each error belongs to
  its own input.
- The model stays loaded in this process, so the next ``compute`` call (the next step
  of a path optimization, say) does not load it again.
- ``output.execution.extras`` records ``batch_size`` and the structure's
  ``batch_index``; ``logs`` and ``wall_time`` are the whole batch's.
"""

import math
import time

from qccompute import compute
from qcdata import CalcType, Model, ProgramInput, Structure


def water(r: float) -> Structure:
    """Water with one O-H bond stretched to r bohr, the other at 1.81 bohr."""
    angle = math.radians(104.5)
    return Structure(
        symbols=["O", "H", "H"],
        geometry=[
            [0.0, 0.0, 0.0],
            [r, 0.0, 0.0],
            [1.81 * math.cos(angle), 1.81 * math.sin(angle), 0.0],
        ],
    )


distances = [1.4 + 0.1 * i for i in range(20)]  # bohr
inputs = [
    ProgramInput(
        program="mace",
        calctype=CalcType.gradient,  # or CalcType.energy
        structure=water(r),
        model=Model(method="mace-off-small"),
        keywords={
            # "device": "cpu",  # or "cuda"; leave out to use a GPU when present
        },
    )
    for r in distances
]

compute(inputs)  # load the model and warm up, so the timings compare evaluation only

start = time.perf_counter()
outputs = compute(inputs)  # one batch of 20
batched = time.perf_counter() - start

start = time.perf_counter()
alone = [compute(inp) for inp in inputs]  # 20 separate calls, for comparison
separate = time.perf_counter() - start

print("r (bohr)  energy (Hartree)  batch_index/batch_size")
for r, output in zip(distances, outputs, strict=True):
    extras = output.execution.extras
    print(
        f"{r:8.2f}  {output.results.energy:16.8f}  "
        f"{extras['batch_index']}/{extras['batch_size']}"
    )
largest = max(
    abs(a.results.energy - b.results.energy)
    for a, b in zip(alone, outputs, strict=True)
)
print(f"Batched: {batched:.3f} s; one at a time: {separate:.3f} s")
print(f"Largest energy difference from computing alone: {largest:.1e} Hartree")