xnn.gnn.featurizers.radial.AgnesiDistanceTransform#

class xnn.gnn.featurizers.radial.AgnesiDistanceTransform(a=1.0805, q=0.9183, p=4.5791)[source]#

Bases: Module

Agnesi radial transform, rescaled per edge by covalent radii.

Maps the interatomic distance r to

T(r) = 1 / (1 + a (r/r0)^q / (1 + (r/r0)^(q-p)))

with r0 the mean covalent radius of the two edge elements, 0.5 (rcov_i + rcov_j). This is the transform of the MACE-MP foundation models (distance_transform="Agnesi"); the functional form comes from the radial transformations of ACEpotentials.jl (J. Chem. Phys. 159, 164101, 2023). The constants a, q, p and the covalent-radii table are registered buffers, so values stored in a trained checkpoint carry over on load_state_dict.

Parameters:

  • a (float, optional) – The transform constants, by default the published MACE values (1.0805, 0.9183, 4.5791).

  • q (float, optional) – The transform constants, by default the published MACE values (1.0805, 0.9183, 4.5791).

  • p (float, optional) – The transform constants, by default the published MACE values (1.0805, 0.9183, 4.5791).

forward(r, atomic_numbers, edge_index)[source]#

Apply the Agnesi transform per edge.

Parameters:

  • r (Tensor) – Interatomic distances, shape (E,).

  • atomic_numbers (Tensor) – Per-node atomic numbers, shape (N,).

  • edge_index (Tensor) – Edge index of shape (2, E).

Returns:

Transformed distances, shape (E,).

Return type:

Tensor