add vectorized jax version of risbo wigner-d recursion

Author: CosmoMatt

docs/api/recursions/index.rst

@@ -89,6 +89,8 @@ Wigner-d recursions
      - Description
    * - :func:`~s2fft.recursions.risbo.compute_full`
      - Compute Wigner-d at argument :math:`\beta` for full plane using Risbo recursion.
+   * - :func:`~s2fft.recursions.risbo_jax.compute_full`
+     - Compute Wigner-d at argument :math:`\beta` for full plane using Risbo recursion (JAX implementation).
 
 .. warning:: 
 

docs/api/recursions/risbo_jax.rst

@@ -0,0 +1,7 @@
+:html_theme.sidebar_secondary.remove:
+
+**************************
+Risbo JAX
+**************************
+.. automodule:: s2fft.recursions.risbo_jax
+   :members: 

s2fft/recursions/__init__.py

@@ -1,5 +1,6 @@
 from . import trapani
 from . import risbo
+from . import risbo_jax
 from . import turok
 from . import turok_jax
 from . import price_mcewen

s2fft/recursions/risbo.py

@@ -56,7 +56,7 @@ def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
         # from l - 1 to l - 1/2.
         dd = np.zeros((2 * el + 2, 2 * el + 2))
         j = 2 * el - 1
-        rj = float(j)  # TODO: is this necessary?
+
         for k in range(0, j):
             sqrt_jmk = np.sqrt(j - k)
             sqrt_kp1 = np.sqrt(k + 1)
@@ -77,7 +77,7 @@ def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
         # 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
-        rj = float(j)  # TODO: is this necessary?
+
         for k in range(0, j):
             sqrt_jmk = np.sqrt(j - k)
             sqrt_kp1 = np.sqrt(k + 1)

s2fft/recursions/risbo_jax.py

@@ -0,0 +1,95 @@
+import jax.numpy as jnp
+from jax import jit
+from functools import partial
+
+
+@partial(jit, static_argnums=(1, 2, 3))
+def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
+    if el == 0:
+        dl = dl.at[el + L - 1, el + L - 1].set(1.0)
+        return dl
+    elif el == 1:
+        cosb = jnp.cos(beta)
+        sinb = jnp.sin(beta)
+
+        coshb = jnp.cos(beta / 2.0)
+        sinhb = jnp.sin(beta / 2.0)
+        sqrt2 = jnp.sqrt(2.0)
+
+        dl = dl.at[L - 2, L - 2].set(coshb**2)
+        dl = dl.at[L - 2, L - 1].set(sinb / sqrt2)
+        dl = dl.at[L - 2, L].set(sinhb**2)
+
+        dl = dl.at[L - 1, L - 2].set(-sinb / sqrt2)
+        dl = dl.at[L - 1, L - 1].set(cosb)
+        dl = dl.at[L - 1, L].set(sinb / sqrt2)
+
+        dl = dl.at[L, L - 2].set(sinhb**2)
+        dl = dl.at[L, L - 1].set(-sinb / sqrt2)
+        dl = dl.at[L, L].set(coshb**2)
+        return dl
+    else:
+        coshb = -jnp.cos(beta / 2.0)
+        sinhb = jnp.sin(beta / 2.0)
+        dd = jnp.zeros((2 * el + 2, 2 * el + 2))
+
+        # First pass
+        j = 2 * el - 1
+        i = jnp.arange(j)
+        k = jnp.arange(j)
+
+        sqrt_jmk = jnp.sqrt(j - k)
+        sqrt_kp1 = jnp.sqrt(k + 1)
+        sqrt_jmi = jnp.sqrt(j - i)
+        sqrt_ip1 = jnp.sqrt(i + 1)
+
+        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
+        )
+        dd = dd.at[:j, 1 : j + 1].add(
+            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
+        )
+        dd = dd.at[1 : j + 1, 1 : j + 1].add(
+            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
+        )
+
+        j = 2 * el
+        i = jnp.arange(j)
+        k = jnp.arange(j)
+
+        # Second pass
+        sqrt_jmk = jnp.sqrt(j - k)
+        sqrt_kp1 = jnp.sqrt(k + 1)
+        sqrt_jmi = jnp.sqrt(j - i)
+        sqrt_ip1 = jnp.sqrt(i + 1)
+
+        dl = dl.at[-el + L - 1 : el + L - 1, -el + L - 1 : el + L - 1].add(
+            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_jmk, optimize=True)
+            * dd[:j, :j]
+            * coshb,
+        )
+        dl = dl.at[-el + L - 1 : el + L - 1, L - el : L + el].add(
+            jnp.einsum("i,k->ki", -sqrt_ip1, sqrt_jmk, optimize=True)
+            * dd[:j, :j]
+            * sinhb,
+        )
+        dl = dl.at[L - el : L + el, -el + L - 1 : el + L - 1].add(
+            jnp.einsum("i,k->ki", sqrt_jmi, sqrt_kp1, optimize=True)
+            * dd[:j, :j]
+            * sinhb,
+        )
+        dl = dl.at[L - el : L + el, L - el : L + el].add(
+            jnp.einsum("i,k->ki", sqrt_ip1, sqrt_kp1, optimize=True)
+            * dd[:j, :j]
+            * coshb,
+        )
+        return dl / ((2 * el) * (2 * el - 1))

tests/test_wigner_recursions.py

@@ -158,12 +158,30 @@ def test_risbo_with_ssht():
 
     # Compare to routines in SSHT, which have been validated extensively.
     dl = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
-    # dl = recursions.trapani.init(dl, L)
+
     for el in range(0, L):
         dl = recursions.risbo.compute_full(dl, beta, L, el)
         np.testing.assert_allclose(dl_array[el, :, :], dl, atol=1e-15)
 
 
+def test_risbo_with_ssht_jax():
+    """Test Risbo JAX computation against ssht"""
+
+    # Test all dl(pi/2) terms up to L.
+    L = 10
+
+    # Compute using SSHT.
+    beta = np.pi / 2.0
+    dl_array = ssht.generate_dl(beta, L)
+
+    # Compare to routines in SSHT, which have been validated extensively.
+    dl = jnp.zeros((2 * L - 1, 2 * L - 1), dtype=jnp.float64)
+
+    for el in range(0, L):
+        dl = recursions.risbo_jax.compute_full(dl, beta, L, el)
+        np.testing.assert_allclose(dl_array[el, :, :], dl, atol=1e-15)
+
+
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("sampling", ["mw", "mwss", "dh", "healpix"])
 def test_turok_with_ssht(L: int, sampling: str):