HealFormers: Mask-Aware HEALPix Transformers

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

HealFormer is a mask-aware HEALPix Transformer for noisy, incomplete sky maps. The package includes public checkpoints, streaming data, training, and reproducible pixel- and harmonic-space evaluation.

Start here: Install · Run inference · Reproduce figures · Data and checkpoints · Paper

Why HealFormers?

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

Released evaluations use fixed masks, NESTED ordering, and no map rotation.

Results

For the masked, noisy Nside=256 example shown above, the paper's Figure 2 reports:

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 reproduces HealFormer and spherical KS on pinned public data, including 100-sample power ratios and cross-correlations with one-sigma bands. Its single-map values differ because the public sample is pinned separately.

At Nside=1024, the public workflow uses the fixed DECaLS-labeled footprint (combined DECaLS+DES) and the projection checkpoint. At Nside=256, the same checkpoint is evaluated on all four fixed survey masks:

[!NOTE] The footprint labeled “DECaLS” in the paper and released artifacts is the combined DECaLS+DES footprint. This is a label-recording offset; the stored mask, calculations, methods, and conclusions are unchanged.

HealFormers ships the HealFormer and spherical KS paths. Use DANTE for the paper's Wiener-filter path.

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 --locked --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, examples/hub_inference.py is a thin wrapper around the same installed command implementation.

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.

Reproduce the diagnostic figures

Open the executable tutorial:

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;
• visible-pixel hist2d plots using the turbo color map; and
• 100-sample power-spectrum ratios and cross-correlation coefficients with shaded one-sigma variation for both public resolutions in one subplot grid.

Nside=256 uses one checkpoint across KiDS, DES, DECaLS+DES, and Planck. Nside=1024 uses only the fixed DECaLS+DES footprint with projection. Published diagnostics never rotate masks.

Fast exact rendering

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. Recomputation 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

One Hugging Face Trainer script supports three methods:

LoRA:

uv run --extra all python examples/train_mass_mapping.py \
  --model lalala404/healformer-nside256-mixed \
  --dataset path/to/dataset-artifact \
  --output-dir lora-adapter

Use --method full for full fine-tuning. Scratch uses --model only as an architecture template:

uv run --extra all python examples/train_mass_mapping.py \
  --method scratch \
  --model path-or-hub-config-template \
  --dataset path/to/dataset-artifact \
  --nside 512 \
  --num-channels-in 4 \
  --num-channels-out 2 \
  --output-dir scratch-checkpoint

Scratch can infer omitted architecture options from the Dataset Artifact. Full and LoRA runs must match the checkpoint shapes.

Set resolution, mask, and noise in the Dataset Artifact recipe. The training script does not regenerate or rotate them. Edge-aware terms follow the recipe's edge_layers; pass --no-add-edge to train without them. 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.