fix bugs from issue 183

Author: CosmoMatt

s2fft/recursions/risbo.py

@@ -5,7 +5,7 @@ def compute_full(dl: np.ndarray, beta: float, 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` (`el`).
+    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.
 
@@ -19,7 +19,7 @@ def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
         el (int): Spherical harmonic degree :math:`\ell`.
 
     Returns:
-        np.ndarray: Plane of Wigner-d for `el` and `beta`, with full plane computed.
+        np.ndarray: Plane of Wigner-d for :math:`\ell` and :math:`\beta`, with full plane computed.
     """
 
     _arg_checks(dl, beta, L, el)
@@ -103,7 +103,7 @@ def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
 
 
 def _arg_checks(dl: np.ndarray, beta: float, L: int, el: int):
-    """Check arguments of Risbo functions.
+    r"""Check arguments of Risbo functions.
 
     Args:
         dl (np.ndarray): Wigner-d plane of which to check shape.

s2fft/transforms/spherical.py

@@ -296,7 +296,7 @@ def f_bwd(res, gtm):
         ftm = ftm.at[:, m_start_ind + m_offset :].multiply(phase_shifts)
 
     # Perform longitundal Fast Fourier Transforms
-    ftm *= (-1) ** spin
+    ftm *= (-1) ** jnp.abs(spin)
     if reality:
         ftm = ftm.at[:, m_offset : L - 1 + m_offset].set(
             jnp.flip(jnp.conj(ftm[:, L - 1 + m_offset + 1 :]), axis=-1)
@@ -657,4 +657,4 @@ def f_bwd(res, glm):
     indices = jnp.repeat(jnp.expand_dims(jnp.arange(L), -1), 2 * L - 1, axis=-1)
     flm = jnp.where(indices < abs(spin), jnp.zeros_like(flm), flm[..., :])
 
-    return flm * (-1) ** spin
+    return flm * (-1) ** jnp.abs(spin)

s2fft/transforms/wigner.py

@@ -259,7 +259,7 @@ def inverse_jax(
     def spherical_loop(n, args):
         fban, flmn, lrenorm, vsign, spins = args
         fban = fban.at[n].add(
-            (-1) ** spins[n]
+            (-1) ** jnp.abs(spins[n])
             * s2fft.inverse_jax(
                 flmn[n],
                 L,

s2fft/utils/rotation.py

@@ -37,7 +37,7 @@ def rotate_flms(
     dl = (
         dl_array
         if dl_array != None
-        else jnp.zeros((2 * L - 1, 2 * L - 1)).astype(jnp.complex128)
+        else jnp.zeros((2 * L - 1, 2 * L - 1)).astype(jnp.float64)
     )
 
     # Perform rotation