From 2fc8ee0087b2e2df7df9b26c8651901fec30017a Mon Sep 17 00:00:00 2001
From: Sacha Guerrini <sacha.guerrini@cea.fr>
Date: Mon, 24 Aug 2026 14:38:34 +0200
Subject: [PATCH 1/3] Update de scale dependent plotting script to read the
 computed rho-tau-stats and plot tomography

---
 .../cosmo_val/psf_systematics.py              | 379 +++++++++++-------
 1 file changed, 240 insertions(+), 139 deletions(-)

diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py
index 6765123e..1a91a3d7 100644
--- a/src/sp_validation/cosmo_val/psf_systematics.py
+++ b/src/sp_validation/cosmo_val/psf_systematics.py
@@ -209,8 +209,10 @@ 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
+
+        # TODO: remove and save the results of the scale-dependent leakage independently
         if not hasattr(self.results[self.versions[0]], "alpha_leak_mean"):
             self.calculate_scale_dependent_leakage()
 
@@ -453,6 +455,62 @@ 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
+
     # --- plotting functions ---
     def plot_rho_stats(
         self,
@@ -920,150 +978,190 @@ def plot_rho_tau_fits(
                 )
         """
 
-    def plot_scale_dependent_leakage(self):
-        if not hasattr(self.results[self.versions[0]], "r_corr_gp"):
-            self.calculate_scale_dependent_leakage()
+    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,
+        )
 
-        theta = []
-        y = []
-        yerr = []
-        labels = []
-        colors = []
-        linestyles = []
-        markers = []
+        # Second plot xi_sys
+        self.plot_xi_sys_from_scale_dependent_leakage(
+            tomography=tomography,
+            cov_type=cov_type,
+            versions=versions,
+            colors=colors,
+            offset=offset,
+            savefig=savefig,
+            show=show,
+            close=close,
+            plot_theta_times_tau=plot_theta_times_tau,
+            fmt=fmt,
+            capsize=capsize,
+        )
 
-        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}")
+    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
 
-        # Plot xi_sys
-        y = []
-        yerr = []
-        colors = []
-        linestyles = []
+        if colors is None:
+            colors = [self.cc[ver]["colour"] for ver in versions]
 
-        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 len(colors) != len(versions):
+            raise ValueError("Colors and versions must have the same length.")
 
-        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,
+        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}"
             )
-            cs_plots.savefig(out_path, close_fig=False)
-            cs_plots.show()
-            self.print_done(f"xi_sys_minus plot saved to {out_path}")
+
+        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
+                )
+
+                jittered_theta = self._get_jittered_theta(
+                    theta, versions.index(ver), len(versions), offset
+                )
+
+                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,
+                )
+
+        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]")
+        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]")
+
+        if tomography:
+            handles, labels = axs[0].get_legend_handles_labels()
+        else:
+            handles, labels = axs.get_legend_handles_labels()
+
+        fig.legend(
+            handles,
+            labels,
+            loc="upper center",
+            bbox_to_anchor=(0.5, -0.02),
+            ncol=3,
+            frameon=False,
+        )
+
+        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}")
+
+        if show:
+            plt.show()
+
+        if close:
+            plt.close()
+
+    def plot_xi_sys_from_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,
+        fmt="",
+        capsize=2,
+    ):
+        pass
 
     def plot_objectwise_leakage(self):
         if not hasattr(self, "leakage_coeff"):
@@ -1490,3 +1588,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)
+
+
+# %%

From 35795aba67f63a5416505f7db05ecc8f39c7a361 Mon Sep 17 00:00:00 2001
From: Sacha Guerrini <sacha.guerrini@cea.fr>
Date: Tue, 25 Aug 2026 09:39:52 +0200
Subject: [PATCH 2/3] Update the plotting script for the scale dependent xi_sys

---
 .../cosmo_val/psf_systematics.py              | 186 ++++++++++++++----
 1 file changed, 146 insertions(+), 40 deletions(-)

diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py
index 1a91a3d7..9d0dd658 100644
--- a/src/sp_validation/cosmo_val/psf_systematics.py
+++ b/src/sp_validation/cosmo_val/psf_systematics.py
@@ -511,6 +511,51 @@ def _get_alpha_leakage(
 
         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,
@@ -959,25 +1004,6 @@ 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"
-                )
-                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,
         tomography=False,
@@ -1009,16 +1035,31 @@ def plot_scale_dependent_leakage(
         )
 
         # Second plot xi_sys
-        self.plot_xi_sys_from_scale_dependent_leakage(
+        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,
-            cov_type=cov_type,
             versions=versions,
             colors=colors,
-            offset=offset,
-            savefig=savefig,
+            savefig=savefig.replace(".png", "_xi_psf_sys.png")
+            if savefig is not None
+            else None,
             show=show,
             close=close,
-            plot_theta_times_tau=plot_theta_times_tau,
+            offset=offset,
+            cov_type=cov_type,
+            times_theta=plot_theta_times_tau,
             fmt=fmt,
             capsize=capsize,
         )
@@ -1115,6 +1156,7 @@ def plot_scale_dependent_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)
@@ -1122,6 +1164,7 @@ def plot_scale_dependent_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)
 
         if tomography:
             handles, labels = axs[0].get_legend_handles_labels()
@@ -1147,22 +1190,6 @@ def plot_scale_dependent_alpha(
         if close:
             plt.close()
 
-    def plot_xi_sys_from_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,
-        fmt="",
-        capsize=2,
-    ):
-        pass
-
     def plot_objectwise_leakage(self):
         if not hasattr(self, "leakage_coeff"):
             self.calculate_objectwise_leakage()
@@ -1547,6 +1574,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)

From a66b48bfcdf342b8d11414eaae4c2bee8bde09e5 Mon Sep 17 00:00:00 2001
From: Sacha Guerrini <sacha.guerrini@cea.fr>
Date: Wed, 26 Aug 2026 18:23:59 +0200
Subject: [PATCH 3/3] Update the objectwise leakage to run in the tomographic
 case.

---
 .../cosmo_val/psf_systematics.py              | 96 ++++++++++---------
 1 file changed, 49 insertions(+), 47 deletions(-)

diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py
index 9d0dd658..a05d7eb9 100644
--- a/src/sp_validation/cosmo_val/psf_systematics.py
+++ b/src/sp_validation/cosmo_val/psf_systematics.py
@@ -211,17 +211,21 @@ def calculate_scale_dependent_leakage(self):
 
     def calculate_objectwise_leakage(self, tomography=False):
         # TODO: Upgrade for tomography
-
-        # TODO: remove and save the results of the scale-dependent leakage independently
-        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()
@@ -230,50 +234,48 @@ def calculate_objectwise_leakage(self, tomography=False):
             # 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)
+
+                suffix = f"tomo_bin_{tomo_bin_id}"
 
-        self.leakage_coeff = leakage_coeff
+                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):