quadrature_jax

s2fft.utils.quadrature_jax.mw_weights(m: int) → float

Compute MW weights given as a function of index m.

MW weights are defined by

\[w(m^\prime) = \int_0^\pi \text{d} \theta \sin \theta \exp(i m^\prime\theta),\]

which can be computed analytically.

Parameters:

m (int) – Harmonic weight index.

Returns:

MW weight.

Return type:

float

s2fft.utils.quadrature_jax.quad_weight_dh_theta_only(theta: float, L: int) → float

Compute DH quadrature weight for \(\theta\) integration (only), for given \(\theta\). JAX implementation of s2fft.quadrature.quad_weights_dh_theta_only().

Parameters:

• theta (float) – \(\theta\) angle for which to compute weight.

• L (int) – Harmonic band-limit.

Returns:

Weight computed for each \(\theta\).

Return type:

float

s2fft.utils.quadrature_jax.quad_weights(L: int = None, sampling: str = 'mw', nside: int = None) → Array

Compute quadrature weights for \(\theta\) and \(\phi\) integration for various sampling schemes. JAX implementation of quad_weights().

Parameters:

• L (int, optional) – Harmonic band-limit. Required if sampling not healpix. Defaults to None.

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

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

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

Raises:

ValueError – Invalid sampling scheme.

Returns:

Quadrature weights for sampling scheme for each \(\theta\) (weights are identical as \(\phi\) varies for given \(\theta\)).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_dh(L: int) → Array

Compute DH quadrature weights for \(\theta\) and \(\phi\) integration. JAX implementation of s2fft.quadrature.quad_weights_dh().

Parameters:

L (int) – Harmonic band-limit.

Returns:

Weights computed for each \(\theta\) (weights are identical as \(\phi\) varies for given \(\theta\)).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_gl(L: int) → Array

Compute GL quadrature weights for \(\theta\) and \(\phi\) integration.

Parameters:

L (int) – Harmonic band-limit.

Returns:

Weights computed for each \(\theta\) (weights are identical as \(\phi\) varies for given \(\theta\)).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_hp(nside: int) → Array

Compute HEALPix quadrature weights for \(\theta\) and \(\phi\) integration. JAX implementation of s2fft.quadrature.quad_weights_hp().

Note

HEALPix weights are identical for all pixels. Nevertheless, an array of weights is returned (with identical values) for consistency of interface across other sampling schemes.

Parameters:

nside (int) – HEALPix Nside resolution parameter.

Returns:

Weights computed for each \(\theta\) (all weights in array are identical).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_mw(L: int) → Array

Compute MW quadrature weights for \(\theta\) and \(\phi\) integration. JAX implementation of s2fft.quadrature.quad_weights_mw().

Parameters:

• L (int) – Harmonic band-limit.

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

Returns:

Weights computed for each \(\theta\) (weights are identical as \(\phi\) varies for given \(\theta\)).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_mw_theta_only(L: int) → Array

Compute MW quadrature weights for \(\theta\) integration (only). JAX implementation of s2fft.quadrature.quad_weights_mw_theta_only().

Parameters:

• L (int) – Harmonic band-limit.

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

Returns:

Weights computed for each \(\theta\).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_mwss(L: int) → Array

Compute MWSS quadrature weights for \(\theta\) and \(\phi\) integration. JAX implementation of s2fft.quadrature.quad_weights_mwss().

Parameters:

• L (int) – Harmonic band-limit.

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

Returns:

Weights computed for each \(\theta\) (weights are identical as \(\phi\) varies for given \(\theta\)).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_mwss_theta_only(L: int) → Array

Compute MWSS quadrature weights for \(\theta\) integration (only). JAX implementation of s2fft.quadrature.quad_weights_mwss_theta_only().

Parameters:

• L (int) – Harmonic band-limit.

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

Returns:

Weights computed for each \(\theta\).

Return type:

jnp.ndarray

s2fft.utils.quadrature_jax.quad_weights_transform(L: int, sampling: str = 'mwss', nside: int = 0) → Array

Compute quadrature weights for \(\theta\) and \(\phi\) integration to use in transform for various sampling schemes. JAX implementation of quad_weights_transform().

Quadrature weights to use in transform for MWSS correspond to quadrature weights are twice the base resolution, i.e. 2 * L.

Parameters:

• L (int) – Harmonic band-limit.

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

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

Raises:

ValueError – Invalid sampling scheme.

Returns:

Quadrature weights to use in transform for sampling scheme for each \(\theta\) (weights are identical as \(\phi\) varies for given \(\theta\)).

Return type:

jnp.ndarray