healpix_analyse.fft_conv

Fast large-kernel convolution on local HEALPix patches.

The input cells are projected with healpix_analyse.fft_local.LocalFFT, convolved on its gnomonic grid with a zero-padded 2-D FFT, then sampled back at the original HEALPix cell centres. All data-path operations are implemented with PyTorch and remain differentiable.

Classes

HealPixFFTConv

Large learned convolution on a local NESTED HEALPix patch.

Module Contents

class healpix_analyse.fft_conv.HealPixFFTConv(level: int, in_channels: int, out_channels: int, kernel_sz: int = 33, *, cell_ids, ellipsoid: str = 'sphere', max_patch_radius_deg: float = 10.0, pixel_size_rad: float | None = None, grid_size: int | None = None, use_norm: bool = False, dtype: torch.dtype = torch.float32, device: str | torch.device | None = None, cache_kernel_fft: bool = True)

Bases: torch.nn.Module

Large learned convolution on a local NESTED HEALPix patch.

Parameters:

• level (int) – Grid4Earth/HEALPix level. Internally, nside = 2**level.

• in_channels, out_channels (int) – Input and output feature channels.

• kernel_sz (int, default 33) – Positive odd side length of the learned 2-D tangent-plane kernel.

• cell_ids (array-like) – Unique NESTED cells forming one local patch. A single gnomonic plane cannot represent a full sphere.

• ellipsoid (str, default "sphere") – Geometry passed to LocalFFT.

• max_patch_radius_deg (float, default 10) – Reject patches extending farther from their spherical centre.

• pixel_size_rad, grid_size (optional) – Projected-grid controls forwarded to LocalFFT.

• use_norm (bool, default False) – Apply GroupNorm followed by ReLU after back-projection.

• dtype, device (optional) – PyTorch execution options. CUDA is selected when available.

• cache_kernel_fft (bool, default True) – Cache the transformed kernel during no-gradient inference. Training always recomputes it so gradients reach the spatial kernel.

Notes

The convolution is linear rather than circular: both the projected data and kernel are zero-padded to at least grid_size + kernel_sz - 1 in each direction. This prevents values wrapping across opposite sides of the local patch. Values near the patch boundary still see zero outside the supplied domain; provide a halo and crop it when boundary accuracy is important.

The learned weight follows the existing HealPixConv channel order and has shape [C_in, C_out, kernel_sz, kernel_sz].

level

nside

in_channels

out_channels

kernel_sz = 33

use_norm = False

dtype = Ellipsis

cache_kernel_fft = True

transform

n_cells

grid_size = 0

fft_shape

weight

bias

property device: torch.device

Current device of weights and projection buffers.

property cell_ids: torch.Tensor

NESTED cell identifiers in input/output order.

train(mode: bool = True)

Set the module in training mode.

This has an effect only on certain modules. See the documentation of particular modules for details of their behaviors in training/evaluation mode, i.e., whether they are affected, e.g. Dropout, BatchNorm, etc.

Args:

mode (bool): whether to set training mode (True) or evaluation

mode (False). Default: True.

Returns:

Module: self

set_kernel(weight, bias=None, requires_grad: bool = False)

Set a spatial kernel using [C_in,C_out,K,K] or flattened form.

forward(x: ArrayLike) → ArrayLike

Apply the FFT convolution to [N], [B,N] or [B,C,N].

extra_repr() → str

Return the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.