FAIRChem v2 UMA for Multidomain Atomistic Simulation across Molecules, Catalysts, Materials, Vibrations, and Molecular Dynamics


In this tutorial, we explore FAIRChem v2 and the UMA universal machine-learning interatomic potential as a unified framework for atomistic simulation across molecular chemistry, catalysis, and inorganic materials. We configure an environment, authenticate with Hugging Face to access the gated UMA model weights, and initialize task-specific calculators for the omol, oc20, and omat domains. We then apply the same pretrained potential to a broad set of computational chemistry workflows, including single-point energy and force prediction, molecular geometry optimization, spin-state comparison, reaction-energy estimation, vibrational analysis, surface adsorption, crystal-cell relaxation, equation-of-state fitting, molecular dynamics, and potential-energy surface scanning. Throughout the tutorial, we integrate FAIRChem with the Atomic Simulation Environment to manage atomic structures, optimizers, constraints, thermodynamic calculations, and trajectory analysis while using GPU acceleration whenever it is available.

import importlib.util, subprocess, sys, os
def _pip(*pkgs):
   subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *pkgs])
if importlib.util.find_spec("fairchem") is None:
   print(">> Installing fairchem-core, ase, and helpers (takes ~2-4 min)...")
   _pip("fairchem-core", "ase", "matplotlib", "huggingface_hub")
   print(">> Installation done.")
else:
   print(">> fairchem already installed.")
from huggingface_hub import login, whoami
def hf_authenticate():
   token = None
   try:
       from google.colab import userdata
       token = userdata.get("HF_TOKEN")
   except Exception:
       pass
   token = token or os.environ.get("HF_TOKEN")
   try:
       whoami()
       print(">> Already authenticated with Hugging Face.")
       return
   except Exception:
       pass
   if token is None:
       from getpass import getpass
       token = getpass("Paste your Hugging Face access token: ").strip()
   login(token=token)
   print(">> Hugging Face login OK.")
hf_authenticate()
import numpy as np
import torch
import matplotlib.pyplot as plt
from fairchem.core import pretrained_mlip, FAIRChemCalculator
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f">> Using device: {DEVICE}")
MODEL = "uma-s-1p2"
predictor = pretrained_mlip.get_predict_unit(MODEL, device=DEVICE)
calc_mol  = FAIRChemCalculator(predictor, task_name="omol")
calc_cat  = FAIRChemCalculator(predictor, task_name="oc20")
calc_mat  = FAIRChemCalculator(predictor, task_name="omat")
print(f">> Loaded {MODEL} with omol / oc20 / omat calculators.")

We install the required FAIRChem, ASE, visualization, and Hugging Face dependencies while ensuring the setup remains safe to rerun in Google Colab. We authenticate with Hugging Face to access the gated UMA model weights and automatically detect whether GPU acceleration is available. We then load the UMA predictor and create separate calculators for molecular, catalysis, and materials simulations.

from ase.build import molecule
from ase import Atoms
print("\n" + "="*70)
print("SECTION 2: Single-point energetics of water (omol task)")
print("="*70)
h2o = molecule("H2O")
h2o.info.update({"charge": 0, "spin": 1})
h2o.calc = calc_mol
E_h2o = h2o.get_potential_energy()
F_h2o = h2o.get_forces()
print(f"E(H2O)            = {E_h2o:.4f} eV")
print(f"Max |force|       = {np.abs(F_h2o).max():.4f} eV/A")
def atom_energy(symbol, spin):
   a = Atoms(symbol, positions=[[0, 0, 0]])
   a.info.update({"charge": 0, "spin": spin})
   a.calc = calc_mol
   return a.get_potential_energy()
E_O = atom_energy("O", spin=3)
E_H = atom_energy("H", spin=2)
E_atomization = -(E_h2o - E_O - 2 * E_H)
print(f"Atomization energy of H2O = {E_atomization:.3f} eV "
     f"(experiment ~ 9.5 eV incl. ZPE effects)")
