Update NP corrections

analysis/tasks/Theory/NPCorrections.py

@@ -11,6 +11,7 @@ import luigi
 import mplhep as hep
 from analysis.framework.quantities import QUANTITIES, ybys_bins
 from analysis.tools.plotting.legend_helpers import merge_legend_handles_labels
+from analysis.tools.plotting.styles import colormap_5_no_grey
 from law.contrib.matplotlib import MatplotlibFormatter
 from law.decorator import safe_output
 
@@ -27,6 +28,8 @@ def json_numpy_obj_hook(dct):
     if isinstance(dct, dict) and "__ndarray__" in dct:
         data = base64.b64decode(dct["__ndarray__"])
         return np.frombuffer(data, dct["dtype"]).reshape(dct["shape"])
+    elif isinstance(dct, dict) and "ndarray" in dct:
+        return np.array(dct["ndarray"])
     return dct
 
 
@@ -49,7 +52,7 @@ class NPCalc(law.ExternalTask):
         targets = {}
         for ybys in ybys_bins:
             targets[ybys] = law.LocalFileTarget(
-                "/work/mhorzela/mc-run/data/PlotNPCorr/"
+                "/work/cverstege/QCD_shared/store/PlotNPCorr/"
                 + f"{self.generator.value}/{self.generator_settings}/"
                 + f"withNP-NPoff-Ratio-Plots/m-ZPt-um-RAW/ZplusJet_3_AK{self.jet_radius}{ybys}.json",
             )
@@ -103,12 +106,46 @@ class NPCorrections(law.Task):
                 [
                     correction_dicts[setting]["ys"]
                     for setting in correction_dicts
-                    if setting in ("sherpa_LHC-ZJetMerging", "herwig_LHC-NLO-ZplusJet")
+                    if setting
+                    in (
+                        "sherpa_LHC-ZJetMerging",
+                        "sherpa_LHC-NLO-ZplusJet",
+                        "herwig_LHC-NLO-ZplusJet",
+                    )
                 ],
             )
-            up = np.max(corrections, axis=0)
-            down = np.min(corrections, axis=0)
+            corrections_all = np.stack(
+                [correction_dicts[setting]["ys"] for setting in correction_dicts],
+            )
+            uncertainties = np.nan_to_num(
+                np.stack(
+                    [
+                        correction_dicts[setting]["yerrs"]
+                        for setting in correction_dicts
+                        if setting
+                        in (
+                            "sherpa_LHC-ZJetMerging",
+                            "sherpa_LHC-NLO-ZplusJet",
+                            "herwig_LHC-NLO-ZplusJet",
+                        )
+                    ],
+                )
+            )
+            uncertainties_all = np.nan_to_num(
+                np.stack(
+                    [
+                        correction_dicts[setting]["yerrs"]
+                        for setting in correction_dicts
+                    ],
+                )
+            )
+
+            up = np.max(corrections + uncertainties, axis=0)
+            down = np.min(corrections - uncertainties, axis=0)
             val = (up + down) / 2
+            up_all = np.max(corrections_all + uncertainties_all, axis=0)
+            down_all = np.min(corrections_all - uncertainties_all, axis=0)
+            val_all = (up_all + down_all) / 2
 
             # plot the corrections
             fig, ax = plt.subplots()
@@ -120,6 +157,7 @@ class NPCorrections(law.Task):
                 yerr=[val - down, up - val],
                 label=r"NP correction $\pm$ uncertainty",
                 edges=False,
+                zorder=-2,
             )
             hep.histplot(
                 val,
@@ -129,8 +167,29 @@ class NPCorrections(law.Task):
                 edges=False,
                 color="black",
                 linewidth=4,
+                zorder=-1,
             )
-            for setting in self.input().keys():
+            hep.histplot(
+                val_all,
+                bins=_d["zpt_binning"],
+                histtype="band",
+                yerr=[val_all - down_all, up_all - val_all],
+                label=r"NP correction (alt) $\pm$ uncertainty",
+                edges=False,
+                hatch="\\\\",
+                zorder=-12,
+            )
+            hep.histplot(
+                val_all,
+                bins=_d["zpt_binning"],
+                histtype="step",
+                label=r"NP correction (alt) $\pm$ uncertainty",
+                edges=False,
+                color="grey",
+                linewidth=3,
+                zorder=-11,
+            )
+            for setting, color in zip(self.input().keys(), colormap_5_no_grey):
                 label = ""
                 if "sherpa" in setting:
                     label += "Sherpa"
@@ -145,10 +204,15 @@ class NPCorrections(law.Task):
                 hep.histplot(
                     correction_dicts[setting]["ys"],
                     bins=_d["zpt_binning"],
-                    histtype="step",
+                    histtype="errorbar",
+                    yerr=correction_dicts[setting]["yerrs"],
                     edges=False,
                     linewidth=2,
                     label=label,
+                    markersize=16,
+                    elinewidth=3,
+                    capsize=4,
+                    color=color,
                 )
             ax.set_ylabel("NP correction")
             ax.set_xlabel(QUANTITIES["zpt"]["label"])
@@ -194,11 +258,37 @@ class NPCorrections(law.Task):
                 len(val),
                 _d["zpt_binning"],
             )
-            for i, (v, u, d) in enumerate(zip(val, up, down), start=1):
+            root_hist_alt = ROOT.TH1D(
+                "h_genzpt_npcorr_alt",
+                f"{ybys}_genzpt_npcorr_alt",
+                len(val),
+                _d["zpt_binning"],
+            )
+            root_hist_up_alt = ROOT.TH1D(
+                "h_genzpt_npcorr_alt_up",
+                f"{ybys}_genzpt_npcorr_alt_up",
+                len(val),
+                _d["zpt_binning"],
+            )
+            root_hist_down_alt = ROOT.TH1D(
+                "h_genzpt_npcorr_alt_down",
+                f"{ybys}_genzpt_npcorr_alt_down",
+                len(val),
+                _d["zpt_binning"],
+            )
+            for i, (v, u, d, va, ua, da) in enumerate(
+                zip(val, up, down, val_all, up_all, down_all), start=1
+            ):
                 root_hist.SetBinContent(i, v)
                 root_hist_up.SetBinContent(i, u)
                 root_hist_down.SetBinContent(i, d)
+                root_hist_alt.SetBinContent(i, va)
+                root_hist_up_alt.SetBinContent(i, ua)
+                root_hist_down_alt.SetBinContent(i, da)
             root_hist.Write()
             root_hist_up.Write()
             root_hist_down.Write()
+            root_hist_alt.Write()
+            root_hist_up_alt.Write()
+            root_hist_down_alt.Write()
         out_file.Close()