xnn.dnn.featurizers.symmetry_functions.AngularSymmetryFunctions#

class xnn.dnn.featurizers.symmetry_functions.AngularSymmetryFunctions(species, cutoff=4.0, etas=(0.5,), zetas=(8.0,), rs=(0.0, 1.5, 3.0), theta_s=(0.0, 1.5708, 3.1416, 4.7124), cos_factor=1.0)[source]#

Bases: Featurizer

ANI-style angular AEV term, resolved by unordered neighbour-species pair.

G^A_i = 2^(1-zeta) * sum_{j,k} (1 + cos(theta_ijk - theta_s))^zeta * exp(-eta ((r_ij + r_ik)/2 - Rs)^2) * fc(r_ij) fc(r_ik)

summed over distinct neighbour pairs (j, k) of centre i, with a grid over (eta, zeta, Rs, theta_s). Contributions are bucketed by the unordered pair of neighbour chemical species, giving an invariant per-atom descriptor. Equivalent to torchani’s 2 * ((1+cos)/2)^zeta * ... form.

Parameters:

  • species (list[int]) – Atomic numbers the descriptor resolves; the angular term is bucketed by the unordered pairs of these species.

  • cutoff (float, optional) – Cutoff radius for the cosine cutoff, by default 4.0.

  • etas (sequence of float, optional) – Radial width parameters, by default (0.5,).

  • zetas (sequence of float, optional) – Angular resolution exponents, by default (8.0,).

  • rs (sequence of float, optional) – Radial shifts Rs, by default (0.0, 1.5, 3.0).

  • theta_s (sequence of float, optional) – Angular shifts (in radians), by default (0.0, 1.5708, 3.1416, 4.7124).

  • cos_factor (float, optional) – Value multiplying cos(theta) before acos. 1.0 (Behler-Parrinello, default) or 0.95 (ANI / NeuroChem / torchani, which keeps acos away from its infinite-gradient endpoints).

Variables:

  • species (list[int]) – The resolved species.

  • cutoff (float) – Cutoff radius.

  • cutoff_fn (CosineCutoff) – Smooth cutoff function applied to each edge.

  • pairs (list[tuple[int, int]]) – Unordered neighbour-species pairs (buckets).

  • pair_to_idx (dict[tuple[int, int], int]) – Mapping from a species pair to its bucket index.

property output_dim: int#

Descriptor length, n_params * n_pairs (grid size over (eta, zeta, Rs, theta_s) per unordered species pair).

Type:

int

forward(data)[source]#

Compute the angular symmetry-function descriptor.

Parameters:

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

Returns:

Per-atom angular descriptor of shape (N, output_dim). Returns all zeros when the graph contains no neighbour triplets.

Return type:

Tensor