Ctrl+K

Note

Go to the end to download the full example code.

UMA Foundation Model: NVE / NVT / NPT Molecular Dynamics#

UMA (Universal Models for Atoms, from fairchem-core) is a multi-task foundation model: a single checkpoint ships heads for molecules (OMol), bulk crystals (OMat), catalysis (OC20), direct air capture (ODAC), and molecular crystals (OMC). UMAWrapper exposes one task head at a time and computes conservative forces (and, for periodic tasks, stress) via autograd.

This example drives UMA through nvalchemi’s velocity-Verlet (NVE), BAOAB Langevin (NVT), or Martyna-Tobias-Klein (NPT) integrator. Observation is delegated to built-in hooks: LoggingHook logs potential energy, temperature, and max force, while EnergyDriftMonitorHook checks NVE energy conservation. Energy drift under NVE is the standard check that the conservative-force path is wired correctly end-to-end; the NPT mode additionally exercises the stress path (the barostat couples the cell to the model’s stress tensor).

Key concepts demonstrated#

  • Loading a UMA checkpoint and selecting a task head via UMAWrapper.from_checkpoint(..., task_name=...).

  • Reaching fairchem’s torch.compile path through the inference_settings argument ("turbo").

  • Driving energy/force dynamics (NVE/NVT) and stress-coupled dynamics (NPT, periodic omat head) from the same wrapper.

  • Evolving a batch of replicas (perturbed starts) in a single shared forward pass per step.

  • That UMA needs no NeighborListHook — the predict unit builds its own graph internally, so the wrapper plugs straight into an integrator.

  • Observing a run with built-in LoggingHook and EnergyDriftMonitorHook instead of a hand-rolled callback.

Setting up UMA#

UMA checkpoints live in the gated facebook/UMA HuggingFace repo. To run this example end-to-end:

  1. Install the optional dependency (its torch pin conflicts with the mace / cuXX extras, so use a dedicated environment):

    uv venv .venv-uma && uv sync --extra uma
# or, with pip:  pip install 'nvalchemi-toolkit[uma]'

  2. Request access at https://huggingface.co/facebook/UMA (one-time approval).

  3. Create a read token at https://huggingface.co/settings/tokens and authenticate, either with:

    huggingface-cli login

    or by exporting it for the shell session:

    export HF_TOKEN=hf_xxx

The first from_checkpoint() call downloads and caches the checkpoint (under ~/.cache/fairchem); subsequent runs reuse the cache. If the checkpoint cannot be loaded (no access, missing token, or fairchem-core not installed) the example prints guidance and skips the trajectory so it stays safe to run in a documentation build.

from __future__ import annotations

import logging
import sys
import time
from typing import Any

import numpy as np
import torch
from ase import Atoms
from ase.build import bulk
from ase.data import atomic_masses as ASE_ATOMIC_MASSES

from nvalchemi.data import AtomicData, Batch
from nvalchemi.dynamics.hooks import EnergyDriftMonitorHook, LoggingHook

# KB_EV is an internal helper used here only to set MB velocities.
from nvalchemi.dynamics.hooks._utils import KB_EV
from nvalchemi.dynamics.integrators.npt import NPT
from nvalchemi.dynamics.integrators.nve import NVE
from nvalchemi.dynamics.integrators.nvt_langevin import NVTLangevin
from nvalchemi.models.uma import UMAWrapper

logging.basicConfig(level=logging.INFO)

# 1 GPa in nvalchemi's internal pressure unit (eV/ų), matching the eV / Å
# energy / length convention used for stress.
_GPA_TO_EV_PER_A3 = 6.241509074e-3

Configuration#

