HealPixFFTConv: large-kernel FFT convolution

This notebook checks the zero-padded FFT convolution against direct conv2d, verifies the identity-kernel projection round trip, tests gradients/CUDA, measures performance, and repeats the geometry test at the North Pole.

%load_ext autoreload
%autoreload 2

import time
import healpix_geo
import numpy as np
import torch
import torch.nn.functional as F

from healpix_analyse import HealPixFFTConv

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.manual_seed(0)
np.random.seed(0)
print("device:", device)

device: cuda

Local NESTED patch

The patch stays below the default ten-degree gnomonic limit. The same ordered cell_ids array defines the last dimension of every input and output.

LEVEL = 8
CENTRE = (30.0, 20.0)
RADIUS_DEG = 5.0
cell_ids, _, _ = healpix_geo.nested.cone_coverage(
    CENTRE, RADIUS_DEG, LEVEL, ellipsoid="sphere"
)
cell_ids = np.asarray(cell_ids, dtype=np.int64)
print("cells:", len(cell_ids))

cells: 1599

KERNEL_SZ = 33
layer = HealPixFFTConv(
    level=LEVEL,
    in_channels=2,
    out_channels=3,
    kernel_sz=KERNEL_SZ,
    cell_ids=cell_ids,
    ellipsoid="sphere",
    device=device,
    dtype=torch.float64,  # strict correctness check; benchmark is converted to float32 later
)
print(layer)
print("projected grid:", layer.transform.grid_shape)
print("padded FFT:", layer.fft_shape)
assert layer.fft_shape[0] >= layer.grid_size + KERNEL_SZ - 1

HealPixFFTConv(
  level=8, nside=256, in=2, out=3, kernel_sz=33, cells=1599, grid=64, fft_shape=(128, 128), radius=5.175 deg
  (transform): LocalFFT(level=8, n_cells=1599, grid_shape=(64, 64), radius=5.175 deg, ellipsoid='sphere', dtype=torch.float64)
)
projected grid: (64, 64)
padded FFT: (128, 128)

FFT result versus direct convolution

Both paths use the same gnomonic projection and HEALPix back-projection. They differ only in the planar convolution implementation. This strict comparison uses float64 because CUDA conv2d may otherwise use cuDNN/TF32 while the FFT path uses complex64 arithmetic.

weight = torch.randn(2, 3, KERNEL_SZ, KERNEL_SZ, device=device, dtype=torch.float64) / (2 * KERNEL_SZ**2) ** 0.5
bias = torch.randn(3, device=device, dtype=torch.float64) * 0.01
layer.set_kernel(weight, bias=bias, requires_grad=True)
x = torch.randn(2, 2, len(cell_ids), device=device, dtype=torch.float64)

y_fft = layer(x)
grid = layer.transform._project_tensor(x)
grid_direct = F.conv2d(
    grid, layer.weight.permute(1, 0, 2, 3), padding=KERNEL_SZ // 2
)
y_direct = layer.transform._unproject_tensor(grid_direct) + layer.bias.view(1, -1, 1)

max_abs = (y_fft - y_direct).abs().max().item()
rel_rms = ((y_fft - y_direct).square().mean().sqrt() / y_direct.square().mean().sqrt().clamp_min(1e-12)).item()
print({"max_abs": max_abs, "relative_rms": rel_rms})
assert rel_rms < 2e-10, f"FFT/direct relative RMS mismatch: {rel_rms:.3e}"
assert max_abs < 2e-9, f"FFT/direct maximum mismatch: {max_abs:.3e}"

{'max_abs': 3.6637359812630166e-15, 'relative_rms': 1.2503260022905908e-15}

Identity kernel

A centred delta must reproduce exactly the project -> unproject result. The remaining difference from the original HEALPix field is the interpolation error documented for LocalFFT.

