Configurability for material plots (#392)

* format material_(plots/scan)_2D.py * refactorization, parseArgs added to material_scan_2D.py, thetaRad as angleDef option added to material_plots_2D.py * outsource creation of 2D hists to function * fix x0max user interface Co-authored-by: Thomas Madlener <[email protected]> * Fix inputFile arg parsing * Remove leftover conflict markers * Complete switch back to x0max * Remove stray conflict marker remnants * bug fix and switch to os.smth to ease searching for usage * make create_histogram function more reusable --------- Co-authored-by: Thomas Madlener <[email protected]>
key4hep · Oct 25, 2024 · 1f414da · 1f414da
1 parent 60243ef
commit 1f414da
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 50 deletions.
diff --git a/utils/material_plots_2D.py b/utils/material_plots_2D.py
@@ -1,17 +1,47 @@
+"""
+This script must be called with python: 'python material_scan_2D.py --{argument} {value}'.
+The output files are saved in data/{outputDir}/name.suffix.
+If no outputDir is specified, it will be data/plots/name.suffix.
+"""
+
 from __future__ import print_function
+
 import argparse
 import math
-
-import sys, os
+import os
+import sys
+from pathlib import Path
 
 sys.path.append(os.path.expandvars("$FCCSW") + "/Examples/scripts")
 from plotstyle import FCCStyle
 import ROOT
 
 
+def create_histogram(
+    name_and_title: str,
+    angle_min: float,
+    angle_max: float,
+    angle_binning: float,
+    n_phi_bins: int,
+) -> ROOT.TH2F:
+    num_bins = int((angle_max - angle_min) / angle_binning)
+    return ROOT.TH2F(
+        name_and_title,
+        name_and_title,
+        num_bins,
+        angle_min,
+        angle_max,
+        n_phi_bins,
+        -math.pi,
+        math.pi,
+    )
+
+
 def main():
     parser = argparse.ArgumentParser(description="Material Plotter")
-    parser.add_argument("--fname", "-f", dest="fname", type=str, help="name of file to read")
+    parser.add_argument(
+        "--inputFile", "--fname", "-f", type=str, help="relative path to the input file"
+    )
     parser.add_argument(
         "--angleMin", dest="angleMin", default=6, type=float, help="minimum eta/theta/cosTheta"
     )
@@ -20,23 +50,27 @@ def main():
     )
     parser.add_argument(
         "--angleDef",
-        dest="angleDef",
         default="eta",
+        choices=["eta", "theta", "cosTheta", "thetaRad"],
         type=str,
-        help="angle definition to use: eta, theta or cosTheta, default: eta",
+        help="Angle definition to use: eta, theta, thetaRad or cosTheta; default: eta",
     )
     parser.add_argument(
         "--angleBinning",
         "-b",
-        dest="angleBinning",
         default=0.05,
         type=float,
-        help="eta/theta/cosTheta bin width",
+        help="Eta/theta/cosTheta bin width",
     )
+    parser.add_argument("--nPhiBins", default=100, type=int, help="Number of bins in phi")
+    parser.add_argument("--x0max", "-x", default=0.0, type=float, help="Max of x0")
     parser.add_argument(
-        "--nPhiBins", dest="nPhiBins", default=100, type=int, help="number of bins in phi"
+        "--outputDir",
+        "-o",
+        type=str,
+        default="plots",
+        help="Directory to store output files in",
     )
-    parser.add_argument("--x0max", "-x", dest="x0max", default=0.0, type=float, help="Max of x0")
     parser.add_argument(
         "--ignoreMats",
         "-i",
@@ -45,45 +79,25 @@ def main():
         default=[],
         help="List of materials that should be ignored",
     )
+
     args = parser.parse_args()
 
+    output_dir = Path("data") / args.outputDir
+    output_dir.mkdir(parents=True, exist_ok=True)  # Create the directory if it doesn't exist
+
     ROOT.gStyle.SetNumberContours(100)
 
-    f = ROOT.TFile.Open(args.fname, "read")
+    f = ROOT.TFile.Open(os.fspath(Path(args.inputFile).with_suffix(".root")), "read")
     tree = f.Get("materials")
-    histDict = {}
 
     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_x0 = create_histogram("h_x0", args.angleMin, args.angleMax, args.angleBinning, args.nPhiBins)
+    h_lambda = create_histogram(
+        "h_lambda", args.angleMin, args.angleMax, args.angleBinning, args.nPhiBins
     )
-    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_depth = create_histogram(
+        "h_depth", args.angleMin, args.angleMax, args.angleBinning, args.nPhiBins
     )
 
     for angleBinning, entry in enumerate(tree):
@@ -103,11 +117,11 @@ def main():
         h_lambda.Fill(tree.angle, tree.phi, entry_lambda)
         h_depth.Fill(tree.angle, tree.phi, entry_depth)
 
