xnn.hybrid.featurizers.radial.ExponentialBernsteinRBF#

class xnn.hybrid.featurizers.radial.ExponentialBernsteinRBF(n_rbf=16, gamma=0.9448630623102202, trainable=True, exp_weighting=False)[source]#

Bases: Module

Exponential Bernstein polynomials, the radial basis of SpookyNet.

Unke et al., Nat. Commun. 12, 7273 (2021), eqs 14 and 15: the distance is mapped to x = exp(-gamma r) in (0, 1] and expanded in the K Bernstein polynomials of degree K - 1,

b_k(x) = binom(K - 1, k) x^k (1 - x)^(K - 1 - k), k = 0, ..., K - 1,

which form a partition of unity and approximate any continuous function of x as K grows. The decay gamma is one positive parameter shared by all functions (kept positive through a softplus), so the rate at which a learned radial function can vary falls off with the distance. The polynomials are evaluated in log space, log binom + k log x + (K - 1 - k) log(1 - x), which keeps the large binomials and small powers apart. A cutoff envelope is not included; the model multiplies it in.

Parameters:

  • n_rbf (int, optional) – Number of basis functions K, by default 16.

  • gamma (float, optional) – Initial decay gamma in inverse length units, by default 0.5 / bohr (0.9449 / Angstrom), the paper’s value.

  • trainable (bool, optional) – Learn gamma, by default True (the paper).

  • exp_weighting (bool, optional) – Multiply every function by x = exp(-gamma r), an option of the reference code that is not part of the paper’s eq 14. Default False.

Variables:

gamma_raw (Tensor) – softplus^-1(gamma), a parameter when trainable.

property gamma: Tensor#

The decay gamma = softplus(gamma_raw).

exact_constants()[source]#

Rebuild float64 log-binomials exactly after a cast; whether they changed.

Return type:

bool

forward(r)[source]#

Expand distances into the basis.

Parameters:

r (Tensor) – Interatomic distances (E,), positive.

Returns:

(E, n_rbf); column k is b_k(exp(-gamma r)).

Return type:

Tensor