torchani.aev#

Atomic Environment Vectors (AEVs) are TorchANI’s name for local atomic features, calculated from the local chemical environment of each atom. This module contains the AEV Computer, 2-body (“radial”), and 3-body (“angular”) AEV terms.

Classes

`AEVComputer`	Base class for modules that compute AEVs
`BaseRadial`	Base class for 2-body expansions
`BaseAngular`	Base class for 3-body expansions
`Radial`
`Angular`
`ANIRadial`	Compute the radial sub-AEV terms given a sequence of atom pair distances
`ANIAngular`	Compute the angular sub-AEV terms of the center atom given neighbor pairs.

class torchani.aev.AEVComputer(radial, angular, num_species, strategy='pyaev', cutoff_fn=None, neighborlist='all_pairs')[source]#

Base class for modules that compute AEVs

Can be used to compute local atomic features (atomic environment vectors or AEVs), given a batch of molecules.

Parameters:

radial ('ani1x' | 'ani2x' | 'ani1ccx' | BaseRadial) – The module used to compute the radial part of the AEVs
angular ('ani1x' | 'ani2x' | 'ani1ccx' | BaseAngular) – The module used to compute the angular part of the AEVs
num_species (Final[int]) – The number of elements this module supports
strategy (str) – The AEV computation strategy. Valid values are: - ‘cuaev’: Use the hand-coded CUDA cuAEV extension - ‘pyaev’: Use pytorch exclusively - ‘auto’: Choose the most performant available option
cutoff_fn ('global' | 'dummy' | 'cosine' | 'smooth' | Cutoff | None) – Cutoff fn to use for the radial and angular terms
neighborlist ('all_pairs' | 'adaptive' | 'cell_list' | 'verlet_cell_list' | 'base' | Neighborlist) – Kind of neighborlist for internal computation.
are (Valid values) – ‘all_pairs’, ‘cell_list’, ‘adaptive’

compute_from_neighbors(elem_idxs, coords, neighbors)[source]#

Compute the AEVs from the result of a neighborlist calculation

Parameters:

elem_idxs (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).
neighbors (Neighbors) – typing.NamedTuple that contains the output of a torchani.neighbors.Neighborlist calculation.

Returns:

A float tensor with local atomic features (AEVs). Shape is (molecules, atoms, num-aev-features).

Return type:

Tensor

forward(elem_idxs, coords=None, cell=None, pbc=None)[source]#

Compute AEVs for a batch of molecules

Parameters:

elem_idxs (Tensor) – An int torch.Tensor that stores the element indices of a batch of molecules, (for example after conversion with torchani.nn.SpeciesConverter). Shape is (molecules, 3).
coords (Tensor | None) – A float tensor with the coordinates of a batch of molecules. Shape is (molecules, atoms, 3). All ANI models use Angstrom.
cell (Tensor | None) – A float tensor with unit cell vectors in its rows. Only use this with PBC. A cell with dimensions 10, 15, 20 (in Angstrom for all ANI models) in the x, y, and z directions is given by torch.tensor([[10., 0., 0.],[0., 15., 0.],[0., 0., 20.]]).
pbc (Tensor | None) – A bool tensor that stores whether periodic boundary conditions (PBC) are enabled for the x, y, z directions. pbc=torch.tensor([True, True, True]) fully enables PBC and pbc=None (or torch.tensor([False, False, False])). fully disables it.

Returns:

A float tensor with local atomic features (AEVs). Shape is (molecules, atoms, num-aev-features).

Return type:

Tensor

classmethod from_constants(radial_cutoff, angular_cutoff, radial_eta, radial_shifts, angular_eta, angular_zeta, angular_shifts, sections, num_species, strategy='pyaev', cutoff_fn='cosine', neighborlist='all_pairs')[source]#

Build an AEVComputer with standard radial and angular terms, from constants

For more detail consult the equations in the original ANI article.

Note

