You can deploy this model and start earning money today!

aryana-haghjoo/
zestimatr

Name: zestimatr
Availability: InStock

Views

No views yet

zestimatr — Spectroscopic Redshift Estimator

A 1D convolutional neural network that predicts galaxy redshifts from high-resolution JWST/NIRSpec spectra, with calibrated uncertainty estimates.

Model Details

Architecture: ZHead1D — 6 residual 1D conv blocks (hidden dim 128, dropout 0.2) + MLP head
Input: Per-spectrum normalized high-resolution flux, 2500 pixels spanning 1–5 µm
Output: Predicted redshift (point estimate) + predicted 1σ uncertainty
Training data: JADES DR4 — 42,105 augmented spectra (zestimatr-jades-dr4)
Loss: Heteroscedastic Gaussian NLL

Performance

Measured on the 502 held-out objects in eval_DR4.npz. These objects are split off at the object level before augmentation, so no augmented copy of an evaluation galaxy appears in training.

Metric	Value
MAE	0.478
RMSE	0.951
Median \|Δz\|	0.138
Median \|Δz\|/(1+z)	0.034
NMAD	0.205
Outlier rate (\|Δz\|/(1+z) > 0.15)	21.7%
Calibration std	1.16
1σ coverage	78.7%
2σ coverage	94.2%

The model is roughly calibrated (calibration std 1.16 against a target of 1.0) but has a substantial outlier population: about one in five galaxies is predicted badly, typically through a catastrophic line misidentification. Treat the predicted σ as meaningful — high-σ predictions are the unreliable ones.

Note on earlier figures. A previous version of this card quoted MAE 0.141 and a 5.7% outlier rate on a "10,530 spectra" validation set. Those came from an earlier data split that augmented before splitting, so augmented copies of validation galaxies leaked into training and the numbers were optimistic. The table above supersedes them.

Usage

python

import numpy as np
from zestimatr import download_pretrained, load_model, predict

ckpt = download_pretrained()          # pulls best_zhead_hires.pth from this repo
zhead, norm_params = load_model(ckpt)

# flux: (N, L) array; wavelength in microns. Resampled to the 2500-pt grid.
out = predict(flux, zhead, norm_params, wavelength=wavelength)
print(out["z_pred"], out["z_uncertainty"])

Full worked example against the public evaluation set:

python

import numpy as np
from huggingface_hub import hf_hub_download
from zestimatr import download_pretrained, load_model, predict

zhead, norm_params = load_model(download_pretrained())

path = hf_hub_download("aryana-haghjoo/zestimatr-jades-dr4",
                       "eval_DR4.npz", repo_type="dataset")
d = np.load(path, allow_pickle=True)

out = predict(d["flux_high"][:5], zhead, norm_params,
              wavelength=d["wavelength_high"])
for zp, zs, zt in zip(out["z_pred"], out["z_uncertainty"], d["z"][:5]):
    print(f"z_pred={zp:.3f} +/- {zs:.3f}   (true {zt:.3f})")

Install with pip install zestimatr.

Limitations

Trained only on JADES GOODS-N / GOODS-S NIRSpec spectra; behavior on other instruments or surveys is untested.
The training set spans z ≈ 0–11, with augmentation extending the nominal range slightly below zero. Predictions are constrained to that interval and should not be trusted outside it.
The ~22% outlier rate means this is not a drop-in replacement for template fitting on individual objects of interest.

aryana-haghjoo/zestimatr