[PWGHF] Cut variation: demote efficiency vector non-monotonicity to warning; fix bug and format (#18447)

Author: lubynets

PWGHF/D2H/Macros/compute_fraction_cutvar.py

@@ -16,12 +16,13 @@
 import numpy as np  # pylint: disable=import-error
 import ROOT  # pylint: disable=import-error
 
-sys.path.insert(0, '..')
+sys.path.insert(0, "..")
 from cut_variation import CutVarMinimiser, MinimisationStatus
 from style_formatter import set_object_style
 
 # pylint: disable=no-member,too-many-locals,too-many-statements
 
+
 class PlotType(IntEnum):
     Rawy = 0
     Eff = auto()
@@ -31,6 +32,7 @@ class PlotType(IntEnum):
     RelUnc = auto()
     N = auto()
 
+
 class ObjectToSave(IntEnum):
     Canvas = 0
     RawYield = auto()
@@ -44,6 +46,7 @@ class ObjectToSave(IntEnum):
     MinimisationStatus = auto()
     N = auto()
 
+
 def main(config):
     """
     Main function
@@ -65,7 +68,7 @@ def main(config):
         hist_rawy_name = cfg["rawyields"]["histoname"]
         hist_rawy.append(infile_rawy.Get(hist_rawy_name))
         if not isinstance(hist_rawy[-1], ROOT.TH1):
-            sys.exit(f"\33[31mFatal error: Histogram with raw yield \"{hist_rawy_name}\" is absent. Exit.\33[0m")
+            sys.exit(f'\33[31mFatal error: Histogram with raw yield "{hist_rawy_name}" is absent. Exit.\33[0m')
         hist_rawy[-1].SetDirectory(0)
         infile_rawy.Close()
 
@@ -75,14 +78,22 @@ def main(config):
         hist_effp.append(infile_eff.Get(hist_effp_name))
         hist_effnp.append(infile_eff.Get(hist_effnp_name))
         if not isinstance(hist_effp[-1], ROOT.TH1):
-            sys.exit(f"\33[31mFatal error: Histogram with efficiency for prompt \"{hist_effp_name}\" is absent. Exit.\33[0m")
+            sys.exit(
+                f'\33[31mFatal error: Histogram with efficiency for prompt "{hist_effp_name}" is absent. Exit.\33[0m'
+            )
         if not isinstance(hist_effnp[-1], ROOT.TH1):
-            sys.exit(f"\33[31mFatal error: Histogram with efficiency for nonprompt \"{hist_effnp}\" is absent. Exit.\33[0m")
+            sys.exit(
+                f'\33[31mFatal error: Histogram with efficiency for nonprompt "{hist_effnp}" is absent. Exit.\33[0m'
+            )
         hist_effp[-1].SetDirectory(0)
         hist_effnp[-1].SetDirectory(0)
         for i_bin in range(1, hist_effp[-1].GetNbinsX() + 1):
-            hist_effp[-1].SetBinContent(i_bin, hist_effp[-1].GetBinContent(i_bin) + effp_shift_nsigma*hist_effp[-1].GetBinError(i_bin))
-            hist_effnp[-1].SetBinContent(i_bin, hist_effnp[-1].GetBinContent(i_bin) + effnp_shift_nsigma*hist_effnp[-1].GetBinError(i_bin))
+            hist_effp[-1].SetBinContent(
+                i_bin, hist_effp[-1].GetBinContent(i_bin) + effp_shift_nsigma * hist_effp[-1].GetBinError(i_bin)
+            )
+            hist_effnp[-1].SetBinContent(
+                i_bin, hist_effnp[-1].GetBinContent(i_bin) + effnp_shift_nsigma * hist_effnp[-1].GetBinError(i_bin)
+            )
             if zero_eff_unc:
                 hist_effp[-1].SetBinError(i_bin, 0.0)
                 hist_effnp[-1].SetBinError(i_bin, 0.0)
@@ -92,7 +103,9 @@ def main(config):
     if not isinstance(pt_bin_to_process, int):
         sys.exit("\33[31mFatal 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("\33[31mFatal error: pt_bin_to_process must be a positive value up to number of bins in raw yield histogram. Exit.")
+        sys.exit(
+            "\33[31mFatal error: pt_bin_to_process must be a positive value up to number of bins in raw yield histogram. Exit."
+        )
 
     is_draw_title = [False] * PlotType.N
     is_draw_title[PlotType.Rawy] = cfg.get("is_draw_title", {}).get("rawy", True)
@@ -117,22 +130,32 @@ def main(config):
     is_save_to_root_file[ObjectToSave.Efficiency] = cfg.get("is_save_to_root_file", {}).get("efficiency", True)
     is_save_to_root_file[ObjectToSave.Fraction] = cfg.get("is_save_to_root_file", {}).get("fraction", True)
     is_save_to_root_file[ObjectToSave.CorrectedYield] = cfg.get("is_save_to_root_file", {}).get("corrected_yield", True)
-    is_save_to_root_file[ObjectToSave.CorrelationMatrix] = cfg.get("is_save_to_root_file", {}).get("correlation_matrix", True)
+    is_save_to_root_file[ObjectToSave.CorrelationMatrix] = cfg.get("is_save_to_root_file", {}).get(
+        "correlation_matrix", True
+    )
     is_save_to_root_file[ObjectToSave.Covariance] = cfg.get("is_save_to_root_file", {}).get("covariance", True)
-    is_save_to_root_file[ObjectToSave.CorrectedFraction] = cfg.get("is_save_to_root_file", {}).get("corrected_fraction", True)
-    is_save_to_root_file[ObjectToSave.MinimisationStatus] = cfg.get("is_save_to_root_file", {}).get("minimisation_status", True)
+    is_save_to_root_file[ObjectToSave.CorrectedFraction] = cfg.get("is_save_to_root_file", {}).get(
+        "corrected_fraction", True
+    )
+    is_save_to_root_file[ObjectToSave.MinimisationStatus] = cfg.get("is_save_to_root_file", {}).get(
+        "minimisation_status", True
+    )
 
     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)
         hist_central_effp_name = cfg["central_efficiency"]["histonames"]["prompt"]
         hist_central_effp = infile_central_eff.Get(hist_central_effp_name)
         if not isinstance(hist_central_effp[-1], ROOT.TH1):
-            sys.exit(f"\33[31mFatal error: Histogram with central efficiency for prompt \"{hist_central_effp_name}\" is absent. Exit.\33[0m")
+            sys.exit(
+                f'\33[31mFatal error: Histogram with central efficiency for prompt "{hist_central_effp_name}" is absent. Exit.\33[0m'
+            )
         hist_central_effnp_name = cfg["central_efficiency"]["histonames"]["nonprompt"]
         hist_central_effnp = infile_central_eff.Get(hist_central_effnp_name)
         if not isinstance(hist_central_effnp[-1], ROOT.TH1):
-            sys.exit(f"\33[31mFatal error: Histogram with central efficiency for nonprompt \"{hist_central_effnp_name}\" is absent. Exit.\33[0m")
+            sys.exit(
+                f'\33[31mFatal error: Histogram with central efficiency for nonprompt "{hist_central_effnp_name}" is absent. Exit.\33[0m'
+            )
         hist_central_effp.SetDirectory(0)
         hist_central_effnp.SetDirectory(0)
         infile_central_eff.Close()
@@ -150,8 +173,18 @@ def main(config):
     hist_corrfrac_prompt = hist_rawy[0].Clone("hCorrFracPrompt")
     hist_corrfrac_nonprompt = hist_rawy[0].Clone("hCorrFracNonPrompt")
     hist_minimisation_status = hist_rawy[0].Clone("hMinimizationStatus")
-    hist_red_chi2  = hist_rawy[0].Clone("hChi2OverNdf")
-    for histo in hist_corry_prompt, hist_corry_nonprompt, hist_covariance_pnp, hist_covariance_pp, hist_covariance_npnp, hist_corrfrac_prompt, hist_corrfrac_nonprompt, hist_minimisation_status, hist_red_chi2:
+    hist_red_chi2 = hist_rawy[0].Clone("hChi2OverNdf")
+    for histo in (
+        hist_corry_prompt,
+        hist_corry_nonprompt,
+        hist_covariance_pnp,
+        hist_covariance_pp,
+        hist_covariance_npnp,
+        hist_corrfrac_prompt,
+        hist_corrfrac_nonprompt,
+        hist_minimisation_status,
+        hist_red_chi2,
+    ):
         histo.Reset()
     hist_corry_prompt.GetYaxis().SetTitle("corrected yields prompt")
     hist_corry_nonprompt.GetYaxis().SetTitle("corrected yields non-prompt")
@@ -164,7 +197,7 @@ def main(config):
     hist_red_chi2.GetYaxis().SetTitle("#chi^{2}/ndf")
     hist_minimisation_status_title = ""
     for min_status in MinimisationStatus:
-        hist_minimisation_status_title += (str(min_status.value) + " = " + min_status.name + ", ")
+        hist_minimisation_status_title += str(min_status.value) + " = " + min_status.name + ", "
     hist_minimisation_status_title = hist_minimisation_status_title[:-2]
     hist_minimisation_status.SetTitle(hist_minimisation_status_title)
     set_object_style(
@@ -221,17 +254,17 @@ def main(config):
     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_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:
+        if pt_bin_to_process != -1 and ipt + 1 != pt_bin_to_process:
             continue
         all_vectors_monotonous = MinimisationStatus.Success
         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}")
+        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)):
@@ -245,16 +278,24 @@ def main(config):
         if cfg["minimisation"]["correlated"]:
             if not (np.all(rawy[1:] > rawy[:-1]) or np.all(rawy[1:] < rawy[:-1])):
                 all_vectors_monotonous = MinimisationStatus.MonotonyViolation
-                print("\0\33[33mWARNING! main(): the raw yield vector is not monotonous. Check the input for stability.\0\33[0m")
+                print(
+                    "\0\33[33m\nWARNING! main(): the raw yield vector is not monotonous. Check the input for stability.\n\0\33[0m"
+                )
                 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])):
                 all_vectors_monotonous = MinimisationStatus.MonotonyViolation
-                print("\0\33[33mWARNING! main(): the raw yield uncertainties vector is not monotonous. Check the input for stability.\0\33[0m")
+                print(
+                    "\0\33[33m\nWARNING! main(): the raw yield uncertainties vector is not monotonous. Check the input for stability.\n\0\33[0m"
+                )
                 print(f"raw yield uncertainties vector elements = {unc_rawy}\n")
             if not (np.all(effp[1:] > effp[:-1]) or np.all(effp[1:] < effp[:-1])):
-                sys.exit("\33[31mFatal error: the prompt efficiency vector is not monotonous. Check the input. Exit.\33[0m")
+                all_vectors_monotonous = MinimisationStatus.MonotonyViolation
+                print("\0\33[33m\nWARNING! the prompt efficiency vector is not monotonous. Check the input.\0\33[0m")
+                print(f"prompt efficiency vector elements = {effp}\n")
             if not (np.all(effnp[1:] > effnp[:-1]) or np.all(effnp[1:] < effnp[:-1])):
-                sys.exit("\33[31mFatal error: the nonprompt efficiency vector is not monotonous. Check the input. Exit.\33[0m")
+                all_vectors_monotonous = MinimisationStatus.MonotonyViolation
+                print("\0\33[33m\nWARNING! the nonprompt efficiency vector is not monotonous. Check the input.\0\33[0m")
+                print(f"nonprompt efficiency vector elements = {effnp}\n")
 
         minimiser = CutVarMinimiser(rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp)
         status = minimiser.minimise_system(cfg["minimisation"]["correlated"])
@@ -290,62 +331,66 @@ def main(config):
                 hist_frac_raw_nonprompt.SetBinContent(ipt + 1, raw_frac_nonprompt[0])
                 hist_frac_raw_nonprompt.SetBinError(ipt + 1, raw_frac_nonprompt[1])
 
-            hist_bin_title = f"bin # {ipt+1}; {pt_axis_title}#in ({pt_min}; {pt_max})"
+            hist_bin_title = f"bin # {ipt + 1}; {pt_axis_title}#in ({pt_min}; {pt_max})"
 
             hist_bin_title_rawy = hist_bin_title if is_draw_title[PlotType.Rawy] else ""
             canv_rawy, histos_rawy, _leg_r = minimiser.plot_result(f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_rawy)
             output.cd()
             if is_save_to_root_file[ObjectToSave.Canvas]:
                 canv_rawy.Write()
             if is_save_to_root_file[ObjectToSave.RawYield]:
-                for _, hist in histos_rawy.values():
+                for hist in histos_rawy.values():
                     hist.Write()
             if is_save_canvas_as_macro[PlotType.Rawy]:
-                canv_rawy.SaveAs(f"canv_rawy_{ipt+1}.C")
+                canv_rawy.SaveAs(f"canv_rawy_{ipt + 1}.C")
 
             hist_bin_title_unc = hist_bin_title if is_draw_title[PlotType.Unc] else ""
-            canv_unc, histos_unc, _leg_unc = minimiser.plot_uncertainties(f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_unc)
+            canv_unc, histos_unc, _leg_unc = minimiser.plot_uncertainties(
+                f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_unc
+            )
             output.cd()
             if is_save_to_root_file[ObjectToSave.Canvas]:
                 canv_unc.Write()
             if is_save_to_root_file[ObjectToSave.Uncertainty]:
-                for _, hist in histos_unc.values():
+                for hist in histos_unc.values():
                     hist.Write()
             if is_save_canvas_as_macro[PlotType.Unc]:
-                canv_unc.SaveAs(f"canv_unc_{ipt+1}.C")
+                canv_unc.SaveAs(f"canv_unc_{ipt + 1}.C")
 
             hist_bin_title_rel_unc = hist_bin_title if is_draw_title[PlotType.RelUnc] else ""
-            canv_rel_unc, histos_rel_unc, _leg_rel_unc = minimiser.plot_relative_uncertainties(f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_rel_unc)
+            canv_rel_unc, histos_rel_unc, _leg_rel_unc = minimiser.plot_relative_uncertainties(
+                f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_rel_unc
+            )
             output.cd()
             if is_save_to_root_file[ObjectToSave.Canvas]:
                 canv_rel_unc.Write()
             if is_save_to_root_file[ObjectToSave.Uncertainty]:
-                for _, hist in histos_rel_unc.values():
+                for hist in histos_rel_unc.values():
                     hist.Write()
             if is_save_canvas_as_macro[PlotType.RelUnc]:
-                canv_rel_unc.SaveAs(f"canv_rel_unc_{ipt+1}.C")
+                canv_rel_unc.SaveAs(f"canv_rel_unc_{ipt + 1}.C")
 
             hist_bin_title_eff = hist_bin_title if is_draw_title[PlotType.Eff] else ""
             canv_eff, histos_eff, _leg_e = minimiser.plot_efficiencies(f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_eff)
             output.cd()
             if is_save_to_root_file[ObjectToSave.Canvas]:
                 canv_eff.Write()
             if is_save_to_root_file[ObjectToSave.Efficiency]:
-                for _, hist in histos_eff.values():
+                for hist in histos_eff.values():
                     hist.Write()
             if is_save_canvas_as_macro[PlotType.Eff]:
-                canv_eff.SaveAs(f"canv_eff_{ipt+1}.C")
+                canv_eff.SaveAs(f"canv_eff_{ipt + 1}.C")
 
             hist_bin_title_frac = hist_bin_title if is_draw_title[PlotType.Frac] else ""
             canv_frac, histos_frac, _leg_f = minimiser.plot_fractions(f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_frac)
             output.cd()
             if is_save_to_root_file[ObjectToSave.Canvas]:
                 canv_frac.Write()
             if is_save_to_root_file[ObjectToSave.Fraction]:
-                for _, hist in histos_frac.values():
+                for hist in histos_frac.values():
                     hist.Write()
             if is_save_canvas_as_macro[PlotType.Frac]:
-                canv_frac.SaveAs(f"canv_frac_{ipt+1}.C")
+                canv_frac.SaveAs(f"canv_frac_{ipt + 1}.C")
 
             hist_bin_title_cov = hist_bin_title if is_draw_title[PlotType.Cov] else ""
             canv_cov, histo_cov = minimiser.plot_cov_matrix(True, f"_pt_{pt_min}_to_{pt_max}", hist_bin_title_cov)
@@ -355,7 +400,7 @@ def main(config):
             if is_save_to_root_file[ObjectToSave.CorrelationMatrix]:
                 histo_cov.Write()
             if is_save_canvas_as_macro[PlotType.Cov]:
-                canv_cov.SaveAs(f"canv_cov_{ipt+1}.C")
+                canv_cov.SaveAs(f"canv_cov_{ipt + 1}.C")
         else:
             print(f"Minimization for pT {pt_min}, {pt_max} not successful")
             hist_minimisation_status.SetBinContent(ipt + 1, MinimisationStatus.Fail)
@@ -377,7 +422,7 @@ def main(config):
         canv_combined.cd(4)
         canv_cov.DrawClonePad()
 
-        output_name_template = output_name_template.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}"