healpix_analyse.large_conv

Memory-efficient large-receptive-field convolution on HEALPix maps.

LargeConv replaces a large fine-resolution stencil by a matched sequence of smooth downsamplings, a compact HealPixConv, and the corresponding upsamplings. The compact kernel is learned end-to-end through the complete multiresolution operator.

Classes

LargeConv

Learn a large effective HEALPix convolution through a compact kernel.

Module Contents

class healpix_analyse.large_conv.LargeConv(level: int, in_channels: int, out_channels: int, kernel_sz: int, *, max_compact_kernel_sz: int = 7, n_gauges: int = 1, gauge_type: str = 'phi', singularity_lonlat: tuple[float, float] | None = None, ref_direction=None, cell_ids=None, nest: bool = True, ellipsoid: str = 'WGS84', radius_deg: float | None = None, sigma_deg: float | None = None, weight_norm: str = 'l1', up_norm: str = 'col_l1', use_norm: bool = False, dtype: torch.dtype = torch.float32, device: str | torch.device | None = None, cache_dir: Path | str | None = _DEFAULT_CACHE_DIR)

Bases: torch.nn.Module

Learn a large effective HEALPix convolution through a compact kernel.

The linear part of the layer is

Up[0] ... Up[L-1] @ Conv_compact @ Down[L-1] ... Down[0].

kernel_sz describes the requested fine-resolution receptive-field diameter. The number of levels and compact odd kernel size are selected automatically so that compact_kernel_sz <= max_compact_kernel_sz. For example, kernel_sz=33 and max_compact_kernel_sz=7 produce three down/up levels and a compact 5 x 5 kernel.

Parameters:

• level (int) – Fine input Grid4Earth/HEALPix level (integer >= 0). Internally, nside = 2**level.

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

• kernel_sz (int) – Requested positive odd fine-resolution kernel diameter.

• max_compact_kernel_sz (int, default 7) – Maximum positive odd kernel diameter at the coarse resolution.

• n_gauges, gauge_type, singularity_lonlat, ref_direction – Passed to the internal HealPixConv.

• cell_ids (array-like, optional) – Fine-resolution NESTED cell identifiers. Partial patches are sorted internally and outputs are restored to the original input order.

• ellipsoid (str, default "WGS84") – Shared by every Down, Up and convolution geometry operation.

• radius_deg, sigma_deg, weight_norm – Identical smoothing parameters used in every matched Down/Up pair.

• up_norm ({"adjoint", "col_l1", "diag_l2"}, default "col_l1") – Upsampling normalisation. col_l1 preserves constant fields when paired with weight_norm="l1".

• use_norm (bool, default False) – Apply GroupNorm and ReLU after the final fine-resolution upsampling.

• dtype, device, cache_dir – PyTorch and HealPixConv execution options.

Notes

The layer learns the compact kernel directly; it does not claim exact equivalence to an arbitrary supplied large kernel. Downsampling removes high-frequency information, so the representable large kernels are those in the multiresolution subspace defined by the selected Down/Up operators.

level

nside

in_channels

out_channels

kernel_sz

max_compact_kernel_sz = 7

coarse_nside

effective_kernel_sz = 1

G

ellipsoid = ''

weight_norm = ''

up_norm = ''

use_norm = False

dtype = Ellipsis

partial

n_pixels

down_layers

conv

up_layers

bias

property device: torch.device

Current device of the layer buffers.

property weight: torch.nn.Parameter

Learnable compact convolution kernel.

set_compact_kernel(weight, bias=None, requires_grad=False)

Set compact weights with the same shapes accepted by HealPixConv.

forward(x: ArrayLike) → ArrayLike

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.