Featurizers#

A featurizer turns an AtomicGraph into model inputs. All subclass Featurizer, an nn.Module with an output_dim property and forward(data), so they train, script and compose like any module and can be used on their own for analysis.

Shared basis functions (xnn.common.featurizers)#

  • GaussianRBF: Gaussian expansion of distances (SchNet), with an optional explicit width gamma.

  • CosineCutoff: smooth cosine envelope.

  • MollifierCutoff: the bump-function envelope exp(-r^2 / (r_c^2 - r^2)), smooth to every order (AIMNet2, SpookyNet).

Descriptors (xnn.dnn.featurizers)#

Invariant per-atom descriptors for the HDNNP and ANI models:

  • RadialSymmetryFunctions and AngularSymmetryFunctions: Behler-Parrinello G2 and angular symmetry functions.

  • AEV: the ANI atomic environment vector.

  • AtomCenteredSymmetryFunctions: the HDNNP descriptor, a list of radial (types 1, 2) and angular (types 3, 8, 9) functions per element, each with its own cutoff function (cosine, hyperbolic tangent, polynomial, hard), with RuNNer’s conventions and feature scaling.

from xnn.dnn.featurizers import AEV

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

Grid featurizers (xnn.cnn.featurizers)#

  • VoxelGrid: one cubic grid per atom (side 2 * cutoff, grid_size voxels per axis) with a density channel per species, each neighbor deposited as a Gaussian under a cosine envelope; differentiable in the positions, the input of the 3D steerable CNN and the conventional 3D CNN.

from xnn.cnn.featurizers import VoxelGrid

vox = VoxelGrid(species=[1, 6, 8], cutoff=4.0, grid_size=17)
grids = vox(graph)           # (N, 3, 17, 17, 17) density fields

  • SphericalGrid: one sphere per atom (radius radius, Driscoll-Healy grid of bandwidth bandwidth) with a channel per species holding the potential of the neighbors of that species, the molecular representation of the spherical CNN paper; differentiable in the positions.

from xnn.cnn.featurizers import SphericalGrid

sph = SphericalGrid(species=[1, 6, 7, 8, 16], cutoff=10.0, radius=0.48, bandwidth=10)
signals = sph(graph)         # (N, 5, 20, 20) spherical signals

Equivariant edge features (xnn.gnn.featurizers)#

  • SphericalHarmonicEdgeEmbedding: edge lengths, real spherical harmonics up to l_max and a radial expansion (NequIP, MACE, Allegro).

  • CartesianAngularBasis: the Cartesian monomials of CACE, spanning the same space without e3nn.

  • BesselRBF and PolynomialCutoff: the trainable Bessel basis and polynomial envelope of NequIP, MACE and DimeNet.

  • SphericalBesselBasis: the 2D spherical Fourier-Bessel basis of DimeNet and GemNet, n_spherical degrees times n_radial Bessel roots over a distance and an angle; plain PyTorch, no scipy.

  • real_spherical_harmonics(): real spherical harmonics as Cartesian polynomials, smooth at the poles; with the radial part above they form GemNet’s 3D basis over a distance and two angles.

from xnn.gnn.featurizers import SphericalHarmonicEdgeEmbedding, CartesianAngularBasis

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

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

Transformer pieces (xnn.transformer)#

ExpNormalSmearing (the exponential-normal radial basis) and EdgeMultiheadAttention are the building blocks of BAMBOO, kept separate for other attention-based models.

Hybrid radial basis (xnn.hybrid.featurizers)#

ExponentialBernsteinRBF: the Bernstein polynomials of exp(-gamma r) with a learned decay, the radial basis of SpookyNet.

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