Local HEALPix FFT on Sentinel-2 L2A data

This notebook applies LocalFFT to the real Sentinel-2 patches prepared by healpix-compress/test/compress_test.ipynb. Each scene is a complete 1024 x 1024 NESTED tile at HEALPix level 20. B04 (red), B08 (near infrared) and NDVI are evaluated.

The source arrays remain in the healpix-compress repository. Set HEALPIX_COMPRESS_REPO when that repository is not a sibling of healpix-analyse.

from pathlib import Path
import os
import time

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

from healpix_analyse.fft_local import LocalFFT

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float32
print("device:", device)

device: cuda

Locate and inspect the benchmark arrays

The domain order is inherited from compress_test.ipynb: urban, water, forest, agriculture, snow/ice, clouds and ocean. Change DOMAIN to test another surface type.

def find_compress_repo():
    configured = os.environ.get("HEALPIX_COMPRESS_REPO")
    candidates = [Path(configured).expanduser()] if configured else []
    cwd = Path.cwd().resolve()
    for root in (cwd, *cwd.parents):
        candidates.extend((root / "healpix-compress", root.parent / "healpix-compress"))
    for candidate in candidates:
        if (candidate / "test" / "compress_test.ipynb").is_file():
            return candidate.resolve()
    raise FileNotFoundError(
        "Cannot find healpix-compress. Set HEALPIX_COMPRESS_REPO to its root."
    )


compress_repo = find_compress_repo()
data_dir = compress_repo / "test" / "data"
required = {
    "cells": data_dir / "all_image_cell.npy",
    "B04": data_dir / "all_image_data.npy",
    "B08": data_dir / "all_image_data_b08.npy",
}
missing = [str(path) for path in required.values() if not path.is_file()]
if missing:
    raise FileNotFoundError(
        "Missing arrays generated by compress_test.ipynb:\n" + "\n".join(missing)
    )

arrays = {name: np.load(path, mmap_mode="r") for name, path in required.items()}
print("healpix-compress:", compress_repo)
for name, array in arrays.items():
    print(name, array.shape, array.dtype)

healpix-compress: /home/jovyan/healpix-compress
cells (7, 1048576) float64
B04 (7, 1048576) float64
B08 (7, 1048576) float64

DOMAINS = ["urban", "water", "forest", "agriculture", "snow_ice", "clouds", "ocean"]
DOMAIN = "urban"
LEVEL = 20
index = DOMAINS.index(DOMAIN)

# cell IDs were saved through a float64 array in compress_test; they are
# below 2**53 and can be converted back to int64 exactly.
cell_ids = np.asarray(arrays["cells"][index], dtype=np.int64)
red = np.asarray(arrays["B04"][index], dtype=np.float32)
nir = np.asarray(arrays["B08"][index], dtype=np.float32)

def fill_non_finite(values):
    values = np.asarray(values, dtype=np.float32).copy()
    finite = np.isfinite(values)
    if not finite.any():
        raise ValueError("The selected Sentinel-2 band has no finite values")
    values[~finite] = np.mean(values[finite], dtype=np.float64)
    return values

red = fill_non_finite(red)
nir = fill_non_finite(nir)
denominator = nir + red
ndvi = np.divide(
    nir - red, denominator, out=np.zeros_like(red), where=np.abs(denominator) > 1e-6
)
ndvi = np.clip(ndvi, -1.0, 1.0)

assert len(cell_ids) == 1024 * 1024
assert np.unique(cell_ids).size == len(cell_ids)
assert np.all(np.diff(np.sort(cell_ids)) == 1), "Expected one complete NESTED child tile"
print(DOMAIN, "cells:", len(cell_ids))
print("B04 range:", float(red.min()), float(red.max()))
print("B08 range:", float(nir.min()), float(nir.max()))
print("NDVI range:", float(ndvi.min()), float(ndvi.max()))

urban cells: 1048576
B04 range: -0.08254926651716232 1.643031120300293
B08 range: -0.025499999523162842 1.6039999723434448
NDVI range: -1.0 1.0

Construct the reusable projection

Geometry construction is performed once. The same transform is then used for every band. Level 20 has a nominal HEALPix scale close to six metres on an Earth-sized sphere.

if device.type == "cuda":
    torch.cuda.synchronize()
start = time.perf_counter()
transform = LocalFFT(
    cell_ids,
    LEVEL,
    ellipsoid="sphere",
    max_patch_radius_deg=10.0,
    dtype=dtype,
    device=device,
)
if device.type == "cuda":
    torch.cuda.synchronize()
