add switch for different wigner precompute modes depending on N/L

Author: CosmoMatt

s2fft/precompute_transforms/alt_construct.py

@@ -84,11 +84,7 @@ def spin_spherical_kernel(
             delta[:, L - 1 - spin],
             1j ** (-spin - m_value[m_start_ind:]),
         )
-        if sampling.lower() in ["mw", "dh"]:
-            temp = np.einsum("am,a->am", temp, np.exp(1j * m_value * thetas[0]))
-        else:
-            temp_new = np.zeros((L + 1, m_dim), dtype=temp.dtype)
-            temp_new[:-1] = temp[L - 1 :]
+        temp = np.einsum("am,a->am", temp, np.exp(1j * m_value * thetas[0]))
         temp = np.fft.irfft(temp[L - 1 :], n=nsamps, axis=0, norm="forward")
 
         dl[:, el] = temp[: len(thetas)]
@@ -173,14 +169,7 @@ def spin_spherical_kernel_jax(
             delta[:, L - 1 - spin],
             1j ** (-spin - m_value[m_start_ind:]),
         )
-        if sampling.lower() in ["mw", "dh"]:
-            temp = jnp.einsum(
-                "am,a->am", temp, jnp.exp(1j * m_value * thetas[0])
-            )
-        else:
-            temp_new = jnp.zeros((L + 1, m_dim), dtype=temp.dtype)
-            temp_new = temp_new.at[:-1].set(temp[L - 1 :])
-
+        temp = jnp.einsum("am,a->am", temp, jnp.exp(1j * m_value * thetas[0]))
         temp = jnp.fft.irfft(temp[L - 1 :], n=nsamps, axis=0, norm="forward")
 
         dl = dl.at[:, el].set(temp[: len(thetas)])
@@ -253,32 +242,35 @@ def wigner_kernel(
         raise ValueError("Sampling in supported list [mw, mwss, dh]")
 
     # Compute Wigner d-functions from their Fourier decomposition.
-    delta = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
+    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)
     dl = np.zeros((n_dim, len(thetas), L, 2 * L - 1), dtype=np.float64)
 
     # Range values which need only be defined once.
     m_value = np.arange(-L + 1, L)
     n = np.arange(n_start_ind - N + 1, N)
 
+    # If N <= L/LogL more efficient to manually compute over FFT
     for el in range(L):
-        delta = recursions.risbo.compute_full(delta, np.pi / 2, L, el)
-        temp = np.einsum(
-            "am,an,m,n->amn",
-            delta,
-            delta[:, L - 1 + n],
-            1j ** (-m_value),
-            1j ** (n),
-        )
-        if sampling.lower() in ["mw", "dh"]:
+        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:
+            delta = recursions.risbo.compute_full(delta, np.pi / 2, L, el)
+            temp = np.einsum(
+                "am,an,m,n->amn",
+                delta,
+                delta[:, L - 1 + n],
+                1j ** (-m_value),
+                1j ** (n),
+            )
             temp = np.einsum(
                 "amn,a->amn", temp, np.exp(1j * m_value * thetas[0])
             )
-
-        else:
-            temp_new = np.zeros((L + 1, 2 * L - 1, len(n)), dtype=temp.dtype)
-            temp_new[:-1] = temp[L - 1 :]
-        temp = np.fft.irfft(temp[L - 1 :], n=nsamps, axis=0, norm="forward")
-        dl[:, :, el] = np.moveaxis(temp[: len(thetas)], -1, 0)
+            temp = np.fft.irfft(temp[L - 1 :], n=nsamps, axis=0, norm="forward")
+            dl[:, :, el] = np.moveaxis(temp[: len(thetas)], -1, 0)
 
     if forward:
         weights = quadrature.quad_weights_transform(L, sampling)
