Merge pull request #76 from astro-informatics/feature/precompute_tran…

…sform

Feature/precompute transform

s2wav/transforms/__init__.py

@@ -1,4 +1,3 @@
 from . import numpy_wavelets
-from . import numpy_scattering
 from . import jax_wavelets
-from . import jax_scattering
+from . import jax_wavelets_precompute

s2wav/transforms/jax_wavelets.py

@@ -63,7 +63,7 @@ def generate_wigner_precomputes(
     return precomps
 
 
-@partial(jit, static_argnums=(2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13))
+# @partial(jit, static_argnums=(2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13))
 def synthesis(
     f_wav: jnp.ndarray,
     f_scal: jnp.ndarray,
@@ -177,7 +177,7 @@ def synthesis(
     return s2fft.inverse_jax(flm, L, spin, nside, sampling, reality)
 
 
-@partial(jit, static_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12))
+# @partial(jit, static_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 12))
 def analysis(
     f: jnp.ndarray,
     L: int,
@@ -193,7 +193,6 @@ def analysis(
     filters: Tuple[jnp.ndarray] = None,
     spmd: bool = False,
     precomps: List[List[jnp.ndarray]] = None,
-    scattering: bool = False,
 ) -> Tuple[jnp.ndarray]:
     r"""Wavelet analysis from pixel space to wavelet space for complex signals.
 
@@ -233,9 +232,6 @@ def analysis(
         precomps (List[jnp.ndarray]): Precomputed list of recursion coefficients. At most
             of length :math:`L^2`, which is a minimal memory overhead.
 
-        scattering (bool, optional): If using for scattering transform return absolute value
-            of scattering coefficients.
-
     Returns:
         f_wav (jnp.ndarray): Array of wavelet pixel-space coefficients
             with shape :math:`[n_{J}, 2N-1, n_{\theta}, n_{\phi}]`.
@@ -293,8 +289,6 @@ def analysis(
             spmd_iter,
             L0j,
         )
-        if scattering:
-            f_wav[j - J_min] = jnp.abs(f_wav[j - J_min])
 
     # Project all harmonic coefficients for each lm onto scaling coefficients
     phi = filters[1][:Ls] * jnp.sqrt(4 * jnp.pi / (2 * jnp.arange(Ls) + 1))
@@ -306,7 +300,7 @@ def analysis(
         f_scal = temp * jnp.sqrt(1 / (4 * jnp.pi))
     else:
         f_scal = s2fft.inverse_jax(temp, Ls, spin, nside, sampling, reality)
-    return f_wav, jnp.real(f_scal) if reality else f_scal
+    return f_wav, f_scal
 
 
 @partial(jit, static_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9, 11))