make is_draw_title and is_save_canvas_as_macro vectors

Author: lubynets

PWGHF/D2H/Macros/compute_fraction_cutvar.py

@@ -14,13 +14,21 @@
 
 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 style_formatter import set_object_style
 
 # pylint: disable=no-member,too-many-locals,too-many-statements
 
+class PlotType(IntEnum):
+    Rawy = 0
+    Eff = auto()
+    Frac = auto()
+    Cov = auto()
+    Unc = auto()
+    N = auto()
 
 def main(config):
     """
@@ -62,17 +70,19 @@ def main(config):
     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.")
 
-    is_draw_title_rawy = cfg.get("is_draw_title", {}).get("rawy", True)
-    is_draw_title_eff = cfg.get("is_draw_title", {}).get("eff", False)
-    is_draw_title_frac = cfg.get("is_draw_title", {}).get("frac", False)
-    is_draw_title_cov = cfg.get("is_draw_title", {}).get("cov", False)
-    is_draw_title_unc = cfg.get("is_draw_title", {}).get("unc", True)
+    is_draw_title = [False] * PlotType.N
+    is_draw_title[PlotType.Rawy] = cfg.get("is_draw_title", {}).get("rawy", True)
+    is_draw_title[PlotType.Eff] = cfg.get("is_draw_title", {}).get("eff", False)
+    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_save_canvas_as_macro_rawy = cfg.get("is_save_canvas_as_macro", {}).get("rawy", False)
-    is_save_canvas_as_macro_eff = cfg.get("is_save_canvas_as_macro", {}).get("eff", False)
-    is_save_canvas_as_macro_frac = cfg.get("is_save_canvas_as_macro", {}).get("frac", False)
-    is_save_canvas_as_macro_cov = cfg.get("is_save_canvas_as_macro", {}).get("cov", False)
-    is_save_canvas_as_macro_unc = cfg.get("is_save_canvas_as_macro", {}).get("unc", False)
+    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)
 
     if cfg["central_efficiency"]["computerawfrac"]:
         infile_name = os.path.join(cfg["central_efficiency"]["inputdir"], cfg["central_efficiency"]["inputfile"])
@@ -240,48 +250,48 @@ 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_rawy else ""
+            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()
             canv_rawy.Write()
             for _, hist in histos_rawy.items():
                 hist.Write()
-            if (is_save_canvas_as_macro_rawy):
+            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_unc else ""
+            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)
             output.cd()
             canv_unc.Write()
             for _, hist in histos_unc.items():
                 hist.Write()
-            if (is_save_canvas_as_macro_unc):
+            if (is_save_canvas_as_macro[PlotType.Unc]):
                 canv_unc.SaveAs(f"canv_unc_{ipt+1}.C")
 
-            hist_bin_title_eff = hist_bin_title if is_draw_title_eff else ""
+            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()
             canv_eff.Write()
             for _, hist in histos_eff.items():
                 hist.Write()
-            if (is_save_canvas_as_macro_eff):
+            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_frac else ""
+            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()
             canv_frac.Write()
             for _, hist in histos_frac.items():
                 hist.Write()
-            if (is_save_canvas_as_macro_frac):
+            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_cov else ""
+            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()
             canv_cov.Write()
             histo_cov.Write()
-            if (is_save_canvas_as_macro_cov):
+            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")

PWGHF/D2H/Macros/cut_variation.py

@@ -11,15 +11,16 @@
 
 import numpy as np  # pylint: disable=import-error
 import ROOT  # pylint: disable=import-error
+from enum import IntEnum, auto
 sys.path.insert(0, '..')
 from style_formatter import set_global_style, set_object_style
-from enum import IntEnum
 
 class MinimisationStatus(IntEnum):
     Undefined = 0
-    Success = 1
-    MonotonyViolation = 2
-    Fail = 3
+    Success = auto()
+    MonotonyViolation = auto()
+    Fail = auto()
+    N = auto()
 
 # pylint: disable=too-many-instance-attributes
 class CutVarMinimiser: