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 | Literal['forces', 'stress'] = 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)

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

  • 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

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

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)