Skip to content

Commit

Permalink
refactorization, parseArgs added to material_scan_2D.py, thetaRad as …
Browse files Browse the repository at this point in the history 
…angleDef option added to material_plots_2D.py
  • Loading branch information
@Victor-Schwan
Victor-Schwan committed Sep 18, 2024
1 parent 5a42648 commit 0c8bf66
Show file tree
Hide file tree
Showing 2 changed files with 187 additions and 51 deletions.
166 changes: 127 additions & 39 deletions utils/material_plots_2D.py
View file
Original file line number Diff line number Diff line change
@@ -1,81 +1,169 @@
"""
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
from os import fspath
from pathlib import Path

import sys,os
print(os.path.dirname(os.path.abspath(__file__)) + "/../../FCCSW/Examples/scripts")
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../../FCCSW/Examples/scripts") # FCC Style from https://github.com/HEP-FCC/FCCSW/blob/master/Examples/scripts/plotstyle.py
from plotstyle import FCCStyle
import ROOT

sys.path.append(
fspath((Path(__file__).resolve().parent / "../../FCCSW/Examples/scripts").resolve())
) # FCC Style from https://github.com/HEP-FCC/FCCSW/blob/master/Examples/scripts/plotstyle.py
print(sys.path[-1]) # Print the last added path for verification
from plotstyle import FCCStyle


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('--angleMin', dest='angleMin', default=6, type=float, help="minimum eta/theta/cosTheta")
parser.add_argument('--angleMax', dest='angleMax', default=6, type=float, help="maximum eta/theta/cosTheta")
parser.add_argument('--angleDef', dest='angleDef', default="eta", type=str, help="angle definition to use: eta, theta or cosTheta, default: eta")
parser.add_argument('--angleBinning', "-b", dest='angleBinning', default=0.05, type=float, help="eta/theta/cosTheta bin width")
parser.add_argument('--nPhiBins', dest='nPhiBins', default=100, type=int, help="number of bins in phi")
parser.add_argument('--x0max', "-x", dest='x0max', default=0.0, type=float, help="Max of x0")
parser = argparse.ArgumentParser(description="Material Plotter")
parser.add_argument(
"--inputFile", "-f", type=str, help="relative path to the input file"
)
parser.add_argument(
"--angleMin",
default=6,
type=float,
help="Minimum eta/theta/cosTheta",
)
parser.add_argument(
"--angleMax",
default=6,
type=float,
help="Maximum eta/theta/cosTheta",
)
parser.add_argument(
"--angleDef",
default="eta",
choices=["eta", "theta", "cosTheta", "thetaRad"],
type=str,
help="Angle definition to use: eta, theta, thetaRad or cosTheta; default: eta",
)
parser.add_argument(
"--angleBinning",
"-b",
default=0.05,
type=float,
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(
"--outputDir",
"-o",
type=str,
default="plots",
help="Directory to store output files in",
)
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(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_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 = 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,
)

for angleBinning, entry in enumerate(tree):
nMat = entry.nMaterials

entry_x0, entry_lambda, entry_depth = 0.0, 0.0, 0.0
for i in range(nMat):
if entry.material.at(i) == "Air": continue
if entry.material.at(i) == "Air":
continue

entry_x0 += entry.nX0.at(i)*100.0
entry_lambda += entry.nLambda.at(i)
entry_depth += entry.matDepth.at(i)
entry_x0 += entry.nX0.at(i) * 100.0
entry_lambda += entry.nLambda.at(i)
entry_depth += entry.matDepth.at(i)

h_x0.Fill(tree.angle,tree.phi,entry_x0)
h_lambda.Fill(tree.angle,tree.phi,entry_lambda)
h_depth.Fill(tree.angle,tree.phi,entry_depth)
h_x0.Fill(tree.angle, tree.phi, entry_x0)
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)):
cv = ROOT.TCanvas("","",800,600)
axis_titles = [
"Material budget x/X_{0} [%]",
"Number of #lambda",
"Material depth [cm]",
]
for i, plot in enumerate(plots):
cv = ROOT.TCanvas("", "", 800, 600)
cv.SetRightMargin(0.18)
histograms[i].Draw("COLZ")

if args.angleDef=="eta":
title="#eta"
elif args.angleDef=="theta":
title="#theta [#circ]"
elif args.angleDef=="cosTheta":
title="cos(#theta)"
if args.angleDef == "eta":
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":
histograms[i].SetMaximum(args.x0max)
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(fspath(output_path.with_suffix(".pdf")))
cv.Print(fspath(output_path.with_suffix(".png")))
cv.SaveAs(fspath(output_path.with_suffix(".root")))


if __name__ == "__main__":
FCCStyle.initialize()
Expand Down
72 changes: 60 additions & 12 deletions utils/material_scan_2D.py
View file
Original file line number Diff line number Diff line change
@@ -1,36 +1,84 @@
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
ApplicationMgr().EvtSel = 'None'
from Gaudi.Configuration import *
from k4FWCore.parseArgs import parser

ApplicationMgr().EvtSel = "None"
ApplicationMgr().EvtMax = 1
ApplicationMgr().OutputLevel = INFO

# 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.OutputLevel = INFO
geoservice.detectors = [compact_file]
geoservice.OutputLevel = INFO
ApplicationMgr().ExtSvc += [geoservice]

# Using material scan from k4SimGeant4: https://github.com/HEP-FCC/k4SimGeant4/tree/main/Detector/DetComponents/src
from Configurables import MaterialScan_2D_genericAngle

# Material scan is done from the interaction point to the end of world volume.
# In order to use other end boundary, please provide the name of a thin, e.g. cylindrical volume.
# 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
materialservice.nPhi = 100 # number of bins in phi
materialservice.nPhi = 100 # number of bins in phi

ApplicationMgr().ExtSvc += [materialservice]


0 comments on commit 0c8bf66

Please sign in to comment.