[PWGHF] Cut variation: add efficiency-related functional

PWGHF/D2H/Macros/compute_fraction_cutvar.py

@@ -11,13 +11,13 @@
 import json
 import os
 import sys
+from enum import IntEnum, auto
 
 import numpy as np  # pylint: disable=import-error
 import ROOT  # pylint: disable=import-error
-from enum import IntEnum, auto
+
 sys.path.insert(0, '..')
-from cut_variation import CutVarMinimiser
-from cut_variation import MinimisationStatus
+from cut_variation import CutVarMinimiser, MinimisationStatus
 from style_formatter import set_object_style
 
 # pylint: disable=no-member,too-many-locals,too-many-statements
@@ -28,6 +28,7 @@ class PlotType(IntEnum):
     Frac = auto()
     Cov = auto()
     Unc = auto()
+    RelUnc = auto()
     N = auto()
 
 class ObjectToSave(IntEnum):
@@ -54,6 +55,10 @@ def main(config):
     with open(config, encoding="utf8") as fil:
         cfg = json.load(fil)
 
+    zero_eff_unc = cfg.get("zero_eff_unc", False)
+    effp_shift_nsigma = cfg.get("effp_shift_nsigma", 0.0)
+    effnp_shift_nsigma = cfg.get("effnp_shift_nsigma", 0.0)
+
     hist_rawy, hist_effp, hist_effnp = ([] for _ in range(3))
     for filename_rawy, filename_eff in zip(cfg["rawyields"]["inputfiles"], cfg["efficiencies"]["inputfiles"]):
         infile_rawy = ROOT.TFile.Open(os.path.join(cfg["rawyields"]["inputdir"], filename_rawy))
@@ -75,6 +80,12 @@ def main(config):
             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))
+            if zero_eff_unc:
+                hist_effp[-1].SetBinError(i_bin, 0.0)
+                hist_effnp[-1].SetBinError(i_bin, 0.0)
         infile_eff.Close()
 
     pt_bin_to_process = cfg.get("pt_bin_to_process", -1)
@@ -89,13 +100,15 @@ def main(config):
     is_draw_title[PlotType.Frac] = cfg.get("is_draw_title", {}).get("frac", False)
     is_draw_title[PlotType.Cov] = cfg.get("is_draw_title", {}).get("cov", False)
     is_draw_title[PlotType.Unc] = cfg.get("is_draw_title", {}).get("unc", True)
+    is_draw_title[PlotType.RelUnc] = cfg.get("is_draw_title", {}).get("relunc", True)
 
     is_save_canvas_as_macro = [False] * PlotType.N
     is_save_canvas_as_macro[PlotType.Rawy] = cfg.get("is_save_canvas_as_macro", {}).get("rawy", False)
     is_save_canvas_as_macro[PlotType.Eff] = cfg.get("is_save_canvas_as_macro", {}).get("eff", False)
     is_save_canvas_as_macro[PlotType.Frac] = cfg.get("is_save_canvas_as_macro", {}).get("frac", False)
     is_save_canvas_as_macro[PlotType.Cov] = cfg.get("is_save_canvas_as_macro", {}).get("cov", False)
     is_save_canvas_as_macro[PlotType.Unc] = cfg.get("is_save_canvas_as_macro", {}).get("unc", False)
+    is_save_canvas_as_macro[PlotType.RelUnc] = cfg.get("is_save_canvas_as_macro", {}).get("relunc", False)
 
     is_save_to_root_file = [False] * ObjectToSave.N
     is_save_to_root_file[ObjectToSave.Canvas] = cfg.get("is_save_to_root_file", {}).get("canvas", True)
@@ -205,14 +218,16 @@ def main(config):
         )
 
     pt_bin_to_process_name_suffix = ""