from ase.optimize import LBFGS
print("\n" + "="*70)
print("SECTION 3: Relaxing a deliberately distorted water molecule")
print("="*70)
h2o_bad = molecule("H2O")
h2o_bad.positions[1] += [0.25, -0.10, 0.05]
h2o_bad.info.update({"charge": 0, "spin": 1})
h2o_bad.calc = calc_mol
opt = LBFGS(h2o_bad, logfile=None)
energies_opt = []
opt.attach(lambda: energies_opt.append(h2o_bad.get_potential_energy()))
opt.run(fmax=0.01, steps=200)
d_OH = h2o_bad.get_distance(0, 1)
ang  = h2o_bad.get_angle(1, 0, 2)
print(f"Converged in {opt.get_number_of_steps()} steps")
print(f"O-H bond length   = {d_OH:.3f} A   (expt ~0.958 A)")
print(f"H-O-H angle       = {ang:.1f} deg (expt ~104.5 deg)")
plt.figure(figsize=(5, 3.2))
plt.plot(energies_opt, "o-")
plt.xlabel("Optimizer step"); plt.ylabel("Energy (eV)")
plt.title("H2O geometry optimization"); plt.tight_layout(); plt.show()

We use the molecular UMA calculator to evaluate the energy, atomic forces, and atomization energy of a water molecule. We define isolated hydrogen and oxygen reference atoms with the correct spin multiplicities to construct the atomization-energy calculation. We then distort the water geometry, relax it with the LBFGS optimizer, and analyze the converged bond length, bond angle, and energy trajectory.

print("\n" + "="*70)
print("SECTION 4: CH2 singlet-triplet gap (UMA is spin-aware!)")
print("="*70)
singlet = molecule("CH2_s1A1d"); singlet.info.update({"charge": 0, "spin": 1})
triplet = molecule("CH2_s3B1d"); triplet.info.update({"charge": 0, "spin": 3})
singlet.calc = FAIRChemCalculator(predictor, task_name="omol")
triplet.calc = FAIRChemCalculator(predictor, task_name="omol")
gap = triplet.get_potential_energy() - singlet.get_potential_energy()
print(f"E(triplet) - E(singlet) = {gap:.3f} eV  "
     f"(negative => triplet ground state; expt ~ -0.39 eV)")
print("\n" + "="*70)
print("SECTION 5: Reaction energy of CH4 + 2 O2 -> CO2 + 2 H2O")
print("="*70)
def relaxed_energy(name, spin=1):
   m = molecule(name)
   m.info.update({"charge": 0, "spin": spin})
   m.calc = FAIRChemCalculator(predictor, task_name="omol")
   LBFGS(m, logfile=None).run(fmax=0.02, steps=200)
   return m.get_potential_energy()
E = {
   "CH4": relaxed_energy("CH4"),
   "O2":  relaxed_energy("O2", spin=3),
   "CO2": relaxed_energy("CO2"),
   "H2O": relaxed_energy("H2O"),
}
dE_rxn = (E["CO2"] + 2*E["H2O"]) - (E["CH4"] + 2*E["O2"])
print(f"Delta E (electronic) = {dE_rxn:.2f} eV = {dE_rxn*96.485:.0f} kJ/mol")
print("Experimental combustion enthalpy ~ -890 kJ/mol (ZPE/thermal not included here)")
from ase.vibrations import Vibrations
print("\n" + "="*70)
print("SECTION 6: Vibrational frequencies of relaxed H2O")
print("="*70)
vib = Vibrations(h2o_bad, name="vib_h2o")
vib.run()
freqs = np.real(vib.get_frequencies())
real_modes = [f for f in freqs if f > 200]
print("Vibrational modes (cm^-1):", ", ".join(f"{f:.0f}" for f in real_modes))
print("Experimental H2O: 1595 (bend), 3657 (sym stretch), 3756 (asym stretch)")
print(f"Zero-point energy = {vib.get_zero_point_energy():.3f} eV")
vib.clean()

We compare the singlet and triplet electronic states of methylene to calculate its spin-state energy gap. We relax methane, oxygen, carbon dioxide, and water before combining their predicted energies to estimate the electronic reaction energy of methane combustion. We also perform a finite-difference vibrational analysis of relaxed water to obtain its normal-mode frequencies and zero-point energy.

from ase.build import fcc100, add_adsorbate
from ase.constraints import FixAtoms
print("\n" + "="*70)
print("SECTION 7: CO/Cu(100) relaxation + adsorption energy (oc20 task)")
print("="*70)
slab = fcc100("Cu", size=(3, 3, 3), vacuum=8.0, periodic=True)
slab.set_constraint(FixAtoms(mask=[a.tag > 1 for a in slab]))
add_adsorbate(slab, molecule("CO"), height=2.0, position="bridge")
slab.calc = calc_cat
opt = LBFGS(slab, logfile=None)
opt.run(fmax=0.05, steps=300)
E_slab_ads = slab.get_potential_energy()
print(f"Relaxed CO/Cu(100) in {opt.get_number_of_steps()} steps, "
     f"E = {E_slab_ads:.3f} eV")
