A6 One model for materials and molecules: Orb and UMA¶

Objectives

Run the same checks on crystals and on molecules with Orb-v3 and OrbMol, and with UMA 1.2 small.
Pass total charge and spin multiplicity to a molecular calculation.
Compare lattice constants, a spin gap and ionisation energies with their reference data, and measure the cost on one MI250X GCD.

Most universal MLIPs learn one potential energy surface at one level of theory. Two open routes cover both crystals and molecules:

Orb [9] (code): one architecture, separate checkpoints. Orb-v3 trained on OMat24 for crystals; OrbMol [10] trained on OMol25 for molecules, with total charge and spin as inputs. Weights are public.
UMA [1] (code): one checkpoint trained on five datasets; a task name selects the level of theory it reproduces. Weights are gated (licence approval on Hugging Face).

Task	Dataset	Reference level	Orb checkpoint	Use
`omat`	OMat24 [2]	PBE and PBE+U (VASP)	`orb-v3-conservative-inf-omat`	inorganic crystals
`omol`	OMol25 [3]	ωB97M-V/def2-TZVPD (ORCA)	`orb-v3-conservative-omol` (OrbMol)	molecules, with charge and spin
`oc20`	OC20	RPBE	none	catalysis, adsorbates on surfaces
`odac`	ODAC	PBE-D3	none	MOFs with guest molecules
`omc`	OMC	PBE-D3	none	molecular crystals

OrbMol and the UMA omol task read the total charge and spin multiplicity from atoms.info["charge"] and atoms.info["spin"], so the same geometry can be evaluated as a neutral singlet, a cation or a triplet. Both use the multiplicity 2S+1 (the orb-models 0.5.5 README example still shows 0; the conditioner treats 0 as unset). The crystal models ignore both keys.

Set-up on LUMI¶

The CSC module provides Python 3.11 and torch 2.7.1+rocm6.2.4.

orb-models 0.5.5 is the last release for Python 3.11 and torch 2.7; it has Orb-v3 and the first OrbMol. Releases from 0.6 need Python 3.12 and torch 2.8 or later. OrbMol-v2 (pretrained.orbmol_v2, orb-models 0.7, long-range electrostatics through nvalchemiops) is therefore not used here.
fairchem-core 2.21.0 is the newest release with uma-s-1p2 that still pins torch 2.8; later releases pin torch 2.13. It is installed without its dependencies, and the rest from a pinned list that keeps the module’s torch.

On a login node, in a copy of this repository at <SCRATCH>/mlip-uma:

module purge && module use /appl/local/csc/modulefiles/ && module load pytorch
python -m venv --system-site-packages .venv && source .venv/bin/activate
pip install --no-deps fairchem-core==2.21.0
pip install -r examples/uma/requirements.txt
export HF_HOME=$PWD/hf-home FAIRCHEM_CACHE_DIR=$PWD/fairchem-cache CACHED_PATH_CACHE_ROOT=$PWD/cached-path
python examples/uma/cli.py prefetch --model orb
sbatch --account=<PROJECT> scripts/lumi-uma.sbatch orb

UMA access¶

UMA is gated. Accept the licence on the model page, create a read token in your Hugging Face settings, log in once on the login node, then prefetch and submit with uma:

hf auth login   # stores the token in ~/.cache/huggingface/token
export HF_TOKEN_PATH=$HOME/.cache/huggingface/token
python examples/uma/cli.py prefetch --model uma
sbatch --account=<PROJECT> scripts/lumi-uma.sbatch uma

Code¶

Each backend exposes calculator(task) for omat or omol. Orb loads one checkpoint per task; UMA loads one predict unit and puts a light ASE calculator on top of it for each task. The tasks do not know which backend they run.

models.py

"""Two backends with one interface: `calculator(task)` for task "omat" (crystals) or "omol" (molecules)."""

from __future__ import annotations

import time

import torch


