xnn.gnn.featurizers.cutoff.PolynomialCutoff#

class xnn.gnn.featurizers.cutoff.PolynomialCutoff(cutoff, p=6)[source]#

Bases: Module

Smooth p-degree polynomial envelope (used by NequIP/MACE), C^2 at rc.

Multiplicative envelope that decays smoothly from 1 at r = 0 to exactly 0 at the cutoff rc, with continuous value and derivatives (C^2) at rc. Applied to radial embeddings so that edge features vanish smoothly as neighbours leave the cutoff sphere.

Parameters:

  • cutoff (float) – Cutoff radius rc; the envelope is zero for r >= cutoff.

  • p (int, optional) – Polynomial degree controlling the smoothness/sharpness of the decay. Default is 6.

forward(r)[source]#

Evaluate the cutoff envelope at the given distances.

Parameters:

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

Returns:

The envelope values, same shape as r, decaying from 1 to 0 and exactly 0 for r >= cutoff.

Return type:

Tensor