Make PlotSignalBG a workflow

analysis/tasks/Plot/PlotSignalBackground.py

@@ -11,11 +11,12 @@ from analysis.framework.utils import DY_Sets, bg_names, karma_env, run_command
 from analysis.tasks.PostProcessing import Histograms, MergeMCShapeUncertainties
 from analysis.tools.plotting.axes_ticker import set_ticklabels
 from law.contrib.matplotlib import MatplotlibFormatter
+from law.workflow.local import LocalWorkflow
 from luigi.util import inherits
 
 
 @inherits(LumberjackConfig)
-class PlotSignalBackground(AnalysisTask):
+class PlotSignalBackground(AnalysisTask, LocalWorkflow):
     dataset = luigi.Parameter(default="data")
     mc_dataset = luigi.EnumParameter(enum=DY_Sets, default=DY_Sets.DYJetsToLL_combined)
 
@@ -24,9 +25,13 @@ class PlotSignalBackground(AnalysisTask):
 
     ratio_range = luigi.TupleParameter(default=(0.75, 1.25))
     scale_to_data = luigi.BoolParameter(default="False")
-    quantities = luigi.ListParameter()
+    quantities = luigi.ListParameter(default=QUANTITIES.keys())
 
-    def requires(self):
+    def create_branch_map(self):
+        branches = {i: q for i, q in enumerate(self.quantities)}
+        return branches
+
+    def workflow_requires(self):
         req = {
             "uncs": MergeMCShapeUncertainties.req(
                 self, dataset=self.mc_dataset.name, task="DATAShape", splitting="ybys"
@@ -108,6 +113,9 @@ class PlotSignalBackground(AnalysisTask):
             ]
         )
         mc_stack = np.sum(mc_hists, axis=0)
+        if np.any(mc_stack <= 0):
+            print(f"MC stack containes negative events for {ybys}/{quantity}!")
+            print(mc_stack)
         hep.histplot(
             mc_hists,
             bins,
@@ -120,7 +128,7 @@ class PlotSignalBackground(AnalysisTask):
         hep.histplot(
             data / mc_stack,
             bins,
-            yerr=data_errors / mc_stack,
+            yerr=data_errors / np.abs(mc_stack),
             histtype="errorbar",
             color="black",
             ax=axratio,
@@ -154,7 +162,7 @@ class PlotSignalBackground(AnalysisTask):
 
         axmain.set_xscale(q_dict.get("scale", "linear"))
         axmain.set_yscale("log")
-        axmain.set_ylim(1e-2, 1e6)
+        # axmain.set_ylim(1e-2, 1e6)
         axmain.set_xlim(q_dict["range"])
         axmain.set_ylabel("Events / GeV")
         axratio.set_xlabel(q_dict["label"])
@@ -168,18 +176,22 @@ class PlotSignalBackground(AnalysisTask):
         handles, labels = hep.sort_legend(axmain, ["Data"] + mc_labels)
         axmain.legend(handles, labels, ncol=2, loc="upper right")
         axratio.legend(ncol=1, loc="upper left")
+        axmain = hep.yscale_legend(axmain)
         fig.tight_layout()
         MatplotlibFormatter.dump(self.output().targets[f"{ybys}/{quantity}"], fig)
+        plt.close(fig)
 
     def run(self):
         for collection in self.output().collections:
             collection.dir.touch()
-        self.data_file = uproot.open(self.input()["data"].abspath)
-        self.unc_file = uproot.open(self.input()["uncs"].abspath)
-        self.mc_files = {bg: uproot.open(self.input()[bg].abspath) for bg in bg_names}
-        self.mc_files["DY"] = uproot.open(self.input()["DY"].abspath)
+        self.data_file = uproot.open(self.workflow_input()["data"].abspath)
+        self.unc_file = uproot.open(self.workflow_input()["uncs"].abspath)
+        self.mc_files = {
+            bg: uproot.open(self.workflow_input()[bg].abspath) for bg in bg_names
+        }
+        self.mc_files["DY"] = uproot.open(self.workflow_input()["DY"].abspath)
+        quantity = self.branch_data
         for ybys in ybys_bins_inclusive:
-            for quantity in self.quantities:
             if self.output().targets[f"{ybys}/{quantity}"].exists():
                 continue
             self.plot_quantity(ybys, quantity, q_dict=QUANTITIES[quantity])