-    if pt_bin_to_process != -1: pt_bin_to_process_name_suffix = "_bin_" + str(pt_bin_to_process)
+    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
+        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)
@@ -237,9 +252,9 @@ def main(config):
                 print("\0\33[33mWARNING! main(): the raw yield uncertainties vector is not monotonous. Check the input for stability.\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(f"\33[31mFatal error: the prompt efficiency vector is not monotonous. Check the input. Exit.\33[0m")
+                sys.exit("\33[31mFatal error: the prompt efficiency vector is not monotonous. Check the input. Exit.\33[0m")
             if not (np.all(effnp[1:] > effnp[:-1]) or np.all(effnp[1:] < effnp[:-1])):
-                sys.exit(f"\33[31mFatal error: the nonprompt efficiency vector is not monotonous. Check the input. Exit.\33[0m")
+                sys.exit("\33[31mFatal error: the nonprompt efficiency vector is not monotonous. Check the input. Exit.\33[0m")
 
         minimiser = CutVarMinimiser(rawy, effp, effnp, unc_rawy, unc_effp, unc_effnp)
         status = minimiser.minimise_system(cfg["minimisation"]["correlated"])
@@ -278,47 +293,69 @@ def main(config):
             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)
+            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.Canvas]:
+                canv_rawy.Write()
             if is_save_to_root_file[ObjectToSave.RawYield]:
-                for _, hist in histos_rawy.items():
+                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")
+            if is_save_canvas_as_macro[PlotType.Rawy]:
+                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():
+                    hist.Write()
+            if is_save_canvas_as_macro[PlotType.Unc]:
+                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)
             output.cd()
-            if is_save_to_root_file[ObjectToSave.Canvas]: canv_unc.Write()
+            if is_save_to_root_file[ObjectToSave.Canvas]:
+                canv_rel_unc.Write()
             if is_save_to_root_file[ObjectToSave.Uncertainty]:
-                for _, hist in histos_unc.items():
+                for _, hist in histos_rel_unc.values():
                     hist.Write()
-            if is_save_canvas_as_macro[PlotType.Unc]: canv_unc.SaveAs(f"canv_unc_{ipt+1}.C")
+            if is_save_canvas_as_macro[PlotType.RelUnc]:
+                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)
+            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.Canvas]:
+                canv_eff.Write()
             if is_save_to_root_file[ObjectToSave.Efficiency]:
-                for _, hist in histos_eff.items():
+                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")
+            if is_save_canvas_as_macro[PlotType.Eff]:
+                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)
+            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.Canvas]:
+                canv_frac.Write()
             if is_save_to_root_file[ObjectToSave.Fraction]:
-                for _, hist in histos_frac.items():
+                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")
+            if is_save_canvas_as_macro[PlotType.Frac]:
+                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)
             output.cd()
-            if is_save_to_root_file[ObjectToSave.Canvas]: canv_cov.Write()
-            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")
+            if is_save_to_root_file[ObjectToSave.Canvas]:
+                canv_cov.Write()
+            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")
         else:
             print(f"Minimization for pT {pt_min}, {pt_max} not successful")
             hist_minimisation_status.SetBinContent(ipt + 1, MinimisationStatus.Fail)
@@ -327,6 +364,7 @@ def main(config):
             canv_frac = ROOT.TCanvas("c_frac_minimization_error", "Minimization error", 500, 500)
             canv_cov = ROOT.TCanvas("c_conv_minimization_error", "Minimization error", 500, 500)
             canv_unc = ROOT.TCanvas("c_unc_minimization_error", "Minimization error", 500, 500)
+            canv_rel_unc = ROOT.TCanvas("c_rel_unc_minimization_error", "Minimization error", 500, 500)
 
         canv_combined = ROOT.TCanvas(f"canv_combined_{ipt}", "", 1000, 1000)
         canv_combined.Divide(2, 2)
@@ -346,6 +384,7 @@ def main(config):
         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_rel_unc_pdf = f"RelUnc_{output_name_template}"
         output_name_pdf = f"{output_name_template}"
 
         if hist_rawy[0].GetNbinsX() == 1 or pt_bin_to_process != -1:
@@ -362,6 +401,7 @@ def main(config):
         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}")
+        canv_rel_unc.Print(f"{os.path.join(cfg['output']['directory'], output_name_rel_unc_pdf)}{print_bracket}")
 
     output.cd()
     if is_save_to_root_file[ObjectToSave.CorrectedYield]:

PWGHF/D2H/Macros/config_cutvar_example.json

@@ -62,14 +62,16 @@
         "frac": false,
         "eff": false,
         "cov": false,
-        "unc": true
+        "unc": true,
+        "relunc": true
     },
     "is_save_canvas_as_macro": {
         "rawy": false,
         "frac": false,
         "eff": false,
         "cov": false,
-        "unc": false
+        "unc": false,
+        "relunc": false
     },
     "is_save_to_root_file": {
         "canvas": true,
@@ -95,5 +97,8 @@
     "output": {
         "directory": ".",
         "file": "CutVarDplus_pp13TeV_MB.root"
-    }
+    },
+    "zero_eff_unc": false,
+    "effp_shift_nsigma": 0,
+    "effnp_shift_nsigma": 0
 }