xnn.gnn.featurizers.spherical.SphericalHarmonicEdgeEmbedding#

class xnn.gnn.featurizers.spherical.SphericalHarmonicEdgeEmbedding(l_max=2, n_rbf=8, cutoff=5.0, p=6, radial_type='bessel', trainable_rbf=False, rbf_prefactor=None)[source]#

Bases: Featurizer

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

For every edge it produces two complementary encodings of the relative displacement vector: the direction encoded as real spherical harmonics Y_l(r_hat) (equivariant edge attributes feeding the tensor-product paths) and the interatomic distance encoded with a smooth radial basis multiplied by a polynomial cutoff envelope (invariant scalar weights).

Parameters:

  • l_max (int, optional) – Maximum spherical-harmonic degree l. The spherical-harmonic irreps are o3.Irreps.spherical_harmonics(l_max). Default is 2.

  • n_rbf (int, optional) – Number of radial basis functions (width of the invariant radial embedding). Default is 8.

  • cutoff (float, optional) – Cutoff radius (in the length units of the coordinates) used by both the radial basis and the cutoff envelope. Default is 5.0.

  • p (int, optional) – Polynomial degree of the PolynomialCutoff envelope. Default is 6.

  • radial_type (str, optional) – Radial basis to use: "bessel" for BesselRBF or "gaussian" for GaussianRBF. Default is "bessel".

  • trainable_rbf (bool, optional) – Make the Bessel frequencies learnable (NequIP’s BesselBasis default). Only meaningful for radial_type="bessel". Default is False.

  • rbf_prefactor (float, optional) – Normalization prefactor of the Bessel basis. None (default) is the DimeNet/MACE convention sqrt(2/cutoff); NequIP uses 2/cutoff. Only meaningful for radial_type="bessel".

Variables:

  • irreps_sh (o3.Irreps) – The e3nn irreps of the spherical-harmonic edge attributes.

  • sph (o3.SphericalHarmonics) – The (normalized, component-normalization) spherical-harmonics module.

  • rbf (torch.nn.Module) – The radial basis module (Bessel or Gaussian).

  • envelope (PolynomialCutoff) – The smooth cutoff envelope applied to the radial embedding.

Raises:
property output_dim: int#

Width of the invariant scalar radial embedding.

Returns:

The number of radial basis functions n_rbf.

Return type:

int

embed(vec)[source]#

TorchScript-compatible core: edge vectors in, geometric features out.

This is the tensor-only path used by the scriptable node_energy model cores (LAMMPS/TorchScript deployment); forward() reuses it.

Parameters:

vec (Tensor) – Edge displacement vectors of shape (E, 3) (already accounting for any periodic cell shifts).

Returns:

(edge_length, edge_sh, edge_radial) – the interatomic distances (E,), the spherical-harmonic edge attributes (E, irreps_sh.dim) and the enveloped radial embedding (E, n_rbf).

Return type:

tuple of (Tensor, Tensor, Tensor)

forward(data)[source]#

Compute the per-edge geometric features.

Parameters:

data (AtomicGraph) – The atomic graph; data.edge_vectors() supplies the relative displacement vector for each of the E edges.

Returns:

Mapping with keys:

  • "edge_vec" : Tensor of shape (E, 3), the relative displacement vectors.

  • "edge_length" : Tensor of shape (E,), the interatomic distances.

  • "edge_sh" : Tensor of shape (E, irreps_sh.dim), the equivariant spherical-harmonic edge attributes Y_l(r_hat).

  • "edge_radial" : Tensor of shape (E, n_rbf), the invariant radial embedding scaled by the cutoff envelope.

Return type:

dict[str, Tensor]