healpix functions

s2fft.utils.healpix_ffts.healpix_fft(f: ndarray, L: int, nside: int, method: str = 'numpy', reality: bool = False) → ndarray

Wrapper function for the Forward Fast Fourier Transform with spectral back-projection in the polar regions to manually enforce Fourier periodicity.

Parameters:

• f (np.ndarray) – HEALPix pixel-space array.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

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

Raises:

ValueError – Deployment method not in {“numpy”, “jax”, “torch”}.

Returns:

Array of Fourier coefficients for all latitudes.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.healpix_fft_jax(f: Array, L: int, nside: int, reality: bool) → Array

Healpix FFT JAX implementation using jax.numpy/numpy stack Computes the Forward Fast Fourier Transform with spectral back-projection in the polar regions to manually enforce Fourier periodicity. JAX specific implementation of healpix_fft_numpy().

Parameters:

• f (jnp.ndarray) – HEALPix pixel-space array.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

Array of Fourier coefficients for all latitudes.

Return type:

jnp.ndarray

s2fft.utils.healpix_ffts.healpix_fft_numpy(f: ndarray, L: int, nside: int, reality: bool) → ndarray

Computes the Forward Fast Fourier Transform with spectral back-projection in the polar regions to manually enforce Fourier periodicity.

Parameters:

• f (np.ndarray) – HEALPix pixel-space array.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

Array of Fourier coefficients for all latitudes.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.healpix_fft_torch(f: tensor, L: int, nside: int, reality: bool) → tensor

Computes the Forward Fast Fourier Transform with spectral back-projection in the polar regions to manually enforce Fourier periodicity. Torch specific implementation of healpix_fft_numpy().

Parameters:

• f (torch.tensor) – HEALPix pixel-space array.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

Array of Fourier coefficients for all latitudes.

Return type:

torch.tensor

s2fft.utils.healpix_ffts.healpix_ifft(ftm: ndarray, L: int, nside: int, method: str = 'numpy', reality: bool = False) → ndarray

Wrapper function for the Inverse Fast Fourier Transform with spectral folding in the polar regions to mitigate aliasing.

Parameters:

• ftm (np.ndarray) – Array of Fourier coefficients for all latitudes.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

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

Raises:

ValueError – Deployment method not in {“numpy”, “jax”, “torch”}.

Returns:

HEALPix pixel-space array.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.healpix_ifft_jax(ftm: Array, L: int, nside: int, reality: bool) → Array

Computes the Inverse Fast Fourier Transform with spectral folding in the polar regions to mitigate aliasing, using JAX. JAX specific implementation of healpix_ifft_numpy().

Parameters:

• ftm (jnp.ndarray) – Array of Fourier coefficients for all latitudes.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

HEALPix pixel-space array.

Return type:

jnp.ndarray

s2fft.utils.healpix_ffts.healpix_ifft_numpy(ftm: ndarray, L: int, nside: int, reality: bool) → ndarray

Computes the Inverse Fast Fourier Transform with spectral folding in the polar regions to mitigate aliasing.

Parameters:

• ftm (np.ndarray) – Array of Fourier coefficients for all latitudes.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

HEALPix pixel-space array.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.healpix_ifft_torch(ftm: tensor, L: int, nside: int, reality: bool) → tensor

Computes the Inverse Fast Fourier Transform with spectral folding in the polar regions to mitigate aliasing. Torch specific implementation of healpix_ifft_numpy().

Parameters:

• ftm (torch.tensor) – Array of Fourier coefficients for all latitudes.

• L (int) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

HEALPix pixel-space array.

Return type:

torch.tensor

s2fft.utils.healpix_ffts.p2phi_rings(t: ndarray, nside: int) → ndarray

Convert index to \(\phi\) angle for HEALPix for all \(\theta\) rings. Vectorised implementation of p2phi_ring().

Parameters:

• t (np.ndarrray) – vector of \(\theta\) ring indicies, i.e. [0,1,…,ntheta-1]

• nside (int) – HEALPix Nside resolution parameter.

Returns:

\(\phi\) offset for each ring index.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.p2phi_rings_jax(t: Array, nside: int) → Array

Convert index to \(\phi\) angle for HEALPix for all \(\theta\) rings. JAX implementation of p2phi_rings().

Parameters:

• t (jnp.ndarrray) – vector of \(\theta\) ring indicies, i.e. [0,1,…,ntheta-1]

