API reference¶

Auto-generated reference for UPET’s Python API

Top-level functions¶

Get a metatomic AtomisticModel for a UPET MLIP.

Parameters:

model (str | None) – name of the UPET model. Required when not using checkpoint_path, or when checkpoint_path has non-standard naming.
size (str | None) – size of the UPET model. Required when not using checkpoint_path, or when checkpoint_path has non-standard naming.
version (str | Version | None) – version of the UPET model.
checkpoint_path (str | Path | None) – path to a checkpoint file to load the model from. If the filename follows standard naming (e.g., “pet-mad-s-v1.0.2.ckpt”), model/size/version are extracted automatically, while the model and version parameters are ignored.

Return type:

AtomisticModel

upet.list_upet(*, model: str | None = None, size: str | None = None, print_summary: bool = True) → List[dict][source]¶

List available UPET models, sizes, and versions.

When called without arguments, returns all available model/size/version combinations. When model is given, filters to that model. When both model and size are given, filters to that specific combination.

Parameters:

model (str | None) – Base model name (e.g., “pet-mad”, “pet-omat”). If None, lists all available models.
size (str | None) – Model size (e.g., “s”, “m”, “l”). If None and model is given, lists all sizes for that model.
print_summary (bool) – Whether to print a human-readable summary to stdout. Defaults to True.

Returns:

A list of dictionaries, each with keys "model", "size", and "version".

Return type:

List[dict]

Save the UPET model to a TorchScript file. These files can be used with LAMMPS and other tools to run simulations without Python.

Parameters:

model (str | None) – name of the UPET model.
size (str | None) – size of the UPET model.
version (str | None) – UPET version to use. Defaults to the latest stable version.
checkpoint_path (str | Path | None) – path to a checkpoint file to load the model from.
output (str | Path | None) – path for the output model. Defaults to “{model}-{size}-v{version}.pt” or “model.pt” for non-standard checkpoint names.

Calculators¶

ASE Calculator for universal MLIPs based on the PET architecture.

Parameters:

model (str | None) –
PET-MLIP model to use. Required when not using checkpoint_path. Can be one of the following:
- ”pet-mad-xs”: PET-MAD model (size “xs”, materials, surfaces and molecules, r2SCAN; versions 1.6.0 and 1.5.0)
- ”pet-mad-s”: PET-MAD model (size “s”, materials, surfaces and molecules, r2SCAN; versions 1.6.0 and 1.5.0)
- ”pet-mad-m”: PET-MAD model (size “m”, materials, surfaces and molecules, r2SCAN; version 1.6.0)
- ”pet-omat-xs”: PET-OMat model (size “xs”, materials, PBE)
- ”pet-omat-s”: PET-OMat model (size “s”, materials, PBE)
- ”pet-omat-m”: PET-OMat model (size “m”, materials, PBE)
- ”pet-omat-l”: PET-OMat model (size “l”, materials, PBE)
- ”pet-omat-xl”: PET-OMat model (size “xl”, materials, PBE)
- ”pet-oam-l”: PET-OAM model (size “l”, materials, Materials-Project-consistent PBE)
- ”pet-oam-xl”: PET-OAM model (size “xl”, materials, Materials-Project-consistent PBE)
- ”pet-omatpes-l”: PET-OMATPES model (size “l”, materials, r2SCAN)
- ”pet-spice-s”: PET-SPICE model (size “s”, molecules, ωB97M-D3)
- ”pet-spice-l”: PET-SPICE model (size “l”, molecules, ωB97M-D3)
- ”pet-omol-s”: PET-OMol model (size “s”, molecules, ωB97M-V)
- ”pet-omol-m”: PET-OMol model (size “m”, molecules, ωB97M-V)
- ”pet-omol-l”: PET-OMol model (size “l”, molecules, ωB97M-V)
- ”pet-mols-s”: PET-MOLS model (size “s”, organic molecular crystals, PBE0+MBD)
version (str | None) –
version of the model to use. Defaults to the latest stable version, which is 1.6.0 for the PET-MAD models.

