Jaxifying forward transform refactored

notebooks/notebook_fwd_jax_vs_gt.py

@@ -0,0 +1,121 @@
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+import sys
+
+sys.path.append("../")
+
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+
+import numpy as np
+from jax.config import config
+import pyssht as ssht
+import s2fft as s2f
+
+config.update("jax_enable_x64", True)
+
+# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
+# Compare All JAX methods to ground truth
+
+# generate spherical harmonics (ground truth)
+DEFAULT_SEED = 8966433580120847635
+rn_gen = np.random.default_rng(DEFAULT_SEED)  # modify to use JAX random key approach?
+
+### input params
+L_in = 6  # input L, redefined to 2*nside if sampling is healpix
+spin = -2
+reality = False
+L_lower = 1
+nside_healpix = 2
+
+# list jax implementations
+list_jax_str = [
+    "jax_vmap_double",
+    "jax_vmap_scan",
+    "jax_vmap_loop",
+    "jax_map_double",
+    "jax_map_scan",
+]
+
+# list of sampling approaches
+list_sampling = ["mw", "mwss", "dh", "healpix"]
+
+# Print inputs common to all sampling methods
+print("-----------------------------------------")
+print("Input params:")
+print(
+    f"L_in = {L_in}, spin = {spin}, reality = {reality}, L_lower = {L_lower}, nside_healpix={nside_healpix}"
+)
+print("-----------------------------------------")
+
+# All JAX methods
+for m_str in list_jax_str:
+
+    for sampling in list_sampling:
+
+        # Groundtruth and starting point f
+        if sampling != "healpix":
+            # Set nside to None if not healpix
+            nside = None
+            L = L_in
+
+            # compute ground truth and starting point f
+            flm_gt = s2f.utils.generate_flm(rn_gen, L, L_lower, spin, reality=reality)
+            f = ssht.inverse(
+                s2f.samples.flm_2d_to_1d(flm_gt, L),
+                L,
+                Method=sampling.upper(),
+                Spin=spin,
+                Reality=False,
+            )
+
+        else:
+            # Set nside to healpix value and redefine L
+            nside = nside_healpix
+            L = 2 * nside  # L redefined to double nside
+
+            # compute ground truth and starting point f
+            flm_gt0 = s2f.utils.generate_flm(rn_gen, L, L_lower, spin, reality=reality)
+            f = s2f.transform._inverse(
+                flm_gt0, L, spin, sampling=sampling, method="direct", nside=nside, reality=reality, L_lower=L_lower
+            )
+
+            # Compute numpy solution if sampling is healpix (used as GT)
+            flm_sov_fft_vec = s2f.transform._forward(
+                f,
+                L,
+                spin,
+                sampling,
+                method="sov_fft_vectorized",
+                nside=nside,
+                reality=reality,
+                L_lower=L_lower,
+            )
+
+        # JAX implementation
+        flm_sov_fft_vec_jax = s2f.transform._forward(
+            f,
+            L,
+            spin,
+            sampling,
+            method=m_str,
+            nside=nside,
+            reality=reality,
+            L_lower=L_lower,
+        )
+
+        # Compare to GT
+        if sampling == "healpix":  # compare to numpy result rather than GT
+            print(
+                f"{m_str} vs numpy ({sampling}, L = {L}): {np.allclose(flm_sov_fft_vec, flm_sov_fft_vec_jax, atol=1e-14, rtol=1e-07)}"
+            )
+        else:
+            print(
+                f"{m_str} vs GT ({sampling}, L = {L}): {np.allclose(flm_gt, flm_sov_fft_vec_jax, atol=1e-14, rtol=1e-07)}" # rtol=1e-07 default in np.testing.assert_allclose
+            )
+
+        # %timeit s2f.transform._forward(f, L, spin, sampling, method=m_str, nside=nside, reality=reality, L_lower=L_lower)
+        # %timeit s2f.transform._forward(f, L, spin, sampling, method=m_str, nside=nside, reality=reality, L_lower=L_lower)
+    print("-----------------------------------------")
+
+
+# %%

s2fft/healpix_ffts.py

@@ -2,6 +2,12 @@
 import numpy.fft as fft
 import s2fft.samples as samples
 
