add rotation function (risbo) and update docs

Author: CosmoMatt

docs/api/utility/augmentation.rst

@@ -0,0 +1,7 @@
+:html_theme.sidebar_secondary.remove:
+
+**************************
+augmentation
+**************************
+.. automodule:: s2fft.utils.augmentation
+   :members: 

docs/api/utility/index.rst

@@ -100,12 +100,22 @@ Utility Functions
    * - :func:`~s2fft.utils.signal_generator.generate_flmn`
      - Generate a 3D set of random Wigner coefficients.
 
-
 .. note::
 
       JAX versions of these functions share an almost identical function trace and 
       are simply accessed by the sub-module :func:`~s2fft.utils.resampling_jax`.
 
+.. list-table:: Augmentation functions
+   :widths: 25 25
+   :header-rows: 1
+
+   * - Function Name
+     - Description
+   * - :func:`~s2fft.utils.augmentation.rotate_flms`
+     - Euler rotates spherical harmonic coefficients by given angle in zyz convention.
+   * - :func:`~s2fft.utils.augmentation.generate_rotate_dls`
+     - Generates an array of all reduced Wigner d-function coefficients for angle beta.
+
 .. toctree::
    :hidden:
    :maxdepth: 2
@@ -118,5 +128,6 @@ Utility Functions
    quadrature_jax
    healpix_ffts
    utils
+   augmentation
    logs
 

s2fft/recursions/risbo_jax.py

