FAIR-Chem¶

FAIR-Chem: Meta's UMA and eSEN machine-learned potentials, through fairchem-core.

The codec calls FAIR-Chem's predict unit, the model interface its ASE calculator wraps, in the calling process; install it next to qccompute with pip install fairchem-core. The pretrained weights are gated on Hugging Face: request access at https://huggingface.co/facebook/UMA (UMA) or https://huggingface.co/facebook/OMol25 (eSEN), then log in with hf auth login, or point model.method at a local checkpoint. Each model loads once per process and is reused.

Inputs map to FAIR-Chem as follows:

model.method is a pretrained checkpoint or a path to a checkpoint file. FAIR-Chem ships three architectures: UMA, one universal model with a head per dataset (uma-s-1p2p1, uma-s-1p2, uma-s-1p1, uma-m-1p1); eSEN, one model per dataset (esen-sm-conserving-all-omol, esen-md-direct-all-omol, and OC25 and ODAC25 variants); and AllScAIP (allscaip-md-conserving-all-omol, ...). fairchem.core.calculate.pretrained_mlip.available_models lists them all. Models have no basis set.
structure.charge and structure.multiplicity reach the model through atoms.info["charge"] and atoms.info["spin"]; only omol heads read them.
keywords:
task: the dataset head to evaluate. UMA models have several: omol (molecules, the default), omat, oc20, odac, and omc. eSEN and AllScAIP checkpoints are trained on one dataset, which is used unless task says otherwise. The head used is recorded in extras["fairchem_task"].
device: "cpu" or "cuda" (CUDA when available by default). 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.
inference_settings: one of fairchem's presets by name ("default", "turbo", "batch", ...) or a dictionary of InferenceSettings fields. The codec's default is {"compile": False, "merge_mole": False, "activation_checkpointing": False}: the model as trained, ready in seconds for any molecule. fairchem's own "default" compiles the model, which pays off only for long runs on one system: compiling takes tens of seconds in every process and again for each new molecule size.
cmdline_args and files are not accepted.

Calculation types: energy, gradient, and hessian (by automatic differentiation; the codec requests it with predict_untrained_hessian). Models predict energies in eV; the codec converts them, and the forces and Hessian, to Hartree and bohr.

FAIR-Chem 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/fairchem/single_point.py

"""FAIR-Chem single point: energy, gradient, or Hessian from UMA or eSEN.

Requires: ``pip install qccompute fairchem-core``, and access to the gated weights:
request it at https://huggingface.co/facebook/UMA (or facebook/OMol25 for eSEN), then
run ``hf auth login``. FAIR-Chem 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. Later
calculations in the same process reuse the loaded model.

How the codec talks to FAIR-Chem (https://fair-chem.github.io/):

- ``model.method`` is a pretrained model ("uma-s-1p2p1", "uma-s-1p1", "uma-m-1p1",
  "esen-sm-conserving-all-omol", ...) or a path to a checkpoint. There is no basis
  set.
- ``task`` picks the UMA head: "omol" (molecules, the default) reads charge and
  multiplicity from ``structure``; "omat", "oc20", "odac", and "omc" are for
  materials and catalysis.
- ``inference_settings`` defaults to the uncompiled model, ready in seconds for any
  molecule. fairchem's own "default" preset compiles the model, which takes tens of
  seconds in every process and again for each new molecule size: worth it only for
  long runs on one system.
- Hessians come from automatic differentiation.
- Energies come back in Hartree and gradients in Hartree/bohr, converted from
  FAIR-Chem'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,
    multiplicity=1,
)

prog_input = ProgramInput(
    program="fairchem",
    calctype=calctype,
    structure=structure,
    model=Model(method="uma-s-1p2p1"),  # or "uma-m-1p1", "esen-sm-conserving-all-omol"
    keywords={
        "task": "omol",
        "device": "cpu",  # or "cuda"; leave out to use a GPU when present
        # A fairchem preset by name, or InferenceSettings fields:
        # "inference_settings": "turbo",
        # "inference_settings": {"compile": False, "base_precision_dtype": "float64"},
    },
)

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("fairchem-core version:", results.provenance.program_version)

Many structures at once: batches¶

examples/fairchem/many_calculations.py

"""FAIR-Chem for many structures at once: a bond scan evaluated in batches.

Requires: what examples/fairchem/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("fairchem").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="fairchem",
        calctype=CalcType.gradient,  # or CalcType.energy
        structure=water(r),
        model=Model(method="uma-s-1p2p1"),
        keywords={
            "task": "omol",
            # "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")