This constructor is not recommended, it is kept for backward compatibility. Consider using either the primary AEVComputer constructor, AEVComputer.like_1x, or AEVComputer.like_2x instead.

Parameters:

radial_cutoff (float) – \(R_C\) in eq. (2) when used at eq. (3)
angular_cutoff (float) – \(R_C\) in eq. (2) when used at eq. (4)
radial_eta (float) – The 1D tensor of \(\eta\) in eq. (3)
radial_shifts (Sequence[float]) – The 1D tensor of \(R_s\) in eq. (3)
angluar_eta – The 1D tensor of \(\eta\) in eq. (4)
angular_zeta (float) – The 1D tensor of \(\zeta\) in eq. (4)
angular_shifts (Sequence[float]) – The 1D tensor of \(R_s\) in eq. (4)
sections (Sequence[float]) – The 1D tensor of \(\theta_s\) in eq. (4)
num_species (int) – Number of supported atom types.
strategy (str) – Compute strategy to use.
cutoff_fn ('global' | 'dummy' | 'cosine' | 'smooth' | Cutoff) – The cutoff function used for the calculation.
neighborlist ('all_pairs' | 'adaptive' | 'cell_list' | 'verlet_cell_list' | 'base' | Neighborlist) – The neighborlist usied for the calculation.

Returns:

The constructed AEVComputer, ready for use.

Return type:

Self

classmethod like_1x(num_species=4, strategy='pyaev', cutoff_fn='cosine', neighborlist='all_pairs', radial_start=0.9, radial_cutoff=5.2, radial_eta=16.0, radial_num_shifts=16, angular_start=0.9, angular_cutoff=3.5, angular_eta=8.0, angular_zeta=32.0, angular_num_shifts=4, angular_num_sections=8)[source]#

Build an AEVComputer with standard radial and angular terms

Uses the same defaults as those in the torchani.models.ANI1x model.

Parameters:

cutoff_fn ('global' | 'dummy' | 'cosine' | 'smooth' | Cutoff) – The cutoff function used for the calculation.
neighborlist ('all_pairs' | 'adaptive' | 'cell_list' | 'verlet_cell_list' | 'base' | Neighborlist) – The neighborlist usied for the calculation.

Returns:

The constructed AEVComputer, ready for use.

Return type:

Self

classmethod like_2x(num_species=7, strategy='pyaev', cutoff_fn='cosine', neighborlist='all_pairs', radial_start=0.8, radial_cutoff=5.1, radial_eta=19.7, radial_num_shifts=16, angular_start=0.8, angular_cutoff=3.5, angular_eta=12.5, angular_zeta=14.1, angular_num_shifts=8, angular_num_sections=4)[source]#

Build an AEVComputer with standard radial and angular terms

Uses the same defaults as those in the torchani.models.ANI2x model.

Parameters:

cutoff_fn ('global' | 'dummy' | 'cosine' | 'smooth' | Cutoff) – The cutoff function used for the calculation.
neighborlist ('all_pairs' | 'adaptive' | 'cell_list' | 'verlet_cell_list' | 'base' | Neighborlist) – The neighborlist usied for the calculation.

Returns:

The constructed AEVComputer, ready for use.

Return type:

Self

class torchani.aev.BaseRadial(cutoff, cutoff_fn='cosine')[source]#

Base class for 2-body expansions

Client classes should only rely on module(distances), compute is meant only to be overriden by users.

compute(distances)[source]#

Compute the radial terms. Output shape is: (pairs, self.num_feats)

Subclasses must implement this method

Note

Don’t call this method directly, instead call, module(neighbors.distances).

class torchani.aev.BaseAngular(cutoff, cutoff_fn='cosine')[source]#

Base class for 3-body expansions

Client classes should only rely on module(tri_distances, tri_vectors), compute_* functions are meant only to be overriden by users.

compute_cos_angles(cos_angles)[source]#

Computes the angular part of the 3-body terms

This function must be overriden by subclasses

Note

Don’t call this method directly, instead call, module(triples.distances, triples.diff_vectors).

