2.1 — EDFS photo-z estimation

Authors: Xiangchong Li, Prakruth Adari, Anja von der Linden LSST Science Pipelines: Weekly 2025_49 Kernel: Python-dp1

Purpose

Run / load 6-band FlexZBoost photo-z for the EDFS field and produce the cached per-source point estimates and PDFs that the cluster notebooks (2_2_edfs_erass_cluster1.ipynb, 2_3_edfs_erass_cluster2.ipynb) consume. No cluster lensing happens here — this notebook is purely photo-z.

Steps

1. Load the EDFS measurement table (rail_data/data_edfs_with_mag.hdf5) and the ECDFS training table.
2. Per-band color histograms (g-r, r-i, i-z, z-y) train vs EDFS data.
3. Show i_mag vs FPFS trace (m00+m20)/m00 to motivate the size cut; apply mag + trace cuts.
4. Either run xlens.catalog.redshift.flexzboostEstimator.get_z on the EDFS source array (toggle the if False: flag) or load the cached data_edfs_redshift.hdf5 + data_edfs_redshift_pdfs.fits.
5. Plot the mean p(z) of the kept sample.

Inputs

• photoz/fzb_6bands/projects/com_cam/data/gold_rubin/model_inform_fzboost.pkl
• rail_data/data_edfs_with_mag.hdf5
• rail_data/train_with_mag.hdf5
• rail_data/data_edfs_redshift.hdf5 (cached point estimates)
• rail_data/data_edfs_redshift_pdfs.fits (cached PDFs)

Outputs

If the regen flag is on: writes data_edfs_redshift.hdf5 + data_edfs_redshift_pdfs.fits. Otherwise inline figures only.

import os

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import scipy.integrate as integrate
import scipy.stats as stats

import fitsio
import xlens

import tables_io

# =============================
# Matplotlib defaults
# =============================
SMALL_SIZE = 8
MEDIUM_SIZE = 10
BIGGER_SIZE = 14
plt.rc("font", size=BIGGER_SIZE)
plt.rc("axes", titlesize=BIGGER_SIZE)
plt.rc("axes", labelsize=BIGGER_SIZE)
plt.rc("xtick", labelsize=BIGGER_SIZE)
plt.rc("ytick", labelsize=BIGGER_SIZE)
plt.rc("legend", fontsize=BIGGER_SIZE)
plt.rc("figure", titlesize=BIGGER_SIZE)

pscratch= "/gpfs02/astro/workarea/xli6/"
Dir = os.path.join(
    pscratch,
    "work/DP1/"
)
rail_pz_h5 = os.path.join(
    Dir, "photoz",
    "fzb_6bands/projects/com_cam/data/gold_rubin/output_estimate_fzboost.hdf5",
)

MODEL_PATH = os.path.join(
    Dir, "photoz",
    "fzb_6bands/projects/com_cam/data/gold_rubin/model_inform_fzboost.pkl",
)

data_fname = os.path.join(Dir, "rail_data", "data_edfs_with_mag.hdf5")
df = tables_io.read(
    data_fname,
    tType="astropyTable",
)
train_fname = os.path.join(Dir, "rail_data", "train_with_mag.hdf5")
df_train = tables_io.read(
    train_fname,
    tType="astropyTable",
)

fig, axs = plt.subplots(1, 4, figsize=(24, 5))
xbins = np.linspace(0, 4, 14)
magbins = np.linspace(14, 26.0, 52)
colorbins = np.linspace(-1, 2.5, 50)
bands = "grizy"
for i in range(4):
    bandi = bands[i]
    bandj = bands[i + 1]
    xi = f"{bandi}_mag_gauss2"
    xj = f"{bandj}_mag_gauss2"
    axs[i].hist(df_train[xi] - df_train[xj], bins=colorbins, alpha=0.35, color="b", density=True)
    axs[i].hist(df[xi] - df[xj], bins=colorbins, alpha=0.8, color="k", density=True, histtype="step", linewidth=2.0)
    axs[i].set_xlabel(f"{bandi}-{bandj}", fontsize=25)
plt.tight_layout()

