diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 6765123e..a05d7eb9 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -209,17 +209,23 @@ def calculate_scale_dependent_leakage(self): self.print_done("Finished scale-dependent leakage calculation.") - def calculate_objectwise_leakage(self): + def calculate_objectwise_leakage(self, tomography=False): # TODO: Upgrade for tomography - if not hasattr(self.results[self.versions[0]], "alpha_leak_mean"): - self.calculate_scale_dependent_leakage() + # Get the tomographic bins + tomo_bins = self._get_tomo_bins_for_versions( + self.versions, tomography=tomography + ) self.print_start("Object-wise leakage:") mix = True order = "lin" + if not hasattr(self, "leakage_coeff"): + self.leakage_coeff = {} for ver in self.versions: self.print_magenta(ver) + self.leakage_coeff.setdefault(ver, {}) + results_obj = self.results_objectwise[ver] results_obj.check_params() results_obj.update_params() @@ -228,50 +234,48 @@ def calculate_objectwise_leakage(self): # 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, - ), - } + # Iterate on the tomographic bins for this version + for tomo_bin_id in tomo_bins[ver]["ids"]: + if tomo_bin_id == "all": + selection = None + else: + selection = self._get_galaxy_mask(ver, tomo_bin_id) - self.leakage_coeff = leakage_coeff + suffix = f"tomo_bin_{tomo_bin_id}" + + out_base = results_obj.get_out_base(mix, order, suffix=suffix) + 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(selection=selection): + try: + results_obj.PSF_leakage(suffix=suffix) + + # Gather coefficients + + 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) + continue + + par_best_fit = results_obj.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) + self.leakage_coeff[ver][f"tomo_bin_{tomo_bin_id}"] = { + "a11": a11, + "a22": a22, + "aii_mean": 0.5 * (a11 + a22), + } # --- utility functions --- def _get_galaxy_mask(self, ver, tomo_bin_id): @@ -453,6 +457,107 @@ def get_xi_psf_sys_samples(self, ver, params, tomo_bin_a, tomo_bin_b): return xi_psf_sys_samples_plus, xi_psf_sys_samples_minus + def _get_alpha_leakage( + self, + rho_stat_handler, + tau_stat_handler, + cov_rho=None, + cov_tau=None, + n_samples=10_000, + ): + """ + Compute the alpha leakage parameter from the rho and tau statistics. + + Parameters + ---------- + rho_stat_handler : RhoStatHandler + The handler for the rho statistics. + tau_stat_handler : TauStatHandler + The handler for the tau statistics. + cov_rho : np.ndarray, optional + The covariance matrix for the rho statistics. If None, it will be computed. + cov_tau : np.ndarray, optional + The covariance matrix for the tau statistics. If None, it will be computed. + + Returns + ------- + alpha_leak : float + The estimated alpha leakage parameter. + """ + if cov_rho is None: + cov_rho = np.diag(rho_stat_handler.rho_stats["varrho_0_p"]) + if cov_tau is None: + cov_tau = np.diag(tau_stat_handler.tau_stats["vartau_0_p"]) + + theta = rho_stat_handler.rho_stats["theta"] + n_bins = len(theta) + alpha = ( + tau_stat_handler.tau_stats["tau_0_p"] + / rho_stat_handler.rho_stats["rho_0_p"] + ) + + # Derive alpha_err by sampling from the covariance matrices of rho and tau statistics + rho_samples = np.random.multivariate_normal( + mean=rho_stat_handler.rho_stats["rho_0_p"], + cov=cov_rho[:n_bins, :n_bins], + size=n_samples, + ) + tau_samples = np.random.multivariate_normal( + mean=tau_stat_handler.tau_stats["tau_0_p"], + cov=cov_tau[:n_bins, :n_bins], + size=n_samples, + ) + + alpha_samples = tau_samples / rho_samples + alpha_err = np.std(alpha_samples, axis=0) + + return theta, alpha, alpha_err + + def _compute_scale_dependent_xi_psf_sys( + self, rho_0, tau_0_a, tau_0_b, cov_rho, cov_tau_a, cov_tau_b, n_samples=10_000 + ): + """ + Compute the scale-dependent xi_psf_sys from the rho and tau statistics. + + Parameters + ---------- + rho_0 : np.ndarray + The rho_0 statistics. + tau_0_a : np.ndarray + The tau_0 statistics for tomographic bin a. + tau_0_b : np.ndarray + The tau_0 statistics for tomographic bin b. + cov_rho : np.ndarray + The covariance matrix for the rho statistics. + cov_tau_a : np.ndarray + The covariance matrix for the tau statistics for tomographic bin a. + cov_tau_b : np.ndarray + The covariance matrix for the tau statistics for tomographic bin b. + """ + # Compute xi_psf_sys for each scale using the formula: + xi_psf_sys = (tau_0_a * tau_0_b) / rho_0 + + # Derive the error bars by sampling the statistics + rho_samples = np.random.multivariate_normal( + mean=rho_0, + cov=cov_rho, + size=n_samples, + ) + tau_samples_a = np.random.multivariate_normal( + mean=tau_0_a, + cov=cov_tau_a, + size=n_samples, + ) + tau_samples_b = np.random.multivariate_normal( + mean=tau_0_b, + cov=cov_tau_b, + size=n_samples, + ) + xi_psf_sys_samples = (tau_samples_a * tau_samples_b) / rho_samples + xi_psf_sys_err = np.std(xi_psf_sys_samples, axis=0) + + return xi_psf_sys, xi_psf_sys_err + # --- plotting functions --- def plot_rho_stats( self, @@ -901,169 +1006,191 @@ def plot_rho_tau_fits( **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" + def plot_scale_dependent_leakage( + self, + tomography=False, + cov_type=None, + versions=None, + colors=None, + offset=0, + savefig=None, + show=True, + close=True, + plot_theta_times_tau=False, + ylim_alpha=False, + fmt="", + capsize=2, + ): + # First plot alpha leakage + self.plot_scale_dependent_alpha( + tomography=tomography, + cov_type=cov_type, + versions=versions, + colors=colors, + offset=offset, + savefig=savefig, + show=show, + close=close, + fmt=fmt, + capsize=capsize, + ylim_alpha=ylim_alpha, + ) + + # Second plot xi_sys + self.plot_2pcf_tomography( + self._scale_dependent_xi_psf_sys_x_y_plot_function, + x_label=r"$\theta$ [arcmin]", + y_label_plus=(r"$\theta$" if plot_theta_times_tau else "") + + r"$\xi^{\rm PSF, sys}_+(\theta)$", + y_label_minus=(r"$\theta$" if plot_theta_times_tau else "") + + r"$\xi^{\rm PSF, sys}_-(\theta)$", + tomo_bin_label_position=(0.05, 0.9) + if not plot_theta_times_tau + else (0.8, 0.95), + extract_text_offset=plot_theta_times_tau, + add_index_version_to_kwargs=True, + x_scale="log", + y_scale="linear" if plot_theta_times_tau else "log", + tomography=tomography, + versions=versions, + colors=colors, + savefig=savefig.replace(".png", "_xi_psf_sys.png") + if savefig is not None + else None, + show=show, + close=close, + offset=offset, + cov_type=cov_type, + times_theta=plot_theta_times_tau, + fmt=fmt, + capsize=capsize, + ) + + def plot_scale_dependent_alpha( + self, + tomography=False, + cov_type=None, + versions=None, + colors=None, + offset=0, + savefig=None, + show=True, + close=True, + fmt="", + capsize=2, + ylim_alpha=False, + ): + 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}" + ) + + 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()) + + out_dir = f"{self.cc['paths']['output']}/rho_tau_stats" + + fig, axs = plt.subplots( + n_tomo_bins_plot, 1, figsize=(8, 3 * n_tomo_bins_plot), sharex=True + ) + + # Iterate upon each version + for ver, color in zip(versions, colors): + label = self.cc[ver]["label"] if "label" in self.cc[ver] else ver + # Iterate upon each tomographic bin + for tomo_bin_id in tomo_bins[ver]["ids"]: + base_rho = self.basename(ver) + base_tau = self.basename(ver, tomo_bin_a=tomo_bin_id) + self.rho_stat_handler.load_rho_stats(f"rho_stats_{base_rho}.fits") + 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) + cov_rho_path = Path(out_dir) / f"cov_rho_{base_rho}_jk.npy" + cov_rho = np.load(cov_rho_path) + else: + cov_tau = None + cov_rho = None + + # Get the error bar sampling from the covariance matrices + theta, alpha, alpha_err = self._get_alpha_leakage( + self.rho_stat_handler, self.tau_stat_handler, cov_rho, cov_tau ) - self.psf_fitter.plot_xi_psf_sys_terms( - ver, mcmc_result[1], out_path, yscale=yscale, show=True + + jittered_theta = self._get_jittered_theta( + theta, versions.index(ver), len(versions), offset ) - self.print_done( - f"{yscale}-scale xi_psf_sys terms plot saved to {out_path}" + + if tomo_bin_id == "all": + ax = axs + else: + ax = axs[tomo_bin_id - 1] + + ax.errorbar( + jittered_theta, + alpha, + yerr=alpha_err, + fmt=fmt, + label=f"{label}", + capsize=capsize, + color=color, ) - """ - def plot_scale_dependent_leakage(self): - if not hasattr(self.results[self.versions[0]], "r_corr_gp"): - self.calculate_scale_dependent_leakage() + if tomography: + for i, ax in enumerate(axs): + ax.set_xscale("log") + ax.set_xlim(self.theta_min_plot, self.theta_max_plot) + if ylim_alpha: + ax.set_ylim(self.ylim_alpha) + ax.set_ylabel(rf"$\alpha_{i + 1}(\theta)$") + if i == len(axs) - 1: + ax.set_xlabel(r"$\theta$ [arcmin]") + ax.text(0.05, 0.9, f"Tomo bin {i + 1}", transform=ax.transAxes) + else: + axs.set_xscale("log") + axs.set_xlim(self.theta_min_plot, self.theta_max_plot) + if ylim_alpha: + axs.set_ylim(self.ylim_alpha) + axs.set_ylabel(r"$\alpha_{\rm all}(\theta)$") + axs.set_xlabel(r"$\theta$ [arcmin]") + axs.text(0.05, 0.9, "All tomographic bins", transform=axs.transAxes) - theta = [] - y = [] - yerr = [] - labels = [] - colors = [] - linestyles = [] - markers = [] + if tomography: + handles, labels = axs[0].get_legend_handles_labels() + else: + handles, labels = axs.get_legend_handles_labels() - for ver in self.versions: - if hasattr(self.results[ver], "r_corr_gp"): - theta.append(self.results[ver].r_corr_gp.meanr) - y.append(self.results[ver].alpha_leak) - yerr.append(self.results[ver].sig_alpha_leak) - labels.append(ver) - colors.append(self.cc[ver]["colour"]) - linestyles.append(self.cc[ver]["ls"]) - markers.append(self.cc[ver]["marker"]) - - if len(theta) > 0: - # Log x - out_path = self._output_path("alpha_leak_log.png") - - title = r"$\alpha$ leakage" - xlabel = r"$\theta$ [arcmin]" - ylabel = r"$\alpha(\theta)$" - cs_plots.plot_data_1d( - theta, - y, - yerr, - title, - xlabel, - ylabel, - out_path=None, - xlog=True, - xlim=[self.theta_min_plot, self.theta_max_plot], - ylim=self.ylim_alpha, - labels=labels, - colors=colors, - linestyles=linestyles, - shift_x=True, - ) - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"Log-scale alpha leakage plot saved to {out_path}") - - # Lin x - out_path = self._output_path("alpha_leak_lin.png") - - title = r"$\alpha$ leakage" - xlabel = r"$\theta$ [arcmin]" - ylabel = r"$\alpha(\theta)$" - cs_plots.plot_data_1d( - theta, - y, - yerr, - title, - xlabel, - ylabel, - out_path=None, - xlog=False, - xlim=[-10, self.theta_max_plot], - ylim=self.ylim_alpha, - labels=labels, - colors=colors, - linestyles=linestyles, - shift_x=False, - ) - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"Lin-scale alpha leakage plot saved to {out_path}") + fig.legend( + handles, + labels, + loc="upper center", + bbox_to_anchor=(0.5, -0.02), + ncol=3, + frameon=False, + ) - # Plot xi_sys - y = [] - yerr = [] - colors = [] - linestyles = [] + if savefig is not None: + plt.savefig(savefig, dpi=300, bbox_inches="tight") + self.print_done(f"Scale-dependent alpha leakage plot saved to {savefig}") - for ver in self.versions: - if hasattr(self.results[ver], "C_sys_p"): - y.append(self.results[ver].C_sys_p) - yerr.append(self.results[ver].C_sys_std_p) - labels.append(ver) - colors.append(self.cc[ver]["colour"]) - linestyles.append(self.cc[ver]["ls"]) - - if len(y) > 0: - xlabel = r"$\theta$ [arcmin]" - ylabel = r"$\xi^{\rm sys}_+(\theta)$" - title = "Cross-correlation leakage" - out_path = self._output_path("xi_sys_p.png") - cs_plots.plot_data_1d( - theta, - y, - yerr, - title, - xlabel, - ylabel, - out_path=None, - labels=labels, - xlog=True, - xlim=[self.theta_min_plot, self.theta_max_plot], - colors=colors, - linestyles=linestyles, - # shift_x=True, - ) - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_sys_plus plot saved to {out_path}") + if show: + plt.show() - y = [] - yerr = [] - for ver in self.versions: - if hasattr(self.results[ver], "C_sys_m"): - y.append(self.results[ver].C_sys_m) - yerr.append(self.results[ver].C_sys_std_m) - - if len(y) > 0: - xlabel = r"$\theta$ [arcmin]" - ylabel = r"$\xi^{\rm sys}_-(\theta)$" - title = "Cross-correlation leakage" - out_path = self._output_path("xi_sys_m.png") - cs_plots.plot_data_1d( - theta, - y, - yerr, - title, - xlabel, - ylabel, - out_path=None, - labels=labels, - xlog=True, - xlim=[self.theta_min_plot, self.theta_max_plot], - ylim=[-1e-7, 1e-6], - colors=colors, - linestyles=linestyles, - # shift_x=True, - ) - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_sys_minus plot saved to {out_path}") + if close: + plt.close() def plot_objectwise_leakage(self): if not hasattr(self, "leakage_coeff"): @@ -1449,6 +1576,85 @@ def plot_axis(ax, ax_type): # Plot the minus axis plot_axis(ax_minus, "minus") + def _scale_dependent_xi_psf_sys_x_y_plot_function( + self, + ax_plus, + ax_minus, + version, + tomo_bin_a, + tomo_bin_b, + idx, + versions, + color, + offset, + cov_type, + times_theta, + fmt, + capsize, + ): + # Load the rho-stats and the tau-stats + base_rho = self.basename(version) + base_tau_a = self.basename(version, tomo_bin_a=tomo_bin_a) + base_tau_b = self.basename(version, tomo_bin_a=tomo_bin_b) + self.rho_stat_handler.load_rho_stats(f"rho_stats_{base_rho}.fits") + + theta = self.rho_stat_handler.rho_stats["theta"] + n_bins = len(theta) + rho_0_p = self.rho_stat_handler.rho_stats["rho_0_p"] + rho_0_m = self.rho_stat_handler.rho_stats["rho_0_m"] + + self.tau_stat_handler.load_tau_stats(f"tau_stats_{base_tau_a}.fits") + + tau_0_p_a = self.tau_stat_handler.tau_stats["tau_0_p"] + tau_0_m_a = self.tau_stat_handler.tau_stats["tau_0_m"] + + self.tau_stat_handler.load_tau_stats(f"tau_stats_{base_tau_b}.fits") + + tau_0_p_b = self.tau_stat_handler.tau_stats["tau_0_p"] + tau_0_m_b = self.tau_stat_handler.tau_stats["tau_0_m"] + + if cov_type is not None: + outdir = f"{self.cc['paths']['output']}/rho_tau_stats" + cov_tau_path = Path(outdir) / f"cov_tau_{base_tau_a}_{cov_type}.npy" + cov_tau_p_a = np.load(cov_tau_path)[:n_bins, :n_bins] + cov_tau_path = Path(outdir) / f"cov_tau_{base_tau_b}_{cov_type}.npy" + cov_tau_p_b = np.load(cov_tau_path)[:n_bins, :n_bins] + cov_rho_path = Path(outdir) / f"cov_rho_{base_rho}_jk.npy" + cov_rho_p = np.load(cov_rho_path)[:n_bins, :n_bins] + else: + cov_tau_p_a = np.diag(self.tau_stat_handler.tau_stats["vartau_0_p"]) + cov_tau_p_b = np.diag(self.tau_stat_handler.tau_stats["vartau_0_p"]) + cov_rho_p = np.diag(self.rho_stat_handler.rho_stats["varrho_0_p"]) + + cov_tau_m_a = np.diag(self.tau_stat_handler.tau_stats["vartau_0_m"]) + cov_tau_m_b = np.diag(self.tau_stat_handler.tau_stats["vartau_0_m"]) + cov_rho_m = np.diag(self.rho_stat_handler.rho_stats["varrho_0_m"]) + + # Compute the scale-dependent xi_psf_sys and its error bars + xi_psf_sys_plus, xi_psf_sys_plus_err = self._compute_scale_dependent_xi_psf_sys( + rho_0_p, tau_0_p_a, tau_0_p_b, cov_rho_p, cov_tau_p_a, cov_tau_p_b + ) + xi_psf_sys_minus, xi_psf_sys_minus_err = ( + self._compute_scale_dependent_xi_psf_sys( + rho_0_m, tau_0_m_a, tau_0_m_b, cov_rho_m, cov_tau_m_a, cov_tau_m_b + ) + ) + + jittered_theta = self._get_jittered_theta(theta, idx, len(versions), offset) + + y_plus = xi_psf_sys_plus * (theta if times_theta else 1) + y_plus_err = xi_psf_sys_plus_err * (theta if times_theta else 1) + y_minus = xi_psf_sys_minus * (theta if times_theta else 1) + y_minus_err = xi_psf_sys_minus_err * (theta if times_theta else 1) + + ax_plus.errorbar( + jittered_theta, y_plus, yerr=y_plus_err, fmt=fmt, capsize=capsize + ) + + ax_minus.errorbar( + jittered_theta, y_minus, yerr=y_minus_err, fmt=fmt, capsize=capsize + ) + def _get_jittered_theta(self, theta, idx, n_versions, offset): """Get the jittered theta values for better visualisation.""" theta_widths = np.diff(theta) @@ -1490,3 +1696,6 @@ def _set_ax_visibility_to_false(self, axs, n_tomo_bins_plot): else: for i in range(n_tomo_bins_plot): axs[i, n_tomo_bins_plot - i].set_visible(False) + + +# %%