Add uncertainty split plots

analysis/tasks/Plot/PlotUncertainties.py

@@ -10,6 +10,7 @@ from analysis.framework.tasks import AnalysisTask, ExcaliburConfig
 from analysis.framework.utils import DY_Sets
 from analysis.tasks.PostProcessing import MergeUncertainties
 from analysis.tools.plotting.axes_ticker import set_ticklabels
+from analysis.tools.plotting.legend_helpers import inline_label
 from law.contrib.matplotlib import MatplotlibFormatter
 from luigi.util import inherits
 
@@ -49,21 +50,24 @@ class PlotUncertainties(AnalysisTask):
     def output(self):
         filelist = {}
         for ybys_bin in ybys_bins:
-            filelist[ybys_bin] = law.LocalFileTarget(
-                self.local_path(
-                    f"{self.excalibur_config_version}_{self.selection_suffix}_cuts",
-                    f"{self.dataset}-{self.input_weight}_unfolded_with_{self.mc_dataset}-{self.unfolding_weight}",
-                    f"{ybys_bin}.pdf",
-                ),
-                fs="local_plots",
-            )
+            splits = {}
+            for split in ("total", "jet", "lum", "sf", "stat"):
+                splits[split] = law.LocalFileTarget(
+                    self.local_path(
+                        f"{self.excalibur_config_version}_{self.selection_suffix}_cuts",
+                        f"{self.dataset}-{self.input_weight}_unfolded_with_{self.mc_dataset}-{self.unfolding_weight}",
+                        f"{ybys_bin}_{split}.pdf",
+                    ),
+                    fs="local_plots",
+                )
+            filelist[ybys_bin] = splits
         return law.SiblingFileCollection(filelist)
 
     def run(self):
         data_file = uproot.open(self.input()["data"].abspath)
         unc_file = uproot.open(self.input()["unc"].abspath)
         uncertainties = {
-            "total": dict(label="Total exp. unc.", color="black", linewidth=3),
+            "total": dict(label="Total", color="black", linewidth=3),
             "stat": dict(color="#999999", label="Stat. unc."),
             "unf": dict(label="MC stat.", color="#377eb8", linestyle="dashdot"),
             "smooth_jec_Total_calc": dict(color="#e41a1c", label="JES"),
@@ -75,54 +79,63 @@ class PlotUncertainties(AnalysisTask):
             "prefiring": dict(color="#4daf4a", label="L1 Prefiring"),
             "bkgsf": dict(color="#f781bf", label="Background"),
             "lumi": dict(color="#984ea3", label="Lumi"),
+            "PileUp": dict(color="#47429d", linestyle="dotted", label="PileUp"),
         }
         self.output().dir.touch()
         hep.style.use("CMS")
         for ybys, _d in ybys_bins.items():
-            fig, ax = plt.subplots()
-            hep.cms.label(
-                label="Preliminary", data=not self.input_is_mc, lumi=138, loc=0, ax=ax
-            )
-            ax.text(
-                0.03,
-                0.95,
-                _d["label"],
-                fontsize="small",
-                verticalalignment="top",
-                fontproperties="Tex Gyre Heros",
-                transform=ax.transAxes,
-            )
-            ax.set_xscale("log")
-            ax.set_xlim(25, 1000)
-            ax.set_ylim(-0.2, 0.4)
-            ax.hlines(0, 25, 1000, color="black", linestyle="dashed")
-            data, bins = data_file[f"{ybys}/h_zpt_{self.unfolding_weight}"].to_numpy()
-            for unc, unc_dict in uncertainties.items():
-                up, _ = unc_file[f"{ybys}/h_zpt_{unc}_up"].to_numpy()
-                down, _ = unc_file[f"{ybys}/h_zpt_{unc}_down"].to_numpy()
-                ax.stairs(
-                    up / data,
-                    edges=bins,
-                    label=unc_dict["label"],
-                    color=unc_dict["color"],
-                    linestyle=unc_dict.get("linestyle", "-"),
-                    linewidth=unc_dict.get("linewidth", 2),
-                    baseline=None,
-                )
-                ax.stairs(
-                    down / data,
-                    edges=bins,
-                    color=unc_dict["color"],
-                    linestyle=unc_dict.get("linestyle", "-"),
-                    linewidth=unc_dict.get("linewidth", 2),
-                    baseline=None,
+            for split in ("total", "jet", "lum", "sf", "stat"):
+                fig, ax = plt.subplots()
+                hep.cms.label(
+                    label="Preliminary",
+                    data=not self.input_is_mc,
+                    lumi=138,
+                    loc=0,
+                    ax=ax,
                 )
-            ax.set_ylabel("Relative uncertainty")
-            ax.set_xlabel(QUANTITIES["zpt"]["label"])
-            set_ticklabels(ax)
-            ax.legend(ncol=2)
-            fig.tight_layout()
-            MatplotlibFormatter.dump(self.output().targets[ybys], fig)
+                inline_label(ax, _d["label"])
+                ax.set_xscale("log")
+                ax.set_xlim(25, 1000)
+                ax.set_ylim(-0.2, 0.4)
+                ax.hlines(0, 25, 1000, color="black", linestyle="dashed")
+                data, bins = data_file[
+                    f"{ybys}/h_zpt_{self.unfolding_weight}"
+                ].to_numpy()
+                for unc, unc_dict in uncertainties.items():
+                    if split == "jet" and unc not in ("smooth_jec_Total_calc", "jer"):
+                        continue
+                    elif split == "lum" and unc not in ("lumi", "bkgsf", "PileUp"):
+                        continue
+                    elif split == "sf" and unc not in ("muon", "prefiring"):
+                        continue
+                    elif split == "stat" and unc not in ("stat", "unf"):
+                        continue
+                    up, _ = unc_file[f"{ybys}/h_zpt_{unc}_up"].to_numpy()
+                    down, _ = unc_file[f"{ybys}/h_zpt_{unc}_down"].to_numpy()
+                    ax.stairs(
+                        up / data,
+                        edges=bins,
+                        label=unc_dict["label"],
+                        color=unc_dict["color"],
+                        linestyle=unc_dict.get("linestyle", "-"),
+                        linewidth=unc_dict.get("linewidth", 2),
+                        baseline=None,
+                    )
+                    ax.stairs(
+                        down / data,
+                        edges=bins,
+                        color=unc_dict["color"],
+                        linestyle=unc_dict.get("linestyle", "-"),
+                        linewidth=unc_dict.get("linewidth", 2),
+                        baseline=None,
+                    )
+                ax.set_ylabel("Relative uncertainty")
+                ax.set_xlabel(QUANTITIES["zpt"]["label"])
+                set_ticklabels(ax)
+                ax.legend(ncol=2)
+                fig.tight_layout()
+                MatplotlibFormatter.dump(self.output().targets[ybys][split], fig)
+                plt.close(fig)
 
 
 @inherits(ExcaliburConfig)