From d02bb4feaa46a6c44d631e06ff3dfd4cd9d0cab9 Mon Sep 17 00:00:00 2001 From: Johan Mathe Date: Wed, 14 Oct 2026 08:37:11 -0700 Subject: [PATCH 1/2] Fix J_0/J_1 accuracy: Miller recurrence for every Bessel order torch.special.bessel_j0/j1 are only ~5e-7 accurate in float64 on x in [2.8, 14.5] (torch 2.11, arm64). bessel_jn used them as base cases for n = 0, 1 and as seeds of the forward recurrence, so every order, every root and the whole disk-harmonic basis inherited the error (root j_{0,5} was off by 1e-6; J_n at roots was ~1e-7, not 0). Evaluate all orders with Miller's backward recurrence, started above max(n, ceil(max|x|)) so it is stable in both the x < n and the oscillatory regime. A single pass now returns several orders (_bessel_jn_multi), which the Newton root finders use to get J_{n-1}, J_n, J_{n+1} in one evaluation. Also fix a convergence pathology in the bracketed Newton solvers: once Newton had converged (3 iterations) the iterate sat on a bracket endpoint, the next step fell outside by rounding, and the solver fell back to ~45 bisection iterations. A converged Newton step now pins the bracket. Net effect: J_n accurate to 1e-15 on [0, 160] for n <= 60, roots to 1e-14, and compute_all_bessel_roots is 3x faster. Tests: mpmath-referenced J_n values (atol 1e-14) for n in {0, 1, 2, 5, 20} on x up to 150, roots to 1e-13, J_0(0) = 1, root residuals < 1e-13; the two tests that asserted equality with torch.special are replaced. --- src/bispectrum/_bessel.py | 154 +++++++++++++++++++++----------------- tests/test_so2_on_disk.py | 100 ++++++++++++++++++++++--- 2 files changed, 173 insertions(+), 81 deletions(-) diff --git a/src/bispectrum/_bessel.py b/src/bispectrum/_bessel.py index 0526bfa..82f4dfd 100644 --- a/src/bispectrum/_bessel.py +++ b/src/bispectrum/_bessel.py @@ -1,7 +1,9 @@ """Internal Bessel function utilities. -Pure torch implementation — no scipy. Used by SO2onDisk for disk harmonic -computations. Not part of the public API. +Pure torch implementation — no scipy, and no ``torch.special.bessel_j0/j1`` +(which are only ~5e-7 accurate in float64 on parts of the real line). All +orders are evaluated with Miller's backward recurrence. Used by SO2onDisk +for disk harmonic computations. Not part of the public API. Provides: bessel_jn — J_n(x) for integer order n >= 0 (torch tensor) @@ -15,13 +17,16 @@ def bessel_jn(n: int, x: torch.Tensor) -> torch.Tensor: - """Compute J_n(x) for integer order n >= 0. + """Compute J_n(x) for integer order n >= 0 to double precision. - Uses the forward recurrence J_{k+1}(x) = (2k/x)*J_k(x) - J_{k-1}(x) - where it is stable (x >= n) and Miller's backward recurrence where - the forward direction diverges (x < n). The forward recurrence - amplifies the Y_n admixture exponentially for x < n, which matters - for disk harmonics: they evaluate J_n(lambda * r) with r near 0. + Every order, including n = 0 and n = 1, is evaluated with Miller's + backward recurrence (see :func:`_bessel_jn_multi`). The recurrence is + started above both ``n`` and ``max|x|`` so it is stable in the + ``x < n`` regime (where the forward recurrence diverges) and in the + oscillatory regime alike. ``torch.special.bessel_j0/j1`` are + deliberately not used: in float64 they are only ~5e-7 accurate on + x in [2.8, 14.5], and any recurrence seeded with them inherits that + error. Args: n: Non-negative integer order. @@ -32,89 +37,91 @@ def bessel_jn(n: int, x: torch.Tensor) -> torch.Tensor: """ if n < 0: raise ValueError(f'Order n must be >= 0, got {n}') + return _bessel_jn_multi((n,), x)[0] - if n == 0: - return torch.special.bessel_j0(x) - if n == 1: - return torch.special.bessel_j1(x) - - needs_backward = x.abs() < n - if not bool(needs_backward.any()): - return _bessel_jn_forward(n, x) - if bool(needs_backward.all()): - return _bessel_jn_miller(n, x) - return torch.where(needs_backward, _bessel_jn_miller(n, x), _bessel_jn_forward(n, x)) - - -def _bessel_jn_forward(n: int, x: torch.Tensor) -> torch.Tensor: - """Forward recurrence for J_n(x); stable only for |x| >= n.""" - j_prev = torch.special.bessel_j0(x) - j_curr = torch.special.bessel_j1(x) - for k in range(1, n): - safe_x = torch.where(x == 0, torch.ones_like(x), x) - j_next = (2.0 * k / safe_x) * j_curr - j_prev - j_next = torch.where(x == 0, torch.zeros_like(j_next), j_next) - j_prev = j_curr - j_curr = j_next - - return j_curr +def _miller_start_order(n_max: int, x_max: float) -> int: + """Even start order for Miller's recurrence covering orders <= n_max and |x| <= x_max.""" + base = max(n_max, int(math.ceil(x_max))) + m_start = base + int(math.sqrt(60.0 * (base + 1))) + 20 + if m_start % 2 == 1: + m_start += 1 + return m_start -def _bessel_jn_miller(n: int, x: torch.Tensor) -> torch.Tensor: - """Miller's backward recurrence for J_n(x); stable for |x| < n. +def _bessel_jn_multi(orders: tuple[int, ...], x: torch.Tensor) -> list[torch.Tensor]: + """Evaluate J_n(x) for several non-negative orders in one Miller pass. - Recurs downward from a start order well above n with an arbitrary seed, then normalizes with - the identity J_0(x) + 2*sum_k J_{2k}(x) = 1. + Recurs ``J_{k-1} = (2k/x) J_k - J_{k+1}`` downward from an arbitrary seed at a start order well + above ``max(orders)`` and ``max|x|``, then normalizes with the identity ``J_0(x) + 2 * + sum_{k>=1} J_{2k}(x) = 1``. The downward direction is minimal-solution stable for every x + below the start order, which gives ~1e-16 relative accuracy throughout. Computation is done in + float64 and cast back to ``x.dtype``. """ - ax = x.abs() - safe_x = torch.where(ax == 0, torch.ones_like(ax), ax) - - m_start = n + int(math.sqrt(60.0 * (n + 1))) + 20 - if m_start % 2 == 1: - m_start += 1 - + if not orders: + return [] + n_max = max(orders) + in_dtype = x.dtype + xw = x.to(torch.float64) + ax = xw.abs() + # |x| below 1e-30 is treated as 0 (J_0 = 1, J_n = 0): the recurrence + # ratio 2k/x would overflow between two rescales otherwise. + is_zero = ax < 1e-30 + inv_x = torch.where(is_zero, torch.ones_like(ax), ax).reciprocal() + x_max = float(ax.max()) if ax.numel() > 0 else 0.0 + m_start = _miller_start_order(n_max, x_max) + + wanted = set(orders) j_up = torch.zeros_like(ax) # J_{k+1} j_k = torch.full_like(ax, 1e-30) # J_k, arbitrary seed normalized away norm_even = torch.zeros_like(ax) # 2 * sum of J_{2k}, k >= 1 - result = torch.zeros_like(ax) + captured: dict[int, torch.Tensor] = {} for k in range(m_start, 0, -1): - j_dn = (2.0 * k / safe_x) * j_k - j_up # J_{k-1} + j_dn = inv_x * j_k + j_dn.mul_(2.0 * k).sub_(j_up) # J_{k-1} = (2k/x) J_k - J_{k+1} j_up = j_k j_k = j_dn - if k - 1 == n: - result = j_k.clone() - if (k - 1) > 0 and (k - 1) % 2 == 0: - norm_even = norm_even + 2.0 * j_k - big = j_k.abs() > 1e250 - if bool(big.any()): - scale = torch.where(big, torch.full_like(j_k, 1e-250), torch.ones_like(j_k)) + km1 = k - 1 + if km1 in wanted: + captured[km1] = j_k + if km1 > 0 and km1 % 2 == 0: + norm_even.add_(j_k, alpha=2.0) + if k % 8 == 0: + # Periodic renormalization keeps the dominant solution below + # overflow; the common factor cancels in the final division. + scale = torch.maximum(j_k.abs(), j_up.abs()).clamp_min(1.0).reciprocal() j_k = j_k * scale j_up = j_up * scale norm_even = norm_even * scale - result = result * scale - - out = result / (j_k + norm_even) # j_k is now J_0 - out = torch.where(ax == 0, torch.zeros_like(out), out) - if n % 2 == 1: - out = torch.where(x < 0, -out, out) + captured = {m: v * scale for m, v in captured.items()} + + norm = j_k + norm_even # j_k is now J_0 + + out: list[torch.Tensor] = [] + for n in orders: + val = captured[n] / norm + at_zero = torch.ones_like(val) if n == 0 else torch.zeros_like(val) + val = torch.where(is_zero, at_zero, val) + if n % 2 == 1: + val = torch.where(xw < 0, -val, val) + out.append(val.to(in_dtype)) return out def _jn_scalar(n: int, x: float) -> float: """Scalar evaluation of J_n(x); delegates to the stable tensor path.""" - if x == 0 and n >= 1: - return 0.0 result: float = bessel_jn(n, torch.tensor(x, dtype=torch.float64)).item() return result def _djn_scalar(n: int, x: float) -> float: """Scalar J_n'(x) = (J_{n-1}(x) - J_{n+1}(x)) / 2.""" + xt = torch.tensor(x, dtype=torch.float64) if n == 0: - return -_jn_scalar(1, x) - return (_jn_scalar(n - 1, x) - _jn_scalar(n + 1, x)) / 2.0 + return -float(_bessel_jn_multi((1,), xt)[0].item()) + j_lo, j_hi = _bessel_jn_multi((n - 1, n + 1), xt) + return float((j_lo - j_hi).item()) / 2.0 def _bisect_newton(n: int, a: float, b: float) -> float: @@ -139,6 +146,8 @@ def _bisect_newton(n: int, a: float, b: float) -> float: if abs(dfx) > 1e-30: x_new = x - fx / dfx + if abs(x_new - x) <= 1e-15 * max(abs(x), 1.0): + return min(max(x_new, a), b) else: x_new = x @@ -169,8 +178,8 @@ def _mcmahon_zeros_j0(num_zeros: int) -> list[float]: beta = math.pi * (s - 0.25) z = beta - 1.0 / (8.0 * beta) for _ in range(10): - fz = bessel_jn(0, z) - dfz = -bessel_jn(1, z) + fz, j1 = _bessel_jn_multi((0, 1), z) + dfz = -j1 safe_dfz = torch.where(dfz.abs() < 1e-30, torch.ones_like(dfz), dfz) dz = fz / safe_dfz dz = torch.where(dfz.abs() < 1e-30, torch.zeros_like(dz), dz) @@ -207,23 +216,30 @@ def _bisect_newton_batch(n: int, a: torch.Tensor, b: torch.Tensor) -> torch.Tens x = mid0.clone() for _ in range(80): - fx = bessel_jn(n, x) if n == 0: - dfx = -bessel_jn(1, x) + fx, j1 = _bessel_jn_multi((0, 1), x) + dfx = -j1 else: - dfx = (bessel_jn(n - 1, x) - bessel_jn(n + 1, x)) / 2.0 + j_lo, fx, j_hi = _bessel_jn_multi((n - 1, n, n + 1), x) + dfx = (j_lo - j_hi) / 2.0 newton_ok = dfx.abs() > 1e-30 safe_dfx = dfx.abs().clamp_min(1e-30).copysign(dfx) x_newton = torch.where(newton_ok, x - fx / safe_dfx, x) + # A Newton step below machine precision means x is converged to a + # simple root. Pin the bracket there: otherwise the iterate sits on + # a bracket endpoint, the next Newton step falls outside by rounding, + # and the solver degrades to ~45 pure bisection iterations. + step_tiny = newton_ok & ((x_newton - x).abs() <= 1e-15 * x.abs().clamp(min=1.0)) in_bracket = (a < x_newton) & (x_newton < b) x = torch.where(in_bracket & newton_ok, x_newton, (a + b) / 2.0) + x = torch.where(step_tiny, torch.minimum(torch.maximum(x_newton, a), b), x) fx = bessel_jn(n, x) # If fx is exactly 0 the sign test below is ill-defined and the # bisection would discard the root; pin the bracket at x instead. - hit = fx == 0 + hit = (fx == 0) | step_tiny a = torch.where(hit, x, a) b = torch.where(hit, x, b) diff --git a/tests/test_so2_on_disk.py b/tests/test_so2_on_disk.py index 5904167..e6267e7 100644 --- a/tests/test_so2_on_disk.py +++ b/tests/test_so2_on_disk.py @@ -15,34 +15,110 @@ compute_all_bessel_roots, ) +# mpmath (40 digits) reference values for J_n(x). Covers x = 3.5, 8, 12 +# where torch.special.bessel_j0/j1 are only ~5e-7 accurate in float64. +_JN_REFERENCE: list[tuple[int, float, float]] = [ + (0, 0.5, 0.9384698072408129), + (0, 3.5, -0.38012773998726338), + (0, 8.0, 0.17165080713755391), + (0, 12.0, 0.047689310796833537), + (0, 40.0, 0.0073668905842372896), + (0, 150.0, -0.00077409037539429125), + (1, 0.5, 0.24226845767487389), + (1, 3.5, 0.13737752736232719), + (1, 8.0, 0.23463634685391462), + (1, 12.0, -0.22344710449062761), + (1, 40.0, 0.126038318037585), + (1, 150.0, -0.06514516365772736), + (2, 0.5, 0.030604023458682641), + (2, 3.5, 0.45862918419430748), + (2, 8.0, -0.11299172042407525), + (2, 12.0, -0.084930494878604805), + (2, 40.0, -0.0010649746823580396), + (2, 150.0, -9.4511806708740224e-5), + (5, 0.5, 8.0536272413574741e-6), + (5, 3.5, 0.080441986647991782), + (5, 8.0, 0.18577477219056331), + (5, 12.0, -0.073470963101658581), + (5, 40.0, 0.12257346597711779), + (5, 150.0, -0.064998631740725847), + (20, 0.5, 3.7272019617047145e-31), + (20, 3.5, 2.5768553102807591e-14), + (20, 8.0, 2.0805829639717028e-7), + (20, 12.0, 0.00025121327024539953), + (20, 40.0, 0.1277939335508489), + (20, 150.0, 0.063447240953861973), +] + +# mpmath besseljzero(n, k) references. +_ROOT_REFERENCE: list[tuple[int, int, float]] = [ + (0, 1, 2.4048255576957728), + (0, 5, 14.930917708487786), + (0, 20, 62.04846919022717), + (1, 1, 3.8317059702075123), + (1, 7, 22.760084380592772), + (2, 3, 11.619841172149059), + (5, 4, 18.980133875179921), + (10, 10, 45.231574103535045), + (20, 8, 51.860019928074574), + (40, 3, 56.658313455430063), +] + class TestBessel: - def test_j0_matches_torch(self): - x = torch.linspace(0.1, 20, 200, dtype=torch.float64) - torch.testing.assert_close(bessel_jn(0, x), torch.special.bessel_j0(x)) + @pytest.mark.parametrize('n, x, expected', _JN_REFERENCE) + def test_jn_matches_mpmath(self, n: int, x: float, expected: float): + """J_n must be accurate to double precision for every order, including n = 0, 1. + + ``torch.special.bessel_j0/j1`` are only ~5e-7 accurate on x in [2.8, 14.5] + (torch 2.11 arm64), so they cannot be used as base cases. + """ + got = bessel_jn(n, torch.tensor([x], dtype=torch.float64)).item() + assert got == pytest.approx(expected, abs=1e-14, rel=1e-13) + + def test_jn_matches_mpmath_batched(self): + """Mixed orders/arguments in one tensor evaluate identically to scalar calls.""" + for n in (0, 1, 2, 5, 20): + rows = [(x, v) for nn, x, v in _JN_REFERENCE if nn == n] + xs = torch.tensor([x for x, _ in rows], dtype=torch.float64) + expected = torch.tensor([v for _, v in rows], dtype=torch.float64) + torch.testing.assert_close(bessel_jn(n, xs), expected, atol=1e-14, rtol=1e-13) - def test_j1_matches_torch(self): - x = torch.linspace(0.1, 20, 200, dtype=torch.float64) - torch.testing.assert_close(bessel_jn(1, x), torch.special.bessel_j1(x)) + def test_j0_at_zero_is_one(self): + x = torch.tensor([0.0, 0.0], dtype=torch.float64) + torch.testing.assert_close(bessel_jn(0, x), torch.ones_like(x), atol=0, rtol=0) def test_jn_at_zero(self): assert bessel_jn(0, torch.tensor([0.0])).item() == pytest.approx(1.0) for n in [1, 2, 5, 10]: assert bessel_jn(n, torch.tensor([0.0])).item() == pytest.approx(0.0) + def test_jn_float32_input(self): + x = torch.tensor([0.5, 3.5, 12.0], dtype=torch.float32) + got = bessel_jn(1, x) + assert got.dtype == torch.float32 + expected = torch.tensor( + [v for n, xx, v in _JN_REFERENCE if n == 1 and xx in (0.5, 3.5, 12.0)], + dtype=torch.float32, + ) + torch.testing.assert_close(got, expected, atol=1e-6, rtol=1e-5) + def test_negative_order_raises(self): with pytest.raises(ValueError, match='Order n must be >= 0'): bessel_jn(-1, torch.tensor([1.0])) + @pytest.mark.parametrize('n, k, expected_root', _ROOT_REFERENCE) + def test_roots_match_mpmath(self, n: int, k: int, expected_root: float): + roots = bessel_jn_zeros(n, k) + assert roots[k - 1].item() == pytest.approx(expected_root, abs=1e-13) + @pytest.mark.parametrize( 'n, expected_first_root', - [(0, 2.4048255577), (1, 3.8317059702), (2, 5.1356223018)], + [(0, 2.4048255576957728), (1, 3.8317059702075123), (2, 5.1356223018406826)], ) def test_known_roots(self, n: int, expected_first_root: float): roots = bessel_jn_zeros(n, 3) - # Root accuracy is limited (~1e-6) by torch.special.bessel_j0/j1, - # which have absolute error ~3e-7 near their zeros. - assert roots[0].item() == pytest.approx(expected_first_root, abs=1e-5) + assert roots[0].item() == pytest.approx(expected_first_root, abs=1e-13) @pytest.mark.parametrize( 'n, x, expected', @@ -91,7 +167,7 @@ def test_jn_odd_symmetry_negative_x(self): ) def test_high_order_roots_match_reference(self, n: int, k: int, expected_root: float): roots = bessel_jn_zeros(n, k) - assert roots[k - 1].item() == pytest.approx(expected_root, abs=1e-5) + assert roots[k - 1].item() == pytest.approx(expected_root, abs=1e-12) def test_roots_interlace(self): """j_{n-1,k} < j_{n,k} < j_{n-1,k+1} for all consecutive orders. @@ -109,7 +185,7 @@ def test_root_residuals_high_orders(self): all_roots = compute_all_bessel_roots(20, 10) for n, roots in all_roots.items(): vals = bessel_jn(n, torch.tensor(roots, dtype=torch.float64)) - assert vals.abs().max().item() < 1e-6, f'Large residual at order {n}' + assert vals.abs().max().item() < 1e-13, f'Large residual at order {n}' @pytest.mark.parametrize('n', [0, 1, 2, 5, 10]) def test_jn_at_zeros_is_zero(self, n: int): From f936db306ac8dd230c2158d39e076feb4c38e5eb Mon Sep 17 00:00:00 2001 From: Johan Mathe Date: Wed, 14 Oct 2026 08:41:30 -0700 Subject: [PATCH 2/2] SO2onDisk: default to the Nyquist band-limit so the DHT is well conditioned The default truncation followed fle_2d: ne = floor(pi L^2 / 4) basis functions, one per disk pixel. That makes the real basis matrix Phi nearly square and admits radial frequencies ~25% above the grid Nyquist rate, so cond(Phi) ~ 1e16 at every L. pinv(rcond=1e-10) then silently discarded directions: DHT roundtrips, inversion and even exact 90-degree pixel rotations were only accurate to O(1), which the tests tolerated (rel_err < 1.5, < 0.5, 85% of coefficients within 20%). Default bandlimit is now pi * L / 2 (new public helper default_bandlimit). cond(Phi) is ~2 for L = 8..64; DHT roundtrip, coefficient inversion and bispectrum roundtrip are exact to 1e-14, and rot90 invariance to 1e-14. The ne-truncation branch and its conjugate-pair patch are removed; an explicit bandlimit still works. Coefficient counts change (L=16: 144 DH coefficients, 76 selective bispectrum entries instead of 105). README table updated. Tests: cond(Phi) < 10 for L in {16, 32, 64}; exact rank; roundtrip tolerances tightened to 1e-9..1e-12; new rot90 test; the arbitrary-angle spatial test now uses band-limited content and bicubic interpolation (its 0.1 tolerance is interpolation-bound, measured ~2e-2). --- README.md | 2 +- src/bispectrum/so2_on_disk.py | 57 ++++++++------- tests/test_so2_on_disk.py | 126 ++++++++++++++++++++-------------- 3 files changed, 102 insertions(+), 83 deletions(-) diff --git a/README.md b/README.md index f2efac8..2bd7af2 100644 --- a/README.md +++ b/README.md @@ -43,7 +43,7 @@ Median wall-clock on a single NVIDIA H100 80 GB GPU (batch=16, `torch.utils.benc | `CnonCn` | C_128 | 128 | 128 | 8,256 | 0.14 | 8.53 | | `TorusOnTorus` | C_32 × C_32 | 1,024 | 1,024 | 524,800 | 0.07 | 0.31 | | `DnonDn` | D_32 | 64 | 245 | — | 0.57 | — | -| `SO2onDisk` | SO(2) | L=16 | 105 | — | 0.22 | — | +| `SO2onDisk` | SO(2) | L=16 | 76 | — | 0.22 | — | | `SO3onS2` | SO(3) | L=16 | 430 | — | 0.48 | — | | `OctaonOcta` | O | 24 | 172 | — | 0.68 | — | diff --git a/src/bispectrum/so2_on_disk.py b/src/bispectrum/so2_on_disk.py index d8259ce..69fc361 100644 --- a/src/bispectrum/so2_on_disk.py +++ b/src/bispectrum/so2_on_disk.py @@ -23,19 +23,31 @@ from bispectrum._bessel import bessel_jn, compute_all_bessel_roots +def default_bandlimit(L: int) -> float: + """Default radial band-limit for an ``L x L`` image: the grid Nyquist rate ``pi * L / 2``. + + The disk harmonic ``J_n(lambda r) e^{i n theta}`` oscillates radially with wavenumber + ``lambda`` on the unit disk, which is sampled with pixel pitch ``2 / L``; the Nyquist bound is + therefore ``lambda <= pi * L / 2``. Up to this frequency the basis is close to orthogonal on + the pixel set (``cond(Phi) ~ 2``), so analysis by least squares is exact to machine precision. + + The fle_2d convention of Marshall et al. (2023), ``ne = floor(pi L^2 / 4)`` basis functions + (one per disk pixel), is *not* used as a default: it admits radial frequencies ~25% beyond + Nyquist, makes ``Phi`` nearly square with ``cond ~ 1e16``, and the truncated pseudo-inverse + then silently drops directions (roundtrip errors of O(1)). Pass ``bandlimit`` explicitly to + reproduce that count. + """ + return math.pi * L / 2 + + def _build_truncated_indices( L: int, bandlimit: float | None, ) -> tuple[list[tuple[int, int]], dict[int, int], int]: """Build the truncated DH coefficient index set. - Following the convention of Marshall et al. (2023, fle_2d), the number - of basis functions equals ne = floor(L^2 * pi / 4), approximately the - number of pixels inside the unit disk. Coefficients are sorted by - ascending Bessel root and the first ne are kept. - - If bandlimit is provided explicitly, it is used as a threshold on the - Bessel root instead. + Keeps every ``(n, k)`` whose Bessel root ``lambda_{n,k}`` is at most ``bandlimit`` + (default: :func:`default_bandlimit`, the grid Nyquist rate). Returns: indices: list of (n, k) pairs sorted by ascending Bessel root, @@ -43,12 +55,7 @@ def _build_truncated_indices( K: dict mapping angular order |n| -> max radial index K_n. N_m: maximum angular frequency (max |n|). """ - if bandlimit is not None: - ne = None - lam_max = bandlimit - else: - ne = int(L * L * math.pi / 4) - lam_max = math.sqrt(math.pi * ne) * 1.5 + 10 + lam_max = default_bandlimit(L) if bandlimit is None else bandlimit n_max_search = int(lam_max) + 2 k_max_search = int(lam_max / math.pi) + 3 @@ -67,18 +74,6 @@ def _build_truncated_indices( nk_lambda.sort(key=lambda t: t[0]) - if ne is not None and len(nk_lambda) >= ne: - nk_lambda = nk_lambda[:ne] - # Ensure complete conjugate pairs: if the last entry has n != 0 - # and its partner (-n, k) was cut off, remove it. - if nk_lambda: - _, last_n, last_k = nk_lambda[-1] - if last_n != 0: - partner = (-last_n, last_k) - partner_present = any(n == partner[0] and k == partner[1] for _, n, k in nk_lambda) - if not partner_present: - nk_lambda.pop() - indices = [(n, k) for _, n, k in nk_lambda] K: dict[int, int] = {} @@ -104,8 +99,10 @@ class SO2onDisk(nn.Module): Args: L: Image side length (input shape is (batch, L, L)). - bandlimit: Maximum Bessel root frequency to include. - Defaults to the fle_2d convention ne = floor(L^2 * pi / 4). + bandlimit: Maximum Bessel root frequency ``lambda_{n,k}`` to include. + Defaults to the grid Nyquist rate ``pi * L / 2`` (see + :func:`default_bandlimit`), which keeps the analysis operator + well conditioned (``cond(Phi) ~ 2``). selective: If True (default), compute O(m) selective bispectrum. If False, raises NotImplementedError. """ @@ -225,9 +222,9 @@ def __init__( self.register_buffer('_bessel_roots', bessel_roots_nonneg) self.register_buffer('_phi', phi) - # Analysis: x_real = pinv(Φ) @ f_disk (real least-squares) - # Regularize by zeroing singular values below 1e-10 * max_sv - # to handle basis functions that are unresolvable on the grid. + # Analysis: x_real = pinv(Φ) @ f_disk (real least-squares). With the + # default Nyquist truncation Φ has cond ~ 2 and the rcond cutoff is + # inactive; it only guards explicit super-Nyquist band-limits. phi_pinv = torch.linalg.pinv(phi, rcond=1e-10) # (d_real, p_disk) self.register_buffer('_phi_pinv', phi_pinv) diff --git a/tests/test_so2_on_disk.py b/tests/test_so2_on_disk.py index e6267e7..e4bd196 100644 --- a/tests/test_so2_on_disk.py +++ b/tests/test_so2_on_disk.py @@ -238,16 +238,31 @@ def test_bisect_newton_degenerate_bracket(self): class TestDiskBasisConditioning: - def test_basis_matrix_near_full_rank(self): - """The DHT basis must be well conditioned. + @pytest.mark.parametrize('L', [16, 32]) + def test_basis_matrix_full_rank(self, L: int): + """The DHT basis must have full column rank. With the unstable forward-only Bessel recurrence the L=16 basis already lost rank (and L=32 collapsed to rank 5 of 804). """ - bsp = SO2onDisk(L=16) + bsp = SO2onDisk(L=L) rank = torch.linalg.matrix_rank(bsp._phi).item() n_cols = bsp._phi.shape[1] - assert rank >= n_cols - 3, f'rank {rank} of {n_cols}' + assert rank == n_cols, f'rank {rank} of {n_cols}' + + @pytest.mark.parametrize('L', [16, 32, 64]) + def test_default_basis_well_conditioned(self, L: int): + """The default truncation must give a well-conditioned analysis operator. + + The former default (fle_2d count ``ne = floor(pi L^2 / 4)``, about one basis function per + disk pixel) made Phi nearly square with cond ~1e16; ``pinv(rcond=1e-10)`` then silently + discarded directions and every roundtrip was only accurate to O(1). Sampling theory bounds + the resolvable radial frequency on an L x L grid by the Nyquist rate pi * L / 2, where the + basis is close to orthogonal on the pixel set. + """ + bsp = SO2onDisk(L=L) + cond = torch.linalg.cond(bsp._phi).item() + assert cond < 10.0, f'L={L}: cond(Phi) = {cond:.3e}' class TestRealComplexConversion: @@ -313,18 +328,20 @@ def test_index_map_types(self): assert k >= 1 @pytest.mark.parametrize( - 'L, expected', - [(8, 27), (16, 105), (28, 315)], + 'L, expected_m, expected_out', + [(8, 34, 19), (16, 144, 76), (28, 462, 238)], ) - def test_coefficient_count_matches_paper(self, L: int, expected: int): - """Coefficient counts match Table 1 of Myers & Miolane 2025.""" + def test_coefficient_count_nyquist(self, L: int, expected_m: int, expected_out: int): + """Default truncation keeps every (n, k) with Bessel root lambda_{n,k} <= pi * L / 2.""" bsp = SO2onDisk(L=L) - assert bsp.output_size == expected + assert bsp._m == expected_m + assert bsp.output_size == expected_out + assert bsp._bessel_roots.max().item() <= math.pi * L / 2 def test_explicit_bandlimit(self): bsp = SO2onDisk(L=8, bandlimit=10.0) assert bsp.output_size > 0 - assert bsp.output_size < 27 + assert bsp.output_size < SO2onDisk(L=8).output_size def test_full_not_implemented(self): bsp = SO2onDisk(L=8, selective=False) @@ -408,40 +425,59 @@ def test_analytical_rotation_invariance(self): torch.testing.assert_close(beta_orig, beta_rot, atol=1e-10, rtol=1e-10) + @pytest.mark.parametrize('L', [16, 28]) + def test_rot90_invariance(self, L: int): + """A 90 degree rotation maps the pixel grid onto itself, so invariance must be exact.""" + bsp = SO2onDisk(L=L) + torch.manual_seed(123) + f = torch.randn(2, L, L, dtype=torch.float64) + beta = bsp(f) + for k in (1, 2, 3): + beta_rot = bsp(torch.rot90(f, k, dims=(1, 2))) + rel_err = (beta - beta_rot).abs().max().item() / beta.abs().max().item() + assert rel_err < 1e-12, f'L={L}, k={k}: rot90 invariance error {rel_err:.2e}' + @pytest.mark.parametrize('L', [16, 28]) def test_spatial_rotation_invariance(self, L: int): - """Verify invariance to actual image rotation (includes discretization error). + """Invariance to an arbitrary-angle image rotation of a smooth, band-limited image. - The error should be moderate for small L and decrease with L. + This is bounded by the bicubic interpolation error of ``grid_sample`` (measured ~2e-2 for + content below a quarter of the Nyquist rate), not by the transform; white-noise images are + not meaningfully rotatable by interpolation at all. """ bsp = SO2onDisk(L=L) - torch.manual_seed(123) - f = torch.randn(1, L, L, dtype=torch.float64) + torch.manual_seed(7) + a = torch.randn(1, bsp._m_nonneg, dtype=torch.complex128) + for j, (n, _k) in enumerate(bsp._nonneg_indices): + if n == 0: + a[:, j] = a[:, j].real + 0j + a[:, bsp._bessel_roots > math.pi * L / 8] = 0 + f = bsp._idht(a).reshape(1, L, L) beta = bsp(f) angle = 0.5 cos_a, sin_a = math.cos(angle), math.sin(angle) theta_mat = torch.tensor( - [[[cos_a, sin_a, 0.0], [-sin_a, cos_a, 0.0]]], dtype=torch.float32 + [[[cos_a, sin_a, 0.0], [-sin_a, cos_a, 0.0]]], dtype=torch.float64 ) grid = F.affine_grid(theta_mat, [1, 1, L, L], align_corners=False) - f_rot = ( - F.grid_sample(f.float().unsqueeze(1), grid, align_corners=False, padding_mode='zeros') - .squeeze(1) - .double() - ) + f_rot = F.grid_sample( + f.unsqueeze(1), grid, mode='bicubic', align_corners=False, padding_mode='zeros' + ).squeeze(1) beta_rot = bsp(f_rot) rel_err = (beta - beta_rot).abs().max().item() / beta.abs().max().item() - assert rel_err < 1.5, f'Rotation invariance error too large: {rel_err:.4f}' + assert rel_err < 0.1, f'Rotation invariance error too large: {rel_err:.4f}' class TestSO2onDiskDHTRoundtrip: - @pytest.mark.parametrize('L', [8, 16]) + @pytest.mark.parametrize('L', [8, 16, 32]) def test_dht_roundtrip_bandlimited(self, L: int): - """DHT roundtrip for a signal synthesized from DH coefficients. + """DHT roundtrip for a signal synthesized from DH coefficients is exact. - Error comes from unresolvable basis functions on the discrete grid. + Every retained basis function is resolvable on the grid (the default truncation is the + Nyquist rate), so analysis after synthesis must return the coefficients to machine + precision. """ bsp = SO2onDisk(L=L) torch.manual_seed(L) @@ -454,10 +490,8 @@ def test_dht_roundtrip_bandlimited(self, L: int): f = bsp._idht(a) a_rec = bsp._dht(f) - # Most coefficients should be recovered well; allow tolerance for - # the 1-2 unresolvable basis functions per grid rel_err = (a - a_rec).abs().max().item() / a.abs().max().item() - assert rel_err < 0.5, f'DHT roundtrip error too large: {rel_err}' + assert rel_err < 1e-12, f'DHT roundtrip error too large: {rel_err:.2e}' @pytest.mark.parametrize('L', [8, 16]) def test_single_harmonic_roundtrip(self, L: int): @@ -473,20 +507,21 @@ def test_single_harmonic_roundtrip(self, L: int): a_rec = bsp._dht(f) # The (0,1) coefficient should be ~1, others ~0 - assert a_rec[0, j].real.item() == pytest.approx(1.0, abs=1e-6) + assert a_rec[0, j].real.item() == pytest.approx(1.0, abs=1e-12) other_mask = torch.ones(bsp._m_nonneg, dtype=torch.bool) other_mask[j] = False - assert a_rec[0, other_mask].abs().max().item() < 1e-4 + assert a_rec[0, other_mask].abs().max().item() < 1e-12 class TestSO2onDiskInvert: - def test_invert_exact_coefficients(self): - """Inversion from exact bispectrum recovers DH magnitude for resolvable coefficients. + @pytest.mark.parametrize('L', [16, 32]) + def test_invert_exact_coefficients(self, L: int): + """Inversion from an exact bispectrum recovers every DH coefficient magnitude. - Some high-frequency coefficients may be unresolvable on the discrete grid (rank deficient - in the real basis matrix). We check that the MAJORITY of well-resolved coefficients match. + Goes through the full pipeline (coefficients -> bispectrum -> invert -> image -> DHT), so + it also exercises the conditioning of the analysis operator. """ - bsp = SO2onDisk(L=16) + bsp = SO2onDisk(L=L) torch.manual_seed(42) a = torch.randn(1, bsp._m_nonneg, dtype=torch.complex128) + 0.5 a = a.clone() @@ -507,21 +542,8 @@ def test_invert_exact_coefficients(self): f_rec = bsp.invert(beta) a_rec = bsp._dht(f_rec) - n_checked = 0 - n_good = 0 - for j, (_n, _k) in enumerate(bsp._nonneg_indices): - orig_mag = a[0, j].abs().item() - if orig_mag > 0.01: - n_checked += 1 - rec_mag = a_rec[0, j].abs().item() - rel_err = abs(orig_mag - rec_mag) / orig_mag - if rel_err < 0.2: - n_good += 1 - - fraction_good = n_good / n_checked if n_checked > 0 else 1.0 - assert fraction_good > 0.85, ( - f'Only {n_good}/{n_checked} coefficients recovered within 20% ({fraction_good:.1%})' - ) + # Inversion is up to a global rotation, so compare magnitudes. + torch.testing.assert_close(a_rec.abs(), a.abs(), atol=1e-9, rtol=0) def test_invert_direct_bispectrum_roundtrip(self): """Inversion at the coefficient level is exact (no DHT involved).""" @@ -598,9 +620,9 @@ def test_invert_not_implemented_full(self): with pytest.raises(NotImplementedError): bsp.invert(torch.randn(2, 27, dtype=torch.complex128)) - @pytest.mark.parametrize('L', [8, 16]) + @pytest.mark.parametrize('L', [8, 16, 32]) def test_invert_bispectrum_roundtrip(self, L: int): - """Bsp(invert(bsp(f))) ≈ bsp(f) within discretization tolerance.""" + """Bsp(invert(bsp(f))) == bsp(f) to machine precision, even for white-noise images.""" bsp = SO2onDisk(L=L) torch.manual_seed(L) f = torch.randn(2, L, L, dtype=torch.float64) @@ -609,4 +631,4 @@ def test_invert_bispectrum_roundtrip(self, L: int): beta_rec = bsp(f_rec) rel_err = (beta - beta_rec).abs().max().item() / beta.abs().max().item() - assert rel_err < 0.5, f'L={L}: bispectrum roundtrip rel err {rel_err:.4f}' + assert rel_err < 1e-10, f'L={L}: bispectrum roundtrip rel err {rel_err:.2e}'