s2fft is currently in an open beta, please provide feedback on GitHub

Wigner Transform#

s2fft.precompute_transforms.wigner.forward(f: ndarray, L: int, N: int, kernel: ndarray | None = None, sampling: str = 'mw', reality: bool = False, method: str = 'jax', nside: int | None = None) ndarray#

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

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) – Directional band-limit.

  • kernel (np.ndarray, optional) – Wigner-d kernel. 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.

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

Returns:

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

Return type:

np.ndarray

Note

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\).

s2fft.precompute_transforms.wigner.forward_transform(f: ndarray, kernel: ndarray, L: int, N: int, sampling: str, reality: bool, nside: int) ndarray#

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

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

  • kernel (np.ndarray) – Wigner-d kernel.

  • L (int) – Harmonic band-limit.

  • N (int) – Directional band-limit.

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

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

Returns:

Wigner space coefficients.

Return type:

np.ndarray

s2fft.precompute_transforms.wigner.forward_transform_jax(f: Array, kernel: Array, L: int, N: int, sampling: str, reality: bool, nside: int) Array#

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

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

  • kernel (jnp.ndarray) – Wigner-d kernel.

  • L (int) – Harmonic band-limit.

  • N (int) – Directional band-limit.

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

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

Returns:

Wigner space coefficients.

Return type:

jnp.ndarray

s2fft.precompute_transforms.wigner.forward_transform_torch(f: tensor, kernel: tensor, L: int, N: int, sampling: str, reality: bool, nside: int) tensor#

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

Parameters:
  • f (torch.tensor) – Signal on the sphere.

  • kernel (torch.tensor) – Wigner-d kernel.

  • L (int) – Harmonic band-limit.

  • N (int) – Directional band-limit.

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

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

Returns:

Wigner space coefficients.

Return type:

torch.tensor

s2fft.precompute_transforms.wigner.inverse(flmn: ndarray, L: int, N: int, kernel: ndarray | None = None, sampling: str = 'mw', reality: bool = False, method: str = 'jax', nside: int | None = None) ndarray#

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

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

  • L (int) – Harmonic band-limit.

  • N (int) – Directional band-limit.

  • kernel (np.ndarray, optional) – Wigner-d kernel. 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.

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

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

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\).

s2fft.precompute_transforms.wigner.inverse_transform(flmn: ndarray, kernel: ndarray, L: int, N: int, sampling: str, reality: bool, nside: int) ndarray#

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

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

  • kernel (np.ndarray) – Wigner-d kernel.

  • L (int) – Harmonic band-limit.

  • N (int) – Directional band-limit.

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

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

Returns:

Pixel-space coefficients.

Return type:

np.ndarray

s2fft.precompute_transforms.wigner.inverse_transform_jax(flmn: Array, kernel: Array, L: int, N: int, sampling: str, reality: bool, nside: int) Array#

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

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

  • kernel (jnp.ndarray) – Wigner-d kernel.

  • L (int) – Harmonic band-limit.

  • N (int) – Directional band-limit.

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

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

Returns:

Pixel-space coefficients.

Return type:

jnp.ndarray

s2fft.precompute_transforms.wigner.inverse_transform_torch(flmn: tensor, kernel: tensor, L: int, N: int, sampling: str, reality: bool, nside: int) tensor#

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

Parameters:
  • flmn (torch.tensor) – Wigner coefficients with shape \([2N-1, L, 2L-1]\).

  • kernel (torch.tensor) – Wigner-d kernel.

  • L (int) – Harmonic band-limit.

  • N (int) – Directional band-limit.

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

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

  • nside (int) – HEALPix Nside resolution parameter. Only required if sampling=”healpix”.

Returns:

Pixel-space coefficients.

Return type:

torch.tensor