xnn.gnn.featurizers.radial.BesselRBF

class xnn.gnn.featurizers.radial.BesselRBF(n_rbf=8, cutoff=5.0, trainable=False, prefactor=None)

Bases: Module

Bessel radial basis (NequIP/DimeNet-style), smoother & fewer functions.

Expands an interatomic distance r into a set of n_rbf invariant radial features using the normalized sinc/Bessel functions prefactor * sin(n*pi*r/rc) / r for n = 1, ..., n_rbf. Compared with a Gaussian basis this is smoother and needs fewer functions to cover the cutoff sphere.

Parameters:

• n_rbf (int, optional) – Number of Bessel basis functions (output width). Default is 8.

• cutoff (float, optional) – Cutoff radius rc used to set the basis frequencies and the normalization. Default is 5.0.

• trainable (bool, optional) – If True, the n * pi frequencies are a learnable torch.nn.Parameter (the NequIP BesselBasis default); otherwise a fixed buffer (the MACE default). Default is False.

• prefactor (float, optional) – Overall normalization factor. None (default) uses the DimeNet/MACE convention sqrt(2 / cutoff); the original NequIP uses 2 / cutoff.

Variables:

• freqs (Tensor) – The angular frequencies n * pi of shape (n_rbf,) – a registered buffer, or a torch.nn.Parameter when trainable.

• norm (float) – The scalar normalization prefactor.

forward(r)

Expand distances into the Bessel radial basis.

Parameters:

r (Tensor) – Interatomic distances of arbitrary shape (...). Values are clamped to a small positive minimum to avoid division by zero.

Returns:

The radial embedding of shape (..., n_rbf).

Return type:

Tensor