SphericalExpansion produces NaN for self-pairs (zero-length displacement vectors)

[sirmarcel]

Problem

When the neighbor list includes self-pairs (R_ij = [0, 0, 0]), SphericalHarmonics produces NaN for all l > 0 components. This propagates through the full expansion and its backward pass. Self-pairs are the standard convention in featomic, LAMMPS, and most atomistic ML frameworks.

The correct result: l = 0 gets a finite contribution (Y_0^0 is constant, radial basis at r = 0 is nonzero for l = 0), and l > 0 gets exactly zero (since r^l → 0).

Root cause

sphericart computes Y_l^m(R/|R|). For R = [0, 0, 0], the normalization is division by zero → NaN for l > 0:

import sphericart.torch
sh = sphericart.torch.SphericalHarmonics(l_max=2)
print(sh.compute(torch.tensor([[0.0, 0.0, 0.0]])))
# tensor([[0.2821, nan, nan, nan, nan, nan, nan, nan, nan]])

This may be a sphericart bug — arguably sphericart should return [Y_0^0, 0, 0, ...] for zero-length vectors. But even if sphericart doesn't fix it, torch-spex can handle it.

Suggested fix

A single torch.where after the computation is not enough — the forward would be clean, but torch.where computes gradients for both branches, so the NaN in the angular embedding would still produce NaN gradients during backward.

Need two torch.where calls — one before sphericart to sanitize the input, one after to zero out the l > 0 result:

is_zero = (r_ij < eps)  # [pair] — detect self-pairs

# 1. Give sphericart a safe input so no NaN enters the autograd graph
safe_R = torch.where(is_zero.unsqueeze(-1), unit_x, R_ij)
angular = sph.compute(safe_R)  # clean forward, clean backward

# 2. Zero out l > 0 for self-pairs (Y_l^m(unit_x) is nonzero, but shouldn't contribute)
for l in range(1, max_angular + 1):
    angular[l] = torch.where(is_zero.unsqueeze(-1), 0.0, angular[l])

The first where ensures sphericart never sees [0, 0, 0], so no NaN enters the computation graph at all. The second ensures the self-pair angular contribution is exactly zero for l > 0 (not whatever Y_l^m(unit_x) happens to give). Backward through both where ops is clean because neither branch contains NaN.

Current workaround

Adding a tiny epsilon offset to self-pair displacements (R_ij += 1e-30 on self-pairs). Works for forward but the gradient for self-pairs is slightly wrong (though physically it should be zero anyway since dR_ii/dpos_i = 0).

---

Filed by Claude Code on behalf of @sirmarcel