xnn.gnn.featurizers.spherical#

Equivariant edge featurization for E(3) GNNs (NequIP / MACE / Allegro).

Computes, per edge, the geometric inputs the equivariant interaction blocks need: the direction encoded as spherical harmonics Y_l(r_hat) (equivariant edge attributes) and the length encoded with a smooth radial basis (invariant weights for the tensor-product paths).

Requires e3nn: pip install “xnn[gnn]”

Classes

SphericalHarmonicEdgeEmbedding([l_max, ...])

Per-edge geometric featurizer for E(3)-equivariant GNNs.