From 8ae222bfb63900098dabb479d1cab6bb0e4cd498 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Tue, 21 Jul 2026 10:18:40 +0200 Subject: [PATCH 01/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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 From aaf75dbc9b3469015dc5e0ece8c078b86e1b3461 Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Thu, 23 Jul 2026 18:52:53 +0200 Subject: [PATCH 11/12] Update the plotting scripts for xi_sys. --- src/sp_validation/cosmo_val/core.py | 28 + src/sp_validation/cosmo_val/pseudo_cl.py | 15 +- .../cosmo_val/psf_systematics.py | 647 +++++++++++++++--- 3 files changed, 598 insertions(+), 92 deletions(-) diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index af2a6701..9c041788 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -691,3 +691,31 @@ def _get_tomo_bins(self, version): else: self.print_cyan(f"Version {version} does not have tomography information.") return None, None + + def _get_tomo_bins_for_versions(self, versions, tomography): + """ + Return a dictionary of tomo_bin_ids and tomo_bin_pairs for each version in versions. + + Parameters + ---------- + versions : list of str + List of catalog version identifiers + tomography : bool + If True, assumes tomography else returns the format for non-tomographic versions. + + Returns + ------- + dict + Dictionary with version as key and a dictionary containing 'ids' and 'pairs' as values. + Example: {version1: {'ids': tomo_bin_ids1, 'pairs': tomo_bin_pairs1}, ...} + """ + tomo_bins = {} + for ver in versions: + if tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + else: + tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] + + tomo_bins[ver] = {"ids": tomo_bin_ids, "pairs": tomo_bin_pairs} + + return tomo_bins diff --git a/src/sp_validation/cosmo_val/pseudo_cl.py b/src/sp_validation/cosmo_val/pseudo_cl.py index 29f87789..72b7aa9a 100644 --- a/src/sp_validation/cosmo_val/pseudo_cl.py +++ b/src/sp_validation/cosmo_val/pseudo_cl.py @@ -1159,16 +1159,11 @@ def get_ell_factor(ell): fmt_dict = {"EE": "o", "BB": "s", "EB": "^", "BE": "v"} # From all the versions, get the maximum number of tomo_bin_ids - tomo_bins = {} - for ver in versions: - if tomography: - tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - else: - tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] - - tomo_bins[ver] = {"ids": tomo_bin_ids, "pairs": tomo_bin_pairs} + tomo_bins = self._get_tomo_bins_for_versions(versions, tomography=tomography) - n_tomo_bins_plot = max(len(bins["ids"]) for bins in tomo_bins.values()) + max_key = max(tomo_bins, key=lambda k: len(tomo_bins[k]["ids"])) + n_tomo_bins_plot = len(tomo_bins[max_key]["ids"]) + reference_tomo_bin_pairs = tomo_bins[max_key]["pairs"] fig, axs = plt.subplots( n_tomo_bins_plot, @@ -1231,7 +1226,7 @@ def get_ell_factor(ell): else: ell_label = r"$\ell(\ell+1) C_\ell$" - for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: + for tomo_bin_a, tomo_bin_b in reference_tomo_bin_pairs: if tomography: ax = axs[tomo_bin_b - 1, tomo_bin_a - 1] else: diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index b3c0be61..2c844e90 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -122,10 +122,12 @@ 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 = {} + if not hasattr(self, "_rho_tau_fits"): + self._rho_tau_fits = {} quantiles = [1 - self.quantile, self.quantile] - self._xi_psf_sys = {} + if not hasattr(self, "_xi_psf_sys"): + self._xi_psf_sys = {} for ver in self.versions: params = self.set_params_rho_tau( ver, self.results[ver]._params, self.cc[ver]["psf"] @@ -411,19 +413,17 @@ def get_samples(self, version, params, tomo_bin_id, track_result=False): ) if track_result: - if not hasattr(self.rho_tau_fits, version): + if version not in self.rho_tau_fits.keys(): self.rho_tau_fits[version] = { - "flat_sample_list": {}, - "result_list": {}, - "q_list": {}, + "flat_samples": {}, + "result": {}, + "quantile": {}, } - 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[version]["flat_samples"][f"tomo_bin_{tomo_bin_id}"] = ( + flat_samples ) - self.rho_tau_fits[version]["q_list"][f"tomo_bin_{tomo_bin_id}"] = q + self.rho_tau_fits[version]["result"][f"tomo_bin_{tomo_bin_id}"] = result + self.rho_tau_fits[version]["quantile"][f"tomo_bin_{tomo_bin_id}"] = q return flat_samples @@ -436,7 +436,6 @@ def get_xi_psf_sys_samples(self, ver, params, tomo_bin_a, tomo_bin_b): 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") @@ -566,11 +565,9 @@ def plot_tau_stats( } # 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} + tomo_bins = self._get_tomo_bins_for_versions( + versions, tomography=tomography + ) n_tomo_bins_plot = max(len(bins["ids"]) for bins in tomo_bins.values()) @@ -611,11 +608,9 @@ def plot_tau_stats( 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 + jittered_theta = self._get_jittered_theta( + theta, file_idx, len(versions), offset + ) factor_theta = ( np.ones_like(jittered_theta) @@ -728,80 +723,188 @@ def plot_tau_stats( + f"{os.path.abspath(self.tau_stat_handler.catalogs._output)}/{savefig}", ) - def plot_rho_tau_fits(self): + def plot_rho_tau_fits( + self, + tomography=False, + versions=None, + colors=None, + savefig_contours=None, + savefig_xi_psf_sys=None, + savefig_format="png", + tomo_bin_label_position=None, + nsamples_to_plot=100, + offset=0, + alpha=0.3, + times_theta=False, + show=True, + close=True, + ): + """ + Plot the Rho/Tau fits and the xi_psf_sys samples. + + Parameters + ---------- + versions : list, optional + List of versions to plot. If None, all versions are plotted. + colors : list, optional + List of colors for each version. If None, default colors are used. + savefig_contours : str, optional + If provided, save the contour plot to this file. + savefig_xi_psf_sys : str, optional + If provided, save the xi_psf_sys samples plot to this file. + nsamples_to_plot : int, optional + Number of xi_psf_sys samples to plot. Default is 100. + offset : float, optional + Offset to apply to the versions for better visualisation. + alpha : float, optional + Alpha value for the xi_psf_sys samples plot. Default is 0.3. + times_theta : bool, optional + If True, plot the xi_psf_sys multiplied by theta. + show : bool, optional + If True, show the figure. + close : bool, optional + If True, close the figure after saving or showing. + """ + if savefig_format not in ["png", "pdf", "jpg", "jpeg", "svg"]: + raise ValueError( + "Invalid savefig_format. Must be one of: 'png', 'pdf', 'jpg', 'jpeg', 'svg'." + ) + out_dir = self.rho_stat_handler.catalogs._output - savefig = f"{out_dir}/contours_tau_stat.png" + versions = versions if versions is not None else self.versions + colors = ( + colors + if colors is not None + else [self.cc[ver]["colour"] for ver in versions] + ) + + # Print the tau-statistics constraints on the PSF error model parameters. + savefig = ( + f"{out_dir}/{savefig_contours}" if savefig_contours is not None else None + ) + self.print_cyan( + "Only plot contours for the non-tomographic case. For tomography, the contours are not plotted." + ) + sample_list = [ + self.rho_tau_fits[ver]["flat_samples"]["tomo_bin_all"] for ver in versions + ] psfleak_plots.plot_contours( - self.rho_tau_fits["flat_sample_list"], + sample_list, names=["x0", "x1", "x2"], labels=[r"\alpha", r"\beta", r"\eta"], - savefig=savefig, - legend_labels=self.versions, + savefig=savefig_contours, + legend_labels=versions, legend_loc="upper right", - contour_colors=self.colors, + contour_colors=colors, markers={"x0": 0, "x1": 1, "x2": 1}, - show=True, - close=True, + show=show, + close=close, ) - self.print_done(f"Tau contours plot saved to {os.path.abspath(savefig)}") + if savefig_contours is not None: + self.print_done(f"Tau contours plot saved to {os.path.abspath(savefig)}") - plt.figure(figsize=(15, 6)) - for mcmc_result, ver, color, flat_sample in zip( - self.rho_tau_fits["result_list"], - self.versions, - self.colors, - self.rho_tau_fits["flat_sample_list"], - ): - self.psf_fitter.load_rho_stat(f"rho_stats_{self.basename(ver)}.fits") - for i in range(100): - self.psf_fitter.plot_xi_psf_sys( - flat_sample[-i + 1], ver, color, alpha=0.1 - ) - self.psf_fitter.plot_xi_psf_sys(mcmc_result[1], ver, color) - plt.legend() - out_path = os.path.abspath(f"{out_dir}/xi_psf_sys_samples.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_psf_sys samples plot saved to {out_path}") + # Plot the xi_psf_sys samples and quantiles. + if tomography: + self.print_cyan("Plot the xi_psf_sys for the tomographic case.") + else: + self.print_cyan("Plot the xi_psf_sys for the non-tomographic case.") - plt.figure(figsize=(15, 6)) - for mcmc_result, ver, color, flat_sample in zip( - self.rho_tau_fits["result_list"], - self.versions, - self.colors, - self.rho_tau_fits["flat_sample_list"], - ): - ls = self.cc[ver]["ls"] - theta = self.psf_fitter.rho_stat_handler.rho_stats["theta"] - xi_psf_sys = self.xi_psf_sys[ver] - plt.plot(theta, xi_psf_sys["mean"], linestyle=ls, color=color) - plt.plot(theta, xi_psf_sys["quantiles"][0], linestyle=ls, color=color) - plt.plot(theta, xi_psf_sys["quantiles"][1], linestyle=ls, color=color) - plt.fill_between( - theta, - xi_psf_sys["quantiles"][0], - xi_psf_sys["quantiles"][1], - color=color, - alpha=0.25, - label=ver, - ) + x_label = r"$\theta$ [arcmin]" + y_label_plus = ( + r"$\theta$" if times_theta else "" + ) + r"$\xi^{\rm PSF, sys}_+(\theta)$" + y_label_minus = ( + r"$\theta$" if times_theta else "" + ) + r"$\xi^{\rm PSF, sys}_-(\theta)$" - plt.xscale("log") - plt.yscale("log") - plt.xlabel(r"$\theta$ [arcmin]") - plt.ylabel(r"$\xi^{\rm PSF}_{\rm sys}$") - plt.title(f"{1 - self.quantile:.1%}, {self.quantile:.1%} quantiles") - plt.legend() - out_path = os.path.abspath(f"{out_dir}/xi_psf_sys_quantiles.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_psf_sys quantiles plot saved to {out_path}") + if tomo_bin_label_position is None: + tomo_bin_label_position = (0.05, 0.9) if times_theta else (0.8, 0.95) + + y_scale = "linear" if times_theta else "log" + + out_path = ( + f"{out_dir}/{savefig_xi_psf_sys}_tomography_{tomography}_sample.{savefig_format}" + if savefig_xi_psf_sys is not None + else None + ) + + kwargs_x_y_plot_function = { + "nsamples_to_plot": nsamples_to_plot, + "alpha": alpha, + "offset": offset, + "times_theta": times_theta, + } + + self.plot_2pcf_tomography( + self._xi_psf_sys_sample_x_y_plot_function, + x_label, + y_label_plus, + y_label_minus, + tomo_bin_label_position, + extract_text_offset=times_theta, + add_index_version_to_kwargs=True, + x_scale="log", + y_scale=y_scale, + tomography=tomography, + versions=versions, + colors=colors, + savefig=out_path, + show=show, + close=close, + **kwargs_x_y_plot_function, + ) + + # Plot the xi_psf_sys mean and std. + x_label = r"$\theta$ [arcmin]" + y_label_plus = ( + r"$\theta$" if times_theta else "" + ) + r"$\xi^{\rm PSF, sys}_+(\theta)$" + y_label_minus = ( + r"$\theta$" if times_theta else "" + ) + r"$\xi^{\rm PSF, sys}_-(\theta)$" + + if tomo_bin_label_position is None: + tomo_bin_label_position = (0.05, 0.9) if times_theta else (0.8, 0.95) + y_scale = "linear" if times_theta else "log" + + out_path = ( + f"{out_dir}/{savefig_xi_psf_sys}_tomography_{tomography}_mean_and_std.{savefig_format}" + if savefig_xi_psf_sys is not None + else None + ) + + kwargs_x_y_plot_function = { + "offset": offset, + "times_theta": times_theta, + "alpha": alpha, + } + + self.plot_2pcf_tomography( + self._xi_psf_sys_mean_and_std_x_y_plot_function, + x_label, + y_label_plus, + y_label_minus, + tomo_bin_label_position, + extract_text_offset=times_theta, + add_index_version_to_kwargs=True, + x_scale="log", + y_scale=y_scale, + tomography=tomography, + versions=versions, + colors=colors, + savefig=out_path, + show=show, + close=close, + **kwargs_x_y_plot_function, + ) + + """ for mcmc_result, ver, flat_sample in zip( self.rho_tau_fits["result_list"], self.versions, - self.rho_tau_fits["flat_sample_list"], + self.rho_tau_fits["flat_samples"], ): self.psf_fitter.load_rho_stat(f"rho_stats_{self.basename(ver)}.fits") for yscale in ("linear", "log"): @@ -814,6 +917,7 @@ def plot_rho_tau_fits(self): self.print_done( f"{yscale}-scale xi_psf_sys terms plot saved to {out_path}" ) + """ def plot_scale_dependent_leakage(self): if not hasattr(self.results[self.versions[0]], "r_corr_gp"): @@ -1006,3 +1110,382 @@ def plot_objectwise_leakage(self): cs_plots.savefig(out_path, close_fig=False) cs_plots.show() self.print_done(f"Object-wise leakage coefficients plot saved to {out_path}") + + # --- utlility functions for plotting --- + def plot_2pcf_tomography( + self, + x_y_plot_function, + x_label, + y_label_plus, + y_label_minus, + tomo_bin_label_position, + extract_text_offset, + add_index_version_to_kwargs, + x_scale=None, + y_scale=None, + tomography=False, + versions=None, + colors=None, + savefig=None, + show=True, + close=True, + **kwargs, + ): + """ + Standard plot function for 2-point correlation functions with tomographic bins. + + Parameters + ---------- + x_y_plot_function : callable + Function to plot the x and y data. Should accept axes for `+' and `-' components, version, tomo_bin_indices, and kwargs. + x_label : str + Label for the x-axis. + y_label_plus : str + Label for the y-axis of the `+' component. + y_label_minus : str + Label for the y-axis of the `-' component. + tomo_bin_label_position : tuple + Position to place the tomographic bin labels in axes coordinates (x, y). + extract_text_offset : bool + If True, extract the y-axis offset text and include it in the y-label. + add_index_version_to_kwargs : bool + If True, add the index of the version to the kwargs for plotting. + x_scale : str, optional + Scale for the x-axis ('linear', 'log', etc.). If None, default scale is used. + y_scale : str, optional + Scale for the y-axis ('linear', 'log', etc.). If None, default scale is used. + tomography : bool, optional + If True, plot the tomographic bins. Default is False. + versions : list, optional + List of versions to plot. If None, all versions are plotted. + colors : list, optional + List of colors for each version. If None, default colors are used. + 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. + kwargs : dict + Additional keyword arguments to pass to the plotting function. + """ + versions = versions if versions is not None else self.versions + colors = ( + colors + if colors is not None + else [self.cc[ver]["colour"] for ver in versions] + ) + + # First get the max number of tomo bins among the versions. + tomo_bins = self._get_tomo_bins_for_versions(versions, tomography=tomography) + + max_key = max(tomo_bins, key=lambda k: len(tomo_bins[k]["ids"])) + n_tomo_bins_plot = len(tomo_bins[max_key]["ids"]) + reference_tomo_bin_pairs = tomo_bins[max_key]["pairs"] + + n_rows = n_tomo_bins_plot + n_cols = n_tomo_bins_plot + 2 + + # First start with the quantiles plot + fig, axs = plt.subplots( + n_rows, + n_cols, + figsize=(5 * n_cols, 5 * n_rows), + sharex=True, + sharey=True, + gridspec_kw={"wspace": 0, "hspace": 0}, + ) + + for idx, (ver, color) in enumerate(zip(versions, colors)): + tomo_bin_pairs = tomo_bins[ver]["pairs"] + + kwargs["color"] = color + if add_index_version_to_kwargs: + kwargs["idx"] = idx + kwargs["versions"] = versions + + for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: + # Plot the nsamples last samples + ax_plus = self._get_ax_plus(axs, tomo_bin_a, tomo_bin_b) + ax_minus = self._get_ax_minus(axs, tomo_bin_a, tomo_bin_b) + + # Apply the x_y_plot_function to plot the data + x_y_plot_function( + ax_plus, ax_minus, ver, tomo_bin_a, tomo_bin_b, **kwargs + ) + + # Draw to extract the y-axis text offset + fig.canvas.draw() + + # Set the visibility to false where necessary + self._set_ax_visibility_to_false(axs, n_tomo_bins_plot) + + # Set the labels and scales for the plots + for tomo_bin_a, tomo_bin_b in reference_tomo_bin_pairs: + ax_plus = self._get_ax_plus(axs, tomo_bin_a, tomo_bin_b) + ax_minus = self._get_ax_minus(axs, tomo_bin_a, tomo_bin_b) + + ax_plus.tick_params( + axis="both", + which="both", + direction="in", + bottom=True, + top=False, + labelbottom=tomo_bin_b == 1 or tomo_bin_b == "all", + left=True, + right=False, + labelleft=tomo_bin_a == 1 or tomo_bin_a == "all", + ) + if x_scale is not None: + ax_plus.set_xscale(x_scale) + + ax_plus.text( + tomo_bin_label_position[0], + tomo_bin_label_position[1], + f"{tomo_bin_a}-{tomo_bin_b}", + transform=ax_plus.transAxes, + verticalalignment="top", + bbox=dict( + boxstyle="square", + facecolor="white", + edgecolor="black", + alpha=0.8, + ), + ) + ax_plus.axhline(0, color="k", linestyle="--") + + if y_scale is not None: + ax_plus.set_yscale(y_scale) + if tomo_bin_b == 1 or tomo_bin_b == "all": + ax_plus.set_xlabel(r"$\theta$ [arcmin]") + if tomo_bin_a == 1 or tomo_bin_a == "all": + text_offset = ( + ax_plus.yaxis.get_offset_text().get_text() + if extract_text_offset + else "" + ) + ax_plus.set_ylabel(y_label_plus + text_offset) + ax_plus.yaxis.get_offset_text().set_visible( + False + ) # Hide the offset text for the plus ax + + # Move the ticks to the right for the minus ax + ax_minus.yaxis.tick_right() + ax_minus.yaxis.set_label_position("right") + ax_minus.tick_params( + axis="both", + which="both", + direction="in", + bottom=True, + top=False, + labelbottom=tomo_bin_b == n_tomo_bins_plot or tomo_bin_b == "all", + left=False, + right=True, + labelleft=False, + labelright=tomo_bin_a == 1 or tomo_bin_a == "all", + ) + if x_scale is not None: + ax_minus.set_xscale(x_scale) + ax_minus.text( + tomo_bin_label_position[0], + tomo_bin_label_position[1], + f"{tomo_bin_a}-{tomo_bin_b}", + transform=ax_minus.transAxes, + verticalalignment="top", + bbox=dict( + boxstyle="square", + facecolor="white", + edgecolor="black", + alpha=0.8, + ), + ) + ax_minus.axhline(0, color="k", linestyle="--") + if y_scale is not None: + ax_minus.set_yscale(y_scale) + if tomo_bin_b == n_tomo_bins_plot or tomo_bin_b == "all": + ax_minus.set_xlabel(x_label) + if tomo_bin_a == 1 or tomo_bin_a == "all": + text_offset = ( + ax_minus.yaxis.get_offset_text().get_text() + if extract_text_offset + else "" + ) + ax_minus.set_ylabel(y_label_minus + text_offset) + ax_minus.yaxis.get_offset_text().set_visible( + False + ) # Hide the offset text for the minus ax + + # Build the legend + handles = [] + for ver, color in zip(versions, colors): + label = self.cc[ver]["label"] if "label" in self.cc[ver] else ver + handles.append(plt.Line2D([0], [0], color=color, lw=2, label=label)) + fig.legend( + handles=handles, + loc="upper center", + ncol=3, + frameon=False, + bbox_to_anchor=(0.5, 0.0), + ) + + if savefig is not None: + plt.savefig(savefig, dpi=300, bbox_inches="tight") + self.print_done(f"Plot saved to {os.path.abspath(savefig)}") + + if show: + plt.show() + + if close: + plt.close() + + def _xi_psf_sys_sample_x_y_plot_function( + self, + ax_plus, + ax_minus, + version, + tomo_bin_a, + tomo_bin_b, + idx, + versions, + color, + offset, + nsamples_to_plot, + times_theta, + alpha, + ): + # Load the rho-stats to compute the xi_psf_sys samples + base_rho = self.basename(version) + self.psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") + + # Get the angular scales for the xi_psf_sys plots + theta = self.psf_fitter.rho_stat_handler.rho_stats["theta"] + + # Add the offset to the theta values for better visualisation + jittered_theta = self._get_jittered_theta(theta, idx, len(versions), offset) + + # Get the parameters for the rho-tau fit + params = self.set_params_rho_tau( + version, self.results[version]._params, self.cc[version]["psf"] + ) + + # Get the xi_psf_sys samples + xi_psf_sys_samples_plus, xi_psf_sys_samples_minus = self.get_xi_psf_sys_samples( + version, params, tomo_bin_a, tomo_bin_b + ) + + y_plus = xi_psf_sys_samples_plus[-nsamples_to_plot:] * ( + theta if times_theta else 1 + ) + y_minus = xi_psf_sys_samples_minus[-nsamples_to_plot:] * ( + theta if times_theta else 1 + ) + + ax_plus.plot( + jittered_theta, + y_plus.T, + color=color, + alpha=alpha, + ) + + ax_minus.plot(jittered_theta, y_minus.T, color=color, alpha=alpha) + + def _xi_psf_sys_mean_and_std_x_y_plot_function( + self, + ax_plus, + ax_minus, + version, + tomo_bin_a, + tomo_bin_b, + idx, + versions, + color, + offset, + times_theta, + alpha, + ): + # Load the rho-stats to compute the xi_psf_sys mean and std. + base_rho = self.basename(version) + self.psf_fitter.load_rho_stat(f"rho_stats_{base_rho}.fits") + + # Get the angular scales for the xi_psf_sys plots + theta = self.psf_fitter.rho_stat_handler.rho_stats["theta"] + + # Add the offset to the theta values for better visualisation + jittered_theta = self._get_jittered_theta(theta, idx, len(versions), offset) + + xi_psf_sys = self.xi_psf_sys[version][ + f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}" + ] + + def plot_axis(ax, ax_type): + """Plot depending on the axis type ('plus' or 'minus').""" + ax.plot( + jittered_theta, + xi_psf_sys[f"mean_{ax_type}"] * (theta if times_theta else 1), + color=color, + ) + ax.plot( + jittered_theta, + xi_psf_sys[f"quantiles_{ax_type}"][0] * (theta if times_theta else 1), + color=color, + ) + ax.plot( + jittered_theta, + xi_psf_sys[f"quantiles_{ax_type}"][1] * (theta if times_theta else 1), + color=color, + ) + ax.fill_between( + jittered_theta, + xi_psf_sys[f"quantiles_{ax_type}"][0] * (theta if times_theta else 1), + xi_psf_sys[f"quantiles_{ax_type}"][1] * (theta if times_theta else 1), + color=color, + alpha=alpha, + ) + + # Plot the plus axis + plot_axis(ax_plus, "plus") + + # Plot the minus axis + plot_axis(ax_minus, "minus") + + def _get_jittered_theta(self, theta, idx, n_versions, offset): + """Get the jittered theta values for better visualisation.""" + theta_widths = np.diff(theta) + theta_widths = np.append(theta_widths, theta_widths[-1]) + jitter_fraction = (idx - (n_versions - 1) / 2) * offset + jittered_theta = theta + jitter_fraction * theta_widths + return jittered_theta + + def _get_ax_plus(self, axs, tomo_bin_a, tomo_bin_b): + if (tomo_bin_a == "all") ^ (tomo_bin_b == "all"): + raise ValueError( + "Invalid combination of tomographic bins: 'all' and a specific bin." + ) + + if tomo_bin_a == "all" and tomo_bin_b == "all": + return axs[0] + + else: + nrows = axs.shape[0] + return axs[nrows - tomo_bin_b, tomo_bin_a - 1] + + def _get_ax_minus(self, axs, tomo_bin_a, tomo_bin_b): + if (tomo_bin_a == "all") ^ (tomo_bin_b == "all"): + raise ValueError( + "Invalid combination of tomographic bins: 'all' and a specific bin." + ) + + if tomo_bin_a == "all" and tomo_bin_b == "all": + return axs[2] + + else: + ncols = axs.shape[1] + return axs[tomo_bin_b - 1, ncols - tomo_bin_a] + + def _set_ax_visibility_to_false(self, axs, n_tomo_bins_plot): + """Set the visibility of empty axes to False for better visualization.""" + if n_tomo_bins_plot == 1: + axs[1].set_visible(False) + else: + for i in range(n_tomo_bins_plot): + axs[i, n_tomo_bins_plot - i].set_visible(False) From 1225bdecba6b5e80a52c56c2841ddd7b87f0482c Mon Sep 17 00:00:00 2001 From: Sacha Guerrini Date: Fri, 24 Jul 2026 09:03:45 +0200 Subject: [PATCH 12/12] Fix bugs identified by Claude --- src/sp_validation/cosmo_val/psf_systematics.py | 3 ++- src/sp_validation/rho_tau.py | 1 - 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 2c844e90..6765123e 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -57,7 +57,7 @@ def psf_fitter(self): @property def rho_tau_fits(self): if not hasattr(self, "_rho_tau_fits"): - self.calculate_rho_tau_fits() + self.calculate_rho_tau_fits(tomography=False) if self.compute_tomography: self.calculate_rho_tau_fits(tomography=True) return self._rho_tau_fits @@ -314,6 +314,7 @@ def set_params_rho_tau(self, ver, params, params_psf): params["dec_units"] = "deg" params["w_col"] = self.cc[ver]["shear"]["w_col"] + params["patch_number"] = self.cc[ver].get("patch_number", 100) return params diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index ae73c0c0..9bc975cc 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -59,7 +59,6 @@ def get_rho_tau_w_cov( mask_gal=None, cov_rho=False, ncov=100, - compute_minus=True, **kwargs, ): """Compute rho/tau statistics and, if requested, their covariance."""