xnn.gnn.models.base.EquivariantGNN#

class xnn.gnn.models.base.EquivariantGNN(species, cutoff, l_max, n_rbf, p=6, radial_type='bessel', trainable_rbf=False, rbf_prefactor=None)[source]#

Bases: GNNPotential

Shared base for the E(3)-equivariant GNN potentials.

Extends GNNPotential with the pieces the spherical-harmonic models share and is subclassed by NequIP, MACE and Allegro: it constructs the shared SphericalHarmonicEdgeEmbedding edge featurizer plus the spherical-harmonic irreps used by the interaction blocks.

Parameters:

species (list[int]) – Atomic numbers of the elements the model supports, in order. The list index defines the element-index (channel) of each species.
cutoff (float) – Cutoff radius passed to the edge featurizer.
l_max (int) – Maximum spherical-harmonic degree; sets irreps_sh via o3.Irreps.spherical_harmonics(l_max).
n_rbf (int) – Number of radial basis functions in the edge featurizer.
p (int, optional) – Polynomial degree of the cutoff envelope, by default 6 (models pass their own convention, e.g. MACE’s num_polynomial_cutoff).
radial_type (str, optional) – Radial basis of the edge featurizer ("bessel"/"gaussian"), by default "bessel".
trainable_rbf (bool, optional) – Learnable Bessel frequencies (the NequIP convention), by default False.
rbf_prefactor (float or None, optional) – Bessel normalization prefactor; None (default) is the MACE/DimeNet sqrt(2/cutoff), NequIP passes 2/cutoff.

Variables:

node_attr_irreps (o3.Irreps) – Irreps of the one-hot species node attributes (n_species scalars).
irreps_sh (o3.Irreps) – Spherical-harmonic irreps of the equivariant edge attributes.
edge_feat (SphericalHarmonicEdgeEmbedding) – The shared edge featurizer.