File size: 6,505 Bytes
37faa46 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | ---
license: mit
tags:
- pytorch
- spherical-cnn
- cmb
- healpix
- astronomy
- cosmology
library_name: pytorch
---
# torch-harmonics-healpix
Spectral CNN models for CMB parameter estimation on the HEALPix sphere, bridging [torch-harmonics](https://github.com/Philippe7427/torch-harmonics) with HEALPix maps.
These models reproduce and improve upon the benchmarks from [Krachmalnicoff & Tomasi (2019)](https://arxiv.org/abs/1902.04083), which originally used the pixel-space [NNhealpix](https://github.com/NToulis/nnhealpix) architecture.
**Source code:** `https://github.com/zonca/torch-harmonics-healpix`
## Model Summary
| Model | File | Task | Input | Output | Error | Params |
|-------|------|------|-------|--------|-------|--------|
| SpectralCNN T1 | `models/test1_v2_fix_noise0.pt` | β_peak estimation | T map | β_peak | 1.27% | 6.4M |
| SpectralCNN T2 | `models/test2_v2_fix_fsky1.0.pt` | β_Ep / β_Bp estimation | Q, U, mask | [β_Ep, β_Bp] | 1.69% / 1.53% | 9.8M |
| SpectralCNN T3 | `models/test3_v2_fix.pt` | Ο estimation | Q, U, mask | Ο | 3.76% | 9.8M |
## Architecture
**SpectralCNN** performs convolution in harmonic space instead of pixel space:
1. **HEALPix β Equiangular** resampling (bilinear interpolation)
2. **SHT** (Spherical Harmonic Transform) via torch-harmonics
3. **Learned spectral weights** β complex-valued 1Γ1 convolutions on (β, m) coefficients
4. **ISHT** (Inverse SHT) back to pixel space
5. **Equiangular β HEALPix** resampling
The network stacks multiple `SpectralConvBlock` layers (SHT β learned weights β ISHT + residual) followed by global average pooling and a linear head.
**Key advantage over pixel-space CNNs:** The spectral prior enforces physical smoothness in harmonic space, which is especially powerful for polarization estimation where E/B modes have characteristic spectral signatures.
### Design Decisions
- **Inpainting for partial sky:** Masked pixels are replaced with the observed-pixel mean before SHT to prevent mode-coupling artifacts
- **Shared mask:** Train/val/test use the same mask geometry; different masks corrupt spectral coefficients
- **Scalar SHT with Q/U stacking:** torch-harmonics v0.8.0 VectorSHT is slow, so Q/U are stacked as independent channels
See [ARCHITECTURE.md](https://github.com/zonca/torch-harmonics-healpix/blob/main/ARCHITECTURE.md) for the full comparison with NNhealpix.
## Benchmark Results
### Test 2 β Polarization (SpectralCNN dominates)
| f_sky | SpectralCNN (β_Ep / β_Bp) | NNhealpix | Improvement |
|-------|---------------------------|-----------|-------------|
| 1.0 | **1.69% / 1.53%** | 2.7% / 2.7% | 37% / 43% |
| 0.5 | **1.95% / 1.91%** | 3.9% / 3.9% | 50% / 51% |
| 0.2 | **2.15% / 2.17%** | 5.3% / 5.3% | 59% / 59% |
| 0.1 | **2.56% / 2.70%** | 6.4% / 6.4% | 60% / 58% |
| 0.05 | **3.01% / 3.11%** | 8.4% / 8.4% | 64% / 63% |
### Test 3 β Optical depth Ο
| Method | Ο % error |
|--------|----------|
| MCMC (paper) | 2.8% |
| **SpectralCNN** | **3.76%** |
| NNhealpix | 4.0% |
### Test 1 β Scalar maps (noise-free only)
| Ο_n | SpectralCNN | NNhealpix |
|-----|------------|-----------|
| 0 | **1.27%** | 1.3% |
| 5 | 3.58% | **2.9%** |
SpectralCNN wins for noise-free data but loses at high noise because SHT spreads local noise globally, while pixel-space convolution naturally filters it.
See [BENCHMARKS.md](https://github.com/zonca/torch-harmonics-healpix/blob/main/BENCHMARKS.md) for full tables including MCMC baselines.
## Usage
### Installation
```bash
uv venv .venv --python 3.11
source .venv/bin/activate
uv pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
uv pip install torch-harmonics==0.8.0 --no-deps
uv pip install healpy h5py scipy huggingface_hub
uv pip install -e "git+https://github.com/zonca/torch-harmonics-healpix#egg=torch-harmonics-healpix"
```
### Download and Load
```python
import torch
import numpy as np
from huggingface_hub import hf_hub_download
from torch_harmonics_healpix.models import SpectralCNN
# Download model weights
model_path = hf_hub_download(
repo_id="zonca/torch-harmonics-healpix",
filename="models/test2_v2_fix_fsky1.0.pt",
)
# Create model with matching architecture
model = SpectralCNN(
in_channels=3, # Test 1: 1, Test 2/3: 3 (Q, U, mask)
out_channels=1, # Test 1/3: 1, Test 2: 2
nside=16,
hidden_channels=32,
num_blocks=3,
inpaint=False, # True for f_sky < 1.0
)
# Load weights
state_dict = torch.load(model_path, map_location="cpu")
model.load_state_dict(state_dict)
model.eval()
# Run inference on a HEALPix Nside=16 map (3072 pixels)
# Stack [Q, U, mask] as 3 channels
input_tensor = torch.from_numpy(
np.stack([q_map, u_map, mask], axis=0).astype(np.float32)
).unsqueeze(0) # [1, 3, 3072]
with torch.no_grad():
prediction = model(input_tensor)
print(f"Predicted parameter: {prediction.item():.4f}")
```
## Training
To retrain from scratch (e.g., for different noise levels or f_sky values):
```bash
# Test 1: β_peak from T maps
python scripts/train_test1_v2.py --noise_std 0 --output results/test1_noise0.json
# Test 2: β_Ep/β_Bp from Q/U maps
python scripts/train_test2_v2.py --f_sky 0.5 --output results/test2_fsky0.5.json
# Test 3: Ο estimation (requires: pip install camb)
python scripts/train_test3_v2.py --f_sky 1.0 --output results/test3.json
```
Each script saves both `results/*.json` (metrics) and `results/*.pt` (model weights).
## Limitations
- **HEALPix Nside=16 only** (3072 pixels) β not tested at higher resolutions
- **torch-harmonics v0.8.0** β VectorSHT too slow; uses scalar SHT with stacked Q/U channels
- **No explicit E/B separation** β relies on spectral prior to learn E/B structure implicitly
- **Noise sensitivity** β SHT spreads local noise globally; pixel-space CNNs are more robust for high-noise scalar maps
- **Full-sky pre-trained models** β partial-sky models require retraining with `inpaint=True`
## Citation
If you use these models, please cite:
```bibtex
@article{krachmalnicoff2019,
title={Convolutional Neural Networks on the {HEALPix} sphere: a pixel-based approach for CMB data analysis},
author={Krachmalnicoff, N. and Tomasi, M.},
journal={Astronomy \& Astrophysics},
volume={624},
pages={A97},
year={2019},
doi={10.1051/0004-6361/201834952},
url={https://arxiv.org/abs/1902.04083}
}
```
## License
MIT
|