• nside (int) – HEALPix Nside resolution parameter.

Returns:

\(\phi\) offset for each ring index.

Return type:

jnp.ndarray

s2fft.utils.healpix_ffts.ring_phase_shifts_hp(L: int, nside: int, forward: bool = False, reality: bool = False) → ndarray

Generates a phase shift vector for HEALPix for all \(\theta\) rings.

Parameters:

• L (int, optional) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

Vector of phase shifts with shape \([n_{\theta}, 2L-1]\).

Return type:

np.ndarray

s2fft.utils.healpix_ffts.ring_phase_shifts_hp_jax(L: int, nside: int, forward: bool = False, reality: bool = False) → Array

Generates a phase shift vector for HEALPix for all \(\theta\) rings. JAX implementation of ring_phase_shifts_hp().

Parameters:

• L (int, optional) – Harmonic band-limit.

• nside (int) – HEALPix Nside resolution parameter.

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

Returns:

Vector of phase shifts with shape \([n_{\theta}, 2L-1]\).

Return type:

jnp.ndarray

s2fft.utils.healpix_ffts.spectral_folding(fm: ndarray, nphi: int, L: int) → ndarray

Folds higher frequency Fourier coefficients back onto lower frequency coefficients, i.e. aliasing high frequencies.

Parameters:

• fm (np.ndarray) – Slice of Fourier coefficients corresponding to ring at latitute t.

• nphi (int) – Total number of pixel space phi samples for latitude t.

• L (int) – Harmonic band-limit.

Returns:

Lower resolution set of aliased Fourier coefficients.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.spectral_folding_jax(fm: Array, nphi: int, L: int) → Array

Folds higher frequency Fourier coefficients back onto lower frequency coefficients, i.e. aliasing high frequencies. JAX specific implementation of spectral_folding().

Parameters:

• fm (jnp.ndarray) – Slice of Fourier coefficients corresponding to ring at latitute t.

• nphi (int) – Total number of pixel space phi samples for latitude t.

• L (int) – Harmonic band-limit.

Returns:

Lower resolution set of aliased Fourier coefficients.

Return type:

jnp.ndarray

s2fft.utils.healpix_ffts.spectral_folding_torch(fm: tensor, nphi: int, L: int) → tensor

Folds higher frequency Fourier coefficients back onto lower frequency coefficients, i.e. aliasing high frequencies. Torch specific implementation of spectral_folding().

Parameters:

• fm (torch.tensor) – Slice of Fourier coefficients corresponding to ring at latitute t.

• nphi (int) – Total number of pixel space phi samples for latitude t.

• L (int) – Harmonic band-limit.

Returns:

Lower resolution set of aliased Fourier coefficients.

Return type:

torch.tensor

s2fft.utils.healpix_ffts.spectral_periodic_extension(fm: ndarray, nphi: int, L: int) → ndarray

Extends lower frequency Fourier coefficients onto higher frequency coefficients, i.e. imposed periodicity in Fourier space.

Parameters:

• fm (np.ndarray) – Slice of Fourier coefficients corresponding to ring at latitute t.

• nphi (int) – Total number of pixel space phi samples for latitude t.

• L (int) – Harmonic band-limit.

Returns:

Higher resolution set of periodic Fourier coefficients.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.spectral_periodic_extension_jax(fm: Array, L: int) → Array

Extends lower frequency Fourier coefficients onto higher frequency coefficients, i.e. imposed periodicity in Fourier space. Based on spectral_periodic_extension(), modified to be JIT-compilable.

Parameters:

• fm (jnp.ndarray) – Slice of Fourier coefficients corresponding to ring at latitute t.

• L (int) – Harmonic band-limit.

Returns:

Higher resolution set of periodic Fourier coefficients.

Return type:

jnp.ndarray

s2fft.utils.healpix_ffts.spectral_periodic_extension_torch(fm: tensor, L: int) → tensor

Extends lower frequency Fourier coefficients onto higher frequency coefficients, i.e. imposed periodicity in Fourier space. Based on spectral_periodic_extension().

Parameters:

• fm (torch.tensor) – Slice of Fourier coefficients corresponding to ring at latitute t.

• L (int) – Harmonic band-limit.

Returns:

Higher resolution set of periodic Fourier coefficients.

Return type:

torch.tensor