xnn.gnn.featurizers.spherical_bessel.SphericalBesselBasis#

class xnn.gnn.featurizers.spherical_bessel.SphericalBesselBasis(n_spherical=7, n_radial=6, cutoff=5.0)[source]#

Bases: Module

The 2D spherical Fourier-Bessel basis a_ln(d, alpha) of DimeNet (eq 6).

Joint representation of a distance d and an angle alpha in the N_SHBF * N_SRBF orthonormal functions sqrt(2 / (c^3 j_{l+1}(z_ln)^2)) j_l(z_ln d / c) Y_l^0(alpha). The basis itself ends with a step at the cutoff; DimeNet multiplies it by the polynomial envelope of eq 8 (PolynomialCutoff) to make it twice continuously differentiable.

The radial and angular factors are exposed separately (radial(), angular()) because in the directional message passing the distance belongs to an edge and the angle to a triplet of atoms, so the radial part is evaluated once per edge and gathered onto the triplets.

Parameters:

  • n_spherical (int, optional) – Number of spherical-harmonic degrees N_SHBF (l = 0 .. N_SHBF-1), by default 7 (the paper’s value).

  • n_radial (int, optional) – Zeros per degree N_SRBF (n = 1 .. N_SRBF), by default 6.

  • cutoff (float, optional) – Cutoff radius c in Angstrom, by default 5.0.

Variables:

  • zeros (Tensor) – Buffer of shape (n_spherical, n_radial) holding z_ln.

  • norm (Tensor) – Buffer of the same shape holding sqrt(2 / (c^3 j_{l+1}(z_ln)^2)).

output_dim#

Number of basis functions n_spherical * n_radial.

radial(r)[source]#

The radial factors sqrt(2 / (c^3 j_{l+1}(z_ln)^2)) j_l(z_ln r / c).

Parameters:

r (Tensor) – Distances, shape (E,).

Returns:

Shape (E, n_spherical, n_radial), in the dtype of r.

Return type:

Tensor

angular(cos_alpha)[source]#

The angular factors Y_l^0(alpha).

Parameters:

cos_alpha (Tensor) – Cosines of the angles, shape (T,).

Returns:

Shape (T, n_spherical).

Return type:

Tensor

forward(r, cos_alpha)[source]#

Evaluate the full basis for distance/angle pairs.

Parameters:

  • r (Tensor) – Distances, shape (T,).

  • cos_alpha (Tensor) – Cosines of the angles paired with them, shape (T,).

Returns:

a_ln(r, alpha) of shape (T, n_spherical * n_radial), the degree l running slowest (column l * n_radial + n - 1).

Return type:

Tensor