lamm-mit/uma

UMA Skill

Run structure relaxation, single-point energy calculations, and phonon property calculations using UMA (Universal Materials Accelerator), a machine-learned interatomic potential from Meta FAIR. UMA is a fast alternative to DFT, achieving near-DFT accuracy at a fraction of the cost.

Built on the fairchem framework.

Prerequisites

• fairchem-core installed (pip install fairchem-core)
• phonopy installed (pip install phonopy) — for phonon calculations
• HF_TOKEN environment variable set (gated model on HuggingFace)
• GPU recommended (CUDA); CPU works but is much slower

Scripts

uma_relax.py — Relax a crystal structure (CLI)

python3 {baseDir}/scripts/uma_relax.py \
  --structure path/to/structure.cif \
  --relax-cell --pressure 150 \
  --format json

| Parameter | Description | |-----------|-------------| | --structure | Path to CIF, POSCAR, or XYZ structure file | | --mp-id | Materials Project ID (fetches structure automatically) | | --model | uma-s-1p1, uma-s-1p2, uma-m-1p1 (default: uma-m-1p1) | | --task | omat (default), omol, omc, oc20, odac | | --device | cuda (default) or cpu | | --relax-cell | Enable full cell relaxation | | --pressure | External pressure in GPa (default: 0). Implies --relax-cell | | --fmax | Force convergence threshold in eV/A (default: 0.05) | | --steps | Max optimizer steps (default: 200) | | --output-cif | Path to save relaxed structure as CIF | | --format | summary | json |

Python API

For batch workflows (screening, phonon), use the Python API directly:

Loading the model (do once, reuse for all calculations)

from fairchem.core import pretrained_mlip, FAIRChemCalculator

predictor = pretrained_mlip.get_predict_unit("uma-m-1p1", device="cuda")
calc = FAIRChemCalculator(predictor, task_name="omat")

Structure relaxation

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

atoms = bulk("Si")  # or: ase.io.read("structure.cif")
atoms.calc = calc

# Relax with external pressure (1 GPa = 1/160.21766208 eV/A^3)
pressure_gpa = 150.0
pressure_ev_per_A3 = pressure_gpa / 160.21766208
opt = FIRE(FrechetCellFilter(atoms, scalar_pressure=pressure_ev_per_A3))
opt.run(fmax=0.05, steps=500)

energy = atoms.get_potential_energy()  # eV
volume = atoms.get_volume()            # A^3

Phonon calculation (finite displacement method)

import phonopy
from phonopy.structure.atoms import PhonopyAtoms
import numpy as np

# Convert ASE Atoms to PhonopyAtoms
ph_atoms = PhonopyAtoms(
    symbols=atoms.get_chemical_symbols(),
    cell=atoms.get_cell(),
    scaled_positions=atoms.get_scaled_positions(),
)

# Create phonopy object with supercell
ph = phonopy.Phonopy(ph_atoms, supercell_matrix=np.diag([2, 2, 2]))
ph.generate_displacements(distance=0.01, is_diagonal=False)

# Compute forces on each displaced supercell using UMA
force_sets = []
for scell in ph.supercells_with_displacements:
    from ase import Atoms as ASEAtoms
    ase_scell = ASEAtoms(
        symbols=scell.symbols, positions=scell.positions,
        cell=scell.cell, pbc=True,
    )
    ase_scell.calc = FAIRChemCalculator(predictor, task_name="omat")
    forces = ase_scell.get_forces()
    forces -= forces.mean(axis=0)  # acoustic sum rule
    force_sets.append(forces)

ph.forces = force_sets
ph.produce_force_constants()
ph.symmetrize_force_constants()

# Get phonon frequencies
ph.run_mesh([8, 8, 8])
frequencies = ph.get_mesh_dict()["frequencies"]  # (n_qpoints, n_bands) in THz

# Check dynamic stability (no imaginary modes)
min_freq = frequencies.min()
dynamically_stable = min_freq > -0.5  # THz threshold

# Thermal properties
ph.run_thermal_properties(t_min=0, t_max=600, t_step=100)
tp = ph.get_thermal_properties_dict()
# tp["temperatures"], tp["free_energy"], tp["entropy"], tp["heat_capacity"]

Auto supercell size

def auto_supercell(n_atoms):
    if n_atoms <= 4: return [3, 3, 3]
    elif n_atoms <= 16: return [2, 2, 2]
    elif n_atoms <= 48: return [2, 2, 1]
    else: return [1, 1, 1]

Structure Generation with pymatgen

Build from spacegroup + Wyckoff positions

from pymatgen.core import Structure, Lattice

# LaH10 clathrate (Fm-3m, SG 225)
structure = Structure.from_spacegroup(
    225, Lattice.cubic(5.1),
    ["La", "H", "H"],
    [[0, 0, 0], [0.25, 0.25, 0.25], [0.118, 0.118, 0.118]],
)

