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/core.py b/src/sp_validation/cosmo_val/core.py index c6204c65..9c041788 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -483,11 +483,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}" @@ -682,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 13f747ad..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, @@ -1206,7 +1201,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) @@ -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 9f8a247e..2c844e90 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -6,10 +6,13 @@ """ 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 from cs_util import plots as cs_plots from shear_psf_leakage import leakage from shear_psf_leakage import plots as psfleak_plots @@ -22,111 +25,25 @@ ) +# TODO: Reorganise the order of functions so it is more readable class PSFSystematicsMixin: - def calculate_rho_tau_stats(self): - 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"): - 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 @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 - def plot_rho_stats(self, abs=False): - filenames = [f"rho_stats_{self.basename(ver)}.fits" for ver in self.versions] - - savefig = "rho_stats.png" - self.rho_stat_handler.plot_rho_stats( - filenames, - self.colors, - self.versions, - savefig=savefig, - legend="outside", - abs=abs, - show=True, - close=True, - ) - - 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] - - 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, - ) - - self.print_done( - "Tau stats plot saved to " - + 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"): @@ -137,150 +54,268 @@ def psf_fitter(self): ) return self._psf_fitter - def calculate_rho_tau_fits(self): + @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")] + + # 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) + 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) + + # 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 + + 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": []} + 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( - self.results[ver]._params, - self.cc[ver]["psf"], - survey=ver, + ver, self.results[ver]._params, self.cc[ver]["psf"] ) - npatch = {"sim": 300, "jk": params["patch_number"]}.get( - self.cov_estimate_method, None - ) + # Get the tomographic bins + if tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - base = self.basename(ver) + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise ValueError( + f"Version {ver} does not have tomography information." + ) - 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, - ) + else: + tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] - 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) + # 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}" + ) + _ = self.get_samples( + ver, params, tomo_bin_id, track_result=track_result + ) - 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), - } + # 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) + ) - @property - def rho_tau_fits(self): - if not hasattr(self, "_rho_tau_fits"): - self.calculate_rho_tau_fits() - return self._rho_tau_fits + if ver not in self._xi_psf_sys.keys(): + self._xi_psf_sys[ver] = {} + + 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 plot_rho_tau_fits(self): - out_dir = self.rho_stat_handler.catalogs._output + 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) + results = self.results[ver] - savefig = f"{out_dir}/contours_tau_stat.png" - psfleak_plots.plot_contours( - self.rho_tau_fits["flat_sample_list"], - names=["x0", "x1", "x2"], - labels=[r"\alpha", r"\beta", r"\eta"], - savefig=savefig, - legend_labels=self.versions, - legend_loc="upper right", - contour_colors=self.colors, - markers={"x0": 0, "x1": 1, "x2": 1}, - show=True, - close=True, - ) - self.print_done(f"Tau contours plot saved to {os.path.abspath(savefig)}") + 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" + ) - 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}") + results.r_corr_gp = treecorr.GGCorrelation(self.treecorr_config) + results.r_corr_gp.read(output_path_ab) - 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, - ) + results.r_corr_pp = treecorr.GGCorrelation(self.treecorr_config) + results.r_corr_pp.read(output_path_aa) - 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}") + else: + results.compute_corr_gp_pp_alpha(output_base_path=output_base_path) - for mcmc_result, ver, flat_sample in zip( - self.rho_tau_fits["result_list"], - self.versions, - self.rho_tau_fits["flat_sample_list"], - ): - self.psf_fitter.load_rho_stat(f"rho_stats_{self.basename(ver)}.fits") - for yscale in ("linear", "log"): - out_path = os.path.abspath( - f"{out_dir}/xi_psf_sys_terms_{yscale}_{ver}.png" - ) - self.psf_fitter.plot_xi_psf_sys_terms( - ver, mcmc_result[1], out_path, yscale=yscale, show=True - ) - self.print_done( - f"{yscale}-scale xi_psf_sys terms plot saved to {out_path}" + results.do_alpha(fast=True) + results.do_xi_sys() + + 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() + + 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") - @property - def xi_psf_sys(self): - if not hasattr(self, "_xi_psf_sys"): - self.calculate_rho_tau_fits() - return self._xi_psf_sys + # 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"]) + 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 set_params_leakage_scale(self, ver): params_in = {} @@ -341,35 +376,548 @@ def set_params_leakage_object(self, ver): 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] + 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 - 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" + 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, + 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: + if version not in self.rho_tau_fits.keys(): + self.rho_tau_fits[version] = { + "flat_samples": {}, + "result": {}, + "quantile": {}, + } + self.rho_tau_fits[version]["flat_samples"][f"tomo_bin_{tomo_bin_id}"] = ( + flat_samples + ) + 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 + + 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"] + + # 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) + + 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_plus, xi_psf_sys_samples_minus + + # --- plotting functions --- + 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.") + + self.rho_stat_handler.plot_rho_stats( + filenames, + colors, + versions, + offset=offset, + savefig=savefig, + legend="outside", + abs=abs, + show=show, + close=close, + ) + + 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, + 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}" + ) + + out_dir = f"{self.cc['paths']['output']}/rho_tau_stats" + + 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 = self._get_tomo_bins_for_versions( + versions, tomography=tomography + ) + + 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] + + jittered_theta = self._get_jittered_theta( + theta, file_idx, len(versions), offset + ) + + 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, ) - results.r_corr_gp = treecorr.GGCorrelation(self.treecorr_config) - results.r_corr_gp.read(output_path_ab) + plt.tight_layout() - results.r_corr_pp = treecorr.GGCorrelation(self.treecorr_config) - results.r_corr_pp.read(output_path_aa) + if savefig is not None: + plt.savefig(savefig, dpi=300, bbox_inches="tight") - else: - results.compute_corr_gp_pp_alpha(output_base_path=output_base_path) + if show: + plt.show() - results.do_alpha(fast=True) - results.do_xi_sys() + if close: + plt.close() - self.print_done("Finished scale-dependent leakage calculation.") + 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 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 + + 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( + sample_list, + names=["x0", "x1", "x2"], + labels=[r"\alpha", r"\beta", r"\eta"], + savefig=savefig_contours, + legend_labels=versions, + legend_loc="upper right", + contour_colors=colors, + markers={"x0": 0, "x1": 1, "x2": 1}, + show=show, + close=close, + ) + if savefig_contours is not None: + self.print_done(f"Tau contours plot saved to {os.path.abspath(savefig)}") + + # 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.") + + 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}_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_samples"], + ): + self.psf_fitter.load_rho_stat(f"rho_stats_{self.basename(ver)}.fits") + for yscale in ("linear", "log"): + out_path = os.path.abspath( + f"{out_dir}/xi_psf_sys_terms_{yscale}_{ver}.png" + ) + self.psf_fitter.plot_xi_psf_sys_terms( + ver, mcmc_result[1], out_path, yscale=yscale, show=True + ) + 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"): @@ -516,69 +1064,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() @@ -625,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) diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 95a720f9..ae73c0c0 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 @@ -16,22 +17,10 @@ def _extract_xip(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"] @@ -40,21 +29,18 @@ 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"] - if survey != "DES": - params["PSF_flag"] = cat["psf"]["PSF_flag"] - params["star_flag"] = cat["psf"]["star_flag"] + 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"] - try: - params["tomo_bin_col"] = cat["shear"]["tomo_bin_col"] - except KeyError: - params["tomo_bin_col"] = None + params["tomo_bin_col"] = cat["shear"].get("tomo_bin_col") params["R11"] = cat["shear"].get("R11") params["R22"] = cat["shear"].get("R22") @@ -66,43 +52,60 @@ def get_rho_tau_w_cov( version, treecorr_config, outdir, - base, + base_rho, + base_tau, 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, 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, + mask_star=mask_star, + mask_gal=mask_gal, ) 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, + 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}_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}") @@ -112,8 +115,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( @@ -128,9 +134,12 @@ def get_rho_tau( version, treecorr_config, outdir, - base, + base_rho, + base_tau, + 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. @@ -145,16 +154,28 @@ def get_rho_tau( TreeCorr configuration (must include 'min_sep', 'max_sep', and 'nbins'). outdir : str Output directory. + 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 + 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() 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( @@ -163,17 +184,15 @@ 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" - rho_stat_handler.build_cat_to_compute_rho( config[version]["psf"]["path"], catalog_id=catalog_id, - mask=mask, + mask=mask_star, hdu=( config[version]["psf"]["hdu"] if config[version]["psf"]["hdu"] is not None @@ -193,7 +212,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, @@ -202,21 +221,19 @@ 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" - # 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, - mask=mask, + mask=mask_star, hdu=( config[version]["psf"]["hdu"] if config[version]["psf"]["hdu"] is not None @@ -229,7 +246,7 @@ def get_rho_tau( config[version]["shear"]["path"], cat_type="gal", catalog_id=version, - mask=mask, + mask=mask_gal, ) # function to extract the tau_+ @@ -247,17 +264,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"] @@ -277,21 +294,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 @@ -300,20 +319,24 @@ def get_jackknife_cov( version, treecorr_config, outdir, - base, + base_rho, + base_tau, 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. Also compute rho and tau-statistics. """ + # 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" - 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" - - 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 @@ -328,13 +351,13 @@ 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 - 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 @@ -352,8 +375,8 @@ def get_jackknife_cov( tau_stat_handler.catalogs.set_params(params, outdir) 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})") @@ -362,7 +385,7 @@ def get_jackknife_cov( rho_stat_handler.build_cat_to_compute_rho( config[version]["psf"]["path"], catalog_id=version + str(i), - mask=False, + mask=mask_star, hdu=config[version]["psf"]["hdu"], ) @@ -375,7 +398,7 @@ def get_jackknife_cov( config[version]["shear"]["path"], cat_type="gal", catalog_id=version + str(i), - mask=False, + mask=mask_gal, ) else: @@ -425,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 @@ -439,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) @@ -448,11 +471,13 @@ def get_jackknife_cov( def get_samples( psf_fitter, - version, - base, + 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. @@ -460,10 +485,10 @@ def get_samples( ---------- psf_fitter : PSFFitter 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 @@ -471,22 +496,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, + 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, + base_rho, + base_tau, cov_type=cov_type, apply_debias=apply_debias, + nsamples=nsamples, ) else: raise ValueError("Sampler must be either 'emcee' or 'lsq'.") @@ -494,8 +526,8 @@ def get_samples( def get_samples_emcee( psf_fitter, - version, - base, + base_rho, + base_tau, nwalkers=124, nsamples=10000, cov_type="jk", @@ -507,10 +539,10 @@ 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 : 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 @@ -521,20 +553,21 @@ 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) + 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) - mcmc_result, q = psf_fitter.get_mcmc_from_samples(base) + 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: 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 @@ -543,16 +576,17 @@ 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_sample) return flat_samples, mcmc_result, q def get_samples_lsq( psf_fitter, - version, - base, + base_rho, + base_tau, + nsamples=10000, apply_debias=None, cov_type="jk", ): @@ -562,36 +596,39 @@ 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. + 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") + 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) + 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) + 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: 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) + npatch=debias_npatch, + apply_debias=(debias_npatch is not None), + n_samples=nsamples, ) - psf_fitter.save_samples(flat_samples, base) + psf_fitter.save_samples(flat_samples, base_sample) return flat_samples, mcmc_result, q 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 )