ASE

Official website

How is metatomic supported?

https://ase-lib.org/

Via the metatomic-ase package

Supported model outputs

How to install the code

The code is available in the metatomic-ase package, which can be installed using pip install metatomic-ase.

How to use the code

metatomic_ase.MetatomicCalculator supports all the features described above, and is designed to be used as a drop-in replacement for any ASE calculator in any ASE workflow. You can also check the corresponding tutorial. To compute rotationally-averaged energies, forces, and stresses with non-equivariant architectures, wrap the model with metatomic.torch.o3.SymmetrizedModel.wrap() before passing it to the calculator.

API documentation

class metatomic_ase.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: bool | Literal['forces', 'stress'] = False, do_gradients_with_energy=True, uncertainty_threshold=0.1, neighbor_skin=2.0)[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 metatomic.torch.AtomisticModel, and used to run simulations using ASE’s MD facilities.

Neighbor lists are computed using the vesin neighbor list library, either on CPU or GPU depending on the device of the model. If nvalchemiops is installed, full neighbor lists on GPU will be computed with it instead.

The calculator also supports the computation of additional properties beyond energy, forces and stress, as long as they are supported by the underlying model. You can either use the additional_outputs argument to the constructor (the additional outputs will be stored in the additional_outputs attribute), or call run_model() directly with the desired outputs.

The calculator can provide the following additional standard quantities as inputs to the model:

As well as the following ASE-specific quantities:

  • per-atom: ase::initial_magmoms, ase::initial_charges, and ase::tags; corresponding to the outputs of ase.Atoms.get_initial_magnetic_moments(), ase.Atoms.get_initial_charges(), and ase.Atoms.get_tags() respectively. Model can also request data from ase.Atoms.arrays with the ase::arrays:: prefix (e.g. ase::arrays::foo will request the foo array from ase.Atoms.arrays).

  • per-system: Model can request data from ase.Atoms.info with the ase::info:: prefix (e.g. ase::info::foo will request the foo entry from ase.Atoms.info).

Parameters:

  • model (str | bytes | PurePath | AtomisticModel) – model to use for the calculation. This can be a file path, a Python instance of metatomic.torch.AtomisticModel, or the output of torch.jit.script() on metatomic.torch.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 (bool | Literal['forces', 'stress']) –

    controls which outputs are obtained directly from the model rather than via autograd on the energy. Accepted values are:

    • False (default): conservative mode; forces and stress are both derived from the gradient of the energy.

    • True: both forces and stress are read directly from the model’s non_conservative_forces and non_conservative_stress outputs.

    • "forces": forces come from the model’s non_conservative_forces output; stress is still obtained via autograd.

    • "stress": stress comes from the model’s non_conservative_stress output; forces are still obtained via autograd.

    Using any value other than False 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.

  • neighbor_skin – skin distance in Å for the Verlet list in the neighbor list calculation. By setting this parameter, you are extending the cutoff radius of the neighbor list calculation by neighbor_skin Å. The extended cutoff helps reduce the frequency of neighbor list updates, because when the maximun atomic displacement is smaller than \((r_{skin} / 2)\), the neighbor list can still ensure that all neighbors within \((r_{cutoff})\) are included. The default value is 2.0 Å.

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.

model() → AtomisticModel[source]

Get the underlying model used by this calculator

Return type:

AtomisticModel

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]

check_state(atoms: Atoms, tol: float = 1e-15) → List[str][source]

Detect system changes, including atoms.info keys used as model inputs.

ASE’s default check_state() only tracks per-atom arrays, cell and pbc, so mutations to atoms.info (e.g. atoms.info["spin"]) would otherwise return a stale cached result. Any watched info key that differs from the last calculation is appended to the change list, forcing a fresh run.

Parameters:
Return type:

List[str]

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