@@ -5,6 +5,25 @@
 
 @partial(jit, static_argnums=(1, 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)
+
+    The Wigner-d plane is computed by recursion over :math:`\ell` (`el`).
+    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 (float): Argument :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 `el` and `beta`, with full plane computed.
+    """
     if el == 0:
         dl = dl.at[el + L - 1, el + L - 1].set(1.0)
         return dl

s2fft/utils/__init__.py

@@ -4,3 +4,4 @@
 from . import resampling_jax
 from . import healpix_ffts
 from . import signal_generator
+from . import augmentation

s2fft/utils/augmentation.py

@@ -0,0 +1,93 @@
+import jax.numpy as jnp
+from jax import jit
+from functools import partial
+from typing import Tuple
+
+from s2fft.recursions.risbo_jax import compute_full
+
+
+@partial(jit, static_argnums=(1, 2))
+def rotate_flms(
+    flm: jnp.ndarray,
+    L: int,
+    rotation: Tuple[float, float, float],
+    dl_array: jnp.ndarray = None,
+) -> jnp.ndarray:
+    """Rotates an array of spherical harmonic coefficients by angle rotation.
+
+    Args:
+        flm (jnp.ndarray): Array of spherical harmonic coefficients.
+        L (int): Harmonic band-limit.
+        rotation  (Tuple[float, float, float]): Rotation on the sphere (alpha, beta, gamma).
+        dl_array (jnp.ndarray, optional): Precomputed array of reduced Wigner d-function
+            coefficients, see :func:~`generate_rotate_dls`. Defaults to None.
+
+    Returns:
+        jnp.ndarray: Rotated spherical harmonic coefficients with shape [L,2L-1].
+    """
+
+    # Split out angles
+    alpha = __exp_array(L, rotation[0])
+    gamma = __exp_array(L, rotation[2])
+    beta = rotation[1]
+
+    # Create empty arrays
+    flm_rotated = jnp.zeros_like(flm)
+
+    dl = (
+        dl_array
+        if dl_array != None
+        else jnp.zeros((2 * L - 1, 2 * L - 1)).astype(jnp.complex128)
+    )
+
+    # Perform rotation
+    for el in range(L):
+        if dl_array is None:
+            dl = compute_full(dl, beta, L, el)
+        n_max = min(el, L - 1)
+
+        m = jnp.arange(-el, el + 1)
+        n = jnp.arange(-n_max, n_max + 1)
+
+        flm_rotated = flm_rotated.at[el, L - 1 + m].add(
+            jnp.einsum(
+                "mn,n->m",
+                jnp.einsum(
+                    "mn,m->mn",
+                    dl[m + L - 1][:, n + L - 1]
+                    if dl_array is None
+                    else dl[el, m + L - 1][:, n + L - 1],
+                    alpha[m + L - 1],
+                    optimize=True,
+                ),
+                gamma[n + L - 1] * flm[el, n + L - 1],
+            )
+        )
+    return flm_rotated
+
+
+@partial(jit, static_argnums=(0, 1))
+def __exp_array(L: int, x: float) -> jnp.ndarray:
+    """Private function to generate rotation arrays for alpha/gamma rotations"""
+    return jnp.exp(-1j * jnp.arange(-L + 1, L) * x)
+
+
+@partial(jit, static_argnums=(0, 1))
+def generate_rotate_dls(L: int, beta: float) -> jnp.ndarray:
+    """Function which recursively generates the complete plane of reduced
+        Wigner d-function coefficients at a given rotation beta.
+
+    Args:
+        L (int): Harmonic band-limit.
+        beta (float): Rotation on the sphere.
+
+    Returns:
+        jnp.ndarray: Complete array of [L, 2L-1,2L-1] Wigner d-function coefficients
+            for a fixed rotation beta.
+    """
+    dl = jnp.zeros((L, 2 * L - 1, 2 * L - 1)).astype(jnp.float64)
+    dl_iter = jnp.zeros((2 * L - 1, 2 * L - 1)).astype(jnp.float64)
+    for el in range(L):
+        dl_iter = compute_full(dl_iter, beta, L, el)
+        dl = dl.at[el].add(dl_iter)
+    return dl

tests/test_utils.py

@@ -1,6 +1,14 @@
+from jax.config import config
+
+config.update("jax_enable_x64", True)
 import pytest
+import pyssht as ssht
 import numpy as np
 from s2fft.sampling import s2_samples as samples
+from s2fft.utils.augmentation import rotate_flms, generate_rotate_dls
+
+L_to_test = [6, 8, 10]
+angles_to_test = [np.pi / 2, np.pi / 6]
 
 
 def test_flm_reindexing_functions(flm_generator):
@@ -47,3 +55,44 @@ def test_flm_reindexing_exceptions(flm_generator):
 
     with pytest.raises(ValueError) as e:
         samples.flm_1d_to_2d(flm_3d, L)
+
+
+@pytest.mark.parametrize("L", L_to_test)
+@pytest.mark.parametrize("alpha", angles_to_test)
+@pytest.mark.parametrize("beta", angles_to_test)
+@pytest.mark.parametrize("gamma", angles_to_test)
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+@pytest.mark.filterwarnings("ignore::FutureWarning")
+def test_rotate_flms(flm_generator, L: int, alpha: float, beta: float, gamma: float):
+    flm = flm_generator(L=L)
+    rot = (alpha, beta, gamma)
+    flm_1d = samples.flm_2d_to_1d(flm, L)
+
+    flm_rot_ssht = samples.flm_1d_to_2d(
+        ssht.rotate_flms(flm_1d, alpha, beta, gamma, L), L
+    )
+    flm_rot_s2fft = rotate_flms(flm, L, rot)
+
+    np.testing.assert_allclose(flm_rot_ssht, flm_rot_s2fft, atol=1e-14)
+
+
+@pytest.mark.parametrize("L", L_to_test)
+@pytest.mark.parametrize("alpha", angles_to_test)
+@pytest.mark.parametrize("beta", angles_to_test)
+@pytest.mark.parametrize("gamma", angles_to_test)
+@pytest.mark.filterwarnings("ignore::RuntimeWarning")
+@pytest.mark.filterwarnings("ignore::FutureWarning")
+def test_rotate_flms_precompute_dls(
+    flm_generator, L: int, alpha: float, beta: float, gamma: float
+):
+    dl = generate_rotate_dls(L, beta)
+    flm = flm_generator(L=L)
+    rot = (alpha, beta, gamma)
+    flm_1d = samples.flm_2d_to_1d(flm, L)
+
+    flm_rot_ssht = samples.flm_1d_to_2d(
+        ssht.rotate_flms(flm_1d, alpha, beta, gamma, L), L
+    )
+    flm_rot_s2fft = rotate_flms(flm, L, rot, dl)
+
+    np.testing.assert_allclose(flm_rot_ssht, flm_rot_s2fft, atol=1e-14)