Add CorrelationMatrix Plot

analysis/tasks/Plot/PlotUncertainties.py

 import matplotlib.pyplot as plt
+import numpy as np
 import uproot
 
 import law
@@ -85,8 +86,7 @@ class PlotUncertainties(AnalysisTask):
             ax.text(
                 0.03,
                 0.95,
-                "${} \\leq y_b < {}$\n".format(*_d["yb_range"])
-                + "${} \\leq y^* < {}$".format(*_d["ys_range"]),
+                _d["label"],
                 fontsize="small",
                 verticalalignment="top",
                 fontproperties="Tex Gyre Heros",
@@ -123,3 +123,73 @@ class PlotUncertainties(AnalysisTask):
             ax.legend(ncol=2)
             fig.tight_layout()
             MatplotlibFormatter.dump(self.output().targets[ybys], fig)
+
+
+@inherits(ExcaliburConfig)
+class PlotCorrelationMatrix(AnalysisTask):
+    dataset = luigi.Parameter(
+        default="data"
+        # default=DY_Sets.DYJetsToLL_combined.name
+    )  # overwrite default from ExcaliburConfig
+    input_is_mc = luigi.BoolParameter(default=False)
+    mc_dataset = luigi.EnumParameter(enum=DY_Sets, default=DY_Sets.DYJetsToLL_combined)
+    input_weight = luigi.Parameter(default="weight")
+    unfolding_weight = luigi.Parameter(default="totalWeight")
+    gen_weight = luigi.Parameter("weight")
+
+    bkg_variation_sfs = luigi.ListParameter(default=[1.5, 0.5])
+
+    selection_suffix = luigi.Parameter()
+    plot_range = luigi.Parameter(default="m0.25,0.35")
+
+    def requires(self):
+        unc = MergeUncertainties.req(self)
+        reravel = unc.requires()["unf"]
+        unfolded = reravel.requires()["Unfold"]
+        return unfolded
+
+    def output(self):
+        return 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}",
+                "correlation_matrix.pdf",
+                name="CorrelationMatrix",
+            ),
+            fs="local_plots",
+        )
+
+    def run(self):
+        quantity = "binIndexzpt"
+        self.output().parent.touch()
+        hep.style.use("CMS")
+        fig, ax = plt.subplots()
+        hep.cms.label(
+            label="Preliminary", data=not self.input_is_mc, lumi=138, loc=0, ax=ax
+        )
+        covariance, xbins, ybins = uproot.open(self.input().abspath)[
+            f"inclusive/h2d_{quantity}_{self.unfolding_weight}_stat_unc"
+        ].to_numpy()
+
+        v = np.sqrt(np.diag(covariance))
+        outer_v = np.outer(v, v)
+        correlation = covariance / outer_v
+        correlation[covariance == 0] = 0
+
+        hist2d = hep.hist2dplot(
+            correlation,
+            xbins=xbins,
+            ybins=ybins,
+            cbarextend=True,
+            flow=None,
+            ax=ax,
+        )
+        hist2d.cbar.set_label("Correlation coefficient")
+        ax.set_xlabel(QUANTITIES[quantity]["label"])
+        ax.set_ylabel(QUANTITIES[quantity]["label"])
+        fig.set_size_inches(
+            fig.get_size_inches()[0] * 1.05,
+            fig.get_size_inches()[1],
+        )
+        fig.tight_layout()
+        MatplotlibFormatter.dump(self.output(), fig)