HealFormers: Mask-Aware HEALPix Transformers

🗺️ Reconstruct weak-lensing convergence directly on the HEALPix sphere.

HealFormer reconstructs weak-lensing convergence directly from noisy, masked shear on the HEALPix sphere. It avoids flat-sky reprojection and uses learnable mask tokens to jointly denoise and inpaint incomplete spherical maps.

At a glance: 0.18 NRMSE · 1 checkpoint → 4 survey masks · Nside=256 & Nside=1024 · 100-sample harmonic diagnostics

The NRMSE is from one representative masked, noisy map in paper Figure 2.

Try it now

python -m pip install healformers
healformers-infer --nside 256 --device cpu --output prediction.npz

This downloads a pinned representative sample, verifies the matching checkpoint, and writes the truth, prediction, and mask to prediction.npz.

Go deeper: Installation ·
Inference · Released diagnostics ·
Data and checkpoints · Paper

Why HealFormer?

Weak-lensing surveys provide noisy shear over incomplete sky footprints. HealFormer handles the sphere and the mask together:

Results

How to read the harmonic panels

Three curves summarize amplitude and phase recovery:

$$ R_\ell=\frac{C_\ell^{\mathrm{rec}}}{C_\ell^{\mathrm{true}}},\qquad r_\ell=\frac{C_\ell^{\mathrm{rec,true}}} {\sqrt{C_\ell^{\mathrm{rec}}C_\ell^{\mathrm{true}}}},\qquad r_\ell^{\max}=\sqrt{\frac{C_\ell^{\mathrm{true}}} {C_\ell^{\mathrm{true}}+N_\ell}}. $$

This follows the paper's $\kappa$-equivalent convention:

$$ N_\ell=\frac{N_\ell^\gamma}{|D_\ell|^2},\qquad D_\ell=-\sqrt{\frac{(\ell-1)(\ell+2)}{\ell(\ell+1)}}. $$

• $R_\ell=1$ means unbiased reconstructed power.
• $r_\ell=1$ means perfect phase agreement with the true map.
• The black max curve is $r_\ell^{\max}$, the full-sky ceiling for positive linear estimators under isotropic, uncorrelated noise.

The simulated shear-to-convergence relation is linear. The result claimed here is specific: HealFormer retains higher small-scale phase correlation than the tested linear estimators by using a learned simulation prior.

Main result: masked, noisy reconstruction

Takeaway: on the representative masked, noisy Nside=256 map in paper Figure 2, HealFormer has the smallest residuals, the tightest pixel-level agreement, nearly unbiased power, and small-scale phase correlation above the black linear-estimator ceiling.

These are Figure 2 values for one representative map, not aggregate metrics. The paper's Wiener result uses DANTE, which is not bundled here.

The notebook reruns HealFormer and spherical KS on one pinned public map and, by default, renders the released checksummed 100-sample harmonic ensemble with one-sigma bands. Set HEALFORMER_RECOMPUTE_ENSEMBLE=1 for an independent 100-sample recomputation. These released diagnostics are not a reproduction of the paper's Figure 2 values; the public representative sample is pinned separately. Evaluations use fixed masks, NESTED ordering, and no map rotation.

Generalizes beyond the training cosmology

Takeaway: a model trained only at the Planck 2018 cosmology maintains $R_\ell\approx1$ across cosmologies sampled from the KiDS-450 posterior and retains small-scale phase correlation above the linear-estimator ceiling.

This is the paper's out-of-training-distribution cosmology stress test. The released mixed checkpoint is not the Planck-only checkpoint used for this figure.

Scales to high resolution

Takeaway: at Nside=1024, HealFormer retains the cleanest residuals, near-unity slope, nearly unbiased power, and the strongest small-scale phase correlation among the compared methods.

The public workflow uses the fixed DECaLS-labeled footprint (combined DECaLS+DES) and the projection checkpoint.

