ASE interface¶

UPET ships an ASE-compatible calculator, UPETCalculator, that wraps any UPET model and exposes the standard ASE get_potential_energy / get_forces / get_stress API. This page covers the full surface of ASE-driven workflows: basic evaluation, loading from a local checkpoint, non-conservative forces and stresses, uncertainty quantification, rotational averaging, empirical dispersion corrections, electronic DOS prediction and dataset visualization.

Basic usage¶

To perform a single-structure evaluation, build an ASE Atoms object and attach a UPETCalculator. The model name is formed by combining the model family and the size, e.g. pet-mad-s, pet-omat-l. See Available models for the full list.

from upet.calculator import UPETCalculator
from ase.build import bulk

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
calculator = UPETCalculator(model="pet-mad-s", version="1.5.0", device="cpu")
atoms.calc = calculator

energy = atoms.get_potential_energy()
forces = atoms.get_forces()

These ASE methods are ideal for single-structure evaluations, but inefficient when you need to run a pre-defined dataset through the model. For efficient batched evaluation, see Batched evaluation with metatrain and the README_BATCHED.md reference.

If the version argument is not specified, the latest available version of the model is downloaded and used by default:

from upet.calculator import UPETCalculator

# uses the latest version of the PET-MAD-S model by default
calculator = UPETCalculator(model="pet-mad-s", device="cpu")

You can list the available versions for any model with upet.list_upet():

from upet import list_upet

list_upet(model="pet-mad", size="s")  # PET-MAD of size S
list_upet(model="pet-mad")             # all PET-MAD sizes
list_upet()                            # all available UPET models

Loading from a local checkpoint¶

You can download a checkpoint from the HuggingFace repository directly:

wget https://huggingface.co/lab-cosmo/upet/resolve/main/models/pet-mad-s-v1.5.0.ckpt

and then load the calculator from that local file (useful for fine-tuned checkpoints or reproducibility):

from upet.calculator import UPETCalculator

calculator = UPETCalculator(checkpoint_path="pet-mad-s-v1.5.0.ckpt", device="cpu")

Non-conservative (direct) forces and stresses¶

By default, UPET models compute conservative forces and stresses as derivatives of the energy via automatic differentiation. UPET models also support the direct prediction of forces and stresses as separate targets, which typically gives a 2–3x speedup. However, as discussed in this preprint, non-conservative forces and stresses require additional care to avoid undesired effects during molecular dynamics.

To use non-conservative forces and stresses, set non_conservative=True when constructing the calculator:

from upet.calculator import UPETCalculator
from ase.build import bulk

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
calculator = UPETCalculator(
    model="pet-mad-s", version="1.5.0", device="cpu", non_conservative=True
)
atoms.calc = calculator

energy = atoms.get_potential_energy()   # computed as usual
forces = atoms.get_forces()             # predicted as a separate target
stresses = atoms.get_stress()           # predicted as a separate target

More details on making direct-force MD simulations reliable are in the Atomistic Cookbook.

Uncertainty quantification¶

UPET models can also estimate the uncertainty of the energy prediction. This is important when probing the model on data that may be substantially different from its training distribution, or when propagating uncertainty to derived observables (phase transition temperatures, diffusion coefficients, and so on).

Note

Uncertainty quantification is only available for a subset of checkpoints (currently pet-mad-s v1.0.2, pet-mad-xs v1.5.0 and pet-mad-s v1.5.0). See Available models for the full list.

Use the get_energy_uncertainty and get_energy_ensemble methods of UPETCalculator:

from upet.calculator import UPETCalculator
from ase.build import bulk

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
calculator = UPETCalculator(model="pet-mad-s", version="1.5.0", device="cpu")
atoms.calc = calculator

energy = atoms.get_potential_energy()
energy_uncertainty = calculator.get_energy_uncertainty(atoms, per_atom=False)
energy_ensemble = calculator.get_energy_ensemble(atoms, per_atom=False)

Both methods accept a per_atom flag, which controls whether the quantity is reported per atom or for the whole system. For the methodology, see the LLPR and shallow ensembles papers.

Rotational averaging¶

By design, UPET models are not exactly equivariant with respect to rotations and inversions. The equivariance error is typically much smaller than the overall model error, but in some cases (geometry optimization or phonon calculations of highly symmetric structures) it can be beneficial to enforce additional rotational averaging.

