Local two-dimensional FFT on HEALPix patches#

healpix_analyse.fft_local computes a conventional two-dimensional FFT after projecting a local set of HEALPix samples onto a square gnomonic grid. It is intended for local flat-sky analysis, image-like diagnostics and local power spectra—not as a replacement for a spherical harmonic transform on large regions or the full sphere.

The implementation is:

  • compatible with NESTED HEALPix identifiers;

  • safe at both poles and across the 0/360-degree meridian;

  • based on PyTorch and usable on CPU or CUDA;

  • differentiable with respect to the input values;

  • reusable: all geometric and interpolation tables are computed once;

  • compatible with NumPy arrays and PyTorch tensors;

  • able to process arbitrary leading batch dimensions [..., N].

Quick start#

import torch
from healpix_analyse.fft_local import LocalFFT

transform = LocalFFT(
    cell_ids,
    level=12,
    ellipsoid="sphere",
    device="cuda" if torch.cuda.is_available() else "cpu",
)

spectrum = transform.fft(data)
reconstructed = transform.ifft(spectrum)

print(transform.grid_shape)
print(transform.patch_radius_deg)

cell_ids contains unique NESTED identifiers at the same level, and the last dimension of data must follow exactly the same order.

For several fields defined on the same cells, construct LocalFFT only once:

spectra = [transform.fft(field) for field in fields]

The geometry construction uses healpix-geo and NumPy on the CPU. Projection, FFT, IFFT and back-projection then run on the device holding the transform’s buffers.

Functional interface#

The module also exposes two convenience functions:

from healpix_analyse.fft_local import fft, ifft

spectrum, transform = fft(
    cell_ids,
    level,
    data,
    return_transform=True,
    device="cuda",
)
reconstructed = ifft(spectrum, transform)

Passing the transform to ifft is recommended. It guarantees that the same projection grid is used and avoids rebuilding the geometry. It is also possible to rebuild it deterministically:

reconstructed = ifft(
    spectrum,
    cell_ids=cell_ids,
    level=level,
    device="cuda",
)

Geometry#

Spherical centre#

The longitude and latitude of every HEALPix cell centre are converted to a three-dimensional unit vector \(\boldsymbol n_i\). The centre of the local patch is the normalised vector mean

\[ \boldsymbol n_0 = \frac{\sum_i \boldsymbol n_i} {\left\|\sum_i \boldsymbol n_i\right\|}. \]

This avoids averaging longitude directly, which would fail for a patch crossing 0/360 degrees.

Tangent frame and pole handling#

Two orthonormal vectors \(\boldsymbol e_x\) and \(\boldsymbol e_y\) are built in the plane perpendicular to \(\boldsymbol n_0\). Away from the poles, the global north direction is used to orient the frame. Close to a pole, a Cartesian reference axis is selected instead. Consequently the construction never divides by cos(latitude) and has no geographic pole singularity.

The gnomonic coordinates of cell \(i\) are

\[ x_i = \frac{\boldsymbol n_i \cdot \boldsymbol e_x} {\boldsymbol n_i \cdot \boldsymbol n_0}, \qquad y_i = \frac{\boldsymbol n_i \cdot \boldsymbol e_y} {\boldsymbol n_i \cdot \boldsymbol n_0}. \]

The coordinates are dimensionless tangent-plane coordinates (tan(angle)). For small patches they are numerically close to angular offsets in radians.

Patch-size guard#

The default maximum angular distance from the computed centre is 10 degrees:

LocalFFT(cell_ids, level, max_patch_radius_deg=10.0)

A larger patch raises ValueError. Gnomonic projection is mathematically defined up to 90 degrees from its centre, but distortion grows rapidly before that limit. Splitting large regions into overlapping local patches is safer for spectral analysis.

Coverage functions can include cells whose centres lie slightly outside the requested region because the cells themselves intersect its boundary. For example, a nominal five-degree cone may have a measured cell-centre radius of approximately 5.3 degrees.

Projected grid#

The default grid spacing is derived from the square root of the HEALPix pixel area:

\[ \Delta = \frac{\sqrt{\pi/3}}{n_{\rm side}}, \qquad n_{\rm side}=2^{\text{level}}. \]

This is a nominal linear scale, not an assertion that HEALPix centres form a Cartesian lattice. It can be overridden with pixel_size_rad.

