healpix_analyse.fft_local#

Fast local 2-D FFTs for data sampled on HEALPix cells.

The HEALPix cell centres are projected onto a square gnomonic grid. The geometry is computed once by LocalFFT; subsequent projections and transforms use differentiable PyTorch scatter/gather operations.

The FFT itself is exactly invertible on the projected grid. The default HEALPix -> grid -> HEALPix round trip is approximate because it uses fast bilinear gridding in both directions.

Classes#

LocalFFT

A reusable local flat-sky FFT for a set of NESTED HEALPix cells.

Functions#

fft(cell_ids, level, data, *[, return_transform])

Build a LocalFFT and compute a local 2-D FFT.

ifft(→ ArrayLike)

Compute a local inverse FFT using cached or newly built geometry.

ps(→ tuple[ArrayLike, ArrayLike])

Compute a circular 1-D power spectrum from a square unshifted 2-D FFT.

Module Contents#

class healpix_analyse.fft_local.LocalFFT(cell_ids: ArrayLike, level: int, *, ellipsoid: str = 'sphere', max_patch_radius_deg: float = 10.0, pixel_size_rad: float | None = None, grid_size: int | None = None, norm: str = 'ortho', dtype: torch.dtype = torch.float32, device: str | torch.device | None = None)[source]#

Bases: torch.nn.Module

A reusable local flat-sky FFT for a set of NESTED HEALPix cells.

Parameters:

  • cell_ids (array-like, shape [N]) – Unique HEALPix cell identifiers in NESTED ordering.

  • level (int) – HEALPix level, with nside = 2**level.

  • ellipsoid (str, default "sphere") – Geometry passed to healpix_geo when locating cell centres.

  • max_patch_radius_deg (float, default 10) – Maximum angular distance of a cell centre from the spherical patch centre. A single gnomonic plane is rejected above this radius.

  • pixel_size_rad (float, optional) – Grid spacing in gnomonic coordinates. The default is the square root of the HEALPix pixel area, sqrt(pi/3) / nside.

  • grid_size (int, optional) – Side of the square projected grid. By default, the smallest power of two that covers the patch with a one-pixel margin is used.

  • norm ({"forward", "backward", "ortho"}, default "ortho") – Normalisation passed to torch.fft.fft2() and ifft2.

  • dtype (torch.dtype, default torch.float32) – Real dtype used for geometry and real-valued input maps.

  • device (torch.device or str, optional) – Device of the cached projection geometry. CUDA is selected when it is available, otherwise CPU.

Notes

The tangent frame is built from three-dimensional unit vectors, not from longitude/latitude differences. It is therefore well defined at both geographic poles and across the 0/360-degree meridian.

For torch inputs, gradients propagate through projection, FFT, IFFT and back-projection with respect to the data. The discrete HEALPix geometry is fixed and is not differentiable with respect to cell_ids.

level#
nside#
ellipsoid = ''#
max_patch_radius_deg#
norm = 'ortho'#
dtype = Ellipsis#
cdtype = Ellipsis#
n_cells#
pixel_size_rad#
grid_size = 0#
grid_shape#
centre_lon_deg#
centre_lat_deg#
patch_radius_deg#
property device: torch.device#

Current device of the cached geometry buffers.

property coverage_mask: torch.Tensor#

Boolean [grid_size, grid_size] mask reached by input cells.

property grid_density: torch.Tensor#

Accumulated interpolation weight on the projected square grid.

project(data: ArrayLike) → ArrayLike[source]#

Bilinearly scatter HEALPix values onto the cached square grid.

unproject(grid: ArrayLike) → ArrayLike[source]#

Bilinearly sample a projected grid at the HEALPix cell centres.

fft(data: ArrayLike) → ArrayLike[source]#

Project HEALPix data and compute its unshifted two-dimensional FFT.

ifft(spectrum: ArrayLike, *, real_output: bool = True) → ArrayLike[source]#

Compute the inverse FFT and sample it at the input HEALPix cells.

This is a fast approximate inverse of fft(). Set real_output=False when transforming genuinely complex-valued data.

ps(spectrum: ArrayLike, *, n_bins: int | None = None, include_corners: bool = False) → tuple[ArrayLike, ArrayLike][source]#

Compute a circularly averaged 1-D power spectrum from a 2-D FFT.

Parameters:

  • spectrum (numpy.ndarray or torch.Tensor, shape [..., grid_size, grid_size]) – Unshifted Fourier coefficients, normally returned by fft().

  • n_bins (int, optional) – Number of equally spaced radial-frequency bins. The default is grid_size // 2 when excluding corners.

  • include_corners (bool, default False) – If False, stop at the axis Nyquist frequency so every annulus is fully contained in the Fourier grid. If True, include the corner modes up to sqrt(2) times the axis Nyquist frequency; outer annuli are then only partially sampled.

Returns:

  • frequency (numpy.ndarray or torch.Tensor, shape [n_bins]) – Bin centres in cycles per gnomonic tangent unit. With the default grid spacing these are approximately cycles per radian.

  • power (numpy.ndarray or torch.Tensor, shape [..., n_bins]) – Mean value of abs(spectrum)**2 in each circular annulus.

Notes

No additional FFT normalisation, mean removal, mask correction or apodization is applied. These choices belong to the spectral-analysis workflow and should be performed before calling this method when required.

extra_repr() → str[source]#

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

healpix_analyse.fft_local.fft(cell_ids: ArrayLike, level: int, data: ArrayLike, *, return_transform: bool = False, **transform_kwargs)[source]#

Build a LocalFFT and compute a local 2-D FFT.

Use return_transform=True to return (spectrum, transform). Reusing that transform for IFFT and additional maps avoids rebuilding geometry.

healpix_analyse.fft_local.ifft(spectrum: ArrayLike, transform: LocalFFT | None = None, *, cell_ids: ArrayLike | None = None, level: int | None = None, real_output: bool = True, **transform_kwargs) → ArrayLike[source]#

Compute a local inverse FFT using cached or newly built geometry.

Prefer passing the LocalFFT returned by fft(). Alternatively provide both cell_ids and level to rebuild the deterministic grid.

healpix_analyse.fft_local.ps(spectrum: ArrayLike, transform: LocalFFT | None = None, *, pixel_size_rad: float | None = None, n_bins: int | None = None, include_corners: bool = False) → tuple[ArrayLike, ArrayLike][source]#

Compute a circular 1-D power spectrum from a square unshifted 2-D FFT.

Pass a LocalFFT to use its grid spacing and shape validation. For a standalone Fourier array, provide pixel_size_rad; it defaults to 1, in which case frequencies are expressed in cycles per grid-spacing unit.