Deprecated model versions:
- ”pet-mad-s-v1.0.2”: PET-MAD-1 model (size “s”, materials and molecules, PBEsol)
- ”pet-omad-xs-v1.0.0”: PET-OMAD model (size “xs”, materials and molecules, PBEsol)
- ”pet-omad-s-v1.0.0”: PET-OMAD model (size “s”, materials and molecules, PBEsol)
- ”pet-omad-l-v0.1.0”: PET-OMAD model (size “l”, materials and molecules, PBEsol)
dtype (dtype | None) – dtype to use for the calculations. If None, we will use the default dtype.
checkpoint_path (str | Path | None) – path to a checkpoint file to load the model from. If the filename follows standard naming (e.g., “pet-mad-s-v1.0.2.ckpt”), model/size/version are extracted automatically, and the model, size, and version parameters are ignored.
variants (Dict[str, str | None] | None) – dictionary specifying which variant to use for each output. This option allows to choose the evaluation head when multiple variants are available for a given output. For example, if both energy/pbe and energy/r2scan variants are available for energy target, one can select which one to use by setting the variants parameter to {"energy": "r2scan"}. If energy is set to a variant also the uncertainty and non-conservative outputs will be taken from this variant. If not provided, the default variant for each output will be used (for example: energy with no variant specification).
rotational_average_order (int | None) – order of the Lebedev-Laikov grid used for averaging the prediction over rotations.
rotational_average_batch_size (int | None) – batch size to use for the rotational averaging. If None, all rotations will be computed at once.
device (str | None) – torch device to use for the calculation. If None, we will try the options in the model’s supported_device in order.
non_conservative (bool | Literal['forces', 'stress']) –
whether to use the non-conservative regime of forces and / or stresses prediction. Available options are:
- False: use the conservative regime (default)
- True: use the non-conservative regime for both forces and stresses
- ”forces”: use the non-conservative regime for forces only
- ”stress”: use the non-conservative regime for stresses only
Defaults to False. Available for all models, except:
- PET-MAD models with version < 1.1.0
- PET-SPICE models
- PET-MOLS models
uncertainty_threshold (float | None) – threshold for the atomic energy uncertainty in eV. The calculator warns whenever the predicted atomic energy uncertainty exceeds this value, which is a useful signal that the structure is far from the training distribution. Only used for models that support atomic uncertainty estimation (see Available models). Defaults to 0.1 eV; set it to None to disable the check.
check_consistency (bool) – whether internal consistency checks should be performed. Mainly for developers, defaults to False.

calculate(atoms: Atoms, properties: List[str], system_changes: List[str]) → None[source]¶

Compute some properties with this calculator, and return them in the format expected by ASE.

This is not intended to be called directly by users, but to be an implementation detail of atoms.get_energy() and related functions. See ase.calculators.calculator.Calculator.calculate() for more information.

If the rotational_average_order parameter is set during initialization, the prediction will be averaged over unique rotations in the Lebedev-Laikov grid of a chosen order.

If the rotational_average_batch_size parameter is set during initialization, averaging will be performed in batches of the given size to avoid out of memory errors.

Parameters:

atoms (Atoms)
properties (List[str])
system_changes (List[str])

Return type:

None

property supports_uncertainty: bool¶: Whether the calculator supports uncertainty quantification.

get_energy_uncertainty(atoms: Atoms | None = None, per_atom: bool = False) → ndarray[source]¶

Calculate the energy uncertainty for a given ase.Atoms object. Note, that the uncertainty is not rotationally averaged, even when the calculator is: it is requested from the base model directly.

Parameters:

atoms (Atoms | None) – ASE atoms object. If None, the last calculated atoms will be used.
per_atom (bool) – Whether to return the energy uncertainty per atom.

Returns:

Energy uncertainty in numpy.ndarray format.

Return type:

ndarray

