zestimatr 0.1.4

Spectroscopic redshift estimation with uncertainty quantification

zestimatr

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

pip install zestimatr

Or install from source:

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.

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

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

# 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

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

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

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

Pretrained Model

A pretrained checkpoint trained on JADES DR4 (52,647 spectra) is available on Hugging Face.

Download it automatically:

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
• Lau, Marie Wingyee
• Michele Woodland