Spherical Harmonic Transform — S2FFT 0.1.dev1 documentation

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 | None = None, _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_cuda”, 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 iterative refinement iterations to use to improve accuracy of forward transform (as an inverse of inverse transform). Primarily of use with HEALPix sampling for which there is not a sampling theorem, and round-tripping through the forward and inverse transforms will introduce an error. If set to None, the default, 3 iterations will be used if sampling == "healpix" and method == "jax_healpy" and zero otherwise. Not used for jax_ssht method.
_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 just-in-time 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, sampling: str = 'mw', reality: bool = False, precomps: list = 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.forward_torch(f: None, 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) → None#

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

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 Torch implementation of forward_numpy().

Parameters:

f (torch.Tensor) – 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[torch.Tensor]) – 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 just-in-time compiling. Defaults to False.

Returns:

Spherical harmonic coefficients

Return type:

torch.Tensor

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(flm: ndarray, L: int, spin: int = 0, nside: int = None, sampling: str = 'mw', method: str = 'numpy', reality: bool = False, precomps: list = 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_cuda”, “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, use_healpix_custom_primitive: bool = False) → 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.
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 just-in-time compiling. Defaults to False.

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, sampling: str = 'mw', reality: bool = False, precomps: list = 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

s2fft.transforms.spherical.inverse_torch(flm: None, 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) → None#

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

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 Torch implementation of inverse_numpy().

Parameters:

flm (torch.Tensor) – 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[torch.Tensor]) – 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 just-in-time compiling. Defaults to False.

Returns:

Signal on the sphere.

Return type:

torch.Tensor

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.

Spherical Harmonic Transform#