get_energy_ensemble(atoms: Atoms | None = None, per_atom: bool = False) → ndarray[source]¶

Calculate the energy ensemble for a given ase.Atoms object. Note, that the ensemble is not rotationally averaged, even when the calculator is: it is requested from the base model directly.

Parameters:

atoms (Atoms | None) – ASE atoms object. If None, the last calculated atoms will be used.
per_atom (bool) – Whether to return the energies per atom.

Returns:

Energy ensemble in numpy.ndarray format.

Return type:

ndarray

get_forces_uncertainty(atoms: Atoms | None = None, non_conservative: bool | None = None) → ndarray[source]¶

Calculate the forces uncertainty for a given ase.Atoms object through a standard deviation of the forces ensemble. Can be calculated in two ways: conservative or non-conservative, where the default is controlled by the non_conservative parameter of the calculator. Optionnaly, the non-conservative forces uncertainty can be requested explicitly for faster evaluation through the non_conservative=True flag, even when the calculator itself is initialized in the conservative regime. Calculating the conservative forces uncertainty for a non-conservative forces calculator is not supported.

Parameters:

atoms (Atoms | None) – ASE atoms object. If None, the last calculated atoms will be used.
non_conservative (bool | None) – whether to use the non-conservative regime of forces uncertainty calculation. If None, the regime of the calculator is used.

Returns:

Forces uncertainty as numpy.ndarray with shape [n_atoms, 3], in eV/Angstrom.

Return type:

ndarray

get_forces_ensemble(atoms: Atoms | None = None, non_conservative: bool | None = None) → ndarray[source]¶

Calculate the forces ensemble for a given ase.Atoms object.

Can be calculated in two ways: conservative or non-conservative, where the default is controlled by the non_conservative parameter of the calculator. Optionnaly, the non-conservative forces ensemble can be requested explicitly for faster evaluation through the non_conservative=True flag, even when the calculator itself is initialized in the conservative regime. If the calculator is in the conservative regime, the non-conservative forces ensemble is centered on the conservative forces values, while keeping the non-conservative ensemble spread. Calculating the conservative forces ensemble for a non-conservative forces calculator is not supported.

Parameters:

atoms (Atoms | None) – ASE atoms object. If None, the last calculated atoms will be used.
non_conservative (bool | None) – whether to use the non-conservative regime of forces ensemble calculation. If None, the regime of the calculator is used.

Returns:

Forces ensemble as numpy.ndarray with shape [n_atoms, 3, n_ensemble], in eV/Angstrom.

Return type:

ndarray

get_stress_uncertainty(atoms: Atoms | None = None, voigt: bool = True, non_conservative: bool | None = None) → ndarray[source]¶

Calculate the stress uncertainty for a given ase.Atoms object through a standard deviation of the stress ensemble.

Can be calculated in two ways: conservative or non-conservative, where the default is controlled by the non_conservative parameter of the calculator. Optionnaly, the non-conservative stress uncertainty can be requested explicitly for faster evaluation through the non_conservative=True flag, even when the calculator itself is initialized in the conservative regime. Calculating the conservative stress uncertainty for a non-conservative stress calculator is not supported.

Parameters:

atoms (Atoms | None) – ASE atoms object. If None, the last calculated atoms will be used.
non_conservative (bool | None) – whether to use the non-conservative regime of stress uncertainty calculation. If None, the regime of the calculator is used.
voigt (bool)

Returns:

Stress uncertainty as numpy.ndarray with shape [6,] if voigt=True or [3, 3] if voigt=False, in eV/Angstrom^3.

Return type:

ndarray

get_stress_ensemble(atoms: Atoms | None = None, voigt: bool = True, non_conservative: bool | None = None) → ndarray[source]¶

Calculate the stress ensemble for a given ase.Atoms object.

Can be calculated in two ways: conservative or non-conservative, where the default is controlled by the non_conservative parameter of the calculator. Optionnaly, the non-conservative stress ensemble can be requested explicitly for faster evaluation through the non_conservative=True flag, even when the calculator itself is initialized in the conservative regime. Calculating the conservative stress uncertainty for a non-conservative stress calculator is not supported.

Parameters:

atoms (Atoms | None) – ASE atoms object. If None, the last calculated atoms will be used.
non_conservative (bool | None) – whether to use the non-conservative regime of stress ensemble calculation. If None, the regime of the calculator is used.
voigt (bool)

Returns:

Stress uncertainty as numpy.ndarray with shape [6, n_ensemble] if voigt=True or [3, 3, n_ensemble] if voigt=False, in eV/Angstrom^3.

Return type:

ndarray

PET-MAD-DOS¶

class upet.ase.dos.PETMADDOSCalculator(version: str = 'latest', model_path: str | None = None, bandgap_model_path: str | None = None, fermi_model_path: str | None = None, *, check_consistency: bool = False, device: str | None = None)[source]¶

PET-MAD DOS Calculator

Parameters:

version (str) – PET-MAD-DOS version to use. Defaults to the latest stable version.
model_path (str | None) – path to a Torch-Scripted model file to load the model from. If provided, the version parameter is ignored.
bandgap_model_path (str | None) – path to a PyTorch checkpoint file with the bandgap model. If provided, the version parameter is ignored.
check_consistency (bool) – should we check the model for consistency when running, defaults to False.
device (str | None) – torch device to use for the calculation. If None, we will try the options in the model’s supported_device in order.
fermi_model_path (str | None)

calculate(atoms: Atoms, properties: Sequence[Literal['dos_raw', 'dos_denoised', 'dos_raw_per_atom', 'bandgap', 'fermi_level']] = ('dos_raw', 'dos_denoised', 'bandgap', 'fermi_level'), system_changes: Sequence[str] = ()) → Dict[str, Tensor][source]¶

Calculate the density of states, bandgap, and Fermi level for a given ase.Atoms object, or a list of ase.Atoms objects.

Parameters:

atoms (Atoms) – ASE atoms object or a list of ASE atoms objects
properties (Sequence[Literal['dos_raw', 'dos_denoised', 'dos_raw_per_atom', 'bandgap', 'fermi_level']]) – List of what needs to be calculated.
system_changes (Sequence[str]) – List of what has changed since last calculation. Currently ignored, but required for compatibility with ASE.

Returns:

Dictionary containing the calculated properties.

Return type:

Dict[str, Tensor]

dos_from_eigenvalues(eigenvalues: Tensor, kweights: Tensor | None = None) → Tuple[Tensor, Tensor][source]¶

Calls the dos_from_eigenvalues function with PET-MAD-DOS default parameters. The function is useful to compute the DOS and mask from eigenvalues and k-point weights from DFT calculations in a way that is consistent with PET-MAD-DOS.

Parameters:

eigenvalues (Tensor) – Tensor of shape (n_kpoints, n_bands) containing the eigenvalues.
kweights (Tensor | None) – Tensor of shape (n_kpoints,) containing the weights of each k-point.

Returns:

DOS

Return type:

Tuple[Tensor, Tensor]

pad_dos(dos: Tensor, mask: Tensor) → Tuple[Tensor, Tensor][source]¶

Pads the input DOS to the length required for PET-MAD-DOS training/ finetuning. It calls the pad_dos utility function with PET-MAD-DOS default parameters. At the end, it replaces the regions where the DOS is not well-defined with zeros.

Parameters:

dos (Tensor) – Tensor containing the density of states values.
mask (Tensor) – Tensor containing the mask values.

Returns:

Padded DOS

Return type:

Tuple[Tensor, Tensor]

Featurizer¶

class upet.ase.explore.PETMADFeaturizer(version: str = 'latest', checkpoint_path: str | None = None, pet_checkpoint_path: str | None = None, *, check_consistency=False, device=None, length_unit='Angstrom', batch_size: int = 1, progress_bar=<class 'tqdm.std.tqdm'>)[source]¶

