Hello World with UMA - FAIR Chemistry Documentation

This tutorial takes you from an ASE structure to a UMA prediction. You will load one pretrained model, use it for two chemistry domains, and learn how to substitute your own structure.

The four-step workflow¶

Every basic calculation follows the same pattern:

Represent the system as an ASE Atoms object.
Load a pretrained UMA model.
Choose the task matching the system and attach a FAIRChemCalculator.
Ask ASE for energies and forces or use an ASE simulation method.

Load UMA once¶

from fairchem.core import FAIRChemCalculator, pretrained_mlip

predictor = pretrained_mlip.get_predict_unit(
    "uma-s-1p2p1",
    device="cuda",
)

Warp 1.18.0 initialized:
   CUDA Toolkit 13.4, Driver 13.1
   Devices:
     "cpu"      : "x86_64"
     "cuda:0"   : "Tesla T4" (16 GiB, sm_75, mempool enabled)
   Kernel cache:
     /home/runner/.cache/warp/1.18.0

The predictor contains the shared UMA model. The calculator created for each system supplies the domain-specific task.

Example 1: calculate a molecular spin gap¶

For the omol task, set the molecule’s total charge and spin multiplicity in atoms.info. Here we compare singlet and triplet states of CH₂.

from ase.build import molecule

singlet = molecule("CH2_s1A1d")
singlet.info.update({"charge": 0, "spin": 1})
singlet.calc = FAIRChemCalculator(predictor, task_name="omol")

triplet = molecule("CH2_s3B1d")
triplet.info.update({"charge": 0, "spin": 3})
triplet.calc = FAIRChemCalculator(predictor, task_name="omol")

spin_gap = triplet.get_potential_energy() - singlet.get_potential_energy()
print(f"Triplet-singlet energy difference: {spin_gap:.3f} eV")

WARNING:root:Model is being compiled this might take a while for the first time

W1007 02:50:32.501000 8074 site-packages/torch/_logging/_internal.py:1345] [0/0] Profiler record function <class 'torch.autograd.profiler.record_function'> will be ignored

W1007 02:50:36.015000 8074 site-packages/torch/_inductor/utils.py:1953] [3/0_1] Not enough SMs to use max_autotune_gemm mode

/home/runner/work/_tool/Python/3.12.15/x64/lib/python3.12/site-packages/torch/_inductor/lowering.py:2352: FutureWarning: `torch._prims_common.check` is deprecated and will be removed in the future. Please use `torch._check*` functions instead.
  check(

WARNING:root:The UMA fast path (merge_mole + compile) is only available for fixed composition, task, charge, and spin. This is optimized for MD applications. Falling back to a less optimized version for subsequent evaluations. Reason: 'Spin differs: tensor([3]) vs tensor([1], device='cuda:0')'.
Use inference_settings='batch' for heterogeneous batched evaluations.

Triplet-singlet energy difference: -0.358 eV

Example 2: relax an inorganic crystal¶

The omat task predicts stress as well as energy and forces, so ASE can relax both the atoms and the unit cell.

from ase.build import bulk
from ase.filters import FrechetCellFilter
from ase.optimize import FIRE

iron = bulk("Fe")
iron.calc = FAIRChemCalculator(predictor, task_name="omat")

optimizer = FIRE(FrechetCellFilter(iron), logfile=None)
optimizer.run(fmax=0.05, steps=100)

print(f"Relaxed energy: {iron.get_potential_energy():.3f} eV")
print("Relaxed cell (Å):")
print(iron.cell)

Relaxed energy: -8.270 eV
Relaxed cell (Å):
Cell([[-1.423450445536107, 1.4234505044951635, 1.4234504600673523], [1.4234504533735954, -1.4234505068412615, 1.4234504702509376], [1.4234504379469857, 1.42345049925214, -1.4234504624134499]])

Try your own structure¶

ASE reads many common chemistry file formats, including XYZ, CIF, POSCAR, and trajectory files. Replace the filename and task below with values appropriate for your system.

from ase.io import read

atoms = read("my-structure.xyz")

# Required for molecules evaluated with the omol task.
atoms.info.update({"charge": 0, "spin": 1})

atoms.calc = FAIRChemCalculator(predictor, task_name="omol")
energy = atoms.get_potential_energy()
forces = atoms.get_forces()

Choose the task by scientific domain, not merely by which task accepts the structure:

Domain	Task	Start here
Molecules and polymers	`omol`	OMol25
Inorganic materials	`omat`	OMat24
Heterogeneous catalysis	`oc20`, `oc22`, or `oc25`	Catalysis datasets
Molecular crystals	`omc`	OMC25
MOFs and direct air capture	`odac`	ODAC datasets

Next steps¶

Explore UMA’s capabilities by domain.
Review task limitations in the UMA model guide.
Learn about inference settings in the ASE calculator guide.
Run NVT and NPT simulations with the molecular dynamics guide.
Try the playground.