xnn.common.featurizers.base.Featurizer

class xnn.common.featurizers.base.Featurizer(*args, **kwargs)

Bases: Module, ABC

Abstract base class for all featurizers.

A featurizer is a plain nn.Module that maps an AtomicGraph to model inputs – either invariant per-atom descriptors or equivariant edge/node embeddings. Subclasses share the output_dim contract so they can be discovered and swapped uniformly.

Notes

Featurizers are independently usable: an instance can be called on a graph without any surrounding model, e.g. to inspect descriptors.

abstract property output_dim: int

Size of the per-atom feature vector.

Returns:

The per-atom descriptor width for invariant featurizers. Equivariant featurizers that return dicts may report the scalar (l=0) channel width here, or raise if not meaningful.

Return type:

int

abstractmethod forward(data)

Compute features for an atomic graph.

Parameters:

data (AtomicGraph) – The input atomic graph (positions, neighbor lists, etc.).

Returns:

A single (N, output_dim) per-atom descriptor tensor for invariant featurizers, or a dict of edge tensors for equivariant featurizers.

Return type:

torch.Tensor or dict[str, torch.Tensor]