On-the-fly (OTF) recursions#

s2fft.transforms.otf_recursions.forward_latitudinal_step(ftm: ndarray, beta: ndarray, L: int, spin: int, nside: int, sampling: str = 'mw', reality: bool = False, precomps: list = None, L_lower: int = 0) → ndarray#

Evaluate the wigner-d recursion forward latitundinal step over \(\theta\). This approach is a heavily engineerd version of the Price & McEwen recursion found in price_mcewen(), which has at most of \(\mathcal{O}(L^2)\) memory footprint.

This latitundinal \(\theta\) step for scalar fields reduces to the associated Legendre transform, however our transform supports arbitrary spin \(s < L\). By construction the Price & McEwen approach recurses over m solely, hence though one must recurse \(\sim L\) times, all \(\theta, \ell\) entries can be computed simultaneously; facilitating GPU/TPU acceleration.

Parameters:

  • ftm (np.ndarray) – Intermediate coefficients with indexing \([\theta, m]\).

  • beta (np.ndarray) – Array of polar angles in radians.

  • 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.otf_recursions.forward_latitudinal_step_jax(ftm_in: Array, beta_in: Array, L: int, spin: int, nside: int, sampling: str = 'mw', reality: bool = False, precomps: list = None, spmd: bool = False, L_lower: int = 0) → Array#

Evaluate the wigner-d recursion forward latitundinal step over \(\theta\). This approach is a heavily engineerd version of the Price & McEwen recursion found in price_mcewen(), which has at most of \(\mathcal{O}(L^2)\) memory footprint. This is a JAX implementation of forward_latitudinal_step().

This latitundinal \(\theta\) step for scalar fields reduces to the associated Legendre transform, however our transform supports arbitrary spin \(s < L\). By construction the Price & McEwen approach recurses over m solely, hence though one must recurse \(\sim L\) times, all \(\theta, \ell\) entries can be computed simultaneously; facilitating GPU/TPU acceleration.

Parameters:

  • ftm (jnp.ndarray) – Intermediate coefficients with indexing \([\theta, m]\).

  • beta (jnp.ndarray) – Array of polar angles in radians.

  • 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:

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.otf_recursions.inverse_latitudinal_step(flm: ndarray, beta: ndarray, L: int, spin: int, nside: int, sampling: str = 'mw', reality: bool = False, precomps: list = None, L_lower: int = 0) → ndarray#

Evaluate the wigner-d recursion inverse latitundinal step over \(\theta\). This approach is a heavily engineerd version of the Price & McEwen recursion found in price_mcewen(), which has at most of \(\mathcal{O}(L^2)\) memory footprint.

This latitundinal \(\theta\) step for scalar fields reduces to the associated Legendre transform, however our transform supports arbitrary spin \(s < L\). By construction the Price & McEwen approach recurses over m solely, hence though one must recurse \(\sim L\) times, all \(\theta, \ell\) entries can be computed simultaneously; facilitating GPU/TPU acceleration.

Parameters:

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

  • beta (np.ndarray) – Array of polar angles in radians.

  • 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:

Coefficients ftm with indexing \([\theta, m]\).

Return type:

np.ndarray

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

Evaluate the wigner-d recursion inverse latitundinal step over \(\theta\). This approach is a heavily engineerd version of the Price & McEwen recursion found in price_mcewen(), which has at most of \(\mathcal{O}(L^2)\) memory footprint. This is a JAX implementation of inverse_latitudinal_step().

This latitundinal \(\theta\) step for scalar fields reduces to the associated Legendre transform, however our transform supports arbitrary spin \(s < L\). By construction the Price & McEwen approach recurses over m solely, hence though one must recurse \(\sim L\) times, all \(\theta, \ell\) entries can be computed simultaneously; facilitating GPU/TPU acceleration.

Parameters:

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

  • beta (jnp.ndarray) – Array of polar angles in radians.

  • 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], optional) – 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:

Coefficients ftm with indexing \([\theta, m]\).

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.