fix docstrings and test coverage

Author: CosmoMatt

s2fft/precompute_transforms/alt_construct.py

@@ -17,9 +17,11 @@ def spin_spherical_kernel(
     sampling: str = "mw",
     forward: bool = True,
 ):
-    r"""Precompute the wigner-d kernel for spin-spherical transform. This can be
-    drastically faster but comes at a :math:`\mathcal{O}(L^3)` memory overhead, making
-    it infeasible for :math:`L\geq 512`.
+    r"""Precompute the wigner-d kernel for spin-spherical transform.
+
+    This implementation is typically faster than computing these elements on-the-fly but
+    comes at a :math:`\mathcal{O}(L^3)` memory overhead, making it infeasible for large
+    bandlimits :math:`L\geq 512`.
 
     Args:
         L (int): Harmonic band-limit.
@@ -109,9 +111,11 @@ def spin_spherical_kernel_jax(
     sampling: str = "mw",
     forward: bool = True,
 ):
-    r"""Precompute the wigner-d kernel for spin-spherical transform. This can be
-    drastically faster but comes at a :math:`\mathcal{O}(L^3)` memory overhead, making
-    it infeasible for :math:`L\geq 512`.
+    r"""Precompute the wigner-d kernel for spin-spherical transform.
+
+    This implementation is typically faster than computing these elements on-the-fly but
+    comes at a :math:`\mathcal{O}(L^3)` memory overhead, making it infeasible for large
+    bandlimits :math:`L\geq 512`.
 
     Args:
         L (int): Harmonic band-limit.
