Metadata-Version: 2.4
Name: healformers
Version: 0.2.1
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)

> 🗺️ 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](#installation) ·
[Run inference](#inference-quickstart) · [Reproduce figures](#reproduce-the-diagnostic-figures) ·
[Data and checkpoints](#public-data-and-checkpoints) · [Paper](#paper)

## Why HealFormers?

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

| Need | What HealFormer provides |
| --- | --- |
| Spherical geometry | Native `NESTED` HEALPix maps; no flat-sky projection |
| Missing sky regions | The survey mask is part of the model input |
| Mask generalization | One `Nside=256` checkpoint for KiDS, DES, DECaLS+DES, and Planck |
| High resolution | A projection checkpoint for the fixed `Nside=1024` footprint |
| Adaptation | Scratch, full fine-tuning, and PEFT LoRA |
| Reproduction | Pinned Hub revisions, checksums, KS baseline, and an executed notebook |

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

## Results

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

For the masked, noisy `Nside=256` example shown above, the
[paper's Figure 2](https://arxiv.org/html/2603.25471v1#S4.SS2) reports:

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

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

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:

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

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

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

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.

## Reproduce the diagnostic figures

Open the executable tutorial:

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

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

| 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-labeled combined DECaLS+DES projection `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

One Hugging Face Trainer script supports three methods:

| Method | Behavior |
| --- | --- |
| `scratch` | Initialize from a config template |
| `full` | Load a checkpoint and update every weight |
| `lora` | Train a PEFT adapter; this is the default |

LoRA:

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

```bash
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](examples/README.md#1-build-deterministic-dataset-splits).

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