# CaH6 sodalite (Im-3m, SG 229)
structure = Structure.from_spacegroup(
    229, Lattice.cubic(3.54),
    ["Ca", "H"],
    [[0, 0, 0], [0.25, 0, 0.5]],
)

Element substitution

from pymatgen.transformations.standard_transformations import SubstitutionTransformation
sub = SubstitutionTransformation({"La": "Y"})
new_structure = sub.apply_transformation(structure)

Convert pymatgen ↔ ASE

from pymatgen.io.ase import AseAtomsAdaptor
atoms = AseAtomsAdaptor.get_atoms(structure)      # pymatgen → ASE
structure = AseAtomsAdaptor.get_structure(atoms)   # ASE → pymatgen

Convex Hull Stability

from mp_api.client import MPRester
from pymatgen.analysis.phase_diagram import PhaseDiagram, PDEntry
from pymatgen.core import Composition
import os

with MPRester(os.environ["MP_API_KEY"]) as mpr:
    entries = mpr.get_entries_in_chemsys("La-H")

my_entry = PDEntry(Composition("LaH10"), total_energy_eV, name="UMA-LaH10")
pd = PhaseDiagram(list(entries) + [my_entry])
e_above_hull = pd.get_e_above_hull(my_entry)  # eV/atom; 0 = on hull

GPU and SLURM

UMA requires a GPU. Structure your code with a GPU check at the top. When no GPU is available, write a SLURM script that re-runs the SAME script on a GPU node. The script file is at agent_scripts/agent_code.py relative to the project root.

import torch, subprocess, os, sys, json

# GPU check — must be at the TOP of the script, before any UMA imports
if not torch.cuda.is_available():
    print("No GPU — submitting to SLURM", file=sys.stderr)
    venv = os.environ.get("VIRTUAL_ENV", "")
    script_path = os.path.abspath(sys.argv[0])  # path to this script
    slurm = f"""#!/bin/bash
#SBATCH --partition=venkvis-h100
#SBATCH --gres=gpu:1
#SBATCH --mem=64G
#SBATCH --time=04:00:00
#SBATCH --output=slurm-%j.out
#SBATCH --error=slurm-%j.err
source {venv}/bin/activate
export HF_TOKEN="{os.environ.get('HF_TOKEN', '')}"
export MP_API_KEY="{os.environ.get('MP_API_KEY', '')}"
cd {os.getcwd()}
{sys.executable} {script_path}
"""
    with open("submit.sh", "w") as f:
        f.write(slurm)
    result = subprocess.run(["sbatch", "submit.sh"], capture_output=True, text=True)
    job_id = result.stdout.strip().split()[-1] if result.returncode == 0 else None
    print(json.dumps({"status": "SUBMITTED_TO_SLURM", "job_id": job_id}))
    sys.exit(0)

# === GPU code below (only runs on GPU node) ===
from fairchem.core import pretrained_mlip, FAIRChemCalculator
# ... rest of computation

For job dependency (ensure job B runs after job A):

sbatch --dependency=afterok:<job_A_id> submit_B.sh

Verification Checklist

Before submitting or executing code that uses UMA, verify:

• [ ] Model loaded with pretrained_mlip.get_predict_unit("uma-m-1p1", device="cuda") — NOT pretrained_mlip("UMA") or pretrained_mlip.load()
• [ ] Calculator created with FAIRChemCalculator(predictor, task_name="omat") — NOT FAIRChemCalculator(model)
• [ ] Cell filter is FrechetCellFilter from ase.filters — NOT ExpCellFilter from ase.constraints
• [ ] Pressure conversion: pressure_gpa / 160.21766208 — NOT * 0.0006242 or other approximations
• [ ] SLURM partition is venkvis-h100 or venkvis-a100 — NOT gpu or standard
• [ ] Environment activation: source <venv>/bin/activate — NOT conda activate or module load
• [ ] Venv path from os.environ.get("VIRTUAL_ENV", "") — NOT hardcoded
• [ ] API keys from os.environ.get("HF_TOKEN") and os.environ.get("MP_API_KEY")
• [ ] Output printed as JSON to stdout, progress/errors to stderr

Models

| Model | Parameters | Speed | Accuracy | |-------|-----------|-------|----------| | uma-s-1p1 | 6.6M active / 150M total | Fast | Good | | uma-s-1p2 | 6.6M active / 290M total | Fast | Better | | uma-m-1p1 | 50M active / 1.4B total | Slower | Best |

uma_screen.py — Batch structure relaxation and stability screening

Relax all CIF files in a directory with UMA at multiple pressures, compute formation energies, and check 0 GPa stability against the Materials Project convex hull. No hardcoded prototypes — structures come from the structure-enumeration skill or any other source.

Typical workflow:

1. Use materials skill to find prototype structures from MP
2. Use structure-enumeration skill to generate candidates by metal substitution
3. Use uma_screen.py to relax all candidates and assess stability