The grid is square, centred on the tangent origin, and includes a one-pixel margin around all projected centres. By default its side is rounded up to the next power of two for efficient FFT execution. A fixed size can be requested with grid_size; an error is raised when it is too small.

Useful geometry attributes are:

Attribute

Meaning

n_cells

Number of input HEALPix cells

grid_size, grid_shape

Side and shape of the square FFT grid

centre_lon_deg, centre_lat_deg

Computed spherical centre

patch_radius_deg

Maximum angular cell-centre distance

pixel_size_rad

Gnomonic grid spacing

projected_x, projected_y

Position of every input cell in the tangent plane

coverage_mask

Boolean grid mask reached by at least one input cell

grid_density

Sum of interpolation weights on every grid point

Bilinear gridding#

Each HEALPix value contributes to the four surrounding grid points with bilinear weights \(w_{ig}\). Overlapping contributions are normalised by their accumulated density:

\[ G_g = \frac{\sum_i w_{ig}d_i}{\sum_i w_{ig}}. \]

Uncovered grid points are zero. The operation is implemented with cached indices and torch.scatter_add, so it is fast, batched and differentiable with respect to \(d_i\).

The projected grid can be inspected directly:

grid = transform.project(data)
back_projected = transform.unproject(grid)
mask = transform.coverage_mask

FFT conventions#

LocalFFT.fft returns the unshifted complex FFT:

spectrum = transform.fft(data)
display_spectrum = torch.fft.fftshift(spectrum, dim=(-2, -1))

The default normalisation is norm="ortho". With this convention, Parseval’s identity on the projected grid has no extra grid-size factor:

\[ \sum_g |G_g|^2 = \sum_k |F_k|^2. \]

"forward" and "backward" are also accepted and are passed unchanged to torch.fft.fft2 and torch.fft.ifft2.

For a grid spacing \(\Delta\), frequency axes in cycles per tangent unit are obtained with:

frequency = torch.fft.fftfreq(
    transform.grid_size,
    d=transform.pixel_size_rad,
    device=transform.device,
)

On a sphere of radius \(R\), a small-patch physical approximation uses d = R * transform.pixel_size_rad, giving frequencies in cycles per metre.

Circular one-dimensional power spectrum#

LocalFFT.ps azimuthally averages the two-dimensional Fourier power in concentric frequency annuli:

\[ P_j = \frac{1}{N_j}\sum_{\boldsymbol k\,\in\,A_j} |F(\boldsymbol k)|^2, \]

where \(A_j\) is radial bin \(j\) and \(N_j\) is its number of Fourier modes.

spectrum = transform.fft(data)
frequency, power = transform.ps(spectrum)

frequency contains the bin centres and power contains the mean power in each circular annulus. For a batched spectrum of shape [..., grid_size, grid_size], the returned power has shape [..., n_bins].

By default, the spectrum stops at the axis Nyquist frequency. Every annulus is then fully contained in the square Fourier grid. Corner modes can be included explicitly:

frequency, power = transform.ps(spectrum, include_corners=True)

The outer annuli are only partially sampled in that mode and should be interpreted with care. The number of bins can be selected with n_bins:

frequency, power = transform.ps(spectrum, n_bins=128)

The standalone interface can use a cached transform or an explicit spacing:

from healpix_analyse.fft_local import ps

frequency, power = ps(spectrum, transform)
frequency, power = ps(spectrum, pixel_size_rad=transform.pixel_size_rad)

No additional FFT normalisation is performed: power is always the circular mean of abs(spectrum)**2. Its amplitude therefore follows the norm used to create the FFT. Mean subtraction, apodization, coverage-mask correction and physical-unit conversion remain explicit analysis choices.

With the default spacing, frequency is expressed in cycles per gnomonic tangent unit, approximately cycles per radian for a small patch. To express the returned axis in cycles per metre on a sphere of radius \(R\):

frequency_per_metre = frequency / R

The radial reduction uses PyTorch scatter operations. Torch inputs remain on their device and gradients propagate from the one-dimensional power back to the Fourier coefficients and input data.

Meaning of IFFT and reconstruction error#

The FFT/IFFT pair is exact, up to floating-point precision, on the projected grid:

grid = transform.project(data)
grid_roundtrip = torch.fft.ifft2(
    torch.fft.fft2(grid, norm=transform.norm),
    norm=transform.norm,
).real

The complete HEALPix round trip is different:

HEALPix values -> bilinear grid -> FFT -> IFFT -> bilinear sampling