Edit these constants to explore different checkpoints, tasks, systems, and ensembles.

  • CHECKPOINT — uma-s-1p1 / uma-s-1p2 (small) or uma-m-1p1 (medium). See fairchem.core.calculate.pretrained_mlip.available_models.

  • TASK — which task head to expose: omat (crystals), omol (molecules), oc20, odac, omc.

  • INFERENCE_SETTINGS — "default" or "turbo". "turbo" enables fairchem’s torch.compile + TF32 + MoLE merge; it is faster after a one-time compilation but assumes a fixed atomic composition for the whole run (true for MD) and shifts numerics slightly.

  • N_REPLICAS — how many copies of the system to evolve in one batch. Replicas start from slightly perturbed positions + independent velocities, so their trajectories diverge while sharing a single batched forward pass.

CHECKPOINT = "uma-s-1p1"
TASK = "omat"  # omat | omol | oc20 | odac | omc
SYSTEM = "bcc-fe"  # bcc-fe | diamond | propane
ENSEMBLE = "nve"  # nve | nvt | npt
INFERENCE_SETTINGS = "turbo"  # "default" | "turbo"
N_REPLICAS = 4
N_STEPS = 300
DT_FS = 0.5
TEMPERATURE_K = 300.0
FRICTION = 0.01  # 1/fs — NVT Langevin only
PRESSURE_GPA = 0.0  # target pressure — NPT only
BAROSTAT_TIME_FS = 1000.0  # τ_P cell-coupling time — NPT only
THERMOSTAT_TIME_FS = 100.0  # τ_T — NVT Nosé-Hoover (inside NPT) only
SEED = 42
LOG_EVERY = 20
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Molecular systems need the OMol head — force it if a periodic task was left set.
if SYSTEM == "propane" and TASK != "omol":
    print(f"[info] forcing TASK='omol' for SYSTEM={SYSTEM!r}")
    TASK = "omol"

# NPT couples the cell to the stress tensor, so it needs a periodic system and
# a task head that produces stress (the periodic tasks, e.g. ``omat``).
if ENSEMBLE == "npt":
    if SYSTEM == "propane":
        print("[info] NPT needs a periodic system — switching SYSTEM='bcc-fe'")
        SYSTEM = "bcc-fe"
    if TASK == "omol":
        print("[info] NPT needs a stress-producing task — switching TASK='omat'")
        TASK = "omat"

Build the batch#

Each replica is the same system with Maxwell-Boltzmann velocities at TEMPERATURE_K (zero net momentum). Replica 0 is the unperturbed reference; the rest get a small position jitter and an independent velocity seed, so the batched trajectories diverge. Batch.from_data_list packs them into one graph batch that shares a single model forward each step.

_PROPANE_POSITIONS = np.array(
    [
        [0.0000, 0.0000, 0.0000],
        [1.5260, 0.0000, 0.0000],
        [2.0330, 1.4360, 0.0000],
        [-0.5093, 1.0222, 0.0000],
        [-0.5093, -0.5111, 0.8853],
        [-0.5093, -0.5111, -0.8853],
        [2.0319, -0.5111, 0.8853],
        [2.0319, -0.5111, -0.8853],
        [3.1193, 1.4360, 0.0000],
        [1.6763, 1.9471, 0.8853],
        [1.6763, 1.9471, -0.8853],
    ]
)
_PROPANE_NUMBERS = [6, 6, 6, 1, 1, 1, 1, 1, 1, 1, 1]


def build_ase_atoms(system: str) -> Atoms:
    """Return an ASE ``Atoms`` object for the requested benchmark system."""
    if system == "bcc-fe":
        return bulk("Fe", "bcc", a=2.87, cubic=True) * (2, 2, 2)
    if system == "diamond":
        return bulk("C", "diamond", a=3.567, cubic=True) * (2, 2, 2)
    if system == "propane":
        atoms = Atoms(numbers=_PROPANE_NUMBERS, positions=_PROPANE_POSITIONS, pbc=False)
        atoms.info["charge"] = 0
        atoms.info["spin"] = 1
        return atoms
    raise ValueError(f"Unknown system {system!r}")


