diff --git a/utils/material_plots_2D.py b/utils/material_plots_2D.py index d67758cf..e6ed7e3d 100644 --- 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,6 +130,8 @@ 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) @@ -123,15 +139,16 @@ def main(): 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 index ab6ee21e..c59a9b75 100644 --- 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