diff --git a/config/columns/shapepipe_v2.yaml b/config/columns/shapepipe_v2.yaml index 0bae5465f..c27f02143 100644 --- a/config/columns/shapepipe_v2.yaml +++ b/config/columns/shapepipe_v2.yaml @@ -28,6 +28,7 @@ galaxy: XWIN_WORLD: SExtractor windowed right ascension [deg]; tile-overlap cut. YWIN_WORLD: SExtractor windowed declination [deg]; tile-overlap cut. TILE_ID: Tile the object was detected on. + NUMBER: Object number within its tile. FLAGS: SExtractor extraction flags. MAG_AUTO: SExtractor magnitude. N_EPOCH: Number of exposures covering the object. diff --git a/cosmo_val/cat_config.yaml b/cosmo_val/cat_config.yaml index 6b69142be..6ad72d1f3 100644 --- a/cosmo_val/cat_config.yaml +++ b/cosmo_val/cat_config.yaml @@ -1292,3 +1292,67 @@ GLASS_mock_validation: e2_col: e2 path: /n17data/UNIONS/WL/v1.6.x/unions_shapepipe_star_2024_v1.6.a.fits patch_number: 150 +SP_v2.0: + host: nibi + blind: unions + subdir: /project/6001537/cdaley/cosmo-val-v2-tomo/catalogue + pipeline: SP + colour: violet + getdist_colour: 0.0, 0.5, 1.0 + ls: solid + marker: '*' + patch_number: 100 + cov_th: + A: 4689.662867153573 + n_e: 6.86756988279185 + sigma_e: 0.28908294714031374 + sigma_e_convention: per_component + n_psf: 0.580817039505617 + tomo: + A: 3814.775466932069 + mask: /project/6001537/cdaley/cosmo-val-v2-tomo/catalogue/footprint-tomo-nside4096-ring.fits + n_e: + - 0.7427992141190386 + - 0.9163698955600097 + - 1.006977964461924 + - 1.3252290285029567 + sigma_e: + - 0.28814703635003425 + - 0.2851918267345306 + - 0.28830578863062367 + - 0.27332974092856765 + mask: /project/6001537/cdaley/cosmo-val-v2-tomo/catalogue/footprint-nside4096-ring.fits + shear: + path: shape_catalog_bound9_tomo.fits + R: 1.0 + ra_col: RA + dec_col: Dec + w_col: w_des + e1_col: e1 + e2_col: e2 + e1_PSF_col: e1_PSF + e2_PSF_col: e2_PSF + # Calibration recipe: cosmo_val/nz/v2_bound9.yaml + redshift_path: /project/def-mjhudson/cdaley/tomo-nz-v2-full/products/nz_v2_bound9_full_tomo.txt + tomo_bin_col: TOM_BIN_ID + psf: + path: stars.hdf5 + ra_col: RA + dec_col: DEC + e1_PSF_col: HSM_G1_PSF + e2_PSF_col: HSM_G2_PSF + e1_star_col: HSM_G1_STAR + e2_star_col: HSM_G2_STAR + PSF_size: HSM_T_PSF + star_size: HSM_T_STAR + PSF_flag: HSM_FLAG_PSF + star_flag: HSM_FLAG_STAR + mask_config: ../config/calibration/mask_v2.0.yaml + mask_cuts: + - MASK_8_Manual + star: + path: psf_catalog_ngmix.fits + ra_col: RA + dec_col: Dec + e1_col: e1 + e2_col: e2 diff --git a/cosmo_val/nz/README.md b/cosmo_val/nz/README.md new file mode 100644 index 000000000..6980bec73 --- /dev/null +++ b/cosmo_val/nz/README.md @@ -0,0 +1,40 @@ +# Provisional v2 n(z) + +`v2_bound9.yaml` records the four-bin SOM calibration recipe for the full-area bound-9 selection. +The output has z in column zero and four densities whose integrals are the retained shear-weight fractions; their sum has unit integral. +The n(z) describes only `TOM_BIN_ID > 0` objects, so a non-tomographic run over all objects pairs an all-object sample with a gold-only n(z). +Source weights come from the 22,200-tile selection (191,016,178 rows), rather than the 19,940-tile catalogue measured by cosmo_val (170.6M rows). +cosmo_val reads the saved file through `shear.redshift_path` and never runs calibration or loads photometry. +The `SP_v2.0` entry uses Nibi catalogue and n(z) paths and declares `host: nibi`. +The candide n(z) is `/n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/calibration/nz_tomo.txt`. +At another site, set `shear.redshift_path` to the local copy. + +This EGS-only calibration is for validation and null tests. +It has no independent-field variance or model of spectroscopic selection and does not support parameter constraints. + +Run from this checkout inside the sp_validation container on a Slurm allocation: + +```bash +python scripts/prepare_nz_photometry.py cosmo_val/nz/v2_bound9.yaml +python scripts/calibrate_nz.py cosmo_val/nz/v2_bound9.yaml --workers 16 +``` + +Photometry preparation is a separate step; reuse its cache for subsequent calibrations. +The calibration uses frozen SOM cells, groups and gold masks, and computes group weights over the full source selection. +`--rebuild-model` retrains the SOM, Ward groups and gold masks from the inputs. +`--output DIRECTORY` selects an output directory; `--figures DIRECTORY` renders diagnostics. +The frozen-delivery reproduction gate lives in `scratch/cdaley/nz_som_reproduction/`. + +The calibration sample, `inputs/matched_calibration.npz`, contains the v1.4.6.3 sample's fiducial spectroscopic `Z_B` as `z`, DR6 BPZ `Z_B` as `photoz`, and detected ugriz magnitudes as `mag`. +The selected shear Parquet and its provenance are `inputs/bound9_full_selection.parquet` and `inputs/selection.json` in the analysis directory named by the recipe. +Photometry matching uses the nearest DR6 position within 0.8 arcsec in the same tile and preserves source row order. +ShapePipe `NUMBER` identifies the source within its tile. +BPZ stays float64 to preserve values on bin edges. + +The recipe fixes a 101×101, 60-epoch batch SOM of all ten ugriz colours plus r magnitude, standardized per feature, and 2,000 Ward groups. +BPZ intervals [0.10, 0.37, 0.55, 0.76, 1.20] include lower edges and exclude upper edges. +Each gold group has in-bin spectra and passes the source-BPZ versus mean-spec-z threshold `max(0.4, 5 NMAD)` on the 600-tile reference sample. +Groups empty in that reference skip the BPZ-versus-spec-z check. +The source shear-weight sum divided by the number of in-bin calibrators weights each spectrum in its group. +Linear deposition on [0, 5] with Δz=0.01 conserves mass and first moment under trapezoidal integration. +Canonical z takes precedence; detected z2 fills missing canonical z after the offset in the recipe. diff --git a/cosmo_val/nz/v2_bound9.yaml b/cosmo_val/nz/v2_bound9.yaml new file mode 100644 index 000000000..a25d6a4f0 --- /dev/null +++ b/cosmo_val/nz/v2_bound9.yaml @@ -0,0 +1,31 @@ +name: provisional_v2_bound9_DR6_BPZ_SOM_EGS_full_z_z2 +limitation: EGS-only provisional calibration for validation and null tests; no independent-field variance or spectroscopic-selection model. +inputs: + # Fiducial spectroscopic Z_B and matched DR6 BPZ/ugriz; 11,114 EGS spectra. + calibrators: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/inputs/matched_calibration.npz + model: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/inputs/som_assignment_model.npz + gold_reference: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/inputs/gold_reference_600tiles.npz +photometry: + selection: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/inputs/bound9_full_selection.parquet + dr6_root: /n17data/UNIONS/WL/photometry/UNIONS_DR6 + output: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/cache/matched_full.npz +som: + side: 101 + epochs: 60 + seed: 42 +groups: 2000 +edges: [0.10, 0.37, 0.55, 0.76, 1.20] +gold_qc: + floor: 0.4 + nmad_factor: 5 +# Canonical z takes precedence over z2. +z2_offset: -0.1104888916015625 +grid: {min: 0, max: 5, step: 0.01} +output: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/calibration +diagnostics: + published_nz: /n17data/sguerrini/UNIONS/WL/nz/v1.4.6.3/nz_SP_v1.4.6.3_B.txt + tile_nz: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_first/products/nz_v2_bound9_tomo.txt + full_label: Full area, z + z2 + tile_label: 600-tile reference + published_label: Published v1.4.6.3 + source_label: 600-tile reference diff --git a/scratch/cdaley/nz_som_reproduction/README.md b/scratch/cdaley/nz_som_reproduction/README.md new file mode 100644 index 000000000..b8236d7fb --- /dev/null +++ b/scratch/cdaley/nz_som_reproduction/README.md @@ -0,0 +1,13 @@ +# Full-area SOM n(z) reproduction + +Run inside the sp_validation container on an allocated candide compute node: + +```bash +python scratch/cdaley/nz_som_reproduction/reproduce.py cosmo_val/nz/v2_bound9.yaml scratch/nz-reproduction --workers 16 +python scratch/cdaley/nz_som_reproduction/reproduce.py cosmo_val/nz/v2_bound9.yaml scratch/nz-rebuild --workers 16 --rebuild-model +``` + +Both commands compare the output against `reference.yaml` and assert zero `TOM_BIN_ID` differences across every source row and per-bin mean redshift differences below 1e-4. +The rebuild command also checks exact equality of the SOM code, feature normalization, groups, edges and gold masks. +`verify.py reference.yaml OUTPUT` can check an existing output independently. +The small cosmo_val reader and SACC round-trip check is in `test_nz_calibration.py`. diff --git a/scratch/cdaley/nz_som_reproduction/figures.py b/scratch/cdaley/nz_som_reproduction/figures.py new file mode 100644 index 000000000..9c36ab6b0 --- /dev/null +++ b/scratch/cdaley/nz_som_reproduction/figures.py @@ -0,0 +1,65 @@ +"""Render calibration diagnostics and saved-delivery z2 coverage.""" + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import yaml + +from sp_validation.nz_diagnostics import COLORS, render + +matplotlib.rcParams.update({"text.usetex": False, "font.family": "DejaVu Sans"}) + +parser = argparse.ArgumentParser(description=__doc__) +parser.add_argument("recipe") +parser.add_argument("output") +parser.add_argument("figures") +args = parser.parse_args() +recipe = yaml.safe_load(Path(args.recipe).read_text()) +output = Path(args.output) +figures = Path(args.figures) +render( + dict( + recipe["diagnostics"], + calibration_arrays=output / "calibration_arrays.npz", + model=output / "som_assignment_model.npz", + summary=output / "summary.json", + source_catalogue=recipe["inputs"]["gold_reference"], + side=recipe["som"]["side"], + workers=16, + ), + figures, +) +# This comparison uses saved delivery counts; production runs only z+z2. +delivery = Path("/n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/products") +retention = json.loads((delivery / "summary.json").read_text())["retention"] +stages = ["all", "complete_photometry", "bpz_bounds", "gold"] +x = np.arange(len(stages)) +fig, ax = plt.subplots(figsize=(8, 4.8), layout="constrained") +width = 0.34 +for j, (key, label, color) in enumerate( + (("zonly", "Canonical z only", COLORS[2]), ("both", "z + z2", COLORS[0])) +): + bars = ax.bar( + x + (j - 0.5) * width, + [100 * retention[key][stage] for stage in stages], + width, + label=label, + color=color, + ) + ax.bar_label(bars, fmt="%.1f%%", fontsize=8, padding=2) +ax.set( + xticks=x, + xticklabels=["All", "Complete\nphotometry", "BPZ\nbounds", "Gold"], + ylabel="Retained bound-9 shear weight (%)", + ylim=(0, 110), + title="Coverage with and without z2 (saved delivery)", +) +ax.legend() +fig.savefig(figures / "coverage_z2.png", dpi=100, bbox_inches="tight") +plt.close(fig) diff --git a/scratch/cdaley/nz_som_reproduction/reference.yaml b/scratch/cdaley/nz_som_reproduction/reference.yaml new file mode 100644 index 000000000..774e1e0a3 --- /dev/null +++ b/scratch/cdaley/nz_som_reproduction/reference.yaml @@ -0,0 +1,3 @@ +reference: + nz: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/products/nz_v2_bound9_full_tomo.txt + tom_bins: /n17data/cdaley/unions/analyses/shear_tomographic/nz_v2_full/products/tom_bin_full.parquet diff --git a/scratch/cdaley/nz_som_reproduction/reproduce.py b/scratch/cdaley/nz_som_reproduction/reproduce.py new file mode 100644 index 000000000..701f431fe --- /dev/null +++ b/scratch/cdaley/nz_som_reproduction/reproduce.py @@ -0,0 +1,33 @@ +"""Reproduce the full-area delivery, optionally rebuilding the assignment model.""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import yaml +from verify import verify + +from sp_validation.nz_pipeline import run + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("recipe") + parser.add_argument("output") + parser.add_argument("--workers", type=int, default=16) + parser.add_argument("--rebuild-model", action="store_true") + args = parser.parse_args() + config = yaml.safe_load(Path(args.recipe).read_text()) + run(config, args.output, workers=args.workers, rebuild=args.rebuild_model) + reference = yaml.safe_load(Path(__file__).with_name("reference.yaml").read_text()) + verify(reference, args.output) + if args.rebuild_model: + new = np.load(Path(args.output) / "som_assignment_model.npz") + old = np.load(config["inputs"]["model"]) + keys = ["code", "center", "scale", "group_map", "edges", "gold"] + result = {key: bool(np.array_equal(new[key], old[key])) for key in keys} + result["max_code_difference"] = float(abs(new["code"] - old["code"]).max()) + result["gold_mismatches"] = int(np.count_nonzero(new["gold"] != old["gold"])) + (Path(args.output) / "rebuild.json").write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2), flush=True) + assert all(result[key] for key in keys), result diff --git a/scratch/cdaley/nz_som_reproduction/verify.py b/scratch/cdaley/nz_som_reproduction/verify.py new file mode 100644 index 000000000..4b1167d7f --- /dev/null +++ b/scratch/cdaley/nz_som_reproduction/verify.py @@ -0,0 +1,60 @@ +"""Check the frozen full-area delivery against a calibration output directory.""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import yaml + + +def verify(config, output): + """Compare delivered full-area n(z) and every sidecar bin ID.""" + output = Path(output) + reference = config["reference"] + actual = np.loadtxt(output / "nz_tomo.txt") + expected = np.loadtxt(reference["nz"]) + z = actual[:, 0] + means = np.trapezoid(z[:, None] * actual[:, 1:], z, axis=0) + means /= np.trapezoid(actual[:, 1:], z, axis=0) + ref_mean = np.trapezoid(expected[:, :1] * expected[:, 1:], expected[:, 0], axis=0) + ref_mean /= np.trapezoid(expected[:, 1:], expected[:, 0], axis=0) + error = abs(means - ref_mean) + assert np.all(error < 1e-4), error + mismatches, count = 0, 0 + # Both sidecars preserve the selected-catalogue row order. + a = pq.ParquetFile(output / "tom_bin.parquet") + b = pq.ParquetFile(reference["tom_bins"]) + assert a.metadata.num_rows == b.metadata.num_rows + for aa, bb in zip( + a.iter_batches(columns=["TILE_ID", "NUMBER", "TOM_BIN_ID"]), + b.iter_batches(columns=["TILE_ID", "NUMBER", "TOM_BIN_ID"]), + ): + assert aa.column(0).equals(bb.column(0)) + assert aa.column(1).equals(bb.column(1)) + mismatches += np.count_nonzero( + aa.column(2).to_numpy() != bb.column(2).to_numpy() + ) + count += len(aa) + assert count == a.metadata.num_rows, (count, a.metadata.num_rows) + assert mismatches == 0, mismatches + result = dict( + means=means.tolist(), + reference_means=ref_mean.tolist(), + absolute_mean_errors=error.tolist(), + rows=count, + tom_bin_mismatches=int(mismatches), + max_density_difference=float(abs(actual - expected).max()), + ) + (output / "verification.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2), flush=True) + return result + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("reference") + parser.add_argument("output") + args = parser.parse_args() + verify(yaml.safe_load(Path(args.reference).read_text()), args.output) diff --git a/scripts/calibrate_nz.py b/scripts/calibrate_nz.py new file mode 100644 index 000000000..38d36036b --- /dev/null +++ b/scripts/calibrate_nz.py @@ -0,0 +1,6 @@ +"""Run a saved SOM n(z) recipe outside cosmo_val (see cosmo_val/nz/README.md).""" + +from sp_validation.nz_pipeline import main + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_nz_photometry.py b/scripts/prepare_nz_photometry.py new file mode 100644 index 000000000..e5b3d3692 --- /dev/null +++ b/scripts/prepare_nz_photometry.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python +"""Prepare the DR6 positional photometry cache used by SOM n(z).""" + +import argparse +from pathlib import Path + +import yaml + +from sp_validation.nz_photometry import prepare_nz_photometry + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("recipe", help="YAML recipe containing a photometry section") + args = parser.parse_args() + photometry = yaml.safe_load(Path(args.recipe).read_text())["photometry"] + selection, dr6_root, output = ( + photometry[key] for key in ("selection", "dr6_root", "output") + ) + prepare_nz_photometry(selection, dr6_root, output) + + +if __name__ == "__main__": + main() diff --git a/src/sp_validation/nz_calibration.py b/src/sp_validation/nz_calibration.py new file mode 100644 index 000000000..dd725e9ac --- /dev/null +++ b/src/sp_validation/nz_calibration.py @@ -0,0 +1,240 @@ +"""Provisional SOM direct calibration of source redshift distributions. + +Photometry is detected when finite and between 0 and 50 magnitudes. +The map uses all ten ugriz colours and r, standardized per feature. +Each BPZ bin has its own gold mask and source-weight/calibrator-count ratios. +The spectroscopic calibration sample's fiducial Z_B is supplied as ``cal_z``; +``photoz`` always means DR6 BPZ Z_B. +""" + +import numpy as np +from scipy.cluster.hierarchy import fcluster, linkage +from scipy.ndimage import gaussian_filter +from scipy.spatial import cKDTree + + +def detected(mag): + """Select detected DR6 magnitudes, excluding -99 and 99 sentinels.""" + return np.isfinite(mag) & (mag > 0) & (mag < 50) + + +def features(mag): + """Return the ten ugriz colours and r magnitude, in the training order.""" + return np.column_stack( + [mag[:, i] - mag[:, j] for i in range(5) for j in range(i + 1, 5)] + [mag[:, 2]] + ) + + +def train(mag, side=101, epochs=60, seed=42, workers=1): + """Train a planar batch SOM with PCA initialization and Gaussian smoothing.""" + x = features(mag) + center, scale = x.mean(0), x.std(0) + scale[scale == 0] = 1 + x = (x - center) / scale + _, s, v = np.linalg.svd(x, full_matrices=False) + grid = np.stack( + np.meshgrid(np.linspace(-2, 2, side), np.linspace(-2, 2, side)), axis=-1 + ).reshape(-1, 2) + code = grid @ (v[:2] * s[:2, None] / np.sqrt(len(x))) + code += np.random.default_rng(seed).normal(0, 0.001, code.shape) + for sigma in np.geomspace(side / 5, 0.5, epochs): + _, cell = cKDTree(code).query(x, workers=workers) + counts = np.bincount(cell, minlength=side * side).reshape(side, side) + den = gaussian_filter(counts.astype(float), sigma, mode="constant") + sums = np.column_stack( + [ + np.bincount(cell, weights=x[:, k], minlength=side * side) + for k in range(x.shape[1]) + ] + ).reshape(side, side, -1) + num = gaussian_filter(sums, (sigma, sigma, 0), mode="constant") + good = den.ravel() > 1e-15 + code[good] = num.reshape(-1, x.shape[1])[good] / den.ravel()[good, None] + return dict(code=code, center=center, scale=scale, side=np.array(side)) + + +def assign(mag, model, workers=1): + """Return map distance and cell for detected five-band magnitudes.""" + x = (features(mag) - model["center"]) / model["scale"] + return cKDTree(model["code"]).query(x, workers=workers) + + +def group_cells(model, n_groups=2000): + """Ward-cluster map code vectors into zero-based calibration groups.""" + return ( + fcluster(linkage(model["code"], method="ward"), n_groups, criterion="maxclust") + - 1 + ) + + +def z_magnitudes(mag, offset, use_z2=True): + """Return ugriz and z-channel IDs (0 missing, 1 canonical z, 2 mapped z2). + + Canonical z takes precedence; z2 + offset enters only when z is undetected. + Five-column inputs contain canonical z only. + """ + out = mag[:, :5].astype("f8") + channel = detected(out[:, 4]).astype(np.int8) + out[channel == 0, 4] = np.nan + if use_z2 and mag.shape[1] == 6: + extra = (channel == 0) & detected(mag[:, 5]) + out[extra, 4] = mag[extra, 5] + offset + channel[extra] = 2 + return out, channel + + +def source_groups( + mag, photoz, weight, model, *, offset, use_z2=True, chunk=1_000_000, workers=1 +): + """Assign eligible sources to BPZ bins and SOM groups in bounded chunks. + + Ineligible sources have preliminary bin and group -1. + BPZ intervals include the lower edge and exclude the upper edge. + """ + edges = model["edges"] + prelim = np.full(len(weight), -1, np.int8) + group = np.full(len(weight), -1, np.int16) + channel = np.zeros(len(weight), np.int8) + tree = cKDTree(model["code"]) + for start in range(0, len(weight), chunk): + stop = min(start + chunk, len(weight)) + m, channel[start:stop] = z_magnitudes(mag[start:stop], offset, use_z2) + p, w = photoz[start:stop], weight[start:stop] + valid = ( + detected(m).all(1) + & np.isfinite(p) + & (p >= edges[0]) + & (p < edges[-1]) + & np.isfinite(w) + & (w > 0) + ) + rows = start + np.flatnonzero(valid) + _, cell = tree.query( + (features(m[valid]) - model["center"]) / model["scale"], workers=workers + ) + group[rows] = model["group_map"][cell] + prelim[rows] = np.searchsorted(edges, p[valid], side="right") - 1 + return prelim, group, channel + + +def gold_threshold(cal_z, cal_photoz, floor=0.4, nmad_factor=5): + """Compute max(floor, factor × NMAD) of unnormalized spec-z − BPZ.""" + residual = cal_z - cal_photoz + good = np.isfinite(residual) & (cal_photoz >= 0) & (cal_photoz < 5) + residual = residual[good] + nmad = 1.4826 * np.median(abs(residual - np.median(residual))) + return float(max(floor, nmad_factor * nmad)) + + +def gold_mask(source_weight, source_photoz_sum, cal_count, cal_z_sum, threshold): + """Keep supported groups whose mean spec-z and source BPZ agree within QC.""" + spec = np.divide( + cal_z_sum, cal_count, out=np.zeros_like(cal_z_sum), where=cal_count > 0 + ) + bpz = np.divide( + source_photoz_sum, + source_weight, + out=np.zeros_like(source_photoz_sum), + where=source_weight > 0, + ) + return (cal_count > 0) & ((source_weight == 0) | (abs(spec - bpz) <= threshold)) + + +def tom_bin_ids(prelim, group, gold): + """Return 1-based gold bin IDs, with zero for every excluded row.""" + result = np.zeros(len(prelim), np.int8) + rows = np.flatnonzero(prelim >= 0) + keep = gold[prelim[rows], group[rows]] + result[rows[keep]] = prelim[rows[keep]] + 1 + return result + + +def deposit(z, weight, grid): + """Linearly deposit mass, conserving its zeroth and first trapezoid moments.""" + if np.any(~np.isfinite(z) | (z < grid[0]) | (z > grid[-1])): + raise ValueError("calibration redshift outside output grid") + lo = np.clip(np.searchsorted(grid, z, side="right") - 1, 0, len(grid) - 2) + fraction = (z - grid[lo]) / (grid[lo + 1] - grid[lo]) + mass = np.bincount(lo, weights=weight * (1 - fraction), minlength=len(grid)) + mass += np.bincount(lo + 1, weights=weight * fraction, minlength=len(grid)) + quadrature = np.r_[grid[1] - grid[0], grid[2:] - grid[:-2], grid[-1] - grid[-2]] / 2 + return mass / quadrature + + +def calibrate( + prelim, + group, + weight, + photoz, + cal_z, + cal_photoz, + cal_group, + edges, + n_groups, + grid, + *, + gold=None, + threshold=None, +): + """Calibrate each bin with source shear weight / in-bin calibrator count. + + With ``gold=None``, compute QC masks from this source sample. + Supplying gold freezes the selection while recalibrating the group weights. + Histograms have unit integral per bin; the writer applies mixture fractions. + """ + weights, masks, bins, hists = [], [], [], [] + for k, (lo, hi) in enumerate(zip(edges[:-1], edges[1:])): + rows = np.flatnonzero(prelim == k) + g, w = group[rows].astype(np.int64), weight[rows] + source = np.bincount(g, weights=w, minlength=n_groups) + cm = np.isfinite(cal_photoz) & (cal_photoz >= lo) & (cal_photoz < hi) + count = np.bincount(cal_group[cm], minlength=n_groups) + if gold is None: + zsum = np.bincount(cal_group[cm], weights=cal_z[cm], minlength=n_groups) + psum = np.bincount(g, weights=w * photoz[rows], minlength=n_groups) + mask = gold_mask(source, psum, count, zsum, threshold) + else: + mask = gold[k] & (count > 0) + ratio = np.divide(source, count, out=np.zeros(n_groups), where=mask) + cw = np.where(cm, ratio[cal_group], 0.0) + norm = cw.sum() + if norm <= 0: + raise ValueError(f"empty calibration support in bin {k + 1}") + kept = mask[g] + bins.append( + dict( + bin=k + 1, + edges=[float(lo), float(hi)], + mean=float(cal_z @ cw / norm), + weight_sum=float(norm), + pre_gold_weight_sum=float(w.sum()), + gold_fraction=float(norm / w.sum()), + n_rows=int(kept.sum()), + calibrators_in_bpZ_bin=int(cm.sum()), + calibrators_with_positive_weight=int((cw > 0).sum()), + ) + ) + weights.append(cw) + masks.append(mask) + hists.append(deposit(cal_z, cw / norm, grid)) + total = sum(b["weight_sum"] for b in bins) + for b in bins: + b["weight_fraction"] = b["weight_sum"] / total + return dict( + bins=bins, + cal_weights=np.array(weights), + gold=np.array(masks), + bin_hists=np.array(hists), + ) + + +def write_nz(path, grid, bin_hists, weight_sums): + """Write z and bin densities integrating to retained shear-weight fractions.""" + fractions = np.asarray(weight_sums) / np.sum(weight_sums) + np.savetxt( + path, + np.column_stack([grid, (fractions[:, None] * bin_hists).T]), + fmt="%.12e", + header="z bin1 bin2 ...; columns integrate to retained " + "shear-weight fractions, total integral=1", + ) diff --git a/src/sp_validation/nz_diagnostics.py b/src/sp_validation/nz_diagnostics.py new file mode 100644 index 000000000..d280fd69e --- /dev/null +++ b/src/sp_validation/nz_diagnostics.py @@ -0,0 +1,352 @@ +"""Figures for the direct SOM redshift calibration.""" + +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LogNorm + +COLORS = ("#146e7c", "#d2791c", "#322d61", "#898989") + + +def _save(fig, path): + fig.savefig(path, dpi=100, bbox_inches="tight", pil_kwargs={"optimize": True}) + plt.close(fig) + + +def _curve(path): + data = np.loadtxt(path) + return data[:, 0], data[:, 1:] + + +def _normalise(y, z): + return y / np.trapezoid(y, z) + + +def _mechanics(path, gm, group, z, weights, grid, side): + ngroup = int(gm.max()) + 1 + ncal = np.bincount(group, minlength=ngroup) + zsum = np.bincount(group, weights=z, minlength=ngroup) + mean_z = np.divide(zsum, ncal, out=np.full(ngroup, np.nan), where=ncal > 0) + fig, axes = plt.subplots(1, 3, figsize=(14, 4.5), layout="constrained") + image = axes[0].imshow( + mean_z[gm].reshape(side, side), origin="lower", cmap="viridis", vmin=0, vmax=2 + ) + axes[0].set( + title="SOM groups colored by mean spec-z", xlabel="SOM cell", ylabel="SOM cell" + ) + axes[0].set_xticks([]) + axes[0].set_yticks([]) + fig.colorbar(image, ax=axes[0], label="Mean spectroscopic z") + for k in range(weights.shape[0]): + count = np.bincount(group, weights=(weights[k] > 0), minlength=ngroup) + group_weight = np.bincount(group, weights=weights[k], minlength=ngroup) + ratio = np.divide( + group_weight, + count, + out=np.full(ngroup, np.nan), + where=(count > 0) & (group_weight > 0), + ) + ratio = ratio[np.isfinite(ratio)] + axes[1].hist( + ratio, + bins=np.geomspace(ratio.min(), ratio.max(), 32), + histtype="step", + lw=1.6, + color=COLORS[k % len(COLORS)], + label=f"Bin {k + 1}", + ) + axes[1].set( + xscale="log", + xlabel="Source shear weight / calibrator", + ylabel="Groups", + title="Group occupancy sets each spectrum's weight", + ) + axes[1].legend(fontsize=8) + raw = np.histogram(z, bins=np.r_[grid, grid[-1] + grid[1] - grid[0]])[0].astype( + float + ) + calibrated = np.zeros(len(grid)) + for k in range(weights.shape[0]): + calibrated += np.histogram( + z, bins=np.r_[grid, grid[-1] + grid[1] - grid[0]], weights=weights[k] + )[0] + raw = raw / np.trapezoid(raw, grid) + calibrated = calibrated / np.trapezoid(calibrated, grid) + axes[2].plot(grid, raw, color=COLORS[3], lw=1.3, label="Unweighted spectra") + axes[2].plot( + grid, calibrated, color=COLORS[0], lw=1.8, label="Source-weighted spectra" + ) + axes[2].set( + xlim=(0, 2), + xlabel="Spectroscopic redshift", + ylabel="Normalized n(z)", + title="Sum the per-group weighted spectra", + ) + axes[2].legend(fontsize=8) + fig.suptitle( + "SOM direct calibration: map similar photometry, then transfer source occupancy to spectra" + ) + _save(fig, path) + + +@matplotlib.rc_context({"text.usetex": False, "font.family": "DejaVu Sans"}) +def render(config, output_dir): + """Render calibration figures from saved products and explicit input paths. + + Required config keys are ``calibration_arrays``, ``model``, ``summary`` + ``side``, ``source_label`` and ``source_catalogue``. The source catalogue NPZ contains ``weight``, + ``photoz`` and ``mag``. An optional ``source_assignment`` NPZ can provide + precomputed ``cell`` values; otherwise cells are assigned in chunks. + Optional keys ``published_nz`` and ``tile_nz`` point to text n(z) files. + """ + out = Path(output_dir) + out.mkdir(parents=True, exist_ok=True) + arrays = np.load(config["calibration_arrays"]) + model = np.load(config["model"]) + summary = json.loads(Path(config["summary"]).read_text()) + source = np.load(config["source_catalogue"]) + side = config["side"] + + z = arrays["z"] + p = arrays["cal_photoz"] + group = arrays["cal_group"] + cell = arrays["cal_cell"] + edges = arrays["edges"] + gold = arrays["gold"] + weights = arrays["cal_weights"] + grid = arrays["grid"] + gm = model["group_map"] + src_w = source["weight"] + src_p = source["photoz"] + src_mag = source["mag"] + usable = ( + np.isfinite(src_mag).all(axis=1) + & (src_mag > 0).all(axis=1) + & (src_mag < 50).all(axis=1) + & (src_w > 0) + & np.isfinite(src_p) + & (src_p >= edges[0]) + & (src_p < edges[-1]) + ) + if config.get("source_assignment"): + cells = np.load(config["source_assignment"])["cell"][usable] + else: + from . import nz_calibration + + valid_mag = src_mag[usable] + cells = np.empty(len(valid_mag), dtype=np.int32) + chunk_size = config.get("assignment_chunk_size", 200_000) + for start in range(0, len(valid_mag), chunk_size): + stop = min(start + chunk_size, len(valid_mag)) + _, cells[start:stop] = nz_calibration.assign( + valid_mag[start:stop], model, workers=config.get("workers", 1) + ) + cells, src_w, src_p = cells, src_w[usable], src_p[usable] + ncell = len(gm) + src_weight_cell = np.bincount(cells, weights=src_w, minlength=ncell) + src_p_cell = np.bincount(cells, weights=src_w * src_p, minlength=ncell) + mean_p = np.divide( + src_p_cell, + src_weight_cell, + out=np.full(ncell, np.nan), + where=src_weight_cell > 0, + ) + cal_count = np.bincount(cell, minlength=ncell) + n_group = int(gm.max()) + 1 + spec_count = np.bincount(group, minlength=n_group) + spec_sum = np.bincount(group, weights=z, minlength=n_group) + spec_sq = np.bincount(group, weights=z * z, minlength=n_group) + spec_mean = np.divide( + spec_sum, spec_count, out=np.full(n_group, np.nan), where=spec_count > 0 + ) + spec_std = np.sqrt( + np.maximum( + 0, + np.divide( + spec_sq, spec_count, out=np.full(n_group, np.nan), where=spec_count > 0 + ) + - spec_mean**2, + ) + ) + spec_std[spec_count < 2] = np.nan + + # SOM plane: sparse calibrators and source occupancy share the trained map. + fig, axes = plt.subplots(2, 3, figsize=(13, 8), layout="constrained") + maps = [ + cal_count, + src_weight_cell, + spec_mean[gm], + spec_std[gm], + mean_p, + spec_count[gm], + ] + titles = [ + "Calibrators per cell", + f"Source shear weight per cell ({config['source_label']})", + "Mean spec-z per group", + "Spec-z scatter per group", + f"Mean source BPZ per cell ({config['source_label']})", + "Calibrators per group", + ] + for ax, values, title in zip(axes.flat, maps, titles): + plane = np.asarray(values).reshape(side, side) + m = np.ma.masked_invalid(plane) + if title.startswith("Calibrators") or title.startswith("Source"): + m = np.ma.masked_where(m <= 0, m) + im = ax.imshow(m, origin="lower", cmap="viridis", norm=LogNorm()) + elif title.startswith("Mean spec-z"): + im = ax.imshow(m, origin="lower", cmap="viridis", vmin=0, vmax=2) + else: + im = ax.imshow(m, origin="lower", cmap="viridis") + ax.set_title(title) + ax.set_xticks([]) + ax.set_yticks([]) + fig.colorbar(im, ax=ax, fraction=0.046, pad=0.03) + _save(fig, out / "som_plane.png") + + # Per-bin gold support. + fig, axes = plt.subplots(1, len(edges) - 1, figsize=(15, 4), layout="constrained") + for k, ax in enumerate(np.atleast_1d(axes)): + passed = gold[k, gm].reshape(side, side) + ax.imshow( + passed, origin="lower", cmap="YlGn", vmin=0, vmax=1, interpolation="nearest" + ) + ax.set_title( + f"Bin {k + 1}: {100 * (1 - summary['bins'][k]['gold_fraction']):.1f}% weight lost" + ) + ax.set_xticks([]) + ax.set_yticks([]) + fig.suptitle("Per-bin gold groups (green) and rejected groups (pale)", y=1.08) + _save(fig, out / "gold_masks.png") + + valid = np.isfinite(p) & (p >= edges[0]) & (p < edges[-1]) + dz = (p[valid] - z[valid]) / (1 + z[valid]) + nmad = 1.4826 * np.median(np.abs(dz - np.median(dz))) + outlier = np.mean(np.abs(dz) > 0.15) + fig, ax = plt.subplots(figsize=(7, 6), layout="constrained") + h = ax.hist2d( + z[valid], + p[valid], + bins=100, + range=[[0, 2], [0, 1.2]], + norm=LogNorm(), + cmap="viridis", + ) + fig.colorbar(h[3], ax=ax, label="calibrators") + ax.plot([0, 2], [0, 2], color="white", lw=1) + for edge in edges: + ax.axhline(edge, color="#d1495b", lw=1) + ax.set( + xlim=(0, 2), + ylim=(0, 1.2), + xlabel="Spectroscopic redshift", + ylabel="DR6 BPZ", + title=f"BPZ-selected calibrators: NMAD {nmad:.3f}, outliers {100 * outlier:.1f}%", + ) + _save(fig, out / "bpz_vs_specz.png") + + fig, axes = plt.subplots(2, 2, figsize=(11, 7.5), sharex=True, layout="constrained") + tile_path = config.get("tile_nz") + if tile_path: + tz, tile = _curve(tile_path) + for k, ax in enumerate(axes.flat): + full_curve = _normalise(arrays["bin_hists"][k], grid) + ax.plot(grid, full_curve, color=COLORS[0], lw=1.7, label=config["full_label"]) + if tile_path: + ax.plot( + tz, + _normalise(tile[:, k], tz), + color=COLORS[1], + lw=1.4, + label=config["tile_label"], + ) + ax.axvspan(edges[k], edges[k + 1], color="0.92", zorder=0) + ax.set( + title=f"Bin {k + 1}: BPZ [{edges[k]:.2f}, {edges[k + 1]:.2f})", + xlabel="Redshift z", + ylabel="n(z)", + ) + ax.set_xlim(0, 2.2) + ax.legend(fontsize=8) + fig.suptitle(f"Per-bin n(z): {config['full_label']} vs {config['tile_label']}") + _save(fig, out / "nz_full_vs_600tiles.png") + + # The sum uses the file columns as delivered; published v1 is normalized + # independently because it represents a different selected population. + if config.get("published_nz"): + pz, published = _curve(config["published_nz"]) + fractions = np.array([b["weight_fraction"] for b in summary["bins"]]) + total = fractions @ arrays["bin_hists"] + total = _normalise(total, grid) + published = _normalise(published[:, 0], pz) + fig, axes = plt.subplots( + 2, + 1, + figsize=(10, 7.5), + gridspec_kw={"height_ratios": [2, 1]}, + layout="constrained", + ) + axes[0].plot( + grid, + total, + color=COLORS[0], + lw=2, + label=f"{config['full_label']}, mean {np.trapezoid(grid * total, grid):.4f}", + ) + if tile_path: + tile_total = _normalise(tile.sum(axis=1), tz) + axes[0].plot( + tz, + tile_total, + color=COLORS[1], + lw=1.7, + label=f"{config['tile_label']}, mean {np.trapezoid(tz * tile_total, tz):.4f}", + ) + axes[0].plot( + pz, + published, + color=COLORS[2], + lw=2, + label=f"{config['published_label']}, mean {np.trapezoid(pz * published, pz):.4f}", + ) + axes[0].set(xlim=(0, 2.2), ylabel="Normalized n(z)") + axes[0].legend() + axes[1].plot( + grid, + np.r_[0, np.cumsum((total[1:] + total[:-1]) * np.diff(grid) / 2)], + color=COLORS[0], + lw=2, + label=config["full_label"], + ) + if tile_path: + axes[1].plot( + tz, + np.r_[ + 0, np.cumsum((tile_total[1:] + tile_total[:-1]) * np.diff(tz) / 2) + ], + color=COLORS[1], + lw=1.7, + label=config["tile_label"], + ) + axes[1].plot( + pz, + np.r_[0, np.cumsum((published[1:] + published[:-1]) * np.diff(pz) / 2)], + color=COLORS[2], + lw=2, + label=config["published_label"], + ) + axes[1].set( + xlim=(0, 2.2), + ylim=(0, 1.02), + xlabel="Redshift z", + ylabel="Cumulative probability", + ) + _save(fig, out / "summed_bins_vs_v1.png") + + _mechanics(out / "som_calibration_mechanics.png", gm, group, z, weights, grid, side) + return sorted(path.name for path in out.glob("*.png")) diff --git a/src/sp_validation/nz_photometry.py b/src/sp_validation/nz_photometry.py new file mode 100644 index 000000000..9d977a96d --- /dev/null +++ b/src/sp_validation/nz_photometry.py @@ -0,0 +1,109 @@ +"""Attach DR6 photometry to selected shear-catalogue rows by sky position.""" + +from pathlib import Path + +import numpy as np +import pyarrow.compute as pc +import pyarrow.parquet as pq +from astropy.io import fits +from scipy.spatial import cKDTree + +BANDS = ("u", "g", "r", "i", "z", "z2") +MATCH_RADIUS_DEG = 0.8 / 3600 +DR6_SUFFIX = "_SP_ugriz_photoz_ext.cat" + + +def _read_dr6_tile(path): + columns = ["ALPHA_J2000", "DELTA_J2000"] + columns += [f"MAG_GAAP_{band}" for band in BANDS] + columns += ["Z_B"] + with fits.open(path, memmap=True) as hdus: + data = hdus[1].data + values = {name: np.asarray(data[name]) for name in columns} + return ( + values["ALPHA_J2000"].astype("f8"), + values["DELTA_J2000"].astype("f8"), + np.column_stack([values[f"MAG_GAAP_{band}"] for band in BANDS]).astype("f4"), + values["Z_B"].astype("f8"), + ) + + +def _match_tile(ra, dec, photo_ra, photo_dec): + """Return nearest DR6 row and separation for each position in one tile.""" + mid_dec = np.median(photo_dec) + cos_dec = np.cos(np.radians(mid_dec)) + ra_origin = np.median(photo_ra) if np.ptp(photo_ra) < 180 else 0 + photo_x = ((photo_ra - ra_origin + 180) % 360 - 180) * cos_dec + query_x = ((ra - ra_origin + 180) % 360 - 180) * cos_dec + return cKDTree(np.c_[photo_x, photo_dec]).query(np.c_[query_x, dec]) + + +def prepare_nz_photometry(selection_path, dr6_root, output_path): + """Write matched magnitudes, BPZ, shear weights and source row keys. + + The Parquet row order is retained. Missing and unreadable DR6 tiles are + reported and leave their selection rows unmatched. + """ + selection = pq.read_table( + selection_path, + columns=["TILE_ID", "NUMBER", "RA", "Dec", "w_des"], + ) + encoded = pc.dictionary_encode(selection["TILE_ID"]).combine_chunks() + tile_names = np.asarray(encoded.dictionary.to_pylist(), dtype=str) + tile_codes = encoded.indices.to_numpy() + name_order = np.argsort(tile_names) + tile_names = tile_names[name_order] + code_remap = np.empty_like(name_order) + code_remap[name_order] = np.arange(len(name_order)) + tile_codes = code_remap[tile_codes] + row_order = np.argsort(tile_codes, kind="stable") + tile_offsets = np.searchsorted( + tile_codes[row_order], np.arange(len(tile_names) + 1) + ) + nrows = len(tile_codes) + ra = selection["RA"].to_numpy().astype("f8") + dec = selection["Dec"].to_numpy().astype("f8") + + mag = np.full((nrows, len(BANDS)), np.nan, dtype="f4") + zb = np.full(nrows, np.nan, dtype="f8") + matched_rows = np.zeros(nrows, dtype=bool) + tile_has_dr6 = np.zeros(len(tile_names), dtype=bool) + root = Path(dr6_root) + + for tile_index, tile in enumerate(tile_names): + path = root / f"UNIONS.{tile}{DR6_SUFFIX}" + if not path.is_file(): + print(f"Missing DR6 tile {tile}: {path}") + continue + try: + photo_ra, photo_dec, photo_mag, photo_zb = _read_dr6_tile(path) + except (OSError, ValueError, fits.verify.VerifyError) as error: + print(f"Unreadable DR6 tile {tile}: {error}") + continue + tile_has_dr6[tile_index] = True + rows = row_order[tile_offsets[tile_index] : tile_offsets[tile_index + 1]] + if not len(rows) or not len(photo_ra): + continue + distance, nearest = _match_tile(ra[rows], dec[rows], photo_ra, photo_dec) + matched = distance < MATCH_RADIUS_DEG + rows = rows[matched] + nearest = nearest[matched] + mag[rows] = photo_mag[nearest] + zb[rows] = photo_zb[nearest] + matched_rows[rows] = True + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + output_path, + tile_names=tile_names, + tile=tile_codes.astype(np.int32), + mag=mag, + zb=zb, + number=selection["NUMBER"].to_numpy(), + weight=selection["w_des"].to_numpy().astype("f8"), + ) + missing = len(tile_names) - tile_has_dr6.sum() + print( + f"Matched {matched_rows.sum()} of {nrows} rows; {missing} tiles missing or unreadable" + ) diff --git a/src/sp_validation/nz_pipeline.py b/src/sp_validation/nz_pipeline.py new file mode 100644 index 000000000..1fc201d47 --- /dev/null +++ b/src/sp_validation/nz_pipeline.py @@ -0,0 +1,199 @@ +"""Run the catalogue's saved n(z) recipe independently of cosmo_val.""" + +import argparse +import json +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq +import yaml + +from . import nz_calibration as nz + + +def rebuild_model(config, workers): + """Train/group the map and derive gold on the recipe's reference sample.""" + cal = np.load(config["inputs"]["calibrators"]) + settings = config["som"] + model = nz.train(cal["mag"], workers=workers, **settings) + model["group_map"] = nz.group_cells(model, config["groups"]) + model["edges"] = np.array(config["edges"]) + source = np.load(config["inputs"]["gold_reference"]) + prelim, group, _ = nz.source_groups( + source["mag"], + source["photoz"], + source["weight"], + model, + offset=config["z2_offset"], + use_z2=False, + workers=workers, + ) + _, cal_cell = nz.assign(cal["mag"], model, workers=workers) + est = nz.calibrate( + prelim, + group, + source["weight"], + source["photoz"], + cal["z"], + cal["photoz"], + model["group_map"][cal_cell], + model["edges"], + config["groups"], + output_grid(config), + threshold=nz.gold_threshold(cal["z"], cal["photoz"], **config["gold_qc"]), + ) + model["gold"] = est["gold"] + return model + + +def output_grid(config): + grid = config["grid"] + return np.linspace( + grid["min"], grid["max"], round((grid["max"] - grid["min"]) / grid["step"]) + 1 + ) + + +def run(config, output, *, workers=1, rebuild=False): + """Recalibrate full-area weights with the frozen SOM, edges and gold masks.""" + output = Path(output) + output.mkdir(parents=True, exist_ok=True) + model = ( + rebuild_model(config, workers) + if rebuild + else dict(np.load(config["inputs"]["model"])) + ) + edges = np.asarray(config["edges"]) + # The recipe and the frozen definition describe the same catalogue selection. + np.testing.assert_array_equal(model["edges"], edges) + np.savez(output / "som_assignment_model.npz", **model) + cal = np.load(config["inputs"]["calibrators"]) + _, cal_cell = nz.assign(cal["mag"], model, workers=workers) + cal_group = model["group_map"][cal_cell] + data = np.load(config["photometry"]["output"]) + mag, photoz, weight = data["mag"], data["zb"], data["weight"] + grid = output_grid(config) + total_weight = weight.sum() + prelim, group, channel = nz.source_groups( + mag, + photoz, + weight, + model, + offset=config["z2_offset"], + use_z2=True, + workers=workers, + ) + print( + f"Assigned: {np.count_nonzero(prelim >= 0)} eligible rows", + flush=True, + ) + est = nz.calibrate( + prelim, + group, + weight, + photoz, + cal["z"], + cal["photoz"], + cal_group, + edges, + config["groups"], + grid, + gold=model["gold"], + ) + tom = nz.tom_bin_ids(prelim, group, est["gold"]) + retained = sum(b["weight_sum"] for b in est["bins"]) + complete = ( + nz.detected(mag[:, :4]).all(1) + & (channel > 0) + & np.isfinite(weight) + & (weight > 0) + ) + retention = dict( + all=1.0, + complete_photometry=float(weight[complete].sum() / total_weight), + bpz_bounds=float(weight[prelim >= 0].sum() / total_weight), + gold=float(retained / total_weight), + ) + for b in est["bins"]: + b["weight_fraction_of_all"] = b["weight_sum"] / float(total_weight) + nz.write_nz( + output / "nz_tomo.txt", + grid, + est["bin_hists"], + [b["weight_sum"] for b in est["bins"]], + ) + tile = pa.DictionaryArray.from_arrays( + pa.array(data["tile"]), pa.array(data["tile_names"].astype(str)) + ).cast(pa.string()) + pq.write_table( + pa.table( + dict( + TILE_ID=tile, + NUMBER=data["number"], + TOM_BIN_ID=tom, + ZBAND=channel, + ) + ), + output / "tom_bin.parquet", + ) + print("TOM_BIN_ID counts:", np.bincount(tom, minlength=len(edges)), flush=True) + np.savez_compressed( + output / "calibration_arrays.npz", + z=cal["z"], + cal_photoz=cal["photoz"], + cal_cell=cal_cell, + cal_group=cal_group, + grid=grid, + edges=edges, + gold=est["gold"], + cal_weights=est["cal_weights"], + bin_hists=est["bin_hists"], + ) + summary = dict( + bins=est["bins"], + retention=retention, + total_weight=float(total_weight), + n_rows=len(weight), + retained_rows=sum(b["n_rows"] for b in est["bins"]), + retained_weight_fraction=retention["gold"], + gold_qc_threshold=nz.gold_threshold( + cal["z"], cal["photoz"], **config["gold_qc"] + ), + method=config["name"], + limitation=config["limitation"], + ) + (output / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + (output / "recipe.yaml").write_text(yaml.safe_dump(config, sort_keys=False)) + return summary + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("recipe") + parser.add_argument("--output", help="override recipe output directory") + parser.add_argument("--workers", type=int, default=1) + parser.add_argument("--rebuild-model", action="store_true") + parser.add_argument("--figures", metavar="DIRECTORY") + args = parser.parse_args(argv) + config = yaml.safe_load(Path(args.recipe).read_text()) + output = Path(args.output or config["output"]) + run(config, output, workers=args.workers, rebuild=args.rebuild_model) + if args.figures: + from .nz_diagnostics import render + + render( + dict( + config["diagnostics"], + calibration_arrays=output / "calibration_arrays.npz", + model=output / "som_assignment_model.npz", + summary=output / "summary.json", + source_catalogue=config["inputs"]["gold_reference"], + side=config["som"]["side"], + workers=args.workers, + ), + args.figures, + ) + + +if __name__ == "__main__": + main() diff --git a/src/sp_validation/tests/test_column_schema.py b/src/sp_validation/tests/test_column_schema.py index 6181c886e..7642b379c 100644 --- a/src/sp_validation/tests/test_column_schema.py +++ b/src/sp_validation/tests/test_column_schema.py @@ -39,6 +39,11 @@ "R", "R11", "R22", + # DR6 photometry columns, outside the ShapePipe schema + "ALPHA_J2000", + "DELTA_J2000", + "MAG_GAAP_{}", + "Z_B", # FITS header keywords "EXTNAME", "R_S11", diff --git a/src/sp_validation/tests/test_config_paths_exist.py b/src/sp_validation/tests/test_config_paths_exist.py index b5f593e97..c6b13eba1 100644 --- a/src/sp_validation/tests/test_config_paths_exist.py +++ b/src/sp_validation/tests/test_config_paths_exist.py @@ -80,6 +80,8 @@ def _non_path_key(key: str) -> bool: def _pathish_value(value: str) -> bool: if not value or any(part in value for part in SKIP_VALUE_PARTS): return False + if value.startswith(("arc:", "vos:", "https://", "http://")): + return False return ( value.startswith(("/", "./", "../", "~")) or "/" in value @@ -91,6 +93,8 @@ def _walk_yaml( value, trail: tuple[str, ...] = (), base_dir: Path | None = None ) -> Iterator[tuple[str, str, Path | None]]: if isinstance(value, dict): + if value.get("host", "candide") != "candide": + return local_base = base_dir subdir = value.get("subdir") if isinstance(subdir, str) and _pathish_value(subdir): @@ -298,6 +302,8 @@ def test_cat_config_columns_exist_on_candide(): for version, entry in config.items(): if not isinstance(entry, dict) or "subdir" not in entry: continue + if entry.get("host", "candide") != "candide": + continue for kind in ("psf", "shear"): block = entry.get(kind) or {} if "path" not in block: @@ -328,3 +334,21 @@ def test_cat_config_columns_exist_on_candide(): for (version, kind), detail in unexpected.items() ) assert not healed, f"known column gaps now pass; drop them: {healed}" + + +def test_yaml_paths_respect_explicit_host(): + entries = { + "local": {"host": "candide", "path": "/data/local.fits"}, + "remote": {"host": "nibi", "path": "/data/remote.fits"}, + "unspecified": {"path": "/project/keep-checking.fits"}, + } + assert [(key, value) for key, value, _ in _walk_yaml(entries)] == [ + ("local.path", "/data/local.fits"), + ("unspecified.path", "/project/keep-checking.fits"), + ] + + +@pytest.mark.parametrize("scheme", ["arc:", "vos:", "https://", "http://"]) +def test_remote_urls_are_not_local_paths(scheme): + assert not _pathish_value(f"{scheme}home/masks/manual.hsp") + assert _pathish_value("/data/masks/manual.hsp") diff --git a/src/sp_validation/tests/test_nz_calibration.py b/src/sp_validation/tests/test_nz_calibration.py new file mode 100644 index 000000000..d695e0376 --- /dev/null +++ b/src/sp_validation/tests/test_nz_calibration.py @@ -0,0 +1,133 @@ +"""Known-answer checks for the provisional SOM n(z) calibration.""" + +import numpy as np +import pytest + +from sp_validation import nz_calibration as nz + + +def test_deposition_conserves_mass_and_mean(): + grid = np.linspace(0, 5, 501) + z = np.array([0.0, 0.007, 0.2851, 4.999, 5.0]) + weight = np.array([1.0, 2.0, 3.0, 0.5, 7.0]) + density = nz.deposit(z, weight, grid) + assert np.trapezoid(density, grid) == pytest.approx(weight.sum(), abs=1e-12) + assert np.trapezoid(grid * density, grid) == pytest.approx(z @ weight, abs=1e-12) + with pytest.raises(ValueError, match="outside output grid"): + nz.deposit(np.array([5.001]), np.ones(1), grid) + + +def test_z2_only_fills_missing_canonical_z(): + mag = np.array( + [ + [22, 21, 20, 19, 18, 17], + [22, 21, 20, 19, 99, 18], + [22, 21, 20, 19, -99, 99], + [22, 21, 20, 19, np.nan, 18], + ], + dtype="f4", + ) + mapped, channel = nz.z_magnitudes(mag, -0.1104888916015625) + np.testing.assert_array_equal(channel, [1, 2, 0, 2]) + assert mapped[0, 4] == 18 + expected = (mag[:, 5] - 0.1104888916015625).astype("f8") + np.testing.assert_array_equal(mapped[[1, 3], 4], expected[[1, 3]]) + assert np.isnan(mapped[2, 4]) + canonical, channel = nz.z_magnitudes(mag, -0.1104888916015625, use_z2=False) + np.testing.assert_array_equal(channel, [1, 0, 0, 0]) + assert np.isnan(canonical[1:, 4]).all() + + +def test_gold_requires_support_and_bpZ_agreement(): + # Supported matching group, unsupported group, discrepant group, empty source. + mask = nz.gold_mask( + np.array([10.0, 10.0, 10.0, 0.0]), + np.array([3.0, 3.0, 3.0, 0.0]), + np.array([2, 0, 1, 1]), + np.array([0.6, 0.0, 1.0, 1.0]), + threshold=0.4, + ) + np.testing.assert_array_equal(mask, [True, False, False, True]) + assert nz.gold_threshold(np.array([0.2, 0.3]), np.array([0.2, 0.3])) == 0.4 + + +def test_bin_edges_and_gold_ids(): + mag = np.tile([22.0, 21.0, 20.0, 19.0, 18.0, 99.0], (7, 1)).astype("f4") + model = dict( + code=nz.features(mag[:1, :5]), + center=np.zeros(11), + scale=np.ones(11), + group_map=np.array([0]), + edges=np.array([0.1, 0.37, 0.55, 0.76, 1.2]), + ) + bpz = np.array([0.099, 0.1, 0.37, 0.55, 0.76, 1.2, 0.4], dtype="f8") + weight = np.ones(7) + weight[-1] = 0 + prelim, group, _ = nz.source_groups(mag, bpz, weight, model, offset=0, chunk=2) + np.testing.assert_array_equal(prelim, [-1, 0, 1, 2, 3, -1, -1]) + ids = nz.tom_bin_ids(prelim, group, np.array([[1], [1], [0], [1]], bool)) + np.testing.assert_array_equal(ids, [0, 1, 2, 0, 4, 0, 0]) + + +def test_direct_weights_and_writer(tmp_path): + grid = np.linspace(0, 5, 501) + # Bin 1 has group weights 6 and 2, bin 2 has weight 4. + result = nz.calibrate( + np.array([0, 0, 1]), + np.array([0, 1, 1]), + np.array([6.0, 2.0, 4.0]), + np.array([0.2, 0.3, 0.8]), + np.array([0.1, 0.3, 0.5, 0.9]), + np.array([0.2, 0.2, 0.3, 0.8]), + np.array([0, 0, 1, 1]), + np.array([0.1, 0.5, 1.2]), + 2, + grid, + gold=np.ones((2, 2), bool), + ) + np.testing.assert_array_equal(result["cal_weights"], [[3, 3, 2, 0], [0, 0, 0, 4]]) + np.testing.assert_allclose([b["mean"] for b in result["bins"]], [0.275, 0.9]) + path = tmp_path / "nz.txt" + nz.write_nz(path, grid, result["bin_hists"], [8, 4]) + text = np.loadtxt(path) + assert text.shape == (501, 3) + np.testing.assert_allclose( + np.trapezoid(text[:, 1:], text[:, 0], axis=0), [2 / 3, 1 / 3], atol=1e-12 + ) + mixture = text[:, 1:].sum(1) + assert np.trapezoid(grid * mixture, grid) == pytest.approx( + (8 * 0.275 + 4 * 0.9) / 12, abs=1e-12 + ) + + +def test_nz_file_cosmo_val_and_sacc_roundtrip(tmp_path): + import sacc + + from sp_validation.cosmo_val import CosmologyValidation + from sp_validation.sacc_io import new_sacc + + grid = np.linspace(0, 2, 201) + hists = np.array( + [ + nz.deposit(np.array([0.2, 0.4]), np.array([0.5, 0.5]), grid), + nz.deposit(np.array([0.8, 1.1]), np.array([0.25, 0.75]), grid), + ] + ) + path = tmp_path / "nz.txt" + nz.write_nz(path, grid, hists, [2, 3]) + expected = np.loadtxt(path) + reader = object.__new__(CosmologyValidation) + reader.cc = {"test": {"shear": {"redshift_path": str(path)}}} + z, densities = reader.read_redshift_distribution("test", is_tomography=True) + np.testing.assert_array_equal(z, expected[:, 0]) + np.testing.assert_array_equal(densities, expected[:, 1:]) + _, summed = reader.read_redshift_distribution("test", is_tomography=False) + np.testing.assert_array_equal(summed, densities.sum(1)) + tracers = new_sacc([(z, d) for d in densities.T]) + fits_path = tmp_path / "nz.fits" + tracers.save_fits(str(fits_path)) + loaded = sacc.Sacc.load_fits(str(fits_path)) + for k in range(densities.shape[1]): + tracer = loaded.get_tracer(f"source_{k}") + np.testing.assert_array_equal(tracer.z, z) + np.testing.assert_array_equal(tracer.nz, densities[:, k]) diff --git a/src/sp_validation/tests/test_nz_diagnostics.py b/src/sp_validation/tests/test_nz_diagnostics.py new file mode 100644 index 000000000..b44f89e38 --- /dev/null +++ b/src/sp_validation/tests/test_nz_diagnostics.py @@ -0,0 +1,86 @@ +"""Diagnostics use the recipe's map and labels and preserve percentage text.""" + +import json + +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.text import Text + +from sp_validation import nz_diagnostics + + +def test_recipe_map_labels_and_percentages(tmp_path, monkeypatch): + grid = np.linspace(0, 2, 21) + cell = np.tile(np.arange(4), 2) + weights = np.tile([1.0, 2.0, 3.0, 4.0], (4, 2)) + paths = { + key: tmp_path / f"{key}.npz" + for key in [ + "calibration_arrays", + "model", + "source_catalogue", + "source_assignment", + ] + } + np.savez( + paths["calibration_arrays"], + z=np.linspace(0.15, 1.1, 8), + cal_photoz=np.linspace(0.15, 1.1, 8), + cal_group=cell, + cal_cell=cell, + edges=[0.1, 0.37, 0.55, 0.76, 1.2], + gold=np.ones((4, 4), bool), + cal_weights=weights, + grid=grid, + bin_hists=np.ones((4, len(grid))), + ) + np.savez(paths["model"], group_map=np.arange(4)) + np.savez( + paths["source_catalogue"], + mag=np.ones((4, 5)) * 20, + weight=np.arange(1, 5), + photoz=np.linspace(0.2, 1.0, 4), + ) + np.savez(paths["source_assignment"], cell=np.arange(4)) + summary = tmp_path / "summary.json" + summary.write_text( + json.dumps( + { + "bins": [ + {"gold_fraction": 0.9, "weight_fraction": 0.25} for _ in range(4) + ] + } + ) + ) + comparison = tmp_path / "nz.txt" + np.savetxt(comparison, np.column_stack([grid, np.ones((len(grid), 4))])) + config = dict( + paths, + summary=summary, + side=2, + source_label="Reference sample", + full_label="Full sample", + tile_label="Small sample", + published_label="Published sample", + tile_nz=comparison, + published_nz=comparison, + ) + texts = [] + + def inspect_and_close(fig, path): + labels = fig.findobj(Text) + assert all(not label.get_usetex() for label in labels) + texts.extend(label.get_text() for label in labels) + plt.close(fig) + + monkeypatch.setattr(nz_diagnostics, "_save", inspect_and_close) + with matplotlib.rc_context({"text.usetex": True}): + nz_diagnostics.render(config, tmp_path / "figures") + assert matplotlib.rcParams["text.usetex"] + assert any("% weight lost" in text for text in texts) + assert any("outliers" in text and "%" in text for text in texts) + assert "Full sample" in texts + assert "Small sample" in texts + assert any(text.startswith("Published sample, mean") for text in texts) + assert any("Reference sample" in text for text in texts) diff --git a/src/sp_validation/tests/test_nz_photometry.py b/src/sp_validation/tests/test_nz_photometry.py new file mode 100644 index 000000000..8a875e397 --- /dev/null +++ b/src/sp_validation/tests/test_nz_photometry.py @@ -0,0 +1,49 @@ +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from sp_validation import nz_photometry + + +def test_prepare_matches_by_position_and_preserves_source_rows(tmp_path, monkeypatch): + selection = tmp_path / "selection.parquet" + pq.write_table( + pa.table( + { + "TILE_ID": ["001.001", "002.001", "001.001"], + "NUMBER": [70, 20, 10], + "RA": [10.0 + 0.2 / 3600, 30.0, 10.0 + 0.7 / 3600], + "Dec": [2.0, 0.0, 2.0], + "w_des": [1.1, 1.2, 1.3], + } + ), + selection, + ) + dr6_root = tmp_path / "dr6" + dr6_root.mkdir() + (dr6_root / "UNIONS.001.001_SP_ugriz_photoz_ext.cat").touch() + + def read_tile(path): + return ( + np.array([10.0]), + np.array([2.0]), + np.array([[-99, 20, 21, 22, 23, 99]], dtype="f4"), + np.array([0.76], dtype="f8"), + ) + + monkeypatch.setattr(nz_photometry, "_read_dr6_tile", read_tile) + output = tmp_path / "matched_full.npz" + nz_photometry.prepare_nz_photometry(selection, dr6_root, output) + + with np.load(output) as cache: + np.testing.assert_array_equal(cache["tile_names"], ["001.001", "002.001"]) + np.testing.assert_array_equal(cache["number"], [70, 20, 10]) + assert np.isnan(cache["mag"][1]).all() + assert np.isnan(cache["zb"][1]) + np.testing.assert_array_equal(cache["tile"], [0, 1, 0]) + np.testing.assert_allclose(cache["weight"], [1.1, 1.2, 1.3]) + assert cache["zb"].dtype == np.dtype("f8") + assert cache["mag"].dtype == np.dtype("f4") + assert cache["zb"][0] == 0.76 + np.testing.assert_array_equal(cache["mag"][0], [-99, 20, 21, 22, 23, 99]) + np.testing.assert_array_equal(cache["mag"][2], cache["mag"][0])