Spherical Harmonic Transform#

s2fft.transforms.spherical.forward(f: ndarray, L: int, spin: int = 0, nside: int | None = None, sampling: str = 'mw', method: str = 'numpy', reality: bool = False, precomps: List | None = None, spmd: bool = False, L_lower: int = 0, iter: int = 3, _ssht_backend: int = 1) → ndarray#

Wrapper for the forward spin-spherical harmonic transform.

Parameters:

  • f (np.ndarray) – Signal on the sphere.

  • L (int) – Harmonic band-limit.

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

  • nside (int, optional) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”. Defaults to None.

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

  • method (str, optional) – Execution mode in {“numpy”, “jax”, “jax_ssht”, “jax_healpy”}. Defaults to “numpy”.

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to False.

  • precomps (List[np.ndarray]) – Precomputed list of recursion coefficients. At most of length \(L^2\), which is a minimal memory overhead.

  • spmd (bool, optional) – Whether to map compute over multiple devices. Currently this only maps over all available devices, and is only valid for JAX implementations. Defaults to False.

  • L_lower (int, optional) – Harmonic lower-bound. Transform will only be computed for \(\texttt{L_lower} \leq \ell < \texttt{L}\). Defaults to 0.

  • iter (int, optional) – Number of subiterations for healpy. Note that iterations increase the precision of the forward transform, but reduce the accuracy of the gradient pass. Between 2 and 3 iterations is a good compromise.

  • _ssht_backend (int, optional, experimental) – Whether to default to SSHT core (set to 0) recursions or pick up ducc0 (set to 1) accelerated experimental backend. Use with caution.

Raises:

ValueError – Transform method not recognised.

Returns:

Spherical harmonic coefficients.

Return type:

np.ndarray

Note

The single-program multiple-data (SPMD) optional variable determines whether the transform is run over a single device or all available devices. For very low harmonic bandlimits L this is inefficient as the I/O overhead for communication between devices is noticable, however as L increases one will asymptotically recover acceleration by the number of devices.

s2fft.transforms.spherical.forward_jax(f: Array, L: int, spin: int = 0, nside: int = None, sampling: str = 'mw', reality: bool = False, precomps: List = None, spmd: bool = False, L_lower: int = 0, use_healpix_custom_primitive: bool = False) → Array#

Compute the forward spin-spherical harmonic transform (JAX).

Uses separation of variables and exploits the Price & McEwen recursion for accelerated and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this function is theoretically \(\mathcal{O}(L^2)\). This is a JAX implementation of forward_numpy().

Parameters:

  • f (jnp.ndarray) – Signal on the sphere

  • L (int) – Harmonic band-limit.

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

  • nside (int, optional) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”. Defaults to None.

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

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to False.

  • precomps (List[jnp.ndarray]) – Precomputed list of recursion coefficients. At most of length \(L^2\), which is a minimal memory overhead.

  • spmd (bool, optional) – Whether to map compute over multiple devices. Currently this only maps over all available devices. Defaults to False.

  • L_lower (int, optional) – Harmonic lower-bound. Transform will only be computed for \(\texttt{L_lower} \leq \ell < \texttt{L}\). Defaults to 0.

  • use_healpix_custom_primitive (bool, optional) – Whether to use a custom CUDA primitive for computing HEALPix fast fourier transform when sampling = “healpix” and running on a cuda compatible gpu device. using a custom primitive reduces long compilation times when jit compiling. defaults to False.

Returns:

Spherical harmonic coefficients

Return type:

jnp.ndarray

Note

The single-program multiple-data (SPMD) optional variable determines whether the transform is run over a single device or all available devices. For very low harmonic bandlimits L this is inefficient as the I/O overhead for communication between devices is noticable, however as L increases one will asymptotically recover acceleration by the number of devices.

s2fft.transforms.spherical.forward_numpy(f: ndarray, L: int, spin: int = 0, nside: int | None = None, sampling: str = 'mw', reality: bool = False, precomps: List | None = None, L_lower: int = 0) → ndarray#

Compute the forward spin-spherical harmonic transform (JAX).

Uses separation of variables and exploits the Price & McEwen recursion for accelerated and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this function is theoretically \(\mathcal{O}(L^2)\).

Parameters:

  • f (np.ndarray) – Signal on the sphere

  • L (int) – Harmonic band-limit.

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

  • nside (int, optional) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”. Defaults to None.

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

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to False.

  • precomps (List[np.ndarray]) – Precomputed list of recursion coefficients. At most of length \(L^2\), which is a minimal memory overhead.

  • L_lower (int, optional) – Harmonic lower-bound. Transform will only be computed for \(\texttt{L_lower} \leq \ell < \texttt{L}\). Defaults to 0.

