Merge pull request #162 from astro-informatics/feature/custom_grads

Feature/custom grads

Author: CosmoMatt

.github/workflows/tests.yml

@@ -25,14 +25,14 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
-    # - name: Install dependencies
-    #   run: |
-    #     python -m pip install --upgrade pip
-    #     pip install -r requirements/requirements-tests.txt
-    #     pip install -r requirements/requirements-core.txt
-    #     pip install .
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements/requirements-tests.txt
+        pip install -r requirements/requirements-core.txt
+        pip install .
     
-    # - name: Run tests
-    #   run: |
-    #     pytest --cov-report term --cov=s2wav --cov-config=.coveragerc 
-    #     codecov --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702
+    - name: Run tests
+      run: |
+        pytest --cov-report term --cov=s2wav --cov-config=.coveragerc 
+        codecov --token 298dc7ee-bb9f-4221-b31f-3576cc6cb702

README.md

@@ -60,7 +60,7 @@ pixels of equal areas, which has many practical advantages.
 
 The Python dependencies for the `S2FFT` package are listed in the file
 `requirements/requirements-core.txt` and will be automatically installed
-into the active python environment by [pip]{.title-ref} when running
+into the active python environment by [pip](https://pypi.org) when running
 
 ``` bash
 pip install .        
@@ -75,7 +75,7 @@ tox -e py38           # for tox
 ```
 
 In the very near future one will be able to install `S2FFT` directly
-from [PyPi]{.title-ref} by `pip install s2fft` but this is not yet
+from [PyPi](https://pypi.org) by `pip install s2fft` but this is not yet
 supported. Note that to run `JAX` on NVIDIA GPUs you will need to follow
 the [guide](https://github.com/google/jax#installation) outlined by
 Google. 

notebooks/custom_gradients.ipynb

@@ -0,0 +1,112 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "cpu\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Specify CUDA device\n",
+    "import os\n",
+    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
+    "os.environ['JAX_CHECK_TRACER_LEAKS'] = 'True'\n",
+    "\n",
+    "from jax.config import config\n",
+    "config.update(\"jax_enable_x64\", True)\n",
+    "\n",
+    "# Check we're running on GPU\n",
+    "from jax.lib import xla_bridge\n",
+    "print(xla_bridge.get_backend().platform)\n",
+    "\n",
+    "import jax\n",
+    "from jax import jit, grad \n",
+    "import jax.numpy as jnp \n",
+    "from jax.test_util import check_grads\n",
+    "import numpy as np \n",
+    "\n",
+    "import s2fft "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "L = 16\n",
+    "sampling = \"mw\"\n",
+    "np.random.seed(1911851)\n",
+    "f_target = np.random.randn(2*L, 2*L-1)+1j*np.random.randn(2*L, 2*L-1)\n",
+    "flm_target = s2fft.forward_jax(f_target, L, sampling=sampling)\n",
+    "f_target = s2fft.inverse_jax(flm_target, L, sampling=sampling)\n",
+    "precomps = s2fft.generate_precomputes_jax(L, forward=True, sampling=sampling)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.random.seed(130672510)\n",
+    "f = np.random.randn(2*L, 2*L-1) + 1j*np.random.randn(2*L, 2*L-1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def func(f):\n",
+    "    flm = s2fft.forward_jax(f, L, reality=False, precomps=precomps,sampling=sampling)\n",
+    "    return jnp.sum(jnp.abs(flm-flm_target)**2)\n",
+    "grad_func = grad(func)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "check_grads(func, (f,), order=1, modes=('rev'))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.9.0 ('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.9.0"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "3425e24474cbe920550266ea26b478634978cc419579f9dbcf479231067df6a3"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

s2fft/recursions/price_mcewen.py

@@ -122,6 +122,7 @@ def generate_precomputes_jax(
     nside: int = None,
     forward: bool = False,
     L_lower: int = 0,
+    betas: jnp.ndarray = None,
 ) -> List[jnp.ndarray]:
     r"""Compute recursion coefficients with :math:`\mathcal{O}(L^2)` memory overhead.
     In practice one could compute these on-the-fly but the memory overhead is
@@ -145,17 +146,22 @@ def generate_precomputes_jax(
         L_lower (int, optional): Harmonic lower-bound. Transform will only be computed
             for :math:`\texttt{L_lower} \leq \ell < \texttt{L}`. Defaults to 0.
 
+        beta (jnp.ndarray): Array of polar angles in radians.
+
     Returns:
         List[jnp.ndarray]: List of precomputed coefficient arrays.
     """
     mm = -spin
     L0 = L_lower
     # Correct for mw to mwss conversion
-    if forward and sampling.lower() in ["mw", "mwss"]:
-        sampling = "mwss"
-        beta = samples.thetas(2 * L, "mwss")[1:-1]
+    if betas is None:
+        if forward and sampling.lower() in ["mw", "mwss"]:
+            sampling = "mwss"
+            beta = samples.thetas(2 * L, "mwss")[1:-1]
+        else:
+            beta = samples.thetas(L, sampling, nside)
     else:
-        beta = samples.thetas(L, sampling, nside)
+        beta = betas
 
     ntheta = len(beta)  # Number of theta samples
     el = jnp.arange(L0, L)
@@ -243,7 +249,7 @@ def renorm_m_loop(i, args):
 
     # Remove redundant nans:
     # - in forward pass these are not accessed, so are irrelevant.
-    # - in backward pass the adjoint computation otherwise accumulates these 
+    # - in backward pass the adjoint computation otherwise accumulates these
     #   nans into grads if not set to zero.
     lrenorm = jnp.nan_to_num(lrenorm, nan=0.0, posinf=0.0, neginf=0.0)
     cpi = jnp.nan_to_num(cpi, nan=0.0, posinf=0.0, neginf=0.0)

s2fft/transforms/otf_recursions.py

@@ -1,7 +1,7 @@
 import numpy as np
 import jax.numpy as jnp
 import jax.lax as lax
-from jax import jit, pmap, local_device_count
+from jax import jit, pmap, local_device_count, custom_vjp
 from functools import partial
 from typing import List
 from s2fft.sampling import s2_samples as samples
@@ -225,7 +225,8 @@ def inverse_latitudinal_step_jax(
 
     mm = -spin  # switch to match convention
     ntheta = len(beta)  # Number of theta samples
-    ftm = jnp.zeros(samples.ftm_shape(L, sampling, nside), dtype=jnp.complex128)
+    m_count = 2 * L if sampling.lower() in ["mwss", "healpix"] else 2 * L - 1
+    ftm = jnp.zeros((ntheta, m_count), dtype=jnp.complex128)
     el = jnp.arange(L_lower, L)
 
     # Trigonometric constant adopted throughout
@@ -239,7 +240,7 @@ def inverse_latitudinal_step_jax(
 
     if precomps is None:
         precomps = generate_precomputes_jax(
-            L, -mm, sampling, nside, L_lower=L_lower
+            L, -mm, sampling, nside, L_lower=L_lower, betas=beta
         )
     lrenorm, vsign, cpi, cp2, indices = precomps
 
@@ -401,6 +402,25 @@ def eval_recursion_step(
                     pm_recursion_step,
                     (ftm, dl_entry, dl_iter, lrenorm, indices, omc, c, s),
                 )
+
+    # Remove south pole singularity
+    m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0
+    if sampling.lower() in ["mw", "mwss"]:
+        ftm = ftm.at[-1].set(0)
+        ftm = ftm.at[-1, L - 1 + spin + m_offset].set(
+            jnp.nansum(
+                (-1) ** abs(jnp.arange(L_lower, L) - spin)
+                * flm[L_lower:, L - 1 + spin]
+            )
+        )
+
+    # Remove north pole singularity
+    if sampling.lower() == "mwss":
+        ftm = ftm.at[0].set(0)
+        ftm = ftm.at[0, L - 1 - spin + m_offset].set(
+            jnp.nansum(flm[L_lower:, L - 1 - spin])
+        )
+
     return ftm
 
 
@@ -562,8 +582,8 @@ def forward_latitudinal_step(
 
 @partial(jit, static_argnums=(2, 4, 5, 6, 8, 9))
 def forward_latitudinal_step_jax(
-    ftm: jnp.ndarray,
-    beta: jnp.ndarray,
+    ftm_in: jnp.ndarray,
+    beta_in: jnp.ndarray,
     L: int,
     spin: int,
     nside: int,
@@ -622,6 +642,16 @@ def forward_latitudinal_step_jax(
         between devices is noticable, however as L increases one will asymptotically
         recover acceleration by the number of devices.
     """
+    # Avoid pole-singularities for MWSS sampling
+    if sampling.lower() == "mwss":
+        ftm = ftm_in[1:-1]
+        beta = beta_in[1:-1]
+    elif sampling.lower() == "mw":
+        ftm = ftm_in[:-1]
+        beta = beta_in[:-1]
+    else:
+        ftm = ftm_in
+        beta = beta_in
 
     mm = -spin  # switch to match convention
     ntheta = len(beta)  # Number of theta samples
@@ -639,7 +669,7 @@ def forward_latitudinal_step_jax(
 
     if precomps is None:
         precomps = generate_precomputes_jax(
-            L, -mm, sampling, nside, True, L_lower
+            L, -mm, sampling, nside, True, L_lower, betas=beta
         )
     lrenorm, vsign, cpi, cp2, indices = precomps
 
@@ -840,4 +870,17 @@ def eval_recursion_step(
                         ),
                     )[0]
                 )
+
+    # Include both pole singularities explicitly
+    m_offset = 1 if sampling.lower() in ["mwss", "healpix"] else 0
+    if sampling.lower() in ["mw", "mwss"]:
+        flm = flm.at[L_lower:, L - 1 + spin].add(
+            (-1) ** abs(jnp.arange(L_lower, L) - spin)
+            * ftm_in[-1, L - 1 + spin + m_offset]
+        )
+
+    if sampling.lower() == "mwss":
+        flm = flm.at[L_lower:, L - 1 - spin].add(
+            ftm_in[0, L - 1 - spin + m_offset]
+        )
     return flm