xnn.gnn.featurizers.radial.SoftDistanceTransform#

class xnn.gnn.featurizers.radial.SoftDistanceTransform(alpha=4.0)[source]#

Bases: Module

Tanh-based soft lower clamp of the interatomic distance.

Smoothly interpolates between a floor of p0 = 3/4 r0 at short range and the identity at long range,

T(r) = p0 + (r - p0) * 1/2 (1 + tanh(alpha' (r - m)))

with r0 the sum of the two elements’ covalent radii, p1 = 4/3 r0, midpoint m = (p0 + p1) / 2 and steepness alpha' = alpha / (p1 - p0). This is the MACE distance_transform="Soft" option.

Parameters:

alpha (float, optional) – Dimensionless steepness of the switch, by default 4.0.

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

Apply the soft clamp 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