clean = fcc100("Cu", size=(3, 3, 3), vacuum=8.0, periodic=True)
clean.set_constraint(FixAtoms(mask=[a.tag > 1 for a in clean]))
clean.calc = FAIRChemCalculator(predictor, task_name="oc20")
LBFGS(clean, logfile=None).run(fmax=0.05, steps=300)
E_clean = clean.get_potential_energy()
co = molecule("CO"); co.info.update({"charge": 0, "spin": 1})
co.calc = FAIRChemCalculator(predictor, task_name="omol")
LBFGS(co, logfile=None).run(fmax=0.02, steps=100)
E_co = co.get_potential_energy()
E_ads = E_slab_ads - E_clean - E_co
print(f"E(clean slab) = {E_clean:.3f} eV, E(CO gas) = {E_co:.3f} eV")
print(f"Adsorption energy (naive cycle) = {E_ads:.3f} eV")
print("(oc20 uses its own DFT reference scheme; for publication-grade numbers")
print(" keep all species within a consistent task/reference framework.)")

We construct a periodic Cu(100) slab, place a carbon monoxide molecule at a bridge adsorption site, and constrain the lower copper layers. We relax the adsorbate–surface system with the OC20 calculator and separately optimize the clean slab and gas-phase carbon monoxide references. We then evaluate a pedagogical adsorption-energy cycle while recognizing that OC20 uses a task-specific energy reference convention.

from ase.build import bulk
from ase.optimize import FIRE
from ase.filters import FrechetCellFilter
from ase.eos import EquationOfState
print("\n" + "="*70)
print("SECTION 8: BCC iron — full cell relaxation and bulk modulus (omat)")
print("="*70)
fe = bulk("Fe", "bcc", a=2.9)
fe.calc = calc_mat
FIRE(FrechetCellFilter(fe), logfile=None).run(fmax=0.02, steps=300)
a_relaxed = fe.cell.cellpar()[0]
print(f"Relaxed BCC Fe lattice constant = {a_relaxed:.3f} A (expt ~2.866 A)")
volumes, energies = [], []
cell0 = fe.get_cell()
for scale in np.linspace(0.94, 1.06, 9):
   s = fe.copy()
   s.set_cell(cell0 * scale, scale_atoms=True)
   s.calc = FAIRChemCalculator(predictor, task_name="omat")
   volumes.append(s.get_volume())
   energies.append(s.get_potential_energy())
eos = EquationOfState(volumes, energies, eos="birchmurnaghan")
v0, e0, B = eos.fit()
from ase.units import GPa as _GPa
B_GPa = B / _GPa
print(f"Equilibrium volume = {v0:.2f} A^3/cell")
print(f"Bulk modulus       = {B_GPa:.0f} GPa (expt ~170 GPa for Fe)")
plt.figure(figsize=(5, 3.2))
plt.plot(volumes, energies, "o", label="UMA points")
vfit = np.linspace(min(volumes), max(volumes), 100)
plt.plot(vfit, [eos.func(v, *eos.eos_parameters) for v in vfit], "-", label="BM fit")
plt.xlabel("Volume (A^3)"); plt.ylabel("Energy (eV)")
plt.title("BCC Fe equation of state"); plt.legend(); plt.tight_layout(); plt.show()
from ase import units
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
print("\n" + "="*70)
print("SECTION 9: 0.5 ps Langevin MD of a water molecule at 300 K")
print("="*70)
md_atoms = molecule("H2O")
md_atoms.info.update({"charge": 0, "spin": 1})
seed = int(np.random.randint(0, np.iinfo(np.int32).max))
md_predictor = pretrained_mlip.get_predict_unit(MODEL, device=DEVICE, seed=seed)
md_atoms.calc = FAIRChemCalculator(md_predictor, task_name="omol")
MaxwellBoltzmannDistribution(md_atoms, temperature_K=300)
dyn = Langevin(md_atoms, timestep=0.5 * units.fs,
              temperature_K=300, friction=0.01 / units.fs)
times, temps, epots, d_oh1 = [], [], [], []
def log_md():
   t = dyn.get_number_of_steps() * 0.5
   times.append(t)
   temps.append(md_atoms.get_temperature())
   epots.append(md_atoms.get_potential_energy())
   d_oh1.append(md_atoms.get_distance(0, 1))
