xnn.gnn.featurizers.spherical_bessel#
The 2D spherical Fourier-Bessel basis of DimeNet (Gasteiger et al., ICLR 2020).
DimeNet represents the distance d and the angle alpha of a pair of
messages jointly in the regular solutions of the Helmholtz equation inside the
cutoff sphere (paper section 5, eq 6),
a_ln(d, alpha) = sqrt(2 / (c^3 j_{l+1}(z_ln)^2)) * j_l(z_ln d / c) * Y_l^0(alpha)
with the spherical Bessel functions j_l, their n-th positive zero
z_ln (so every function vanishes at the cutoff c), and the m = 0
spherical harmonics Y_l^0. The functions are orthonormal on the cutoff
sphere, l = 0 .. N_SHBF - 1 and n = 1 .. N_SRBF; the radial-only
member of the family (l = 0, zeros n pi) is the Bessel radial basis
BesselRBF (eq 7).
Everything here is evaluated in plain torch: the spherical Bessel functions
by the power series close to the origin and the upward recurrence beyond
(the switch keeps both exact to double precision), their zeros by bisection
between the interlacing zeros of the previous order, and Y_l^0 through
the Legendre recurrence. The internal arithmetic is float64 whatever the
model dtype, so a float32 model gets correctly rounded basis values.
Functions
|
Real spherical harmonics |
|
Spherical Bessel functions of the first kind, orders 0 to |
|
The first |
|
The |
Classes
|
The 2D spherical Fourier-Bessel basis |