HealPixConv — gauge-equivariant convolution#

HealPixConv is a PyTorch nn.Module that convolves a signal on a HEALPix sphere with a learned kernel, equivariantly: the same stencil is applied at every cell, transported into a local frame, so the result does not depend on how the pixel grid happens to be oriented there.

It works on full-sky maps and on partial-sky patches, takes NumPy arrays or Torch tensors, and is differentiable end to end.

This page is the guide. See HealPixConv — API reference for the parameters and shapes, and HealPixConv — how the geometry is built for how the geometry is built.

What it does#

HealPixConv implements the spherical convolution described in the gauge-equivariant framework: a fixed kernel_sz × kernel_sz stencil is defined at the North Pole, then rotated to each target pixel by a composition of spherical rotations. The same learned kernel is applied at every pixel, in every gauge orientation, making the layer equivariant to the chosen gauge group.

Key design choices:

The kernel never moves. It is the data that is pulled toward the fixed kernel via bilinear interpolation of the input map, not the kernel that is warped per pixel.
All heavy geometry is precomputed once at construction time and stored as registered buffers. The forward pass is purely index-and-multiply.
Fully vectorised: no Python loops over pixels or gauges in either the constructor or the forward pass.
Differentiable: every operation in the hot path (index_select, einsum, arithmetic) is supported by PyTorch autograd.

Quick start#

import numpy as np
import torch
from healpix_analyse.convol import HealPixConv

level      = 6                         # Grid4Earth resolution level
nside      = 2**level                  # internal HEALPix resolution
npix       = 12 * nside ** 2           # 49 152 pixels
in_ch, out_ch = 3, 16

# --- build the layer (geometry precomputed here) ---
conv = HealPixConv(
    level       = level,
    in_channels = in_ch,
    out_channels= out_ch,
    kernel_sz   = 3,          # 3×3 = 9 stencil points
    n_gauges    = 4,          # G=4 gauge orientations
    gauge_type  = "projected_ref",
    singularity_lonlat = (84.0, 28.0),   # Himalayas + antipodal Pacific
)

# --- forward pass (differentiable) ---
x = torch.randn(8, in_ch, npix)         # batch of 8 maps
y = conv(x)                             # [8, G*out_ch, npix] = [8, 64, 49152]

# --- inspect singularity placement ---
print(conv.singularity_info())

Gauges and singularities#

The hairy-ball theorem states that every smooth tangent vector field on S² must have at least one zero. Equivalently, the total index of all singularities of a gauge field on S² must equal exactly 2.

Each built-in gauge places those unavoidable bad points differently.

`"phi"` — meridian-aligned#

α_base(k) = 0

The kernel is always aligned with the local meridian. Singularities are fixed at the geographic North and South Poles (index +1 each, total = 2).

Best for data that is well-behaved away from the poles and where interpretability matters. Computationally cheapest.

`"cosmo"` — cosmological convention#

α_base(k) = -φ_k   (Northern hemisphere, θ ≤ π/2)
α_base(k) = +φ_k   (Southern hemisphere, θ > π/2)

Same singularity locations as "phi" but the gauge angle flips sign across the equator to match the CMB/cosmological convention used by healpy. Useful when comparing results with legacy healpy-based pipelines.

`"projected_ref"` — one freely placed antipodal pair#

A reference vector r ∈ ℝ³ is projected onto the tangent plane at each pixel k:

r_proj(k) = r − (r · n_k) · n_k
α_base(k) = atan2( r_proj · e_φ,  r_proj · e_θ )

The gauge is undefined exactly where r_proj = 0, i.e. where r ∥ n_k. This occurs at the two antipodal points:

singularity₁ = (lon_s,        lat_s)       ← direction of r
singularity₂ = (lon_s + 180°, −lat_s)      ← antipode, forced

To place the first singularity at a desired location, pass:

singularity_lonlat = (lon_s, lat_s)

The reference vector is computed automatically:

r = [cos(lat_s)·cos(lon_s),  cos(lat_s)·sin(lon_s),  sin(lat_s)]

Constraint: the second singularity is always the antipode of the first. You cannot move them independently with this gauge type — use "two_ref" for that.

`"two_ref"` — two freely placed singularity pairs#

Two independent reference vectors r₁ and r₂ define the gauge angle via the complex product of their tangent-plane projections:

z_j(k) = (r_j_proj · e_θ)  +  i · (r_j_proj · e_φ)      j = 1, 2

