Local HEALPix FFT: projection and reconstruction tests

This notebook validates the gnomonic projection, the fast FFT/IFFT round trip, pole handling, the 0/360-degree meridian and PyTorch autograd. The HEALPix-to-grid reprojection is bilinear, so the final reconstruction is expected to be approximate.

import numpy as np
import torch
import matplotlib.pyplot as plt
import healpix_geo

from healpix_analyse.fft_local import LocalFFT, fft, ifft, ps

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

device: cuda

<torch._C.Generator at 0x7e5c04b09290>

def make_patch(lon_deg, lat_deg, *, level=7, radius_deg=5.0):
    cell_ids, _, _ = healpix_geo.nested.cone_coverage(
        (lon_deg, lat_deg), radius_deg, level, ellipsoid="sphere"
    )
    return np.asarray(cell_ids, dtype=np.int64)


def relative_rms(reference, estimate):
    reference = torch.as_tensor(reference)
    estimate = torch.as_tensor(estimate)
    return (
        torch.linalg.vector_norm(estimate - reference)
        / torch.linalg.vector_norm(reference)
    ).item()

Geometry and FFT/IFFT reconstruction

The four cases exercise a regular equatorial patch, a patch crossing the longitude discontinuity, and both poles.

cases = {
    "equator": (45.0, 0.0),
    "meridian_0": (359.0, 20.0),
    "north_pole": (0.0, 90.0),
    "south_pole": (0.0, -90.0),
}

results = []
transforms = {}
for name, centre in cases.items():
    cell_ids = make_patch(*centre)
    transform = LocalFFT(
        cell_ids, level=7, device=device, dtype=torch.float64
    )
    transforms[name] = transform

    # A smooth signal defined directly in the local tangent plane.
    x = transform.projected_x
    y = transform.projected_y
    data = torch.sin(25.0 * x) + 0.4 * torch.cos(17.0 * y)
    spectrum = transform.fft(data)
    reconstructed = transform.ifft(spectrum)

    constant = torch.ones_like(data)
    constant_reconstructed = transform.ifft(transform.fft(constant))
    results.append({
        "case": name,
        "cells": transform.n_cells,
        "grid": transform.grid_size,
        "radius_deg": transform.patch_radius_deg,
        "smooth_rel_rms": relative_rms(data, reconstructed),
        "constant_max_abs": (constant_reconstructed - constant).abs().max().item(),
    })

for result in results:
    print(result)

{'case': 'equator', 'cells': 434, 'grid': 32, 'radius_deg': 5.345987578998197, 'smooth_rel_rms': 0.009681833550991233, 'constant_max_abs': 4.440892098500626e-16}
{'case': 'meridian_0', 'cells': 428, 'grid': 32, 'radius_deg': 5.341431499616129, 'smooth_rel_rms': 0.010064204688898055, 'constant_max_abs': 3.3306690738754696e-16}
{'case': 'north_pole', 'cells': 420, 'grid': 32, 'radius_deg': 5.118462097772193, 'smooth_rel_rms': 0.016842954587225012, 'constant_max_abs': 4.440892098500626e-16}
{'case': 'south_pole', 'cells': 420, 'grid': 32, 'radius_deg': 5.118462097772193, 'smooth_rel_rms': 0.01684295458722501, 'constant_max_abs': 4.440892098500626e-16}

# Visual inspection of the most singular geographic case.
transform = transforms["north_pole"]
data = torch.sin(25.0 * transform.projected_x) + 0.4 * torch.cos(17.0 * transform.projected_y)
grid = transform.project(data).detach().cpu()
reconstructed = transform.ifft(transform.fft(data)).detach().cpu()

fig, axes = plt.subplots(1, 3, figsize=(15, 4))
axes[0].scatter(transform.projected_x.cpu(), transform.projected_y.cpu(), c=data.cpu(), s=8)
axes[0].set_title("HEALPix centres in tangent plane")
axes[0].set_aspect("equal")
axes[1].imshow(grid, origin="lower", cmap="viridis")
axes[1].set_title("Projected square grid")
axes[2].scatter(data.cpu(), reconstructed, s=8, alpha=0.6)
limits = [float(data.min().cpu()), float(data.max().cpu())]
axes[2].plot(limits, limits, "k--")
axes[2].set_xlabel("input")
axes[2].set_ylabel("reconstructed")
axes[2].set_title("FFT/IFFT round trip")
plt.tight_layout()

Circular one-dimensional power spectrum

transform = transforms["equator"]
data = torch.sin(25.0 * transform.projected_x) + 0.4 * torch.cos(17.0 * transform.projected_y)
spectrum = transform.fft(data)
frequency, power = transform.ps(spectrum)

assert frequency.shape == (transform.grid_size // 2,)
assert power.shape == frequency.shape
assert torch.isfinite(power).all()
assert torch.all(frequency[1:] > frequency[:-1])
assert frequency[-1] < 0.5 / transform.pixel_size_rad

# A Fourier plane with unit power must have unit power in every annulus.
flat_spectrum = torch.ones_like(spectrum)
_, flat_power = transform.ps(flat_spectrum)
assert torch.allclose(flat_power, torch.ones_like(flat_power))
corner_frequency, corner_power = transform.ps(spectrum, include_corners=True)
assert corner_frequency.numel() > frequency.numel()
assert corner_power.shape == corner_frequency.shape
plt.loglog(frequency.cpu(), power.detach().cpu())
plt.xlabel("spatial frequency [cycles per tangent unit]")
plt.ylabel("circularly averaged power")
plt.grid(True, which="both", alpha=0.25)

Batched tensors, CUDA and autograd

transform = transforms["equator"].to(device)
batch = torch.randn(3, transform.n_cells, device=device, dtype=torch.float64, requires_grad=True)
spectrum = transform.fft(batch)
reconstructed = transform.ifft(spectrum)
frequency, batch_power = transform.ps(spectrum)
loss = reconstructed.square().mean() + torch.nan_to_num(batch_power).mean()
loss.backward()

print("batch spectrum shape:", tuple(spectrum.shape))
print("batch power shape:", tuple(batch_power.shape))
print("batch reconstruction shape:", tuple(reconstructed.shape))
print("finite gradient:", bool(torch.isfinite(batch.grad).all()))
print("gradient norm:", batch.grad.norm().item())

batch spectrum shape: (3, 32, 32)
batch power shape: (3, 16)
batch reconstruction shape: (3, 434)
finite gradient: True
gradient norm: 0.06853051345715866

Functional API and radius guard

cell_ids = make_patch(10.0, 30.0, level=7, radius_deg=4.0)
data = np.random.default_rng(0).standard_normal(len(cell_ids))
spectrum, transform = fft(
    cell_ids, 7, data, return_transform=True, device=device
)
reconstructed = ifft(spectrum, transform)
frequency, power = ps(spectrum, transform)
print("functional API relative RMS:", relative_rms(data, reconstructed))
print("functional PS shapes:", frequency.shape, power.shape)

large_patch = make_patch(0.0, 0.0, level=6, radius_deg=12.0)
try:
    LocalFFT(large_patch, level=6, max_patch_radius_deg=10.0, device=device)
except ValueError as error:
    print("expected radius error:", error)
else:
    raise AssertionError("A patch larger than 10 degrees should be rejected")

functional API relative RMS: 0.5946385112854626
functional PS shapes: (16,) (16,)
expected radius error: The patch radius is 12.6701 degrees, greater than max_patch_radius_deg=10. Split the cells into smaller local patches.