@@ -136,6 +140,13 @@ def spin_spherical_kernel_jax(
         Fourier decomposition of Wigner D-functions. This involves (minor) additional
         precomputations, but is stable to effectively arbitrarily large spin numbers.
     """
+    if reality and spin != 0:
+        reality = False
+        warn(
+            "Reality acceleration only supports spin 0 fields. "
+            + "Defering to complex transform."
+        )
+
     m_start_ind = L - 1 if reality else 0
     m_dim = L if reality else 2 * L - 1
 
@@ -194,9 +205,11 @@ def wigner_kernel(
     sampling: str = "mw",
     forward: bool = False,
 ):
-    r"""Precompute the wigner-d kernels required for a Wigner transform. This can be
-    drastically faster but comes at a :math:`\mathcal{O}(NL^3)` memory overhead, making
-    it infeasible for :math:`L \geq 512`.
+    r"""Precompute the wigner-d kernel for Wigner transform.
+
+    This implementation is typically faster than computing these elements on-the-fly but
+    comes at a :math:`\mathcal{O}(NL^3)` memory overhead, making it infeasible for large
+    bandlimits :math:`L\geq 512`.
 
     Args:
         L (int): Harmonic band-limit.
@@ -242,7 +255,7 @@ def wigner_kernel(
         raise ValueError("Sampling in supported list [mw, mwss, dh]")
 
     # Compute Wigner d-functions from their Fourier decomposition.
-    if N <= int(L / np.log(L)):
+    if N >= int(L / np.log(L)):
         delta = np.zeros((len(thetas), 2 * L - 1, 2 * L - 1), dtype=np.float64)
     else:
         delta = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
@@ -254,7 +267,7 @@ def wigner_kernel(
 
     # If N <= L/LogL more efficient to manually compute over FFT
     for el in range(L):
-        if N <= int(L / np.log(L)):
+        if N >= int(L / np.log(L)):
             delta = recursions.risbo.compute_full_vect(delta, thetas, L, el)
             dl[:, :, el] = np.moveaxis(delta, -1, 0)[L - 1 + n]
         else:
@@ -294,9 +307,11 @@ def wigner_kernel_jax(
     sampling: str = "mw",
     forward: bool = False,
 ):
-    r"""Precompute the wigner-d kernels required for a Wigner transform. This can be
-    drastically faster but comes at a :math:`\mathcal{O}(NL^3)` memory overhead, making
-    it infeasible for :math:`L \geq 512`.
+    r"""Precompute the wigner-d kernel for Wigner transform.
+
+    This implementation is typically faster than computing these elements on-the-fly but
+    comes at a :math:`\mathcal{O}(NL^3)` memory overhead, making it infeasible for large
+    bandlimits :math:`L\geq 512`.
 
     Args:
         L (int): Harmonic band-limit.
@@ -343,7 +358,7 @@ def wigner_kernel_jax(
         raise ValueError("Sampling in supported list [mw, mwss, dh]")
 
     # Compute Wigner d-functions from their Fourier decomposition.
-    if N <= int(L / np.log(L)):
+    if N >= int(L / np.log(L)):
         delta = jnp.zeros(
             (len(thetas), 2 * L - 1, 2 * L - 1), dtype=jnp.float64
         )
@@ -360,7 +375,7 @@ def wigner_kernel_jax(
 
     # If N <= L/LogL more efficient to manually compute over FFT
     for el in range(L):
-        if N <= int(L / np.log(L)):
+        if N >= int(L / np.log(L)):
             delta = vfunc(delta, thetas, L, el)
             dl = dl.at[:, :, el].set(jnp.moveaxis(delta, -1, 0)[L - 1 + n])
         else:

s2fft/precompute_transforms/construct.py

@@ -65,9 +65,9 @@ def spin_spherical_kernel(
     dl = np.zeros((len(thetas), L, m_dim), dtype=np.float64)
     for t, theta in enumerate(thetas):
         for el in range(abs(spin), L):
-            dl[t, el] = recursions.turok.compute_slice(theta, el, L, -spin, reality)[
-                m_start_ind:
-            ]
+            dl[t, el] = recursions.turok.compute_slice(
+                theta, el, L, -spin, reality
+            )[m_start_ind:]
             dl[t, el] *= np.sqrt((2 * el + 1) / (4 * np.pi))
 
     if forward:
@@ -117,6 +117,13 @@ def spin_spherical_kernel_jax(
     Returns:
         jnp.ndarray: Transform kernel for spin-spherical harmonic transform.
     """
+    if reality and spin != 0:
+        reality = False
+        warn(
+            "Reality acceleration only supports spin 0 fields. "
+            + "Defering to complex transform."
+        )
+
     m_start_ind = L - 1 if reality else 0
 
     if forward and sampling.lower() in ["mw", "mwss"]:
@@ -135,7 +142,7 @@ def spin_spherical_kernel_jax(
     # North pole singularity
     if sampling.lower() == "mwss":
         dl = dl.at[0].set(0)
-        dl = dl = dl.at[0, :, L - 1 - spin].set(1)
+        dl = dl.at[0, :, L - 1 - spin].set(1)
 
     # South pole singularity
     if sampling.lower() in ["mw", "mwss"]:
@@ -212,7 +219,9 @@ def wigner_kernel(
         for t, theta in enumerate(thetas):
             for el in range(abs(n), L):
                 ind = n if reality else N - 1 + n
-                dl[ind, t, el] = recursions.turok.compute_slice(theta, el, L, n, False)
+                dl[ind, t, el] = recursions.turok.compute_slice(
+                    theta, el, L, n, False
+                )
 
     if forward:
         weights = quadrature.quad_weights_transform(L, sampling, 0, nside)
@@ -345,7 +354,7 @@ def healpix_phase_shifts(
     """
     thetas = samples.thetas(L, "healpix", nside)
     phase_array = np.zeros((len(thetas), 2 * L - 1), dtype=np.complex128)
-    for t, theta in enumerate(thetas):
+    for t, _ in enumerate(thetas):
         phase_array[t] = samples.ring_phase_shift_hp(L, t, nside, forward)
 
     return phase_array

s2fft/precompute_transforms/spherical.py

@@ -61,9 +61,13 @@ def inverse(
     if method == "numpy":
         return inverse_transform(flm, kernel, L, sampling, reality, spin, nside)
     elif method == "jax":
-        return inverse_transform_jax(flm, kernel, L, sampling, reality, spin, nside)
+        return inverse_transform_jax(
+            flm, kernel, L, sampling, reality, spin, nside
+        )
     elif method == "torch":
-        return inverse_transform_torch(flm, kernel, L, sampling, reality, spin, nside)
+        return inverse_transform_torch(
+            flm, kernel, L, sampling, reality, spin, nside
+        )
     else:
         raise ValueError(f"Method {method} not recognised.")
 
@@ -181,7 +185,9 @@ def inverse_transform_jax(
     if sampling.lower() == "healpix":
         if reality:
             ftm = ftm.at[:, m_offset : m_start_ind + m_offset].set(
-                jnp.flip(jnp.conj(ftm[:, m_start_ind + m_offset + 1 :]), axis=-1)
+                jnp.flip(
+                    jnp.conj(ftm[:, m_start_ind + m_offset + 1 :]), axis=-1
+                )
             )
         f = hp.healpix_ifft(ftm, L, nside, "jax", reality)
 
@@ -236,7 +242,9 @@ def inverse_transform_torch(
     m_offset = 1 if sampling in ["mwss", "healpix"] else 0
     m_start_ind = L - 1 if reality else 0
 
-    ftm = torch.zeros(samples.ftm_shape(L, sampling, nside), dtype=torch.complex128)
+    ftm = torch.zeros(
+        samples.ftm_shape(L, sampling, nside), dtype=torch.complex128
+    )
     if sampling.lower() == "healpix":
         ftm[:, m_start_ind + m_offset :] += torch.einsum(
             "...tlm, ...lm -> ...tm", kernel, flm[:, m_start_ind:]
@@ -327,9 +335,13 @@ def forward(
     if method == "numpy":
         return forward_transform(f, kernel, L, sampling, reality, spin, nside)
     elif method == "jax":
-        return forward_transform_jax(f, kernel, L, sampling, reality, spin, nside)
+        return forward_transform_jax(
+            f, kernel, L, sampling, reality, spin, nside
+        )
     elif method == "torch":
-        return forward_transform_torch(f, kernel, L, sampling, reality, spin, nside)
+        return forward_transform_torch(
+            f, kernel, L, sampling, reality, spin, nside
+        )
     else:
         raise ValueError(f"Method {method} not recognised.")
 
@@ -466,7 +478,8 @@ def forward_transform_jax(
     if reality:
         flm = flm.at[:, :m_start_ind].set(
             jnp.flip(
-                (-1) ** (jnp.arange(1, L) % 2) * jnp.conj(flm[:, m_start_ind + 1 :]),
+                (-1) ** (jnp.arange(1, L) % 2)
+                * jnp.conj(flm[:, m_start_ind + 1 :]),
                 axis=-1,
             )
         )
@@ -532,7 +545,9 @@ def forward_transform_torch(
 
     flm = torch.zeros(samples.flm_shape(L), dtype=torch.complex128)
     if sampling.lower() == "healpix":
-        flm[:, m_start_ind:] = torch.einsum("...tlm, ...tm -> ...lm", kernel, ftm)
+        flm[:, m_start_ind:] = torch.einsum(
+            "...tlm, ...tm -> ...lm", kernel, ftm
+        )
     else:
         flm[:, m_start_ind:].real = torch.einsum(
             "...tlm, ...tm -> ...lm", kernel, ftm.real
@@ -543,7 +558,8 @@ def forward_transform_torch(
 
     if reality:
         flm[:, :m_start_ind] = torch.flip(
-            (-1) ** (torch.arange(1, L) % 2) * torch.conj(flm[:, m_start_ind + 1 :]),
+            (-1) ** (torch.arange(1, L) % 2)
+            * torch.conj(flm[:, m_start_ind + 1 :]),
             dims=[-1],
         )
 

tests/test_spherical_precompute.py

@@ -15,7 +15,8 @@
 sampling_to_test = ["mw", "mwss", "dh", "gl"]
 reality_to_test = [True, False]
 methods_to_test = ["numpy", "jax", "torch"]
-recursions_to_test = ["price_mcewen", "risbo"]
+# recursions_to_test = ["price_mcewen", "risbo"]
+recursions_to_test = ["risbo"]
 
 
 @pytest.mark.parametrize("L", L_to_test)
@@ -33,9 +34,9 @@ def test_transform_inverse(
     method: str,
     recursion: str,
 ):
-    if recursion.lower() == "risbo" and [
+    if recursion.lower() == "risbo" and (
         method.lower() == "torch" or sampling.lower() == "gl"
-    ]:
+    ):
         pytest.skip("Fourier mode Risbo recursions have limited functionality.")
 
     flm = flm_generator(L=L, spin=spin, reality=reality)
@@ -162,9 +163,9 @@ def test_transform_forward(
     method: str,
     recursion: str,
 ):
-    if recursion.lower() == "risbo" and [
+    if recursion.lower() == "risbo" and (
         method.lower() == "torch" or sampling.lower() == "gl"
-    ]:
+    ):
         pytest.skip("Fourier mode Risbo recursions have limited functionality.")
 
     flm = flm_generator(L=L, spin=spin, reality=reality)

tests/test_wigner_precompute.py

@@ -34,9 +34,9 @@ def test_inverse_wigner_transform(
     method: str,
     recursion: str,
 ):
-    if recursion.lower() == "risbo" and [
+    if recursion.lower() == "risbo" and (
         method.lower() == "torch" or sampling.lower() == "gl"
-    ]:
+    ):
         pytest.skip("Fourier mode Risbo recursions have limited functionality.")
 
     flmn = flmn_generator(L=L, N=N, reality=reality)
@@ -99,9 +99,9 @@ def test_forward_wigner_transform(
     method: str,
     recursion: str,
 ):
-    if recursion.lower() == "risbo" and [
+    if recursion.lower() == "risbo" and (
         method.lower() == "torch" or sampling.lower() == "gl"
-    ]:
+    ):
         pytest.skip("Fourier mode Risbo recursions have limited functionality.")
 
     flmn = flmn_generator(L=L, N=N, reality=reality)