def device() -> str:
    return "cuda" if torch.cuda.is_available() else "cpu"


class Backend:
    """Loads checkpoints on first use and records the load time of each."""

    checkpoints: dict[str, str] = {}

    def __init__(self, dtype: str = "float64"):
        self.dtype = dtype
        self.load_s: dict[str, float] = {}
        self._calcs = {}

    def checkpoint(self, task: str) -> str:
        return self.checkpoints[task]

    def calculator(self, task: str):
        if task not in self._calcs:
            start = time.perf_counter()
            self._calcs[task] = self._load(task)
            self.load_s.setdefault(self.checkpoint(task), time.perf_counter() - start)
        return self._calcs[task]

    def _load(self, task: str):
        raise NotImplementedError


class Orb(Backend):
    """Orb-v3 conservative (OMat24) for crystals, OrbMol (OMol25, charge and spin) for molecules."""

    name = "orb"
    checkpoints = {"omat": "orb-v3-conservative-inf-omat", "omol": "orb-v3-conservative-omol"}

    def _load(self, task: str):
        from orb_models.forcefield import pretrained
        from orb_models.forcefield.calculator import ORBCalculator

        precision = "float64" if self.dtype == "float64" else "float32-highest"
        loader = getattr(pretrained, self.checkpoint(task).replace("-", "_"))
        model = loader(device=device(), precision=precision, compile=False)
        return ORBCalculator(model, device=device())


class UMA(Backend):
    """UMA 1.2 small: one predict unit shared by all tasks."""

    name = "uma"
    checkpoints = {"omat": "uma-s-1p2", "omol": "uma-s-1p2"}

    def _load(self, task: str):
        from fairchem.core import FAIRChemCalculator, pretrained_mlip
        from fairchem.core.units.mlip_unit.api.inference import inference_settings_default

        if not hasattr(self, "unit"):
            settings = inference_settings_default()
            settings.base_precision_dtype = getattr(torch, self.dtype)
            self.unit = pretrained_mlip.get_predict_unit(self.checkpoint(task), inference_settings=settings,
                                                         device=device())
        return FAIRChemCalculator(self.unit, task_name=task)


BACKENDS = {"orb": Orb, "uma": UMA}
PACKAGES = {"orb": ("orb-models",), "uma": ("fairchem-core", "e3nn")}


def versions(name: str) -> dict:
    from importlib.metadata import version

    out = {p: version(p) for p in ("torch", "ase", "numpy", *PACKAGES[name])}
    out["checkpoints"] = sorted(set(BACKENDS[name].checkpoints.values()))
    out["gpu"] = torch.cuda.get_device_name(0) if torch.cuda.is_available() else None
    return out

tasks.py

"""The three checks, for any backend: lattice constants (omat), charge and spin (omol), cost per force call."""

from __future__ import annotations

import csv
import json
import time
from pathlib import Path

import numpy as np
from ase.filters import FrechetCellFilter
from ase.optimize import BFGS

import structures

RESULTS = Path(__file__).resolve().parent / "results"
FMAX = 0.005


def outdir(model) -> Path:
    path = RESULTS / model.name
    path.mkdir(parents=True, exist_ok=True)
    return path


def save_csv(model, name: str, rows: list[dict]) -> None:
    with open(outdir(model) / name, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def save_json(model, name: str, data: dict) -> None:
    (outdir(model) / name).write_text(json.dumps(data, indent=2) + "\n")


def attach(atoms, calc) -> None:
    """ASE compares positions and cell but not atoms.info, so drop cached results before a new charge or spin."""
    calc.reset()
    atoms.calc = calc


def relax(atoms, calc, cell: bool = False):
    attach(atoms, calc)
    target = FrechetCellFilter(atoms) if cell else atoms
    BFGS(target, logfile=None).run(fmax=FMAX, steps=500)
    return atoms


def lattice(model) -> list[dict]:
    rows = []
    for name, spec in structures.CRYSTALS.items():
        atoms = relax(structures.crystal(name), model.calculator("omat"), cell=True)
        a = float(np.cbrt(atoms.get_volume()))
        rows.append(dict(crystal=name, model=model.checkpoint("omat"), a_A=round(a, 4), pbe_A=spec["pbe"], experiment_A=spec["experiment"],
                         vs_pbe_pct=round(100 * (a / spec["pbe"] - 1), 2),
                         vs_experiment_pct=round(100 * (a / spec["experiment"] - 1), 2)))
    save_csv(model, "lattice.csv", rows)
    return rows


def state(atoms) -> dict:
    return dict(energy_eV=atoms.get_potential_energy(), bond_A=round(atoms.get_distance(0, 1), 4))


def oxygen(model) -> dict:
    """Triplet ground state, closed-shell singlet, and the doublet cation at the triplet geometry."""
    calc = model.calculator("omol")
    triplet = relax(structures.oxygen(spin=3), calc)
    states = dict(triplet=state(triplet), singlet=state(relax(structures.oxygen(spin=1), calc)))
    cation = structures.oxygen(spin=2, charge=1)
    cation.positions = triplet.positions
    ie = single_point(cation, calc) - states["triplet"]["energy_eV"]
    gap = states["singlet"]["energy_eV"] - states["triplet"]["energy_eV"]
    return dict(states=states, gap_eV=round(gap, 3), experiment_eV=structures.O2_GAP_EV,
                ie_vertical_eV=round(ie, 3), ie_experiment_eV=structures.O2_IE_EV)


def single_point(atoms, calc) -> float:
    attach(atoms, calc)
    return atoms.get_potential_energy()


def geometry(atoms) -> dict:
    return dict(oh_A=round(atoms.get_distance(0, 1), 4), hoh_deg=round(atoms.get_angle(1, 0, 2), 1))


def water(model) -> dict:
    calc = model.calculator("omol")
    neutral = relax(structures.water(), calc)
    e_neutral = neutral.get_potential_energy()
    cation = structures.water(charge=1, spin=2)
    cation.positions = neutral.positions
    vertical = single_point(cation, calc) - e_neutral
    adiabatic = relax(cation, calc).get_potential_energy() - e_neutral
    return dict(vertical_eV=round(vertical, 3), adiabatic_eV=round(adiabatic, 3), experiment_eV=structures.WATER_IE_EV,
                neutral=geometry(neutral), cation=geometry(cation))


def molecules(model) -> dict:
    out = dict(model=model.checkpoint("omol"), o2=oxygen(model), water=water(model))
    save_json(model, "molecules.json", out)
    return out


def seconds_per_call(atoms, calc, calls: int = 20) -> tuple[float, float]:
    """First call (includes warm-up), then the median of later calls on rattled copies."""
    times = []
    for i in range(calls + 1):
        trial = atoms.copy()
        trial.rattle(0.01, seed=i)
        trial.calc = calc
        start = time.perf_counter()
        trial.get_forces()
        times.append(time.perf_counter() - start)
    return times[0], float(np.median(times[1:]))


def timing(model) -> list[dict]:
    systems = {
        "Si 216 atoms (omat)": (structures.crystal("Si").repeat(3), "omat"),
        "Cu 256 atoms (omat)": (structures.crystal("Cu").repeat(4), "omat"),
        "H2O 3 atoms (omol)": (structures.water(), "omol"),
    }
    rows = []
    for label, (atoms, task) in systems.items():
        row = dict(system=label, model=model.checkpoint(task), dtype=model.dtype, load_s="", first_call_s="",
                   per_call_s="")
        try:
            first, median = seconds_per_call(atoms, model.calculator(task))
            row.update(load_s=round(model.load_s[model.checkpoint(task)], 1), first_call_s=round(first, 3),
                       per_call_s=round(median, 4), status="ok")
        except RuntimeError as err:
            row["status"] = str(err).splitlines()[0][:80]
        rows.append(row)
    save_csv(model, f"timing_{model.dtype}.csv", rows)
    return rows

What the job computes¶

omat: relax the conventional cells of Si (diamond) and Cu (fcc), positions and cell together, with BFGS to fmax 0.005 eV/Å.
omol, spin: relax O₂ as a triplet (spin=3, the ground state) and as a singlet (spin=1); the gap is the difference of the two minima. As a control for the charge input, evaluate the O₂⁺ doublet (charge=1, spin=2) at the triplet geometry.
omol, charge: relax neutral water (charge=0, spin=1), then evaluate the cation (charge=1, spin=2) at that geometry for the vertical ionisation energy, and relax it for the adiabatic one.
Cost: model load time, the first force call and the median of 20 later calls, in float32 and float64, one process per precision.

ASE decides whether to recompute from positions, cell and atomic numbers, not from atoms.info. A new charge or spin on the same geometry would return the cached energy, so tasks.attach resets the calculator first.

Stress on a periodic cell uses the cell determinant (torch.det), which fails in float32 on LUMI’s torch 2.7.1+rocm6.2.4 (see A3 MatGL: relaxation, MD, lattice benchmark and training), so the science runs in float64. The float32 timing run records which calls fail rather than stopping.

Reference values¶

Quantity	Reference	Value
Si lattice constant	experiment, 22.5 °C (CODATA 2022) [4]	5.431 Å
Si lattice constant	PBE, all-electron, static lattice [5]	5.470 Å
Cu lattice constant	experiment, 25 °C [6]	3.615 Å
Cu lattice constant	PBE, all-electron, static lattice [5]	3.629 Å
O₂ a¹Δ_g minus X³Σ_g⁻, T_e	experiment [7]	7918.1 cm⁻¹ = 0.982 eV
O₂ first ionisation energy	experiment, evaluated, adiabatic [11]	12.07 eV
H₂O first ionisation energy	experiment, evaluated [8]	12.621 eV (photoelectron vertical value 12.62 eV)

The PBE lattice constants come from the equilibrium volumes of the WIEN2k reference set (data): a = (8V₀)^1/3 for Si and (4V₀)^1/3 for Cu. The crystal models reproduce PBE, so compare them first with PBE; the difference from experiment is mostly the functional, plus thermal expansion and zero-point motion, which a static relaxation leaves out.

The singlet needs care. The a¹Δ_g state of O₂ is not a single closed-shell determinant, and OMol25 ran main-group singlets in restricted Kohn-Sham unless bonds were expected to break [3]. The OMol25 singlet is therefore a closed-shell ωB97M-V state, not the experimental a¹Δ_g, and the reference method itself overestimates the 0.982 eV gap; a model trained on it should too.

Results¶

One MI250X GCD on dev-g, torch 2.7.1+rocm6.2.4, float64; orb-models 0.5.5 and fairchem-core 2.21.0. Each job (both checks and both timing runs) took about 4 min: 3 min 39 s for Orb, 3 min 58 s for UMA.

Crystals¶

Crystal	Orb-v3 (Å)	UMA `omat` (Å)	PBE (Å)	Experiment (Å)
Si	5.4603 (−0.18 % vs PBE)	5.4607 (−0.17 %)	5.470	5.431
Cu	3.6134 (−0.43 %)	3.6154 (−0.37 %)	3.629	3.615

Both models land within 0.5 % of PBE, slightly below it, and within 0.002 Å of each other. Cu then matches experiment, but only because the model error and the PBE overestimate cancel; Si keeps most of the PBE error (+0.5 % against experiment).

Molecules¶

Quantity	OrbMol	UMA `omol`	Reference
O₂ singlet minus triplet (eV)	1.544	1.504	0.982 (experiment, a¹Δ_g)
O₂ bond, triplet / singlet (Å)	1.188 / 1.196	1.197 / 1.196
O₂ ionisation energy, vertical (eV)	12.66	12.71	12.07 (experiment, adiabatic)
H₂O ionisation energy, vertical (eV)	6.07	12.64	12.62 (experiment)
H₂O ionisation energy, adiabatic (eV)	6.06	12.55	12.621 (experiment)

O₂ spin: both models put the singlet 1.5 eV above the triplet, 0.5 eV more than experiment, as expected from a closed-shell singlet reference. The two agree to 0.04 eV, which points at the reference method rather than the models.
O₂ charge: both vertical ionisation energies are about 0.6 eV above the adiabatic experimental value, a plausible size for a vertical value.
Water: UMA is within 0.02 eV of the experimental vertical value, and its cation relaxes (O–H 0.962 to 0.999 Å), lowering the energy by 0.08 eV. OrbMol gives half the ionisation energy and a cation that barely moves. The charge input reaches OrbMol (O₂ is right), so this checkpoint does not describe the water cation. OrbMol-v2 adds explicit electrostatics aimed at charge states like this but needs a newer software stack than the CSC module.

A model that takes charge and spin as inputs will return a number for any combination. Check each charge and spin state against a known reference before using it.

Cost on one MI250X GCD¶

System	Orb load / first / per call (s)	UMA load / first / per call (s)
Si, 216 atoms	8.8 / 1.18 / 0.058	22.4 / 1.82 / 0.137
Cu, 256 atoms	8.8 / 0.10 / 0.097	22.4 / 0.25 / 0.248
H₂O, 3 atoms	1.1 / 0.027 / 0.022	22.4 / 0.078 / 0.061

Orb loads one checkpoint per task (Orb-v3 9 s, OrbMol 1 s); UMA loads one for all tasks (22 s). These are warm loads; the first load in a fresh job step took about 30 s for Orb-v3 and 68 s for UMA.
Per force call, Orb-v3 is 2.3 to 2.6 times faster than UMA on these crystals.
The first call includes a one-off GPU warm-up, paid once per process (Si ran first; Cu reuses it).
Three atoms cost 20 to 60 ms per call: small molecules are bound by kernel launch and Python overhead, not by arithmetic.
In float32 the crystal calls fail with CUDA driver error: 209 (the cell determinant) for both models. OrbMol on the non-periodic molecule runs in float32 at 21 ms per call; for UMA the water call failed too, after the crystal failure in the same process.

Keypoints

Orb covers crystals and molecules with two checkpoints; UMA with one checkpoint and a task name. Either way, each model reproduces its own reference level: PBE for crystals, ωB97M-V for molecules.
Charge and spin multiplicity are inputs in atoms.info; ASE caching ignores them, so reset the calculator between states.
Crystal results agree closely between Orb-v3 and UMA. For molecules, a charge-aware model can still be wrong for one charge state: OrbMol halves the water ionisation energy, UMA reproduces it.

References¶

B. M. Wood et al., UMA: a family of universal models for atoms. arXiv:2506.23971
L. Barroso-Luque et al., OMat24. arXiv:2410.12771
D. S. Levine et al., OMol25. arXiv:2505.08762
CODATA 2022, lattice parameter of silicon, NIST. physics.nist.gov
K. Lejaeghere et al., Reproducibility in density functional theory calculations of solids, Science 351, aad3000 (2016). doi:10.1126/science.aad3000
M. E. Straumanis and L. S. Yu, Acta Cryst. A 25, 676 (1969). doi:10.1107/S0567739469001549
NIST Chemistry WebBook, O₂, constants of diatomic molecules. webbook.nist.gov
NIST Chemistry WebBook, H₂O, gas-phase ion energetics. webbook.nist.gov
B. Rhodes et al., Orb-v3. arXiv:2504.06231
Orbital Materials, OrbMol: extending Orb to molecular systems (2025). orbitalindustries.com
NIST Chemistry WebBook, O₂, gas-phase ion energetics. webbook.nist.gov