updated tests

Author: azrael417

tests/test_attention.py

@@ -42,7 +42,7 @@
 from torch_harmonics.attention._attention_utils import _neighborhood_s2_attention_torch
 from torch_harmonics.attention import cuda_kernels_is_available, optimized_kernels_is_available
 
-from testutils import set_seed, compare_tensors
+from testutils import disable_tf32, set_seed, compare_tensors
 
 if not optimized_kernels_is_available():
     print(f"Warning: Couldn't import optimized disco convolution kernels")
@@ -93,6 +93,10 @@ def test_custom_implementation(self, batch_size, channels, channels_out, heads,
         if (self.device.type == "cuda") and (not cuda_kernels_is_available()):
             raise unittest.SkipTest("skipping test because CUDA kernels are not available")
 
+        # disable tf32
+        disable_tf32()
+
+        # set seed
         set_seed(333)
 
         nlat_in, nlon_in = in_shape
@@ -159,6 +163,10 @@ def test_device_vs_cpu(self, batch_size, channels, heads, in_shape, out_shape, g
             # comparing CPU with itself does not make sense
             return
 
+        # disable tf32
+        disable_tf32()
+
+        # set seed
         set_seed(333)
 
         nlat_in, nlon_in = in_shape
@@ -228,6 +236,10 @@ def test_neighborhood_global_equivalence(self, batch_size, channels, channels_ou
         if (self.device.type == "cuda") and (not cuda_kernels_is_available()):
             raise unittest.SkipTest("skipping test because CUDA kernels are not available")
 
+        # disable tf32
+        disable_tf32()
+
+        # set seed
         set_seed(333)
 
         nlat_in, nlon_in = in_shape
@@ -333,6 +345,10 @@ def test_perf(self, batch_size, channels, heads, in_shape, out_shape, grid_in, g
         if (self.device.type == "cuda") and (not cuda_kernels_is_available()):
             raise unittest.SkipTest("skipping test because CUDA kernels are not available")
 
+        # disable tf32
+        disable_tf32()
+
+        # set seed
         set_seed(333)
 
         # extract some parameters

tests/test_sht.py

@@ -35,14 +35,29 @@
 import torch
 from torch.autograd import gradcheck
 import torch_harmonics as th
+from torch_harmonics.quadrature import precompute_latitudes
 
-from testutils import set_seed, compare_tensors
+from testutils import disable_tf32, set_seed, compare_tensors
 
 _devices = [(torch.device("cpu"),)]
 if torch.cuda.is_available():
     _devices.append((torch.device("cuda"),))
 
 
+def random_sht_coeffs(batch_size, lmax, mmax, device, zero_l0=False):
+    """Random scalar SHT coefficients with proper structure:
+    m=0 column real (real-valued field), triangular support (m <= l),
+    and optionally l=0 row zeroed (needed when testing gradient/curl)."""
+    c = torch.randn(batch_size, lmax, mmax, dtype=torch.complex128, device=device)
+    c[:, :, 0] = c[:, :, 0].real
+    for l in range(lmax):
+        if l + 1 < mmax:
+            c[:, l, l + 1:] = 0.0
+    if zero_l0:
+        c[:, 0, :] = 0.0
+    return c
+
+
 class TestLegendrePolynomials(unittest.TestCase):
     """Test the associated Legendre polynomials (CPU/CUDA if available)."""
     def setUp(self):
@@ -112,6 +127,10 @@ def test_forward_inverse(self, nlat, nlon, batch_size, norm, grid, atol, rtol, v
         if verbose:
             print(f"Testing real-valued SHT on {nlat}x{nlon} {grid} grid with {norm} normalization on {self.device.type} device")
 
+        # disable tf32
+        disable_tf32()
+
+        # set seed
         set_seed(333)
 
         testiters = [1, 2, 4, 8, 16]
@@ -173,6 +192,10 @@ def test_grads(self, nlat, nlon, batch_size, norm, grid, atol, rtol, verbose=Fal
         if verbose:
             print(f"Testing gradients of real-valued SHT on {nlat}x{nlon} {grid} grid with {norm} normalization")
 
+        # disable tf32
+        disable_tf32()
+
+        # set seed
         set_seed(333)
 
         if grid == "equiangular":
@@ -203,6 +226,73 @@ def test_grads(self, nlat, nlon, batch_size, norm, grid, atol, rtol, verbose=Fal
         test_result = gradcheck(err_handle, grad_input, eps=1e-6, atol=atol, rtol=rtol)
         self.assertTrue(test_result)
 
+    @parameterized.expand(
+        [
+            [32, 64, 32, "ortho",   "equiangular",   1e-9, 1e-9],
+            [32, 64, 32, "ortho",   "legendre-gauss", 1e-9, 1e-9],
+            [32, 64, 32, "ortho",   "lobatto",        1e-9, 1e-9],
+            [32, 64, 32, "four-pi", "equiangular",   1e-9, 1e-9],
+            [32, 64, 32, "four-pi", "legendre-gauss", 1e-9, 1e-9],
+            [32, 64, 32, "four-pi", "lobatto",        1e-9, 1e-9],
+            [32, 64, 32, "schmidt", "equiangular",   1e-9, 1e-9],
+            [32, 64, 32, "schmidt", "legendre-gauss", 1e-9, 1e-9],
+            [32, 64, 32, "schmidt", "lobatto",        1e-9, 1e-9],
+        ],
+        skip_on_empty=True,
+    )
+    def test_parseval(self, nlat, nlon, batch_size, norm, grid, atol, rtol, verbose=False):
+        """Parseval's theorem: the spatial L2 norm of isht(c) equals a weighted spectral norm
+        of c.  The spectral weights W_{l,m} depend on the normalization convention:
+
+          ortho:   W_{l,m} = w_m                        (w_m = 1 for m=0, 2 for m>0)
+          four-pi: W_{l,m} = w_m / (4*pi)               (coefficients are sqrt(4*pi) larger)
+          schmidt: W_{l,m} = w_m * (2*l+1) / (4*pi)     (coefficients are sqrt(4*pi/(2*l+1)) larger)
+
+        In all cases: ||f||^2_{S^2} = sum_{l,m} W_{l,m} * |c_{l,m}|^2
+        """
+        if verbose:
+            print(f"Testing Parseval's theorem on {nlat}x{nlon} {grid} grid with {norm} normalization on {self.device.type}")
+
+        # disable tf32
+        disable_tf32()
+
+        # set seed
+        set_seed(333)
+
+        isht = th.InverseRealSHT(nlat, nlon, grid=grid, norm=norm).to(self.device)
+        lmax = isht.lmax
+        mmax = isht.mmax
+
+        with torch.no_grad():
+            c = random_sht_coeffs(batch_size, lmax, mmax, self.device)
+            f = isht(c)  # (batch, nlat, nlon)
+
+        # Spatial L2 norm via spherical quadrature: integral of f^2 over S^2
+        _, w_lat = precompute_latitudes(nlat, grid=grid)
+        w_lat = w_lat.to(device=self.device, dtype=torch.float64)
+        dlon = 2.0 * math.pi / nlon
+        spatial_norm_sq = torch.einsum("bnl,n->b", f ** 2, w_lat) * dlon  # (batch,)
+
+        # Build the (lmax, mmax) spectral weight matrix W_{l,m}.
+        # w_m accounts for the ±m folding in the real irfft (m=0: weight 1, m>0: weight 2).
+        w_m = torch.ones(mmax, dtype=torch.float64, device=self.device)
+        w_m[1:] = 2.0
+
+        if norm == "ortho":
+            # c_lm^ortho are the orthonormal coefficients; W_{l,m} = w_m
+            W = w_m.unsqueeze(0).expand(lmax, mmax)
+        elif norm == "four-pi":
+            # c_lm^{four-pi} = sqrt(4*pi) * c_lm^ortho  =>  W_{l,m} = w_m / (4*pi)
+            W = w_m.unsqueeze(0).expand(lmax, mmax) / (4.0 * math.pi)
+        elif norm == "schmidt":
+            # c_lm^{schmidt} = sqrt(4*pi / (2*l+1)) * c_lm^ortho  =>  W_{l,m} = w_m * (2*l+1) / (4*pi)
+            l_vals = torch.arange(lmax, dtype=torch.float64, device=self.device)
+            W = torch.outer(2.0 * l_vals + 1.0, w_m) / (4.0 * math.pi)
+
+        spectral_norm_sq = torch.einsum("blm,lm->b", c.abs() ** 2, W)  # (batch,)
+
+        self.assertTrue(compare_tensors("Parseval's theorem", spatial_norm_sq, spectral_norm_sq, atol=atol, rtol=rtol, verbose=verbose))
+
     @parameterized.expand(
         [
             # even-even
@@ -217,6 +307,10 @@ def test_device_instantiation(self, nlat, nlon, norm, grid, atol, rtol, verbose=
         if verbose:
             print(f"Testing device instantiation of real-valued SHT on {nlat}x{nlon} {grid} grid with {norm} normalization")
 
+        # disable tf32
+        disable_tf32()
+
+        # set seed
         set_seed(333)
 
         # init on cpu
@@ -232,5 +326,207 @@ def test_device_instantiation(self, nlat, nlon, norm, grid, atol, rtol, verbose=
         self.assertTrue(compare_tensors(f"isht weights", isht_host.pct.cpu(), isht_device.pct.cpu(), atol=atol, rtol=rtol, verbose=verbose))
 
 
+@parameterized_class(("device"), _devices)
+class TestSphericalHarmonicsY(unittest.TestCase):
+    """Test fundamental properties of the real spherical harmonic basis functions.
+
+    InverseRealSHT with norm="ortho" synthesizes orthonormal basis functions on
+    the sphere.  Setting a single complex coefficient c_{l,m} = 1 synthesizes
+
+      f_{l,0}      = Y_l^0(theta, phi)                   for m = 0
+      f_{l,m,cos}  ~ P_l^m(cos theta) * cos(m*phi)       for m > 0  (c_{l,m} = 1+0j)
+      f_{l,m,sin}  ~ P_l^m(cos theta) * sin(m*phi)       for m > 0  (c_{l,m} = 0+1j)
+
+    With ortho normalization the sphere inner products satisfy:
+      <f_{l,0},     f_{l',0}    > = delta_{ll'}
+      <f_{l,m,cos}, f_{l',m',*}> = 2 * delta_{ll'} * delta_{mm'}   for m, m' > 0
+      <f_{l,m,sin}, f_{l',m',*}> = 2 * delta_{ll'} * delta_{mm'}   for m, m' > 0
+
+    The factor of 2 for m > 0 arises because the real irfft folds the +m and -m
+    modes together, doubling the amplitude of each mode.
+    """
+
+    @parameterized.expand(
+        [
+            [12, 24, "legendre-gauss", 1e-9, 1e-9],
+            [12, 24, "equiangular",    1e-9, 1e-9],
+            [12, 24, "lobatto",        1e-9, 1e-9],
+        ],
+        skip_on_empty=True,
+    )
+    def test_orthogonality(self, nlat, nlon, grid, atol, rtol, verbose=False):
+        """Verify that isht(norm="ortho") synthesizes mutually orthogonal basis
+        functions and that the self inner-products equal 1 (m=0) or 2 (m>0)."""
+        if verbose:
+            print(f"Testing Y_lm orthogonality on {nlat}x{nlon} {grid} grid on {self.device.type}")
+
+        # disable tf32
+        disable_tf32()
+
+        # set seed
+        set_seed(333)
+
+        if grid == "equiangular":
+            lmax = mmax = nlat // 2
+        elif grid == "lobatto":
+            lmax = mmax = nlat - 1
+        else:
+            lmax = mmax = nlat
+
+        isht = th.InverseRealSHT(nlat, nlon, lmax=lmax, mmax=mmax, grid=grid, norm="ortho").to(self.device)
+
+        # Build one coefficient tensor per real basis function.
+        # For m = 0: one tensor with c[l, 0] = 1+0j (real mode only).
+        # For m > 0: two tensors — c[l, m] = 1+0j (cos) and c[l, m] = 0+1j (sin).
+        basis_list = []
+        expected_diag = []
+        for l in range(lmax):
+            for m in range(min(l + 1, mmax)):
+                c = torch.zeros(lmax, mmax, dtype=torch.complex128)
+                c[l, m] = 1.0 + 0.0j
+                basis_list.append(c)
+                expected_diag.append(1.0 if m == 0 else 2.0)
+                if m > 0:
+                    c = torch.zeros(lmax, mmax, dtype=torch.complex128)
+                    c[l, m] = 0.0 + 1.0j
+                    basis_list.append(c)
+                    expected_diag.append(2.0)
+
+        coeffs = torch.stack(basis_list).to(self.device)  # (N, lmax, mmax)
+
+        with torch.no_grad():
+            funcs = isht(coeffs)  # (N, nlat, nlon), real-valued
+
+        # Gram matrix via spherical quadrature: G[i,j] = integral of f_i * f_j over S^2
+        # Weights from precompute_latitudes are in the cos(theta) domain and integrate
+        # over [-1, 1], so the full measure is w_lat[k] * dlon.
+        _, w_lat = precompute_latitudes(nlat, grid=grid)
+        w_lat = w_lat.to(device=self.device, dtype=torch.float64)
+        dlon = 2.0 * math.pi / nlon
+
+        weighted = funcs * (dlon * w_lat).unsqueeze(-1)  # (N, nlat, nlon)
+        gram = torch.einsum("inl,jnl->ij", weighted, funcs)  # (N, N)
+
+        expected = torch.diag(
+            torch.tensor(expected_diag, dtype=torch.float64, device=self.device)
+        )
+        self.assertTrue(compare_tensors("Gram matrix", gram, expected, atol=atol, rtol=rtol, verbose=verbose))
+
+
+@parameterized_class(("device"), _devices)
+class TestVectorSphericalHarmonicTransform(unittest.TestCase):
+    """Tests for the consistency between the scalar SHT and the vector SHT.
+
+    RealVectorSHT includes a 1/(l*(l+1)) normalization in its quadrature weights
+    so that the spheroidal/toroidal spectral coefficients relate directly to the
+    scalar SHT coefficients of the generating potential:
+
+      Gradient:  vsht(ivsht([c, 0]))[spheroidal] = c,  [toroidal] = 0  (l > 0)
+      Curl:      vsht(ivsht([0, c]))[spheroidal] = 0,  [toroidal] = c  (l > 0)
+
+    The l = 0 mode is zero in both vsht and ivsht because the gradient and curl
+    of a constant field (Y_0^0) vanish identically on the sphere.
+
+    These tests catch swapped spheroidal/toroidal channels, wrong signs in the
+    dP/dtheta or P/sin(theta) terms, and incorrect l*(l+1) normalization.
+    """
+
+    @parameterized.expand(
+        [
+            [32, 64, 16, "ortho",   "legendre-gauss", 1e-7, 1e-7],
+            [32, 64, 16, "ortho",   "equiangular",    1e-7, 1e-7],
+            [32, 64, 16, "ortho",   "lobatto",        1e-7, 1e-7],
+            [32, 64, 16, "four-pi", "legendre-gauss", 1e-7, 1e-7],
+            [32, 64, 16, "four-pi", "equiangular",    1e-7, 1e-7],
+            [32, 64, 16, "four-pi", "lobatto",        1e-7, 1e-7],
+            # [32, 64, 16, "schmidt", "legendre-gauss", 1e-7, 1e-7],
+            # [32, 64, 16, "schmidt", "equiangular",    1e-7, 1e-7],
+            # [32, 64, 16, "schmidt", "lobatto",        1e-7, 1e-7],
+        ],
+        skip_on_empty=True,
+    )
+    def test_gradient_consistency(self, nlat, nlon, batch_size, norm, grid, atol, rtol, verbose=True):
+        """ivsht([c, 0]) synthesizes the surface gradient ∇_S f of a scalar field
+        f = isht(c).  Applying vsht to this gradient field must recover c in the
+        spheroidal channel and zero in the toroidal channel, because a gradient
+        field is curl-free (purely spheroidal).
+        """
+        if verbose:
+            print(f"Testing gradient consistency on {nlat}x{nlon} {grid} grid with {norm} norm on {self.device.type}")
+
+        # disable tf32
+        disable_tf32()
+
+        # set seed
+        set_seed(333)
+
+        vsht  = th.RealVectorSHT        (nlat, nlon, grid=grid, norm=norm).to(self.device)
+        ivsht = th.InverseRealVectorSHT (nlat, nlon, grid=grid, norm=norm).to(self.device)
+        lmax, mmax = vsht.lmax, vsht.mmax
+
+        with torch.no_grad():
+            c     = random_sht_coeffs(batch_size, lmax, mmax, self.device, zero_l0=True)
+            zeros = torch.zeros_like(c)
+
+            # synthesize gradient field: ivsht([c, 0]) = ∇_S f
+            grad_f = ivsht(torch.stack([c, zeros], dim=-3))  # (batch, 2, nlat, nlon)
+
+            # analyse: vsht(∇_S f) must give [c, 0]
+            st = vsht(grad_f)   # (batch, 2, lmax, mmax)
+            s  = st[..., 0, :, :]  # spheroidal
+            t  = st[..., 1, :, :]  # toroidal
+
+        self.assertTrue(compare_tensors("spheroidal coefficients", s, c,     atol=atol, rtol=rtol, verbose=verbose))
+        self.assertTrue(compare_tensors("toroidal coefficients",   t, zeros, atol=atol, rtol=rtol, verbose=verbose))
+
+    @parameterized.expand(
+        [
+            [32, 64, 16, "ortho",   "legendre-gauss", 1e-7, 1e-7],
+            [32, 64, 16, "ortho",   "equiangular",    1e-7, 1e-7],
+            [32, 64, 16, "ortho",   "lobatto",        1e-7, 1e-7],
+            [32, 64, 16, "four-pi", "legendre-gauss", 1e-7, 1e-7],
+            [32, 64, 16, "four-pi", "equiangular",    1e-7, 1e-7],
+            [32, 64, 16, "four-pi", "lobatto",        1e-7, 1e-7],
+            # [32, 64, 16, "schmidt", "legendre-gauss", 1e-9, 1e-9],
+            # [32, 64, 16, "schmidt", "equiangular",    1e-9, 1e-9],
+            # [32, 64, 16, "schmidt", "lobatto",        1e-9, 1e-9],
+        ],
+        skip_on_empty=True,
+    )
+    def test_curl_consistency(self, nlat, nlon, batch_size, norm, grid, atol, rtol, verbose=False):
+        """ivsht([0, c]) synthesizes the surface curl ê_r × ∇_S f of a scalar field
+        f = isht(c).  Applying vsht to this curl field must recover c in the
+        toroidal channel and zero in the spheroidal channel, because a surface
+        curl field is divergence-free (purely toroidal).
+        """
+        if verbose:
+            print(f"Testing curl consistency on {nlat}x{nlon} {grid} grid with {norm} norm on {self.device.type}")
+
+        # disable tf32
+        disable_tf32()
+
+        # set seed
+        set_seed(333)
+
+        vsht  = th.RealVectorSHT        (nlat, nlon, grid=grid, norm=norm).to(self.device)
+        ivsht = th.InverseRealVectorSHT (nlat, nlon, grid=grid, norm=norm).to(self.device)
+        lmax, mmax = vsht.lmax, vsht.mmax
+
+        with torch.no_grad():
+            c     = random_sht_coeffs(batch_size, lmax, mmax, self.device, zero_l0=True)
+            zeros = torch.zeros_like(c)
+
+            # synthesize curl field: ivsht([0, c]) = ê_r × ∇_S f
+            curl_f = ivsht(torch.stack([zeros, c], dim=-3))  # (batch, 2, nlat, nlon)
+
+            # analyse: vsht(ê_r × ∇_S f) must give [0, c]
+            st = vsht(curl_f)   # (batch, 2, lmax, mmax)
+            s  = st[..., 0, :, :]  # spheroidal
+            t  = st[..., 1, :, :]  # toroidal
+
+        self.assertTrue(compare_tensors("spheroidal coefficients", s, zeros, atol=atol, rtol=rtol, verbose=verbose))
+        self.assertTrue(compare_tensors("toroidal coefficients",   t, c,     atol=atol, rtol=rtol, verbose=verbose))
+
+
 if __name__ == "__main__":
     unittest.main()