Converts structures into low-dimensional projections using PET-MAD features, with dimensionality reduction based on sketch-map.

Usage example:

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

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

>>> # Create visualization
>>> chemiscope.explore(frames, featurize=PETMADFeaturizer(version="latest"))

Parameters:

version (str) – PET-MAD Explorer version to use. Supported versions are “latest”, “1.0.0”.
checkpoint_path (str | None) – path to a checkpoint file to load the exploration model from. If provided, the version parameter is ignored.
pet_checkpoint_path (str | None) – path to a petmad checkpoint file to use for the model from. If not provided, the latest checkpoint is fetched from HuggingFace.
check_consistency – should we check the model for consistency when running, defaults to False.
device – torch device to use for the calculation. If None, we will try the options in the model’s supported_device in order.
length_unit – unit of length used in the structures
batch_size (int)

nvalchemi-toolkit integration¶

Requires the optional nvalchemi extra (pip install "upet[nvalchemi]"); see Installation.

class upet.nvalchemi.UPETWrapper(atomic_types: Sequence[int], hypers: dict[str, Any], composition_energy: Tensor, scale_energy: Tensor)[source]¶

nvalchemi.models.base.BaseModelMixin wrapper around UPET/PET.

Builds and owns a metatrain.pet.modules.backend.PETBackend (from hypers and atomic_types) and drives its three building blocks. Handles:

translating a nvalchemi.data.Batch into the concatenated plain tensors consumed by PETBackend.preprocess (adapt_input());
enabling gradients on positions when autograd outputs are active, and wiring the affine strain trick for stress;
applying the flat composition / scaler buffers decoded from the checkpoint at load time;
producing nvalchemi._typing.ModelOutputs with energy, forces, and stress.

Parameters:

atomic_types (Sequence[int]) – Atomic numbers in species-index order.
hypers (dict[str, Any]) – PET hyper-parameters (see upet.nvalchemi.utils.REQUIRED_HYPERS).
composition_energy (Tensor) – Per-species reference energy, shape [num_species], indexed by species index (not atomic number).
scale_energy (Tensor) – Scalar (0-dim) tensor used as the global energy scale.

property embedding_shapes: dict[str, tuple[int, ...]]¶

Node/graph embedding shapes.

Embeddings concatenate the per-layer node features with the cutoff-weighted, neighbor-summed per-layer edge features (see compute_embeddings()), so the dimension is num_readout_layers * (d_node + d_pet).

property cutoff: float¶: Interaction cutoff in Angstroms.

adapt_input(data: AtomicData | Batch, **_kwargs: Any) → dict[str, Tensor][source]¶

Translate a nvalchemi.data.Batch into backend input tensors.

Produces the concatenated, plain-tensor structure representation that PETBackend.preprocess consumes. All the edge manipulation (NEF reshaping, adaptive cutoffs, reversed-neighbor indexing) then happens inside the backend.

Handles AtomicData -> Batch promotion and gradient enabling on positions when an autograd output is active. Strain handling (for stress) is done by forward() before calling this method, so that the scaled positions/cell flow through the full featurisation.

Parameters:

data (AtomicData | Batch) – Input data. AtomicData inputs are promoted to a single-graph Batch.
_kwargs (Any)

Returns:

Keyword arguments for PETBackend.preprocess: positions, centers, neighbors, species, cells, cell_shifts, system_indices.

Return type:

dict[str, Tensor]

Map raw PET outputs to the standard ModelOutputs layout.

Parameters:

raw_output (dict[str, Tensor | None]) – Dict with optional energy, forces, stress tensors.
data (AtomicData | Batch) – Original input batch (forwarded to BaseModelMixin.adapt_output).

Returns:

Ordered dict keyed by the wrapper’s active outputs.

Return type:

Run the PET backend and return energy / forces / stress.

