Radial Basis Functions

Sets of radial basis functions are used to expand the distances (or other molecular features) between atoms on a fixed-sized vector. For instance, this is the main transformation of the distances in the SchNet model.

class mlcg.nn.radial_basis.GaussianBasis(cutoff, num_rbf=50, trainable=False)

Class that generates a set of equidistant 1-D gaussian basis functions scattered between a specified lower and upper cutoff:

\[f_n = \exp{ \left( -\gamma(r-c_n)^2 \right) }\]

Parameters:

• cutoff (Union[int, float, _Cutoff]) – Defines the smooth cutoff function. If a float is provided, it will be interpreted as an upper cutoff and an IdentityCutoff will be used between 0 and the provided float. Otherwise, a chosen _Cutoff instance can be supplied.

• num_rbf (int) – The number of gaussian functions in the basis set.

• trainable (bool) – If True, the parameters of the gaussian basis (the centers and widths of each function) are registered as optimizable parameters that will be updated during backpropagation. If False, these parameters will be instead fixed in an unoptimizable buffer.

class mlcg.nn.radial_basis.ExpNormalBasis(cutoff, num_rbf=50, trainable=True)

Class for generating a set of exponential normal radial basis functions, as described in [Physnet] . The functions have the following form:

\[f_n(r_{ij};\alpha, r_{low},r_{high}) = f_{cut}(r_{ij},r_{low},r_{high}) \times \exp\left[-\beta_n \left(e^{\alpha (r_{ij} -r_{high}) } - \mu_n \right)^2\right]\]

where

\[\alpha = 5.0/(r_{high} - r_{low})\]

is a distance rescaling factor, and, by default

\[f_{cut} ( r_{ij},r_{low},r_{high} ) = \cos{\left( r_{ij} \times \pi / r_{high}\right)} + 1.0\]

represents a cosine cutoff function.

Parameters:

• cutoff (Union[int, float, _Cutoff]) – Defines the smooth cutoff function. If a float is provided, it will be interpreted as an upper cutoff and a CosineCutoff will be used between 0 and the provided float. Otherwise, a chosen _Cutoff instance can be supplied.

• num_rbf (int) – The number of functions in the basis set.

• trainable (bool) – If True, the parameters of the basis (the centers and widths of each function) are registered as optimizable parameters that will be updated during backpropagation. If False, these parameters will be instead fixed in an unoptimizable buffer.

class mlcg.nn.radial_basis.RIGTOBasis(cutoff, nmax=5, lmax=5, sigma=0.4, mesh_size=300)

This radial basis is the effective basis when expanding an atomic density smeared by Gaussains of width \(\sigma\) on a set of \(n_{max}\) orthonormal Gaussian Type Orbitals (GTOs) and \(l_{max}+1\) Spherical Harmonics (SPHs) namely. This radial basis set is interpolated using natural cubic splines for efficiency and the cutoff is included into the splined functions.

The basis is defined as

\[R_{nl}(r) = f_c(r) \mathcal{N}_n \frac{\Gamma(\frac{n+l+3}{2})}{\Gamma(l+\frac{3}{2})} c^l r^l(c+b_n)^{-\frac{(n+l+3)}{2}} {}_1F_1\left(\frac{n+l+3}{2},l+\frac{3}{2};\frac{c^2 r^2}{c+b_n}\right),\]