dyn.attach(log_md, interval=5)
dyn.run(steps=1000)
print(f"MD done.  = {np.mean(temps[10:]):.0f} K, "
     f" = {np.mean(d_oh1):.3f} A")
fig, ax = plt.subplots(1, 3, figsize=(12, 3.2))
ax[0].plot(times, temps);  ax[0].set_title("Temperature (K)")
ax[1].plot(times, epots);  ax[1].set_title("Potential energy (eV)")
ax[2].plot(times, d_oh1);  ax[2].set_title("O-H bond length (A)")
for a in ax: a.set_xlabel("time (fs)")
plt.tight_layout(); plt.show()

We relax the atomic positions and simulation cell of BCC iron using the materials-domain UMA calculator and a Frechet cell filter. We sample energies across a range of compressed and expanded volumes, fit a Birch–Murnaghan equation of state, and estimate the equilibrium volume and bulk modulus. We also run Langevin molecular dynamics for water at 300 K and track its temperature, potential energy, and O–H bond length over time.

print("\n" + "="*70)
print("SECTION 10: O-H bond stretch scan in water")
print("="*70)
distances = np.linspace(0.7, 2.5, 25)
pes = []
for d in distances:
   a = molecule("H2O")
   a.info.update({"charge": 0, "spin": 1})
   vec = a.positions[1] - a.positions[0]
   a.positions[1] = a.positions[0] + vec / np.linalg.norm(vec) * d
   a.calc = FAIRChemCalculator(predictor, task_name="omol")
   pes.append(a.get_potential_energy())
pes = np.array(pes) - min(pes)
plt.figure(figsize=(5.5, 3.4))
plt.plot(distances, pes, "o-")
plt.axvline(0.958, ls="--", c="gray", label="expt r_e")
plt.xlabel("O-H distance (A)"); plt.ylabel("Relative energy (eV)")
plt.title("O-H stretch PES from UMA"); plt.legend()
plt.tight_layout(); plt.show()
print("\n" + "="*70)
print("TUTORIAL COMPLETE!")
print("="*70)
print("""
Next steps to explore:
 * Swap MODEL to "uma-m-1p1" for higher accuracy (needs more GPU memory).
 * Try task_name="odac" with a MOF CIF, or "omc" for molecular crystals.
 * Larger MD: fairchem supports multi-GPU inference via workers=N
   (pip install fairchem-core[extras]).
 * Docs: https://fair-chem.github.io/
""")

We scan the potential-energy surface of water by systematically stretching one O–H bond across a selected distance range. We evaluate the molecular energy at each geometry, normalize the energies relative to the minimum, and visualize the resulting dissociation profile. We conclude the tutorial by identifying higher-accuracy UMA models, additional chemical domains, and larger multi-GPU simulations as possible extensions.

In conclusion, we built a complete atomistic simulation workflow around FAIRChem v2 and demonstrated how UMA provides a shared learned potential across chemically distinct domains without requiring a separate model for every task. We used molecular calculations to evaluate energies, forces, atomization behavior, spin gaps, reaction energetics, vibrational modes, and bond-stretch profiles; we used the catalysis domain to relax CO on a Cu(100) surface and examine adsorption energetics; and we used the materials domain to relax BCC iron and estimate its bulk modulus from an equation-of-state fit. We also ran Langevin molecular dynamics to observe finite-temperature structural and energetic fluctuations over time. By combining UMA inference with ASE structure builders, optimizers, filters, vibrational tools, and molecular-dynamics utilities, we established a reusable foundation for extending the workflow to larger molecules, catalytic interfaces, crystalline materials, metal-organic frameworks, molecular crystals, and higher-accuracy UMA model variants.


Check out the Full Code hereAlso, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us


Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.



Source link

  • Related Posts

    Induction Labs Photon-1 Simulates Desktops, Plays Checkers, and Models Billiard Physics From One Pretraining Run

    Most agents that learn from video need to know what action produced each frame. Induction Labs is arguing that this requirement is the bottleneck. Last week, they released imagination models,…

    Sakana AI Releases Fugu-Cyber: An Orchestration Model Reporting 86.9% on CyberGym and 72.1% on CTI-REALM

    Sakana AI has released Fugu-Cyber (model ID is fugu-cyber-v1.0), a cybersecurity-specialized addition to its Fugu orchestration family. It is not just a new frontier model.…

    Leave a Reply

    Your email address will not be published. Required fields are marked *