xnn.gnn.featurizers.radial#

Radial pieces of the equivariant GNN edge embedding.

The Bessel basis (NequIP/MACE/DimeNet-style) plus the chemistry-aware distance transforms some force fields warp their radial coordinate with before the basis expansion: the Agnesi transform (Batatia et al., MACE-MP-0, arXiv:2401.00096; the transform itself is from the radial transformations of ACEpotentials.jl, J. Chem. Phys. 159, 164101, 2023) and a tanh-based soft lower clamp. Both rescale each edge by the covalent radii of its two elements, so short bonds between small atoms and long bonds between large atoms map to comparable coordinates.

Classes

AgnesiDistanceTransform([a, q, p])

Agnesi radial transform, rescaled per edge by covalent radii.

BesselRBF([n_rbf, cutoff, trainable, prefactor])

Bessel radial basis (NequIP/DimeNet-style), smoother & fewer functions.

IdentityDistanceTransform(*args, **kwargs)

The trivial distance transform: distances pass through unchanged.

SoftDistanceTransform([alpha])

Tanh-based soft lower clamp of the interatomic distance.
Show Source