Inpaints masked pixels at patch boundaries

Inner-edge pixels are the masked pixels in mixed 16×16 patches that contain both visible and masked pixels.

Takeaway: the scatter plots use only masked pixels from mixed patches, not the easier visible region. For point-like holes, HealFormer reaches NRMSE 0.26, $\rho=0.97$, and slope $m=0.96$. For the extended hole, it remains the strongest compared method with NRMSE 0.52, $\rho=0.85$, and $m=0.82$.

This directly tests the mask-token design on pixels without observed shear.

One checkpoint across survey footprints

Takeaway: one released Nside=256 mixed-mask checkpoint is evaluated unchanged across KiDS, DES, DECaLS+DES, and Planck.

Installation

Python 3.11 and 3.12 are supported.

pip

Core inference, Hub streaming, Dataset Artifacts, HEALPix diagnostics, and KS:

python -m pip install healformers

Notebook and plotting:

python -m pip install "healformers[notebook]"

Everything, including training and LoRA export:

python -m pip install "healformers[all]"

Extras are intentionally limited to notebook and all.

uv

Add it to a project:

uv add healformers
uv add "healformers[notebook]"  # notebook + scientific figures
uv add "healformers[all]"       # every optional dependency

Locked source environment:

uv sync --frozen --extra notebook

Editable pip install:

python -m pip install -e ".[notebook]"

Inference quickstart

After the pip install above, stream one public sample, verify the matching checkpoint, and save a map with the installed command:

healformers-infer \
  --nside 256 \
  --device cuda \
  --output healformer-nside256-prediction.npz

The output contains kappa_true, kappa_healformer, and mask_npix. CPU works for Nside=256; CUDA is recommended. Use --nside 1024 for the fixed-mask projection model; CUDA is strongly recommended. From a source checkout, run the same healformers-infer command through uv run.

Python API:

from healformers import get_public_release

release = get_public_release(256)
sample = release.load_representative_sample()
pipeline = release.load_pipeline(device="cuda")

kappa_pred = pipeline(
    sample.shear[0],
    sample.shear[1],
    sample.mask_npix,
)[0, 0]

PublicRelease pins revisions, verifies files, and checks mask/projection rules.

Released diagnostic figures

Open the executable tutorial to rerun representative inference and render the released harmonic diagnostics:

uv run --frozen jupyter lab examples/quickstart_mass_mapping.ipynb

The notebook shows:

• input shear and fixed-mask maps;
• HealFormer and spherical KS reconstructions;
• normalized absolute residual maps, abs((kappa_true - kappa_pred) / sigma_kappa), displayed on [0, 1] with the blue la_blue palette;
• patch-aware inner-edge-pixel hist2d plots using the turbo color map; and
• 100-sample power-spectrum ratios and cross-correlation coefficients with shaded one-sigma variation for the selected public resolution (256: four masks; 1024: fixed DECaLS+DES only).

The 100-sample artifact does not include the per-sky truth spectra required for an ensemble-mean $r_\ell^{\max}$, so the notebook does not overlay a single-sky ceiling on those mean curves.

Published diagnostics never rotate masks.

Fast rendering of released arrays

Default mode renders the published, checksummed 100-sample arrays:

HEALFORMER_NSIDE=256 \
HEALFORMER_DEVICE=cuda \
HEALFORMER_OUTPUT_DIR=artifacts/notebook-nside256 \
  uv run --frozen jupyter nbconvert \
    --to notebook --execute examples/quickstart_mass_mapping.ipynb \
    --output quickstart_mass_mapping.nside256.executed.ipynb

Use HEALFORMER_NSIDE=1024 and a new output directory for high resolution.

Independent 100-sample recomputation

Set one flag to rerun both methods on the release's sealed 100-sample plan (seed 20260813):