geometry_seconds = time.perf_counter() - start

earth_radius_m = 6_371_008.8
pixel_size_m = earth_radius_m * transform.pixel_size_rad
print(transform)
print(f"centre: ({transform.centre_lon_deg:.6f}, {transform.centre_lat_deg:.6f}) deg")
print(f"radius: {transform.patch_radius_deg:.6f} deg")
print(f"nominal grid spacing: {pixel_size_m:.3f} m")
print(f"geometry construction: {geometry_seconds:.3f} s")
if device.type == "cuda":
    print(f"allocated CUDA memory: {torch.cuda.memory_allocated() / 2**20:.1f} MiB")

LocalFFT(level=20, n_cells=1048576, grid_shape=(2048, 2048), radius=0.05825 deg, ellipsoid='sphere', dtype=torch.float32)
centre: (2.290723, 48.727590) deg
radius: 0.058248 deg
nominal grid spacing: 6.218 m
geometry construction: 0.506 s
allocated CUDA memory: 96.0 MiB

Reconstruction accuracy and timing

The FFT/IFFT is exact on the projected grid. The reported HEALPix error also includes the fast bilinear back-projection and is therefore expected to be non-zero.

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


def reconstruction_metrics(name, values):
    data = torch.as_tensor(values, dtype=dtype, device=device)
    synchronize()
    start = time.perf_counter()
    spectrum = transform.fft(data)
    reconstructed = transform.ifft(spectrum)
    synchronize()
    elapsed = time.perf_counter() - start

    error = reconstructed - data
    rms = torch.sqrt(torch.mean(error.square()))
    signal_rms = torch.sqrt(torch.mean(data.square()))
    result = {
        "field": name,
        "relative_rms": float((rms / signal_rms).cpu()),
        "rmse": float(rms.cpu()),
        "mean_bias": float(error.mean().cpu()),
        "max_abs": float(error.abs().max().cpu()),
        "fft_ifft_seconds": elapsed,
    }
    return result, reconstructed.detach()


metrics = []
reconstructions = {}
for name, values in {"B04": red, "B08": nir, "NDVI": ndvi}.items():
    result, reconstructed = reconstruction_metrics(name, values)
    metrics.append(result)
    reconstructions[name] = reconstructed
    print(result)

constant = torch.ones(transform.n_cells, dtype=dtype, device=device)
constant_error = (transform.ifft(transform.fft(constant)) - constant).abs().max()
print("constant max absolute error:", float(constant_error.cpu()))

red_tensor = torch.as_tensor(red, dtype=dtype, device=device)
grid = transform.project(red_tensor)
grid_roundtrip = torch.fft.ifft2(
    torch.fft.fft2(grid, norm=transform.norm), norm=transform.norm
).real
print("projected-grid max error:", float((grid_roundtrip - grid).abs().max().cpu()))

{'field': 'B04', 'relative_rms': 0.0489577054977417, 'rmse': 0.005347579717636108, 'mean_bias': -1.5713681023044046e-08, 'max_abs': 0.315108060836792, 'fft_ifft_seconds': 0.02131364308297634}
{'field': 'B08', 'relative_rms': 0.08192791789770126, 'rmse': 0.020443975925445557, 'mean_bias': -3.7551444620476104e-08, 'max_abs': 0.5749712586402893, 'fft_ifft_seconds': 0.001677841879427433}
{'field': 'NDVI', 'relative_rms': 0.2056197077035904, 'rmse': 0.09963735193014145, 'mean_bias': -6.129448593128473e-08, 'max_abs': 1.420716643333435, 'fft_ifft_seconds': 0.0016176397912204266}
constant max absolute error: 7.152557373046875e-07
projected-grid max error: 5.960464477539062e-07

# Visualise B04 and its round-trip error on the regular tangent grid.
red_grid = transform.project(torch.as_tensor(red, dtype=dtype, device=device))
reconstructed_grid = transform.project(reconstructions["B04"])
mask = transform.coverage_mask

def masked_image(values):
    image = values.detach().cpu().numpy().copy()
    image[~mask.cpu().numpy()] = np.nan
    return image

