Custom Operations#

s2fft.precompute_transforms.custom_ops.s2_to_wigner_subset(fs: ndarray, spins: ndarray, DW: Tuple[ndarray, ndarray], L: int, sampling: str = 'mw') → ndarray#

Transforms from a collection of arbitrary spin signals on the sphere to the corresponding collection of their harmonic coefficients.

This function takes a multimodal collection of spin spherical harmonic signals on the sphere and transforms them into their spin spherical harmonic coefficients. These coefficients may then be combined into a subset of Wigner coefficients for downstream analysis. In this way one may combine input features across a variety of spins into a unified representation.

Parameters:

  • fs (np.ndarray) – Collection of spin signal maps on the sphere with shape \([batch, n_s, n_\theta,n_\phi, channels]\).

  • spins (np.ndarray) – Spins of each field in fs with shape \([n_s]\).

  • DW (Tuple[np.ndarray, np.ndarray]) – Fourier coefficients of the reduced Wigner d-functions and the corresponding upsampled quadrature weights.

  • L (int) – Harmonic band-limit.

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

Returns:

Collection of spin spherical harmonic coefficients with shape \([batch, n_s, L, 2L-1, channels]\).

Return type:

np.ndarray

s2fft.precompute_transforms.custom_ops.s2_to_wigner_subset_jax(fs: Array, spins: Array, DW: Tuple[Array, Array], L: int, sampling: str = 'mw') → Array#

Transforms from a collection of arbitrary spin signals on the sphere to the corresponding collection of their harmonic coefficients (JAX).

This function takes a multimodal collection of spin spherical harmonic signals on the sphere and transforms them into their spin spherical harmonic coefficients. These coefficients may then be combined into a subset of Wigner coefficients for downstream analysis. In this way one may combine input features across a variety of spins into a unified representation.

Parameters:

  • fs (jnp.ndarray) – Collection of spin signal maps on the sphere with shape \([batch, n_s, n_\theta,n_\phi, channels]\).

  • spins (jnp.ndarray) – Spins of each field in fs with shape \([n_s]\).

  • DW (Tuple[jnp.ndarray, jnp.ndarray]) – Fourier coefficients of the reduced Wigner d-functions and the corresponding upsampled quadrature weights.

  • L (int) – Harmonic band-limit.

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

Returns:

Collection of spin spherical harmonic coefficients with shape \([batch, n_s, L, 2L-1, channels]\).

Return type:

jnp.ndarray

s2fft.precompute_transforms.custom_ops.so3_to_wigner_subset(f: ndarray, spins: ndarray, DW: Tuple[ndarray, ndarray], L: int, N: int, sampling: str = 'mw') → ndarray#

Transforms a signal on the rotation group to an arbitrary subset of its Wigner coefficients.

This function takes a signal on the rotation group SO(3) and computes a subset of spin spherical harmonic coefficients corresponding to slices across the requested spin numbers. These spin numbers can be arbitrarily chosen such that their absolute value is less than or equal to the azimuthal band-limit \(N\leq L\).

Parameters:

  • f (np.ndarray) – Signal on the rotation group with shape \([batch, n_\gamma, n_\theta,n_\phi, channels]\).

  • spins (np.ndarray) – Spins of each field in fs with shape \([n_s]\).

  • DW (Tuple[np.ndarray, np.ndarray]) – Fourier coefficients of the reduced Wigner d-functions and the corresponding upsampled quadrature weights.

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

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

Returns:

Collection of spin spherical harmonic coefficients with shape \([batch, n_s, L, 2L-1, channels]\).

Return type:

np.ndarray

s2fft.precompute_transforms.custom_ops.so3_to_wigner_subset_jax(f: Array, spins: Array, DW: Tuple[Array, Array], L: int, N: int, sampling: str = 'mw') → Array#

Transforms a signal on the rotation group to an arbitrary subset of its Wigner coefficients (JAX).

This function takes a signal on the rotation group SO(3) and computes a subset of spin spherical harmonic coefficients corresponding to slices across the requested spin numbers. These spin numbers can be arbitrarily chosen such that their absolute value is less than or equal to the azimuthal band-limit \(N\leq L\).

Parameters:

  • f (jnp.ndarray) – Signal on the rotation group with shape \([batch, n_\gamma, n_\theta,n_\phi, channels]\).

  • spins (jnp.ndarray) – Spins of each field in fs with shape \([n_s]\).

  • DW (Tuple[jnp.ndarray, jnp.ndarray]) – Fourier coefficients of the reduced Wigner d-functions and the corresponding upsampled quadrature weights.

  • L (int) – Harmonic band-limit.

  • N (int) – Azimuthal band-limit.

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

Returns:

Collection of spin spherical harmonic coefficients

with shape \([batch, n_s, L, 2L-1, channels]\).

Return type:

jnp.ndarray

s2fft.precompute_transforms.custom_ops.wigner_subset_to_s2(flmn: ndarray, spins: ndarray, DW: Tuple[ndarray, ndarray], L: int, sampling: str = 'mw') → ndarray#

Transforms an arbitrary subset of Wigner coefficients onto a subset of spin signals on the sphere.

This function takes a collection of spin spherical harmonic coefficients each with a different (though not necessarily unique) spin and maps them back to their corresponding pixel-space representations.

Parameters:

  • flmn (np.ndarray) – Collection of spin spherical harmonic coefficients with shape \([batch, n_s, L, 2L-1, channels]\).

  • spins (np.ndarray) – Spins of each field in fs with shape \([n_s]\).

  • DW (Tuple[np.ndarray, np.ndarray]) – Fourier coefficients of the reduced Wigner d-functions and the corresponding upsampled quadrature weights.

  • L (int) – Harmonic band-limit.

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

Raises:

  • ValueError – If sampling scheme is not recognised.

  • ValueError – If the number of spins does not match the number of Wigner coefficients.

Returns:

A collection of spin signals with

shape \([batch, n_s, n_\theta, n_\phi, channels]\).

Return type:

np.ndarray

s2fft.precompute_transforms.custom_ops.wigner_subset_to_s2_jax(flmn: Array, spins: Array, DW: Tuple[Array, Array], L: int, sampling: str = 'mw') → Array#

Transforms an arbitrary subset of Wigner coefficients onto a subset of spin signals on the sphere (JAX).

This function takes a collection of spin spherical harmonic coefficients each with a different (though not necessarily unique) spin and maps them back to their corresponding pixel-space representations.

Parameters:

  • flmn (jnp.ndarray) – Collection of spin spherical harmonic coefficients with shape \([batch, n_s, L, 2L-1, channels]\).

  • spins (jnp.ndarray) – Spins of each field in fs with shape \([n_s]\).

  • DW (Tuple[jnp.ndarray, jnp.ndarray]) – Fourier coefficients of the reduced Wigner d-functions and the corresponding upsampled quadrature weights.

  • L (int) – Harmonic band-limit.

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

Raises:

  • ValueError – If sampling scheme is not recognised.

  • ValueError – If the number of spins does not match the number of Wigner coefficients.

Returns:

A collection of spin signals with shape \([batch, n_s, n_\theta, n_\phi, channels]\).

Return type:

jnp.ndarray