from typing import Dict, List, Optional
import torch
from ..documentation import Labels, TensorBlock
[docs]
class System:
"""
A System contains all the information about an atomistic system; and should be used
as the input of metatensor atomistic models.
"""
def __init__(
self,
species: torch.Tensor,
positions: torch.Tensor,
cell: torch.Tensor,
):
"""
You can create a :py:class:`System` with ``species``, ``positions`` and ``cell``
tensors, or convert data from other libraries.
.. admonition:: Converting data to metatensor `System`
Some external packages provides ways to create :py:class:`System` using data
from other libraries:
- `rascaline`_ has the :py:func:`rascaline.torch.systems_to_torch()`
function that can convert from ASE, chemfiles and PySCF.
.. _rascaline: https://luthaf.fr/rascaline/latest/index.html
:param species: 1D tensor of integer representing the particles identity. For
atoms, this is typically their atomic numbers.
:param positions: 2D tensor of shape (len(species), 3) containing the Cartesian
positions of all particles in the system.
:param cell: 2D tensor of shape (3, 3), describing the bounding box/unit cell of
the system. Each row should be one of the bounding box vector; and columns
should contain the x, y, and z components of these vectors (i.e. the cell
should be given in row-major order). Systems are assumed to obey periodic
boundary conditions, non-periodic systems should set the cell to 0.
"""
def __len__(self) -> int:
...
@property
def species(self) -> torch.Tensor:
"""Tensor of 32-bit integers representing the particles identity"""
@property
def positions(self) -> torch.Tensor:
"""
Tensor of floating point values containing the particles cartesian coordinates
"""
@property
def cell(self) -> torch.Tensor:
"""Tensor of floating point values containing bounding box/cell of the system"""
[docs]
def add_neighbors_list(
self,
options: "NeighborsListOptions",
neighbors: TensorBlock,
):
"""
Add a new neighbors list in this system corresponding to the given ``options``.
The neighbors list should have the following samples: ``"first_atom"``,
``"second_atom"``, ``"cell_shift_a"``, ``"cell_shift_b"``, ``"cell_shift_c"``,
containing the index of the first and second atoms (matching the "atom" sample
in the positions); and the number of cell vector a/b/c to add to the positions
difference to get the pair vector.
The neighbors should also have a single component ``"xyz"`` with values ``[0, 1,
2]``; and a single property ``"distance"`` with value 0.
The neighbors values must contain the distance vector from the first to the
second atom, i.e. ``positions[second_atom] - positions[first_atom] +
cell_shift_a * cell_a + cell_shift_b * cell_b + cell_shift_c * cell_c``.
:param options: options of the neighbors list
:param neighbors: list of neighbors stored in a :py:class:`TensorBlock`
"""
[docs]
def get_neighbors_list(
self,
options: "NeighborsListOptions",
) -> TensorBlock:
"""
Retrieve a previously stored neighbors list with the given ``options``, or throw
an error if no such neighbors list exists.
:param options: options of the neighbors list to retrieve
"""
[docs]
def known_neighbors_lists(self) -> List["NeighborsListOptions"]:
"""
Get all the neighbors lists options registered with this :py:class:`System`
"""
[docs]
def add_data(self, name: str, data: TensorBlock):
"""
Add custom data to this system, stored as :py:class:`TensorBlock`.
This is intended for experimentation with models that need more data as input,
and moved into a field of ``System`` later.
:param name: name of the custom data
:param data: values of the custom data
"""
[docs]
def get_data(self, name: str) -> TensorBlock:
"""
Retrieve custom data stored in this System with the given ``name``, or throw
an error if no data can be found.
:param name: name of the custom data to retrieve
"""
[docs]
def known_data(self) -> List[str]:
"""
Get the name of all the custom data registered with this :py:class:`System`
"""
[docs]
class NeighborsListOptions:
"""Options for the calculation of a neighbors list"""
def __init__(self, model_cutoff: float, full_list: bool, requestor: str = ""):
"""
:param model_cutoff: spherical cutoff radius for the neighbors list, in the
model units
:param full_list: should the list be a full or half neighbors list
:param requestor: who requested this neighbors list, you can add additional
requestors later using :py:meth:`add_requestor`
"""
@property
def model_cutoff(self) -> float:
"""Spherical cutoff radius for this neighbors list in model units"""
@property
def engine_cutoff(self) -> float:
"""
Spherical cutoff radius for this neighbors list in engine units. This defaults
to the same value as ``model_cutoff`` until :py:meth:`set_engine_unit()` is
called.
"""
[docs]
def set_engine_unit(self, conversion):
"""Set the conversion factor from the model units to the engine units"""
@property
def full_list(self) -> bool:
"""
Should the list be a full neighbors list (contains both the pair ``i->j`` and
``j->i``) or a half neighbors list (contains only the pair ``i->j``)
"""
[docs]
def requestors(self) -> List[str]:
"""Get the list of modules requesting this neighbors list"""
[docs]
def add_requestor(self, requestor: str):
"""
Add another ``requestor`` to the list of modules requesting this neighbors list
"""
def __repr__(self) -> str:
...
def __str__(self) -> str:
...
[docs]
def __eq__(self, other: "NeighborsListOptions") -> bool:
...
[docs]
def __ne__(self, other: "NeighborsListOptions") -> bool:
...
[docs]
class ModelOutput:
"""Description of one of the quantity a model can compute."""
def __init__(
self,
quantity: str = "",
unit: str = "",
per_atom: bool = False,
explicit_gradients: List[str] = [], # noqa B006
):
...
quantity: str
"""
Quantity of the output (e.g. energy, dipole, …). If this is an empty
string, no unit conversion will be performed.
"""
unit: str
"""
Unit of the output. If this is an empty string, no unit conversion will
be performed.
"""
per_atom: bool
"""Is the output defined per-atom or for the overall structure"""
explicit_gradients: List[str]
"""
Which gradients should be computed eagerly and stored inside the output
:py:class:`TensorMap`.
"""
[docs]
class ModelCapabilities:
"""Description of a model capabilities, i.e. everything a model can do."""
def __init__(
self,
length_unit: str = "",
species: List[int] = [], # noqa B006
outputs: Dict[str, ModelOutput] = {}, # noqa B006
):
...
length_unit: str
"""unit of lengths the model expects as input"""
species: List[int]
"""which atomic species the model can handle"""
outputs: Dict[str, ModelOutput]
"""
All possible outputs from this model and corresponding settings.
During a specific run, a model might be asked to only compute a subset of these
outputs.
"""
[docs]
class ModelEvaluationOptions:
"""
Options requested by the simulation engine/evaluation code when doing a single model
evaluation.
"""
def __init__(
self,
length_unit: str = "",
outputs: Dict[str, ModelOutput] = {}, # noqa B006
selected_atoms: Optional[Labels] = None,
):
...
length_unit: str
"""unit of lengths the engine uses for the model input"""
outputs: Dict[str, ModelOutput]
"""requested outputs for this run and corresponding settings"""
selected_atoms: Optional[Labels]
"""
Only run the calculation for a selected subset of atoms.
If this is set to ``None``, run the calculation on all atoms. If this is a set of
:py:class:`metatensor.torch.Labels`, it will have two dimensions named ``"system"``
and ``"atom"``, containing the 0-based indices of all the atoms in the selected
subset.
"""
[docs]
def check_atomistic_model(path: str):
"""
Check that the file at ``path`` contains an exported metatensor atomistic model, and
that this model can be loaded in the current process.
This function should be called before :py:func:`torch.jit.load()` when loading an
existing model.
"""
def register_autograd_neighbors(
system: System, neighbors: TensorBlock, check_consistency: bool
):
"""
Register a new autograd node going from (``system.positions``, ``system.cell``) to
the ``neighbors`` distance vectors.
This does not recompute the distance vectors, but work as-if all the data in
``neighbors.values`` was computed directly from ``system.positions`` and
``system.cell``, allowing downstream models to use it directly with full autograd
integration.
:param system: system containing the positions and cell used to compute the
neighbors list
:param system: neighbors list, following the same format as
:py:meth:`System.add_neighbors_list`
:param check_consistency: can be set to ``True`` to run a handful of additional
checks in case the data in neighbors does not follow what's expected.
"""