.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/advanced/09_uma_nve.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_advanced_09_uma_nve.py: 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). :class:`~nvalchemi.models.uma.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** :class:`~nvalchemi.hooks.NeighborListHook` — the predict unit builds its own graph internally, so the wrapper plugs straight into an integrator. * Observing a run with built-in :class:`~nvalchemi.dynamics.hooks.LoggingHook` and :class:`~nvalchemi.dynamics.hooks.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 :meth:`~nvalchemi.models.uma.UMAWrapper.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. .. GENERATED FROM PYTHON SOURCE LINES 83-113 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 114-130 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. .. GENERATED FROM PYTHON SOURCE LINES 130-163 .. code-block:: Python 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" .. GENERATED FROM PYTHON SOURCE LINES 164-171 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. .. GENERATED FROM PYTHON SOURCE LINES 171-270 .. code-block:: Python _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}" ) .. GENERATED FROM PYTHON SOURCE LINES 271-278 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. .. GENERATED FROM PYTHON SOURCE LINES 278-298 .. code-block:: Python 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} Å)" ) .. GENERATED FROM PYTHON SOURCE LINES 299-310 Run with built-in observation + drift-monitor hooks --------------------------------------------------- Rather than hand-rolling observation, register two built-in hooks: * :class:`~nvalchemi.dynamics.hooks.LoggingHook` — logs per-step potential energy, temperature, and max force; here a custom writer prints them. * :class:`~nvalchemi.dynamics.hooks.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. .. GENERATED FROM PYTHON SOURCE LINES 310-373 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 374-376 Summary ------- .. GENERATED FROM PYTHON SOURCE LINES 376-393 .. code-block:: Python 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)) .. _sphx_glr_download_examples_advanced_09_uma_nve.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 09_uma_nve.ipynb <09_uma_nve.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 09_uma_nve.py <09_uma_nve.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 09_uma_nve.zip <09_uma_nve.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_