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