def replica_data(
    atoms: Atoms, temperature_k: float, device: torch.device, seed: int, jitter: float
) -> AtomicData:
    """One replica: positions optionally jittered, fresh MB velocities."""
    n = len(atoms)
    numbers_np = np.asarray(atoms.get_atomic_numbers())
    numbers = torch.as_tensor(numbers_np, dtype=torch.long, device=device)
    masses = torch.as_tensor(
        ASE_ATOMIC_MASSES[numbers_np], dtype=torch.float32, device=device
    )

    g = torch.Generator(device="cpu").manual_seed(seed)
    pos = torch.as_tensor(np.asarray(atoms.positions), dtype=torch.float32)
    if jitter > 0.0:
        pos = pos + jitter * torch.randn(pos.shape, generator=g)
    pos = pos.to(device)

    # Maxwell-Boltzmann velocities (sigma = sqrt(kB T / m)), zero net momentum.
    sigma = torch.sqrt(torch.as_tensor(KB_EV * temperature_k) / masses.cpu())
    vel = (torch.randn(n, 3, generator=g) * sigma.unsqueeze(-1)).to(device)
    vel -= vel.mean(dim=0, keepdim=True)

    kwargs: dict[str, Any] = {
        "positions": pos,
        "atomic_numbers": numbers,
        "atomic_masses": masses,
        "velocities": vel,
        "forces": torch.zeros_like(pos),
        "energy": torch.zeros(1, 1, device=device, dtype=torch.float32),
    }
    if bool(np.any(atoms.pbc)):
        kwargs["cell"] = torch.as_tensor(
            np.asarray(atoms.cell.array), dtype=torch.float32, device=device
        ).unsqueeze(0)
        kwargs["pbc"] = torch.as_tensor(
            np.asarray(atoms.pbc), dtype=torch.bool, device=device
        ).reshape(1, 3)
        # Stress placeholder: the dynamics loop copy_()s model outputs into
        # existing batch fields, so NPT's barostat needs this slot up front.
        kwargs["stress"] = torch.zeros(1, 3, 3, device=device, dtype=torch.float32)

    return AtomicData(**kwargs)


atoms = build_ase_atoms(SYSTEM)
# Replica 0 is unperturbed; replicas 1..N-1 get a small position jitter so the
# batched trajectories diverge from independent starts.
batch = Batch.from_data_list(
    [
        replica_data(
            atoms, TEMPERATURE_K, DEVICE, seed=SEED + i, jitter=0.0 if i == 0 else 0.05
        )
        for i in range(N_REPLICAS)
    ]
)
n_per = batch.num_nodes // batch.num_graphs
print(
    f"UMA {ENSEMBLE.upper()} | system={SYSTEM} | task={TASK} | "
    f"checkpoint={CHECKPOINT} | inference_settings={INFERENCE_SETTINGS} | "
    f"device={DEVICE}"
)
print(
    f"Batch: {batch.num_graphs} replicas × {n_per} atoms = {batch.num_nodes} total, "
    f"pbc={bool(torch.any(batch.pbc)) if batch.pbc is not None else False}"
)

Load the UMA model#

from_checkpoint resolves the registered name, downloads the checkpoint via HuggingFace Hub (gated — see the setup notes above), and pins the task head. inference_settings="turbo" routes through fairchem’s torch.compile path. If the checkpoint can’t be loaded we exit early with the setup hint, so the rest of the script needn’t be guarded.

try:
    load_t0 = time.perf_counter()
    model = UMAWrapper.from_checkpoint(
        CHECKPOINT,
        task_name=TASK,
        device=str(DEVICE),
        inference_settings=INFERENCE_SETTINGS,
    )
except Exception as exc:  # noqa: BLE001
    sys.exit(
        f"Could not load UMA ({exc}). Install 'nvalchemi-toolkit[uma]', request "
        "access to the gated 'facebook/UMA' repo, and authenticate via "
        "'huggingface-cli login' or HF_TOKEN."
    )
print(
    f"Loaded UMAWrapper in {time.perf_counter() - load_t0:.1f}s "
    f"(cutoff={model.cutoff:.2f} Å)"
)