msk = (
    (df["u_mag_gauss2"] < 26.4) & 
    (df["g_mag_gauss2"] < 27.8) & 
    (df["r_mag_gauss2"] < 27.1) & 
    (df["i_mag_gauss2"] < 24.5) & 
    (df["z_mag_gauss2"] < 25.8) & 
    (df["y_mag_gauss2"] < 24.6) &
    ((df["i_fpfs1_m00"] + df["i_fpfs1_m20"]) / df["i_fpfs1_m00"] > 0.1) &
    ((df["i_fpfs1_m00"] + df["i_fpfs1_m20"]) / df["i_fpfs1_m00"] < 1.5)
)

size = (df["i_fpfs1_m00"] + df["i_fpfs1_m20"]) / df["i_fpfs1_m00"]
# remove stars in the catalog with resolution cut
plt.scatter(df["i_mag_gauss2"][msk], size[msk], s=0.1)
plt.xlim(18, 25.0)

(18.0, 25.0)

zfname = os.path.join(Dir, "rail_data", "data_edfs_redshift.hdf5")
pzfname = os.path.join(Dir, "rail_data", "data_edfs_redshift_pdfs.fits")
if False:
    import pickle
    with open(MODEL_PATH, "rb") as f:
        pz_obj = pickle.load(f)
    fzbobj = xlens.catalog.redshift.flexzboostEstimator(pz_obj)
    redshift = fzbobj.get_z(
        df.as_array(),
        mag_zero=31.4,
        flux_name="gauss2",
        bands="ugrizy",
        ref_band="i",
        return_pdfs=True,
    )
    pdfs = redshift["pdfs"]
    fitsio.write(pzfname, pdfs)
    d2 = {k: v for k, v in redshift.items() if k != "pdfs"}
    tables_io.write(d2, zfname)
else:
    redshift = tables_io.read(
        zfname,
        tType="astropyTable",
    )
    pdfs = fitsio.read(pzfname)

# Shear-perturbed photo-z (used downstream for R_sel selection-bias correction).
#   1p / 1m : comp=1, dg = ±0.02   (perturbs g1)
#   2p / 2m : comp=2, dg = ±0.02   (perturbs g2)
zfname_1p = os.path.join(Dir, "rail_data", "data_edfs_redshift_1p.hdf5")
zfname_1m = os.path.join(Dir, "rail_data", "data_edfs_redshift_1m.hdf5")
zfname_2p = os.path.join(Dir, "rail_data", "data_edfs_redshift_2p.hdf5")
zfname_2m = os.path.join(Dir, "rail_data", "data_edfs_redshift_2m.hdf5")

if False:
    import pickle
    with open(MODEL_PATH, "rb") as f:
        pz_obj = pickle.load(f)
    fzbobj = xlens.catalog.redshift.flexzboostEstimator(pz_obj)

    x = df.as_array()
    perturbed = {}
    for tag, comp, dg, fname in [
        ("1p", 1, +0.02, zfname_1p),
        ("1m", 1, -0.02, zfname_1m),
        ("2p", 2, +0.02, zfname_2p),
        ("2m", 2, -0.02, zfname_2m),
    ]:
        print(f"[info] running fzboost for {tag} (comp={comp}, dg={dg:+.2f})")
        rz = fzbobj.get_z(
            x,
            mag_zero=31.4,
            flux_name="gauss2",
            bands="ugrizy",
            ref_band="i",
            comp=comp,
            dg=dg,
            return_pdfs=False,
        )
        rz = {k: v for k, v in rz.items() if k != "pdfs"}
        tables_io.write(rz, fname)
        perturbed[tag] = rz
    redshift_1p = perturbed["1p"]
    redshift_1m = perturbed["1m"]
    redshift_2p = perturbed["2p"]
    redshift_2m = perturbed["2m"]
else:
    redshift_1p = tables_io.read(zfname_1p, tType="astropyTable")
    redshift_1m = tables_io.read(zfname_1m, tType="astropyTable")
    redshift_2p = tables_io.read(zfname_2p, tType="astropyTable")
    redshift_2m = tables_io.read(zfname_2m, tType="astropyTable")

mm = np.sum(pdfs[:, 0:70], axis=1) * 0.01 < 0.25
plt.plot(np.average(pdfs[mm], axis=0))
plt.grid()
np.sum(np.average(pdfs[mm], axis=0)[0: 70]) * 0.01

np.float64(0.08282436482082149)