Price & McEwen#

s2fft.recursions.price_mcewen.compute_all_slices(beta: ndarray, L: int, spin: int, precomps=None) → ndarray#

Compute a particular slice \(m^{\prime}\), denoted mm, of the complete Wigner-d matrix for all sampled polar angles \(\beta\) and all \(\ell\) using Price & McEwen recursion.

The Wigner-d slice for all \(\ell\) (el) and \(\beta\) is computed recursively over \(m\) labelled ‘m’ at a specific \(m^{\prime}\). The Price & McEwen recursion is analytically correct from \(-\ell < m < \ell\) however numerically it can become unstable for \(m > 0\). To avoid this we compute \(d_{m, m^{\prime}}^{\ell}(\beta)\) for negative \(m\) and then evaluate \(d_{m, -m^{\prime}}^{\ell}(\beta) = (-1)^{m-m^{\prime}} d_{-m, m^{\prime}}^{\ell}(\beta)\) which we can again evaluate using the same recursion.

On-the-fly renormalisation is implemented to avoid potential over/under-flows, within any given iteration of the recursion the iterants are \(\sim \mathcal{O}(1)\).

The Wigner-d slice \(d^\ell_{m, m^{\prime}}(\beta)\) is indexed for \(-L < m < L\) by dl[L - 1 - m, beta, ell]. This implementation has computational scaling \(\mathcal{O}(L)\) and typically requires \(\sim 2L\) operations.

Parameters:

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

  • L (int) – Harmonic band-limit.

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

  • precomps (List[np.ndarray]) – Precomputed recursion coefficients with memory overhead \(\mathcal{O}(L^2)\), which is minimal.

Returns:

Wigner-d matrix mm slice of dimension \([2L-1, n_{\theta}, n_{\ell}]\).

Return type:

np.ndarray

s2fft.recursions.price_mcewen.compute_all_slices_jax(beta: Array, L: int, spin: int, sampling: str = 'mw', forward: bool = False, nside: int = None, precomps=None) → Array#

Compute a particular slice \(m^{\prime}\), denoted mm, of the complete Wigner-d matrix for all sampled polar angles \(\beta\) and all \(\ell\) using Price & McEwen recursion.

The Wigner-d slice for all \(\ell\) (el) and \(\beta\) is computed recursively over \(m\) labelled ‘m’ at a specific \(m^{\prime}\). The Price & McEwen recursion is analytically correct from \(-\ell < m < \ell\) however numerically it can become unstable for \(m > 0\). To avoid this we compute \(d_{m, m^{\prime}}^{\ell}(\beta)\) for negative \(m\) and then evaluate \(d_{m, -m^{\prime}}^{\ell}(\beta) = (-1)^{m-m^{\prime}} d_{-m, m^{\prime}}^{\ell}(\beta)\) which we can again evaluate using the same recursion.

On-the-fly renormalisation is implemented to avoid potential over/under-flows, within any given iteration of the recursion the iterants are \(\sim \mathcal{O}(1)\).

The Wigner-d slice \(d^\ell_{m, m^{\prime}}(\beta)\) is indexed for \(-L < m < L\) by dl[L - 1 - m, beta, ell]. This implementation has computational scaling \(\mathcal{O}(L)\) and typically requires \(\sim 2L\) operations.

Parameters:

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

  • L (int) – Harmonic band-limit.

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

  • precomps (List[np.ndarray]) – Precomputed recursion coefficients with memory overhead \(\mathcal{O}(L^2)\), which is minimal.

Returns:

Wigner-d matrix mm slice of dimension \([2L-1, n_{\theta}, n_{\ell}]\).

Return type:

jnp.ndarray

s2fft.recursions.price_mcewen.generate_precomputes(L: int, spin: int = 0, sampling: str = 'mw', nside: int | None = None, forward: bool = False, L_lower: int = 0) → List[ndarray]#

Compute recursion coefficients with \(\mathcal{O}(L^3)\) memory overhead.

The corresponding JAX implementation in generate_precomputes_jax() has a \(\mathcal{O}(L^2)\) overhead and so should generally be preferred, particularly for large \(L\).

Parameters:

  • L (int) – Harmonic band-limit.

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

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

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

  • forward (bool, optional) – Whether to provide forward or inverse shift. 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:

List of precomputed coefficient arrays.

Return type:

List[np.ndarray]

Note

TODO: this function should be optimised.

s2fft.recursions.price_mcewen.generate_precomputes_jax(L: int, spin: int = 0, sampling: str = 'mw', nside: int = None, forward: bool = False, L_lower: int = 0, betas: Array = None) → List[Array]#

Compute recursion coefficients with \(\mathcal{O}(L^2)\) memory overhead. In practice one could compute these on-the-fly but the memory overhead is negligible and well worth the acceleration. JAX implementation of generate_precomputes().

Parameters:

  • L (int) – Harmonic band-limit.

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

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

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

  • forward (bool, optional) – Whether to provide forward or inverse shift. 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.

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

Returns:

List of precomputed coefficient arrays.

Return type:

List[jnp.ndarray]

s2fft.recursions.price_mcewen.generate_precomputes_wigner(L: int, N: int, sampling: str = 'mw', nside: int | None = None, forward: bool = False, reality: bool = False, L_lower: int = 0) → List[List[ndarray]]#

Compute recursion coefficients with \(\mathcal{O}(L^2)\) memory overhead. In practice one could compute these on-the-fly but the memory overhead is negligible and well worth the acceleration. This is a wrapped extension of generate_precomputes() for the case of multiple spins, i.e. the Wigner transform over SO(3).

Parameters:

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal bandlimit

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

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

  • forward (bool, optional) – Whether to provide forward or inverse shift. Defaults to False.

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. 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:

2N-1 length List of Lists of precomputed coefficient arrays.

Return type:

List[List[np.ndarray]]

Note

TODO: this function should be optimised.

s2fft.recursions.price_mcewen.generate_precomputes_wigner_jax(L: int, N: int, sampling: str = 'mw', nside: int = None, forward: bool = False, reality: bool = False, L_lower: int = 0) → List[List[Array]]#

Compute recursion coefficients with \(\mathcal{O}(L^2)\) memory overhead. In practice one could compute these on-the-fly but the memory overhead is negligible and well worth the acceleration. This is a wrapped extension of generate_precomputes() for the case of multiple spins, i.e. the Wigner transform over SO(3). JAX implementation of generate_precomputes_wigner().

Parameters:

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal bandlimit

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

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

  • forward (bool, optional) – Whether to provide forward or inverse shift. Defaults to False.

  • reality (bool, optional) – Whether the signal on the sphere is real. If so, conjugate symmetry is exploited to reduce computational costs. 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:

2N-1 length List of Lists of precomputed coefficient arrays.

Return type:

List[List[jnp.ndarray]]