xnn.gnn.models.base#
Shared bases for the GNN potentials (NequIP / MACE / Allegro / CACE).
GNNPotential holds the pieces every graph potential in this package
needs regardless of its angular basis: species bookkeeping (atomic-number <->
element-index table, one-hot node attributes) and the per-element reference
energy atom_ref. EquivariantGNN adds what the spherical-harmonic
models share on top: the
SphericalHarmonicEdgeEmbedding edge featurizer
and the e3nn irreps of the node/edge attributes.
Factored out of the individual model modules so nequip.py, mace.py and
allegro.py can all subclass EquivariantGNN while cace.py
(Cartesian basis, no spherical harmonics) subclasses GNNPotential
directly (see the package-structure convention: code common to several models
of one family lives in a shared module of that family).
Classes
|
Shared base for the E(3)-equivariant GNN potentials. |
|
Species bookkeeping shared by every GNN potential in this package. |