HEALFORMER_RECOMPUTE_ENSEMBLE=1 \
HEALFORMER_NSIDE=256 \
HEALFORMER_DEVICE=cuda \
HEALFORMER_OUTPUT_DIR=artifacts/notebook-nside256-recomputed \
  uv run --frozen jupyter nbconvert \
    --to notebook --execute examples/quickstart_mass_mapping.ipynb \
    --output quickstart_mass_mapping.nside256.recomputed.ipynb

Each run writes four PNGs, metrics.json, environment-lock.json, and manifest.sha256. Re-computation also saves the 100-sample arrays. Use an empty output directory. Public mode fails closed on unavailable or invalid Hub inputs. Set HEALFORMER_OFFLINE=1 explicitly for the tiny smoke workflow.

Public data and checkpoints

The stream provides physical (gamma1, gamma2, kappa) maps. Evaluation adds fixed masks and deterministic shape noise. Shards stream on demand. The current dataset, code, and model checkpoints use Apache-2.0.

To inspect a public validation sample without loading a checkpoint:

from healformers import get_public_release

sample = next(iter(get_public_release(256).stream_validation()))
print(sample.sample_id, sample.gamma1.shape, sample.kappa.shape)

Input conventions

• Maps use NESTED HEALPix ordering.
• gamma1, gamma2, and kappa use physical shear/convergence units.
• mask_npix uses 0 = visible, 1 = edge to reconstruct, and 2 = unseen or excluded.
• An input map contains 12 * nside**2 pixels.
• All maps in one batch must use the same mask.

For custom local arrays, see examples/run_inference.py. For the KS comparison, see examples/compare_kaiser_squires.py. Dataset preparation, normalization, inference, and evaluation examples are collected in examples/README.md.

Training

healformers-train is the single public training entry point. It runs transformers.Trainer with a versioned Training Recipe, Schedule-Free RAdam, deterministic resume, validation-selected checkpoints, and sealed provenance.

For source-derived Nside=256 LoRA fine-tuning from the released checkpoint:

uv run --extra all healformers-train fit \
  recipes/training/nside256_source_derived.json \
  path/to/dataset-artifact \
  training-run \
  --model-repo lalala404/healformer-nside256-mixed \
  --revision abc4bbefc2cc391bf12cd45500636e1bd49c28ce

Set adaptation to scratch to use the model source only as a config template, or to full for full fine-tuning. Resume an interrupted epoch-boundary checkpoint with:

uv run --extra all healformers-train resume training-run.incomplete

The Nside=1024 source-derived recipe requires the sealed best checkpoint from a completed parent run. The Nside=256 recipe does not.

Set resolution, mask, and noise in the Dataset Artifact recipe. The training runner does not regenerate or rotate them. The completed Training Run Artifact contains best/, resumable last/, epoch metrics, checksums, and the exact Training Recipe. See the examples guide.

Paper

Yihe Wang and Yu Yu, “Advancing weak lensing mass mapping with a mask-aware HEALPix transformer,” Physical Review D 113, 043553 (2026).

• DOI: 10.1103/kc9z-jllp
• arXiv: 2603.25471
• PDF: arXiv:2603.25471

Use CITATION.cff for software metadata and cite the paper as:

@article{wang2026advancing,
  title={Advancing weak lensing mass mapping with a mask-aware HEALPix transformer},
  author={Wang, Yihe and Yu, Yu},
  journal={Physical Review D},
  volume={113},
  number={4},
  pages={043553},
  year={2026},
  publisher={APS},
  doi={10.1103/kc9z-jllp}
}

License and development disclosure

HealFormers code, current public dataset, and model checkpoints are licensed under Apache-2.0. See LICENSE, NOTICE, and the dataset card for terms and attribution.

AI coding assistants, including Claude and GPT-based tools, were used for code drafting, refactoring, debugging, and documentation. The scientific questions, core ideas, methodological innovations, experimental design, interpretation, and conclusions were developed and decided by the human authors, who remain responsible for the work.