Fix SphericalBesselFunction(smooth=True): float32 basis constants, dtype-dependent output, dead lru_cache

[toohardtogetname]

Affects every model constructed with use_smooth=True (M3GNet, TensorNet and
QET all expose the flag; the TensorNet MatPES PES checkpoints ship with it
enabled). Reproduced on matgl 4.0.3 and on main as of 2026-08-22, torch 2.9.1,
sympy 1.14.0.

float32 output does not change at all (0 of 200,001 sampled points move, 0.00
ULP), so this is not a "your results were wrong" report and nothing needs
retraining. What it costs today is reproducibility, float64 accuracy, and about
220 ms of sympy.simplify per model constructed.

Symptom 1: matgl's two implementations of the same basis disagree

matgl contains two implementations of the smooth spherical-Bessel basis:
SphericalBesselFunction(smooth=True) (sympy, lambdified) and the hand-written
spherical_bessel_smooth(). Evaluated on the same points in float64 they differ:

import torch, matgl
from matgl.layers._basis import SphericalBesselFunction, spherical_bessel_smooth

matgl.float_th = torch.float64
r = torch.linspace(0.3, 5.0, 4097, dtype=torch.float64)
a = SphericalBesselFunction(3, 3, 5.0, smooth=True)(r)
b = spherical_bessel_smooth(r, cutoff=5.0, max_n=3)
print((a - b).abs().max())        # -> 1.804e-08

spherical_bessel_smooth is the correct one: it spells its constants out
exactly (sqrt2 = 1.4142135623730951, pi_local = 3.141592653589793).

This matters beyond a numerical curiosity, because matgl already treats the two
as interchangeable. matgl.ext._lammps._SmoothSBFExpansion substitutes
spherical_bessel_smooth for BondExpansion(rbf_type='SphericalBessel', smooth=True) when exporting a model for LAMMPS, on the stated grounds that it
"is mathematically identical". It should be — and after this fix it is. Today a
LAMMPS-exported model and the Python model evaluate slightly different bases.

Symptom 2: the basis depends on the ambient default dtype

import torch, inspect
from matgl.layers._basis import SphericalBesselFunction

for dt in (torch.float32, torch.float64):
    torch.set_default_dtype(dt)
    f = SphericalBesselFunction(3, 3, 5.0, smooth=True).funcs[0]
    print(dt, inspect.getsource(f).strip().splitlines()[-1].strip())

torch.float32  return (0.565685439521744*sin(...) + 0.282842719760872*sin(...))/r
torch.float64  return (0.565685424949238*sin(...) + 0.282842712474619*sin(...))/r

The same model definition and the same checkpoint therefore produce different
radial features depending on whether something earlier in the process called
torch.set_default_dtype. Nothing warns you.

The exact value is 2*sqrt(2)/5 = 0.565685424949238; the float32 path is high by
2.576e-08 relative, and every coefficient of every n is off by that same
factor.

Cause

__init__ stores the cutoff as a buffer, so its dtype follows the global default:

self.register_buffer("cutoff", torch.tensor(cutoff))

and the smooth branch passes that tensor into the symbolic builder unchanged:

@lru_cache(maxsize=128)
def _calculate_smooth_symbolic_funcs(self) -> list:
    return _get_lambda_func(max_n=self.max_n, cutoff=self.cutoff)

_get_lambda_func evaluates cutoff**1.5 while assembling the basis prefactor.
With a float32 buffer that is a float32 op, so the prefactor is rounded once and
the rounding scales every generated coefficient.

The non-smooth branch of the same __init__ already avoids this — it writes
factor = sqrt(2.0 / float(cutoff) ** 3).

Symptom 3: the symbolic functions are rebuilt for every module

_get_lambda_func is lru_cached, but a tensor argument hashes by identity, so
the cache never hits and each construction re-runs sympy.simplify:

	measured
_get_lambda_func on a cache miss	~220 ms (592 ms on the very first call)
_get_lambda_func on a cache hit	0.9 µs
constructing 5 modules, today	1105 ms
constructing 5 modules, with the fix	217 ms

Separately, the @lru_cache on _calculate_smooth_symbolic_funcs keys on
self. The method is called exactly once per instance, so that cache can never
hit either — while holding a strong reference to every module constructed, up to
128 of them. Eight modules dropped and garbage-collected leave eight alive.
_calculate_symbolic_funcs (the non-smooth branch) carries the same decorator
and the same problem.

Fix

See fix.diff: pass float(self.cutoff), and drop the two method-level
lru_cache decorators that cannot hit.

Backward compatibility: float32 output is bit-identical

The correction is smaller than a float32 ULP everywhere, so no existing float32
result changes and no retraining is needed. Over 200,001 points on
r in [0.3, 5.0]:

	before vs after
float32 elements that change	0 / 200,001
float32 max difference	0.00 ULP
float64 max difference, normalised to the basis amplitude	2.576e-08

Agreement with the closed form
sqrt(2)(2 sin(pi r/5) + sin(2 pi r/5))/(5 r) improves from 2.576e-08 to
1.268e-15, and the two implementations above come into agreement.

Who this actually affects

• float32 inference and MD: nobody. Output is bit-identical.
• Cross-implementation validation in float64. Anyone checking a
  reimplementation (a LAMMPS pair style, a rewritten kernel, another framework)
  against matgl meets a 2.6e-08 floor that is not in their own code. This is
  how the bug was found.
• Workflows that set torch.set_default_dtype(torch.float64) — common for
  phonons and finite-difference work — silently get a different basis than
  workflows that do not.
• Anyone constructing more than one model per process: ensembles, sweeps,
  test suites, all pay ~220 ms of avoidable sympy work per model.

Tests

test_smooth_sbf_basis.py adds three regression tests: closed-form agreement,
default-dtype independence, and cache reuse. All three fail on 4.0.3 and pass
with the fix.

🤖 Generated with Claude Code

[shyuep]

Thanks!