Transforms#

Spin spherical harmonic and Wigner transforms, through s2fft, with a choice of backend.

The backends are:

  • "jax": s2fft’s JAX transforms with on-the-fly Wigner recursions. They are JIT compiled, differentiable and run on GPUs, with memory \(\mathcal{O}(L^2)\), but their recursions are accurate only for spins and azimuthal band-limits below RECURSION_LIMIT.

  • "precompute": s2fft’s precompute transforms. The Wigner functions are computed once, by a recursion that is stable at every spin, and cached, so that the transforms are accurate for any spin and azimuthal band-limit at the price of \(\mathcal{O}(L^3)\) (spherical) or \(\mathcal{O}(NL^3)\) (Wigner) memory. They are JIT compiled, differentiable and run on GPUs.

  • "jax_ssht": s2fft’s differentiable JAX frontend to the SSHT C library, which is accurate at every spin but runs on the CPU only (requires pyssht and ducc0).

  • "numpy": s2fft’s NumPy reference transforms.

All transforms are exact, up to floating-point error, for band-limited signals. All backends except "numpy" support reverse-mode differentiation. Forward mode (jax.jvp) is supported by "precompute", and by "jax" with s2fft 1.5.0 or later (the JAX transforms of s2fft 1.4.0 define only a custom VJP); "jax_ssht" supports reverse mode only.

s2conv.transforms.METHODS = ('jax', 'precompute', 'jax_ssht', 'numpy')#

s2fft’s JAX, precompute, SSHT and NumPy transforms.

Type:

Transform backends

s2conv.transforms.RECURSION_LIMIT = 8#

Spins and azimuthal band-limits from which the recursive backends lose accuracy.

s2conv.transforms.spherical_kernel(L: int, spin: int = 0, forward: bool = True, sampling: str = 'mw', reality: bool = False, nside: int = None)[source]#

Cached kernel of the precomputed spin spherical harmonic transform.

Parameters:

  • L (int) – Harmonic band-limit.

  • spin (int, optional) – Spin. Defaults to 0.

  • forward (bool, optional) – Kernel of the forward (True) or inverse transform. Defaults to True.

  • sampling (str, optional) – Sampling scheme. Defaults to “mw”.

  • reality (bool, optional) – Whether the signal is real. Defaults to False.

  • nside (int, optional) – HEALPix resolution, if sampling="healpix".

Returns:

The kernel, of size \(\mathcal{O}(L^3)\).

Return type:

jnp.ndarray

s2conv.transforms.wigner_kernel(L: int, N: int, forward: bool = True, sampling: str = 'mw', reality: bool = False, nside: int = None)[source]#

Cached kernel of the precomputed Wigner transform.

Parameters:

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

  • forward (bool, optional) – Kernel of the forward (True) or inverse transform. Defaults to True.

  • sampling (str, optional) – Sampling scheme. Defaults to “mw”.

  • reality (bool, optional) – Whether the signal is real. Defaults to False.

  • nside (int, optional) – HEALPix resolution, if sampling="healpix".

Returns:

The kernel, of size \(\mathcal{O}(NL^3)\).

Return type:

jnp.ndarray

s2conv.transforms.forward(f, L: int, spin: int = 0, sampling: str = 'mw', method: str = 'jax', reality: bool = False, nside: int = None)[source]#

Forward spin spherical harmonic transform.

Parameters:

  • f – Samples of a spin-\(s\) field on the sphere.

  • L (int) – Harmonic band-limit.

  • spin (int, optional) – Spin. Defaults to 0.

  • sampling (str, optional) – Sampling scheme, one of {“mw”, “mwss”, “dh”, “gl”, “healpix”}. Defaults to “mw”.

  • method (str, optional) – Backend, one of METHODS. Defaults to “jax”.

  • reality (bool, optional) – Whether the signal is real. Defaults to False.

  • nside (int, optional) – HEALPix resolution, if sampling="healpix".

Returns:

Harmonic coefficients flm with shape [L, 2L - 1].

s2conv.transforms.inverse(flm, L: int, spin: int = 0, sampling: str = 'mw', method: str = 'jax', reality: bool = False, nside: int = None)[source]#

Inverse spin spherical harmonic transform.

Parameters:

  • flm – Harmonic coefficients with shape [L, 2L - 1].

  • L (int) – Harmonic band-limit.

  • spin (int, optional) – Spin. Defaults to 0.

  • sampling (str, optional) – Sampling scheme. Defaults to “mw”.

  • method (str, optional) – Backend, one of METHODS. Defaults to “jax”.

  • reality (bool, optional) – Whether the signal is real. Defaults to False.

  • nside (int, optional) – HEALPix resolution, if sampling="healpix".

Returns:

Samples of the spin-\(s\) field on the sphere.

s2conv.transforms.wigner_forward(f, L: int, N: int, sampling: str = 'mw', method: str = 'jax', reality: bool = False, nside: int = None)[source]#

Forward Wigner transform, \(F^{\ell}_{mn} = \int F\, D^{\ell}_{mn}\,\mathrm{d}\mu\) (Eq. 7).

Parameters:

  • f – Samples on the rotation group with shape [2N - 1, ...], the first axis indexing the third Euler angle \(\gamma_j = 2\pi j/(2N-1)\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

  • sampling (str, optional) – Sampling scheme. Defaults to “mw”.

  • method (str, optional) – Backend, one of METHODS. Defaults to “jax”.

  • reality (bool, optional) – Whether the signal is real. Defaults to False.

  • nside (int, optional) – HEALPix resolution, if sampling="healpix".

Returns:

Wigner coefficients flmn with shape [2N - 1, L, 2L - 1].

s2conv.transforms.wigner_inverse(flmn, L: int, N: int, sampling: str = 'mw', method: str = 'jax', reality: bool = False, nside: int = None)[source]#

Inverse Wigner transform, \(F = \sum_\ell \frac{2\ell+1}{8\pi^2} \sum_{m,n} F^{\ell}_{mn} D^{\ell *}_{mn}\) (Eq. 7).

Parameters:

  • flmn – Wigner coefficients with shape [2N - 1, L, 2L - 1].

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

  • sampling (str, optional) – Sampling scheme. Defaults to “mw”.

  • method (str, optional) – Backend, one of METHODS. Defaults to “jax”.

  • reality (bool, optional) – Whether the signal is real. Defaults to False.

  • nside (int, optional) – HEALPix resolution, if sampling="healpix".

Returns:

Samples on the rotation group with shape [2N - 1, ...].

s2conv.transforms.sphere_angles(L: int, sampling: str = 'mw')[source]#

Colatitudes and longitudes of an equiangular sampling of the sphere.

Parameters:

  • L (int) – Harmonic band-limit.

  • sampling (str, optional) – Equiangular sampling scheme, one of {“mw”, “mwss”, “dh”, “gl”}. Defaults to “mw”.

Returns:

Arrays of colatitudes and longitudes.

Return type:

tuple