FFT-accelerated convolution on local HEALPix patches#

HealPixFFTConv applies a learned, potentially very large convolution kernel to data stored on a local NESTED HEALPix patch. It reuses the pole-safe gnomonic projection from LocalFFT and evaluates the planar convolution with torch.fft.rfft2/irfft2.

HEALPix values → gnomonic projection → zero-padded FFT convolution
               → inverse FFT → HEALPix back-projection

The geometry is built once. Projection, convolution and back-projection run on CPU or CUDA and are differentiable with respect to the input and kernel.

Quick start#

import torch
from healpix_analyse import HealPixFFTConv

layer = HealPixFFTConv(
    level=12,
    in_channels=4,
    out_channels=8,
    kernel_sz=65,
    cell_ids=cell_ids,
    ellipsoid="sphere",
    device="cuda" if torch.cuda.is_available() else "cpu",
)

x = torch.randn(2, 4, len(cell_ids), device=layer.device)
y = layer(x)
print(y.shape)  # [2, 8, len(cell_ids)]

cell_ids must contain unique NESTED identifiers at level and must describe one local patch. The default maximum cell-centre radius is 10 degrees, as for LocalFFT. The implementation is safe at the poles and across the longitude wrap because the tangent frame is constructed from three-dimensional vectors.

Interface#

HealPixFFTConv(
    level,
    in_channels,
    out_channels,
    kernel_sz=33,
    cell_ids=cell_ids,
    ellipsoid="sphere",
    max_patch_radius_deg=10.0,
    pixel_size_rad=None,
    grid_size=None,
    use_norm=False,
    dtype=torch.float32,
    device=None,
    cache_kernel_fft=True,
)

Parameter	Meaning
`level`	Grid4Earth resolution; internally `nside = 2**level`
`in_channels`, `out_channels`	Input and output feature channels
`kernel_sz`	Positive odd tangent-grid kernel size
`cell_ids`	Local NESTED HEALPix domain and data order
`ellipsoid`	Geometry used by `healpix_geo`
`max_patch_radius_deg`	Gnomonic distortion guard
`pixel_size_rad`, `grid_size`	Optional `LocalFFT` grid overrides
`use_norm`	Apply GroupNorm and ReLU at the HEALPix output
`cache_kernel_fft`	Cache the kernel spectrum during no-gradient inference

Accepted input shapes are [N], [B, N] for a single input channel, and [B, C_in, N]. The output contains C_out channels and follows the same NumPy/Torch return-type convention as the other operators.

Useful attributes include:

layer.transform.patch_radius_deg
layer.grid_size
layer.fft_shape
layer.weight.shape  # [C_in, C_out, K, K]
layer.cell_ids

Linear convolution rather than circular wrap#

A raw FFT product evaluates a periodic convolution. That would incorrectly connect the west/east and north/south sides of the local tangent grid. HealPixFFTConv instead pads both operands to

\[ N_{\rm linear} = N_{\rm grid} + K - 1 \]

and rounds this size up to the next power of two. After the inverse FFT, the central N_grid × N_grid region is selected. The result therefore matches a standard 2-D cross-correlation with padding=K//2 and zeros outside the grid.

For input spectra \(X_c\) and flipped kernel spectra \(H_{oc}\), the multi-channel operation is

\[ Y_o(k_x,k_y) = \sum_c X_c(k_x,k_y)H_{oc}(k_x,k_y). \]

The channel contraction is performed directly in Fourier space. Real FFTs are used, so only the non-redundant half-spectrum is stored.

Kernel convention and initialisation#

The learned spatial parameter has shape [C_in, C_out, kernel_sz, kernel_sz], consistent with the channel ordering of HealPixConv. set_kernel accepts this shape or the flattened equivalent [C_in, C_out, kernel_sz**2]:

kernel = torch.zeros(1, 1, 65, 65)
kernel[0, 0, 32, 32] = 1
layer.set_kernel(kernel, requires_grad=False)

The layer implements the neural-network cross-correlation convention. The internal spatial flip is handled automatically before transforming the kernel.

During training, the kernel FFT is recomputed on every forward pass so autograd can update the spatial weights. During eval() under torch.no_grad() or torch.inference_mode(), its spectrum is cached and reused until the weight changes.

Accuracy and boundaries#

The FFT convolution itself agrees with direct conv2d to floating-point precision on the projected grid. The complete HEALPix operator also contains the approximate interpolation sequence

HEALPix → grid → HEALPix.

Consequently, even an identity kernel returns the same approximate round trip as LocalFFT.project followed by LocalFFT.unproject, rather than an exact copy of arbitrary HEALPix data.

Uncovered grid points and values beyond the supplied patch are zero. For scientific output near a boundary, supply a halo at least half the kernel width and crop the halo after convolution. A smooth apodization can also reduce edge discontinuities.

Performance#

For a projected grid of side \(N\) and kernel side \(K\):

direct convolution costs approximately \(O(C_{in}C_{out}N^2K^2)\);
FFT convolution costs approximately \(O((C_{in}+C_{out})M^2\log M + C_{in}C_{out}M^2)\), where \(M\) is the padded FFT side.

Small kernels such as 3×3 or 5×5 are normally faster with HealPixConv or torch.conv2d. FFT convolution becomes attractive for large kernels, especially for batches and repeated inference on a GPU.

`HealPixFFTConv` or `LargeConv`?#

Property	`HealPixFFTConv`	`LargeConv`
Domain	One local gnomonic patch	Full sphere or partial NESTED patch
Large-kernel representation	Full learned planar kernel	Compact learned multiresolution kernel
High-frequency kernel detail	Retained on projected grid	Reduced by Down operations
Main approximation	Projection/back-projection	Multiresolution subspace
Best use	Very large kernels on local imagery	Wide spherical context and memory control

HealPixFFTConv has one global tangent-frame orientation for the patch. It is not a gauge-equivariant spherical convolution. Use HealPixConv when gauge equivariance is required.