Metadata-Version: 2.4
Name: zestimatr
Version: 0.1.3
Summary: Spectroscopic redshift estimation with uncertainty quantification
License-Expression: MIT
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Operating System :: OS Independent
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: torch
Requires-Dist: matplotlib
Requires-Dist: tqdm
Requires-Dist: huggingface_hub
Requires-Dist: pandas
Provides-Extra: train
Requires-Dist: wandb; extra == "train"
Requires-Dist: pyyaml; extra == "train"
Provides-Extra: preprocess
Requires-Dist: astropy; extra == "preprocess"
Requires-Dist: scipy; extra == "preprocess"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Dynamic: license-file

# zestimatr

[![PyPI](https://img.shields.io/pypi/v/zestimatr)](https://pypi.org/project/zestimatr/)
[![Python](https://img.shields.io/pypi/pyversions/zestimatr)](https://pypi.org/project/zestimatr/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
[![Documentation](https://readthedocs.org/projects/zestimatr/badge/?version=latest)](https://zestimatr.readthedocs.io/)
[![Hugging Face](https://img.shields.io/badge/Model-Hugging%20Face-orange)](https://huggingface.co/aryana-haghjoo/zestimatr)
[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.20835578.svg)](https://doi.org/10.5281/zenodo.20835578)
[![arXiv](https://img.shields.io/badge/arXiv-2603.18357-b31b1b)](https://arxiv.org/abs/2603.18357)
[![A rectangular badge, half black half purple containing the text made at Code Astro](https://img.shields.io/badge/Made%20at-Code/Astro-blueviolet.svg)](https://semaphorep.github.io/codeastro/)

Spectroscopic redshift estimation from high-resolution galaxy spectra, with uncertainty quantification.

`zestimatr` is a Python package that uses a residual 1D convolutional neural network with an MLP head to predict redshifts directly from high-resolution spectral flux. The model outputs both a point estimate and a calibrated uncertainty (predicted standard deviation) for each spectrum.

## Installation

```bash
pip install zestimatr
```

Or install from source:

```bash
git clone https://github.com/aryana-haghjoo/zestimatr.git
cd zestimatr
pip install -e .
```

## Quick Start

You can run the following command below to predict the redshift for a given spectrum npz file.

```
python run_zestimatr.py $PATH_to_data_file
```

The code below can also be pasted into a notebook or .py file to predict the redshift for a spectrum. Just pass your flux array (and optionally the corresponding wavelength array) — normalization and resampling are handled automatically.

```python
import numpy as np
import zestimatr

# Load a spectrum
data = np.load("galaxy300_spectrum.npz")
flux = data["flux_high"]
wavelength = data["wavelength_high"]

# Download and load pretrained model from Hugging Face
checkpoint_path = zestimatr.download_pretrained()
zhead, norm_params = zestimatr.load_model(checkpoint_path)

# Predict — pass wavelength to automatically resample to the training grid
result = zestimatr.predict(flux, zhead, norm_params, wavelength=wavelength)

print(f"Predicted: z = {result['z_pred']:.4f} +/- {result['z_uncertainty']:.4f}")
```

If your spectrum is already on the training wavelength grid (2500 points, 1–5 μm), you can omit the `wavelength` argument. For batch prediction, pass a 2-D flux array of shape `(N, L)`:

```python
# Batch prediction
flux_batch = data["flux_high"]  # shape (N, L)
results = zestimatr.predict(flux_batch, zhead, norm_params, wavelength=wavelength)
print(results["z_pred"])  # shape (N,)
```

## Emission Line Detection

After estimating a redshift, you can detect and visualize emission lines in the spectrum:

```python
# Detect lines — returns a pandas DataFrame
lines = zestimatr.detect_emission_lines(wavelength, flux, result["z_pred"])
print(lines)

# Plot spectrum with emission lines marked as dashed vertical lines
zestimatr.plot_spectrum(wavelength, flux, z=result["z_pred"])
```

The built-in catalog includes 16 common rest-frame lines (Ly-alpha, H-alpha, H-beta, [O II], [O III], [N II], [S II], and more). Detection uses a local peak-finding approach with a configurable `sigma_thresh` (default 3.0).

## Metrics

`zestimatr` provides two evaluation functions:

- **`compute_metrics(z_pred, z_true)`** -- accuracy metrics:

  - MAE, RMSE, NMAD
  - Median |dz|/(1+z)
  - Outlier rate (fraction with |dz|/(1+z) > 0.15)
- **`compute_calibration_metrics(z_pred, z_true, z_uncertainty)`** -- uncertainty calibration:

  - Calibration std and mean of normalized residuals
  - 1/2/3-sigma coverage fractions
  - Median predicted uncertainty

## For Developers

Most users only need the base install above. The following is for retraining or rebuilding the dataset from raw JADES data.

### Extra dependencies

```bash
# Training only
pip install -e ".[train]"

# Data preprocessing only (astropy, scipy)
pip install -e ".[preprocess]"

# Everything
pip install -e ".[train,preprocess]"
```

### Data preprocessing

Prepare train/eval datasets from JADES DR4 FITS files. The pipeline splits at the object level before augmentation to prevent data leakage:

```bash
python scripts/prepare_dataset.py --jades_dir /path/to/JADES_data/DR4
```

This produces `data/train_DR4.npz` (augmented training set) and `data/eval_DR4.npz` (held-out evaluation set).

### Training

```bash
python scripts/train.py --train_data data/train_DR4.npz --eval_data data/eval_DR4.npz --wandb_mode online
```

Key training options:

| Flag             | Default | Description                              |
| ---------------- | ------- | ---------------------------------------- |
| `--hidden_dim` | 128     | Conv block hidden dimension              |
| `--num_blocks` | 6       | Number of residual conv blocks           |
| `--dropout`    | 0.2     | Dropout rate                             |
| `--epochs`     | 200     | Training epochs                          |
| `--lr`         | 3e-4    | Learning rate                            |
| `--batch_size` | 32      | Batch size                               |
| `--wandb_mode` | online  | `online`, `offline`, or `disabled` |

Training logs and plots are synced to [Weights &amp; Biases](https://wandb.ai/).

## Pretrained Model

A pretrained checkpoint trained on JADES DR4 (52,647 spectra) is available on [Hugging Face](https://huggingface.co/aryana-haghjoo/zestimatr).

Download it automatically:

```python
import zestimatr

path = zestimatr.download_pretrained()
zhead, norm_params = zestimatr.load_model(path)
```

| Metric             | Value |
| ------------------ | ----- |
| MAE                | 0.141 |
| RMSE               | 0.323 |
| Median\|dz\|/(1+z) | 0.012 |
| Outlier rate       | 5.7%  |
| Calibration std    | 0.84  |

## Data Format

Input `.npz` files should contain:

- `flux_high` -- high-resolution spectral flux, shape `(N, L)` for datasets or `(L,)` for a single spectrum
- `z` -- ground truth redshifts, shape `(N,)` or scalar (also accepts keys: `redshift`, `z_true`, `z_spec`)
- `wavelength_high` -- wavelength array, shape `(L,)`. If provided (along with `wavelength` to `predict()`), spectra are automatically resampled onto the model's training grid (2500 points, 1–5 μm). If omitted, flux is assumed to already be on the training grid.

Optional keys: `flux_high_err`, `id`, `ra`, `dec`.

## Project Structure

```
zestimatr/
├── src/zestimatr/          # Package source
│   ├── model.py             # ZHead1D network + loss function
│   ├── dataset.py           # PyTorch dataset for training
│   ├── metrics.py           # Accuracy and calibration metrics
│   ├── inference.py         # Model loading + prediction pipeline
│   ├── plotting.py          # Validation plots
│   └── emission_lines.py    # Emission line detection + spectrum plotting
├── scripts/
│   ├── prepare_dataset.py   # Data extraction, quality cuts, split, augmentation
│   └── train.py             # Training CLI (not part of the package)
│   └── run_zestimatr.py.    # Predict redshift for a given spectrum
├── tests/
│   └── test_metrics.py      # Unit tests
├── tutorials/
│   ├── quickstart_single_spectrum.ipynb  # Single spectrum tutorial
│   ├── quickstart_batch_spectra.ipynb    # Batch (200 spectra) tutorial
│   ├── galaxy300_spectrum.npz            # Example single spectrum
│   └── sample_200_spectra.npz            # Example 200-spectrum sample
├── pyproject.toml           # Package metadata
└── README.md
```

## License

MIT

## Authors

Code/Astro 2026 Group 2:

- Aryana Haghjoo [![ORCID](https://img.shields.io/badge/ORCID-0009--0006--3071--7143-green?logo=orcid)](https://orcid.org/0009-0006-3071-7143)
- Lau, Marie Wingyee [![ORCID](https://img.shields.io/badge/ORCID-0000--0001--9755--9406-green?logo=orcid)](https://orcid.org/0000-0001-9755-9406)
- Michele Woodland
