Repository: astro-informatics/s2fft (language: python)

Issue:
Design method interface



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 7653bf6c..5eebf196 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -1,9 +1,9 @@
 name: Docs
 
-# on:
-#   push:
-#     branches:        
-#     - main
+on:
+  push:
+    branches:
+    - main
     # Add following in for testing
     # - feature/deploy_docs
 
diff --git a/README.rst b/README.rst
index f56bdab9..43a91093 100644
--- a/README.rst
+++ b/README.rst
@@ -21,6 +21,20 @@
 This is a very loose framework for people to start python projects from. To get up and running, go through the code carefully replacing ``Project-name`` with your 
 desired project name (check all documents!), don't forget to change the directory s2fft as well! You will also need to update the links in all badges!
 
+Interface
+=========
+
+Temporary notes on interface to be updated.
+
+flm = forward_transform(f, L, sampling, reality, implementation)
+f = inverse_transform(flm, sampling, reality, implementation, nside=None)
+
+sampling = {"mw", "mwss", "healpix"}; default = mw
+reality = {"real", "complex"}; default = complex
+implementation = {"loopy", "vectorized", "jax"}; default = jax
+nside default = None
+
+
 
 Auto-formatting code
 ====================
diff --git a/s2fft/wigner/__init__.py b/s2fft/wigner/__init__.py
index 1400764d..1d58f330 100644
--- a/s2fft/wigner/__init__.py
+++ b/s2fft/wigner/__init__.py
@@ -1,1 +1,2 @@
 from . import trapani
+from . import risbo
diff --git a/s2fft/wigner/risbo.py b/s2fft/wigner/risbo.py
new file mode 100644
index 00000000..5897ff4c
--- /dev/null
+++ b/s2fft/wigner/risbo.py
@@ -0,0 +1,110 @@
+import numpy as np
+import logs
+
+
+def compute_full(dl: np.ndarray, beta: float, L: int, el: int) -> np.ndarray:
+    """TODO"""
+
+    _arg_checks(dl, L, el)
+
+    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((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)
+
+            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
+        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)
+
+            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
+
+
+def _arg_checks(dl: np.ndarray, L: int, el: int):
+    """TODO: Refactor into general utils for specialfuncs
+
+    Also need to check theta.
+
+    Args:
+
+        dl: Wigner-d plane to check shape of.
+
+        L: Harmonic band-limit.
+
+        ell: Spherical harmonic degree :math:`\ell`.
+    """
+
+    assert 0 <= el < L  # should be < not <= once have init routine
+    assert dl.shape[0] == dl.shape[1] == 2 * L - 1


Test patch:
diff --git a/tests/test_wigner.py b/tests/test_wigner.py
index 938a1877..dd4f71a6 100644
--- a/tests/test_wigner.py
+++ b/tests/test_wigner.py
@@ -73,3 +73,21 @@ def test_trapani_checks():
     # Check throws exception if don't init
 
     return
+
+
+def test_risbo_with_ssht():
+    """Test Risbo 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 = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
+    # dl = wigner.trapani.init(dl, L)
+    for el in range(0, L):
+        dl = wigner.risbo.compute_full(dl, beta, L, el)
+        np.testing.assert_allclose(dl_array[el, :, :], dl, atol=1e-15)