Radial basis: `bessel_basis`

Role

A radial basis turns an interatomic distance into a vector of smooth invariant features \(\bm\varphi(r)\). Every distance dependence in the network — radial filters of convolutions, attention scores, encoder amplitudes — is a learned function of these features, so the basis sets the radial resolution of the model.

Built-in members of this kind: the Gaussian soft one-hot basis of the flat MANDALA network and gaussian_basis (fixed or trainable centres), and the sine basis with polynomial cutoff cutoff_radial_expansion.

The new layer

Spherical Bessel functions of order zero are the eigenfunctions of the Laplacian in a ball with Dirichlet boundary at \(r_c\):

\[\varphi_n(r)=\sqrt{\frac{2}{r_c}}\;\frac{\sin(n\pi r/r_c)}{r}\;u\!\left(\frac{r}{r_c}\right), \qquad n=1,\dots,N .\]

They are orthogonal on \([0,r_c]\) with weight \(r^2\), so different basis functions carry non-redundant radial information. The polynomial envelope

\[u(x)=1-\frac{(p+1)(p+2)}{2}x^{p}+p(p+2)x^{p+1}-\frac{p(p+1)}{2}x^{p+2}, \qquad u(1)=u'(1)=u''(1)=0,\]

brings every feature smoothly to zero at the cutoff, so the predicted operator varies smoothly when an atom enters or leaves the neighbourhood — important for molecular dynamics and for differentiable geometry. Optional trainable frequencies let the model adapt the radial resolution.

Onsite edges. At \(r=0\), \(\sin(n\pi r/r_c)/r\) is finite (\(n\pi/r_c\)) but its direct evaluation divides by zero. Writing it as

\[\frac{\sin(n\pi r/r_c)}{r}=\frac{n\pi}{r_c}\,\operatorname{sinc}\!\left(\frac{nr}{r_c}\right), \qquad \operatorname{sinc}(x)=\frac{\sin\pi x}{\pi x},\]

gives exact values and gradients everywhere (checklist item 7).

Symmetry. Lengths are invariant, so any function of them is invariant.

Implementation

class BesselBasis(nn.Module):
    r"""``phi_n(r) = sqrt(2/r_c) * sin(n pi r / r_c) / r * u(r / r_c)``.

    ``u`` is the C^2 polynomial envelope of order ``p``. The ``sin(x)/r`` factor
    is evaluated through ``sinc`` so that onsite edges (``r = 0``) are finite.
    """

    def __init__(self, cutoff, n_basis=8, envelope_exponent=6, trainable=False):
        super().__init__()
        self.cutoff = float(cutoff)
        self.p = int(envelope_exponent)
        frequencies = torch.arange(1, n_basis + 1, dtype=torch.get_default_dtype())
        if trainable:
            self.frequencies = nn.Parameter(frequencies)
        else:
            self.register_buffer("frequencies", frequencies)

    def envelope(self, x):
        p = self.p
        u = (
            1
            - (p + 1) * (p + 2) / 2 * x**p
            + p * (p + 2) * x ** (p + 1)
            - p * (p + 1) / 2 * x ** (p + 2)
        )
        return torch.where(x < 1, u, torch.zeros_like(u))

    def forward(self, lengths):
        r = lengths.reshape(-1, 1)
        x = r / self.cutoff
        # sin(n pi x) / r == (n pi / r_c) * sinc(n x)
        radial = (
            self.frequencies * math.pi / self.cutoff * torch.sinc(self.frequencies * x)
        )
        return math.sqrt(2 / self.cutoff) * radial * self.envelope(x)

The factory requires one length per edge and declares N invariant scalars on the same support:

@register_transform("bessel_basis")
def _bessel_basis(options, inputs, mapper):
    (lengths,) = inputs
    if lengths.domain != "edge" or lengths.dim != 1:
        raise ValueError("bessel_basis expects one edge length per edge")
    module = BesselBasis(**options)
    width = len(module.frequencies)
    return module, [
        FeatureSpec("edge", irreps=f"{width}x0e", width=width, support=lengths.support)
    ]

Usage

- id: radial
  transform: bessel_basis
  inputs: [edge_lengths]
  outputs: [radial]
  options: {cutoff: 6.0, n_basis: 8, envelope_exponent: 6, trainable: false}

radial is a 8x0e edge feature and can be consumed by any layer that expects invariant radial features, including the built-in radial_tensor_product and projected_attention.

Variations

Replace the envelope by a learned monotone function to let the network choose its effective range.
Scale the argument per species pair by a characteristic bond length \(r_0^{(ab)}\) (as pair_distance_normalization: pair_r0 does for the Gaussian basis) to give every bond type the same resolution near its equilibrium distance.

Radial basis: bessel_basis