Returns:

Spherical harmonic coefficients

Return type:

np.ndarray

s2fft.transforms.spherical.inverse(flm: ndarray, L: int, spin: int = 0, nside: int | None = None, sampling: str = 'mw', method: str = 'numpy', reality: bool = False, precomps: List | None = None, spmd: bool = False, L_lower: int = 0, _ssht_backend: int = 1) → ndarray#

Wrapper for the inverse spin-spherical harmonic transform.

Parameters:

  • flm (np.ndarray) – Spherical harmonic coefficients.

  • L (int) – Harmonic band-limit.

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

  • nside (int, optional) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”. Defaults to None.

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

  • method (str, optional) – Execution mode in {“numpy”, “jax”, “jax_ssht”, “jax_healpy”}. Defaults to “numpy”.

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to False.

  • precomps (List[np.ndarray]) – Precomputed list of recursion coefficients. At most of length \(L^2\), which is a minimal memory overhead.

  • spmd (bool, optional) – Whether to map compute over multiple devices. Currently this only maps over all available devices, and is only valid for JAX implementations. Defaults to False.

  • L_lower (int, optional) – Harmonic lower-bound. Transform will only be computed for \(\texttt{L_lower} \leq \ell < \texttt{L}\). Defaults to 0.

  • _ssht_backend (int, optional, experimental) – Whether to default to SSHT core (set to 0) recursions or pick up ducc0 (set to 1) accelerated experimental backend. Use with caution.

Raises:

ValueError – Transform method not recognised.

Returns:

Signal on the sphere.

Return type:

np.ndarray

Note

The single-program multiple-data (SPMD) optional variable determines whether the transform is run over a single device or all available devices. For very low harmonic bandlimits L this is inefficient as the I/O overhead for communication between devices is noticable, however as L increases one will asymptotically recover acceleration by the number of devices.

s2fft.transforms.spherical.inverse_jax(flm: Array, L: int, spin: int = 0, nside: int = None, sampling: str = 'mw', reality: bool = False, precomps: List = None, spmd: bool = False, L_lower: int = 0) → Array#

Compute the inverse spin-spherical harmonic transform (JAX).

Uses separation of variables and exploits the Price & McEwen recursion for accelerated and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this function is theoretically \(\mathcal{O}(L^2)\). This is a JAX implementation of inverse_numpy().

Parameters:

  • flm (jnp.ndarray) – Spherical harmonic coefficients.

  • L (int) – Harmonic band-limit.

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

  • nside (int, optional) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”. Defaults to None.

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

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to False.

  • precomps (List[jnp.ndarray]) – Precomputed list of recursion coefficients. At most of length \(L^2\), which is a minimal memory overhead.

  • spmd (bool, optional) – Whether to map compute over multiple devices. Currently this only maps over all available devices. Defaults to False.

  • L_lower (int, optional) – Harmonic lower-bound. Transform will only be computed for \(\texttt{L_lower} \leq \ell < \texttt{L}\). Defaults to 0.

Returns:

Signal on the sphere.

Return type:

jnp.ndarray

Note

The single-program multiple-data (SPMD) optional variable determines whether the transform is run over a single device or all available devices. For very low harmonic bandlimits L this is inefficient as the I/O overhead for communication between devices is noticable, however as L increases one will asymptotically recover acceleration by the number of devices.

s2fft.transforms.spherical.inverse_numpy(flm: ndarray, L: int, spin: int = 0, nside: int | None = None, sampling: str = 'mw', reality: bool = False, precomps: List | None = None, L_lower: int = 0) → ndarray#

Compute the inverse spin-spherical harmonic transform (numpy).

Uses separation of variables and exploits the Price & McEwen recursion for accelerated and numerically stable Wiger-d on-the-fly recursions. The memory overhead for this function is theoretically \(\mathcal{O}(L^2)\).

Parameters:

  • flm (np.ndarray) – Spherical harmonic coefficients.

  • L (int) – Harmonic band-limit.

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

  • nside (int, optional) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”. Defaults to None.

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

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. Defaults to False.

  • precomps (List[np.ndarray]) – Precomputed list of recursion coefficients. At most of length \(L^2\), which is a minimal memory overhead.

  • L_lower (int, optional) – Harmonic lower-bound. Transform will only be computed for \(\texttt{L_lower} \leq \ell < \texttt{L}\). Defaults to 0.

Returns:

Signal on the sphere.

Return type:

np.ndarray