Metadata-Version: 2.4
Name: healformers
Version: 0.2.3
Summary: Mask-Aware HEALPix Transformers
Author: lalalabox
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/lalalabox/healformers
Project-URL: Repository, https://github.com/lalalabox/healformers
Project-URL: Documentation, https://github.com/lalalabox/healformers#readme
Project-URL: Changelog, https://github.com/lalalabox/healformers/blob/main/CHANGELOG.md
Project-URL: Issues, https://github.com/lalalabox/healformers/issues
Keywords: healpix,weak lensing,mass mapping,transformers,astronomy
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Astronomy
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: <3.13,>=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: datasets<5,>=4
Requires-Dist: healpy<1.20,>=1.18
Requires-Dist: huggingface-hub<1,>=0.25
Requires-Dist: numpy<2.3,>=1.26
Requires-Dist: pandas<2.3,>=2.2
Requires-Dist: pyarrow<21,>=18
Requires-Dist: scipy<1.16,>=1.13
Requires-Dist: torch<2.7,>=2.5
Requires-Dist: transformers==4.45.2
Provides-Extra: notebook
Requires-Dist: contourpy<1.3,>=1.2; extra == "notebook"
Requires-Dist: jupyterlab<5,>=4; extra == "notebook"
Requires-Dist: matplotlib<4,>=3.9; extra == "notebook"
Requires-Dist: pillow<11,>=10; extra == "notebook"
Requires-Dist: pyzmq<27,>=26; extra == "notebook"
Provides-Extra: all
Requires-Dist: accelerate<2,>=0.34; extra == "all"
Requires-Dist: astropy<8,>=6; extra == "all"
Requires-Dist: contourpy<1.3,>=1.2; extra == "all"
Requires-Dist: jupyterlab<5,>=4; extra == "all"
Requires-Dist: matplotlib<4,>=3.9; extra == "all"
Requires-Dist: peft==0.13.2; extra == "all"
Requires-Dist: pillow<11,>=10; extra == "all"
Requires-Dist: pyccl<4,>=3; extra == "all"
Requires-Dist: pyzmq<27,>=26; extra == "all"
Requires-Dist: schedulefree<2,>=1.4; extra == "all"
Dynamic: license-file

# HealFormers: Mask-Aware HEALPix Transformers

![HealFormer model architecture](imgs/healformer_architecture.jpg)

