Atomic Simulation Environment (ASE) integration

The code in metatomic.torch.ase_calculator defines a class that allows using a AtomisticModel which predicts the energy and forces of a system as an ASE calculator; enabling the use of machine learning interatomic potentials to drive calculations compatible with ASE calculators.

Additionally, it allows using arbitrary models with prediction targets which are not just the energy, through the ase_calculator.MetatomicCalculator.run_model() function.

class metatomic.torch.ase_calculator.MetatomicCalculator(model: str | bytes | PurePath | AtomisticModel, *, additional_outputs: Dict[str, ModelOutput] | None = None, extensions_directory=None, check_consistency=False, device=None, variants: Dict[str, str | None] | None = None, non_conservative=False, do_gradients_with_energy=True, uncertainty_threshold=0.1)[source]

Bases: Calculator

The MetatomicCalculator class implements ASE’s ase.calculators.calculator.Calculator API using metatomic models to compute energy, forces and any other supported property.

This class can be initialized with any AtomisticModel, and used to run simulations using ASE’s MD facilities.

Neighbor lists are computed using the fast vesin neighbor list library, unless the system has mixed periodic and non-periodic boundary conditions (which are not yet supported by vesin), in which case the slower ASE neighbor list is used.

Parameters:

  • model (str | bytes | PurePath | AtomisticModel) – model to use for the calculation. This can be a file path, a Python instance of AtomisticModel, or the output of torch.jit.script() on AtomisticModel.

  • additional_outputs (Dict[str, ModelOutput] | None) – Dictionary of additional outputs to be computed by the model. These outputs will always be computed whenever the calculate() function is called (e.g. by ase.Atoms.get_potential_energy(), ase.optimize.optimize.Dynamics.run(), etc.) and stored in the additional_outputs attribute. If you want more control over when and how to compute specific outputs, you should use run_model() instead.

  • extensions_directory – if the model uses extensions, we will try to load them from this directory

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

  • variants (Dict[str, str | None] | None) – dictionary mapping output names to a variant that should be used for the calculations (e.g. {"energy": "PBE"}). If "energy" is set to a variant also the uncertainty and non conservative outputs will be taken from this variant. This behaviour can be overriden by setting the corresponding keys explicitly to None or to another value (e.g. {"energy_uncertainty": "r2scan"}).

  • non_conservative – if True, the model will be asked to compute non-conservative forces and stresses. This can afford a speed-up, potentially at the expense of physical correctness (especially in molecular dynamics simulations).

  • do_gradients_with_energy – if True, this calculator will always compute the energy gradients (forces and stress) when the energy is requested (e.g. through atoms.get_potential_energy()). Because the results of a calculation are cached by ASE, this means future calls to atom.get_forces() will return immediately, without needing to execute the model again. If you are mainly interested in the energy, you can set this to False and enjoy a faster model. Forces will still be calculated if requested with atoms.get_forces().

  • uncertainty_threshold – threshold for the atomic energy uncertainty in eV. This will only be used if the model supports atomic uncertainty estimation (https://pubs.acs.org/doi/full/10.1021/acs.jctc.3c00704). Set this to None to disable uncertainty quantification even if the model supports it.

additional_outputs: Dict[str, TensorMap]

Additional outputs computed by calculate() are stored in this dictionary.

The keys will match the keys of the additional_outputs parameters to the constructor; and the values will be the corresponding raw metatensor.torch.TensorMap produced by the model.

todict()[source]

Obtain a dictionary of parameter information

metadata() → ModelMetadata[source]

Get the metadata of the underlying model

Return type:

ModelMetadata

run_model(atoms: Atoms | List[Atoms], outputs: Dict[str, ModelOutput], selected_atoms: Labels | None = None) → Dict[str, TensorMap][source]

Run the model on the given atoms, computing the requested outputs and only these.

The output of the model is returned directly, and as such the blocks’ values will be torch.Tensor.

This is intended as an easy way to run metatensor models on ase.Atoms when the model can compute outputs not supported by the standard ASE’s calculator interface.

All the parameters have the same meaning as the corresponding ones in metatomic.torch.ModelInterface.forward().

Parameters:
Return type:

Dict[str, TensorMap]

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.

Parameters:
Return type:

None

compute_energy(atoms: Atoms | List[Atoms], compute_forces_and_stresses: bool = False, *, per_atom: bool = False) → Dict[str, float | ndarray | List[float | ndarray]][source]

Compute the energy of the given atoms.

Energies are computed in eV, forces in eV/Å, and stresses in 3x3 tensor format and in units of eV/Å^3.

Parameters:

  • atoms (Atoms | List[Atoms]) – ase.Atoms, or list of ase.Atoms, on which to run the model

  • compute_forces_and_stresses (bool) – if True, the model will also compute forces and stresses. IMPORTANT: stresses will only be computed if all provided systems have periodic boundary conditions in all directions.

  • per_atom (bool) – if True, the per-atom energies will also be computed.

Returns:

A dictionary with the computed properties. The dictionary will contain the energy as a float. If compute_forces_and_stresses is True, the forces and stress will also be included as numpy arrays. If per_atom is True, the energies key will also be present, containing the per-atom energies as a numpy array. In case of a list of ase.Atoms, the dictionary values will instead be lists of the corresponding properties, in the same format.

Return type:

Dict[str, float | ndarray | List[float | ndarray]]

class metatomic.torch.ase_calculator.SymmetrizedCalculator(base_calculator: MetatomicCalculator, *, l_max: int = 3, batch_size: int | None = None, include_inversion: bool = True, apply_space_group_symmetry: bool = False, store_rotational_std: bool = False)[source]

Bases: Calculator

Take a MetatomicCalculator and average its predictions to make it (approximately) equivariant. Only predictions for energy, forces and stress are supported.

The default is to average over a quadrature of the orthogonal group O(3) composed this way:

  • Lebedev quadrature of the unit sphere (S^2)

  • Equispaced sampling of the unit circle (S^1)

  • Both proper and improper rotations are taken into account by including the

    inversion operation (if include_inversion=True)

Parameters:

  • base_calculator (MetatomicCalculator) – the MetatomicCalculator to be symmetrized

  • l_max (int) – the maximum spherical harmonic degree that the model is expected to be able to represent. This is used to choose the quadrature order. If 0, no rotational averaging will be performed (it can be useful to average only over the space group, see apply_group_symmetry).

  • batch_size (int | None) – number of rotated systems to evaluate at once. If None, all systems will be evaluated at once (this can lead to high memory usage).

  • include_inversion (bool) – if True, the inversion operation will be included in the averaging. This is required to average over the full orthogonal group O(3).

  • apply_space_group_symmetry (bool) – if True, the results will be averaged over discrete space group of rotations for the input system. The group operations are computed with spglib, and the average is performed after the O(3) averaging (if any). This has no effect for non-periodic systems.

  • store_rotational_std (bool) – if True, the results will contain the standard deviation over the different rotations for each property (e.g., energy_std).

Basic calculator implementation.

restart: str

Prefix for restart file. May contain a directory. Default is None: don’t restart.

ignore_bad_restart_file: bool

Deprecated, please do not use. Passing more than one positional argument to Calculator() is deprecated and will stop working in the future. Ignore broken or missing restart file. By default, it is an error if the restart file is missing or broken.

directory: str or PurePath

Working directory in which to read and write files and perform calculations.

label: str

Name used for all files. Not supported by all calculators. May contain a directory, but please use the directory parameter for that instead.

atoms: Atoms object

Optional Atoms object to which the calculator will be attached. When restarting, atoms will get its positions and unit-cell updated from file.

implemented_properties: List[str] = ['energy', 'energies', 'forces', 'stress', 'stresses']

Properties calculator can handle (energy, forces, …)

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

Perform the calculation for the given atoms and properties.

Parameters:

  • atoms (Atoms) – the ase.Atoms on which to perform the calculation

  • properties (List[str]) – list of properties to compute, among energy, forces, and stress

  • system_changes (List[str]) – list of changes to the system since the last call to calculate

Return type:

None