The energy comes from PETBackend.preprocess -> PETBackend.calculate_features -> PETBackend.predict (the latter already sums the node and cutoff-weighted edge contributions over all readout layers). The flat scaler / composition buffers are then applied.

Conservative forces are derived via torch.autograd.grad() of the total energy with respect to positions. Stresses use the affine-strain trick from nvalchemi.models._utils.

Parameters:

data (AtomicData | Batch) – Input batch.
kwargs (Any) – Forwarded to adapt_input().

Returns:

Dict with the active output keys populated.

Return type:

compute_embeddings(data: AtomicData | Batch, **kwargs: Any) → AtomicData | Batch[source]¶

Compute node and graph embeddings without autograd.

The node embedding is the concatenation of the per-layer node features with the cutoff-weighted, neighbor-summed per-layer edge features:

node = cat(node_features_list, dim=1)
edge = (cat(edge_features_list, dim=2) * cutoff_factors).sum(neighbors)
feats = cat([node, edge], dim=1)

Writes node_embeddings ([N, num_readout_layers*(d_node+d_pet)]) and graph_embeddings ([B, ...], sum-pooled over atoms) into data and returns it. Does not mutate model_config.

Parameters:

data (AtomicData | Batch) – Input data.
kwargs (Any) – Forwarded to adapt_input().

Returns:

The same batch with node_embeddings and graph_embeddings attached.

Return type:

AtomicData | Batch

Load a UPET/PET checkpoint from disk or HuggingFace.

Either checkpoint_path (a local file) or model (a name to fetch from HuggingFace, e.g. "pet-mad-s", optionally with version) must be given. The full list of available models and versions can be listed programmatically via upet.list_upet().

UPETWrapper supports compiling the PET backend building blocks (preprocess / calculate_features / predict) via torch.compile, controlled by compile_model. Models using the 'grid' adaptive-cutoff method (e.g. pet-mad <= v1.5.0) cannot be compiled because of a break in the autograd backward; use a 'solver'-method checkpoint (pet-mad >= v1.6.0) instead.

Parameters:

model (str | None) – Model name to fetch from HuggingFace, either a combined <model>-<size> name (e.g. "pet-mad-s") or a bare base name (e.g. "pet-mad"). Used when checkpoint_path is None; ignored otherwise.
version (str | None) – Model version to fetch, or "latest" / None for the newest available. Ignored when checkpoint_path is given. Defaults to "latest".
dtype (dtype | None) – If set, cast the backend and composition/scaler buffers to this dtype before returning.
checkpoint_path (Path | str | None) – Path to a local PET checkpoint file (.ckpt / .pt). If None, model must be given instead.
device (device | None) – Target device. Defaults to CPU.
compile_model (bool) – torch.compile the three backend building blocks. Sets eval mode and freezes parameters; the model is inference-only after this step.
compile_kwargs (Any) – Forwarded verbatim to each torch.compile call (e.g. fullgraph=True, mode=..., dynamic=...).

Returns:

The loaded wrapper.

Raises:

ValueError – When compile_model is requested for a 'grid' adaptive-cutoff model, or when neither checkpoint_path nor model is given.
FileNotFoundError – When checkpoint_path is neither an existing local file nor a parseable named model.

Return type:

UPETWrapper

export_model(path: Path, as_state_dict: bool = False) → None[source]¶

Serialize the wrapper to disk in a pure-torch layout.

Writes a plain dict containing the backend state_dict, the hyper-parameters, the atomic-type list, and the composition/scaler buffers. The output is not a metatrain / metatomic checkpoint — it is a self-contained snapshot that can be reloaded by constructing UPETWrapper(atomic_types, hypers, ...) and calling load_state_dict on its backend.

Parameters:

path (Path) – Output path.
as_state_dict (bool) – If True, save only the backend’s state_dict. Defaults to False (saves the full snapshot).

Return type:

None

extra_repr() → str¶

Format the model config for nn.Module.__repr__.

Parameters:: self (Any)
Return type:: str