[PWGHF] Cut variation: make analysis more flexible and output more informative (#11985)

Author: lubynets

PWGHF/D2H/Macros/compute_fraction_cutvar.py

@@ -55,6 +55,12 @@ def main(config):
         hist_effnp[-1].SetDirectory(0)
         infile_eff.Close()
 
+    pt_bin_to_process = cfg.get("pt_bin_to_process", -1)
+    if not isinstance(pt_bin_to_process, int):
+        sys.exit("Fatal error: pt_bin_to_process must be an integer value. Exit.")
+    if (pt_bin_to_process != -1 and pt_bin_to_process < 1) or pt_bin_to_process > hist_rawy[0].GetNbinsX():
+        sys.exit("Fatal error: pt_bin_to_process must be a positive value up to number of bins in raw yield histogram. Exit.")
+
     if cfg["central_efficiency"]["computerawfrac"]:
         infile_name = os.path.join(cfg["central_efficiency"]["inputdir"], cfg["central_efficiency"]["inputfile"])
         infile_central_eff = ROOT.TFile.Open(infile_name)
@@ -80,6 +86,8 @@ def main(config):
     hist_covariance = hist_rawy[0].Clone("hCovPromptNonPrompt")
     hist_corrfrac_prompt = hist_rawy[0].Clone("hCorrFracPrompt")
     hist_corrfrac_nonprompt = hist_rawy[0].Clone("hCorrFracNonPrompt")
+    for histo in hist_corry_prompt, hist_corry_nonprompt, hist_covariance, hist_corrfrac_prompt, hist_corrfrac_nonprompt:
+        histo.Reset()
     hist_corry_prompt.GetYaxis().SetTitle("corrected yields prompt")
     hist_corry_nonprompt.GetYaxis().SetTitle("corrected yields non-prompt")
     hist_covariance.GetYaxis().SetTitle("#sigma(prompt, non-prompt)")
@@ -113,6 +121,8 @@ def main(config):
     if cfg["central_efficiency"]["computerawfrac"]:
         hist_frac_raw_prompt = hist_rawy[0].Clone("hRawFracPrompt")
         hist_frac_raw_nonprompt = hist_rawy[0].Clone("hRawFracNonPrompt")
+        for histo in hist_frac_raw_prompt, hist_frac_raw_nonprompt:
+            histo.Reset()
         hist_frac_raw_prompt.GetYaxis().SetTitle("raw fraction prompt")
         hist_frac_raw_nonprompt.GetYaxis().SetTitle("raw fraction non-prompt")
         set_object_style(
@@ -129,12 +139,20 @@ def main(config):
             markerstyle=ROOT.kFullSquare,
         )
 
-    output = ROOT.TFile(os.path.join(cfg["output"]["directory"], cfg["output"]["file"]), "recreate")
+    pt_bin_to_process_name_suffix = ""
+    if pt_bin_to_process != -1:
+        pt_bin_to_process_name_suffix = "_bin_" + str(pt_bin_to_process)
+
+    output_name_template = cfg['output']['file'].replace(".root", "") + pt_bin_to_process_name_suffix + ".root"
+    output = ROOT.TFile(os.path.join(cfg["output"]["directory"], output_name_template), "recreate")
     n_sets = len(hist_rawy)
     pt_axis_title = hist_rawy[0].GetXaxis().GetTitle()
     for ipt in range(hist_rawy[0].GetNbinsX()):
+        if pt_bin_to_process !=-1 and ipt+1 != pt_bin_to_process:
+            continue
         pt_min = hist_rawy[0].GetXaxis().GetBinLowEdge(ipt + 1)
         pt_max = hist_rawy[0].GetXaxis().GetBinUpEdge(ipt + 1)
+        print(f"\n\nINFO: processing pt range {ipt+1} from {pt_min} to {pt_max} {pt_axis_title}")
 
         rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp = (np.zeros(n_sets) for _ in range(6))
         for iset, (hrawy, heffp, heffnp) in enumerate(zip(hist_rawy, hist_effp, hist_effnp)):
@@ -145,6 +163,14 @@ def main(config):
             unc_effp[iset] = heffp.GetBinError(ipt + 1)
             unc_effnp[iset] = heffnp.GetBinError(ipt + 1)
 
+        if cfg["minimisation"]["correlated"]:
+            if not (np.all(rawy[1:] > rawy[:-1]) or np.all(rawy[1:] < rawy[:-1])):
+                print("WARNING! main(): the raw yield vector is not monotonous. Check the input for stability.")
+                print(f"raw yield vector elements = {rawy}\n")
+            if not (np.all(unc_rawy[1:] > unc_rawy[:-1]) or np.all(unc_rawy[1:] < unc_rawy[:-1])):
+                print("WARNING! main(): the raw yield uncertainties vector is not monotonous. Check the input for stability.")
+                print(f"raw yield uncertainties vector elements = {unc_rawy}\n")
+
         minimiser = CutVarMinimiser(rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp)
         minimiser.minimise_system(cfg["minimisation"]["correlated"])
 
@@ -180,6 +206,12 @@ def main(config):
         for _, hist in histos_rawy.items():
             hist.Write()
 
+        canv_unc, histos_unc, leg_unc = minimiser.plot_uncertainties(f"_pt{pt_min}_{pt_max}", hist_bin_title)
+        output.cd()
+        canv_unc.Write()
+        for _, hist in histos_unc.items():
+            hist.Write()
+
         canv_eff, histos_eff, leg_e = minimiser.plot_efficiencies(f"_pt{pt_min}_{pt_max}", hist_bin_title)
         output.cd()
         canv_eff.Write()
@@ -208,13 +240,16 @@ def main(config):
         canv_combined.cd(4)
         canv_cov.DrawClonePad()
 
-        output_name_rawy_pdf = f"Distr_{cfg['output']['file'].replace('.root', '.pdf')}"
-        output_name_eff_pdf = f"Eff_{cfg['output']['file'].replace('.root', '.pdf')}"
-        output_name_frac_pdf = f"Frac_{cfg['output']['file'].replace('.root', '.pdf')}"
-        output_name_covmat_pdf = f"CovMatrix_{cfg['output']['file'].replace('.root', '.pdf')}"
-        output_name_pdf = f"{cfg['output']['file'].replace('.root', '.pdf')}"
+        output_name_template = output_name_template.replace('.root', '.pdf')
+
+        output_name_rawy_pdf = f"Distr_{output_name_template}"
+        output_name_eff_pdf = f"Eff_{output_name_template}"
+        output_name_frac_pdf = f"Frac_{output_name_template}"
+        output_name_covmat_pdf = f"CovMatrix_{output_name_template}"
+        output_name_unc_pdf = f"Unc_{output_name_template}"
+        output_name_pdf = f"{output_name_template}"
 
-        if hist_rawy[0].GetNbinsX() == 1:
+        if hist_rawy[0].GetNbinsX() == 1 or pt_bin_to_process != -1:
             print_bracket = ""
         elif ipt == 0:
             print_bracket = "("
@@ -227,6 +262,7 @@ def main(config):
         canv_frac.Print(f"{os.path.join(cfg['output']['directory'], output_name_frac_pdf)}{print_bracket}")
         canv_cov.Print(f"{os.path.join(cfg['output']['directory'], output_name_covmat_pdf)}{print_bracket}")
         canv_combined.Print(f"{os.path.join(cfg['output']['directory'], output_name_pdf)}{print_bracket}")
+        canv_unc.Print(f"{os.path.join(cfg['output']['directory'], output_name_unc_pdf)}{print_bracket}")
 
     output.cd()
     hist_corry_prompt.Write()

PWGHF/D2H/Macros/cut_variation.py

@@ -136,7 +136,7 @@ def minimise_system(self, correlated=True, precision=1.0e-8, max_iterations=100)
         self.m_eff = np.matrix(self.m_eff)
         m_corr_yields_old = np.zeros(shape=(2, 1))
 
-        for _ in range(max_iterations):
+        for iteration in range(max_iterations):
             for i_row, (rw_unc_row, effp_unc_row, effnp_unc_row) in enumerate(
                 zip(self.unc_raw_yields, self.unc_eff_prompt, self.unc_eff_nonprompt)
             ):
@@ -189,6 +189,13 @@ def minimise_system(self, correlated=True, precision=1.0e-8, max_iterations=100)
 
             m_corr_yields_old = np.copy(self.m_corr_yields)
 
+        print(f"INFO: number of processed iterations = {iteration+1}\n")
+        if correlated:
+            m_cov_sets_diag = np.diag(self.m_cov_sets)
+            if not (np.all(m_cov_sets_diag[1:] > m_cov_sets_diag[:-1]) or np.all(m_cov_sets_diag[1:] < m_cov_sets_diag[:-1])):
+                print("WARNING! minimise_system(): the residual vector uncertainties elements are not monotonous. Check the input for stability.")
+                print(f"residual vector uncertainties elements = {np.sqrt(m_cov_sets_diag)}\n")
+
         # chi2
         self.chi_2 = float(np.transpose(self.m_res) * self.m_weights * self.m_res)
 
@@ -432,6 +439,8 @@ def plot_result(self, suffix="", title=""):
         -----------------------------------------------------
         - suffix: str
             suffix to be added in the name of the output objects
+        - title: str
+            title to be written at the top margin of the output objects
 
         Returns
         -----------------------------------------------------
@@ -553,6 +562,8 @@ def plot_cov_matrix(self, correlated=True, suffix="", title=""):
             correlation between cut sets
         - suffix: str
             suffix to be added in the name of the output objects
+        - title: str
+            title to be written at the top margin of the output objects
 
         Returns
         -----------------------------------------------------
@@ -612,6 +623,8 @@ def plot_efficiencies(self, suffix="", title=""):
         -----------------------------------------------------
         - suffix: str
             suffix to be added in the name of the output objects
+        - title: str
+            title to be written at the top margin of the output objects
 
         Returns
         -----------------------------------------------------
@@ -701,6 +714,8 @@ def plot_fractions(self, suffix="", title=""):
         -----------------------------------------------------
         - suffix: str
             suffix to be added in the name of the output objects
+        - title: str
+            title to be written at the top margin of the output objects
 
         Returns
         -----------------------------------------------------
@@ -774,3 +789,84 @@ def plot_fractions(self, suffix="", title=""):
         canvas.Update()
 
         return canvas, histos, leg
+
+    # pylint: disable=no-member
+    def plot_uncertainties(self, suffix="", title=""):
+        """
+        Helper function to plot uncertainties as a function of cut set
+
+        Parameters
+        -----------------------------------------------------
+        - suffix: str
+            suffix to be added in the name of the output objects
+        - title: str
+            title to be written at the top margin of the output objects
+
+        Returns
+        -----------------------------------------------------
+        - canvas: ROOT.TCanvas
+            canvas with plot
+        - histos: dict
+            dictionary of ROOT.TH1F with uncertainties distributions for
+            raw yield and residual vector
+        - leg: ROOT.TLegend
+            needed otherwise it is destroyed
+        """
+
+        set_global_style(padleftmargin=0.16, padbottommargin=0.12, padtopmargin=0.075, titleoffsety=1.6)
+
+        hist_raw_yield_unc = ROOT.TH1F(
+            f"hRawYieldUncVsCut{suffix}",
+            ";cut set;runc.",
+            self.n_sets,
+            -0.5,
+            self.n_sets - 0.5,
+        )
+
+        hist_residual_unc = ROOT.TH1F(
+            f"hResidualUncVsCut{suffix}",
+            ";cut set;unc.",
+            self.n_sets,
+            -0.5,
+            self.n_sets - 0.5,
+        )
+
+        m_cov_sets_diag = np.diag(self.m_cov_sets)
+        m_cov_sets_diag = np.sqrt(m_cov_sets_diag)
+
+        for i_bin, (unc_rawy, unc_res) in enumerate(zip(self.unc_raw_yields, m_cov_sets_diag)):
+            hist_raw_yield_unc.SetBinContent(i_bin + 1, unc_rawy)
+            hist_residual_unc.SetBinContent(i_bin+1, unc_res)
+
+        set_object_style(hist_raw_yield_unc, color=ROOT.kRed + 1, fillstyle=0)
+        set_object_style(hist_residual_unc, color=ROOT.kAzure + 4, fillstyle=0)
+
+        canvas = ROOT.TCanvas(f"cUncVsCut{suffix}", "", 500, 500)
+        canvas.DrawFrame(
+            -0.5,
+            0.0,
+            self.n_sets - 0.5,
+            hist_residual_unc.GetMaximum() * 1.2,
+            ";cut set;unc.",
+        )
+        leg = ROOT.TLegend(0.6, 0.75, 0.8, 0.85)
+        leg.SetBorderSize(0)
+        leg.SetFillStyle(0)
+        leg.SetTextSize(0.04)
+        leg.AddEntry(hist_raw_yield_unc, "raw yield", "l")
+        leg.AddEntry(hist_residual_unc, "residual vector", "l")
+        leg.Draw()
+        hist_raw_yield_unc.Draw("histsame")
+        hist_residual_unc.Draw("histsame")
+        tex = ROOT.TLatex()
+        tex.SetTextSize(0.04)
+        tex.DrawLatexNDC(0.05, 0.95, title)
+        canvas.Modified()
+        canvas.Update()
+
+        histos = {
+            "rawy": hist_raw_yield_unc,
+            "residual": hist_residual_unc,
+        }
+
+        return canvas, histos, leg