@@ -351,32 +343,42 @@ def wigner_kernel_jax(
         raise ValueError("Sampling in supported list [mw, mwss, dh]")
 
     # Compute Wigner d-functions from their Fourier decomposition.
-    delta = jnp.zeros((2 * L - 1, 2 * L - 1), dtype=jnp.float64)
+    if N <= int(L / np.log(L)):
+        delta = jnp.zeros(
+            (len(thetas), 2 * L - 1, 2 * L - 1), dtype=jnp.float64
+        )
+        vfunc = jax.vmap(
+            recursions.risbo_jax.compute_full, in_axes=(0, 0, None, None)
+        )
+    else:
+        delta = jnp.zeros((2 * L - 1, 2 * L - 1), dtype=jnp.float64)
     dl = jnp.zeros((n_dim, len(thetas), L, 2 * L - 1), dtype=jnp.float64)
 
     # Range values which need only be defined once.
     m_value = jnp.arange(-L + 1, L)
     n = jnp.arange(n_start_ind - N + 1, N)
 
+    # If N <= L/LogL more efficient to manually compute over FFT
     for el in range(L):
-        delta = recursions.risbo_jax.compute_full(delta, jnp.pi / 2, L, el)
-        temp = jnp.einsum(
-            "am,an,m,n->amn",
-            delta,
-            delta[:, L - 1 + n],
-            1j ** (-m_value),
-            1j ** (n),
-        )
-        if sampling.lower() in ["mw", "dh"]:
+        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:
+            delta = recursions.risbo_jax.compute_full(delta, jnp.pi / 2, L, el)
+            temp = jnp.einsum(
+                "am,an,m,n->amn",
+                delta,
+                delta[:, L - 1 + n],
+                1j ** (-m_value),
+                1j ** (n),
+            )
             temp = jnp.einsum(
                 "amn,a->amn", temp, jnp.exp(1j * m_value * thetas[0])
             )
-
-        else:
-            temp_new = jnp.zeros((L + 1, 2 * L - 1, len(n)), dtype=temp.dtype)
-            temp_new = temp_new.at[:-1].set(temp[L - 1 :])
-        temp = jnp.fft.irfft(temp[L - 1 :], n=nsamps, axis=0, norm="forward")
-        dl = dl.at[:, :, el].set(jnp.moveaxis(temp[: len(thetas)], -1, 0))
+            temp = jnp.fft.irfft(
+                temp[L - 1 :], n=nsamps, axis=0, norm="forward"
+            )
+            dl = dl.at[:, :, el].set(jnp.moveaxis(temp[: len(thetas)], -1, 0))
 
     if forward:
         weights = quadrature_jax.quad_weights_transform(L, sampling)

s2fft/recursions/risbo.py

@@ -88,7 +88,9 @@ def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
 
                 ddj = dd[i, k] / j
 
-                dl[k - el + L - 1, i - el + L - 1] += sqrt_jmi * sqrt_jmk * ddj * coshb
+                dl[k - el + L - 1, i - el + L - 1] += (
+                    sqrt_jmi * sqrt_jmk * ddj * coshb
+                )
                 dl[k - el + L - 1, i + 1 - el + L - 1] -= (
                     sqrt_ip1 * sqrt_jmk * ddj * sinhb
                 )
@@ -116,3 +118,105 @@ def _arg_checks(dl: np.ndarray, beta: float, L: int, el: int):
     assert 0 <= el < L  # Should be < not <= once have init routine.
     assert dl.shape[0] == dl.shape[1] == 2 * L - 1
     assert 0 <= beta <= np.pi
+
+
+def compute_full_vect(
+    dl: np.ndarray, beta: np.ndarray, L: int, el: int
+) -> np.ndarray:
+    r"""Compute Wigner-d at argument :math:`\beta` for full plane using
+    Risbo recursion.
+
+    The Wigner-d plane is computed by recursion over :math:`\ell`.
+    Thus, for :math:`\ell > 0` the plane must be computed already for
+    :math:`\ell - 1`. At present, for :math:`\ell = 0` the recusion is initialised.
+
+    Args:
+        dl (np.ndarray): Wigner-d plane for :math:`\ell - 1` at :math:`\beta`.
+
+        beta (np.ndarray): Arguments :math:`\beta` at which to compute Wigner-d plane.
+
+        L (int): Harmonic band-limit.
+
+        el (int): Spherical harmonic degree :math:`\ell`.
+
+    Returns:
+        np.ndarray: Plane of Wigner-d for :math:`\ell` and :math:`\beta`, with full plane computed.
+    """
+    if el == 0:
+        el = 0
+        dl[:, el + L - 1, el + L - 1] = 1.0
+
+    elif el == 1:
+        cosb = np.cos(beta)
+        sinb = np.sin(beta)
+
+        coshb = np.cos(beta / 2.0)
+        sinhb = np.sin(beta / 2.0)
+        sqrt2 = np.sqrt(2.0)
+
+        dl[:, -1 + L - 1, -1 + L - 1] = coshb**2
+        dl[:, -1 + L - 1, 0 + L - 1] = sinb / sqrt2
+        dl[:, -1 + L - 1, 1 + L - 1] = sinhb**2
+
+        dl[:, 0 + L - 1, -1 + L - 1] = -sinb / sqrt2
+        dl[:, 0 + L - 1, 0 + L - 1] = cosb
+        dl[:, 0 + L - 1, 1 + L - 1] = sinb / sqrt2
+
+        dl[:, 1 + L - 1, -1 + L - 1] = sinhb**2
+        dl[:, 1 + L - 1, 0 + L - 1] = -sinb / sqrt2
+        dl[:, 1 + L - 1, 1 + L - 1] = coshb**2
+
+    else:
+        coshb = -np.cos(beta / 2.0)
+        sinhb = np.sin(beta / 2.0)
+
+        # Initialise the plane of the dl-matrix to 0.0 for the recursion
+        # from l - 1 to l - 1/2.
+        dd = np.zeros((dl.shape[0], 2 * el + 2, 2 * el + 2))
+        j = 2 * el - 1
+
+        for k in range(0, j):
+            sqrt_jmk = np.sqrt(j - k)
+            sqrt_kp1 = np.sqrt(k + 1)
+
+            for i in range(0, j):
+                sqrt_jmi = np.sqrt(j - i)
+                sqrt_ip1 = np.sqrt(i + 1)
+
+                dlj = dl[:, k - (el - 1) + L - 1, i - (el - 1) + L - 1] / j
+
+                dd[:, i, k] += sqrt_jmi * sqrt_jmk * dlj * coshb
+                dd[:, i + 1, k] -= sqrt_ip1 * sqrt_jmk * dlj * sinhb
+                dd[:, i, k + 1] += sqrt_jmi * sqrt_kp1 * dlj * sinhb
+                dd[:, i + 1, k + 1] += sqrt_ip1 * sqrt_kp1 * dlj * coshb
+
+        # Having constructed the d^(l+1/2) matrix in dd, do the second
+        # half-step recursion from dd to dl. Start by initilalising
+        # the plane of the dl-matrix to 0.0.
+        dl[:, -el + L - 1 : el + 1 + L - 1, -el + L - 1 : el + 1 + L - 1] = 0.0
+        j = 2 * el
+
+        for k in range(0, j):
+            sqrt_jmk = np.sqrt(j - k)
+            sqrt_kp1 = np.sqrt(k + 1)
+
+            for i in range(0, j):
+                sqrt_jmi = np.sqrt(j - i)
+                sqrt_ip1 = np.sqrt(i + 1)
+
+                ddj = dd[:, i, k] / j
+
+                dl[:, k - el + L - 1, i - el + L - 1] += (
+                    sqrt_jmi * sqrt_jmk * ddj * coshb
+                )
+                dl[:, k - el + L - 1, i + 1 - el + L - 1] -= (
+                    sqrt_ip1 * sqrt_jmk * ddj * sinhb
+                )
+                dl[:, k + 1 - el + L - 1, i - el + L - 1] += (
+                    sqrt_jmi * sqrt_kp1 * ddj * sinhb
+                )
+                dl[:, k + 1 - el + L - 1, i + 1 - el + L - 1] += (
+                    sqrt_ip1 * sqrt_kp1 * ddj * coshb
+                )
+
+    return dl

s2fft/recursions/risbo_jax.py

@@ -3,7 +3,7 @@
 from functools import partial
 
 
-@partial(jit, static_argnums=(1, 2, 3))
+@partial(jit, static_argnums=(2, 3))
 def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
     r"""Compute Wigner-d at argument :math:`\beta` for full plane using
     Risbo recursion (JAX implementation)
@@ -13,7 +13,7 @@ def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
     :math:`\ell - 1`. At present, for :math:`\ell = 0` the recusion is initialised.
 
     Args:
-        dl (np.ndarray): Wigner-d plane for :math:`\ell - 1` at :math:`\beta`.
+        dl (jnp.ndarray): Wigner-d plane for :math:`\ell - 1` at :math:`\beta`.
 
         beta (float): Argument :math:`\beta` at which to compute Wigner-d plane.
 
@@ -22,7 +22,7 @@ def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
         el (int): Spherical harmonic degree :math:`\ell`.
 
     Returns:
-        np.ndarray: Plane of Wigner-d for `el` and `beta`, with full plane computed.
+        jnp.ndarray: Plane of Wigner-d for `el` and `beta`, with full plane computed.
     """
 
     if el == 0:
@@ -66,21 +66,29 @@ def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
         dlj = dl[k - (el - 1) + L - 1][:, i - (el - 1) + L - 1]
 
         dd = dd.at[:j, :j].add(
-            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_jmk, optimize=True) * dlj * coshb
+            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_jmk, optimize=True)
+            * dlj
+            * coshb
         )
         dd = dd.at[:j, 1 : j + 1].add(
-            jnp.einsum("i,k->ki", -sqrt_ip1, sqrt_jmk, optimize=True) * dlj * sinhb
+            jnp.einsum("i,k->ki", -sqrt_ip1, sqrt_jmk, optimize=True)
+            * dlj
+            * sinhb
         )
         dd = dd.at[1 : j + 1, :j].add(
-            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_kp1, optimize=True) * dlj * sinhb
+            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_kp1, optimize=True)
+            * dlj
+            * sinhb
         )
         dd = dd.at[1 : j + 1, 1 : j + 1].add(
-            jnp.einsum("i,k->ki", sqrt_ip1, sqrt_kp1, optimize=True) * dlj * coshb
+            jnp.einsum("i,k->ki", sqrt_ip1, sqrt_kp1, optimize=True)
+            * dlj
+            * coshb
         )
 
-        dl = dl.at[-el + L - 1 : el + 1 + L - 1, -el + L - 1 : el + 1 + L - 1].multiply(
-            0.0
-        )
+        dl = dl.at[
+            -el + L - 1 : el + 1 + L - 1, -el + L - 1 : el + 1 + L - 1
+        ].multiply(0.0)
 
         j = 2 * el
         i = jnp.arange(j)