where \({}_1F_1\) is the confluent hypergeometric function, \(\Gamma\) is the gamma function, \(f_c\) is a cutoff function, \(b_n=\frac{1}{2\sigma_n^2}\), \(c= 1 / (2\sigma^2\), \(\sigma_n = r_\text{cut} \max(\sqrt{n},1)/n_{max}\) and \(\mathcal{N}_n^2 = \frac{2}{\sigma_n^{2n + 3}\Gamma(n + 3/2)}\).

For more details on the derivation, refer to appendix A.

Parameters:

• nmax (int) – number of radial basis

• lmax (int) – maximum spherical order (lmax included so there are lmax+1 orders)

• sigma (float) – smearing of the atomic density

• cutoff (Union[int, float, _Cutoff]) – Defines the smooth cutoff function. If a float is provided, it will be interpreted as an upper cutoff and a CosineCutoff will be used between 0 and the provided float. Otherwise, a chosen _Cutoff instance can be supplied.

• mesh_size (int) – number of points used to interpolate with splines the radial basis spanning uniformly the range difined by the cutoff \([0, r_c]\).

class mlcg.nn.radial_basis.SpacedExpBasis(cutoff, sigma_min=0.25, sigma_factor=2.0, mean_spacing=2.0, trainable=True)

Class for generating a set of exponential normal radial basis functions, with means and standard deviations designed to capture the physics around 2 \(\AA\). The functions have an exponential form with the following means and std:

\[\begin{split}&\mu_0 = 0 &,\quad &\sigma_0 = \sigma_f/s \\ &\mu_1 = \sigma_f &,\quad &\sigma_1 = \sigma_{min} \\ &\mu_2 = \mu_1 + s\,\sigma_1 &,\quad &\sigma_2 = \sigma_f \sigma_1 \\ &\vdots &\vdots \; &\vdots \\ &\mu_n = \mu_{n-1} + s\,\sigma_{n-1} &,\quad &\sigma_n = \sigma_f \sigma_{n-1}\end{split}\]

with \(\sigma_{min}\) the minimum sigma, \(\sigma_f\) the sigma factor and \(s\) the mean spacing

Parameters:

• cutoff (Union[int, float, _Cutoff]) – Defines the smooth cutoff function. If a float is provided, it will be interpreted as an IdentityCutoff ranging over [0, cutoff]. Otherwise, a chosen _Cutoff instance can be supplied.

• sigma_min (float) – \(\sigma_{min}\), width of first non-zero-centered basis function

• sigma_factor (float) – \(\sigma_f\), location of first non-zero-centered basis function and multiplicative factor to spread std of each new peak by

• mean_spacing (float) – \(s\) this time previous sigma indicates how much to distance the mean of subsequent gaussian by

• trainable (bool) – If True, the parameters of the basis (the centers and widths of each function) are registered as optimizable parameters that will be updated during backpropagation. If False, these parameters will be instead fixed in an unoptimizable buffer.

class mlcg.nn.angular_basis.SphericalHarmonics(lmax, normalize=False, normalization='integral', irreps_in=None)

JITable module version of real Spherical harmonics

Polynomials defined on the 3d space \(Y^l: \mathbb{R}^3 \longrightarrow \mathbb{R}^{2l+1}\)

Usually restricted on the sphere (with normalize=True) \(Y^l: S^2 \longrightarrow \mathbb{R}^{2l+1}\)

who satisfies the following properties:

• are polynomials of the cartesian coordinates x, y, z

• is equivariant \(Y^l(R x) = D^l(R) Y^l(x)\)

• are orthogonal \(\int_{S^2} Y^l_m(x) Y^j_n(x) dx = \text{cste} \; \delta_{lj} \delta_{mn}\)

The value of the constant depends on the choice of normalization.

It obeys the following property:

\[ \begin{align}\begin{aligned}Y^{l+1}_i(x) &= \text{cste}(l) \; & C_{ijk} Y^l_j(x) x_k\\\partial_k Y^{l+1}_i(x) &= \text{cste}(l) \; (l+1) & C_{ijk} Y^l_j(x)\end{aligned}\end{align} \]

Where \(C\) are the Wigner 3-j symbols.

Note

This function matches with the following table of standard real spherical harmonics from Wikipedia when normalize=True, normalization='integral' and is called with the argument in the order y,z,x (instead of x,y,z).

Parameters:

• l (int or list of int) – degree of the spherical harmonics.

• x (torch.Tensor) – tensor \(x\) of shape (..., 3).

• normalize (bool) – whether to normalize the x to vectors that lie on the unit sphere before projecting onto the spherical harmonics

• normalization ({'integral', 'component', 'norm'}) – normalization of the output tensors — note that this option is independent of normalize, which controls the processing of the input, rather than the output. Valid options: * component: \(\|Y^l(x)\|^2 = 2l+1, x \in S^2\) * norm: \(\|Y^l(x)\| = 1, x \in S^2\), component / sqrt(2l+1) * integral: \(\int_{S^2} Y^l_m(x)^2 dx = 1\), component / sqrt(4pi)

Returns:

a list of tensors of shapes [lmax][n_neighbor, m] where \(m<=l\)

\[Y^l(x)\]

Return type:

List[torch.Tensor]