Run with built-in observation + drift-monitor hooks#

Rather than hand-rolling observation, register two built-in hooks:

  • LoggingHook — logs per-step potential energy, temperature, and max force; here a custom writer prints them.

  • EnergyDriftMonitorHook — for NVE, warns when total-energy drift exceeds the per-atom-per-step budget.

No neighbor-list hook is needed (UMA builds its own graph). NPT couples the cell to the model’s stress tensor, exercising the periodic stress path.

if ENSEMBLE == "nve":
    integrator = NVE(model=model, dt=DT_FS)
elif ENSEMBLE == "nvt":
    integrator = NVTLangevin(
        model=model,
        dt=DT_FS,
        temperature=TEMPERATURE_K,
        friction=FRICTION,
        random_seed=SEED,
    )
else:  # npt — needs forces + stress from the model
    integrator = NPT(
        model=model,
        dt=DT_FS,
        temperature=TEMPERATURE_K,
        pressure=PRESSURE_GPA * _GPA_TO_EV_PER_A3,
        barostat_time=BAROSTAT_TIME_FS,
        thermostat_time=THERMOSTAT_TIME_FS,
        pressure_coupling="isotropic",
    )

# LoggingHook emits one row per replica (graph) at each logged step; the
# custom writer prints energy / temperature / fmax with the replica index.
header = f"{'step':>6} {'rep':>4} {'PE(eV)':>16} {'T(K)':>9} {'fmax(eV/Å)':>12}"
log_hook = LoggingHook(
    backend="custom",
    frequency=LOG_EVERY,
    writer_fn=lambda step, rows: print(
        "\n".join(
            f"{step:6d} {int(r['graph_idx']):4d} {r['energy']:16.6f} "
            f"{r['temperature']:9.2f} {r['fmax']:12.4f}"
            for r in rows
        )
    ),
)
integrator.register_hook(log_hook)

# Drift is only meaningful under NVE (NVT/NPT exchange energy by design).
if ENSEMBLE == "nve":
    integrator.register_hook(
        EnergyDriftMonitorHook(
            threshold=1e-4, metric="per_atom_per_step", action="warn"
        )
    )

# NPT evolves the cell; record per-replica volumes up front (always periodic).
vol_initial = (
    torch.det(batch.cell.reshape(-1, 3, 3)).abs() if ENSEMBLE == "npt" else None
)

print(f"\nRunning {N_STEPS} {ENSEMBLE.upper()} steps …")
print(header)
print("-" * len(header))
if DEVICE.type == "cuda":
    torch.cuda.synchronize(DEVICE)
t0 = time.perf_counter()
with log_hook:  # context manager flushes the async logging stream on exit
    batch = integrator.run(batch, n_steps=N_STEPS)
if DEVICE.type == "cuda":
    torch.cuda.synchronize(DEVICE)
wall_s = time.perf_counter() - t0

Summary#

print(
    f"\nwall time     : {wall_s:.2f} s ({wall_s * 1e3 / max(1, N_STEPS):.2f} ms/step)"
)
if ENSEMBLE == "nve":
    print(
        "drift         : monitored by EnergyDriftMonitorHook (warns above 1e-4 eV/atom/step)"
    )
elif ENSEMBLE == "nvt":
    print("drift         : not gated for NVT (thermostat absorbs drift)")
else:  # npt — per-replica cell response confirms the stress path ran
    vol_final = torch.det(batch.cell.reshape(-1, 3, 3)).abs()
    pct = (100.0 * (vol_final - vol_initial) / vol_initial).tolist()
    print(
        f"cell volume   : mean {vol_initial.mean():.3f} → {vol_final.mean():.3f} ų "
        f"@ {PRESSURE_GPA:.3f} GPa target"
    )
    print("per-replica Δ : " + "  ".join(f"{p:+.2f}%" for p in pct))

Gallery generated by Sphinx-Gallery

On this page
Show Source