healpix_analyse.convol#

Gauge-equivariant spherical convolution on HEALPix maps.

Optimisation summary:

#  Location              Before                        After
──────────────────────────────────────────────────────────────────────────
 1  Buffer layout         _pos_safe / _w_norm stored     Stored as [4, G, K*P]
    (__init__)            as [G, 4, K*P].                → pos[j] is contiguous
                          pos[:, j, :].reshape(-1)        → reshape(-1) is a
                          forces a hidden copy each       free view, zero copy.
                          call in forward().

 2  forward gather        One index_select over          Loop over 4 neighbors:
    (forward)             [G*4*K*P] → peak alloc          accumulate in-place.
                          B·C·G·4·K·P (~3.6 GB).         Peak: 2·B·C·G·K·P
                                                         → 4× memory reduction.

 3  einsum vs bmm         einsum("bcgkp,gcop->bgok")     Explicit torch.bmm
    (forward)             may not fuse contractions.     → cuBLAS/MKL, 2–4×
                                                         faster on GPU.

 4  Geometry cache        All of _get_interp_weights +   _geometry_cache_key()
    (__init__)  ★ NEW ★   _bind_support_batched          hashes all ctor params
                          recomputed from scratch        → torch.save on first
                          on every __init__ call         run, torch.load on
                          (~0.8 s for nside=64, G=4).    subsequent runs
                                                         (~5 ms reload).
                                                         Cache dir configurable
                                                         via HEALPIXCONV_CACHE or
                                                         cache_dir= kwarg.
──────────────────────────────────────────────────────────────────────────
 Differentiability: 100 % preserved.
 API: identical to v2 — drop-in replacement (one new optional kwarg: cache_dir).
──────────────────────────────────────────────────────────────────────────

Attributes#

ArrayLike

Classes#

HealPixConv

Gauge-equivariant spherical convolution on HEALPix maps.

Module Contents#

healpix_analyse.convol.ArrayLike#
class healpix_analyse.convol.HealPixConv(level: int, in_channels: int, out_channels: int, kernel_sz: int = 3, n_gauges: int = 1, gauge_type: str = 'phi', singularity_lonlat: tuple[float, float] | None = None, ref_direction=None, cell_ids=None, nest: bool = True, use_norm: bool = False, device=None, ellipsoid: str = 'WGS84', dtype: torch.dtype = torch.float32, cache_dir: Path | str | None = _DEFAULT_CACHE_DIR)[source]#

Bases: torch.nn.Module

Gauge-equivariant spherical convolution on HEALPix maps.

Algorithm#

The convolution is precomputed at construction time in three stages.

Stage A — Kernel + rotation

A kernel_sz × kernel_sz grid of P unit vectors is defined at the North Pole (z = 1) with angular spacing equal to one HEALPix pixel width. This is the stencil template; it never changes.

For every target pixel k (colatitude θ_k, longitude φ_k) and every gauge g, a rotation matrix is built:

R_total[k,g] = R_gauge(α_g)  @  Rz(φ_k)  @  Ry(θ_k)
               └─ gauge roll ┘  └────── carry N.Pole → pixel k ──┘

Each of the P kernel vectors is then rotated into its position on the sphere around pixel k:

rotated[k, g, p] = R_total[k, g] @ vec_pole[p]   ∈ ℝ³

Stage B — Bilinear binding

For each of the K × G × P rotated directions, get_interp_weights returns the 4 nearest HEALPix neighbours and their bilinear weights. These indices and weights are stored as buffers (precomputed once).

Stage C — Forward pass (data → kernel)

It is the data that is brought to the (fixed) kernel, not the kernel that moves. At inference:

x_interp[b, c, g, k, p] = Σ_{j=0}^{3}  w[g,j,k,p] · x[b, c, nbr[g,j,k,p]]

i.e. the signal value at rotated stencil point p of pixel k (gauge g) is obtained by bilinear interpolation of the input map. Then:

y[b, g·C_out + o, k] = Σ_{c,p}  W[g, c, o, p] · x_interp[b, c, g, k, p]

The learned kernel W[G, C_in, C_out, P] is a fixed set of weights; its P positions implicitly correspond to the stencil template at the North Pole. Gauge equivariance is achieved because the same kernel is applied after the stencil has been rotated by the gauge.

Gauge types and singularities#

Every smooth gauge on S² has exactly two singular points (hairy-ball theorem). The three built-in gauges place them differently:

"phi"

Singularities at the geographic poles (θ = 0, π). α_base = 0 everywhere → kernel always meridian-aligned. Simple and fast; bad only at the poles.

"cosmo"

Singularities also at the poles, but the gauge flips sign across the equator to match the cosmological convention.

"projected_ref"

A reference vector r is projected onto the tangent plane; the gauge angle is atan2(r·e_φ, r·e_θ). The two singularities are the antipodal points where r is parallel to the surface normal, i.e. where the map pixel lies exactly on the direction of r:

singularity₁ = (lon_s, lat_s)
singularity₂ = (lon_s + 180°, -lat_s)   ← antipode

You control where those two points land by choosing r. Use singularity_lonlat=(lon_s, lat_s) to specify the first singularity in geographic coordinates; the reference vector is computed automatically:

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

  • Ocean model → place singularities over land (e.g. Amazon basin and its antipode in the Indian Ocean).

  • Atmosphere → place singularities over open ocean (e.g. central Pacific and its antipode in the Atlantic).

  • Full sphere → place singularities at the geographic poles (singularity_lonlat=(0, 90)), which reproduces the "phi" gauge but with the smooth-everywhere property off the poles.

