diff --git a/pyproject.toml b/pyproject.toml index 94676da..87048c7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,10 @@ classifiers = [ ] dependencies = [ - "torch>=2.10", + # torch-harmonics 0.9.x ships C extensions built against the torch<2.12 + # ABI; importing it under torch>=2.12 fails with undefined symbols. + # Lift the upper bound once torch-harmonics releases compatible wheels. + "torch>=2.10,<2.12", "numpy>=1.22.4", "torch-harmonics>=0.9", ] diff --git a/src/bispectrum/_bessel.py b/src/bispectrum/_bessel.py index 3732bb1..0526bfa 100644 --- a/src/bispectrum/_bessel.py +++ b/src/bispectrum/_bessel.py @@ -15,13 +15,13 @@ def bessel_jn(n: int, x: torch.Tensor) -> torch.Tensor: - """Compute J_n(x) for integer order n >= 0 via forward recurrence. + """Compute J_n(x) for integer order n >= 0. - Uses torch.special.bessel_j0 and bessel_j1 as base cases and the - standard recurrence J_{k+1}(x) = (2k/x)*J_k(x) - J_{k-1}(x). - - Forward recurrence is stable when x >= n, which holds for our use - case (evaluating at Bessel root * r where r in [0, 1]). + 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. Args: n: Non-negative integer order. @@ -38,6 +38,16 @@ def bessel_jn(n: int, x: torch.Tensor) -> torch.Tensor: 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) @@ -51,26 +61,53 @@ def bessel_jn(n: int, x: torch.Tensor) -> torch.Tensor: return j_curr -def _jn_scalar(n: int, x: float) -> float: - """Fast scalar evaluation of J_n(x) using raw math.""" - if n == 0: - return torch.special.bessel_j0(torch.tensor(x, dtype=torch.float64)).item() - if n == 1: - return torch.special.bessel_j1(torch.tensor(x, dtype=torch.float64)).item() - - xt = torch.tensor(x, dtype=torch.float64) - j_prev = torch.special.bessel_j0(xt).item() - j_curr = torch.special.bessel_j1(xt).item() +def _bessel_jn_miller(n: int, x: torch.Tensor) -> torch.Tensor: + """Miller's backward recurrence for J_n(x); stable for |x| < n. - if x == 0: - return 0.0 + 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. + """ + 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 + + 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) + + for k in range(m_start, 0, -1): + j_dn = (2.0 * k / safe_x) * j_k - j_up # 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)) + 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) + return out - for k in range(1, n): - j_next = (2.0 * k / x) * j_curr - j_prev - j_prev = j_curr - j_curr = j_next - return j_curr +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: @@ -161,7 +198,13 @@ def _bisect_newton_batch(n: int, a: torch.Tensor, b: torch.Tensor) -> torch.Tens exact_b = fb.abs() < 1e-15 no_sign_change = fa * fb > 0 - x = (a + b) / 2.0 + # Collapse brackets whose endpoint is already a root so the bisection + # below cannot walk away from it (fa*fx < 0 is False when fa == 0). + b = torch.where(exact_a, a, b) + a = torch.where(exact_b, b, a) + + mid0 = (a + b) / 2.0 + x = mid0.clone() for _ in range(80): fx = bessel_jn(n, x) @@ -171,25 +214,32 @@ def _bisect_newton_batch(n: int, a: torch.Tensor, b: torch.Tensor) -> torch.Tens dfx = (bessel_jn(n - 1, x) - bessel_jn(n + 1, x)) / 2.0 newton_ok = dfx.abs() > 1e-30 - x_newton = torch.where(newton_ok, x - fx / dfx.clamp_min(1e-30).copysign(dfx), x) + safe_dfx = dfx.abs().clamp_min(1e-30).copysign(dfx) + x_newton = torch.where(newton_ok, x - fx / safe_dfx, x) in_bracket = (a < x_newton) & (x_newton < b) x = torch.where(in_bracket & newton_ok, x_newton, (a + b) / 2.0) 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 + a = torch.where(hit, x, a) + b = torch.where(hit, x, b) + go_left = fa * fx < 0 - b = torch.where(go_left, x, b) - fb = torch.where(go_left, fx, fb) - a = torch.where(~go_left, x, a) - fa = torch.where(~go_left, fx, fa) + keep = ~hit + b = torch.where(keep & go_left, x, b) + fb = torch.where(keep & go_left, fx, fb) + a = torch.where(keep & ~go_left, x, a) + fa = torch.where(keep & ~go_left, fx, fa) converged = (b - a) < 1e-14 * a.abs().clamp(min=1.0) if converged.all(): break result = (a + b) / 2.0 - result = torch.where(exact_a, a, result) - result = torch.where(exact_b, b, result) - result = torch.where(no_sign_change & ~exact_a & ~exact_b, (a + b) / 2.0, result) + result = torch.where(no_sign_change & ~exact_a & ~exact_b, mid0, result) return result diff --git a/src/bispectrum/_cg.py b/src/bispectrum/_cg.py index 37a3651..c82389a 100644 --- a/src/bispectrum/_cg.py +++ b/src/bispectrum/_cg.py @@ -495,7 +495,6 @@ def compute_sparse_cg_entry( _ensure_log_fact(max_n) lf = np.array(_LOG_FACT[: max_n + 1]) - 2 * l2 + 1 sqrt_2l1 = math.sqrt(2 * l_val + 1) log_tri = ( lf[l1 + l2 - l_val] + lf[l1 - l2 + l_val] + lf[-l1 + l2 + l_val] - lf[l1 + l2 + l_val + 1] diff --git a/src/bispectrum/dn_on_dn.py b/src/bispectrum/dn_on_dn.py index 5089ffd..4f39b01 100644 --- a/src/bispectrum/dn_on_dn.py +++ b/src/bispectrum/dn_on_dn.py @@ -445,6 +445,10 @@ def forward(self, f: torch.Tensor) -> torch.Tensor: """ if not self.selective: raise NotImplementedError('Full bispectrum not yet implemented for DnonDn.') + if f.is_complex(): + raise TypeError('f must be a real-valued tensor, got complex dtype.') + if f.ndim != 2 or f.shape[-1] != 2 * self.n: + raise ValueError(f'Expected shape (batch, {2 * self.n}), got {tuple(f.shape)}') n3 = self._n3 batch = f.shape[0] @@ -488,50 +492,36 @@ def forward(self, f: torch.Tensor) -> torch.Tensor: dim=-1, ) - def invert(self, beta: torch.Tensor, **kwargs: object) -> torch.Tensor: - """Recover a signal from its selective bispectrum. - - Implements Algorithm 3 (Sec. 4.1.3) from Mataigne et al., ICML 2024. - Reconstruction has O(2) indeterminacy (continuous rotations and - reflections), so the recovered signal matches the original up to - a D_n group action. + def _chain_extract( + self, + beta: torch.Tensor, + F0: torch.Tensor, + S: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Sequential 2D-coefficient recovery from beta_{rho1, rho_k}. - Args: - beta: Selective bispectrum, shape ``(batch, output_size)``. + Given seeds F(rho_0) = *F0* and F(rho_1) = *S*, extracts F(rho_{k+1}) + from beta_{rho1, rho_k} for k = 1..n3. Only forward-chain 2D blocks + (label k+1) and the rho_01 entry are written: backward blocks + (label k-1) and the 1D rho_02/rho_03 fold entries are gauge-dependent + re-extractions that must not overwrite exact seeds. Returns: - Reconstructed real signal, shape ``(batch, 2n)``. + (fhat, last_block) where *last_block* is the raw 4x4 + block-diagonal matrix of the final chain step (contains the + fold entries), or None if n3 == 0. """ - if not self.selective: - raise NotImplementedError('Inversion only implemented for selective bispectrum.') - n2d = self._n2d n3 = self._n3 batch = beta.shape[0] - device = beta.device dtype = beta.dtype + device = beta.device fhat = torch.zeros(batch, 2, 2, n2d + 1, device=device, dtype=dtype) - - # Step 1 — F(rho_0) from beta_{rho0,rho0} = F(rho0)^3 - b00 = beta[:, 0] - fhat[:, 0, 0, 0] = torch.sign(b00) * torch.abs(b00) ** (1.0 / 3.0) - - # Step 2 — F(rho_1) via eigendecomposition of beta_{rho0,rho1}/F(rho0) - b01 = beta[:, 1:5].reshape(batch, 2, 2) - F0 = fhat[:, 0, 0, 0] - M = b01 / F0[:, None, None] # = F1^T @ F1, positive semi-definite - eigvals_M, eigvecs_M = torch.linalg.eigh(M) - eigvals_M = torch.clamp(eigvals_M, min=0.0) - S = eigvecs_M @ torch.diag_embed(torch.sqrt(eigvals_M)) @ eigvecs_M.transpose(-1, -2) - + fhat[:, 0, 0, 0] = F0 fhat[:, :, :, 1] = S - # Step 3 — sequential recovery from beta_{rho1, rho_k}. - # The extraction gives Fourier coefficients that are "twisted" - # by the O(2) ambiguity in F_1. Fourier coefficient magnitudes - # (Frobenius norms) are always exact; the bispectrum of the - # recovered signal matches the original only up to O(2). + last_block: torch.Tensor | None = None offset = 5 for m in range(n3): k_prev = m + 1 @@ -539,10 +529,7 @@ def invert(self, beta: torch.Tensor, **kwargs: object) -> torch.Tensor: offset += 16 C = self._cg_matrices[m].to(dtype) - F1 = fhat[:, :, :, 1] - Fkp = fhat[:, :, :, k_prev] - - A = _batched_kron_2x2(F1, Fkp) + A = _batched_kron_2x2(fhat[:, :, :, 1], fhat[:, :, :, k_prev]) A_inv = torch.linalg.inv(A) # beta = C (oplus F^T) C^T A => oplus F = [C^T beta A^{-1} C]^T @@ -551,28 +538,155 @@ def invert(self, beta: torch.Tensor, **kwargs: object) -> torch.Tensor: temp = torch.matmul(temp, C) block_diag = temp.transpose(-1, -2) - decomp = self._decompositions[m] - for block in decomp: + for block in self._decompositions[m]: if block.block_type == '2d': - r0, r1 = block.rows k_label = block.label assert isinstance(k_label, int) + if k_label != k_prev + 1: + continue + r0, r1 = block.rows fhat[:, :, :, k_label] = block_diag[:, r0 : r1 + 1, r0 : r1 + 1] - else: - r = block.rows[0] - lbl = block.label - assert isinstance(lbl, str) - val = block_diag[:, r, r] - if lbl == 'rho0': - fhat[:, 0, 0, 0] = val - elif lbl == 'rho01': - fhat[:, 1, 0, 0] = val - elif lbl == 'rho02': - fhat[:, 0, 1, 0] = val - elif lbl == 'rho03': - fhat[:, 1, 1, 0] = val - - return self._inverse_dft(fhat) + elif block.label == 'rho01': + fhat[:, 1, 0, 0] = block_diag[:, block.rows[0], block.rows[0]] + last_block = block_diag + + return fhat, last_block + + @staticmethod + def _givens_batch(theta: torch.Tensor) -> torch.Tensor: + """Batched 2x2 rotation matrices R(theta). + + Shape (batch, 2, 2). + """ + c, s = torch.cos(theta), torch.sin(theta) + return torch.stack([torch.stack([c, -s], -1), torch.stack([s, c], -1)], -2) + + def invert(self, beta: torch.Tensor, **kwargs: object) -> torch.Tensor: + """Recover a signal from its selective bispectrum. + + Implements Algorithm 3 (Sec. 4.1.3) from Mataigne et al., ICML 2024, + with an explicit resolution of the O(2) gauge ambiguity. + + The symmetric square root S of beta_{rho0,rho1}/F(rho0) determines + F(rho_1) only up to an O(2) twist Q (F_1 = Q S). The chain + extraction of the 2D coefficients is consistent in any fixed twist + frame, but the *fold* entries — rho_02/rho_03 for even n (from + rho_1 x rho_{n/2 - 1}) and the folded rho_1 block at n = 3 — mix + components by a rotation R(h*theta) that depends on the twist angle. + This method measures that rotation from the fold block, corrects the + seed S by a finite candidate set of gauge angles (rotations and + reflection), and selects the candidate whose reconstruction + reproduces *beta* exactly (verified via :meth:`forward`). + + For odd n >= 5 there is no fold and a single chain pass is exact. + + The reconstruction satisfies ``forward(invert(beta)) == beta`` to + machine precision for generic signals. The recovered signal is + determined up to the O(2) twist; degenerate inputs (F(rho_0) = 0 or + singular F(rho_1)) are outside the algorithm's domain and may + reconstruct inaccurately. + + Args: + beta: Selective bispectrum, shape ``(batch, output_size)``. + + Returns: + Reconstructed real signal, shape ``(batch, 2n)``. + """ + if not self.selective: + raise NotImplementedError('Inversion only implemented for selective bispectrum.') + + n = self.n + batch = beta.shape[0] + dtype = beta.dtype + device = beta.device + eps = 1e-12 + + # Step 1 — F(rho_0) from beta_{rho0,rho0} = F(rho0)^3 + b00 = beta[:, 0] + F0 = torch.sign(b00) * torch.abs(b00) ** (1.0 / 3.0) + + # Step 2 — F(rho_1) up to O(2): symmetric sqrt of beta_{rho0,rho1}/F0 + b01 = beta[:, 1:5].reshape(batch, 2, 2) + M = b01 / F0[:, None, None] # = F1^T @ F1, positive semi-definite + eigvals_M, eigvecs_M = torch.linalg.eigh(M) + eigvals_M = torch.clamp(eigvals_M, min=0.0) + S = eigvecs_M @ torch.diag_embed(torch.sqrt(eigvals_M)) @ eigvecs_M.transpose(-1, -2) + + has_fold = n % 2 == 0 or n == 3 + if not has_fold: + # Odd n >= 5: chain extraction is exact in the S frame. + fhat, _ = self._chain_extract(beta, F0, S) + return self._inverse_dft(fhat) + + decomp_last = self._decompositions[self._n3 - 1] + if n % 2 == 0: + fold_1d = [ + b for b in decomp_last if b.block_type == '1d' and b.label in ('rho02', 'rho03') + ] + fold_idx = torch.tensor([b.rows[0] for b in fold_1d], device=device) + fold_labels = [b.label for b in fold_1d] + fold_scale = float(n // 2) + n_shift = 2 + else: # n == 3: the rho_1 block folds back onto itself + fold_2d = next(b for b in decomp_last if b.block_type == '2d' and b.label == 1) + fold_scale = float(n) + n_shift = 6 + + D_refl = torch.tensor([[1.0, 0.0], [0.0, -1.0]], device=device, dtype=dtype) + + best_resid = torch.full((batch,), float('inf'), device=device, dtype=dtype) + best_f = torch.zeros(batch, 2 * n, device=device, dtype=dtype) + + def consider(f_rec: torch.Tensor) -> None: + nonlocal best_resid, best_f + resid = (self.forward(f_rec) - beta).norm(dim=-1) + better = resid < best_resid + best_f = torch.where(better.unsqueeze(-1), f_rec, best_f) + best_resid = torch.minimum(best_resid, resid) + + for refl in (False, True): + S_base = (D_refl @ S) if refl else S + _, lb = self._chain_extract(beta, F0, S_base) + assert lb is not None + + # Measure the fold rotation angle in this frame. + if n % 2 == 0: + W = lb[:, fold_idx][:, :, fold_idx] + # W = R(h*theta) @ diag(v): column norms are |v|, the + # normalized columns give the rotation (up to column signs, + # covered by the shift/sign candidates below). + cnorm = W.norm(dim=1).clamp_min(eps) + O = W / cnorm.unsqueeze(1) + phi = torch.atan2(O[:, 1, 0], O[:, 0, 0]) + else: + r0, r1 = fold_2d.rows + W = lb[:, r0 : r1 + 1, r0 : r1 + 1] + Mw = W @ torch.linalg.inv(S_base) + U, _sv, Vh = torch.linalg.svd(Mw) + det = torch.linalg.det(U @ Vh) + U2 = U.clone() + U2[:, :, 1] = U2[:, :, 1] * det.unsqueeze(-1) + R = U2 @ Vh + phi = torch.atan2(R[:, 1, 0], R[:, 0, 0]) + + for mshift in range(n_shift): + for sign in (1.0, -1.0): + theta_hat = sign * (-(phi + torch.pi * mshift) / fold_scale) + S2 = self._givens_batch(theta_hat) @ S_base + fhat2, lb2 = self._chain_extract(beta, F0, S2) + assert lb2 is not None + if n % 2 == 0: + # Fold block is now diagonal: read rho_02/rho_03 off. + W2 = lb2[:, fold_idx][:, :, fold_idx] + for i, lbl in enumerate(fold_labels): + val = W2[:, i, i] + if lbl == 'rho02': + fhat2[:, 0, 1, 0] = val + else: + fhat2[:, 1, 1, 0] = val + consider(self._inverse_dft(fhat2)) + + return best_f @property def output_size(self) -> int: diff --git a/src/bispectrum/octa_on_octa.py b/src/bispectrum/octa_on_octa.py index d7c8d9b..eb4a3d3 100644 --- a/src/bispectrum/octa_on_octa.py +++ b/src/bispectrum/octa_on_octa.py @@ -502,7 +502,6 @@ def _build_fplus( Fplus = torch.zeros(batch, d, d, dtype=dtype, device=device) for irrep_k, r0, r1 in block_info: - r1 - r0 block = fhat[irrep_k] padded = torch.nn.functional.pad(block, (r0, d - r1, r0, d - r1)) Fplus = Fplus + padded @@ -735,7 +734,6 @@ def _lm_step( Uses (J^T J + mu I)^{-1} J^T r with mu adapted per sample to ensure the residual decreases. """ - f.shape[0] target_real = beta_target.real def fwd_single(x: torch.Tensor) -> torch.Tensor: diff --git a/src/bispectrum/so3_on_s2.py b/src/bispectrum/so3_on_s2.py index 2875953..1d44d2d 100644 --- a/src/bispectrum/so3_on_s2.py +++ b/src/bispectrum/so3_on_s2.py @@ -916,7 +916,11 @@ def _ensure_precomputed_on_device(self, device: torch.device, dtype: torch.dtype fl_i = self._sparse_fl_abs eid = self._sparse_entry_ids - perm = torch.randperm(bi.numel()) + # Use a private generator: this cache build must not perturb the + # user's global RNG stream, and the layout should be reproducible. + gen = torch.Generator() + gen.manual_seed(0x5CA77E12) + perm = torch.randperm(bi.numel(), generator=gen) bi_perm = bi[perm] self._sc_bi_fl1 = fl1_i[bi_perm].to(device=device) self._sc_bi_fl2 = fl2_i[bi_perm].to(device=device) diff --git a/tests/test_cn_on_cn.py b/tests/test_cn_on_cn.py index 059eeaa..cdbef24 100644 --- a/tests/test_cn_on_cn.py +++ b/tests/test_cn_on_cn.py @@ -107,6 +107,13 @@ def test_output_dtype_float32(self): out = bsp(f) assert out.is_complex() + def test_complex_input_accepted(self): + """Complex signals are part of the signal model: invert() returns + complex reconstructions that must be re-checkable via forward().""" + bsp = CnonCn(n=8) + out = bsp(torch.randn(2, 8, dtype=torch.complex128)) + assert out.dtype == torch.complex128 + def test_output_dtype_float64(self): bsp = CnonCn(n=8) f = torch.randn(2, 8, dtype=torch.float64) diff --git a/tests/test_dn_on_dn.py b/tests/test_dn_on_dn.py index a5c8e36..b64cf10 100644 --- a/tests/test_dn_on_dn.py +++ b/tests/test_dn_on_dn.py @@ -161,18 +161,61 @@ def test_dft_roundtrip_various_n(self, n: int): torch.testing.assert_close(f_rec, f, atol=1e-10, rtol=1e-10) +def _group_action(f: torch.Tensor, n: int, rot: int, refl: bool) -> torch.Tensor: + """Apply the D_n element a^rot x^refl to a batch of signals.""" + idx = [] + for m in (0, 1): + for pos in range(n): + if not refl: + lp, mp = (pos - rot) % n, m + else: + lp, mp = (rot - pos) % n, (m + 1) % 2 + idx.append(mp * n + lp) + return f[:, idx] + + class TestDnonDnInvert: + @pytest.mark.parametrize('n', [3, 4, 5, 6, 7, 8, 9, 10, 12, 16]) + def test_roundtrip_full_bispectrum(self, n: int): + """Forward(invert(beta)) == beta for the FULL selective bispectrum. + + This is the completeness contract. The previous implementation failed it for even n and n=3 + (relative errors of order 1). + """ + torch.manual_seed(n + 42) + bsp = DnonDn(n=n) + f = torch.randn(8, 2 * n, dtype=torch.float64) + beta = bsp(f) + f_rec = bsp.invert(beta) + beta_rec = bsp(f_rec) + rel = ((beta_rec - beta).norm(dim=-1) / beta.norm(dim=-1)).max().item() + assert rel < 1e-6, f'n={n}: full bispectrum roundtrip rel err {rel:.3e}' + + @pytest.mark.parametrize('n', [4, 6, 8, 10, 12]) + def test_orbit_recovery_even_n(self, n: int): + """For even n the reconstruction lies in the D_n orbit of the input.""" + torch.manual_seed(n + 7) + bsp = DnonDn(n=n) + f = torch.randn(4, 2 * n, dtype=torch.float64) + f_rec = bsp.invert(bsp(f)) + + best = None + for refl in (False, True): + for rot in range(n): + d = (_group_action(f_rec, n, rot, refl) - f).norm(dim=-1) + best = d if best is None else torch.minimum(best, d) + rel = (best / f.norm(dim=-1)).max().item() + assert rel < 1e-6, f'n={n}: orbit distance {rel:.3e}' + @pytest.mark.parametrize('n', [3, 4, 5, 7, 8]) def test_roundtrip_bispectrum(self, n: int): - """β_{00} and β_{01} must roundtrip exactly; β_{1k} up to O(2).""" + """β_{00} and β_{01} blocks roundtrip (subset of the full check).""" torch.manual_seed(n + 42) bsp = DnonDn(n=n) f = torch.randn(4, 2 * n, dtype=torch.float64) beta = bsp(f) f_rec = bsp.invert(beta) beta_rec = bsp(f_rec) - # β_{ρ0,ρ0} and β_{ρ0,ρ1} are determined by F_0 and F_1^T F_1 - # which are recovered exactly (the O(2) ambiguity cancels). torch.testing.assert_close(beta[:, :5], beta_rec[:, :5], atol=ATOL, rtol=RTOL) @pytest.mark.parametrize('n', [3, 4, 5, 7, 8]) @@ -210,3 +253,17 @@ def test_invert_not_implemented_full(self): bsp = DnonDn(n=8, selective=False) with pytest.raises(NotImplementedError): bsp.invert(torch.randn(2, 4)) + + +class TestDnonDnValidation: + def test_complex_input_raises(self): + bsp = DnonDn(n=4) + with pytest.raises(TypeError, match='real-valued'): + bsp(torch.randn(2, 8, dtype=torch.complex128)) + + def test_wrong_shape_raises(self): + bsp = DnonDn(n=4) + with pytest.raises(ValueError, match='Expected shape'): + bsp(torch.randn(2, 7)) + with pytest.raises(ValueError, match='Expected shape'): + bsp(torch.randn(8)) diff --git a/tests/test_so2_on_disk.py b/tests/test_so2_on_disk.py index b18610b..5904167 100644 --- a/tests/test_so2_on_disk.py +++ b/tests/test_so2_on_disk.py @@ -40,11 +40,77 @@ def test_negative_order_raises(self): ) def test_known_roots(self, n: int, expected_first_root: float): roots = bessel_jn_zeros(n, 3) - # Forward recurrence has cancellation near J_n roots for n >= 2, - # giving ~1e-6 precision. Roots are self-consistent (bessel_jn - # at the root is ~0), which is what matters for the DHT. + # 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) + @pytest.mark.parametrize( + 'n, x, expected', + [ + # Reference values from mpmath (50 digits). The old forward-only + # recurrence returned ~1e+11 for J_20(0.5). + (10, 0.5, 2.6131773608228023e-13), + (20, 0.5, 3.7272019617047145e-31), + (20, 10.0, 1.1513369247813403e-05), + (30, 1.0, 3.4828697942514834e-42), + (40, 2.0, 1.1960774581136798e-48), + (60, 30.0, 9.8075576431286213e-14), + ], + ) + def test_jn_small_argument_reference_values(self, n: int, x: float, expected: float): + """Backward (Miller) recurrence is accurate in the x < n regime.""" + got = bessel_jn(n, torch.tensor([x], dtype=torch.float64)).item() + assert got == pytest.approx(expected, rel=1e-12) + + def test_jn_continuous_across_regime_boundary(self): + """Forward/backward switch at |x| = n must not create a jump.""" + for n in (5, 12, 25): + x = torch.linspace(n - 0.5, n + 0.5, 101, dtype=torch.float64) + vals = bessel_jn(n, x) + jumps = (vals[1:] - vals[:-1]).abs() + assert jumps.max().item() < 1e-2 + assert torch.isfinite(vals).all() + + def test_jn_odd_symmetry_negative_x(self): + """J_n(-x) = (-1)^n J_n(x).""" + x = torch.linspace(0.1, 10.0, 50, dtype=torch.float64) + for n in (3, 8): + sign = -1.0 if n % 2 == 1 else 1.0 + torch.testing.assert_close(bessel_jn(n, -x), sign * bessel_jn(n, x)) + + @pytest.mark.parametrize( + 'n, k, expected_root', + [ + # mpmath besseljzero references, including cases where the old + # batch bisection walked away from an already-converged root + # (e.g. j_{6,5} came back as 25.43 instead of 23.586). + (6, 5, 23.586084435581391), + (10, 9, 42.004190236671805), + (20, 8, 51.860019928074567), + ], + ) + 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) + + def test_roots_interlace(self): + """j_{n-1,k} < j_{n,k} < j_{n-1,k+1} for all consecutive orders. + + The old solver returned roots of J_{n-1} in J_n slots, which violates strict interlacing. + """ + all_roots = compute_all_bessel_roots(20, 10) + for n in range(1, 21): + prev = all_roots[n - 1] + curr = all_roots[n] + for k in range(min(len(curr), len(prev) - 1)): + assert prev[k] < curr[k] < prev[k + 1], f'Interlacing violated at n={n}, k={k + 1}' + + 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}' + @pytest.mark.parametrize('n', [0, 1, 2, 5, 10]) def test_jn_at_zeros_is_zero(self, n: int): roots = bessel_jn_zeros(n, 5) @@ -95,6 +161,19 @@ def test_bisect_newton_degenerate_bracket(self): assert root == pytest.approx(0.3, abs=1e-6) +class TestDiskBasisConditioning: + def test_basis_matrix_near_full_rank(self): + """The DHT basis must be well conditioned. + + 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) + 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}' + + class TestRealComplexConversion: """Direct tests for _real_to_complex / _complex_to_real roundtrip.""" diff --git a/tests/test_so3_on_s2.py b/tests/test_so3_on_s2.py index 907301d..84b7589 100644 --- a/tests/test_so3_on_s2.py +++ b/tests/test_so3_on_s2.py @@ -125,6 +125,20 @@ def test_large_lmax_computes_analytically(self): bsp = SO3onS2(lmax=6, nlat=32, nlon=64) assert bsp.output_size > 0 + def test_forward_does_not_perturb_global_rng(self): + """Building the internal device cache must not consume global RNG state.""" + bsp = SO3onS2(lmax=3, nlat=32, nlon=64) + f = torch.randn(1, 32, 64, dtype=torch.float64) + + torch.manual_seed(1234) + expected = torch.randn(5) + + torch.manual_seed(1234) + bsp(f) # first forward triggers the lazy cache build (randperm) + observed = torch.randn(5) + + torch.testing.assert_close(observed, expected) + def test_index_map_structure(self): bsp = SO3onS2(lmax=3, nlat=32, nlon=64) for l1, l2, l in bsp.index_map: