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”, “cuda”}. 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”, “cuda”}.

Returns:

Array of Fourier coefficients for all latitudes.

Return type:

np.ndarray

s2fft.utils.healpix_ffts.healpix_fft_cuda(f: Array, L: int, nside: int, reality: bool, norm: str = 'backward') → Array#

Healpix FFT JAX implementation using custom CUDA primitive.

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

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_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_cuda(ftm: Array, L: int, nside: int, reality: bool, norm: str = 'forward') → Array#

Healpix IFFT JAX implementation using custom CUDA primitive.

Computes the inverse fast Fourier transform with spectral folding in the polar regions to mitigate aliasing.

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