Wigner Transform#

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

Wrapper for the forward Wigner transform, i.e. Fourier transform on \(SO(3)\).

Importantly, the convention adopted for storage of f is \([\gamma, \beta, \alpha]\), for Euler angles \((\alpha, \beta, \gamma)\) following the \(zyz\) Euler convention, in order to simplify indexing for internal use. For a given \(\gamma\) we thus recover a signal on the sphere indexed by \([\theta, \phi]\), i.e. we associate \(\beta\) with \(\theta\) and \(\alpha\) with \(\phi\).

Should the user select method = “jax_ssht” they will be restricted to deployment on CPU using our custom JAX frontend for the SSHT C library [1]. In many cases this approach may be desirable to mitigate e.g. memory i/o cost.

Parameters:

  • f (np.ndarray) – Signal on the on \(SO(3)\) with shape \([n_{\gamma}, n_{\beta}, n_{\alpha}]\), where \(n_\xi\) denotes the number of samples for angle \(\xi\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

  • 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”}. 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.

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

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

Return type:

np.ndarray

Note

[1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”

IEEE Transactions on Signal Processing 59 (2011): 5876-5887.

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

Compute the forward Wigner 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}(NL^2)\).

Importantly, the convention adopted for storage of f is \([\gamma, \beta, \alpha]\), for Euler angles \((\alpha, \beta, \gamma)\) following the \(zyz\) Euler convention, in order to simplify indexing for internal use. For a given \(\gamma\) we thus recover a signal on the sphere indexed by \([\theta, \phi]\), i.e. we associate \(\beta\) with \(\theta\) and \(\alpha\) with \(\phi\).

