From 8ae222bfb63900098dabb479d1cab6bb0e4cd498 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Tue, 21 Jul 2026 10:18:40 +0200 Subject: [PATCH 01/10] Start updating rho-tau-stats file for tomography --- src/sp_validation/cosmo_val/core.py | 13 +- .../cosmo_val/psf_systematics.py | 6 +- src/sp_validation/rho_tau.py | 117 ++++++++---------- 3 files changed, 71 insertions(+), 65 deletions(-) diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index ee65f658..de819751 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -465,11 +465,20 @@ def results_objectwise(self): self._results_objectwise = self.init_results(objectwise=True) return self._results_objectwise - def basename(self, version, treecorr_config=None, npatch=None): + def basename( + self, + version, + tomo_bin_a="all", + tomo_bin_b=None, + treecorr_config=None, + npatch=None, + ): cfg = treecorr_config or self.treecorr_config patches = npatch or self.npatch + tomo_bin_a_str = f"tomo_bin_{tomo_bin_a}" + tomo_bin_b_str = f"_tomo_bin_{tomo_bin_b}" if tomo_bin_b is not None else "" return ( - f"{version}_minsep={cfg['min_sep']}" + f"{version}_{tomo_bin_a_str}{tomo_bin_b_str}_minsep={cfg['min_sep']}" f"_maxsep={cfg['max_sep']}" f"_nbins={cfg['nbins']}" f"_npatch={patches}" diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 9f8a247e..327c74bb 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -23,7 +23,7 @@ class PSFSystematicsMixin: - def calculate_rho_tau_stats(self): + def calculate_rho_tau_stats(self, tomography=True): out_dir = f"{self.cc['paths']['output']}/rho_tau_stats" if not os.path.exists(out_dir): os.mkdir(out_dir) @@ -50,12 +50,16 @@ def calculate_rho_tau_stats(self): def rho_stat_handler(self): if not hasattr(self, "_rho_stat_handler"): self.calculate_rho_tau_stats() + if self.compute_tomography: + self.calculate_rho_tau_stats(tomography=True) return self._rho_stat_handler @property def tau_stat_handler(self): if not hasattr(self, "_tau_stat_handler"): self.calculate_rho_tau_stats() + if self.compute_tomography: + self.calculate_rho_tau_stats(tomography=True) return self._tau_stat_handler def plot_rho_stats(self, abs=False): diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 0a311432..81b40aa6 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -1,3 +1,4 @@ +import gc import os import time from pathlib import Path @@ -10,37 +11,16 @@ # `from sp_validation.rho_tau import SquareRootScale` keeps working. from sp_validation.plots import SquareRootScale # noqa: F401 -not_square_size = [ - "DES", - "SP_v1.3_LFmask_8k", - "SP_v1.3_LFmask_8k_no_alpha", - "SP_v1.3_LFmask_8k_li_2024", - "SP_v1.3_LFmask_8k_SN8", - "SP_v1.3_LFmask_8k_F2", -] - def _extract_xip(correlations): """Return flattened array of xip values from a list of correlations.""" return np.array([corr.xip for corr in correlations]).flatten() -def get_params_rho_tau(cat, survey="other"): +def get_params_rho_tau(cat): # Set parameters params = {} - # TODO to yaml file - if survey == "DES": - params["patch_number"] = 120 - print("DES, jackknife patch number = 120") - elif survey == "SP_axel_v0.0": - params["patch_number"] = 120 - print("SP_Axel_v0.0, jackknife patch number =120") - elif survey == "SP_v1.4-P3" or survey == "SP_v1.4-P3_LFmask": - params["patch_number"] = 120 - print("SP_v1.4, jackknife patch number =120") - else: - params["patch_number"] = 150 params["ra_PSF_col"] = cat["psf"]["ra_col"] params["dec_PSF_col"] = cat["psf"]["dec_col"] params["e1_PSF_col"] = cat["psf"]["e1_PSF_col"] @@ -49,18 +29,19 @@ def get_params_rho_tau(cat, survey="other"): params["e2_star_col"] = cat["psf"]["e2_star_col"] params["PSF_size"] = cat["psf"]["PSF_size"] params["star_size"] = cat["psf"]["star_size"] - params["square_size"] = survey not in not_square_size - if survey != "DES": - params["PSF_flag"] = cat["psf"]["PSF_flag"] - params["star_flag"] = cat["psf"]["star_flag"] + params["square_size"] = cat["psf"].get("square_size", False) + params["PSF_flag"] = cat["psf"].get("PSF_flag") + params["star_flag"] = cat["psf"].get("star_flag") params["ra_units"] = "deg" params["dec_units"] = "deg" + params["patch_number"] = cat.get("patch_number", 100) # Default patch number to 100 params["ra_col"] = cat["shear"].get("ra_col", "RA") params["dec_col"] = cat["shear"].get("dec_col", "Dec") params["w_col"] = cat["shear"]["w_col"] params["e1_col"] = cat["shear"]["e1_col"] params["e2_col"] = cat["shear"]["e2_col"] + params["tomo_bin_id_col"] = cat["shear"].get("tomo_bin_id_col") params["R11"] = cat["shear"].get("R11") params["R22"] = cat["shear"].get("R22") @@ -74,12 +55,17 @@ def get_rho_tau_w_cov( outdir, base, method, + mask_star=None, + mask_gal=None, cov_rho=False, - npatch=None, + ncov=100, + compute_minus=True, + **kwargs, ): """Compute rho/tau statistics and, if requested, their covariance.""" if method == "th": - nbin_ang, nbin_rad = 100, 200 + nbin_ang, nbin_rad = kwargs.get("nbin_ang", 100), kwargs.get("nbin_rad", 200) + compute_minus = kwargs.get("compute_minus", True) rho_stat_handler, tau_stat_handler = get_rho_tau( config, version, @@ -96,17 +82,13 @@ def get_rho_tau_w_cov( base, nbin_ang=nbin_ang, nbin_rad=nbin_rad, + compute_minus=compute_minus, + mask_star=mask_star, + mask_gal=mask_gal, ) return rho_stat_handler, tau_stat_handler elif method == "jk": - return get_jackknife_cov( - config, - version, - treecorr_config, - outdir, - base, - npatch=npatch, - ) + return get_jackknife_cov(config, version, treecorr_config, outdir, base) elif method == "sim": tau_cov_path = Path(outdir) / f"cov_tau_{base}_th.npy" @@ -135,8 +117,10 @@ def get_rho_tau( treecorr_config, outdir, base, + mask_star=None, + mask_gal=None, cov_rho=False, - npatch=None, + force_run=False, ): """ Compute rho and tau statistics for a given version of the catalogue. @@ -151,9 +135,19 @@ def get_rho_tau( TreeCorr configuration (must include 'min_sep', 'max_sep', and 'nbins'). outdir : str Output directory. + base : str + Base name for output files. + mask_star : array-like, optional + Boolean mask for stars. If None, no masking is applied. + mask_gal : array-like, optional + Boolean mask for galaxies. If None, no masking is applied. + cov_rho : bool, optional + If True, compute the covariance of rho statistics. + force_run : bool, optional + If True, force the computation even if output files already exist. """ - params = get_params_rho_tau(config[version], survey=version) + params = get_params_rho_tau(config[version]) print("Compute Rho and Tau statistics for the version: ", version) start_time = time.time() @@ -169,19 +163,18 @@ def get_rho_tau( rho_stats_exists = rho_path.exists() cov_exists = True if not cov_rho else cov_rho_path.exists() - need_compute = (not rho_stats_exists) or (not cov_exists) + need_compute = (not rho_stats_exists) or (not cov_exists) or force_run if need_compute: rho_stat_handler.catalogs.set_params(params, outdir) - mask = version != "DES" square_size = params["square_size"] rho_stat_handler.build_cat_to_compute_rho( config[version]["psf"]["path"], catalog_id=catalog_id, square_size=square_size, - mask=mask, + mask=mask_star, hdu=( config[version]["psf"]["hdu"] if config[version]["psf"]["hdu"] is not None @@ -210,14 +203,12 @@ def get_rho_tau( verbose=True, ) - if tau_path.exists(): + if tau_path.exists() and not force_run: print(f"Skipping tau statistics computation, file {tau_path} already exists.") tau_stat_handler.load_tau_stats(tau_path.name) else: tau_stat_handler.catalogs.set_params(params, outdir) - mask = version != "DES" - square_size = params["square_size"] # Build the different catalogs if necessary @@ -227,7 +218,7 @@ def get_rho_tau( cat_type="psf", catalog_id=version, square_size=square_size, - mask=mask, + mask=mask_star, hdu=( config[version]["psf"]["hdu"] if config[version]["psf"]["hdu"] is not None @@ -241,7 +232,7 @@ def get_rho_tau( cat_type="gal", catalog_id=version, square_size=square_size, - mask=mask, + mask=mask_gal, ) # function to extract the tau_+ @@ -259,17 +250,17 @@ def get_theory_cov( base, nbin_ang=100, nbin_rad=100, + compute_minus=True, + mask_star=None, # TODO add the masking in the theory covariance + mask_gal=None, ): """ Compute an analytical estimate of the covariance matrix of rho and tau-statistics. """ - params = get_params_rho_tau(config[version], survey=version) + params = get_params_rho_tau(config[version]) info = config[version] - A = info["cov_th"]["A"] * 60 * 60 - n_e = info["cov_th"]["n_e"] - n_psf = info["cov_th"]["n_psf"] path_gal = info["shear"]["path"] path_psf = info["psf"]["path"] @@ -289,21 +280,23 @@ def get_theory_cov( path_psf=path_psf, hdu_psf=hdu_psf, treecorr_config=treecorr_config, - A=A, - n_e=n_e, - n_psf=n_psf, params=params, + mask_star=mask_star, + mask_gal=mask_gal, ) elapsed = time.time() - start_time print(f"--- Rho/tau statistics for covariance computed in {elapsed:.2f}s ---") - cov = cov_tau_th.build_cov(nbin_ang=nbin_ang, nbin_rad=nbin_rad) + cov = cov_tau_th.build_cov( + nbin_ang=nbin_ang, nbin_rad=nbin_rad, compute_minus=compute_minus + ) print(f"--- Covariance matrix assembled in {time.time() - start_time:.2f}s ---") target_cov.parent.mkdir(parents=True, exist_ok=True) np.save(target_cov, cov) print("Saved covariance matrix of version: ", version) del cov_tau_th + gc.collect() return @@ -315,6 +308,9 @@ def get_jackknife_cov( base, npatch, ncov=100, + mask_star=None, + mask_gal=None, + force_run=False, ): """ Compute the covariance matrix of rho and tau-statistics using the jackknife method. @@ -325,7 +321,7 @@ def get_jackknife_cov( tau_filename = f"tau_stats_{base}.fits" tau_cov_path = Path(outdir) / f"cov_tau_{base}_jk.npy" - if tau_cov_path.exists(): + if tau_cov_path.exists() and not force_run: print(f"Skipping covariance computation, file {tau_cov_path} already exists.") rho_stat_handler = RhoStat( output=outdir, treecorr_config=treecorr_config, verbose=False @@ -340,9 +336,9 @@ def get_jackknife_cov( rho_path = Path(outdir) / rho_filename tau_path = Path(outdir) / tau_filename - if rho_path.exists(): + if rho_path.exists() and not force_run: rho_stat_handler.load_rho_stats(rho_path.name) - if tau_path.exists(): + if tau_path.exists() and not force_run: tau_stat_handler.load_tau_stats(tau_path.name) return rho_stat_handler, tau_stat_handler @@ -377,7 +373,7 @@ def get_jackknife_cov( config[version]["psf"]["path"], catalog_id=version + str(i), square_size=square_size, - mask=False, + mask=mask_star, hdu=config[version]["psf"]["hdu"], ) @@ -391,7 +387,7 @@ def get_jackknife_cov( cat_type="gal", catalog_id=version + str(i), square_size=square_size, - mask=False, + mask=mask_gal, ) else: @@ -567,7 +563,6 @@ def get_samples_emcee( def get_samples_lsq( psf_fitter, - version, base, apply_debias=None, cov_type="jk", @@ -578,8 +573,6 @@ def get_samples_lsq( ---------- psf_fitter : PSFFitter PSF fitter instance managing rho/tau statistics and covariances. - version : str - Catalog identifier whose rho/tau statistics are sampled. base : str Precomputed basename for locating statistics/covariance files. apply_debias : int or None, optional From 0ce5416a2a3d37561d35766252658df0e3694c57 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Tue, 21 Jul 2026 10:26:09 +0200 Subject: [PATCH 02/10] Bring parent branch sp_validation-extend-to-tomography to the child branch for rho-tau-statistics --- src/sp_validation/rho_tau.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index f6569104..be32fa98 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -359,6 +359,8 @@ def get_jackknife_cov( tau_stat_handler.catalogs.set_params(params, outdir) + square_size = params["square_size"] + for i in range(ncov): tau_chunk = outdir + f"/cov_tau_{version}{i}.npy" rho_chunk = outdir + f"/cov_rho_{version}{i}.npy" From 4912b500fd6deef1c8e4edb05892e79c4f90400d Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Tue, 21 Jul 2026 11:50:16 +0200 Subject: [PATCH 03/10] Update calculation function to work in tomographic context. --- .../cosmo_val/psf_systematics.py | 271 +++++++++++------- src/sp_validation/rho_tau.py | 124 ++++---- 2 files changed, 240 insertions(+), 155 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 327c74bb..30957774 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -10,6 +10,7 @@ import matplotlib.pyplot as plt import numpy as np import treecorr +from astropy.io import fits from cs_util import plots as cs_plots from shear_psf_leakage import leakage from shear_psf_leakage import plots as psfleak_plots @@ -22,30 +23,9 @@ ) +# TODO: Reorganise the order of functions so it is more readable class PSFSystematicsMixin: - def calculate_rho_tau_stats(self, tomography=True): - out_dir = f"{self.cc['paths']['output']}/rho_tau_stats" - if not os.path.exists(out_dir): - os.mkdir(out_dir) - - self.print_start("Rho stats") - for ver in self.versions: - base = self.basename(ver) - rho_stat_handler, tau_stat_handler = get_rho_tau_w_cov( - self.cc, - ver, - self.treecorr_config, - out_dir, - base, - method=self.cov_estimate_method, - cov_rho=self.compute_cov_rho, - npatch=self.npatch, - ) - self.print_done("Rho stats finished") - - self._rho_stat_handler = rho_stat_handler - self._tau_stat_handler = tau_stat_handler - + # --- property definitions --- @property def rho_stat_handler(self): if not hasattr(self, "_rho_stat_handler"): @@ -62,6 +42,152 @@ def tau_stat_handler(self): self.calculate_rho_tau_stats(tomography=True) return self._tau_stat_handler + @property + def psf_fitter(self): + if not hasattr(self, "_psf_fitter"): + self._psf_fitter = PSFErrorFit( + self.rho_stat_handler, + self.tau_stat_handler, + self.rho_stat_handler.catalogs._output, + ) + return self._psf_fitter + + @property + def rho_tau_fits(self): + if not hasattr(self, "_rho_tau_fits"): + self.calculate_rho_tau_fits() + if self.compute_tomography: + self.calculate_rho_tau_fits(tomography=True) + return self._rho_tau_fits + + @property + def xi_psf_sys(self): + if not hasattr(self, "_xi_psf_sys"): + self.calculate_rho_tau_fits() + return self._xi_psf_sys + + # --- calculate functions --- + def calculate_rho_tau_stats(self, tomography=True): + out_dir = f"{self.cc['paths']['output']}/rho_tau_stats" + if not os.path.exists(out_dir): + os.mkdir(out_dir) + + self.print_start("Rho stats") + for ver in self.versions: + # Get the tomographic bins + if tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise ValueError( + f"Version {ver} does not have tomography information." + ) + + else: + tomo_bin_ids, tomo_bin_pairs = "all", [("all", "all")] + + for tomo_bin_id in tomo_bin_ids: + base_rho = self.basename(ver) + base_tau = self.basename(ver, tomo_bin_a=tomo_bin_id) + + # Get the selection for galaxies + mask_gal = self._get_galaxy_mask(ver, tomo_bin_id) + + # Get the selection for stars + mask_star = self._get_star_mask(ver) + + # Compute rho and tau statistics + rho_stat_handler, tau_stat_handler = get_rho_tau_w_cov( + self.cc, + ver, + self.treecorr_config, + out_dir, + base_rho, + base_tau, + mask_star=mask_star, + mask_gal=mask_gal, + method=self.cov_estimate_method, + cov_rho=self.compute_cov_rho, + npatch=self.npatch, + ) + self.print_done("Rho stats finished") + + self._rho_stat_handler = rho_stat_handler + self._tau_stat_handler = tau_stat_handler + + # --- utility functions --- + def _get_galaxy_mask(self, ver, tomo_bin_id): + cat_gal = fits.getdata(self.cc["shear"]["path"]) + if tomo_bin_id != "all": + gal_mask = cat_gal[self.cc[ver]["shear"]["tomo_bin_col"]] == tomo_bin_id + else: + gal_mask = np.ones(len(cat_gal), dtype=bool) + return gal_mask + + def _get_star_mask(self, ver): + cat_star = fits.getdata( + self.cc[ver]["psf"]["path"], hdu=self.cc[ver]["psf"]["hdu"] + ) + PSF_flag = self.cc[ver]["psf"].get("PSF_flag") + star_flag = self.cc[ver]["psf"].get("star_flag") + if PSF_flag is not None: + if star_flag is not None: + star_mask = (cat_star[PSF_flag] == 0) & (cat_star[star_flag] == 0) + else: + star_mask = cat_star[PSF_flag] == 0 + else: + star_mask = np.ones(len(cat_star), dtype=bool) + return star_mask + + def set_params_rho_tau(self, ver, params, params_psf): + params = {**params} + + params["ra_PSF_col"] = params_psf["ra_col"] + params["dec_PSF_col"] = params_psf["dec_col"] + params["e1_PSF_col"] = params_psf["e1_PSF_col"] + params["e2_PSF_col"] = params_psf["e2_PSF_col"] + params["e1_star_col"] = params_psf["e1_star_col"] + params["e2_star_col"] = params_psf["e2_star_col"] + params["PSF_size"] = params_psf["PSF_size"] + params["star_size"] = params_psf["star_size"] + params["PSF_flag"] = params_psf.get("PSF_flag") + params["star_flag"] = params_psf.get("star_flag") + params["ra_units"] = "deg" + params["dec_units"] = "deg" + + params["w_col"] = self.cc[ver]["shear"]["w_col"] + + return params + + def get_samples(self, version, params, tomo_bin_id, track_result=False): + npatch = params["patch_number"] if self.cov_estimate_method == "jk" else None + + base_rho = self.basename(version) + base_tau = self.basename(version, tomo_bin_a=tomo_bin_id) + flat_samples, result, q = get_samples( + self.psf_fitter, + version, + base_rho, + base_tau, + cov_type=self.cov_estimate_method, + apply_debias=npatch, + sampler=self.rho_tau_method, + ) + + if track_result: + self.rho_tau_fits["flat_sample_list"].append(flat_samples) + self.rho_tau_fits["result_list"].append(result) + self.rho_tau_fits["q_list"].append(q) + + self.psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") + nbins = self.psf_fitter.rho_stat_handler._treecorr_config["nbins"] + xi_psf_sys_samples = np.array( + [self.psf_fitter.compute_xi_psf_sys(sample) for sample in flat_samples] + ).reshape(-1, nbins) + + return xi_psf_sys_samples + + # --- plotting functions --- def plot_rho_stats(self, abs=False): filenames = [f"rho_stats_{self.basename(ver)}.fits" for ver in self.versions] @@ -102,46 +228,7 @@ def plot_tau_stats(self, plot_tau_m=False): + f"{os.path.abspath(self.tau_stat_handler.catalogs._output)}/{savefig}", ) - def set_params_rho_tau(self, params, params_psf, survey="other"): - params = {**params} - if survey in ("DES", "SP_axel_v0.0", "SP_axel_v0.0_repr"): - params["patch_number"] = 120 - print("DES, jackknife patch number = 120") - elif survey in ("SP_v1.4-P3", "SP_v1.4-P3_LFmask"): - params["patch_number"] = 120 - print("SP_v1.4, jackknife patch number =120") - else: - params["patch_number"] = 150 - - params["ra_PSF_col"] = params_psf["ra_col"] - params["dec_PSF_col"] = params_psf["dec_col"] - params["e1_PSF_col"] = params_psf["e1_PSF_col"] - params["e2_PSF_col"] = params_psf["e2_PSF_col"] - params["e1_star_col"] = params_psf["e1_star_col"] - params["e2_star_col"] = params_psf["e2_star_col"] - params["PSF_size"] = params_psf["PSF_size"] - params["star_size"] = params_psf["star_size"] - if survey != "DES": - params["PSF_flag"] = params_psf["PSF_flag"] - params["star_flag"] = params_psf["star_flag"] - params["ra_units"] = "deg" - params["dec_units"] = "deg" - - params["w_col"] = self.cc[survey]["shear"]["w_col"] - - return params - - @property - def psf_fitter(self): - if not hasattr(self, "_psf_fitter"): - self._psf_fitter = PSFErrorFit( - self.rho_stat_handler, - self.tau_stat_handler, - self.rho_stat_handler.catalogs._output, - ) - return self._psf_fitter - - def calculate_rho_tau_fits(self): + def calculate_rho_tau_fits(self, tomography=True): assert self.rho_tau_method != "none" # this initializes the rho_tau_fits attribute @@ -151,47 +238,29 @@ def calculate_rho_tau_fits(self): self._xi_psf_sys = {} for ver in self.versions: params = self.set_params_rho_tau( - self.results[ver]._params, - self.cc[ver]["psf"], - survey=ver, - ) - - npatch = {"sim": 300, "jk": params["patch_number"]}.get( - self.cov_estimate_method, None + ver, self.results[ver]._params, self.cc[ver]["psf"] ) - base = self.basename(ver) + # Get the tomographic bins + if tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - flat_samples, result, q = get_samples( - self.psf_fitter, - ver, - base, - cov_type=self.cov_estimate_method, - apply_debias=npatch, - sampler=self.rho_tau_method, - ) + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise ValueError( + f"Version {ver} does not have tomography information." + ) - self.rho_tau_fits["flat_sample_list"].append(flat_samples) - self.rho_tau_fits["result_list"].append(result) - self.rho_tau_fits["q_list"].append(q) + else: + tomo_bin_ids, tomo_bin_pairs = "all", [("all", "all")] - self.psf_fitter.load_rho_stat(f"rho_stats_{self.basename(ver)}.fits") - nbins = self.psf_fitter.rho_stat_handler._treecorr_config["nbins"] - xi_psf_sys_samples = np.array( - [self.psf_fitter.compute_xi_psf_sys(sample) for sample in flat_samples] - ).reshape(-1, nbins) - - self._xi_psf_sys[ver] = { - "mean": np.mean(xi_psf_sys_samples, axis=0), - "var": np.var(xi_psf_sys_samples, axis=0), - "quantiles": np.quantile(xi_psf_sys_samples, quantiles, axis=0), - } + for tomo_bin_id in tomo_bin_ids: + xi_psf_sys_samples = self.get_samples(ver, params, tomo_bin_id) - @property - def rho_tau_fits(self): - if not hasattr(self, "_rho_tau_fits"): - self.calculate_rho_tau_fits() - return self._rho_tau_fits + self._xi_psf_sys[ver][f"tomo_bin_{tomo_bin_id}"] = { + "mean": np.mean(xi_psf_sys_samples, axis=0), + "var": np.var(xi_psf_sys_samples, axis=0), + "quantiles": np.quantile(xi_psf_sys_samples, quantiles, axis=0), + } def plot_rho_tau_fits(self): out_dir = self.rho_stat_handler.catalogs._output @@ -280,12 +349,6 @@ def plot_rho_tau_fits(self): f"{yscale}-scale xi_psf_sys terms plot saved to {out_path}" ) - @property - def xi_psf_sys(self): - if not hasattr(self, "_xi_psf_sys"): - self.calculate_rho_tau_fits() - return self._xi_psf_sys - def set_params_leakage_scale(self, ver): params_in = {} diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index be32fa98..73ec4406 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -53,7 +53,8 @@ def get_rho_tau_w_cov( version, treecorr_config, outdir, - base, + base_rho, + base_tau, method, mask_star=None, mask_gal=None, @@ -71,15 +72,18 @@ def get_rho_tau_w_cov( version, treecorr_config, outdir, - base, + base_rho, + base_tau, cov_rho=cov_rho, + mask_star=mask_star, + mask_gal=mask_gal, ) get_theory_cov( config, version, treecorr_config, outdir, - base, + base_tau, nbin_ang=nbin_ang, nbin_rad=nbin_rad, compute_minus=compute_minus, @@ -88,9 +92,11 @@ def get_rho_tau_w_cov( ) return rho_stat_handler, tau_stat_handler elif method == "jk": - return get_jackknife_cov(config, version, treecorr_config, outdir, base) + return get_jackknife_cov( + config, version, treecorr_config, outdir, base_rho, base_tau + ) elif method == "sim": - tau_cov_path = Path(outdir) / f"cov_tau_{base}_th.npy" + tau_cov_path = Path(outdir) / f"cov_tau_{base_tau}_th.npy" if tau_cov_path.exists(): print(f"Found existing covariance at {tau_cov_path}") @@ -100,8 +106,11 @@ def get_rho_tau_w_cov( version, treecorr_config, outdir, - base, + base_rho, + base_tau, cov_rho=cov_rho, + mask_star=mask_star, + mask_gal=mask_gal, ) else: raise ValueError( @@ -116,7 +125,8 @@ def get_rho_tau( version, treecorr_config, outdir, - base, + base_rho, + base_tau, mask_star=None, mask_gal=None, cov_rho=False, @@ -153,8 +163,8 @@ def get_rho_tau( start_time = time.time() outdir_path = Path(outdir) - rho_path = outdir_path / f"rho_stats_{base}.fits" - catalog_id = f"{base}_jk" if cov_rho else base + rho_path = outdir_path / f"rho_stats_{base_rho}.fits" + catalog_id = f"{base_rho}_jk" if cov_rho else base_rho cov_rho_path = outdir_path / f"cov_rho_{catalog_id}.npy" if cov_rho else None rho_stat_handler = RhoStat( @@ -194,7 +204,7 @@ def get_rho_tau( print(f"Skipping rho statistics computation, file {rho_path} already exists.") rho_stat_handler.load_rho_stats(rho_path.name) - tau_path = outdir_path / f"tau_stats_{base}.fits" + tau_path = outdir_path / f"tau_stats_{base_tau}.fits" tau_stat_handler = TauStat( catalogs=rho_stat_handler.catalogs, @@ -305,7 +315,8 @@ def get_jackknife_cov( version, treecorr_config, outdir, - base, + base_rho, + base_tau, npatch, ncov=100, mask_star=None, @@ -316,10 +327,10 @@ def get_jackknife_cov( Compute the covariance matrix of rho and tau-statistics using the jackknife method. Also compute rho and tau-statistics. """ - - rho_filename = f"rho_stats_{base}.fits" - tau_filename = f"tau_stats_{base}.fits" - tau_cov_path = Path(outdir) / f"cov_tau_{base}_jk.npy" + # TODO: Reorganise this to avoid recomputing unnecessarily the rho-stats multiple times. + rho_filename = f"rho_stats_{base_rho}.fits" + tau_filename = f"tau_stats_{base_tau}.fits" + tau_cov_path = Path(outdir) / f"cov_tau_{base_tau}_jk.npy" if tau_cov_path.exists() and not force_run: print(f"Skipping covariance computation, file {tau_cov_path} already exists.") @@ -342,7 +353,7 @@ def get_jackknife_cov( tau_stat_handler.load_tau_stats(tau_path.name) return rho_stat_handler, tau_stat_handler - params = get_params_rho_tau(config[version], survey=version) + params = get_params_rho_tau(config[version]) rho_stat_handler = RhoStat( output=outdir, treecorr_config=treecorr_config, verbose=False @@ -362,8 +373,8 @@ def get_jackknife_cov( square_size = params["square_size"] for i in range(ncov): - tau_chunk = outdir + f"/cov_tau_{version}{i}.npy" - rho_chunk = outdir + f"/cov_rho_{version}{i}.npy" + tau_chunk = outdir + f"/cov_tau_{base_tau}{i}.npy" + rho_chunk = outdir + f"/cov_rho_{base_rho}{i}.npy" if not (os.path.exists(tau_chunk) and os.path.exists(rho_chunk)): print(f"Computing rho-statistics for {version} (patch {i + 1}/{ncov})") @@ -437,13 +448,13 @@ def get_jackknife_cov( ) tau_dict[f"gal_{version}{i + 1}"] = tau_dict.pop(f"gal_{version}{i}") - cov_tau_loc = np.zeros_like(np.load(outdir + f"/cov_tau_{version}0.npy")) - cov_rho_loc = np.zeros_like(np.load(outdir + f"/cov_rho_{version}0.npy")) + cov_tau_loc = np.zeros_like(np.load(outdir + f"/cov_tau_{base_tau}0.npy")) + cov_rho_loc = np.zeros_like(np.load(outdir + f"/cov_rho_{base_rho}0.npy")) for i in range(ncov): - cov_tau_loc += np.load(outdir + f"/cov_tau_{version}{i}.npy") - cov_rho_loc += np.load(outdir + f"/cov_rho_{version}{i}.npy") - os.remove(outdir + f"/cov_tau_{version}{i}.npy") - os.remove(outdir + f"/cov_rho_{version}{i}.npy") + cov_tau_loc += np.load(outdir + f"/cov_tau_{base_tau}{i}.npy") + cov_rho_loc += np.load(outdir + f"/cov_rho_{base_rho}{i}.npy") + os.remove(outdir + f"/cov_tau_{base_tau}{i}.npy") + os.remove(outdir + f"/cov_rho_{base_rho}{i}.npy") cov_tau = cov_tau_loc / ncov cov_rho = cov_rho_loc / ncov @@ -451,7 +462,7 @@ def get_jackknife_cov( tau_cov_path.parent.mkdir(parents=True, exist_ok=True) np.save(tau_cov_path, cov_tau) - cov_rho_path = Path(outdir) / f"cov_rho_{base}_jk.npy" + cov_rho_path = Path(outdir) / f"cov_rho_{base_rho}_jk.npy" cov_rho_path.parent.mkdir(parents=True, exist_ok=True) np.save(cov_rho_path, cov_rho) @@ -461,7 +472,8 @@ def get_jackknife_cov( def get_samples( psf_fitter, version, - base, + base_rho, + base_tau, cov_type="jk", apply_debias=None, sampler="emcee", @@ -474,8 +486,10 @@ def get_samples( PSF fitter instance that provides ``load_*`` helpers. version : str Catalog identifier whose rho/tau statistics are sampled. - base : str - Precomputed basename (e.g. ``SP_v1.4_minsep=…``) used for filenames. + base_rho : str + Precomputed basename (e.g. ``SP_v1.4_minsep=…``) used for rho-stat filenames. + base_tau : str + Precomputed basename (e.g. ``SP_v1.4_minsep=…``) used for tau-stat filenames. cov_type : str, optional Covariance label (``'jk'``, ``'th'``, or ``'sim'``). Defaults to ``'jk'``. apply_debias : int or None, optional @@ -488,7 +502,8 @@ def get_samples( return get_samples_emcee( psf_fitter, version, - base, + base_rho, + base_tau, cov_type=cov_type, apply_debias=apply_debias, ) @@ -496,7 +511,8 @@ def get_samples( return get_samples_lsq( psf_fitter, version, - base, + base_rho, + base_tau, cov_type=cov_type, apply_debias=apply_debias, ) @@ -507,7 +523,8 @@ def get_samples( def get_samples_emcee( psf_fitter, version, - base, + base_rho, + base_tau, nwalkers=124, nsamples=10000, cov_type="jk", @@ -521,8 +538,10 @@ def get_samples_emcee( PSF fitter instance managing rho/tau statistics and covariances. version : str Catalog identifier whose rho/tau statistics are sampled. - base : str - Precomputed basename for locating statistics/covariance files. + base_rho : str + Precomputed basename for locating rho statistics/covariance files. + base_tau : str + Precomputed basename for locating tau statistics/covariance files. nwalkers : int, optional Number of walkers for the MCMC run (default ``124``). nsamples : int, optional @@ -533,20 +552,20 @@ def get_samples_emcee( Jackknife patch count applied during debiasing. Disabled when ``None``. """ # Load rho and tau stats - psf_fitter.load_rho_stat(f"rho_stats_{base}.fits") - psf_fitter.load_tau_stat(f"tau_stats_{base}.fits") + psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") + psf_fitter.load_tau_stat(f"tau_stats_{base_tau}.fits") # Check if the path exists (use base for cache key to account for different TreeCorr configs) - sample_file_path = psf_fitter.get_sample_file_path(base) + sample_file_path = psf_fitter.get_sample_file_path(base_tau) if os.path.exists(sample_file_path): print(f"Skipping sampling; {sample_file_path} exists.") - flat_samples = psf_fitter.load_samples(base) - mcmc_result, q = psf_fitter.get_mcmc_from_samples(base) + flat_samples = psf_fitter.load_samples(base_tau) + mcmc_result, q = psf_fitter.get_mcmc_from_samples(base_tau) print(mcmc_result) # Or run MCMC else: print("MCMC sampling") - cov_filename = f"cov_tau_{base}_{cov_type}.npy" + cov_filename = f"cov_tau_{base_tau}_{cov_type}.npy" psf_fitter.load_covariance(cov_filename, cov_type="tau") debias_npatch = apply_debias if (apply_debias is not None) else None @@ -555,15 +574,16 @@ def get_samples_emcee( nsamples=nsamples, npatch=debias_npatch, apply_debias=debias_npatch is not None, - savefig="mcmc_samples_" + version + ".png", + savefig="mcmc_samples_" + base_tau + ".png", ) - psf_fitter.save_samples(flat_samples, base) + psf_fitter.save_samples(flat_samples, base_tau) return flat_samples, mcmc_result, q def get_samples_lsq( psf_fitter, - base, + base_rho, + base_tau, apply_debias=None, cov_type="jk", ): @@ -573,34 +593,36 @@ def get_samples_lsq( ---------- psf_fitter : PSFFitter PSF fitter instance managing rho/tau statistics and covariances. - base : str - Precomputed basename for locating statistics/covariance files. + base_rho : str + Precomputed basename for locating rho statistics/covariance files. + base_tau : str + Precomputed basename for locating tau statistics/covariance files. apply_debias : int or None, optional Jackknife patch count applied during debiasing. Disabled when ``None``. cov_type : str, optional Covariance label (defaults to ``'jk'``). """ # Load rho and tau stats - psf_fitter.load_rho_stat(f"rho_stats_{base}.fits") - psf_fitter.load_tau_stat(f"tau_stats_{base}.fits") + psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") + psf_fitter.load_tau_stat(f"tau_stats_{base_tau}.fits") # Check if the path exists (use base for cache key to account for different TreeCorr configs) - sample_file_path = psf_fitter.get_sample_path(base) + sample_file_path = psf_fitter.get_sample_path(base_tau) if os.path.exists(sample_file_path): print(f"Skipping sampling; {sample_file_path} exists.") - flat_samples = psf_fitter.load_samples(base) + flat_samples = psf_fitter.load_samples(base_tau) mcmc_result, q = psf_fitter.get_mcmc_from_samples(flat_samples) print(mcmc_result) # Or run MCMC else: print("Least square sampling") - tau_covariance = f"cov_tau_{base}_{cov_type}.npy" - rho_covariance = f"cov_rho_{base}_jk.npy" + tau_covariance = f"cov_tau_{base_tau}_{cov_type}.npy" + rho_covariance = f"cov_rho_{base_rho}_jk.npy" psf_fitter.load_covariance(tau_covariance, cov_type="tau") psf_fitter.load_covariance(rho_covariance, cov_type="rho") debias_npatch = apply_debias if (apply_debias is not None) else None flat_samples, mcmc_result, q = psf_fitter.get_least_squares_params_samples( npatch=debias_npatch, apply_debias=(debias_npatch is not None) ) - psf_fitter.save_samples(flat_samples, base) + psf_fitter.save_samples(flat_samples, base_tau) return flat_samples, mcmc_result, q From d8b1507552049870e5195a28a646c7d74d9d03e7 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Wed, 22 Jul 2026 14:30:41 +0200 Subject: [PATCH 04/10] Update plot scripts for tau-statistics to handle tomography. Validate the code running and GLASS mocks. Upgrade the sampling of error parameters. --- cosmo_val/cat_config.yaml | 2 +- src/sp_validation/cosmo_val/pseudo_cl.py | 2 +- .../cosmo_val/psf_systematics.py | 328 ++++++++++++++++-- src/sp_validation/rho_tau.py | 50 ++- 4 files changed, 320 insertions(+), 62 deletions(-) diff --git a/cosmo_val/cat_config.yaml b/cosmo_val/cat_config.yaml index ea36a81d..5f0676bf 100644 --- a/cosmo_val/cat_config.yaml +++ b/cosmo_val/cat_config.yaml @@ -1280,7 +1280,7 @@ GLASS_mock_validation: star_flag: FLAG_STAR_HSM star_size: SIGMA_STAR_HSM hdu: 1 - path: unions_shapepipe_psf_2024_v1.6.a.fits + path: ../../../../n17data/UNIONS/WL/v1.6.x/unions_shapepipe_psf_2024_v1.6.a.fits ra_col: RA dec_col: Dec e1_PSF_col: E1_PSF_HSM diff --git a/src/sp_validation/cosmo_val/pseudo_cl.py b/src/sp_validation/cosmo_val/pseudo_cl.py index 13f747ad..29f87789 100644 --- a/src/sp_validation/cosmo_val/pseudo_cl.py +++ b/src/sp_validation/cosmo_val/pseudo_cl.py @@ -1206,7 +1206,7 @@ def get_ell_factor(ell): # Better jittering: symmetric around original ell values jitter_fraction = (j - (len(versions) - 1) / 2) * offset jittered_ell = ell + jitter_fraction * ell_widths - ell_factor_ = get_ell_factor(jittered_ell) + ell_factor_ = get_ell_factor(ell) for pol in pol_list: pol_color = self.get_pol_color(ver_color, pol, pol_list) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 30957774..c1df321e 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -6,8 +6,10 @@ """ import os +from pathlib import Path import matplotlib.pyplot as plt +import matplotlib.ticker as mticker import numpy as np import treecorr from astropy.io import fits @@ -84,9 +86,10 @@ def calculate_rho_tau_stats(self, tomography=True): ) else: - tomo_bin_ids, tomo_bin_pairs = "all", [("all", "all")] + tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] for tomo_bin_id in tomo_bin_ids: + self.print_cyan(f"Computing for the tomographic bin: {tomo_bin_id}") base_rho = self.basename(ver) base_tau = self.basename(ver, tomo_bin_a=tomo_bin_id) @@ -117,7 +120,7 @@ def calculate_rho_tau_stats(self, tomography=True): # --- utility functions --- def _get_galaxy_mask(self, ver, tomo_bin_id): - cat_gal = fits.getdata(self.cc["shear"]["path"]) + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) if tomo_bin_id != "all": gal_mask = cat_gal[self.cc[ver]["shear"]["tomo_bin_col"]] == tomo_bin_id else: @@ -159,19 +162,40 @@ def set_params_rho_tau(self, ver, params, params_psf): return params + def set_psf_parameter_sampling_method(self, rho_tau_method): + if rho_tau_method not in ["emcee", "lsq"]: + raise ValueError("Invalid PSF parameter sampling method.") + self.rho_tau_method = rho_tau_method + + def set_psf_parameter_nsamples(self, nsamples): + self.psf_error_nsamples = nsamples + + def set_psf_parameter_nwalkers(self, nwalkers): + self.psf_error_nwalkers = nwalkers + def get_samples(self, version, params, tomo_bin_id, track_result=False): npatch = params["patch_number"] if self.cov_estimate_method == "jk" else None base_rho = self.basename(version) base_tau = self.basename(version, tomo_bin_a=tomo_bin_id) + + # Set the number of samples and walkers and fallback to defaults if not attributed. + n_samples = ( + self.psf_error_nsamples if hasattr(self, "psf_error_nsamples") else 10_000 + ) + n_walkers = ( + self.psf_error_nwalkers if hasattr(self, "psf_error_nwalkers") else 124 + ) + flat_samples, result, q = get_samples( self.psf_fitter, - version, base_rho, base_tau, cov_type=self.cov_estimate_method, apply_debias=npatch, sampler=self.rho_tau_method, + nsamples=n_samples, + nwalkers=n_walkers, ) if track_result: @@ -188,45 +212,279 @@ def get_samples(self, version, params, tomo_bin_id, track_result=False): return xi_psf_sys_samples # --- plotting functions --- - def plot_rho_stats(self, abs=False): - filenames = [f"rho_stats_{self.basename(ver)}.fits" for ver in self.versions] + def plot_rho_stats( + self, + versions=None, + colors=None, + abs=False, + offset=0, + savefig=None, + show=True, + close=True, + ): + """ + Plot the Rho statistics. + + Parameters + ---------- + versions : list, optional + List of versions to plot. If None, all versions are plotted. + abs : bool, optional + If True, plot the absolute values of the Rho statistics. + offset : float, optional + Offset to apply to the versions for better visualisation. + savefig : str, optional + If provided, save the figure to this file. + show : bool, optional + If True, show the figure. + close : bool, optional + If True, close the figure after saving or showing. + """ + if versions is None: + versions = self.versions + + filenames = [f"rho_stats_{self.basename(ver)}.fits" for ver in versions] + + if colors is None: + colors = [self.cc[ver]["colour"] for ver in versions] + + if len(colors) != len(versions): + raise ValueError("Colors and versions must have the same length.") - savefig = "rho_stats.png" self.rho_stat_handler.plot_rho_stats( filenames, - self.colors, - self.versions, + colors, + versions, + offset=offset, savefig=savefig, legend="outside", abs=abs, - show=True, - close=True, + show=show, + close=close, ) - self.print_done( - "Rho stats plot saved to " - + f"{os.path.abspath(self.rho_stat_handler.catalogs._output)}/{savefig}", - ) + if savefig is not None: + self.print_done( + "Rho stats plot saved to " + + f"{os.path.abspath(self.rho_stat_handler.catalogs._output)}/{savefig}", + ) - def plot_tau_stats(self, plot_tau_m=False): - filenames = [f"tau_stats_{self.basename(ver)}.fits" for ver in self.versions] + def plot_tau_stats( + self, + tomography=False, + cov_type=None, + versions=None, + colors=None, + offset=0, + savefig=None, + show=True, + close=True, + plot_tau_m=False, + plot_theta_times_tau=False, + fmt="", + capsize=2, + ): + if versions is None: + versions = self.versions + + if colors is None: + colors = [self.cc[ver]["colour"] for ver in versions] + + if len(colors) != len(versions): + raise ValueError("Colors and versions must have the same length.") + + if cov_type is None: + self.print_cyan("Using the error bars from the tau-statistics files") + else: + self.print_cyan( + f"Using the error bars from the covariance files of type: {cov_type}" + ) - savefig = "tau_stats.png" - self.tau_stat_handler.plot_tau_stats( - filenames, - self.colors, - self.versions, - savefig=savefig, - legend="outside", - plot_tau_m=plot_tau_m, - show=True, - close=True, - ) + out_dir = f"{self.cc['paths']['output']}/rho_tau_stats" - self.print_done( - "Tau stats plot saved to " - + f"{os.path.abspath(self.tau_stat_handler.catalogs._output)}/{savefig}", - ) + if tomography: + # Write the whole script for the tomography. It does not exist in shear_psf_leakage + e_obs = r"e^\mathrm{obs}" + e_psf = r"e^\mathrm{PSF}" + delta_e_psf = r"\delta e^\mathrm{PSF}" + delta_T_psf = r"\delta T^\mathrm{PSF}" + + factor_theta_label = r"\theta" if plot_theta_times_tau else r"" + + titles = [ + rf"$\tau_0 = \langle {e_obs} {e_psf} \rangle$", + rf"${factor_theta_label} \tau_2 = {factor_theta_label} \langle {e_obs} {delta_e_psf} \rangle$", + rf"${factor_theta_label} \tau_5 = {factor_theta_label} \langle {e_obs} {delta_T_psf} \rangle$", + ] + + dict_index_tau = { + 0: "0", + 1: "2", + 2: "5", + } + + # From all the versions, get the maximum number of tomo_bin_ids + tomo_bins = {} + for ver in versions: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + + tomo_bins[ver] = {"ids": tomo_bin_ids, "pairs": tomo_bin_pairs} + + n_tomo_bins_plot = max(len(bins["ids"]) for bins in tomo_bins.values()) + + n_rows = n_tomo_bins_plot * (1 + plot_tau_m) + + fig = plt.figure(figsize=(20 * (1 + plot_tau_m), 10 * (1 + plot_tau_m))) + gs = fig.add_gridspec(n_rows, 3, wspace=0.1, hspace=0) + all_axs = gs.subplots(sharex="col") + + for k in range(n_rows): + axs = all_axs[k] + + tomo_bin_id = k // 2 + 1 if plot_tau_m else k + 1 + is_m_component = k % 2 if plot_tau_m else 0 + + for file_idx, (ver, color) in enumerate(zip(versions, colors)): + # Check if the tomo bin is valid for this version + if tomo_bin_id not in tomo_bins[ver]["ids"]: + continue + + # Load the tau-stats in the tau_stat_handler for easier read + base_tau = self.basename(ver, tomo_bin_a=tomo_bin_id) + + self.tau_stat_handler.load_tau_stats(f"tau_stats_{base_tau}.fits") + + if cov_type is not None: + cov_tau_path = ( + Path(out_dir) / f"cov_tau_{base_tau}_{cov_type}.npy" + ) + cov_tau = np.load(cov_tau_path) + + # Plot the different tau-stats per row + for i in range(3): + p_or_m = "m" if is_m_component else "p" + p_or_m_label = "-" if is_m_component else "+" + + # Get the jittered angular scale for the x-axis + theta = self.tau_stat_handler.tau_stats["theta"] + num_theta_bins = theta.shape[0] + + theta_widths = np.diff(theta) + theta_widths = np.append(theta_widths, theta_widths[-1]) + + jitter_fraction = (file_idx - (len(versions) - 1) / 2) * offset + jittered_theta = theta + jitter_fraction * theta_widths + + factor_theta = ( + np.ones_like(jittered_theta) + if (i == 0) or not plot_theta_times_tau + else theta + ) + + y_plot = ( + self.tau_stat_handler.tau_stats[ + f"tau_{dict_index_tau[i]}_{p_or_m}" + ] + * factor_theta + ) + + if cov_type is None or p_or_m == "m": + cov_diag = self.tau_stat_handler.tau_stats[ + "vartau_" + dict_index_tau[i] + "_" + p_or_m + ] + else: + cov_diag = np.diag( + cov_tau[ + i * num_theta_bins : (i + 1) * num_theta_bins, + i * num_theta_bins : (i + 1) * num_theta_bins, + ] + ) + + yerr_plot = np.sqrt(cov_diag) * factor_theta + + ver_label = ( + self.cc[ver]["label"] if "label" in self.cc[ver] else ver + ) + axs[i].errorbar( + jittered_theta, + y_plot, + yerr=yerr_plot, + fmt=fmt, + label=ver_label, + capsize=capsize, + color=color, + ) + + # Set the style of the plot + for i in range(3): + axs[i].set_xscale("log") + axs[i].set_xlim(theta.min() * 0.9, theta.max() * 1.1) + if i == 0: + axs[i].set_ylabel(f"Bin {tomo_bin_id}\n `{p_or_m_label}' comp.") + if k == n_rows - 1: + axs[i].set_xlabel(r"$\theta$ [arcmin]") + if k == 0: + axs[i].set_title(titles[i]) + + # --- Force scientific notation and scaling --- + axs[i].ticklabel_format(axis="y", style="sci", scilimits=(0, 0)) + axs[i].yaxis.offsetText.set_visible(True) + + # --- Make it look clean --- + axs[i].yaxis.set_major_locator( + mticker.MaxNLocator(nbins=5) + ) # Fewer, rounded ticks + # axs[i].yaxis.get_offset_text().set_fontsize(10) # Smaller ×10⁻⁴ label + axs[i].yaxis.get_offset_text().set_position( + (0, 1.02) + ) # Move scale factor slightly above axis + axs[i].yaxis.get_offset_text().set_ha("left") + + if k == n_rows - 1: + axs[1].legend( + loc="upper center", + bbox_to_anchor=(0.5, -0.5), # (x, y) relative to the axes + ncol=3, # number of columns + frameon=False, + ) + + plt.tight_layout() + + if savefig is not None: + plt.savefig(savefig, dpi=300, bbox_inches="tight") + + if show: + plt.show() + + if close: + plt.close() + + else: + filenames = [f"tau_stats_{self.basename(ver)}.fits" for ver in versions] + cov_paths = [ + f"cov_tau_{self.basename(ver)}_{cov_type}.npy" for ver in versions + ] + self.tau_stat_handler.plot_tau_stats( + filenames, + colors, + versions, + cov_paths=cov_paths, + offset=offset, + savefig=savefig, + legend="outside", + plot_tau_m=plot_tau_m, + plot_theta_times_tau=plot_theta_times_tau, + show=show, + close=close, + fmt=fmt, + capsize=capsize, + ) + + if savefig is not None: + self.print_done( + "Tau stats plot saved to " + + f"{os.path.abspath(self.tau_stat_handler.catalogs._output)}/{savefig}", + ) def calculate_rho_tau_fits(self, tomography=True): assert self.rho_tau_method != "none" @@ -251,11 +509,17 @@ def calculate_rho_tau_fits(self, tomography=True): ) else: - tomo_bin_ids, tomo_bin_pairs = "all", [("all", "all")] + tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] for tomo_bin_id in tomo_bin_ids: + self.print_cyan( + f"Sample PSF error parameters for tomographic bin {tomo_bin_id}" + ) xi_psf_sys_samples = self.get_samples(ver, params, tomo_bin_id) + if ver not in self._xi_psf_sys.keys(): + self._xi_psf_sys[ver] = {} + self._xi_psf_sys[ver][f"tomo_bin_{tomo_bin_id}"] = { "mean": np.mean(xi_psf_sys_samples, axis=0), "var": np.var(xi_psf_sys_samples, axis=0), diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 73ec4406..db4863c0 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -29,7 +29,6 @@ def get_params_rho_tau(cat): params["e2_star_col"] = cat["psf"]["e2_star_col"] params["PSF_size"] = cat["psf"]["PSF_size"] params["star_size"] = cat["psf"]["star_size"] - params["square_size"] = cat["psf"].get("square_size", False) params["PSF_flag"] = cat["psf"].get("PSF_flag") params["star_flag"] = cat["psf"].get("star_flag") params["ra_units"] = "deg" @@ -178,12 +177,9 @@ def get_rho_tau( if need_compute: rho_stat_handler.catalogs.set_params(params, outdir) - square_size = params["square_size"] - rho_stat_handler.build_cat_to_compute_rho( config[version]["psf"]["path"], catalog_id=catalog_id, - square_size=square_size, mask=mask_star, hdu=( config[version]["psf"]["hdu"] @@ -219,15 +215,12 @@ def get_rho_tau( else: tau_stat_handler.catalogs.set_params(params, outdir) - square_size = params["square_size"] - # Build the different catalogs if necessary if f"psf_{version}" not in tau_stat_handler.catalogs.catalogs_dict.keys(): tau_stat_handler.build_cat_to_compute_tau( config[version]["psf"]["path"], cat_type="psf", catalog_id=version, - square_size=square_size, mask=mask_star, hdu=( config[version]["psf"]["hdu"] @@ -241,7 +234,6 @@ def get_rho_tau( config[version]["shear"]["path"], cat_type="gal", catalog_id=version, - square_size=square_size, mask=mask_gal, ) @@ -370,8 +362,6 @@ def get_jackknife_cov( tau_stat_handler.catalogs.set_params(params, outdir) - square_size = params["square_size"] - for i in range(ncov): tau_chunk = outdir + f"/cov_tau_{base_tau}{i}.npy" rho_chunk = outdir + f"/cov_rho_{base_rho}{i}.npy" @@ -383,7 +373,6 @@ def get_jackknife_cov( rho_stat_handler.build_cat_to_compute_rho( config[version]["psf"]["path"], catalog_id=version + str(i), - square_size=square_size, mask=mask_star, hdu=config[version]["psf"]["hdu"], ) @@ -397,7 +386,6 @@ def get_jackknife_cov( config[version]["shear"]["path"], cat_type="gal", catalog_id=version + str(i), - square_size=square_size, mask=mask_gal, ) @@ -471,12 +459,13 @@ def get_jackknife_cov( def get_samples( psf_fitter, - version, base_rho, base_tau, cov_type="jk", apply_debias=None, sampler="emcee", + nsamples=10000, + nwalkers=124, ): """Return (alpha, beta, eta) samples using ``emcee`` or least squares. @@ -484,8 +473,6 @@ def get_samples( ---------- psf_fitter : PSFFitter PSF fitter instance that provides ``load_*`` helpers. - version : str - Catalog identifier whose rho/tau statistics are sampled. base_rho : str Precomputed basename (e.g. ``SP_v1.4_minsep=…``) used for rho-stat filenames. base_tau : str @@ -497,24 +484,29 @@ def get_samples( sampler : str, optional ``'emcee'`` for MCMC sampling, ``'lsq'`` for least squares (default ``'emcee'``). + nsamples : int, optional + Number of samples to draw (default ``10000``). + nwalkers : int, optional + Number of walkers for the MCMC run (default ``124``). """ if sampler == "emcee": return get_samples_emcee( psf_fitter, - version, base_rho, base_tau, cov_type=cov_type, apply_debias=apply_debias, + nsamples=nsamples, + nwalkers=nwalkers, ) elif sampler == "lsq": return get_samples_lsq( psf_fitter, - version, base_rho, base_tau, cov_type=cov_type, apply_debias=apply_debias, + nsamples=nsamples, ) else: raise ValueError("Sampler must be either 'emcee' or 'lsq'.") @@ -522,7 +514,6 @@ def get_samples( def get_samples_emcee( psf_fitter, - version, base_rho, base_tau, nwalkers=124, @@ -536,8 +527,6 @@ def get_samples_emcee( ---------- psf_fitter : PSFFitter PSF fitter instance managing rho/tau statistics and covariances. - version : str - Catalog identifier whose rho/tau statistics are sampled. base_rho : str Precomputed basename for locating rho statistics/covariance files. base_tau : str @@ -556,11 +545,12 @@ def get_samples_emcee( psf_fitter.load_tau_stat(f"tau_stats_{base_tau}.fits") # Check if the path exists (use base for cache key to account for different TreeCorr configs) - sample_file_path = psf_fitter.get_sample_file_path(base_tau) + base_sample = f"{base_tau}_sampler_emcee_cov_tau_type_{cov_type}" + sample_file_path = psf_fitter.get_sample_path(base_sample) if os.path.exists(sample_file_path): print(f"Skipping sampling; {sample_file_path} exists.") - flat_samples = psf_fitter.load_samples(base_tau) - mcmc_result, q = psf_fitter.get_mcmc_from_samples(base_tau) + flat_samples = psf_fitter.load_samples(base_sample) + mcmc_result, q = psf_fitter.get_mcmc_from_samples(flat_samples) print(mcmc_result) # Or run MCMC else: @@ -576,7 +566,7 @@ def get_samples_emcee( apply_debias=debias_npatch is not None, savefig="mcmc_samples_" + base_tau + ".png", ) - psf_fitter.save_samples(flat_samples, base_tau) + psf_fitter.save_samples(flat_samples, base_sample) return flat_samples, mcmc_result, q @@ -584,6 +574,7 @@ def get_samples_lsq( psf_fitter, base_rho, base_tau, + nsamples=10000, apply_debias=None, cov_type="jk", ): @@ -606,11 +597,12 @@ def get_samples_lsq( psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") psf_fitter.load_tau_stat(f"tau_stats_{base_tau}.fits") + base_sample = f"{base_tau}_sampler_lsq_cov_tau_type_{cov_type}" # Check if the path exists (use base for cache key to account for different TreeCorr configs) - sample_file_path = psf_fitter.get_sample_path(base_tau) + sample_file_path = psf_fitter.get_sample_path(base_sample) if os.path.exists(sample_file_path): print(f"Skipping sampling; {sample_file_path} exists.") - flat_samples = psf_fitter.load_samples(base_tau) + flat_samples = psf_fitter.load_samples(base_sample) mcmc_result, q = psf_fitter.get_mcmc_from_samples(flat_samples) print(mcmc_result) # Or run MCMC @@ -622,7 +614,9 @@ def get_samples_lsq( psf_fitter.load_covariance(rho_covariance, cov_type="rho") debias_npatch = apply_debias if (apply_debias is not None) else None flat_samples, mcmc_result, q = psf_fitter.get_least_squares_params_samples( - npatch=debias_npatch, apply_debias=(debias_npatch is not None) + npatch=debias_npatch, + apply_debias=(debias_npatch is not None), + n_samples=nsamples, ) - psf_fitter.save_samples(flat_samples, base_tau) + psf_fitter.save_samples(flat_samples, base_sample) return flat_samples, mcmc_result, q From 1e1a251cf550a8b7461ebefbeef4df7656ecc368 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Wed, 22 Jul 2026 14:46:39 +0200 Subject: [PATCH 05/10] Fix typo in tomo bins col name --- src/sp_validation/rho_tau.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index db4863c0..30d4e8a1 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -40,7 +40,7 @@ def get_params_rho_tau(cat): params["w_col"] = cat["shear"]["w_col"] params["e1_col"] = cat["shear"]["e1_col"] params["e2_col"] = cat["shear"]["e2_col"] - params["tomo_bin_id_col"] = cat["shear"].get("tomo_bin_id_col") + params["tomo_bin_col"] = cat["shear"].get("tomo_bin_col") params["R11"] = cat["shear"].get("R11") params["R22"] = cat["shear"].get("R22") From 813e4485345f82a65aa29e614c0059b0f95733fe Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Wed, 22 Jul 2026 14:59:10 +0200 Subject: [PATCH 06/10] Account for Copilot comments --- src/sp_validation/cosmo_val/psf_systematics.py | 10 +++++----- src/sp_validation/rho_tau.py | 18 +++++++++++++++--- 2 files changed, 20 insertions(+), 8 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index c1df321e..483271b1 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -31,7 +31,7 @@ class PSFSystematicsMixin: @property def rho_stat_handler(self): if not hasattr(self, "_rho_stat_handler"): - self.calculate_rho_tau_stats() + self.calculate_rho_tau_stats(tomography=False) if self.compute_tomography: self.calculate_rho_tau_stats(tomography=True) return self._rho_stat_handler @@ -39,7 +39,7 @@ def rho_stat_handler(self): @property def tau_stat_handler(self): if not hasattr(self, "_tau_stat_handler"): - self.calculate_rho_tau_stats() + self.calculate_rho_tau_stats(tomography=False) if self.compute_tomography: self.calculate_rho_tau_stats(tomography=True) return self._tau_stat_handler @@ -88,6 +88,9 @@ def calculate_rho_tau_stats(self, tomography=True): else: tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] + # Get the selection for stars + mask_star = self._get_star_mask(ver) + for tomo_bin_id in tomo_bin_ids: self.print_cyan(f"Computing for the tomographic bin: {tomo_bin_id}") base_rho = self.basename(ver) @@ -96,9 +99,6 @@ def calculate_rho_tau_stats(self, tomography=True): # Get the selection for galaxies mask_gal = self._get_galaxy_mask(ver, tomo_bin_id) - # Get the selection for stars - mask_star = self._get_star_mask(ver) - # Compute rho and tau statistics rho_stat_handler, tau_stat_handler = get_rho_tau_w_cov( self.cc, diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 30d4e8a1..ae73c0c0 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -91,8 +91,18 @@ def get_rho_tau_w_cov( ) return rho_stat_handler, tau_stat_handler elif method == "jk": + npatch = kwargs.get("npatch", 100) return get_jackknife_cov( - config, version, treecorr_config, outdir, base_rho, base_tau + config, + version, + treecorr_config, + outdir, + base_rho, + base_tau, + npatch=npatch, + ncov=ncov, + mask_star=mask_star, + mask_gal=mask_gal, ) elif method == "sim": tau_cov_path = Path(outdir) / f"cov_tau_{base_tau}_th.npy" @@ -144,8 +154,10 @@ def get_rho_tau( TreeCorr configuration (must include 'min_sep', 'max_sep', and 'nbins'). outdir : str Output directory. - base : str - Base name for output files. + base_rho : str + Base name for the rho output files. + base_tau : str + Base name for the tau output files. mask_star : array-like, optional Boolean mask for stars. If None, no masking is applied. mask_gal : array-like, optional From e26c37dd446e82603edc21f9fd0e0f0d18732fee Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Wed, 22 Jul 2026 15:01:27 +0200 Subject: [PATCH 07/10] Clean pseudo-cl test with the modified parameter initialisation --- src/sp_validation/tests/test_pseudo_cl.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/sp_validation/tests/test_pseudo_cl.py b/src/sp_validation/tests/test_pseudo_cl.py index c07afd00..eb4eb533 100644 --- a/src/sp_validation/tests/test_pseudo_cl.py +++ b/src/sp_validation/tests/test_pseudo_cl.py @@ -181,7 +181,7 @@ def cv(tmp_path): def cat_and_params(cv): """(catalog ndarray, params dict) for the synthetic version.""" ver = cv._test_version - params = get_params_rho_tau(cv.cc[ver], survey=ver) + params = get_params_rho_tau(cv.cc[ver]) cat_gal = fits.getdata(cv.cc[ver]["shear"]["path"]) return cat_gal, params @@ -734,7 +734,7 @@ def test_calculate_pseudo_cl_catalog_end_to_end(cv, tmp_path): # The end-to-end catalog EE matches the primitive get_pseudo_cls_catalog EE # (same computation, FITS round-trip) -- consistency, not an independent pin. cat_gal = fits.getdata(cv.cc[ver]["shear"]["path"]) - params = get_params_rho_tau(cv.cc[ver], survey=ver) + params = get_params_rho_tau(cv.cc[ver]) _, cl_prim, _ = cv.get_pseudo_cls_catalog( catalog=cat_gal, params=params, tomo_bin_a="all", tomo_bin_b="all" ) @@ -821,7 +821,7 @@ def test_calculate_pseudo_cl_catalog_end_to_end_tomo(cv, tmp_path): # The end-to-end catalog EE matches the primitive get_pseudo_cls_catalog EE # (same computation, FITS round-trip) -- consistency, not an independent pin. cat_gal = fits.getdata(cv.cc[ver]["shear"]["path"]) - params = get_params_rho_tau(cv.cc[ver], survey=ver) + params = get_params_rho_tau(cv.cc[ver]) _, cl_prim, _ = cv.get_pseudo_cls_catalog( catalog=cat_gal, params=params, tomo_bin_a=tomo_bin_a, tomo_bin_b=tomo_bin_b ) From 4559975f95bc9a7391b11cd3c67169369cd4c50c Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Wed, 22 Jul 2026 15:39:01 +0200 Subject: [PATCH 08/10] Reorganise functions in psf_systematics.py --- .../cosmo_val/psf_systematics.py | 384 +++++++++--------- 1 file changed, 192 insertions(+), 192 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 483271b1..79347029 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -118,6 +118,139 @@ def calculate_rho_tau_stats(self, tomography=True): self._rho_stat_handler = rho_stat_handler self._tau_stat_handler = tau_stat_handler + def calculate_rho_tau_fits(self, tomography=True): + assert self.rho_tau_method != "none" + + # this initializes the rho_tau_fits attribute + self._rho_tau_fits = {"flat_sample_list": [], "result_list": [], "q_list": []} + quantiles = [1 - self.quantile, self.quantile] + + self._xi_psf_sys = {} + for ver in self.versions: + params = self.set_params_rho_tau( + ver, self.results[ver]._params, self.cc[ver]["psf"] + ) + + # Get the tomographic bins + if tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise ValueError( + f"Version {ver} does not have tomography information." + ) + + else: + tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] + + for tomo_bin_id in tomo_bin_ids: + self.print_cyan( + f"Sample PSF error parameters for tomographic bin {tomo_bin_id}" + ) + xi_psf_sys_samples = self.get_samples(ver, params, tomo_bin_id) + + if ver not in self._xi_psf_sys.keys(): + self._xi_psf_sys[ver] = {} + + self._xi_psf_sys[ver][f"tomo_bin_{tomo_bin_id}"] = { + "mean": np.mean(xi_psf_sys_samples, axis=0), + "var": np.var(xi_psf_sys_samples, axis=0), + "quantiles": np.quantile(xi_psf_sys_samples, quantiles, axis=0), + } + + def calculate_scale_dependent_leakage(self): + self.print_start("Calculating scale-dependent leakage:") + for ver in self.versions: + self.print_magenta(ver) + results = self.results[ver] + + output_base_path = self._output_path(f"leakage_{ver}/xi_for_leak_scale") + output_path_ab = f"{output_base_path}_a_b.txt" + output_path_aa = f"{output_base_path}_a_a.txt" + with self.results[ver].temporarily_read_data(): + if os.path.exists(output_path_ab) and os.path.exists(output_path_aa): + self.print_green( + f"Skipping computation, reading {output_path_ab} and " + f"{output_path_aa} instead" + ) + + results.r_corr_gp = treecorr.GGCorrelation(self.treecorr_config) + results.r_corr_gp.read(output_path_ab) + + results.r_corr_pp = treecorr.GGCorrelation(self.treecorr_config) + results.r_corr_pp.read(output_path_aa) + + else: + results.compute_corr_gp_pp_alpha(output_base_path=output_base_path) + + results.do_alpha(fast=True) + results.do_xi_sys() + + self.print_done("Finished scale-dependent leakage calculation.") + + def calculate_objectwise_leakage(self): + if not hasattr(self.results[self.versions[0]], "alpha_leak_mean"): + self.calculate_scale_dependent_leakage() + + self.print_start("Object-wise leakage:") + mix = True + order = "lin" + for ver in self.versions: + self.print_magenta(ver) + + results_obj = self.results_objectwise[ver] + results_obj.check_params() + results_obj.update_params() + results_obj.prepare_output() + + # Skip read_data() and copy catalogue from scale leakage instance instead + # results_obj._dat = self.results[ver].dat_shear + + out_base = results_obj.get_out_base(mix, order) + out_path = f"{out_base}.pkl" + if os.path.exists(out_path): + self.print_green( + f"Skipping object-wise leakage, file {out_path} exists" + ) + results_obj.par_best_fit = leakage.read_from_file(out_path) + else: + self.print_cyan("Computing object-wise leakage regression") + + # Run + with results_obj.temporarily_read_data(): + try: + results_obj.PSF_leakage() + except KeyError as e: + print(f"{e}\nExpected key is missing from catalog.") + # remove the results object for this version + self.results_objectwise.pop(ver) + + # Gather coefficients + leakage_coeff = {} + for ver in self.results_objectwise: + results = self.results[ver] + par_best_fit = self.results_objectwise[ver].par_best_fit + + # Object-wise leakage + a11 = ufloat(par_best_fit["a11"].value, par_best_fit["a11"].stderr) + a22 = ufloat(par_best_fit["a22"].value, par_best_fit["a22"].stderr) + leakage_coeff[ver] = { + "a11": a11, + "a22": a22, + "aii_mean": 0.5 * (a11 + a22), + # Scale-dependent leakage: mean + "alpha_mean": ufloat(results.alpha_leak_mean, results.alpha_leak_std), + # Scale-dependent leakage: value at smallest scale + "alpha_1": ufloat(results.alpha_leak[0], results.sig_alpha_leak[0]), + # Scale-dependent leakage: value extrapolated to 0 using affine model + "alpha_0": ufloat( + results.alpha_affine_best_fit["c"].value, + results.alpha_affine_best_fit["c"].stderr, + ), + } + + self.leakage_coeff = leakage_coeff + # --- utility functions --- def _get_galaxy_mask(self, ver, tomo_bin_id): cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) @@ -162,6 +295,65 @@ def set_params_rho_tau(self, ver, params, params_psf): return params + def set_params_leakage_scale(self, ver): + params_in = {} + + # Set parameters + params_in["input_path_shear"] = self.cc[ver]["shear"]["path"] + params_in["input_path_PSF"] = self.cc[ver]["star"]["path"] + params_in["dndz_path"] = ( + f"{self.cc['nz']['dndz']['path']}_{self.cc[ver]['pipeline']}_{self.cc['nz']['dndz']['blind']}.txt" + ) + params_in["output_dir"] = f"{self.cc['paths']['output']}/leakage_{ver}" + + # Note: for SP these are calibrated shear estimates + params_in["e1_col"] = self.cc[ver]["shear"]["e1_col"] + params_in["e2_col"] = self.cc[ver]["shear"]["e2_col"] + params_in["w_col"] = self.cc[ver]["shear"]["w_col"] + params_in["R11"] = None if ver != "DES" else self.cc[ver]["shear"]["R11"] + params_in["R22"] = None if ver != "DES" else self.cc[ver]["shear"]["R22"] + + params_in["ra_star_col"] = self.cc[ver]["star"]["ra_col"] + params_in["dec_star_col"] = self.cc[ver]["star"]["dec_col"] + params_in["e1_PSF_star_col"] = self.cc[ver]["star"]["e1_col"] + params_in["e2_PSF_star_col"] = self.cc[ver]["star"]["e2_col"] + + params_in["theta_min_amin"] = self.theta_min + params_in["theta_max_amin"] = self.theta_max + params_in["n_theta"] = self.nbins + + params_in["verbose"] = False + + return params_in + + def set_params_leakage_object(self, ver): + params_in = {} + + # Set parameters + params_in["input_path_shear"] = self.cc[ver]["shear"]["path"] + params_in["output_dir"] = f"{self.cc['paths']['output']}/leakage_{ver}" + + # Note: for SP these are calibrated shear estimates + params_in["e1_col"] = self.cc[ver]["shear"]["e1_col"] + params_in["e2_col"] = self.cc[ver]["shear"]["e2_col"] + params_in["w_col"] = self.cc[ver]["shear"]["w_col"] + + if ( + "e1_PSF_col" in self.cc[ver]["shear"] + and "e2_PSF_col" in self.cc[ver]["shear"] + ): + params_in["e1_PSF_col"] = self.cc[ver]["shear"]["e1_PSF_col"] + params_in["e2_PSF_col"] = self.cc[ver]["shear"]["e2_PSF_col"] + else: + raise KeyError( + "Keys 'e1_PSF_col' and 'e2_PSF_col' not found in" + + f" shear yaml entry for version {ver}" + ) + + params_in["verbose"] = False + + return params_in + def set_psf_parameter_sampling_method(self, rho_tau_method): if rho_tau_method not in ["emcee", "lsq"]: raise ValueError("Invalid PSF parameter sampling method.") @@ -486,46 +678,6 @@ def plot_tau_stats( + f"{os.path.abspath(self.tau_stat_handler.catalogs._output)}/{savefig}", ) - def calculate_rho_tau_fits(self, tomography=True): - assert self.rho_tau_method != "none" - - # this initializes the rho_tau_fits attribute - self._rho_tau_fits = {"flat_sample_list": [], "result_list": [], "q_list": []} - quantiles = [1 - self.quantile, self.quantile] - - self._xi_psf_sys = {} - for ver in self.versions: - params = self.set_params_rho_tau( - ver, self.results[ver]._params, self.cc[ver]["psf"] - ) - - # Get the tomographic bins - if tomography: - tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - - if tomo_bin_ids is None or tomo_bin_pairs is None: - raise ValueError( - f"Version {ver} does not have tomography information." - ) - - else: - tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] - - for tomo_bin_id in tomo_bin_ids: - self.print_cyan( - f"Sample PSF error parameters for tomographic bin {tomo_bin_id}" - ) - xi_psf_sys_samples = self.get_samples(ver, params, tomo_bin_id) - - if ver not in self._xi_psf_sys.keys(): - self._xi_psf_sys[ver] = {} - - self._xi_psf_sys[ver][f"tomo_bin_{tomo_bin_id}"] = { - "mean": np.mean(xi_psf_sys_samples, axis=0), - "var": np.var(xi_psf_sys_samples, axis=0), - "quantiles": np.quantile(xi_psf_sys_samples, quantiles, axis=0), - } - def plot_rho_tau_fits(self): out_dir = self.rho_stat_handler.catalogs._output @@ -613,95 +765,6 @@ def plot_rho_tau_fits(self): f"{yscale}-scale xi_psf_sys terms plot saved to {out_path}" ) - def set_params_leakage_scale(self, ver): - params_in = {} - - # Set parameters - params_in["input_path_shear"] = self.cc[ver]["shear"]["path"] - params_in["input_path_PSF"] = self.cc[ver]["star"]["path"] - params_in["dndz_path"] = ( - f"{self.cc['nz']['dndz']['path']}_{self.cc[ver]['pipeline']}_{self.cc['nz']['dndz']['blind']}.txt" - ) - params_in["output_dir"] = f"{self.cc['paths']['output']}/leakage_{ver}" - - # Note: for SP these are calibrated shear estimates - params_in["e1_col"] = self.cc[ver]["shear"]["e1_col"] - params_in["e2_col"] = self.cc[ver]["shear"]["e2_col"] - params_in["w_col"] = self.cc[ver]["shear"]["w_col"] - params_in["R11"] = None if ver != "DES" else self.cc[ver]["shear"]["R11"] - params_in["R22"] = None if ver != "DES" else self.cc[ver]["shear"]["R22"] - - params_in["ra_star_col"] = self.cc[ver]["star"]["ra_col"] - params_in["dec_star_col"] = self.cc[ver]["star"]["dec_col"] - params_in["e1_PSF_star_col"] = self.cc[ver]["star"]["e1_col"] - params_in["e2_PSF_star_col"] = self.cc[ver]["star"]["e2_col"] - - params_in["theta_min_amin"] = self.theta_min - params_in["theta_max_amin"] = self.theta_max - params_in["n_theta"] = self.nbins - - params_in["verbose"] = False - - return params_in - - def set_params_leakage_object(self, ver): - params_in = {} - - # Set parameters - params_in["input_path_shear"] = self.cc[ver]["shear"]["path"] - params_in["output_dir"] = f"{self.cc['paths']['output']}/leakage_{ver}" - - # Note: for SP these are calibrated shear estimates - params_in["e1_col"] = self.cc[ver]["shear"]["e1_col"] - params_in["e2_col"] = self.cc[ver]["shear"]["e2_col"] - params_in["w_col"] = self.cc[ver]["shear"]["w_col"] - - if ( - "e1_PSF_col" in self.cc[ver]["shear"] - and "e2_PSF_col" in self.cc[ver]["shear"] - ): - params_in["e1_PSF_col"] = self.cc[ver]["shear"]["e1_PSF_col"] - params_in["e2_PSF_col"] = self.cc[ver]["shear"]["e2_PSF_col"] - else: - raise KeyError( - "Keys 'e1_PSF_col' and 'e2_PSF_col' not found in" - + f" shear yaml entry for version {ver}" - ) - - params_in["verbose"] = False - - return params_in - - def calculate_scale_dependent_leakage(self): - self.print_start("Calculating scale-dependent leakage:") - for ver in self.versions: - self.print_magenta(ver) - results = self.results[ver] - - output_base_path = self._output_path(f"leakage_{ver}/xi_for_leak_scale") - output_path_ab = f"{output_base_path}_a_b.txt" - output_path_aa = f"{output_base_path}_a_a.txt" - with self.results[ver].temporarily_read_data(): - if os.path.exists(output_path_ab) and os.path.exists(output_path_aa): - self.print_green( - f"Skipping computation, reading {output_path_ab} and " - f"{output_path_aa} instead" - ) - - results.r_corr_gp = treecorr.GGCorrelation(self.treecorr_config) - results.r_corr_gp.read(output_path_ab) - - results.r_corr_pp = treecorr.GGCorrelation(self.treecorr_config) - results.r_corr_pp.read(output_path_aa) - - else: - results.compute_corr_gp_pp_alpha(output_base_path=output_base_path) - - results.do_alpha(fast=True) - results.do_xi_sys() - - self.print_done("Finished scale-dependent leakage calculation.") - def plot_scale_dependent_leakage(self): if not hasattr(self.results[self.versions[0]], "r_corr_gp"): self.calculate_scale_dependent_leakage() @@ -847,69 +910,6 @@ def plot_scale_dependent_leakage(self): cs_plots.show() self.print_done(f"xi_sys_minus plot saved to {out_path}") - def calculate_objectwise_leakage(self): - if not hasattr(self.results[self.versions[0]], "alpha_leak_mean"): - self.calculate_scale_dependent_leakage() - - self.print_start("Object-wise leakage:") - mix = True - order = "lin" - for ver in self.versions: - self.print_magenta(ver) - - results_obj = self.results_objectwise[ver] - results_obj.check_params() - results_obj.update_params() - results_obj.prepare_output() - - # Skip read_data() and copy catalogue from scale leakage instance instead - # results_obj._dat = self.results[ver].dat_shear - - out_base = results_obj.get_out_base(mix, order) - out_path = f"{out_base}.pkl" - if os.path.exists(out_path): - self.print_green( - f"Skipping object-wise leakage, file {out_path} exists" - ) - results_obj.par_best_fit = leakage.read_from_file(out_path) - else: - self.print_cyan("Computing object-wise leakage regression") - - # Run - with results_obj.temporarily_read_data(): - try: - results_obj.PSF_leakage() - except KeyError as e: - print(f"{e}\nExpected key is missing from catalog.") - # remove the results object for this version - self.results_objectwise.pop(ver) - - # Gather coefficients - leakage_coeff = {} - for ver in self.results_objectwise: - results = self.results[ver] - par_best_fit = self.results_objectwise[ver].par_best_fit - - # Object-wise leakage - a11 = ufloat(par_best_fit["a11"].value, par_best_fit["a11"].stderr) - a22 = ufloat(par_best_fit["a22"].value, par_best_fit["a22"].stderr) - leakage_coeff[ver] = { - "a11": a11, - "a22": a22, - "aii_mean": 0.5 * (a11 + a22), - # Scale-dependent leakage: mean - "alpha_mean": ufloat(results.alpha_leak_mean, results.alpha_leak_std), - # Scale-dependent leakage: value at smallest scale - "alpha_1": ufloat(results.alpha_leak[0], results.sig_alpha_leak[0]), - # Scale-dependent leakage: value extrapolated to 0 using affine model - "alpha_0": ufloat( - results.alpha_affine_best_fit["c"].value, - results.alpha_affine_best_fit["c"].stderr, - ), - } - - self.leakage_coeff = leakage_coeff - def plot_objectwise_leakage(self): if not hasattr(self, "leakage_coeff"): self.calculate_objectwise_leakage() From 197d54cb93b6f5e58a5c58c2b43bd3952cdc1242 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Wed, 22 Jul 2026 16:52:22 +0200 Subject: [PATCH 09/10] Update the rho_tau_fits script to compute the tomographic PSF additive bias. --- .../cosmo_val/psf_systematics.py | 75 +++++++++++++++---- 1 file changed, 62 insertions(+), 13 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 79347029..2b60716e 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -118,11 +118,12 @@ def calculate_rho_tau_stats(self, tomography=True): self._rho_stat_handler = rho_stat_handler self._tau_stat_handler = tau_stat_handler - def calculate_rho_tau_fits(self, tomography=True): + def calculate_rho_tau_fits(self, tomography=True, track_result=True): assert self.rho_tau_method != "none" # this initializes the rho_tau_fits attribute - self._rho_tau_fits = {"flat_sample_list": [], "result_list": [], "q_list": []} + self._rho_tau_fits = {} + version_init_dict = {"flat_sample_list": {}, "result_list": {}, "q_list": {}} quantiles = [1 - self.quantile, self.quantile] self._xi_psf_sys = {} @@ -143,22 +144,41 @@ def calculate_rho_tau_fits(self, tomography=True): else: tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] + # Get the samples for each tomographic bin for tomo_bin_id in tomo_bin_ids: self.print_cyan( f"Sample PSF error parameters for tomographic bin {tomo_bin_id}" ) - xi_psf_sys_samples = self.get_samples(ver, params, tomo_bin_id) + _ = self.get_samples( + ver, params, tomo_bin_id, track_result=track_result + ) + + # Get the xi_psf_sys for each tomographic bin pairs + for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: + xi_psf_sys_samples_plus, xi_psf_sys_samples_minus = ( + self.get_xi_psf_sys_samples(ver, params, tomo_bin_a, tomo_bin_b) + ) if ver not in self._xi_psf_sys.keys(): self._xi_psf_sys[ver] = {} - self._xi_psf_sys[ver][f"tomo_bin_{tomo_bin_id}"] = { - "mean": np.mean(xi_psf_sys_samples, axis=0), - "var": np.var(xi_psf_sys_samples, axis=0), - "quantiles": np.quantile(xi_psf_sys_samples, quantiles, axis=0), + self._xi_psf_sys[ver][ + f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}" + ] = { + "mean_plus": np.mean(xi_psf_sys_samples_plus, axis=0), + "var_plus": np.var(xi_psf_sys_samples_plus, axis=0), + "quantiles_plus": np.quantile( + xi_psf_sys_samples_plus, quantiles, axis=0 + ), + "mean_minus": np.mean(xi_psf_sys_samples_minus, axis=0), + "var_minus": np.var(xi_psf_sys_samples_minus, axis=0), + "quantiles_minus": np.quantile( + xi_psf_sys_samples_minus, quantiles, axis=0 + ), } def calculate_scale_dependent_leakage(self): + # TODO: Upgrade for tomography self.print_start("Calculating scale-dependent leakage:") for ver in self.versions: self.print_magenta(ver) @@ -189,6 +209,7 @@ def calculate_scale_dependent_leakage(self): self.print_done("Finished scale-dependent leakage calculation.") def calculate_objectwise_leakage(self): + # TODO: Upgrade for tomography if not hasattr(self.results[self.versions[0]], "alpha_leak_mean"): self.calculate_scale_dependent_leakage() @@ -391,17 +412,45 @@ def get_samples(self, version, params, tomo_bin_id, track_result=False): ) if track_result: - self.rho_tau_fits["flat_sample_list"].append(flat_samples) - self.rho_tau_fits["result_list"].append(result) - self.rho_tau_fits["q_list"].append(q) + if not hasattr(self.rho_tau_fits, version): + self.rho_tau_fits[version] = { + "flat_sample_list": {}, + "result_list": {}, + "q_list": {}, + } + self.rho_tau_fits[ver]["flat_sample_list"][f"tomo_bin_{tomo_bin_id}"] = ( + flat_samples + ) + self.rho_tau_fits[ver]["result_list"][f"tomo_bin_{tomo_bin_id}"] = result + self.rho_tau_fits[ver]["q_list"][f"tomo_bin_{tomo_bin_id}"] = q + + return flat_samples + def get_xi_psf_sys_samples(self, ver, params, tomo_bin_a, tomo_bin_b): + base_rho = self.basename(ver) self.psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") nbins = self.psf_fitter.rho_stat_handler._treecorr_config["nbins"] - xi_psf_sys_samples = np.array( - [self.psf_fitter.compute_xi_psf_sys(sample) for sample in flat_samples] + + # Get the samples for the given tomographic bins + flat_samples_a = self.get_samples(ver, params, tomo_bin_a, track_result=False) + flat_samples_b = self.get_samples(ver, params, tomo_bin_b, track_result=False) + + # TODO: update the function compute_xi_psf_sys to handle two tomographic bins. For now, we will just use the samples from the first tomographic bin. + xi_psf_sys_samples_plus = np.array( + [ + self.psf_fitter.compute_xi_psf_sys(sample_a, sample_b, p_or_m="p") + for (sample_a, sample_b) in zip(flat_samples_a, flat_samples_b) + ] + ).reshape(-1, nbins) + + xi_psf_sys_samples_minus = np.array( + [ + self.psf_fitter.compute_xi_psf_sys(sample_a, sample_b, p_or_m="m") + for (sample_a, sample_b) in zip(flat_samples_a, flat_samples_b) + ] ).reshape(-1, nbins) - return xi_psf_sys_samples + return xi_psf_sys_samples_plus, xi_psf_sys_samples_minus # --- plotting functions --- def plot_rho_stats( From 038f507c15d4c5402b8abeab5c2747367de2ff2b Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Wed, 22 Jul 2026 16:53:22 +0200 Subject: [PATCH 10/10] Update the rho_tau_fits script to compute the tomographic PSF additive bias. --- src/sp_validation/cosmo_val/psf_systematics.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 2b60716e..b3c0be61 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -123,7 +123,6 @@ def calculate_rho_tau_fits(self, tomography=True, track_result=True): # this initializes the rho_tau_fits attribute self._rho_tau_fits = {} - version_init_dict = {"flat_sample_list": {}, "result_list": {}, "q_list": {}} quantiles = [1 - self.quantile, self.quantile] self._xi_psf_sys = {} @@ -418,11 +417,13 @@ def get_samples(self, version, params, tomo_bin_id, track_result=False): "result_list": {}, "q_list": {}, } - self.rho_tau_fits[ver]["flat_sample_list"][f"tomo_bin_{tomo_bin_id}"] = ( - flat_samples + self.rho_tau_fits[version]["flat_sample_list"][ + f"tomo_bin_{tomo_bin_id}" + ] = flat_samples + self.rho_tau_fits[version]["result_list"][f"tomo_bin_{tomo_bin_id}"] = ( + result ) - self.rho_tau_fits[ver]["result_list"][f"tomo_bin_{tomo_bin_id}"] = result - self.rho_tau_fits[ver]["q_list"][f"tomo_bin_{tomo_bin_id}"] = q + self.rho_tau_fits[version]["q_list"][f"tomo_bin_{tomo_bin_id}"] = q return flat_samples