Parameters:: cos_angles (Tensor) – Cosine of the angle for a triple ijk
Returns:: A float torch.Tensor of shape (triples, feats). By design this function does not sum over atoms.
Return type:: Tensor

compute_radial(distances_ji, distances_jk)[source]#

Computes the radial part of the 3-body terms

This function must be overriden by subclasses

Note

Don’t call this method directly, instead call, module(triple_distances, triple_vectors).

Parameters:

distances_ji (Tensor) – Distances from the center atom, ‘j’, to a side atom, ‘i’.
distances_jk – Distances from the center atom, ‘j’, to a side atom, ‘k’.

Returns:

A float torch.Tensor of shape (triples, feats). By design this function does not sum over atoms.

Return type:

Tensor

class torchani.aev.Radial(cutoff, trainable=(), cutoff_fn='cosine', **kwargs)[source]#

class torchani.aev.Angular(cutoff, trainable=(), cutoff_fn='cosine', **kwargs)[source]#

class torchani.aev.ANIRadial(eta, shifts, cutoff, cutoff_fn='cosine')[source]#

Compute the radial sub-AEV terms given a sequence of atom pair distances

This correspond to equation (3) in the ANI paper. This function just computes the terms. The sum in the equation is not computed. The input tensor has shape (conformations, atoms, N), where N is the number of neighbor atoms within the cutoff radius and the output tensor should have shape (conformations, atoms, self.num_feats)

compute(distances)[source]#

Computes the ANI terms associated with a group of pairs

Note

Don’t call this method directly, instead call the module, module(distances).

Parameters:: distances (Tensor) – A float tensor with the distances between pairs of atoms in a system. Shape is (pairs,).
Returns:: A float torch.Tensor of shape (pairs, shifts). By design this function does not sum over atoms.
Return type:: Tensor

classmethod cover_linearly(start=0.9, cutoff=5.2, eta=19.7, num_shifts=16, cutoff_fn='cosine')[source]#

Builds angular terms by linearly dividing space radially up to a cutoff

num_shifts are created, starting from start until cutoff, excluding it. This is the way angular and radial shifts were originally created in the ANI paper.

class torchani.aev.ANIAngular(eta, zeta, shifts, sections, cutoff, cutoff_fn='cosine')[source]#

Compute the angular sub-AEV terms of the center atom given neighbor pairs.

This correspond to equation (4) in the ANI paper. This function just compute the terms. The sum is not computed. The input tensor has shape (conformations, atoms, N), where N is the number of neighbor atom pairs within the cutoff radius and the output tensor should have shape (conformations, atoms, self.num_feats)

compute_cos_angles(cos_angles)[source]#

Computes the angular part of the ANI 3-body terms

Note

Don’t call this method directly, instead call, module(triples.distances, triples.diff_vectors).

Parameters:: cos_angles (Tensor) – Cosine of the angle for a triple ijk
Returns:: A float torch.Tensor of shape (triples, feats). By design this function does not sum over atoms.
Return type:: Tensor

compute_radial(distances_ji, distances_jk)[source]#

Computes the radial part of the ANI 3-body terms

Note

Don’t call this method directly, instead call, module(triple_distances, triple_vectors).

Parameters:

distances_ji (Tensor) – Distances from the center atom, ‘j’, to a side atom, ‘i’.
distances_jk – Distances from the center atom, ‘j’, to a side atom, ‘k’.

Returns:

A float torch.Tensor of shape (triples, feats). By design this function does not sum over atoms.

Return type:

Tensor

classmethod cover_linearly(start=0.9, cutoff=3.5, eta=12.5, zeta=14.1, num_shifts=8, num_sections=4, cutoff_fn='cosine')[source]#

Builds angular terms by dividing angular and radial coords, up to a cutoff

The divisions are equally spaced “num_shifts” are created, starting from “start” until “cutoff”, excluding it. “num_sections” does a similar thing for the angles. This is the way angular and radial shifts were originally created in ANI.