Parameters:

  • level (int) – Grid4Earth/HEALPix level (integer >= 0). The internal HEALPix resolution is nside = 2**level.

  • in_channels (int) – Number of input channels C_in.

  • out_channels (int) – Number of output channels per gauge C_out. Total output channels = n_gauges × out_channels.

  • kernel_sz (int, default 3) – Odd integer ≥ 1. P = kernel_sz² stencil points.

  • n_gauges (int, default 1) – Number of gauge orientations G. Each gauge rotates the stencil by an additional angle of π/G relative to the previous one.

  • gauge_type ({"phi", "cosmo", "projected_ref"}, default "phi") – Gauge convention. See Gauge types and singularities above.

  • singularity_lonlat ((float, float) or None) – Only used when gauge_type="projected_ref". Geographic coordinates (longitude_deg, latitude_deg) of the first singularity point. The second singularity is placed automatically at the antipodal point. Overrides ref_direction when both are provided. Example — singularities over the Himalayas and the central Pacific:

    singularity_lonlat=(84.0, 28.0)   # Kathmandu region

  • ref_direction (array-like (3,) or None) – Low-level alternative to singularity_lonlat: pass the reference 3-D unit vector r directly. Ignored when singularity_lonlat is provided. Default [1, 0, 0] → singularities at (lon=0°, lat=0°) and (lon=180°, lat=0°).

  • cell_ids (array-like or None) – Pixel indices (NESTED) for partial-sky. None = full sphere.

  • nest (bool, default True) – NESTED pixel ordering.

  • use_norm (bool, default False) – Apply GroupNorm + ReLU after convolution.

  • ellipsoid (str, default "WGS84") – Reference ellipsoid for healpix_geo coordinate conversions.

  • device, dtype (torch device / dtype.)

Examples

Ocean model — singularities over Africa and its antipode (Pacific):

>>> conv = HealPixConv(
...     level=6, in_channels=1, out_channels=16,
...     gauge_type="projected_ref",
...     singularity_lonlat=(20.0, 5.0),   # Gulf of Guinea coast
... )

Atmosphere — singularities over the Pacific and Indian oceans:

>>> conv = HealPixConv(
...     level=6, in_channels=1, out_channels=16,
...     gauge_type="projected_ref",
...     singularity_lonlat=(-160.0, 0.0),  # central Pacific
... )
device#
dtype = Ellipsis#
ellipsoid = ''#
level#
nside#
in_channels#
out_channels#
kernel_sz = 3#
G#
P = 9#
nest = True#
gauge_type = 'phi'#
partial#
K#
weight#
bias#
use_norm = False#
set_kernel(W, bias=None, requires_grad=False)[source]#

Replace the learnable kernel with a fixed (or re-initialised) array.

Parameters:

  • W (array-like) – Shape [C_in, C_out, P] — same kernel broadcast over all G. Shape [G, C_in, C_out, P] — per-gauge kernels.

  • bias (array-like or None) – Shape [G * C_out]. None resets bias to zero.

  • requires_grad (bool, default False)

forward(x)[source]#

Apply gauge-equivariant spherical convolution.

Parameters:

x (array-like, shape [N], [B, N] or [B, C_in, N])

Returns:

y (same type, shape [G*C_out, N] or [B, G*C_out, N])

Notes

Optimisation vs v1:

v1 (single giant gather):

pos_flat  = pos.reshape(-1)              # [G*4*K*P]
vals_flat = t_sorted.index_select(2, pos_flat)  # [B, C_in, G*4*K*P]
vals      = vals_flat.view(B, C_in, G, 4, K, P)
gathered  = (vals * w).sum(dim=3)        # [B, C_in, G, K, P]
y         = einsum("bcgkp,gcop->bgok", gathered, W)

v2 (4-neighbor accumulation + bmm):

# OPT 2: accumulate over 4 neighbors instead of one giant alloc
gathered = zeros(B, C_in, G, K, P)
for j in range(4):                       # tiny loop, only 4 iters
    pj = pos[j].reshape(-1)              # [G*K*P]  — contiguous (OPT 1)
    vf = t_sorted.index_select(2, pj)    # [B, C_in, G*K*P]
    gathered += vf.view(...) * wn[j]     # in-place accumulate

# OPT 3: bmm dispatches to cuBLAS instead of generic einsum
g_mat = gathered.permute(2,0,3,1,4).reshape(G, B*K, C_in*P)
W_mat = W.permute(0,1,3,2).reshape(G, C_in*P, C_out)
y = bmm(g_mat, W_mat).reshape(G,B,K,C_out).permute(1,0,3,2)

Peak intermediate memory: B·C_in·G·K·P (v2) vs B·C_in·G·4·K·P (v1) — a 4x reduction.

singularity_info() → str[source]#

Return a human-readable description of the gauge singularities.

For “phi” and “cosmo”: singularities are always at the geographic poles. For “projected_ref”: two antipodal singularities at plus/minus r. For “two_ref”: four singularities at {plus/minus r1, plus/minus r2} (index +1 each) plus index-(-1) singularities at the geographic poles.

extra_repr()[source]#

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.