The last bilinear sampling is a fast approximate inverse, not the matrix pseudoinverse of the forward gridding operation. Therefore transform.ifft(transform.fft(data)) is generally approximate. A constant field is preserved, while spatially varying fields are slightly smoothed.

Example results for smooth level-7 patches in float64 are:

Location

Relative RMS error

Constant-field maximum error

Equator

0.968%

\(4.44\,10^{-16}\)

Across 0/360 degrees

1.006%

\(4.44\,10^{-16}\)

North pole

1.684%

\(4.44\,10^{-16}\)

South pole

1.684%

\(4.44\,10^{-16}\)

These figures validate the coordinate construction but are not universal accuracy guarantees. Reconstruction error depends on the field bandwidth, HEALPix level, patch shape and pixel_size_rad.

For genuinely complex input data, request a complex reconstruction:

reconstructed = transform.ifft(spectrum, real_output=False)

Power-spectrum recommendations#

The dominant systematic effect in a local power spectrum is normally the HEALPix-to-grid interpolation, followed by the finite patch boundary. The following checks are recommended:

  1. remove the mean before computing the spectrum;

  2. inspect coverage_mask and avoid interpreting zero padding as data;

  3. compare several pixel_size_rad values;

  4. test plane waves of increasing spatial frequency;

  5. estimate the projection transfer function with white noise or simulations;

  6. compare equivalent signals at the equator and near the poles;

  7. use a window for power-spectrum estimation when boundary leakage matters.

A Hann window reduces edge leakage but changes amplitudes and makes the windowed transform unsuitable for direct map reconstruction:

grid = transform.project(data)
window_1d = torch.hann_window(
    transform.grid_size,
    periodic=False,
    device=transform.device,
    dtype=grid.dtype,
)
window = window_1d[:, None] * window_1d[None, :]
windowed_spectrum = torch.fft.fft2(
    (grid - grid[transform.coverage_mask].mean()) * window,
    norm="ortho",
)

Window-energy and mask corrections must be included before comparing absolute power between patches.

NumPy, batches, CUDA and autograd#

NumPy input produces NumPy output. Torch input produces Torch output on the transform’s device:

# [N]
spectrum = transform.fft(single_map)

# [batch, channel, N]
spectrum = transform.fft(batch)
# -> [batch, channel, grid_size, grid_size]

The transform is an nn.Module; its cached buffers move normally between devices:

transform = transform.to("cuda")

Gradients propagate to Torch input values:

data = torch.randn(
    4, transform.n_cells,
    device=transform.device,
    requires_grad=True,
)
loss = transform.fft(data).abs().square().mean()
loss.backward()
assert torch.isfinite(data.grad).all()

Cell identifiers and projection geometry are discrete precomputed quantities; gradients are not defined with respect to them.

Sentinel-2 example#

Notebooks/fft_sentinel2_test.ipynb loads the real Sentinel-2 arrays produced by healpix-compress/test/compress_test.ipynb:

  • B04 red reflectance;

  • B08 near-infrared reflectance;

  • one complete 1024 x 1024 NESTED tile at level 20 per domain;

  • urban, water, forest, agriculture, snow/ice, clouds and ocean scenes.

Set HEALPIX_COMPRESS_REPO if the repositories are not siblings:

export HEALPIX_COMPRESS_REPO=/path/to/healpix-compress

The notebook reports reconstruction errors, verifies constant preservation, displays the tangent-plane projection and compares radial spectra for B04, B08 and NDVI.

Which power-spectrum helper?#

LocalFFT.ps() (also exposed as local_ps) is the production-ready 1D isotropic power spectrum for a local patch, described above. The separate powerspectra/powerspectra_lonlat functions (built on AlmTransform in alm.py) offer a similar-looking API but are still under development — see Flat-sky power spectra: powerspectra, powerspectra_lonlat, ps before reaching for them.

Errors and limitations#

  • cell_ids must be unique, valid NESTED identifiers at one level.

  • A single transform assumes one connected local patch.

  • The default 10-degree radius is deliberately conservative.

  • Missing cells and irregular patch boundaries appear as a spatial mask and affect the spectrum.

  • The current IFFT is a fast interpolating inverse, not a least-squares reconstruction.

  • float32 is recommended for throughput; float64 is useful for numerical validation.

  • The local tangent-plane spectrum should not be interpreted as a full-sky spherical \(C_\ell\) without an explicit flat-sky-to-spherical calibration.