Fix smooth spherical-Bessel basis: build constants in float64, make the symbolic-function cache effective (#826)

SphericalBesselFunction(smooth=True) fed its cutoff BUFFER into the
symbolic builder, so cutoff**1.5 was evaluated in the ambient default
dtype: under float32 the basis prefactor is rounded once and every
generated coefficient is scaled by 1 + 2.6e-8. The basis therefore
disagreed with matgl's own spherical_bessel_smooth() (which the LAMMPS
export path substitutes as 'mathematically identical') and silently
changed with torch.set_default_dtype. Passing float(self.cutoff) fixes
both, and float32 output is bit-identical before/after (the correction is
below one float32 ULP; nothing needs retraining).

The same change makes the module-level lru_cache on _get_lambda_func
actually hit (a tensor argument hashes by identity), saving ~220 ms of
sympy.simplify per constructed module. The two method-level lru_cache
decorators are removed: keyed on self they can never hit, while pinning up
to 128 dropped modules against garbage collection.

Three regression tests: closed-form agreement, default-dtype independence,
cache reuse. All three fail on 4.0.3 and pass with this change.

Co-authored-by: Longyun Shen <longyun.shen@nus.edu>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

src/matgl/layers/_basis.py

@@ -35,7 +35,6 @@
 
 from __future__ import annotations
 
-from functools import lru_cache
 from math import pi, sqrt
 
 import sympy
@@ -129,7 +128,6 @@ def __init__(self, max_l: int, max_n: int = 5, cutoff: float = 5.0, smooth: bool
             self.register_buffer("roots_slice", roots_slice, persistent=False)
             self.register_buffer("inv_norm", inv_norm, persistent=False)
 
-    @lru_cache(maxsize=128)
     def _calculate_symbolic_funcs(self) -> list:
         """Generate spherical basis functions based on Rayleigh formula.
 
@@ -140,9 +138,13 @@ def _calculate_symbolic_funcs(self) -> list:
         funcs = [sympy.expand_func(sympy.functions.special.bessel.jn(i, x)) for i in range(self.max_l + 1)]
         return [sympy.lambdify(x, func, torch) for func in funcs]
 
-    @lru_cache(maxsize=128)
     def _calculate_smooth_symbolic_funcs(self) -> list:
-        return _get_lambda_func(max_n=self.max_n, cutoff=self.cutoff)
+        # ``self.cutoff`` is a buffer whose dtype follows the ambient default;
+        # feeding the tensor into the symbolic builder makes ``cutoff**1.5``
+        # a float32 op under the default dtype, rounding the basis prefactor
+        # once and scaling every generated coefficient by 1 + 2.6e-8. The
+        # non-smooth branch already guards this with ``float(cutoff)``.
+        return _get_lambda_func(max_n=self.max_n, cutoff=float(self.cutoff))
 
     def forward(self, r: torch.Tensor) -> torch.Tensor:
         """Compute the spherical Bessel function values.

tests/layers/test_basis.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import math
+
 import numpy as np
 import pytest
 import torch
@@ -15,6 +17,7 @@
     spherical_bessel_smooth,
 )
 from matgl.layers._three_body import combine_sbf_shf
+from matgl.utils.maths import _get_lambda_func
 
 
 def test_gaussian():
@@ -147,3 +150,44 @@ def test_fourier_expansion(learnable):
         assert fe.frequencies.requires_grad
     else:
         assert not fe.frequencies.requires_grad
+
+
+@pytest.fixture
+def restore_dtype():
+    old = torch.get_default_dtype()
+    yield
+    torch.set_default_dtype(old)
+
+
+def test_smooth_sbf_matches_closed_form(restore_dtype):
+    """The n=0 basis function is sqrt(2)(2 sin(pi r/5) + sin(2 pi r/5))/(5 r)."""
+    torch.set_default_dtype(torch.float32)
+    sbf = SphericalBesselFunction(max_l=3, max_n=3, cutoff=5.0, smooth=True)
+    r = torch.linspace(0.3, 5.0, 257, dtype=torch.float64)
+    want = math.sqrt(2.0) * (2 * torch.sin(math.pi * r / 5) + torch.sin(2 * math.pi * r / 5)) / (5 * r)
+    got = sbf(r)[:, 0]
+    # `want` has an exact zero at r == cutoff, so the tolerance is normalised to
+    # the amplitude of the basis function rather than applied pointwise.
+    scale = want.abs().max()
+    assert (got - want).abs().max() <= 1e-12 * scale
+
+
+def test_smooth_sbf_independent_of_default_dtype(restore_dtype):
+    """The basis must not change with the ambient default dtype."""
+    r = torch.linspace(0.3, 5.0, 257, dtype=torch.float64)
+    out = {}
+    for dtype in (torch.float32, torch.float64):
+        torch.set_default_dtype(dtype)
+        out[dtype] = SphericalBesselFunction(3, 3, 5.0, smooth=True)(r)
+    assert torch.equal(out[torch.float32], out[torch.float64])
+
+
+def test_smooth_sbf_lambda_cache_is_reused(restore_dtype):
+    """Identical modules must share the cached symbolic functions."""
+    torch.set_default_dtype(torch.float32)
+    _get_lambda_func.cache_clear()
+    for _ in range(4):
+        SphericalBesselFunction(3, 3, 5.0, smooth=True)
+    info = _get_lambda_func.cache_info()
+    assert info.currsize == 1, f"cache did not coalesce: {info}"
+    assert info.hits == 3, f"cache never hit: {info}"