+import jax
+import jax.numpy as jnp
+import jax.numpy.fft as jfft
+
+from functools import partial
+
 
 def spectral_folding(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
     """Folds higher frequency Fourier coefficients back onto lower frequency
@@ -35,41 +41,34 @@ def spectral_folding(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
     return ftm_slice
 
 
-def spectral_periodic_extension(fm: np.ndarray, nphi: int, L: int) -> np.ndarray:
+def spectral_periodic_extension(fm: np.ndarray, L: int, numpy_module=np) -> np.ndarray: 
     """Extends lower frequency Fourier coefficients onto higher frequency
-    coefficients, i.e. imposed periodicity in Fourier space.
+    coefficients, i.e. imposed periodicity in Fourier space. 
+    Based on `spectral_periodic_extension`, modified to be JIT-compilable.
 
     Args:
         fm (np.ndarray): Slice of Fourier coefficients corresponding to ring at latitute t.
 
-        nphi (int): Total number of pixel space phi samples for latitude t.
-
         L (int): Harmonic band-limit.
+
+        numpy_module: JAX's Numpy-like API or Numpy. Default Numpy.
 
     Returns:
         np.ndarray: Higher resolution set of periodic Fourier coefficients.
     """
-    assert nphi <= 2 * L
-
-    slice_start = L - nphi // 2
-    slice_stop = slice_start + nphi
-    fm_full = np.zeros(2 * L, dtype=np.complex128)
-    fm_full[slice_start:slice_stop] = fm
-
-    idx = 1
-    while slice_start - idx >= 0:
-        fm_full[slice_start - idx] = fm[-idx % nphi]
-        idx += 1
-    idx = 0
-    while slice_stop + idx < len(fm_full):
-        fm_full[slice_stop + idx] = fm[idx % nphi]
-        idx += 1
-
-    return fm_full
-
-
-def healpix_fft(f: np.ndarray, L: int, nside: int) -> np.ndarray:
-    """Computes the Forward Fast Fourier Transform with spectral back-projection
+    nphi = fm.shape[0] 
+    return numpy_module.concatenate(
+        ( 
+            fm[-numpy_module.arange(L - nphi // 2, 0, -1) % nphi], 
+            fm, 
+            fm[numpy_module.arange(L - (nphi + 1) // 2) % nphi] 
+        ) 
+    )
+
+
+def healpix_fft(f: np.ndarray, L: int, nside: int, numpy_module=np) -> np.ndarray:
+    '''
+    Computes the Forward Fast Fourier Transform with spectral back-projection
     in the polar regions to manually enforce Fourier periodicity.
 
     Args:
@@ -79,24 +78,24 @@ def healpix_fft(f: np.ndarray, L: int, nside: int) -> np.ndarray:
 
         nside (int): HEALPix Nside resolution parameter.
 
+        numpy_module: JAX's Numpy-like API or Numpy. Default Numpy.
+
     Returns:
         np.ndarray: Array of Fourier coefficients for all latitudes.
-    """
-    assert L >= 2 * nside
+    '''
 
+    assert L >= 2 * nside
+    ntheta = samples.ntheta(L, "healpix", nside)
     index = 0
-    ftm = np.zeros(samples.ftm_shape(L, "healpix", nside), dtype=np.complex128)
-    ntheta = ftm.shape[0]
+    ftm_rows = []
     for t in range(ntheta):
         nphi = samples.nphi_ring(t, nside)
-        fm_chunk = fft.fftshift(fft.fft(f[index : index + nphi], norm="backward"))
-        ftm[t] = (
-            fm_chunk
-            if nphi == 2 * L
-            else spectral_periodic_extension(fm_chunk, nphi, L)
+        fm_chunk = numpy_module.fft.fftshift(
+            numpy_module.fft.fft(f[index : index + nphi], norm="backward")
         )
+        ftm_rows.append(spectral_periodic_extension(fm_chunk, L, numpy_module))
         index += nphi
-    return ftm
+    return numpy_module.stack(ftm_rows)
 
 
 def healpix_ifft(ftm: np.ndarray, L: int, nside: int) -> np.ndarray:

s2fft/samples.py

@@ -1,5 +1,6 @@
 import numpy as np
-
+import jax.numpy as jnp
+import jax.lax as lax
 
 def ntheta(L: int = None, sampling: str = "mw", nside: int = None) -> int:
     r"""Number of :math:`\theta` samples for sampling scheme at specified resolution.
@@ -332,9 +333,50 @@ def phis_ring(t: int, nside: int) -> np.ndarray:
     return p2phi_ring(t, p, nside)
 
 
-def p2phi_ring(t: int, p: int, nside: int) -> np.ndarray:
+def p2phi_ring(t: int, p: int, nside: int, numpy_module=np):
     r"""Convert index to :math:`\phi` angle for HEALPix for given :math:`\theta` ring.
 
+    See 'p2phi_ring_jax_lax' for an alternative JAX-only implementation using nested lax.cond's,
+    and 'p2phi_ring_np' for an alternative Numpy-only implementation. 
+
+    Args:
+        t (int): :math:`\theta` index of ring.
+
+        p (int): :math:`\phi` index within ring.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+        numpy_module: JAX's Numpy-like API or Numpy. Default Numpy.
+
+    Returns:
+        np.ndarray: :math:`\phi` angle.
+    """
+
+    # shift per region (2 regions)
+    shift = numpy_module.where(
+        (t + 1 >= nside) & (t + 1 <= 3 * nside), (1 / 2) * ((t - nside + 2) % 2), 1 / 2
+    )
+
+    # factor per region (3 regions)
+    factor_reg_1 = numpy_module.where(
+        (t + 1 >= nside) & (t + 1 <= 3 * nside), numpy_module.pi / (2 * nside), 1
+    )
+    factor_reg_2 = numpy_module.where(t + 1 > 3 * nside, numpy_module.pi / (2 * (4 * nside - t - 1)), 1)
+    factor_reg_3 = numpy_module.where(
+        (factor_reg_1 == 1) & (factor_reg_2 == 1), numpy_module.pi / (2 * (t + 1)), 1
+    )
+    factor = (
+        factor_reg_1 * factor_reg_2 * factor_reg_3
+    ) 
+
+    return factor * (p + shift)
+
+def p2phi_ring_np(t: int, p: int, nside: int):
+    r"""Convert index to :math:`\phi` angle for HEALPix for given :math:`\theta` ring - Numpy-only implementation.
+
+    See 'p2phi_ring_jax_lax' for an alternative JAX implementation using nested lax.cond's, and 'p2phi_ring' for a 
+    JAX/Numpy implementation using jnp.where/np.where.
+
     Args:
         t (int): :math:`\theta` index of ring.
 
@@ -350,11 +392,48 @@ def p2phi_ring(t: int, p: int, nside: int) -> np.ndarray:
     if (t + 1 >= nside) & (t + 1 <= 3 * nside):
         shift *= (t - nside + 2) % 2
         factor = np.pi / (2 * nside)
-        return factor * (p + shift)
     elif t + 1 > 3 * nside:
         factor = np.pi / (2 * (4 * nside - t - 1))
     else:
         factor = np.pi / (2 * (t + 1))
+
+    return factor * (p + shift)
+
+def p2phi_ring_jax_lax(t: int, p: int, nside: int) -> jnp.ndarray:
+    r"""Convert index to :math:`\phi` angle for HEALPix for given :math:`\theta` ring - JAX-only implementation.
+
+    Uses nested lax.cond from JAX. See 'p2phi_ring' for an alternative JAX/Numpy implementation using 
+    jnp.where/np.where, and 'p2phi_ring_np' for a Numpy-only implementation.
+
+    Args:
+        t (int): :math:`\theta` index of ring.
+
+        p (int): :math:`\phi` index within ring.
+
+        nside (int): HEALPix Nside resolution parameter.
+
+    Returns:
+        jnp.ndarray: :math:`\phi` angle.
+    """
+
+    # using nested lax.cond
+    shift = lax.cond((t + 1 >= nside) & (t + 1 <= 3 * nside),
+                    lambda t: (1 / 2) * ((t - nside + 2) % 2),
+                    lambda t: 1/2,
+                    t)
+
+    factor = lax.cond(
+        (t + 1 >= nside) & (t + 1 <= 3 * nside),
+        lambda x: jnp.pi / (2 * nside),
+        lambda x: lax.cond(
+            x + 1 > 3 * nside,
+            lambda x: jnp.pi / (2 * (4 * nside - x - 1)),
+            lambda x: jnp.pi / (2 * (x + 1)),
+            x),
+        t)
+
+
+
     return factor * (p + shift)
 
 
@@ -417,8 +496,8 @@ def p2phi_equiang(L: int, p: int, sampling: str = "mw") -> np.ndarray:
 
 
 def ring_phase_shift_hp(
-    L: int, t: int, nside: int, forward: bool = False
-) -> np.ndarray:
+    L: int, t: int, nside: int, forward: bool = False, numpy_module=np
+):
     r"""Generates a phase shift vector for HEALPix for a given :math:`\theta` ring.
 
     Args:
@@ -431,13 +510,14 @@ def ring_phase_shift_hp(
         forward (bool, optional): Whether to provide forward or inverse shift.
             Defaults to False.
 
+        numpy_module: JAX's Numpy-like API or Numpy. Default Numpy.
+
     Returns:
         np.ndarray: Vector of phase shifts with shape :math:`[2L-1]`.
     """
-    phi_offset = p2phi_ring(t, 0, nside)
+    phi_offset = p2phi_ring(t, 0, nside, numpy_module)
     sign = -1 if forward else 1
-    return np.exp(sign * 1j * np.arange(-L + 1, L) * phi_offset)
-
+    return numpy_module.exp(sign * 1j * numpy_module.arange(-L + 1, L) * phi_offset)
 
 def f_shape(L: int = None, sampling: str = "mw", nside: int = None) -> tuple:
     """Shape of spherical signal.