xnn.dnn.featurizers.symmetry_functions.RadialSymmetryFunctions#

class xnn.dnn.featurizers.symmetry_functions.RadialSymmetryFunctions(species, cutoff=6.0, etas=(0.05, 0.5, 2.0, 8.0), rs=(0.0,), prefactor=1.0)[source]#

Bases: Featurizer

Behler G2 radial symmetry functions, resolved by neighbour species.

Computes G2_i = pref * sum_j exp(-eta (r_ij - Rs)^2) fc(r_ij) for each central atom i, over a grid of (eta, Rs) parameters and bucketed by the chemical species of the neighbour j, giving an invariant per-atom descriptor with chemical awareness.

Parameters:

  • species (list[int]) – Atomic numbers the descriptor resolves neighbours into (one bucket per species).

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

  • etas (sequence of float, optional) – Width parameters of the radial Gaussians, by default (0.05, 0.5, 2.0, 8.0).

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

  • prefactor (float, optional) – Constant multiplying every term. 1.0 (Behler-Parrinello, default) or 0.25 (ANI / NeuroChem / torchani convention).

Variables:

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

  • cutoff (float) – Cutoff radius.

  • z_to_idx (dict[int, int]) – Mapping from atomic number to its species bucket index.

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

property output_dim: int#

Descriptor length, n_params * n_species (n_eta * n_rs parameters per species bucket).

Type:

int

forward(data)[source]#

Compute the radial symmetry-function descriptor.

Parameters:

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

Returns:

Per-atom radial descriptor of shape (N, output_dim).

Return type:

Tensor