xnn.common.featurizers.base

The Featurizer abstraction.

A Featurizer turns an AtomicGraph into model inputs – either invariant per-atom descriptors (symmetry functions, AEV) or equivariant edge/node embeddings (spherical-harmonic edge attributes). Featurizers are plain nn.Module subclasses and are independently usable: you can instantiate one and call it on a graph without any model, e.g. to inspect descriptors or to build your own model on top.

Two output conventions are used in this package, distinguished by what a model needs:

• Invariant featurizers (RadialSymmetryFunctions, AngularSymmetryFunctions, AEV) return a single (N, output_dim) per-atom descriptor tensor.

• Equivariant featurizers (SphericalHarmonicEdgeEmbedding) return a dict of edge tensors (geometry + spherical harmonics + radial embedding) consumed by the GNN interaction blocks.

Both subclass Featurizer so they share the output_dim contract and can be discovered / swapped uniformly.

Classes

Featurizer(*args, **kwargs)

Abstract base class for all featurizers.