explicitly handle singularities in risbo JAX recursion

Author: CosmoMatt

s2fft/recursions/risbo_jax.py

@@ -24,10 +24,24 @@ def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
     Returns:
         np.ndarray: Plane of Wigner-d for `el` and `beta`, with full plane computed.
     """
+
+    if beta == 0:
+        dl = dl.at[-el + L - 1, el + L - 1].set(1.0)
+
     if el == 0:
         dl = dl.at[el + L - 1, el + L - 1].set(1.0)
         return dl
-    elif el == 1:
+
+    if beta == jnp.pi:
+        dl = dl.at[:, :].multiply(0)
+        ind = jnp.arange(-el + L - 1, el + L)
+        dl = jnp.flip(dl.at[ind, ind].add(1.0), axis=0)
+        dl = jnp.einsum(
+            "nm,m->nm", dl, (-1) ** (el + jnp.arange(-L + 1, L)), optimize=True
+        )
+        return dl
+
+    if el == 1:
         cosb = jnp.cos(beta)
         sinb = jnp.sin(beta)
 

tests/test_wigner_recursions.py

@@ -136,16 +136,6 @@ def test_trapani_interfaces():
         recursions.trapani.compute_full(dl_jax, L, el, implementation="unexpected")
 
 
-def test_trapani_checks():
-    # TODO
-
-    # Check throws exception if arguments wrong
-
-    # Check throws exception if don't init
-
-    return
-
-
 def test_risbo_with_ssht():
     """Test Risbo computation against ssht"""
 
@@ -171,15 +161,16 @@ def test_risbo_with_ssht_jax():
     L = 10
 
     # Compute using SSHT.
-    beta = np.pi / 2.0
-    dl_array = ssht.generate_dl(beta, L)
+    betas = [0, np.pi / 2.0, np.pi]
+    for beta in betas:
+        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)
+        # 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)
+        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)