[![PyPI](https://img.shields.io/pypi/v/healformers.svg)](https://pypi.org/project/healformers/)
[![Python](https://img.shields.io/pypi/pyversions/healformers.svg)](https://pypi.org/project/healformers/)
[![CI](https://github.com/lalalabox/healformers/actions/workflows/python-publish.yml/badge.svg)](https://github.com/lalalabox/healformers/actions/workflows/python-publish.yml)
[![License](https://img.shields.io/pypi/l/healformers.svg)](LICENSE)
[![DOI](https://img.shields.io/badge/DOI-10.1103%2Fkc9z--jllp-blue.svg)](https://doi.org/10.1103/kc9z-jllp)

> 🗺️ 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

```bash
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](#installation) ·  
[Inference](#inference-quickstart) · [Released diagnostics](#released-diagnostic-figures) ·  
[Data and checkpoints](#public-data-and-checkpoints) · [Paper](#paper)

## Why HealFormer?

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

| 🚩 Research challenge | 🎯 HealFormer advantage |
| --- | --- |
| Distortion from planar projections | ✅ Native `NESTED` HEALPix operations—no flat-sky reprojection |
| Noise-dominated small scales | ✅ Nearly unbiased power and higher small-scale phase correlation than linear estimators |
| Mask boundaries and holes | ✅ Patch-aware mask tokens inpaint inner-edge pixels instead of zero-filling them |
| Cosmology and footprint generalization | ✅ Planck-only training generalize across unseen cosmologies; one mixed checkpoint covers four survey masks |
| High-resolution spherical maps | ✅ `Nside=256` to `Nside=1024` scaling with projection and parameter-efficient LoRA |

## 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

![Noisy mass-mapping comparison at Nside 256](imgs/noise_effect_nside256_maskDECaLS_noiseTrue.jpg)

**Takeaway:** on the representative masked, noisy `Nside=256` map in
[paper Figure 2](https://arxiv.org/html/2603.25471v1#S4.SS2), HealFormer has the
smallest residuals, the tightest pixel-level agreement, nearly unbiased power,
and small-scale phase correlation above the black linear-estimator ceiling.

| Metric | HealFormer | Wiener filter | Spherical KS |
| --- | ---: | ---: | ---: |
| NRMSE | 0.18 | 0.51 | 0.63 |
| Pearson correlation | 0.98 | 0.86 | 0.84 |
| Multiplicative slope | 0.98 | 0.84 | 1.17 |
| Jensen--Shannon divergence | 0.0004 | 0.0079 | 0.0061 |

These are Figure 2 values for one representative map, not aggregate metrics.
The paper's Wiener result uses [DANTE](https://github.com/doogesh/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

![Generalization from Planck 2018 training cosmology to unseen cosmologies](imgs/generalization_multiverse_nside256.jpg)

**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

![High-resolution HealFormer reconstruction at Nside 1024](imgs/denoising_nside1024_maskDECaLS_noiseTrue.jpg)

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

![Inner-edge inpainting at Nside 1024](imgs/inpainting_inner_edge_nside1024.jpg)

**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

![Residual comparison across four survey masks](imgs/compare_allMask_residual_nside256.jpg)

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

> [!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.

## Installation

Python 3.11 and 3.12 are supported.

### pip

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

```bash
python -m pip install healformers
```

Notebook and plotting:

```bash
python -m pip install "healformers[notebook]"
```

Everything, including training and LoRA export:

```bash
python -m pip install "healformers[all]"
```

Extras are intentionally limited to `notebook` and `all`.

### uv

Add it to a project:

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

Locked source environment:

```bash
uv sync --frozen --extra notebook
```

Editable pip install:

```bash
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:

```bash
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:

```python
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:

```bash
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:

```bash
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`):

```bash
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

| Artifact | Hugging Face repository | Immutable revision |
| --- | --- | --- |
| Streaming physical-map triples, fixed masks, and evaluation bundles for `nside256` and `nside1024` | [`lalala404/healformer-multiverse-maps`](https://huggingface.co/datasets/lalala404/healformer-multiverse-maps) | `a030c368a2839037710274950065fa735b2d7857` |
| Fixed mixed-mask `Nside=256` checkpoint used unchanged across four fixed masks | [`lalala404/healformer-nside256-mixed`](https://huggingface.co/lalala404/healformer-nside256-mixed) | `abc4bbefc2cc391bf12cd45500636e1bd49c28ce` |
| Fixed DECaLS+DES mask `Nside=1024` checkpoint | [`lalala404/healformer-nside1024-decals`](https://huggingface.co/lalala404/healformer-nside1024-decals) | `d80d07062c0478aaad5dfaa95e2d47fecbc8278c` |

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:

```python
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`](examples/run_inference.py).
For the KS comparison, see
[`examples/compare_kaiser_squires.py`](examples/compare_kaiser_squires.py).
Dataset preparation, normalization, inference, and evaluation examples are
collected in [`examples/README.md`](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.

| Recipe `adaptation` | Behavior |
| --- | --- |
| `scratch` | Initialize from a config template |
| `full` | Load a checkpoint and update every weight |
| `source_lora` | Train the checkpoint-compatible PEFT policy |

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

```bash
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:

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

The [`Nside=1024` source-derived recipe](recipes/training/nside1024_source_derived.json)
requires the sealed best checkpoint from a completed parent run. The
[`Nside=256` recipe](recipes/training/nside256_source_derived.json) 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](examples/README.md#5-train-with-hugging-face-trainer).

## 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](https://doi.org/10.1103/kc9z-jllp)
- arXiv: [2603.25471](https://arxiv.org/abs/2603.25471)
- PDF: [arXiv:2603.25471](https://arxiv.org/pdf/2603.25471)

Use [`CITATION.cff`](CITATION.cff) for software metadata and cite the paper as:

```bibtex
@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`](LICENSE), [`NOTICE`](NOTICE), and the
[dataset card](release/cards/dataset/README.md) 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.
