API reference

Auto-generated reference for UPET’s Python API

Top-level functions

upet.get_upet(*, model: str | None = None, size: str | None = None, version: str | Version | None = 'latest', checkpoint_path: str | None = None) → AtomisticModel[source]

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

upet.save_upet(*, model: str | None = None, size: str | None = None, version: str | None = 'latest', checkpoint_path: str | None = None, output: str | None = None)[source]

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 | None) – path to a checkpoint file to load the model from.

  • output (str | None) – path for the output model. Defaults to “{model}-{size}-v{version}.pt” or “model.pt” for non-standard checkpoint names.

Calculators

class upet.calculator.UPETCalculator(model: str | None = None, version: str | None = 'latest', dtype: dtype | None = None, checkpoint_path: str | None = None, variants: Dict[str, str | None] | None = None, rotational_average_order: int | None = None, rotational_average_batch_size: int | None = None, *, device: str | None = None, non_conservative: bool = False, check_consistency: bool = False)[source]

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-1.5 model (size “xs”, materials and molecules, r2SCAN)

    • ”pet-mad-s”: PET-MAD-1.5 model (size “s”, materials and molecules, r2SCAN)

    • ”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)

  • version (str | None) –

    version of the model to use. Defaults to the latest stable version. 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 | 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) –

    whether to use the non-conservative regime of forces and stresses prediction. Defaults to False. Available for all models, except:

    • PET-MAD models with version < 1.1.0

    • PET-SPICE models

  • 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:
Return type:

None

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

Get the energy uncertainty for a given ase.Atoms object.

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]

Get the ensemble of energies for a given ase.Atoms object.

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 uncertainty in numpy.ndarray format.

Return type:

ndarray

class upet.calculator.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.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 pet_mad.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)