fig, axes = plt.subplots(1, 3, figsize=(17, 5))
im0 = axes[0].imshow(masked_image(red_grid), origin="lower", cmap="gray")
axes[0].set_title(f"{DOMAIN}: Sentinel-2 B04")
plt.colorbar(im0, ax=axes[0], shrink=0.8)
im1 = axes[1].imshow(masked_image(reconstructed_grid), origin="lower", cmap="gray")
axes[1].set_title("FFT/IFFT reconstruction")
plt.colorbar(im1, ax=axes[1], shrink=0.8)
im2 = axes[2].imshow(
    masked_image(reconstructed_grid - red_grid), origin="lower", cmap="RdBu_r"
)
axes[2].set_title("Reconstruction minus input")
plt.colorbar(im2, ax=axes[2], shrink=0.8)
plt.tight_layout()

Apodized two-dimensional and radial power spectra

For spectral estimation, the projected mean is removed and a Hann window is applied to reduce leakage at the finite tile boundary. The window is for power-spectrum estimation only; it must not be used when testing map reconstruction. Frequencies below use the small-patch spherical conversion to cycles per metre.

def apodized_spectrum(values):
    data = torch.as_tensor(values, dtype=dtype, device=device)
    grid = transform.project(data)
    mask = transform.coverage_mask
    mean = grid[mask].mean()
    window_1d = torch.hann_window(
        transform.grid_size, periodic=False, dtype=dtype, device=device
    )
    window = window_1d[:, None] * window_1d[None, :]
    weighted = (grid - mean) * window * mask
    spectrum = torch.fft.fft2(weighted, norm="ortho")
    window_energy = torch.mean(window[mask].square())
    spectrum = spectrum / torch.sqrt(window_energy)
    return spectrum, spectrum.abs().square()


power_spectra = {}
radial_spectra = {}
for name, values in {"B04": red, "B08": nir, "NDVI": ndvi}.items():
    spectrum, power = apodized_spectrum(values)
    power_spectra[name] = power
    frequency, radial_power = transform.ps(spectrum)
    radial_spectra[name] = (
        (frequency / earth_radius_m).detach().cpu().numpy(),
        radial_power.detach().cpu().numpy(),
    )

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
display_power = torch.fft.fftshift(power_spectra["B04"]).detach().cpu().numpy()
axes[0].imshow(np.log10(display_power + np.finfo(np.float32).tiny), origin="lower")
axes[0].set_title("B04 log10 2D power")
for name, (frequency, power) in radial_spectra.items():
    valid = (frequency > 0) & np.isfinite(power) & (power > 0)
    axes[1].loglog(frequency[valid], power[valid], label=name)
axes[1].set_xlabel("spatial frequency [cycles m$^{-1}$]")
axes[1].set_ylabel("azimuthally averaged power")
axes[1].grid(True, which="both", alpha=0.25)
axes[1].legend()
axes[1].set_title(f"{DOMAIN}: radial spectra")
plt.tight_layout()

Optional resolution sweep

Run this cell to quantify the trade-off between interpolation error, grid sparsity, memory and FFT time. Factors below multiply the nominal HEALPix spacing. Smaller factors create finer grids but can expose gaps between projected samples.

nominal = np.sqrt(np.pi / 3.0) / (2**LEVEL)
for factor in [1.25, 1.0, 0.75, 0.5]:
    trial = LocalFFT(
        cell_ids, LEVEL, pixel_size_rad=factor * nominal, dtype=dtype, device=device
    )
    values = torch.as_tensor(red, dtype=dtype, device=device)
    synchronize()
    start = time.perf_counter()
    reconstructed = trial.ifft(trial.fft(values))
    synchronize()
    relative_rms = torch.linalg.vector_norm(reconstructed - values) / torch.linalg.vector_norm(values)
    covered_fraction = trial.coverage_mask.float().mean()
    print({
        "factor": factor,
        "grid": trial.grid_size,
        "covered_fraction": float(covered_fraction.cpu()),
        "relative_rms": float(relative_rms.cpu()),
        "seconds": time.perf_counter() - start,
    })
    del trial, reconstructed
    if device.type == "cuda":
        torch.cuda.empty_cache()

{'factor': 1.25, 'grid': 2048, 'covered_fraction': 0.160658597946167, 'relative_rms': 0.06721141934394836, 'seconds': 0.009381711017340422}
{'factor': 1.0, 'grid': 2048, 'covered_fraction': 0.25067734718322754, 'relative_rms': 0.0489577054977417, 'seconds': 0.0021828352473676205}
{'factor': 0.75, 'grid': 4096, 'covered_fraction': 0.11125421524047852, 'relative_rms': 0.027795156463980675, 'seconds': 0.012319224886596203}
{'factor': 0.5, 'grid': 4096, 'covered_fraction': 0.21428292989730835, 'relative_rms': 0.006400515325367451, 'seconds': 0.005952510051429272}