Featurizers#

A featurizer turns an AtomicGraph into model inputs. All featurizers subclass Featurizer (an nn.Module with an output_dim property and a forward(data) method), so they can be trained, scripted, and composed like any other module, and used standalone for analysis.

Shared basis functions (xnn.common.featurizers)#

  • GaussianRBF: Gaussian radial basis expansion of distances (used by SchNet); the optional gamma fixes the width explicitly (SchNet’s gamma = 10 Å-2) instead of tying it to the center spacing.

  • CosineCutoff: smooth cosine cutoff envelope.

Descriptor featurizers (xnn.dnn.featurizers)#

Invariant per-atom descriptors for HDNNP/ANI-style models:

  • RadialSymmetryFunctions: Behler–Parrinello G2 radial symmetry functions.

  • AngularSymmetryFunctions: angular symmetry functions over atomic triplets (built with build_triplets()).

  • AEV: the ANI atomic environment vector (radial + angular parts, per species pair).

from xnn.dnn.featurizers import AEV

aev = AEV(species=[1, 6, 8])
descriptor = aev(graph)      # (N, aev.output_dim), rotation invariant

Equivariant featurizers (xnn.gnn.featurizers)#

Edge attributes for the GNN models:

  • SphericalHarmonicEdgeEmbedding: the standard NequIP/MACE/Allegro edge embedding, with edge lengths, real spherical harmonics \(Y_{lm}(\hat r_{ij})\) up to l_max, and a radial expansion.

  • CartesianAngularBasis: the CACE angular basis, i.e. the Cartesian monomials \(x^{l_x} y^{l_y} z^{l_z}\) up to l_max, spanning the same space as the spherical harmonics per total \(l\) without e3nn.

  • BesselRBF: (trainable) Bessel radial basis.

  • PolynomialCutoff: the polynomial cutoff envelope of NequIP/MACE.

from xnn.gnn.featurizers import SphericalHarmonicEdgeEmbedding

embed = SphericalHarmonicEdgeEmbedding(l_max=2)
lengths, edge_sh, edge_radial = embed.embed(graph.edge_vectors())

from xnn.gnn.featurizers import CartesianAngularBasis

basis = CartesianAngularBasis(l_max=3)
unit = graph.edge_vectors()
angular = basis(unit / unit.norm(dim=-1, keepdim=True))   # (E, 20)

Writing your own featurizer is a small task; see Extending xnn.