outsource creation of 2D hists to function

key4hep · Oct 1, 2024 · a494994 · a494994
1 parent ba64ad0
commit a494994
Showing 1 changed file with 21 additions and 30 deletions.
diff --git a/utils/material_plots_2D.py b/utils/material_plots_2D.py
@@ -18,6 +18,24 @@
 import ROOT
 
 
+def create_histogram(
+    name_and_title: str, argument_name_space: argparse.Namespace
+) -> ROOT.TH2F:
+    return ROOT.TH2F(
+        name_and_title,
+        name_and_title,
+        int(
+            (argument_name_space.angleMax - argument_name_space.angleMin)
+            / argument_name_space.angleBinning
+        ),
+        argument_name_space.angleMin,
+        argument_name_space.angleMax,
+        argument_name_space.nPhiBins,
+        -math.pi,
+        math.pi,
+    )
+
+
 def main():
     parser = argparse.ArgumentParser(description="Material Plotter")
     parser.add_argument("--inputFile", "-f", type=str, help="relative path to the input file")
@@ -76,36 +94,9 @@ def main():
 
     ROOT.gROOT.SetBatch(1)
 
-    h_x0 = ROOT.TH2F(
-        "h_x0",
-        "h_x0",
-        int((args.angleMax - args.angleMin) / args.angleBinning),
-        args.angleMin,
-        args.angleMax,
-        args.nPhiBins,
-        -math.pi,
-        math.pi,
-    )
-    h_lambda = ROOT.TH2F(
-        "h_lambda",
-        "h_lambda",
-        int((args.angleMax - args.angleMin) / args.angleBinning),
-        args.angleMin,
-        args.angleMax,
-        args.nPhiBins,
-        -math.pi,
-        math.pi,
-    )
-    h_depth = ROOT.TH2F(
-        "h_depth",
-        "h_depth",
-        int((args.angleMax - args.angleMin) / args.angleBinning),
-        args.angleMin,
-        args.angleMax,
-        args.nPhiBins,
-        -math.pi,
-        math.pi,
-    )
+    h_x0 = create_histogram("h_x0", args)
+    h_lambda = create_histogram("h_lambda", args)
+    h_depth = create_histogram("h_depth", args)
 
     for angleBinning, entry in enumerate(tree):
         nMat = entry.nMaterials