α_base(k) = arg( z₁(k) · z₂(k) )
           = atan2( Re(z₁)·Im(z₂) + Im(z₁)·Re(z₂),
                    Re(z₁)·Re(z₂) − Im(z₁)·Im(z₂) )

Using the complex product avoids two independent atan2 calls (which would each wrap and accumulate phase jumps) and computes arg(z₁) + arg(z₂) in a single numerically stable operation.

Singularity structure (Poincaré–Hopf budget):

Location	Count	Index	Origin
`+r₁`, `−r₁`	2	+1 each	zeros of z₁
`+r₂`, `−r₂`	2	+1 each	zeros of z₂
N-Pole, S-Pole	2	−1 each	base-frame side-effect
Total		4−2 = 2	✓

The four user-controlled points are index +1 (vortex-like), which is the well-behaved type to place over unimportant regions. The geographic poles become index −1 (hyperbolic saddle), a sharper singularity — keep them outside the domain of interest or over regions where accuracy is not required.

To place singularities, pass a list of two (lon, lat) pairs:

singularity_lonlat = [(lon_1, lat_1), (lon_2, lat_2)]

Each pair controls one user-specified singularity; its antipodal point appears automatically. The full set of four bad points is thus:

{(lon₁, lat₁),  (lon₁+180°, −lat₁),  (lon₂, lat₂),  (lon₂+180°, −lat₂)}

Choosing where the singularities go#

Domain	Strategy
Ocean model	Place all 4 bad points over land masses. Example: Amazon + Borneo (antipodal pair) and Africa + central Pacific (second pair).
Atmosphere	Place all 4 bad points over open ocean. Example: central Pacific + Indian Ocean (pair 1) and South Atlantic + Maritime Continent (pair 2).
Full sphere / neutral	Use `"phi"` (poles). For `"projected_ref"`, passing `singularity_lonlat=(0, 90)` reproduces the pole placement with a smoother field off the poles.
Poles outside domain	Any `"two_ref"` config naturally moves the geographic-pole index-−1 singularities to a fixed location — confirm they are harmless for your domain.

Recipes#

Full-sphere convolution, `"phi"` gauge#

conv = HealPixConv(level=6, in_channels=1, out_channels=32,
                   kernel_sz=3, n_gauges=1, gauge_type="phi")
y = conv(x)   # x: [B, 1, 49152]  →  y: [B, 32, 49152]

Multi-gauge equivariant layer#

conv = HealPixConv(level=6, in_channels=16, out_channels=16,
                   kernel_sz=3, n_gauges=4, gauge_type="phi")
# output has G * C_out = 64 channels

Ocean model — singularities over two land masses#

# Singularity pair 1: Amazon basin + Borneo (its antipode)
# Singularity pair 2: Africa + central Pacific (its antipode)
conv = HealPixConv(
    level=6, in_channels=3, out_channels=16,
    gauge_type="two_ref",
    singularity_lonlat=[(-55.0, -10.0), (20.0, 5.0)],
)
print(conv.singularity_info())

Atmospheric model — singularities over open ocean#

conv = HealPixConv(
    level=6, in_channels=5, out_channels=32,
    gauge_type="projected_ref",
    singularity_lonlat=(-160.0, 0.0),   # central Pacific + Indian Ocean antipode
)

Partial-sky patch#

import healpix_geo

level = 6
nside = 2**level
cell_ids, _, _ = healpix_geo.nested.cone_coverage(
    (0.0, 45.0), 20.0, level, ellipsoid="WGS84"
)
conv = HealPixConv(
    level=level, in_channels=1, out_channels=8,
    cell_ids=cell_ids,
)
x_patch = torch.randn(4, 1, len(cell_ids))
y_patch = conv(x_patch)

Fixed (non-learnable) isotropic kernel#

conv = HealPixConv(level=5, in_channels=1, out_channels=1, kernel_sz=3)
W = np.zeros((1, 1, 9), dtype=np.float32)
W[0, 0, 4]              = 0.5       # centre
W[0, 0, [0,1,2,3,5,6,7,8]] = 0.5/8 # ring
conv.set_kernel(W, requires_grad=False)

With GroupNorm + ReLU#

conv = HealPixConv(
    level=6, in_channels=16, out_channels=16,
    n_gauges=4, use_norm=True,
)
# The output is already passed through GroupNorm then ReLU inside forward()

NumPy in / NumPy out#

x_np = np.random.randn(49152).astype(np.float32)
y_np = conv(x_np)   # returns np.ndarray, same shape policy as input