Parameters:

  • f (jnp.ndarray) – Signal on the on \(SO(3)\) with shape \([n_{\gamma}, n_{\beta}, n_{\alpha}]\), where \(n_\xi\) denotes the number of samples for angle \(\xi\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

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

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

Returns:

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

Return type:

jnp.ndarray

s2fft.transforms.wigner.forward_jax_ssht(f: Array, L: int, N: int, L_lower: int = 0, sampling: str = 'mw', reality: bool = False, _ssht_backend: int = 1) → Array#

Compute the forward Wigner transform (SSHT JAX).

SSHT is a C library which implements the spin-spherical harmonic transform outlined in McEwen & Wiaux 2011 [1]. We make use of their python bindings for which we provide custom JAX frontends, hence providing support for automatic differentiation. Currently these transforms can only be deployed on CPU, which is a limitation of the SSHT C package.

Parameters:

  • f (jnp.ndarray) – Signal on the on \(SO(3)\) with shape \([n_{\gamma}, n_{\beta}, n_{\alpha}]\), where \(n_\xi\) denotes the number of samples for angle \(\xi\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

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

  • sampling (str, optional) – Sampling scheme. Supported sampling schemes include {“mw”, “mwss”, “dh”, “gl”}. 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.

  • _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.

Returns:

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

Return type:

jnp.ndarray

Note

[1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”

IEEE Transactions on Signal Processing 59 (2011): 5876-5887.

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

Compute the forward Wigner 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}(NL^2)\).

Importantly, the convention adopted for storage of f is \([\gamma, \beta, \alpha]\), for Euler angles \((\alpha, \beta, \gamma)\) following the \(zyz\) Euler convention, in order to simplify indexing for internal use. For a given \(\gamma\) we thus recover a signal on the sphere indexed by \([\theta, \phi]\), i.e. we associate \(\beta\) with \(\theta\) and \(\alpha\) with \(\phi\).

Parameters:

  • f (np.ndarray) – Signal on the on \(SO(3)\) with shape \([n_{\gamma}, n_{\beta}, n_{\alpha}]\), where \(n_\xi\) denotes the number of samples for angle \(\xi\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

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

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

Return type:

np.ndarray

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

Wrapper for the inverse Wigner transform, i.e. inverse Fourier transform on \(SO(3)\).

Importantly, the convention adopted for storage of f is \([\gamma, \beta, \alpha]\), for Euler angles \((\alpha, \beta, \gamma)\) following the \(zyz\) Euler convention, in order to simplify indexing for internal use. For a given \(\gamma\) we thus recover a signal on the sphere indexed by \([\theta, \phi]\), i.e. we associate \(\beta\) with \(\theta\) and \(\alpha\) with \(\phi\).

Should the user select method = “jax_ssht” they will be restricted to deployment on CPU using our custom JAX frontend for the SSHT C library [1]. In many cases this approach may be desirable to mitigate e.g. memory i/o cost.

Parameters:

  • flmn (np.ndarray) – Wigner coefficients with shape \([2N-1, L, 2L-1]\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit. For high precision beyond \(N \approx 8\), one should use method=”jax_ssht”.

  • 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”}. 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.

  • 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 on \(SO(3)\) with shape \([n_{\gamma}, n_{\beta}, n_{\alpha}]\), where \(n_\xi\) denotes the number of samples for angle \(\xi\).

Return type:

np.ndarray

Note

[1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”

IEEE Transactions on Signal Processing 59 (2011): 5876-5887.

s2fft.transforms.wigner.inverse_jax(flmn: Array, L: int, N: int, nside: int = None, sampling: str = 'mw', reality: bool = False, precomps: List = None, L_lower: int = 0) → Array#

Compute the inverse Wigner 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}(NL^2)\).

Importantly, the convention adopted for storage of f is \([\gamma, \beta, \alpha]\), for Euler angles \((\alpha, \beta, \gamma)\) following the \(zyz\) Euler convention, in order to simplify indexing for internal use. For a given \(\gamma\) we thus recover a signal on the sphere indexed by \([\theta, \phi]\), i.e. we associate \(\beta\) with \(\theta\) and \(\alpha\) with \(\phi\).

Parameters:

  • flmn (jnp.ndarray) – Wigner coefficients with shape \([2N-1, L, 2L-1]\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit. Recursive transform will have lower precision beyond \(N \approx 8\).

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

  • 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

s2fft.transforms.wigner.inverse_jax_ssht(flmn: Array, L: int, N: int, L_lower: int = 0, sampling: str = 'mw', reality: bool = False, _ssht_backend: int = 1) → Array#

Compute the inverse Wigner transform (SSHT JAX).

SSHT is a C library which implements the spin-spherical harmonic transform outlined in McEwen & Wiaux 2011 [1]. We make use of their python bindings for which we provide custom JAX frontends, hence providing support for automatic differentiation. Currently these transforms can only be deployed on CPU, which is a limitation of the SSHT C package.

Parameters:

  • flmn (jnp.ndarray) – Wigner coefficients with shape \([2N-1, L, 2L-1]\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

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

  • sampling (str, optional) – Sampling scheme. Supported sampling schemes include {“mw”, “mwss”, “dh”, “gl”}. 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.

  • _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.

Returns:

Signal on the sphere.

Return type:

np.ndarray

Note

[1] McEwen, Jason D. and Yves Wiaux. “A Novel Sampling Theorem on the Sphere.”

IEEE Transactions on Signal Processing 59 (2011): 5876-5887.

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

Compute the inverse Wigner 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}(NL^2)\).

Importantly, the convention adopted for storage of f is \([\gamma, \beta, \alpha]\), for Euler angles \((\alpha, \beta, \gamma)\) following the \(zyz\) Euler convention, in order to simplify indexing for internal use. For a given \(\gamma\) we thus recover a signal on the sphere indexed by \([\theta, \phi]\), i.e. we associate \(\beta\) with \(\theta\) and \(\alpha\) with \(\phi\).

Parameters:

  • flmn (np.ndarray) – Wigner coefficients with shape \([2N-1, L, 2L-1]\).

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit. Recursive transform will have lower precision beyond \(N \approx 8\).

  • 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