Spherical Harmonic Functions (sphtfunc)

Spherical Harmonic Functions (sphtfunc)#

This module provides functions for spherical harmonic transforms, enabling conversion between pixel-domain maps and spherical harmonic coefficients.

Note

These functions require the s2fft package to be installed. Install with:

pip install jax-healpy[recommended]
jax_healpy.sphtfunc.alm2map(alms, nside, lmax=None, mmax=None, pixwin=False, fwhm=0.0, sigma=None, pol=True, inplace=False, verbose=True, healpy_ordering=False)[source]#

Computes a Healpix map given the alm.

The alm are given as a complex array. You can specify lmax and mmax, or they will be computed from array size (assuming lmax==mmax).

Parameters:

  • alms (complex, array or sequence of arrays) – A complex array or a sequence of complex arrays. Each array must have a size of the form: mmax * (2 * lmax + 1 - mmax) / 2 + lmax + 1

  • nside (int, scalar) – The nside of the output map.

  • lmax (None or int, scalar, optional) – Explicitly define lmax (needed if mmax!=lmax)

  • mmax (None or int, scalar, optional) – Explicitly define mmax (needed if mmax!=lmax)

  • pixwin (bool, optional) – Smooth the alm using the pixel window functions. Default: False.

  • fwhm (float, scalar, optional) – The fwhm of the Gaussian used to smooth the map (applied on alm) [in radians]

  • sigma (float, scalar, optional) – The sigma of the Gaussian used to smooth the map (applied on alm) [in radians]

  • pol (bool, optional) – If True, assumes input alms are TEB. Output will be TQU maps. (input must be 1 or 3 alms) If False, apply spin 0 harmonic transform to each alm. (input can be any number of alms) If there is only one input alm, it has no effect. Default: True.

  • inplace (bool, optional) – If True, input alms may be modified by pixel window function and beam smoothing (if alm(s) are complex128 contiguous arrays). Otherwise, input alms are not modified. A copy is made if needed to apply beam smoothing or pixel window.

  • ordering (healpy) – True if the input alms follow the healpy ordering. By default, the s2fft ordering is assumed.

  • healpy_ordering (bool)

Returns:

maps – A Healpix map in RING scheme at nside or a list of T,Q,U maps (if polarized input)

Return type:

array or list of arrays

Notes