Pass a rotational_average_order to the calculator:

from upet.calculator import UPETCalculator
from ase.build import bulk

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
calculator = UPETCalculator(
    model="pet-mad-s",
    version="1.5.0",
    device="cpu",
    rotational_average_order=3,
)
atoms.calc = calculator
energy = atoms.get_potential_energy()
forces = atoms.get_forces()
stresses = atoms.get_stress()

Predictions are averaged over a Lebedev quadrature of the O(3) group at the requested order (here 3). Higher orders are more accurate but costlier.

By default, all transformed structures are evaluated in a single batch, which can be memory-intensive for large systems. Reduce memory pressure by setting rotational_average_batch_size to a smaller value:

from upet.calculator import UPETCalculator

calculator = UPETCalculator(
    model="pet-mad-s",
    version="1.5.0",
    device="cpu",
    rotational_average_order=3,
    rotational_average_batch_size=8,
)

After the evaluation, rotational averaging error statistics are available on the calculator’s results dictionary:

energy_rot_std = calculator.results["energy_rot_std"]
forces_rot_std = calculator.results["forces_rot_std"]
stresses_rot_std = calculator.results["stresses_rot_std"]

Empirical dispersion corrections (D3)¶

You can combine the UPET calculator with the torch-based D3 dispersion correction of pfnet-research, torch-dftd. Install it inside your UPET environment:

pip install torch-dftd

Then combine both calculators with ase.calculators.mixing.SumCalculator:

import torch
from ase.calculators.mixing import SumCalculator
from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
from upet.calculator import UPETCalculator

device = "cuda" if torch.cuda.is_available() else "cpu"

calc_upet = UPETCalculator(model="pet-mad-s", version="1.5.0", device=device)
dft_d3 = TorchDFTD3Calculator(device=device, xc="r2scan", damping="bj")

combined_calc = SumCalculator([calc_upet, dft_d3])
atoms.calc = combined_calc

DOS, Fermi levels and bandgaps (PET-MAD-DOS)¶

The UPET package also exposes the PET-MAD-DOS model for predicting electronic density of states of materials, their Fermi levels and bandgaps, via PETMADDOSCalculator:

from ase.build import bulk
from upet.calculator import PETMADDOSCalculator

atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
pet_mad_dos_calculator = PETMADDOSCalculator(version="latest", device="cpu")

energies, dos = pet_mad_dos_calculator.calculate_dos(atoms)

Per-atom DOS and batched evaluation on a list of structures are also supported:

# per-atom DOS
energies, dos_per_atom = pet_mad_dos_calculator.calculate_dos(atoms, per_atom=True)

# multiple structures
atoms_1 = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
atoms_2 = bulk("C", cubic=True, a=3.55, crystalstructure="diamond")
energies, dos = pet_mad_dos_calculator.calculate_dos(
    [atoms_1, atoms_2], per_atom=False
)

Bandgap and Fermi level predictions are available via dedicated methods:

bandgap = pet_mad_dos_calculator.calculate_bandgap(atoms)
fermi_level = pet_mad_dos_calculator.calculate_efermi(atoms)

You can also re-use a previously computed DOS:

bandgap = pet_mad_dos_calculator.calculate_bandgap(atoms, dos=dos)
fermi_level = pet_mad_dos_calculator.calculate_efermi(atoms, dos=dos)

The same is supported for a list of Atoms:

bandgaps = pet_mad_dos_calculator.calculate_bandgap([atoms_1, atoms_2], dos=dos)
fermi_levels = pet_mad_dos_calculator.calculate_efermi([atoms_1, atoms_2], dos=dos)

Dataset visualization with the PET-MAD featurizer¶

You can use the last-layer features of the PET-MAD model, together with a pre-trained sketch-map dimensionality reduction, to obtain 2D and 3D representations of a dataset (e.g. to identify structural or chemical motifs). It can be used as a stand-alone feature builder or combined with the chemiscope viewer for interactive visualization.

import ase.io
import chemiscope
from upet.explore import PETMADFeaturizer

featurizer = PETMADFeaturizer(version="latest")

# Load structures
frames = ase.io.read("dataset.xyz", ":")

# Compute features only
features = featurizer(frames, None)

# Or create an interactive Jupyter visualization
chemiscope.explore(frames, featurize=featurizer)