xnn.transformer.featurizers.radial.ExpNormalSmearing#

class xnn.transformer.featurizers.radial.ExpNormalSmearing(n_rbf=32, cutoff=5.0, cutoff_lower=0.0, trainable=True)[source]#

Bases: Module

Exponential-normal radial basis expansion (TorchMD-Net-style).

Expands a scalar distance r onto n_rbf Gaussians in exponential distance space, multiplied by a cosine cutoff envelope:

phi_k(r) = cos_cutoff(r) *
           exp(-beta_k * (exp(alpha * (cutoff_lower - r)) - mu_k) ** 2)

with alpha = 5 / (cutoff_upper - cutoff_lower). The centers mu_k are spaced linearly from exp(cutoff_lower - cutoff_upper) to 1 and the widths beta_k are initialised to the TorchMD-Net/PhysNet default (2 / n_rbf * (1 - exp(cutoff_lower - cutoff_upper))) ** -2.

Parameters:

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

  • cutoff (float, optional) – Upper cutoff cutoff_upper in the same units as the distances. Default is 5.0.

  • cutoff_lower (float, optional) – Lower cutoff, by default 0.0.

  • trainable (bool, optional) – If True (the default) the centers means and widths betas are learnable torch.nn.Parameter; otherwise fixed buffers.

Variables:

  • means (Tensor) – The Gaussian centers mu_k of shape (n_rbf,) (parameter or buffer depending on trainable).

  • betas (Tensor) – The Gaussian inverse-widths beta_k of shape (n_rbf,).

  • alpha (float) – The exponential-space scale 5 / (cutoff_upper - cutoff_lower).

reset_parameters()[source]#

Fill means and betas with their default values in place.

The defaults follow the TorchMD-Net/PhysNet convention: in exponential distance space the centers run linearly from exp(cutoff_lower - cutoff_upper) (the image of the upper cutoff) up to 1 (the image of the lower cutoff), and every Gaussian starts with the same width, matched to the center spacing.

Return type:

None

forward(r)[source]#

Expand distances onto the exponential-normal basis.

Parameters:

r (Tensor) – Interatomic distances of arbitrary shape (...).

Returns:

The radial embedding of shape (..., n_rbf), smoothly zero at and beyond the cutoff.

Return type:

Tensor