ATLAS Offline Software
plot_material.py
Go to the documentation of this file.
1 #!/usr/bin/env python
2 from __future__ import print_function
3 import ROOT
4 import argparse
5 import math
6 import os
7 import sys
8 
9 # sys.path.append(os.path.dirname(__file__))
10 from geo_id import GeometryID
11 
12 
13 def style(h):
14  h.SetFillColor(ROOT.kBlue)
15 
16 def mkdir(path, recursive=True):
17  cmd = ["mkdir"]
18  if recursive: cmd += ["-p"]
19  cmd += [path]
20  os.system(" ".join(cmd))
21 
22 def draw_eta_lines(hist, etas = [0, 1, 1.5, 2, 2.5, 2, 3, 4], sym=True):
23 
24  etas = etas + [-eta for eta in etas if eta != 0.]
25 
26  xmin = hist.GetXaxis().GetBinLowEdge(0)
27  xmax = hist.GetXaxis().GetBinUpEdge(hist.GetXaxis().GetNbins())
28 
29  xext = min(map(abs, (xmin, xmax)))
30  xmin = -xext
31  xmax = xext
32 
33  ymax = hist.GetYaxis().GetBinUpEdge(hist.GetYaxis().GetNbins())
34 
35  l = ROOT.TLine()
36  l.SetLineWidth(1)
37  l.SetLineColor(ROOT.kRed)
38  tl = ROOT.TLatex()
39  tl.SetTextSize(0.025)
40  tl.SetTextColor(ROOT.kRed)
41 
42  for eta in etas:
43  theta = 2 * math.atan(math.exp(-eta))
44  y = ymax
45  x = y / math.tan(theta)
46 
47  m = y/x
48  if x > xmax:
49  y = m*xmax
50  x = y/m
51  if x < xmin:
52  y = m*xmin
53  x = y/m
54 
55  lab_offset = ymax * tl.GetTextSize()*0.5
56  l.DrawLine(0, 0, x, y)
57  if eta != 0:
58  if x < 0:
59  tl.SetTextAlign(11)
60  else:
61  tl.SetTextAlign(31)
62  else:
63  tl.SetTextAlign(21)
64  tl.DrawLatex(x, y + lab_offset, "#eta = {:.1f}".format(eta))
65 
66 
67 class HistManager:
68  def __init__(self):
69  self._hist2D = {}
70  self._hist2D_count = {}
71 
72  def _key(self, k):
73  if type(k) == tuple:
74  return "_".join(str(v) for v in k)
75  else:
76  return k
77 
78  def addHist2D(self, name, binx, xmin, xmax, biny, ymin, ymax, xlab="", ylab="", zlab="", average=False):
79  key = self._key(name)
80  root_name = os.path.basename(name)
81 
82  assert key not in self._hist2D, "hist2D with key {} exists".format(key)
83  self._hist2D[key] = ROOT.TH2D(root_name, root_name, binx, xmin, xmax, biny, ymin, ymax)
84  self._hist2D[key].GetXaxis().SetTitle(xlab)
85  self._hist2D[key].GetYaxis().SetTitle(ylab)
86  self._hist2D[key].GetZaxis().SetTitle(zlab)
87 
88 
89 
90  if average:
91  self._hist2D_count[key] = ROOT.TH2D(root_name+"_count", root_name+"_count", binx, xmin, xmax, biny, ymin, ymax)
92 
93  def hist2DExists(self, name):
94  key = self._key(name)
95  return key in self._hist2D
96 
97  def fillHist2D(self, name, x, y, *args):
98  key = self._key(name)
99  assert key in self._hist2D, "{} not in histograms".format(key)
100  self._hist2D[key].Fill(x, y, *args)
101  if key in self._hist2D_count:
102  self._hist2D_count[key].Fill(x, y)
103 
104  def hist2D(self, name):
105  key = self._key(name)
106  return self._hist2D[key]
107 
108  def hist2Diter(self):
109  for name, hist in self._hist2D.items():
110  if name not in self._hist2D_count:
111  yield name, hist
112  else:
113  hist_divided = hist.Clone()
114  hist_divided.Divide(self._hist2D_count[name])
115  yield name, hist_divided
116 
117 def axislabels(h, x, y):
118  h.GetXaxis().SetTitle(x)
119  h.GetYaxis().SetTitle(y)
120 
121 
122 def main():
123 
124  p = argparse.ArgumentParser()
125  p.add_argument("geant")
126  p.add_argument("mapped")
127  p.add_argument("--events", "-n", type=int)
128  p.add_argument("outdir")
129 
130  args = p.parse_args()
131 
132 
133  ROOT.gROOT.SetBatch()
134  # ROOT.gROOT.SetStyle("ATLAS")
135  ROOT.gROOT.LoadMacro("../AtlasStyle.C")
136  ROOT.SetAtlasStyle()
137 
138 
139  ROOT.gROOT.LoadMacro(os.path.join(os.path.dirname(__file__), "plot_helper.C"))
140  geom_graph = ROOT.fill_geometry_hist_rz("excells_charged.root", 5000)
141 
142  c = ROOT.TCanvas("c", "c", 800, 600)
143 
144  geom_graph.Draw("ap")
145  c.SaveAs(os.path.join(args.outdir, "geom_graph.png"))
146 
147  geant_file = ROOT.TFile.Open(args.geant)
148  geant_tree = geant_file.Get("MaterialTracks")
149 
150  mapped_file = ROOT.TFile.Open(args.mapped)
151  mapped_tree = mapped_file.Get("MaterialTracks")
152 
153 
154  colors = [
155  ROOT.kOrange,
156  ROOT.kAzure,
157  ROOT.kRed,
158  ROOT.kViolet,
159  ROOT.kGreen,
160  ROOT.kPink,
161  ROOT.kBlue,
162  ROOT.kMagenta,
163  ROOT.kCyan,
164  ROOT.kTeal,
165  ROOT.kSpring,
166  ]
167  # colors = [
168  # "#e6194b",
169  # "#3cb44b",
170  # "#ffe119",
171  # "#0082c8",
172  # "#f58231",
173  # "#911eb4",
174  # "#46f0f0",
175  # "#f032e6",
176  # "#d2f53c",
177  # "#fabebe",
178  # "#008080",
179  # "#e6beff",
180  # "#aa6e28",
181  # "#fffac8",
182  # "#800000",
183  # "#aaffc3",
184  # "#808000",
185  # "#ffd8b1",
186  # "#000080"
187  # ]
188 
189  vol_lay_col_map = {}
190  vol_6_lay_idxs = range(2, 12+2, 2)
191  vol_7_lay_idxs = range(2, 8+2, 2)
192  vol_11_lay_idxs = range(2, 18+2, 2)
193  vol_12_lay_idxs = range(2, 8+2, 2)
194 
195  for i, l in enumerate(vol_6_lay_idxs):
196  vol_lay_col_map[(6, l)] = colors[i%len(colors)]
197 
198  for l1, l2 in zip(vol_6_lay_idxs, reversed(vol_6_lay_idxs)):
199  vol_lay_col_map[(8, l1)] = vol_lay_col_map[(6, l2)]
200 
201  for i, l in enumerate(vol_11_lay_idxs):
202  vol_lay_col_map[(11, l)] = colors[i%len(colors)]
203 
204  for l1, l2 in zip(vol_11_lay_idxs, reversed(vol_11_lay_idxs)):
205  vol_lay_col_map[(13, l1)] = vol_lay_col_map[(11, l2)]
206 
207  for i, l in enumerate(vol_7_lay_idxs):
208  vol_lay_col_map[(7, l)] = colors[i%len(colors)]
209  for i, l in enumerate(vol_12_lay_idxs):
210  vol_lay_col_map[(12, l)] = colors[i%len(colors)]
211 
212 
213  # vol_lay_col_map[(6, 2)] = colors[0]
214  # vol_lay_col_map[(6, 4)] = colors[1]
215  # vol_lay_col_map[(6, 6)] = colors[2]
216  # vol_lay_col_map[(6, 8)] = colors[3]
217  # vol_lay_col_map[(6, 10)] = colors[4]
218  # vol_lay_col_map[(6, 12)] = colors[5]
219 
220  # vol_lay_col_map[(8, 2)] = col_lay_col_map[(6, 12)]
221  # vol_lay_col_map[(8, 4)] = col_lay_col_map[(6, 10)]
222  # vol_lay_col_map[(8, 6)] = col_lay_col_map[(6, 8)]
223  # vol_lay_col_map[(8, 8)] = col_lay_col_map[(6, 6)]
224  # vol_lay_col_map[(8, 10)] = col_lay_col_map[(6, 4)]
225  # vol_lay_col_map[(8, 12)] = col_lay_col_map[(6, 2)]
226 
227  def hex2rgb(h):
228  h = h[1:]
229  return tuple(int(h[i:i+2], 16) for i in (0, 2 ,4))
230 
231 
232  layers_geant = do_step(c, geant_tree, os.path.join(args.outdir, "geant"), args.events, geom_graph)
233  layers_mapped = do_step(c, mapped_tree, os.path.join(args.outdir, "mapped"), args.events, geom_graph)
234 
235  # c.SetWindowSize(1000, 600)
236  c_all = ROOT.TCanvas("c_all", "c_all", 1000, 600)
237 
238  keys_geant = set(layers_geant.keys())
239  keys_mapped = set(layers_mapped.keys())
240  keys = keys_geant & keys_mapped
241 
242  for idx, key in enumerate(keys):
243  geo_id = GeometryID(key)
244  print(geo_id)
245 
246  h2_geant = layers_geant[key]
247  h2_mapped = layers_mapped[key]
248 
249  def clip(h, val):
250  for bx in range(h.GetNbinsX()):
251  for by in range(h.GetNbinsY()):
252  bc = h.GetBinContent(bx, by)
253  if bc > 0:
254  h.SetBinContent(bx, by, val)
255 
256  clip(h2_geant, 1)
257  clip(h2_mapped, 2)
258 
259  # for g in (g_geant, g_mapped):
260  # g.SetMarkerSize(0.5)
261 
262  # color_idx = geo_id.lay_id / 2
263  # print("col_idx:", color_idx)
264 
265  # color_a = colors[color_idx%len(colors)]
266 
267  color_a = vol_lay_col_map[(geo_id.vol_id, geo_id.lay_id)]
268 
269  # r, g, b = hex2rgb(colors[idx%len(colors)])
270  # color_a = ROOT.TColor(r, g, b, 1.0)
271  # color_b = ROOT.TColor(r, g, b, 0.6)
272 
273  # basecolnum, _ = colors[idx%len(colors)]
274  # basecol = ROOT.gROOT.GetColor(basecolnum)
275  # r, g, b = hex2rgb(basecol.AsHexString())
276 
277 
278 
279  # color_a = ROOT.TColor.GetColor(r, g, b)
280  color_b = ROOT.TColor.GetColorTransparent(color_a, 0.6)
281 
282  # color_a = ROOT.TColor(r, g, b, 1.0)
283  # color_b = ROOT.TColor(r, g, b, 0.6)
284 
285 
286 
287  # h2_geant.SetLineColor(color_b)
288  h2_geant.SetLineColor(color_a)
289  h2_mapped.SetLineColor(color_a)
290  h2_mapped.SetLineStyle(2)
291  h2_mapped.SetFillColor(color_a)
292 
293  for h in (h2_geant, h2_mapped):
294  h.SetLineWidth(1)
295 
296  opt = "BOX"
297 
298  c.cd()
299  h2_geant.Draw(opt)
300  h2_mapped.Draw(opt+" sames")
301 
302  c_all.cd()
303 
304  if idx == 0:
305  h2_geant.Draw(opt)
306  h2_mapped.Draw(opt+" sames")
307  else:
308  opt = "BOX sames"
309  h2_geant.Draw(opt)
310  h2_mapped.Draw(opt)
311 
312 
313  lay_str = "vol_{:03d}_lay_{:03d}_app_{:03d}.pdf".format(geo_id.vol_id, geo_id.lay_id, geo_id.app_id)
314  path = os.path.join(args.outdir, "assoc", "{}.pdf".format(lay_str))
315  mkdir(os.path.dirname(path))
316  c.SaveAs(path)
317 
318  path = os.path.join(args.outdir, "assoc.pdf")
319  mkdir(os.path.dirname(path))
320  c_all.SaveAs(path)
321 
322 
323 
324 
325 
326 def do_step(c, tree, outdir, events, geom_graph):
327  hm = HistManager()
328 
329  hists = {}
330 
331  def mkprofile(name, xlab, ylab, *args):
332  hists[name] = ROOT.TProfile(name, name, *args)
333  axislabels(hists[name], xlab, ylab)
334 
335  mkprofile("theta_x0", '\\theta', "X_{0}", 40, 0, math.pi)
336  mkprofile("theta_dInX0", '\\theta', "t/X_{0}", 40, 0, math.pi)
337  mkprofile("theta_t", '\\theta', "t", 40, 0, math.pi)
338  mkprofile("theta_tTot", '\\theta', "t_{tot}", 40, 0, math.pi)
339 
340  mkprofile("phi_x0", "\\phi", "X_{0}", 40, -math.pi, math.pi)
341  mkprofile("phi_dInX0", '\\phi', "t/X_{0}", 40, -math.pi, math.pi)
342  mkprofile("phi_t", '\\phi', 't', 40, -math.pi, math.pi)
343  mkprofile("phi_tTot", '\\phi', 't_{tot}', 40, -math.pi, math.pi)
344 
345  mkprofile("eta_x0", '\\eta', 'X_{0}', 40, -4, 4)
346  mkprofile("eta_dInX0", '\\eta', 't/X_{0}', 40, -4, 4)
347  mkprofile("eta_t", '\\eta', 't', 40, -4, 4)
348  mkprofile("eta_tTot", '\\eta', 't_{tot}', 40, -4, 4)
349 
350  nBinsEta = 80
351  nBinsPhi = 80
352  nBinsR = 200
353  nBinsZ = 500
354 
355 
356  hm.addHist2D("eta_phi/eta_phi_x0", nBinsEta, -4, 4, nBinsPhi, -math.pi, math.pi, "\\eta", "\\phi", average=True)
357  hm.addHist2D("eta_phi/eta_phi_dInX0", nBinsEta, -4, 4, nBinsPhi, -math.pi, math.pi, "\\eta", "\\phi", average=True)
358  hm.addHist2D("eta_phi/eta_phi_tTot", nBinsEta, -4, 4, nBinsPhi, -math.pi, math.pi, "\\eta", "\\phi", average=True)
359 
360  hm.addHist2D("z_r/z_r_x0", nBinsZ, -3000, 3000, nBinsR, 0, 1200, "z [mm]", "r [mm]", "X_{0}", average=True)
361 
362  for k, v in hists.iteritems():
363  style(v)
364 
365  layers = {}
366 
367 
368  for idx, event in enumerate(tree):
369 
370  if events:
371  if idx > events: break
372  if idx % 1000 == 0:
373  print("{} / {}".format(idx, tree.GetEntries()))
374 
375  # X0 = event.X0
376  theta = event.theta
377  phi = event.phi
378 
379  eta = - math.log(math.tan(theta/2))
380 
381 
382  tTot = 0
383 
384  dInX0_tot = event.dInX0
385 
386  totalX0 = 0
387 
388  zipped = zip(event.step_t,
389  event.step_X0,
390  event.step_dInX0,
391  event.step_z,
392  event.step_r,
393  event.step_geo_id)
394  for t, x0, dInX0, z, r, _geo_id in zipped:
395  geo_id = GeometryID(_geo_id)
396  # tX0 = t/x0
397  tTot += t
398  totalX0 += x0
399 
400  # if geo_id.value() not in layers:
401  # layers.append(geo_id.value())
402 
403  lay_str = "vol_{}_lay_{}_app_{}".format(geo_id.vol_id, geo_id.lay_id, geo_id.app_id)
404  if _geo_id not in layers:
405  hname = lay_str+"_"+outdir+"_z_r"
406  layers[_geo_id] = ROOT.TH2D(hname, hname, nBinsZ/2, -3000, 3000, nBinsR/2, 0, 1200)
407 
408  h2 = layers[_geo_id]
409  # g.SetPoint(g.GetN(), z, r)
410  h2.Fill(z, r)
411 
412  hm.fillHist2D("z_r/z_r_x0", z, r, x0)
413 
414  # hname = "z_r/lay/z_r_"+lay_str+"_x0"
415  # if not hm.hist2DExist(hname):
416  # hm.hist2D
417 
418 
419  hists["theta_x0"].Fill(theta, totalX0)
420  hists["phi_x0"].Fill(phi, totalX0)
421  hists["eta_x0"].Fill(eta, totalX0)
422 
423  hm.fillHist2D("eta_phi/eta_phi_x0", eta, phi, totalX0)
424  hm.fillHist2D("eta_phi/eta_phi_tTot", eta, phi, tTot)
425  hm.fillHist2D("eta_phi/eta_phi_dInX0", eta, phi, dInX0_tot)
426 
427  hists["theta_dInX0"].Fill(theta, dInX0_tot)
428  hists["phi_dInX0"].Fill(phi, dInX0_tot)
429  hists["eta_dInX0"].Fill(eta, dInX0_tot)
430 
431  hists["theta_tTot"].Fill(theta, tTot)
432  hists["phi_tTot"].Fill(phi, tTot)
433  hists["eta_tTot"].Fill(eta, tTot)
434 
435 
436  hists["phi_dInX0"].GetYaxis().SetRangeUser(0, hists["phi_dInX0"].GetMaximum()*1.2)
437  hists["phi_x0"].GetYaxis().SetRangeUser(0, hists["phi_x0"].GetMaximum()*1.2)
438  hists["phi_t"].GetYaxis().SetRangeUser(0, hists["phi_t"].GetMaximum()*1.2)
439  hists["phi_tTot"].GetYaxis().SetRangeUser(0, hists["phi_tTot"].GetMaximum()*1.2)
440 
441  os.system("mkdir -p {}".format(outdir))
442 
443  for k, v in hists.iteritems():
444  v.Draw("hist")
445  c.SaveAs("{}/{}.pdf".format(outdir, k))
446 
447  c.SetRightMargin(0.2)
448  c.SetTopMargin(0.1)
449 
450  for name, h2 in hm.hist2Diter():
451 
452  parts = name.split("/")
453  print(parts)
454  if "z_r" in parts:
455  h2.Draw("colz")
456  draw_eta_lines(h2)
457  filename = "{}.pdf".format(os.path.join(outdir, name))
458  mkdir(os.path.dirname(filename))
459  c.SaveAs(filename)
460 
461  h2.Draw("colz")
462  geom_graph.Draw("p+same")
463  draw_eta_lines(h2)
464  filename = "{}_geom.pdf".format(os.path.join(outdir, name))
465  mkdir(os.path.dirname(filename))
466  c.SaveAs(filename)
467 
468  else:
469  h2.Draw("colz")
470  filename = "{}.pdf".format(os.path.join(outdir, name))
471  mkdir(os.path.dirname(filename))
472  c.SaveAs(filename)
473 
474  # for h2 in hist2D_zr:
475  # h2.Divide(hist2D_zr_count)
476  # h2.Draw("colz")
477 
478  # h2.GetXaxis().SetTitle("z [mm]")
479  # h2.GetYaxis().SetTitle("r [mm]")
480 
481  # draw_eta_lines(h2)
482  # c.SaveAs("{}/{}.pdf".format(args.outdir, h2.GetName()))
483 
484  # geom_graph.Draw("p+same")
485  # c.SaveAs("{}/{}_w_geom.pdf".format(args.outdir, h2.GetName()))
486 
487  # for h2 in hist2D_eta_phi:
488  # h2.Divide(hist2D_eta_phi_count)
489  # h2.Draw("colz")
490  # c.SaveAs("{}/{}.pdf".format(args.outdir, h2.GetName()))
491 
492  return layers
493 
494 
495 if "__main__" == __name__:
496  main()
vtune_athena.format
format
Definition: vtune_athena.py:14
plot_material.mkdir
def mkdir(path, recursive=True)
Definition: plot_material.py:16
CaloCellPos2Ntuple.int
int
Definition: CaloCellPos2Ntuple.py:24
plot_material.HistManager
Definition: plot_material.py:67
plot_material.HistManager.fillHist2D
def fillHist2D(self, name, x, y, *args)
Definition: plot_material.py:97
plot_material.do_step
def do_step(c, tree, outdir, events, geom_graph)
Definition: plot_material.py:326
plot_material.HistManager.addHist2D
def addHist2D(self, name, binx, xmin, xmax, biny, ymin, ymax, xlab="", ylab="", zlab="", average=False)
Definition: plot_material.py:78
plot_material.axislabels
def axislabels(h, x, y)
Definition: plot_material.py:117
plot_material.HistManager.hist2Diter
def hist2Diter(self)
Definition: plot_material.py:108
plotBeamSpotVxVal.range
range
Definition: plotBeamSpotVxVal.py:195
plot_material.HistManager._hist2D
_hist2D
Definition: plot_material.py:69
plot_material.HistManager.hist2D
def hist2D(self, name)
Definition: plot_material.py:104
plot_material.HistManager._key
def _key(self, k)
Definition: plot_material.py:72
CxxUtils::set
constexpr std::enable_if_t< is_bitmask_v< E >, E & > set(E &lhs, E rhs)
Convenience function to set bits in a class enum bitmask.
Definition: bitmask.h:232
min
#define min(a, b)
Definition: cfImp.cxx:40
TCS::join
std::string join(const std::vector< std::string > &v, const char c=',')
Definition: Trigger/TrigT1/L1Topo/L1TopoCommon/Root/StringUtils.cxx:10
TrigJetMonitorAlgorithm.items
items
Definition: TrigJetMonitorAlgorithm.py:79
Run3DQTestingDriver.GeometryID
GeometryID
Definition: Run3DQTestingDriver.py:82
python.CaloScaleNoiseConfig.type
type
Definition: CaloScaleNoiseConfig.py:78
plot_material.main
def main()
Definition: plot_material.py:122
plot_material.HistManager.hist2DExists
def hist2DExists(self, name)
Definition: plot_material.py:93
plot_material.HistManager.__init__
def __init__(self)
Definition: plot_material.py:68
str
Definition: BTagTrackIpAccessor.cxx:11
plot_material.style
def style(h)
Definition: plot_material.py:13
dbg::print
void print(std::FILE *stream, std::format_string< Args... > fmt, Args &&... args)
Definition: SGImplSvc.cxx:70
plot_material.draw_eta_lines
def draw_eta_lines(hist, etas=[0, 1, 1.5, 2, 2.5, 2, 3, 4], sym=True)
Definition: plot_material.py:22
plot_material.HistManager._hist2D_count
_hist2D_count
Definition: plot_material.py:70
Generated on Thu Nov 7 2024 21:23:57 for ATLAS Offline Software by   doxygen 1.8.18