Uma

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

from fairchem.core import pretrained_mlip, FAIRChemCalculator

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

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

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"]

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]

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]],
)

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

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

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

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

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

# 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