Merge pull request #253 from jeyserma/master

Add option for 2D and 3D histograms in final stage
HEP-FCC · Mar 15, 2023 · aa9de28 · aa9de28
2 parents d2c2ffa + de0d887
commit aa9de28
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/config/FCCAnalysisRun.py b/config/FCCAnalysisRun.py
@@ -874,8 +874,28 @@ def runFinal(rdfModule):
             histos = []
 
             for v in histoList:
-                model = ROOT.RDF.TH1DModel(v, ";{};".format(histoList[v]["title"]), histoList[v]["bin"], histoList[v]["xmin"], histoList[v]["xmax"])
-                histos.append(df_cut.Histo1D(model,histoList[v]["name"]))
+                if "name" in histoList[v]: # default 1D histogram
+                    model = ROOT.RDF.TH1DModel(v, ";{};".format(histoList[v]["title"]), histoList[v]["bin"], histoList[v]["xmin"], histoList[v]["xmax"])
+                    histos.append(df_cut.Histo1D(model,histoList[v]["name"]))
+                elif "cols" in histoList[v]: # multi dim histogram (1, 2 or 3D)
+                    cols = histoList[v]['cols']
+                    bins = histoList[v]['bins']
+                    bins_unpacked = tuple([i for sub in bins for i in sub])
+                    if len(bins) != len(cols):
+                        print ('----> Amount of columns should be equal to the amount of bin configs.')
+                        sys.exit(3)
+                    if len(cols) == 1:
+                        histos.append(df_cut.Histo1D((v, "", *bins_unpacked), cols[0]))
+                    elif len(cols) == 2:
+                        histos.append(df_cut.Histo2D((v, "", *bins_unpacked), cols[0], cols[1]))
+                    elif len(cols) == 3:
+                        histos.append(df_cut.Histo3D((v, "", *bins_unpacked), cols[0], cols[1], cols[2]))
+                    else:
+                        print ('----> Only 1, 2 or 3D histograms supported.')
+                        sys.exit(3)
+                else:
+                    print ('----> Error parsing the histogram config. Provide either name or cols.')
+                    sys.exit(3)
             histos_list.append(histos)
 
             if doTree:
@@ -1029,7 +1049,7 @@ def setup_run_parser(parser):
     publicOptions.add_argument("--rerunfailed", action='store_true', help="Rerun failed jobs", default=False)
     publicOptions.add_argument("--jobdir", help="Specify the batch job directory", type=str, default="output.root")
     publicOptions.add_argument("--eloglevel", help="Specify the RDataFrame ELogLevel", type=str, default="kUnset", choices = ['kUnset','kFatal','kError','kWarning','kInfo','kDebug'])
-
+    
     internalOptions = parser.add_argument_group('\033[4m\033[1m\033[91m Internal options, NOT FOR USERS\033[0m')
     internalOptions.add_argument("--batch", action='store_true', help="Submit on batch", default=False)
 

diff --git a/examples/FCCee/higgs/mH-recoil/mumu/analysis_final.py b/examples/FCCee/higgs/mH-recoil/mumu/analysis_final.py
@@ -42,4 +42,7 @@
     "leptonic_recoil_m_zoom4":{"name":"Zcand_recoil_m","title":"Z leptonic recoil [GeV]","bin":800,"xmin":120,"xmax":140},
     "leptonic_recoil_m_zoom5":{"name":"Zcand_recoil_m","title":"Z leptonic recoil [GeV]","bin":2000,"xmin":120,"xmax":140},
     "leptonic_recoil_m_zoom6":{"name":"Zcand_recoil_m","title":"Z leptonic recoil [GeV]","bin":100,"xmin":130.3,"xmax":132.5},
+    "mz_1D":{"cols":["Zcand_m"],"title":"m_{Z} [GeV]", "bins": [(40,80,100)]}, # 1D histogram (alternative syntax)
+    "mz_recoil_2D":{"cols":["Zcand_m", "Zcand_recoil_m"],"title":"m_{Z} - leptonic recoil [GeV]", "bins": [(40,80,100), (100,120,140)]}, # 2D histogram
+    "mz_recoil_3D":{"cols":["Zcand_m", "Zcand_recoil_m", "Zcand_recoil_m"],"title":"m_{Z} - leptonic recoil - leptonic recoil [GeV]", "bins": [(40,80,100), (100,120,140), (100,120,140)]}, # 3D histogram
 }