-    # go through the
+    # go through the plots
     plots = ["x0", "lambda", "depth"]
     histograms = [h_x0, h_lambda, h_depth]
     axis_titles = ["Material budget x/X_{0} [%]", "Number of #lambda", "Material depth [cm]"]
-    for i in range(len(plots)):
+    for i, plot in enumerate(plots):
         cv = ROOT.TCanvas("", "", 800, 600)
         cv.SetRightMargin(0.18)
         histograms[i].Draw("COLZ")
@@ -116,22 +130,25 @@ def main():
             title = "#eta"
         elif args.angleDef == "theta":
             title = "#theta [#circ]"
+        elif args.angleDef == "thetaRad":
+            title = "#theta [rad]"
         elif args.angleDef == "cosTheta":
             title = "cos(#theta)"
         histograms[i].GetXaxis().SetTitle(title)
         histograms[i].GetYaxis().SetTitle("#phi")
 
         histograms[i].GetZaxis().SetTitle(axis_titles[i])
 
-        if args.x0max != 0.0 and plots[i] == "x0":
+        if args.x0max != 0.0 and plot == "x0":
             histograms[i].SetMaximum(args.x0max)
 
         histograms[i].GetXaxis().SetRangeUser(args.angleMin, args.angleMax)
 
         ROOT.gStyle.SetPadRightMargin(0.5)
-        cv.Print(plots[i] + ".pdf")
-        cv.Print(plots[i] + ".png")
-        cv.SaveAs(plots[i] + ".root")
+        output_path = output_dir / plot
+        cv.Print(os.fspath(output_path.with_suffix(".pdf")))
+        cv.Print(os.fspath(output_path.with_suffix(".png")))
+        cv.SaveAs(os.fspath(output_path.with_suffix(".root")))
 
 
 if __name__ == "__main__":

diff --git a/utils/material_scan_2D.py b/utils/material_scan_2D.py
@@ -1,7 +1,15 @@
-import os
-from Gaudi.Configuration import *
+"""
+This script must be called with k4run: 'k4run material_scan_2D.py --{argument} {value}'.
+The output files are saved in 'data/{outputDir}/{outputFileBase}.root'.
+If no outputDir is specified, it will be 'data/{outputFileBase}.root'.
+"""
+
+from os import environ, fspath
+from pathlib import Path
 
 from Configurables import ApplicationMgr
+from Gaudi.Configuration import *
+from k4FWCore.parseArgs import parser
 
 ApplicationMgr().EvtSel = "None"
 ApplicationMgr().EvtMax = 1
@@ -10,9 +18,42 @@
 # DD4hep geometry service
 from Configurables import GeoSvc
 
+parser.add_argument(
+    "--compactFile",
+    help="Compact detector file to use",
+    type=str,
+    default=fspath(
+        Path(environ["k4geo_DIR"]) / "ILD" / "compact" / "ILD_sl5_v02" / "ILD_l5_v02.xml"
+    ),
+)
+parser.add_argument(
+    "--outputFileBase",
+    help="Base name of all the produced output files",
+    default="out_material_scan",
+)
+parser.add_argument(
+    "--angleDef",
+    help="angle definition to use: eta, theta, cosTheta or thetaRad, default: eta",
+    choices=["eta", "theta", "cosTheta", "thetaRad"],
+    default="eta",
+)
+parser.add_argument(
+    "--outputDir",
+    "-o",
+    type=str,
+    default="",
+    help="Directory to store the output file in",
+)
+
+reco_args = parser.parse_known_args()[0]
+compact_file = reco_args.compactFile
+angle_def = reco_args.angleDef
+output_dir = "data" / Path(reco_args.outputDir)
+output_dir.mkdir(parents=True, exist_ok=True)  # Create the directory if it doesn't exist
+
 ## parse the given xml file
 geoservice = GeoSvc("GeoSvc")
-geoservice.detectors = ["IDEA_o1_v02.xml"]
+geoservice.detectors = [compact_file]
 geoservice.OutputLevel = INFO
 ApplicationMgr().ExtSvc += [geoservice]
 
@@ -24,9 +65,9 @@
 # For instance adding envelopeName="BoundaryPostCalorimetry" will perform the scan only till the end of calorimetry.
 # BoundaryPostCalorimetry is defined in Detector/DetFCChhECalInclined/compact/envelopePreCalo.xml
 materialservice = MaterialScan_2D_genericAngle("GeoDump")
-materialservice.filename = "out_material_scan.root"
+materialservice.filename = fspath(output_dir / Path(reco_args.outputFileBase).with_suffix(".root"))
 
-materialservice.angleDef = "eta"  # eta, theta, cosTheta or thetaRad
+materialservice.angleDef = angle_def  # eta, theta, cosTheta or thetaRad
 materialservice.angleBinning = 0.05
 materialservice.angleMax = 3.0
 materialservice.angleMin = -3.0