rename augmentation to rotation, add notebook, update docs

Author: CosmoMatt

docs/api/utility/index.rst

@@ -105,15 +105,15 @@ Utility Functions
       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
+.. list-table:: Rotation functions
    :widths: 25 25
    :header-rows: 1
 
    * - Function Name
      - Description
-   * - :func:`~s2fft.utils.augmentation.rotate_flms`
+   * - :func:`~s2fft.utils.rotation.rotate_flms`
      - Euler rotates spherical harmonic coefficients by given angle in zyz convention.
-   * - :func:`~s2fft.utils.augmentation.generate_rotate_dls`
+   * - :func:`~s2fft.utils.rotation.generate_rotate_dls`
      - Generates an array of all reduced Wigner d-function coefficients for angle beta.
 
 .. toctree::
@@ -128,6 +128,6 @@ Utility Functions
    quadrature_jax
    healpix_ffts
    utils
-   augmentation
+   rotation
    logs
 

docs/api/utility/augmentation.rst renamed to docs/api/utility/rotation.rst

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

docs/tutorials/index.rst

@@ -67,3 +67,4 @@ the case for many other methods that scale linearly with spin).
 
    spherical_harmonic/spherical_harmonic_transform.nblink
    wigner/wigner_transform.nblink
+   rotation/rotation.nblink

docs/tutorials/rotation/rotation.nblink

@@ -0,0 +1,3 @@
+{
+    "path": "../../../notebooks/spherical_rotation.ipynb"
+}

notebooks/spherical_rotation.ipynb

@@ -0,0 +1,129 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# __Rotate a signal__\n",
+    "\n",
+    "---\n",
+    "\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "This tutorial demonstrates how to use `S2FFT` to rotate a signal on the sphere. A signal can be rotated in pixel space, however this can introduce artifacts. The best way to perform a rotation is through spherical harmonic space using the Wigner d-matrices, that is:\n",
+    "\n",
+    "- forward spherical harmonic transform\n",
+    "- rotation on the flm coefficients\n",
+    "- inverse spherical harmonic transform\n",
+    "\n",
+    "Specifically, we will adopt the sampling scheme of [McEwen & Wiaux (2012)](https://arxiv.org/abs/1110.6298). For our purposes here we'll just generate a random bandlimited signal."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from jax.config import config\n",
+    "config.update(\"jax_enable_x64\", True)\n",
+    "\n",
+    "import numpy as np\n",
+    "import s2fft \n",
+    "import plotting_functions\n",
+    "\n",
+    "L = 128\n",
+    "sampling = \"mw\"\n",
+    "rng = np.random.default_rng(12346161)\n",
+    "flm = s2fft.utils.signal_generator.generate_flm(rng, L)\n",
+    "f = s2fft.inverse(flm, L)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Execute the rotation steps\n",
+    "\n",
+    "---\n",
+    "\n",
+    "First, we will run the JAX function to compute the spherical harmonic transform of our signal"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "flm = s2fft.forward_jax(f, L, reality=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now apply the rotation (here pi/2 in each of alpha, beta, gamma) on the harmonic coefficients flm "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "flm_rotated = s2fft.rotate_flms(flm, L, (np.pi/2, np.pi/2,np.pi/2))"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Finally, we will run the JAX function to compute the inverse spherical harmonic transform to get back to pixel space"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_rotated = s2fft.inverse_jax(flm_rotated, L, reality=True)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.4 ('s2fft')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "3425e24474cbe920550266ea26b478634978cc419579f9dbcf479231067df6a3"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

s2fft/__init__.py

@@ -7,6 +7,7 @@
     generate_precomputes_wigner,
     generate_precomputes_wigner_jax,
 )
+from .utils.rotation import rotate_flms, generate_rotate_dls
 
 import logging
 from jax.config import config

s2fft/recursions/risbo_jax.py

@@ -25,22 +25,9 @@ def compute_full(dl: jnp.ndarray, beta: float, L: int, el: int) -> jnp.ndarray:
         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
-
-    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)

s2fft/utils/__init__.py

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

s2fft/utils/rotation.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

@@ -5,7 +5,7 @@
 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
+from s2fft.utils.rotation import rotate_flms, generate_rotate_dls
 import jax.numpy as jnp
 from jax.test_util import check_grads