Large receptive-field convolution#

LargeConv provides large effective convolution kernels without constructing the expensive fine-resolution interpolation stencil required by a direct HealPixConv.

It applies a matched multiresolution operator:

\[ y = U_1 U_2 \cdots U_L\; C_{W_c}\;D_L\cdots D_2D_1 x + b, \]

where:

  • \(D_i\) is a smooth HealPixDown;

  • \(C_{W_c}\) is a compact trainable HealPixConv;

  • \(U_i\) is the HealPixUp constructed from the same smoothing parameters;

  • \(b\) is added after returning to the fine resolution.

The compact kernel is learned end-to-end through the complete chain. Its effective fine-resolution kernel therefore includes the exact Down/Up operators used at inference.

Quick start#

import torch
from healpix_analyse import LargeConv

layer = LargeConv(
    level=8,  # nside = 256
    in_channels=8,
    out_channels=16,
    kernel_sz=33,
    max_compact_kernel_sz=7,
    cell_ids=cell_ids,
    ellipsoid="sphere",
    device="cuda" if torch.cuda.is_available() else "cpu",
)

x = torch.randn(4, 8, len(cell_ids), device=layer.device)
y = layer(x)
print(y.shape)  # [4, 16, len(cell_ids)] for one gauge

Selecting the resolution and compact kernel#

Kernel size is converted using its radius rather than its diameter. For a requested fine kernel \(K_f\):

\[ r_f = \frac{K_f-1}{2}, \qquad r_c(L) = \left\lceil\frac{r_f}{2^L}\right\rceil, \qquad K_c(L)=2r_c(L)+1. \]

The smallest \(L\) satisfying

\[K_c(L) \leq K_{\max}\]

is selected automatically.

For kernel_sz=33 and max_compact_kernel_sz=7:

Down levels

Coarse radius

Compact kernel

Accepted?

0

16

33

no

1

8

17

no

2

4

9

no

3

2

5

yes

Therefore a requested 33×33 receptive field uses three Down operations, a 5×5 compact convolution, and three Up operations. A 7×7 compact kernel at that level would have a nominal fine diameter of 49×49.

The selected plan is exposed through:

print(layer.n_levels)                 # 3
print(layer.compact_kernel_sz)        # 5
print(layer.coarse_nside)             # nside // 8
print(layer.effective_kernel_sz)      # 33

When the requested kernel already fits within the configured maximum, no resolution change is performed and LargeConv reduces to one HealPixConv plus its fine-resolution bias and optional normalisation.

Why the compact kernel is learned directly#

Naively resizing a 33×33 array to 5×5 would ignore:

  • the smoothing transfer function of HealPixDown;

  • the transpose and normalisation used by HealPixUp;

  • HEALPix geometry and bilinear stencil interpolation;

  • gauge orientation;

  • partial-patch boundaries.

Learning \(W_c\) through the whole chain optimises the actual operator \(U^LC_{W_c}D^L\). Gradients pass through every sparse matrix multiplication, the compact convolution, and each interpolation operation.

The representation is intentionally lower rank than an arbitrary direct large kernel. Downsampling removes high-frequency information, so an oscillatory fine kernel cannot in general be reproduced exactly. Large smooth kernels and learned contextual filters are the intended use cases.

Matched Down and Up operators#

Every pair receives the same:

  • ellipsoid;

  • radius_deg;

  • sigma_deg;

  • weight_norm;

  • floating-point dtype and device.

In addition, every HealPixUp is constructed with paired_down=down_layer. It transposes the already stored sparse Down matrix instead of rebuilding an operator on a potentially different set of fine children. The pairing is therefore exact even for irregular partial patches.

Only HealPixDown(mode="smooth") is used. Max pooling is nonlinear and cannot define a linear effective convolution kernel.

up_norm="col_l1" with weight_norm="l1" is the default combination. It preserves constant fields and is a suitable default for learned features.

layer = LargeConv(
    level=7,
    in_channels=4,
    out_channels=8,
    kernel_sz=65,
    weight_norm="l1",
    up_norm="col_l1",
)

Other supported combinations are inherited from HealPixDown and HealPixUp:

  • up_norm="adjoint" gives the raw transpose of Down;

  • weight_norm="l2", up_norm="diag_l2" favours local energy consistency.

Bias, normalisation and activation#

The internal compact convolution has a fixed zero bias. LargeConv.bias is added only after the final Up operation, so it represents a true fine-resolution bias even when partial-patch boundary normalisation is not uniform.

When use_norm=True, GroupNorm and ReLU are also applied at the final fine resolution:

Down × L → compact convolution → Up × L → bias → GroupNorm → ReLU

Applying these nonlinear operations at the coarse level would define a different operator and would prevent interpretation as one effective large convolution.

Input and output shapes#

Input

Interpretation

Output

[N]

one single-channel map

[N] when total output channels = 1, otherwise [C_out, N]

[B, N]

batch of single-channel maps

[B, G*C_out, N]

[B, C_in, N]

general batch

[B, G*C_out, N]

NumPy input returns NumPy output. Torch input returns Torch output and retains autograd connectivity.

Partial HEALPix patches#

cell_ids must be unique NESTED identifiers at the fine level. LargeConv:

  1. sorts identifiers internally to match the sparse Down matrices;

  2. builds the parent cell set at every coarse level;

  3. builds each Up operator from the exact corresponding parent set;

  4. crops extra children produced by partial-patch upsampling;

  5. restores the original cell_ids order before returning.

layer = LargeConv(
    level=6,
    in_channels=1,
    out_channels=2,
    kernel_sz=33,
    cell_ids=cell_ids,
)

As with every local convolution, values close to the patch boundary have less context. Supply a halo around the requested output region and crop the final result when accurate boundary responses are important.

Gauges#

The internal convolution supports the same gauge options as HealPixConv:

  • "phi";

  • "cosmo";

  • "projected_ref";

  • "two_ref".

n_gauges, singularity_lonlat and ref_direction are forwarded unchanged. The output contains n_gauges * out_channels channels.

Kernel access and initialisation#

The compact trainable parameter is available as layer.weight and has shape

[n_gauges, in_channels, out_channels, compact_kernel_sz**2]

A fixed or custom initial compact kernel can be installed with:

layer.set_compact_kernel(weight, bias=bias, requires_grad=True)

The provided array describes the compact coarse kernel, not a fine 33×33 kernel. Fitting a compact kernel from a pre-existing large fine kernel requires an explicit operator-matching optimisation and is outside the current API.

Computational benefit#

A direct convolution stores and gathers values proportional to

\[K_{\text{fine pixels}}\,K_f^2.\]

After \(L\) levels, the dominant compact convolution scales approximately as

\[\frac{K_{\text{fine pixels}}}{4^L}\,K_c^2,\]

plus sparse Down/Up operations. For 33×33 → 5×5 with three levels, the convolution gather term is reduced nominally by

\[ 4^3\frac{33^2}{5^2} \approx 2788. \]

Actual speedup depends on batch size, channels, patch shape, device and the sparse-matrix overhead.

Limitations#

  • The effective kernel is approximate in the sense of representational bandwidth; it is exactly the operator learned through the selected chain.

  • nside must remain at least one after all selected downsamplings.

  • Only NESTED ordering is currently supported.

  • Geometry construction can be expensive; the internal HealPixConv cache is enabled by default.

  • Full-sphere operation at very high nside still requires substantial memory for Down/Up sparse matrices.

  • Partial patches should include a halo for reliable edge values.

See Notebooks/large_conv_test.ipynb for configuration, constant-field, ordering, gradient, receptive-field and training tests.