Back to home page

EIC code displayed by LXR
 
 

    

File indexing completed on 2025-07-01 07:56:15 

0001 #!/usr/bin/env python
0002 # Copyright 2023, Christopher Dilks
0003 # Subject to the terms in the LICENSE file found in the top-level directory.
0004 
0005 import ROOT as r
0006 import sys, getopt, pathlib, math
0007 
0008 # suppress graphics
0009 r.gROOT.SetBatch(True)
0010 
0011 # GLOBAL SETTINGS
0012 ################################################################
0013 RESID_MAX = 10 # cherenkov angle |residual| maximum
0014 RADIATORS = {
0015         'Aerogel': { 'p_max': 30, 'p_rebin': 12, 'rindex_ref': 1.0190},
0016         'Gas':     { 'p_max': 70, 'p_rebin': 30, 'rindex_ref': 1.00076},
0017         'Merged':  { 'p_max': 70, 'p_rebin': 30, },
0018         }
0019 PARTICLE_MASS = {
0020         'e':  0.00051,
0021         'pi': 0.13957,
0022         'K':  0.49368,
0023         'p':  0.93827,
0024         }
0025 
0026 # ARGUMENTS
0027 ################################################################
0028 
0029 ana_file_name = 'out/ana.edm4hep.root'
0030 output_dir    = 'out/ana.plots'
0031 
0032 helpStr = f'''
0033 {sys.argv[0]} [OPTIONS]
0034 
0035     -i <input file>: specify an input file, e.g., hepmc
0036        default: {ana_file_name}
0037 
0038     -o <output dir>: specify an output directory
0039        default: {output_dir}
0040     
0041     -h: show this usage guide
0042 
0043     '''
0044 
0045 try:
0046     opts, args = getopt.getopt(sys.argv[1:], 'i:o:h')
0047 except getopt.GetoptError:
0048     print('\n\nERROR: invalid argument\n', helpStr, file=sys.stderr)
0049     sys.exit(2)
0050 for opt, arg in opts:
0051     if(opt == '-i'): ana_file_name = arg.lstrip()
0052     if(opt == '-o'): output_dir    = arg.lstrip()
0053     if(opt == '-h'):
0054         print(helpStr)
0055         sys.exit(2)
0056 print(f'''
0057 ana_file_name = {ana_file_name}
0058 output_dir    = {output_dir}
0059 ''')
0060 
0061 
0062 # PLOTTING
0063 ################################################################
0064 
0065 # make canvases
0066 ana_file = r.TFile.Open(ana_file_name, "READ")
0067 def make_canv(name, nx, ny):
0068     canv = r.TCanvas(name, name, nx*600, ny*400)
0069     canv.Divide(nx,ny)
0070     return canv
0071 canv_dict = {
0072     "photon_spectra": make_canv("photon_spectra", 2, 2),
0073     "digitization":   make_canv("digitization",   2, 2),
0074     "pidAerogel":     make_canv("pidAerogel",     5, 3),
0075     "pidGas":         make_canv("pidGas",         5, 3),
0076     "pidMerged":      make_canv("pidMerged",      2, 3),
0077 }
0078 
0079 # helper functions
0080 ### minimum momentum for Cherenkov radiation
0081 def calculate_mom_min(mass,n):
0082     return mass / math.sqrt(n**2-1)
0083 ### Cherenkov angle at saturation
0084 def calculate_theta_max(n):
0085     return 1000 * math.acos(1/n)
0086 ### execute `op` for every pad of `canv`
0087 def loop_pads(canv, op):
0088     for pad in canv.GetListOfPrimitives():
0089         if pad.InheritsFrom(r.TVirtualPad.Class()):
0090             op(pad)
0091 ### draw a profile on 2D histogram `hist`
0092 def draw_profile(hist, rad, style="BOX"):
0093     prof = hist.ProfileX("_pfx", 1, -1, "i")
0094     prof.SetMarkerStyle(r.kFullCircle)
0095     prof.SetMarkerColor(r.kRed)
0096     prof.SetMarkerSize(2)
0097     hist_rebinned = hist.Clone(f'{hist.GetName()}_rebin')
0098     if 'vs_p' in hist.GetName():
0099         hist_rebinned.RebinX(RADIATORS[rad]['p_rebin'])
0100     else:
0101         hist_rebinned.RebinX(4)
0102     hist_rebinned.SetLineColor(r.kBlue)
0103     hist_rebinned.SetFillColor(r.kBlue)
0104     for p in [hist_rebinned, prof]:
0105         if 'vs_p' in hist.GetName():
0106             p.GetXaxis().SetRangeUser(0, RADIATORS[rad]['p_max'])
0107         if 'rindex_ref' in RADIATORS[rad] and 'theta_vs_p' in hist.GetName():
0108             p.GetYaxis().SetRangeUser(0, 1.5*calculate_theta_max(RADIATORS[rad]['rindex_ref']))
0109         if 'thetaResid_vs' in hist.GetName():
0110             p.GetYaxis().SetRangeUser(-RESID_MAX, RESID_MAX)
0111     hist_rebinned.Draw(style)
0112     prof.Draw("SAME")
0113 
0114 # Cherenkov angle vs. p functions
0115 for rad, radHash in RADIATORS.items():
0116     if rad == "Merged":
0117         continue
0118     radHash['cherenkov_curves'] = []
0119     n = radHash['rindex_ref']
0120     for part, mass in PARTICLE_MASS.items():
0121         ftn = r.TF1(
0122                 f'ftn_theta_{part}_{rad}',
0123                 f'1000*TMath::ACos(TMath::Sqrt(x^2+{mass}^2)/({n}*x))',
0124                 calculate_mom_min(mass, n),
0125                 radHash['p_max']
0126                 )
0127         ftn.SetLineColor(r.kBlack)
0128         ftn.SetLineWidth(1)
0129         radHash['cherenkov_curves'].append(ftn)
0130 
0131 # draw photon spectra
0132 canv = canv_dict["photon_spectra"]
0133 def canv_op(pad):
0134     pad.SetGrid(1,1)
0135     pad.SetLogy()
0136 loop_pads(canv, canv_op)
0137 canv.cd(1)
0138 ana_file.Get("phot/phot_spectrum_sim").Draw()
0139 canv.cd(2)
0140 ana_file.Get("phot/nphot_dist").Draw()
0141 canv.cd(3)
0142 ana_file.Get("digi/phot_spectrum_rec").Draw()
0143 canv.cd(4)
0144 ana_file.Get("digi/nhits_dist").Draw()
0145 
0146 # draw digitization
0147 canv = canv_dict["digitization"]
0148 def canv_op(pad):
0149     pad.SetGrid(1,1)
0150     if pad.GetNumber() < 3:
0151         pad.SetLogy()
0152     else:
0153         pad.SetLogz()
0154 loop_pads(canv, canv_op)
0155 canv.cd(1)
0156 ana_file.Get("digi/adc_dist").Draw()
0157 canv.cd(2)
0158 ana_file.Get("digi/tdc_dist").Draw()
0159 canv.cd(3)
0160 ana_file.Get("digi/tdc_vs_adc").Draw("COLZ")
0161 
0162 # draw CherenkovPID for each radiator
0163 for rad in ["Aerogel", "Gas"]:
0164     pid_name = f'pid{rad}'
0165     canv = canv_dict[pid_name]
0166     def canv_op(pad):
0167         pad.SetGrid(1,1)
0168     loop_pads(canv, canv_op)
0169     canv.cd(1)
0170     ana_file.Get(f'{pid_name}/npe_dist_{rad}').Draw()
0171     canv.cd(2)
0172     ana_file.Get(f'{pid_name}/theta_dist_{rad}').Draw()
0173     canv.cd(3)
0174     hist = ana_file.Get(f'{pid_name}/thetaResid_dist_{rad}')
0175     hist.DrawCopy()
0176     canv.cd(4)
0177     hist.SetTitle(hist.GetTitle() + " - ZOOM")
0178     hist.GetXaxis().SetRangeUser(-RESID_MAX, RESID_MAX)
0179     hist.Draw()
0180     resid_mode = hist.GetBinCenter(hist.GetMaximumBin())
0181     resid_dev  = hist.GetStdDev()
0182     nsigma = 3
0183     hist.Fit("gaus", "", "", resid_mode - nsigma*resid_dev, resid_mode + nsigma*resid_dev)
0184     if hist.GetFunction("gaus"):
0185         hist.GetFunction("gaus").SetLineWidth(5)
0186     canv.cd(5)
0187     ana_file.Get(f'{pid_name}/highestWeight_dist_{rad}').Draw()
0188     canv.cd(6)
0189     hist = ana_file.Get(f'{pid_name}/npe_vs_p_{rad}')
0190     draw_profile(hist, rad)
0191     canv.cd(7)
0192     hist = ana_file.Get(f'{pid_name}/theta_vs_p_{rad}')
0193     draw_profile(hist, rad)
0194     if 'cherenkov_curves' in RADIATORS[rad]:
0195         for ftn in RADIATORS[rad]['cherenkov_curves']:
0196             ftn.Draw("SAME")
0197     canv.cd(8)
0198     hist = ana_file.Get(f'{pid_name}/thetaResid_vs_p_{rad}')
0199     hist.GetYaxis().SetRangeUser(-RESID_MAX, RESID_MAX)
0200     draw_profile(hist, rad)
0201     canv.cd(9)
0202     ana_file.Get(f'{pid_name}/photonTheta_vs_photonPhi_{rad}').Draw("COLZ")
0203     canv.cd(10)
0204     ana_file.Get(f'{pid_name}/mcRindex_{rad}').Draw()
0205     canv.cd(11)
0206     hist = ana_file.Get(f'{pid_name}/npe_vs_eta_{rad}')
0207     draw_profile(hist, rad)
0208     canv.cd(12)
0209     hist = ana_file.Get(f'{pid_name}/theta_vs_eta_{rad}')
0210     draw_profile(hist, rad)
0211     canv.cd(13)
0212     hist = ana_file.Get(f'{pid_name}/thetaResid_vs_eta_{rad}')
0213     hist.GetYaxis().SetRangeUser(-RESID_MAX, RESID_MAX)
0214     draw_profile(hist, rad)
0215 
0216 # draw CherenkovPID for merged radiators
0217 pid_name = f'pidMerged'
0218 canv = canv_dict[pid_name]
0219 def canv_op(pad):
0220     pad.SetGrid(1,1)
0221 loop_pads(canv, canv_op)
0222 canv.cd(1)
0223 ana_file.Get(f'{pid_name}/npe_dist_Merged').Draw()
0224 canv.cd(2)
0225 ana_file.Get(f'{pid_name}/highestWeight_dist_Merged').Draw()
0226 canv.cd(3)
0227 hist = ana_file.Get(f'{pid_name}/npe_vs_p_Merged')
0228 draw_profile(hist, 'Merged')
0229 canv.cd(4)
0230 ana_file.Get(f'{pid_name}/photonTheta_vs_photonPhi_Merged').Draw("COLZ")
0231 canv.cd(5)
0232 hist = ana_file.Get(f'{pid_name}/npe_vs_eta_Merged')
0233 draw_profile(hist, 'Merged')
0234 
0235 # FINISH
0236 ################################################################
0237 
0238 pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)
0239 for name, canvas in canv_dict.items():
0240     canvas.SaveAs(f'{output_dir}/{name}.png')
0241 ana_file.Close()
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.3.7 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!