Garden's Almanac of Matter Models MACE-MH-1
January 2025
| Checkpoint | Params | Polaris | Sophia | Perlmutter | Delta | Frontier | Della |
|---|---|---|---|---|---|---|---|
| mace-mh-1 | ● | ● | — | ● | ● | ● | |
| mace:custom | ● | ● | — | — | ● | — |
● verified, last 30 days ○ installed, not recently verified — not installed
No cluster has installed any checkpoint of this model yet.
Running this model
This model has task heads. Select which output layer is active at inference time with
setup_kwargs={"head": "omat_pbe"} —
there is no default, and the calculator raises an error if no head is given.
| Task head | Trained against |
|---|---|
| omat_pbe | OMat24 — inorganic crystals |
| mp_pbe_refit_add | — |
| omol | OMol25 — molecules |
| spice_wB97M | SPICE — molecules |
| oc20_usemppbe | OC20 — surfaces/interfaces |
| matpes_r2scan | MatPES — inorganic crystals |
1# from a job or interactive session on a supported cluster:
2from rootstock import RootstockCalculator
3
4with RootstockCalculator(
5 cluster=YOUR_CLUSTER_ID, # eg, "polaris", "sophia"
6 checkpoint="mace-mh-1",
7 device="cuda",
8 setup_kwargs={"head": "omat_pbe"}, # required: select a task head
9) as calc:
10 # model is now running in subprocess on compute node
11 atoms.calc = calc
12 atoms.get_potential_energy()
Environments
Rootstock runs each model family inside an isolated Python environment defined by a single file. This environment file includes the specific dependencies needed, plus a setup() function that loads the model and returns an ASE calculator. These files are usually almost identical for a given model family, but because of cluster-specific quirks (eg, an old CUDA driver) the dependencies and setup code can vary a bit.
mace_env.py
1# /// script
2# requires-python = ">=3.11"
3# dependencies = [
4# # 0.3.15+ needed for the mh-1 registry entry (matpes needs 0.3.13,
5# # omol needs 0.3.14, mpa-0 needs 0.3.10).
6# "mace-torch>=0.3.15",
7# "ase>=3.22",
8# # 2.4.1 is explicitly unsupported by mace-torch.
9# "torch>=2.4.0,!=2.4.1,<2.10",
10# ]
11# ///
12"""MACE env — hosts MACE-MP-0, MACE-OFF23, MPA-0, MATPES, MH-1, and OMOL checkpoints.
13
14All ship in the same `mace-torch` package, so they share an environment.
15Upstream-string routing in CHECKPOINTS: an `off:` prefix routes to mace_off()
16and an `omol:` prefix to mace_omol() (float64, molecules only); an `mh:`
17prefix marks a multi-head model (float64, per the MACE-MH-1 model card).
18
19Multi-head checkpoints select a head via the `head` kwarg on setup()
20(setup_kwargs={"head": ...} / --kwarg head=...), named by upstream's training
21corpus — see MH1_HEADS; omat_pbe is the default.
22
23The OMOL checkpoint expects `charge` and `spin` in `atoms.info`.
24"""
25
26CHECKPOINTS = {
27 "mace-mp-0-small": "small",
28 "mace-mp-0-medium": "medium",
29 "mace-mp-0-large": "large",
30 "mace-off23-small": "off:small",
31 "mace-off23-medium": "off:medium",
32 "mace-off23-large": "off:large",
33 # Only a medium MPA-0 has been released, but upstream names the weights
34 # file mace-mpa-0-medium.model — keep the size explicit like mace-mp-0.
35 "mace-mpa-0-medium": "medium-mpa-0",
36 "mace-matpes-r2scan-0": "mace-matpes-r2scan-0",
37 # One entry per weights file: MH-1's heads are selected by setup(head=...).
38 "mace-mh-1": "mh:mh-1",
39 # Only the extra-large OMOL model has been released.
40 "mace-omol-0-extra-large": "omol:extra_large",
41 # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).
42 "mace:custom": None,
43}
44
45# MH-1's heads, named by training corpus. The released weights file is the
46# authority (mace_select_head --list_heads; the model card also lists a
47# rgd1_b3lyp head, but that shipped only in mh-0 — ACEsuit/mace#1462).
48# Validated here because upstream only warns on an unknown head and silently
49# falls back to the last one.
50MH1_HEADS = (
51 "omat_pbe",
52 "omol",
53 "spice_wB97M",
54 "oc20_usemppbe",
55 "mp_pbe_refit_add",
56 "matpes_r2scan",
57)
58
59
60def setup(checkpoint: str, device: str = "cuda", head: str | None = None):
61 arg = CHECKPOINTS[checkpoint]
62 if arg.startswith("mh:"):
63 head = head or "omat_pbe"
64 if head not in MH1_HEADS:
65 raise ValueError(f"unknown head {head!r}; expected one of {', '.join(MH1_HEADS)}")
66 from mace.calculators import mace_mp
67
68 return mace_mp(model=arg[3:], device=device, default_dtype="float64", head=head)
69 if head is not None:
70 raise ValueError(f"'head' selects a head of a multi-head model; {checkpoint} has one head")
71 if arg.startswith("off:"):
72 from mace.calculators import mace_off
73
74 return mace_off(model=arg[4:], device=device, default_dtype="float32")
75 if arg.startswith("omol:"):
76 from mace.calculators import mace_omol
77
78 return mace_omol(model=arg[5:], device=device, default_dtype="float64")
79 from mace.calculators import mace_mp
80
81 return mace_mp(model=arg, device=device, default_dtype="float32")
82
83
84def setup_from_path(path: str, device: str = "cuda", default_dtype: str = "float32", head: str | None = None):
85 # Custom checkpoints (`:custom` ids with user weights): fine-tunes load through
86 # MACECalculator directly — the mp/off dispatch in setup() only exists
87 # to pick which pretrained file to download. `head` is for fine-tunes that
88 # keep multiple heads; single-head weights load without it.
89 from mace.calculators import MACECalculator
90
91 return MACECalculator(model_paths=path, device=device, default_dtype=default_dtype, head=head)
92
Built on Polaris: 2026-08-03
Couldn't load the current environments from Rootstock.