Running map2alm then alm2map will not return exactly the same map if the discretized field you construct on the sphere is not band-limited (for example, if you have a map containing pixel-based noise rather than beam-smoothed noise). If you need a band-limited map, you have to start with random numbers in lm space and transform these via alm2map. With such an input, the accuracy of map2alm->alm2map should be quite good, depending on your choices of lmax, mmax and nside (for some typical values, see e.g., section 5.1 of https://arxiv.org/pdf/1010.2084).

jax_healpy.sphtfunc.map2alm(maps, lmax=None, mmax=None, iter=3, pol=True, use_weights=False, datapath=None, gal_cut=0, use_pixel_weights=False, verbose=True, healpy_ordering=False)[source]#

Computes the alm of a Healpix map. The input maps must all be in ring ordering.

For recommendations about how to set lmax, iter, and weights, see the Anafast documentation

Pixel values are weighted before applying the transform:

  • when you don’t specify any weights, the uniform weight value 4*pi/n_pix is used

  • with ring weights enabled (use_weights=True), pixels in every ring are weighted with a uniform value similar to the one above, ring weights are included in healpy

  • with pixel weights (use_pixel_weights=True), every pixel gets an individual weight

Pixel weights provide the most accurate transform, so you should always use them if possible. However they are not included in healpy and will be automatically downloaded and cached in ~/.astropy the first time you compute a trasform at a specific nside.

If datapath is specified, healpy will first check that local folder before downloading the weights. The easiest way to setup the folder is to clone the healpy-data repository:

git clone –depth 1 healpy/healpy-data cd healpy-data bash download_weights_8192.sh

and set datapath to the root of the repository.

Parameters:

  • maps (array-like, shape (Npix,) or (n, Npix)) – The input map or a list of n input maps. Must be in ring ordering.

  • lmax (int, scalar, optional) – Maximum l of the power spectrum. Default: 3*nside-1

  • mmax (int, scalar, optional) – Maximum m of the alm. Default: lmax

  • iter (int, scalar, optional) – Number of iteration (default: 3)

  • pol (bool, optional) – If True, assumes input maps are TQU. Output will be TEB alm’s. (input must be 1 or 3 maps) If False, apply spin 0 harmonic transform to each map. (input can be any number of maps) If there is only one input map, it has no effect. Default: True.

  • use_weights (bool, scalar, optional) – If True, use the ring weighting. Default: False.

  • datapath (None or str, optional) – If given, the directory where to find the pixel weights. See in the docstring above details on how to set it up.

  • gal_cut (float [degrees]) – pixels at latitude in [-gal_cut;+gal_cut] are not taken into account

  • use_pixel_weights (bool, optional) – If True, use pixel by pixel weighting, healpy will automatically download the weights, if needed

  • verbose (bool, optional) – Deprecated, has not effect.

  • healpy_ordering (bool, optional) – By default, we follow the s2fft ordering for the alms. To use healpy ordering, set it to True.

Returns:

alms – alm or a tuple of 3 alm (almT, almE, almB) if polarized input.

Return type:

array or tuple of array

Notes

The pixels which have the special UNSEEN value are replaced by zeros before spherical harmonic transform. They are converted back to UNSEEN value, so that the input maps are not modified. Each map have its own, independent mask.

Spherical Harmonic Transforms#

Core functions for forward and inverse spherical harmonic transforms:

jax_healpy.map2alm(maps, lmax=None, mmax=None, iter=3, pol=True, use_weights=False, datapath=None, gal_cut=0, use_pixel_weights=False, verbose=True, healpy_ordering=False)[source]#

Computes the alm of a Healpix map. The input maps must all be in ring ordering.

For recommendations about how to set lmax, iter, and weights, see the Anafast documentation

Pixel values are weighted before applying the transform:

  • when you don’t specify any weights, the uniform weight value 4*pi/n_pix is used

  • with ring weights enabled (use_weights=True), pixels in every ring are weighted with a uniform value similar to the one above, ring weights are included in healpy

  • with pixel weights (use_pixel_weights=True), every pixel gets an individual weight

Pixel weights provide the most accurate transform, so you should always use them if possible. However they are not included in healpy and will be automatically downloaded and cached in ~/.astropy the first time you compute a trasform at a specific nside.

If datapath is specified, healpy will first check that local folder before downloading the weights. The easiest way to setup the folder is to clone the healpy-data repository:

git clone –depth 1 healpy/healpy-data cd healpy-data bash download_weights_8192.sh

and set datapath to the root of the repository.

Parameters:

  • maps (array-like, shape (Npix,) or (n, Npix)) – The input map or a list of n input maps. Must be in ring ordering.

  • lmax (int, scalar, optional) – Maximum l of the power spectrum. Default: 3*nside-1

  • mmax (int, scalar, optional) – Maximum m of the alm. Default: lmax

  • iter (int, scalar, optional) – Number of iteration (default: 3)

  • pol (bool, optional) – If True, assumes input maps are TQU. Output will be TEB alm’s. (input must be 1 or 3 maps) If False, apply spin 0 harmonic transform to each map. (input can be any number of maps) If there is only one input map, it has no effect. Default: True.

  • use_weights (bool, scalar, optional) – If True, use the ring weighting. Default: False.

  • datapath (None or str, optional) – If given, the directory where to find the pixel weights. See in the docstring above details on how to set it up.

  • gal_cut (float [degrees]) – pixels at latitude in [-gal_cut;+gal_cut] are not taken into account

  • use_pixel_weights (bool, optional) – If True, use pixel by pixel weighting, healpy will automatically download the weights, if needed

  • verbose (bool, optional) – Deprecated, has not effect.

  • healpy_ordering (bool, optional) – By default, we follow the s2fft ordering for the alms. To use healpy ordering, set it to True.

Returns:

alms – alm or a tuple of 3 alm (almT, almE, almB) if polarized input.

Return type:

array or tuple of array

Notes

The pixels which have the special UNSEEN value are replaced by zeros before spherical harmonic transform. They are converted back to UNSEEN value, so that the input maps are not modified. Each map have its own, independent mask.

Convert a HEALPix map to spherical harmonic coefficients.

This function performs the forward spherical harmonic transform, converting a map defined on the HEALPix pixelization into spherical harmonic coefficients \(a_{\ell m}\).

Parameters#

marray_like

Input HEALPix map

lmaxint, optional

Maximum multipole moment. If not specified, lmax = 3*nside - 1

mmaxint, optional

Maximum azimuthal quantum number. If not specified, mmax = lmax

iterint, optional

Number of iterations for the iterative map2alm algorithm

use_weightsbool, optional

Whether to use ring weights for improved accuracy

Returns#

almarray_like

Complex spherical harmonic coefficients

jax_healpy.alm2map(alms, nside, lmax=None, mmax=None, pixwin=False, fwhm=0.0, sigma=None, pol=True, inplace=False, verbose=True, healpy_ordering=False)[source]#

Computes a Healpix map given the alm.

The alm are given as a complex array. You can specify lmax and mmax, or they will be computed from array size (assuming lmax==mmax).

Parameters:

  • alms (complex, array or sequence of arrays) – A complex array or a sequence of complex arrays. Each array must have a size of the form: mmax * (2 * lmax + 1 - mmax) / 2 + lmax + 1

  • nside (int, scalar) – The nside of the output map.

  • lmax (None or int, scalar, optional) – Explicitly define lmax (needed if mmax!=lmax)

  • mmax (None or int, scalar, optional) – Explicitly define mmax (needed if mmax!=lmax)

  • pixwin (bool, optional) – Smooth the alm using the pixel window functions. Default: False.

  • fwhm (float, scalar, optional) – The fwhm of the Gaussian used to smooth the map (applied on alm) [in radians]

  • sigma (float, scalar, optional) – The sigma of the Gaussian used to smooth the map (applied on alm) [in radians]

  • pol (bool, optional) – If True, assumes input alms are TEB. Output will be TQU maps. (input must be 1 or 3 alms) If False, apply spin 0 harmonic transform to each alm. (input can be any number of alms) If there is only one input alm, it has no effect. Default: True.

  • inplace (bool, optional) – If True, input alms may be modified by pixel window function and beam smoothing (if alm(s) are complex128 contiguous arrays). Otherwise, input alms are not modified. A copy is made if needed to apply beam smoothing or pixel window.

  • ordering (healpy) – True if the input alms follow the healpy ordering. By default, the s2fft ordering is assumed.

  • healpy_ordering (bool)

Returns:

maps – A Healpix map in RING scheme at nside or a list of T,Q,U maps (if polarized input)

Return type:

array or list of arrays

Notes

Running map2alm then alm2map will not return exactly the same map if the discretized field you construct on the sphere is not band-limited (for example, if you have a map containing pixel-based noise rather than beam-smoothed noise). If you need a band-limited map, you have to start with random numbers in lm space and transform these via alm2map. With such an input, the accuracy of map2alm->alm2map should be quite good, depending on your choices of lmax, mmax and nside (for some typical values, see e.g., section 5.1 of https://arxiv.org/pdf/1010.2084).

Convert spherical harmonic coefficients to a HEALPix map.

This function performs the inverse spherical harmonic transform, converting spherical harmonic coefficients \(a_{\ell m}\) into a map defined on the HEALPix pixelization.

Parameters#

almarray_like

Input spherical harmonic coefficients

nsideint

HEALPix resolution parameter

lmaxint, optional

Maximum multipole moment

mmaxint, optional

Maximum azimuthal quantum number

pixwinbool, optional

Whether to apply pixel window function correction

Returns#

marray_like

Output HEALPix map

Examples#

Basic spherical harmonic transform:

import jax.numpy as jnp
import jax_healpy as hp

# Create a test map
nside = 64
npix = hp.nside2npix(nside)
test_map = jnp.random.normal(0, 1, npix)

# Forward transform: map -> alm
alm = hp.map2alm(test_map, lmax=128)

# Inverse transform: alm -> map
reconstructed_map = hp.alm2map(alm, nside=nside)

# Check reconstruction accuracy
rms_error = jnp.sqrt(jnp.mean((test_map - reconstructed_map)**2))
print(f"RMS reconstruction error: {rms_error}")

Working with specific multipole ranges:

# Transform with specific lmax and mmax
lmax, mmax = 64, 32
alm = hp.map2alm(test_map, lmax=lmax, mmax=mmax)

# Reconstruct with the same parameters
reconstructed = hp.alm2map(alm, nside=nside, lmax=lmax, mmax=mmax)

Batch processing multiple maps:

# Process multiple maps simultaneously
batch_maps = jnp.random.normal(0, 1, (10, npix))

# Vectorized transform
batch_alm = jax.vmap(lambda m: hp.map2alm(m, lmax=64))(batch_maps)

# Vectorized inverse transform
batch_reconstructed = jax.vmap(lambda a: hp.alm2map(a, nside=nside))(batch_alm)

Mathematical Background#

The spherical harmonic transform decomposes a function \(f(\theta, \phi)\) on the sphere into spherical harmonic coefficients:

\[a_{\ell m} = \int f(\theta, \phi) Y_{\ell m}^*(\theta, \phi) d\Omega\]

where \(Y_{\ell m}(\theta, \phi)\) are the spherical harmonic basis functions.

The inverse transform reconstructs the function:

\[f(\theta, \phi) = \sum_{\ell=0}^{\ell_{max}} \sum_{m=-\ell}^{\ell} a_{\ell m} Y_{\ell m}(\theta, \phi)\]

HEALPix pixelization provides an efficient discretization of the sphere that enables fast spherical harmonic transforms while maintaining good accuracy for most astronomical applications.

Performance Notes#

  • Transforms are automatically JIT-compiled for optimal performance

  • GPU acceleration is supported when JAX is configured with CUDA/ROCm

  • Memory usage scales as O(npix) for maps and O(lmax * mmax) for coefficients

  • For large transforms, consider using batch processing to manage memory usage

Contents