xnn.dnn.featurizers.aev.AEV#

class xnn.dnn.featurizers.aev.AEV(species, radial_cutoff=5.2, angular_cutoff=3.5, radial_etas=(16.0,), radial_rs=(0.9, 1.16875, 1.4375, 1.7062499999999998, 1.975, 2.24375, 2.5124999999999997, 2.78125, 3.05, 3.3187499999999996, 3.5875, 3.8562499999999997, 4.125, 4.39375, 4.6625, 4.93125), angular_etas=(8.0,), angular_zetas=(32.0,), angular_rs=(0.9, 1.55, 2.2, 2.85), angular_theta_s=(0.19634954084936207, 0.5890486225480862, 0.9817477042468103, 1.3744467859455345, 1.7671458676442586, 2.1598449493429825, 2.552544031041707, 2.945243112740431), radial_prefactor=0.25, angular_cos_factor=0.95)[source]#

Bases: Featurizer

Atomic Environment Vector: concatenated radial + angular descriptors.

Composes RadialSymmetryFunctions and AngularSymmetryFunctions and concatenates their per-atom outputs to form the invariant descriptor ANI feeds to its per-element networks. The two sub-featurizers may use different cutoffs; the model’s neighbour list uses the larger of the two.

The defaults are the ANI / NeuroChem conventions (radial prefactor 0.25, cosine factor 0.95), so with matching grids this reproduces torchani.AEVComputer. Prefer the ani1x() / ani1() classmethods for the two published grids.

Parameters:

  • species (list[int]) – Atomic numbers the descriptor resolves.

  • radial_cutoff (float, optional) – Cutoff radius for the radial symmetry functions, by default 5.2.

  • angular_cutoff (float, optional) – Cutoff radius for the angular symmetry functions, by default 3.5.

  • radial_etas (sequence of float, optional) – Width parameters for the radial term, by default (16.0,).

  • radial_rs (sequence of float, optional) – Radial shifts for the radial term. Defaults to the 16 ANI-1x shifts.

  • angular_etas (sequence of float, optional) – Width parameters for the angular term, by default (8.0,).

  • angular_zetas (sequence of float, optional) – Angular resolution exponents, by default (32.0,).

  • angular_rs (sequence of float, optional) – Radial shifts for the angular term. Defaults to the 4 ANI-1x shifts.

  • angular_theta_s (sequence of float, optional) – Angular shifts (radians). Defaults to the 8 ANI-1x shifts.

  • radial_prefactor (float, optional) – Constant multiplying the radial term, by default 0.25 (ANI / NeuroChem). Use 1.0 for the Behler-Parrinello convention.

  • angular_cos_factor (float, optional) – Value multiplying cos(theta) before acos in the angular term, by default 0.95 (ANI / NeuroChem). Use 1.0 for Behler-Parrinello.

Variables:
classmethod ani1x(species=[1, 6, 7, 8])[source]#

AEV with the exact ANI-1x constants shipped by torchani.

Radial cutoff 5.2 A (16 shifts), angular cutoff 3.5 A (4 radial x 8 angular shifts). For the four ANI elements this is a 384-length AEV that reproduces torchani.AEVComputer element-for-element.

Parameters:

species (list[int], optional) – Atomic numbers, by default [1, 6, 7, 8] (H, C, N, O).

Returns:

The ANI-1x-configured featurizer.

Return type:

AEV

classmethod ani2x(species=[1, 6, 7, 8, 16, 9, 17])[source]#

AEV with the exact ANI-2x constants shipped by torchani.

Radial cutoff 5.1 A (16 shifts), angular cutoff 3.5 A (8 radial x 4 angular shifts), with the shift grids starting at 0.8 A and the ANI-2x widths (eta 19.7 / 12.5, zeta 14.1). For the seven ANI-2x elements (H, C, N, O, S, F, Cl) this is a 1008-length AEV that reproduces torchani.AEVComputer element-for-element.

Parameters:

species (list[int], optional) – Atomic numbers, by default [1, 6, 7, 8, 16, 9, 17] (H, C, N, O, S, F, Cl, in torchani’s order).

Returns:

The ANI-2x-configured featurizer.

Return type:

AEV

classmethod ani1(species=[1, 6, 7, 8])[source]#

AEV with the original ANI-1 parameterisation (Smith et al. 2017).

Radial cutoff 4.6 A (32 shifts), angular cutoff 3.1 A (8 radial x 8 angular shifts). For the four ANI elements this is the 768-length AEV described in the paper, built with its evenly-spaced-shift recipe.

Parameters:

species (list[int], optional) – Atomic numbers, by default [1, 6, 7, 8] (H, C, N, O).

Returns:

The ANI-1-configured featurizer.

Return type:

AEV

property output_dim: int#

Descriptor length, the sum of the radial and angular symmetry-function output dimensions.

Type:

int

forward(data)[source]#

Compute the AEV by concatenating radial and angular descriptors.

Parameters:

data (AtomicGraph) – Atomic graph providing atomic numbers, edge index and edge vectors.

Returns:

Per-atom AEV of shape (N, output_dim), the concatenation of the radial and angular symmetry functions along the last dimension.

Return type:

Tensor