identity_layer = HealPixFFTConv(
    LEVEL, 1, 1, KERNEL_SZ, cell_ids=cell_ids, ellipsoid="sphere", device=device
)
delta = torch.zeros(1, 1, KERNEL_SZ, KERNEL_SZ, device=device)
delta[0, 0, KERNEL_SZ // 2, KERNEL_SZ // 2] = 1
identity_layer.set_kernel(delta)
signal = torch.randn(len(cell_ids), device=device)
identity_output = identity_layer(signal)
reference = identity_layer.transform.unproject(identity_layer.transform.project(signal))
identity_max_abs = (identity_output - reference).abs().max().item()
print("identity versus project/unproject max abs:", identity_max_abs)
assert torch.allclose(identity_output, reference, rtol=3e-5, atol=3e-6)

identity versus project/unproject max abs: 5.960464477539062e-07

Autograd

The spatial kernel is transformed inside the graph during training, so gradients must reach both the HEALPix input and the kernel.

x_grad = torch.randn(2, 2, len(cell_ids), device=device, requires_grad=True)
loss = layer(x_grad).square().mean()
loss.backward()
print({
    "loss": loss.item(),
    "input_grad_norm": x_grad.grad.norm().item(),
    "kernel_grad_norm": layer.weight.grad.norm().item(),
})
assert torch.isfinite(x_grad.grad).all()
assert torch.isfinite(layer.weight.grad).all()

{'loss': 0.17717390947697245, 'input_grad_norm': 0.00801227055490017, 'kernel_grad_norm': 0.38059604792559915}

Timing against direct conv2d

This compares complete HEALPix-to-HEALPix paths using the same projection. Results depend strongly on device, batch size, grid and kernel size. Increase KERNEL_SZ and batch size to locate the crossover on the target GPU.

def synchronize():
    if device.type == "cuda":
        torch.cuda.synchronize(device)

def benchmark(function, repetitions=20):
    for _ in range(3):
        function()
    synchronize()
    start = time.perf_counter()
    for _ in range(repetitions):
        function()
    synchronize()
    return 1e3 * (time.perf_counter() - start) / repetitions

layer = layer.float()
x = x.float()
layer.eval()
with torch.inference_mode():
    fft_ms = benchmark(lambda: layer(x))

    def direct_path():
        projected = layer.transform._project_tensor(x)
        convolved = F.conv2d(
            projected, layer.weight.permute(1, 0, 2, 3), padding=KERNEL_SZ // 2
        )
        return layer.transform._unproject_tensor(convolved)

    direct_ms = benchmark(direct_path)
print({"fft_ms": fft_ms, "direct_ms": direct_ms, "speedup": direct_ms / fft_ms})

{'fft_ms': 0.2798809204250574, 'direct_ms': 1.5176092041656375, 'speedup': 5.422338907063876}

North Pole geometry

A final identity test verifies that constructing the tangent frame does not depend on a longitude direction at the pole.

polar_ids, _, _ = healpix_geo.nested.cone_coverage(
    (0.0, 90.0), 3.0, LEVEL, ellipsoid="sphere"
)
polar_ids = np.asarray(polar_ids, dtype=np.int64)
polar = HealPixFFTConv(
    LEVEL, 1, 1, 17, cell_ids=polar_ids, ellipsoid="sphere", device=device
)
polar_delta = torch.zeros(1, 1, 17, 17, device=device)
polar_delta[0, 0, 8, 8] = 1
polar.set_kernel(polar_delta)
polar_x = torch.randn(len(polar_ids), device=device)
polar_y = polar(polar_x)
polar_reference = polar.transform.unproject(polar.transform.project(polar_x))
print({
    "centre": (polar.transform.centre_lon_deg, polar.transform.centre_lat_deg),
    "radius_deg": polar.transform.patch_radius_deg,
    "max_abs": (polar_y - polar_reference).abs().max().item(),
})
assert torch.allclose(polar_y, polar_reference, rtol=3e-5, atol=3e-6)

{'centre': (92.96285521726368, 90.0), 'radius_deg': 3.1069851092541816, 'max_abs': 4.76837158203125e-07}