# Relax all CIFs in the default enumeration directory
python3 {baseDir}/scripts/uma_screen.py --format json

# From a custom directory
python3 {baseDir}/scripts/uma_screen.py \
  --structures-dir ./my_candidates \
  --pressures 0,150 \
  --format json

If no GPU is available, the script auto-submits to SLURM (venkvis-h100).

Screening Parameters

| Parameter | Description | |-----------|-------------| | --structures-dir | Directory of CIF files to relax (default: ~/.scienceclaw/enumerated_structures) | | --pressures | Pressures in GPa (default: 0,150) | | --after-job | SLURM job ID to wait for before starting (ensures prior job completes first) | | --model | UMA checkpoint (default: uma-m-1p1) | | --device | cuda (default) or cpu | | --fmax | Force convergence threshold (default: 0.05) | | --steps | Max optimizer steps (default: 200) | | --output-dir | Directory for relaxed CIFs (default: ./uma_screen_output) | | --format | json or summary | | --dry-run | Show plan without running |

Pipeline Steps

1. Scan --structures-dir for CIF files
2. Identify metal elements in each structure
3. Relax elemental references + H2 at 0 GPa
4. Relax each structure at each pressure
5. Compute formation energy per atom
6. At 0 GPa, query MP convex hull for energy above hull
7. Rank by formation energy, save relaxed CIFs

Output (JSON)

{
  "status": "COMPLETED",
  "model": "uma-m-1p1",
  "ranking": [
    {
      "rank": 1,
      "formula": "LaH10",
      "prototype": "LaH10-type",
      "formation_energy_eV_per_atom": -0.1234,
      "converged": true,
      "e_above_hull_eV_per_atom_0GPa": 0.045
    }
  ],
  "candidates": [ ... ],
  "reference_energies": { ... }
}

---

Python API (for scripting)

You can also use UMA directly in Python without the CLI script:

from fairchem.core import pretrained_mlip, FAIRChemCalculator
from ase.build import bulk
from ase.optimize import FIRE
from ase.filters import FrechetCellFilter

# Load model (do this once, reuse for many structures)
predictor = pretrained_mlip.get_predict_unit("uma-m-1p1", device="cuda")
calc = FAIRChemCalculator(predictor, task_name="omat")

# Build or load structure
atoms = bulk("Si")  # or: ase.io.read("structure.cif")
atoms.calc = calc

# Relax with optional external pressure
pressure_gpa = 150.0
pressure_ev_per_A3 = pressure_gpa / 160.21766208
filtered = FrechetCellFilter(atoms, scalar_pressure=pressure_ev_per_A3)
opt = FIRE(filtered)
opt.run(fmax=0.05, steps=500)

# Read results
energy = atoms.get_potential_energy()        # eV
forces = atoms.get_forces()                   # eV/A
volume = atoms.get_volume()                   # A^3

Typical Workflow with HPC

Since UMA requires a GPU, a typical workflow on Artemis is:

1. Prepare your structure files (CIF/POSCAR) on the login node
2. Write a SLURM script that runs uma_relax.py (see hpc skill for script format)
3. Submit with sbatch to venkvis-h100 or venkvis-a100
4. Check status with squeue -j <job_id>
5. Read results from the output JSON after completion

Structure Generation with pymatgen

To enumerate candidate structures for screening, use pymatgen directly:

from pymatgen.core import Structure, Lattice

# Build a structure from spacegroup + Wyckoff positions
structure = Structure.from_spacegroup(
    225,                              # Fm-3m
    Lattice.cubic(5.1),               # lattice parameter
    ["La", "H", "H"],                 # species at each site
    [[0, 0, 0],                       # 4a: metal
     [0.25, 0.25, 0.25],             # 8c: H
     [0.118, 0.118, 0.118]],         # 32f: H (clathrate cage)
)
structure.to("LaH10.cif", fmt="cif")

To substitute metals in a prototype:

from pymatgen.transformations.standard_transformations import SubstitutionTransformation

# Replace La with Y in a prototype
sub = SubstitutionTransformation({"La": "Y"})
new_structure = sub.apply_transformation(structure)

Convex Hull Stability Analysis

To check thermodynamic stability at 0 GPa, compare against the Materials Project convex hull:

from mp_api.client import MPRester
from pymatgen.analysis.phase_diagram import PhaseDiagram, PDEntry
from pymatgen.core import Composition

with MPRester("YOUR_MP_API_KEY") as mpr:
    # Fetch all known phases in the La-H system
    entries = mpr.get_entries_in_chemsys("La-H")

# Add your UMA-computed entry
my_entry = PDEntry(Composition("LaH10"), total_energy_eV, name="UMA-LaH10")
all_entries = list(entries) + [my_entry]

# Build phase diagram and check stability
pd = PhaseDiagram(all_entries)
e_above_hull = pd.get_e_above_hull(my_entry)  # eV/atom; 0 = on hull (stable)