healpix_analyse.convol ====================== .. py:module:: healpix_analyse.convol .. autoapi-nested-parse:: convol.py (optimised v3) ========================= 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 ---------- .. autoapisummary:: healpix_analyse.convol.ArrayLike Classes ------- .. autoapisummary:: healpix_analyse.convol.HealPixConv Module Contents --------------- .. py:data:: ArrayLike .. py:class:: HealPixConv(level: int, in_channels: int, out_channels: int, kernel_sz: int = 3, n_gauges: int = 1, gauge_type: str = 'phi', singularity_lonlat: Optional[tuple[float, float]] = 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) Bases: :py:obj:`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** (:py:class:`int`) -- Grid4Earth/HEALPix level (integer >= 0). The internal HEALPix resolution is ``nside = 2**level``. * **in_channels** (:py:class:`int`) -- Number of input channels C_in. * **out_channels** (:py:class:`int`) -- Number of output channels per gauge C_out. Total output channels = n_gauges × out_channels. * **kernel_sz** (:py:class:`int`, *default* ``3``) -- Odd integer ≥ 1. P = kernel_sz² stencil points. * **n_gauges** (:py:class:`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** (:py:class:`(float`, :py:class:`float)` or :py:obj:`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** (:py:class:`array-like (3,)` or :py:obj:`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** (:py:class:`array-like` or :py:obj:`None`) -- Pixel indices (NESTED) for partial-sky. None = full sphere. * **nest** (:py:class:`bool`, *default* :py:obj:`True`) -- NESTED pixel ordering. * **use_norm** (:py:class:`bool`, *default* :py:obj:`False`) -- Apply GroupNorm + ReLU after convolution. * **ellipsoid** (:py:class:`str`, *default* ``"WGS84"``) -- Reference ellipsoid for healpix_geo coordinate conversions. * **device, dtype** (:py:class:`torch device / dtype.`) .. admonition:: 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 ... ) .. py:attribute:: device .. py:attribute:: dtype :value: Ellipsis .. py:attribute:: ellipsoid :value: '' .. py:attribute:: level .. py:attribute:: nside .. py:attribute:: in_channels .. py:attribute:: out_channels .. py:attribute:: kernel_sz :value: 3 .. py:attribute:: G .. py:attribute:: P :value: 9 .. py:attribute:: nest :value: True .. py:attribute:: gauge_type :value: 'phi' .. py:attribute:: partial .. py:attribute:: K .. py:attribute:: weight .. py:attribute:: bias .. py:attribute:: use_norm :value: False .. py:method:: set_kernel(W, bias=None, requires_grad=False) Replace the learnable kernel with a fixed (or re-initialised) array. :Parameters: * **W** (:py:class:`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** (:py:class:`array-like` or :py:obj:`None`) -- Shape [G * C_out]. None resets bias to zero. * **requires_grad** (:py:class:`bool`, *default* :py:obj:`False`) .. py:method:: forward(x) Apply gauge-equivariant spherical convolution. :Parameters: **x** (:py:class:`array-like`, :py:class:`shape [N]`, :py:class:`[B`, :py:class:`N]` or :py:class:`[B`, :py:class:`C_in`, :py:class:`N]`) :returns: **y** (:py:class:`same type`, :py:class:`shape [G*C_out`, :py:class:`N]` or :py:class:`[B`, :py:class:`G*C_out`, :py:class:`N]`) .. admonition:: 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. .. py:method:: singularity_info() -> str 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. .. py:method:: extra_repr() 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.