SND@LHC Software
Loading...
Searching...
No Matches
flux_map.py
Go to the documentation of this file.
1#!/usr/bin/env python2
2from __future__ import division
3import argparse
4import numpy as np
5import ROOT as r
6# Fix https://root-forum.cern.ch/t/pyroot-hijacks-help/15207 :
7r.PyConfig.IgnoreCommandLineOptions = True
8import shipunit as u
9import rootUtils as ut
10import logger as log
11
12
13def main():
14 parser = argparse.ArgumentParser(description='Script to create flux maps.')
15 parser.add_argument(
16 'inputfile',
17 help='''Simulation results to use as input. '''
18 '''Supports retrieving files from EOS via the XRootD protocol.''')
19 parser.add_argument(
20 '-o',
21 '--outputfile',
22 default='flux_map.root',
23 help='''File to write the flux maps to. '''
24 '''Will be recreated if it already exists.''')
25 args = parser.parse_args()
26 f = r.TFile.Open(args.outputfile, 'recreate')
27 f.cd()
28 maxpt = 10. * u.GeV
29 maxp = 360. * u.GeV
30 h = {}
31
32 # Define histograms
33 for nplane in range(0, 23):
34 ut.bookHist(h, 'NuTauMu_all_{}'.format(nplane),
35 'Rpc_{};x[cm];y[cm]'.format(
36 nplane), 100, -300, +300, 100, -300,
37 300)
38 ut.bookHist(h, 'NuTauMu_mu_{}'.format(nplane),
39 'Rpc_{};x[cm];y[cm]'.format(
40 nplane), 100, -300, +300, 100, -300,
41 300)
42 for suffix, title in [('mu', '#mu#pm hits'), ('all', 'All hits')]:
43 ut.bookHist(h, 'muon_tiles_{}'.format(suffix),
44 '{};x[cm];y[cm]'.format(title), 200, -1000, +1000, 90,
45 -900, 900)
46 ut.bookHist(h, 'muon_bars_x_{}'.format(suffix),
47 '{};x[cm];y[cm]'.format(title), 2, -300, +300, 240, -600,
48 600)
49 ut.bookHist(h, 'muon_bars_y_{}'.format(suffix),
50 '{};x[cm];y[cm]'.format(title), 120, -300, +300, 4, -600,
51 600)
52 ut.bookHist(h, 'timing_{}'.format(suffix),
53 '{};x[cm];y[cm]'.format(title), 3, -252, +252, 167, -501,
54 501)
55 ut.bookHist(h, 'TargetTracker_{}'.format(suffix),
56 '{};x[cm];y[cm]'.format(title), 120, -60, 60, 120, -60,
57 60)
58 ut.bookHist(h, 'TargetTracker_yvsz_{}'.format(suffix),
59 '{};z[cm];y[cm]'.format(
60 title), 400, -3300, -2900, 120, -60,
61 60)
62 ut.bookHist(h, 'TargetTracker_xvsz_{}'.format(suffix),
63 '{};z[cm];x[cm]'.format(
64 title), 400, -3300, -2900, 120, -60,
65 60)
66 ut.bookHist(h, 'NuTauMu_{}'.format(suffix),
67 '{};x[cm];y[cm]'.format(title), 100, -300, +300, 100, -300,
68 300)
69 ut.bookHist(h, 'NuTauMu_yvsz_{}'.format(suffix),
70 '{};z[cm];y[cm]'.format(
71 title), 200, -2680, -2480, 600, -300,
72 300)
73 ut.bookHist(h, 'NuTauMu_xvsz_{}'.format(suffix),
74 '{};z[cm];x[cm]'.format(
75 title), 200, -2680, -2480, 600, -300,
76 300)
77 ut.bookHist(h, 'ECAL_TP_{}'.format(suffix),
78 '{};x[cm];y[cm]'.format(title), 167, -501, +501, 334,
79 -1002, 1002)
80 ut.bookHist(h, 'ECAL_Alt_{}'.format(suffix),
81 '{};x[cm];y[cm]'.format(title), 50, -500, +500, 100, -1000,
82 1000)
83 ut.bookHist(h, 'SBT_Liquid_{}'.format(suffix),
84 '{};z[cm];#phi'.format(title), 100, -3000, +3000, 100,
85 -r.TMath.Pi(), r.TMath.Pi())
86 ut.bookHist(h, 'SBT_Plastic_{}'.format(suffix),
87 '{};z[cm];#phi'.format(title), 100, -3000, +3000, 100,
88 -r.TMath.Pi(), r.TMath.Pi())
89 for station in range(1, 5):
90 ut.bookHist(h, 'T{}_{}'.format(station, suffix),
91 '{};x[cm];y[cm]'.format(title), 10, -250, +250, 20,
92 -500, 500)
93
94 ut.bookHist(h, 'NuTauMu_mu_p', '#mu#pm;p[GeV];', 100, 0, maxp)
95 ut.bookHist(h, 'NuTauMu_mu_pt', '#mu#pm;p_t[GeV];', 100, 0,
96 maxpt)
97 ut.bookHist(h, 'NuTauMu_mu_ppt', '#mu#pm;p[GeV];p_t[GeV];',
98 100, 0, maxp, 100, 0, maxpt)
99 ut.bookHist(h, 'NuTauMu_all_p', '#mu#pm;p[GeV];', 100, 0, maxp)
100 ut.bookHist(h, 'NuTauMu_all_pt', '#mu#pm;p_t[GeV];', 100, 0,
101 maxpt)
102 ut.bookHist(h, 'NuTauMu_all_ppt', '#mu#pm;p[GeV];p_t[GeV];',
103 100, 0, maxp, 100, 0, maxpt)
104
105 for suffix in ['', '_original']:
106 ut.bookHist(h, 'mu_p{}'.format(suffix), '#mu#pm;p[GeV];', 100, 0, maxp)
107 ut.bookHist(h, 'mu_pt{}'.format(suffix), '#mu#pm;p_t[GeV];', 100, 0,
108 maxpt)
109 ut.bookHist(h, 'mu_ppt{}'.format(suffix), '#mu#pm;p[GeV];p_t[GeV];',
110 100, 0, maxp, 100, 0, maxpt)
111 ut.bookHist(h, 'ECAL_TP_e', 'e#pm with E#geq 250 MeV;x[cm];y[cm]', 167,
112 -501, +501, 334, -1002, 1002)
113 ut.bookHist(h, 'ECAL_Alt_e', 'e#pm with E#geq 250 MeV;x[cm];y[cm]', 50,
114 -500, +500, 100, -1000, 1000)
115 ut.bookHist(h, 'ECAL_TP_gamma', '#gamma;x[cm];y[cm]', 167, -501, +501, 334,
116 -1002, 1002)
117 ut.bookHist(h, 'ECAL_Alt_gamma', '#gamma;x[cm];y[cm]', 50, -500, +500, 100,
118 -1000, 1000)
119 ut.bookHist(h, 'ECAL_e_E', 'e#pm;E[GeV/c^{2}];', 100, 0, 1)
120 ut.bookHist(h, 'ECAL_gamma_E', '#gamma;E[GeV/c^{2}];', 100, 0, 1)
121 ch = r.TChain('cbmsim')
122 ch.Add(args.inputfile)
123 n = ch.GetEntries()
124 log.info(n)
125 i = 0
126 for event in ch:
127 if i % 10000 == 0:
128 log.info('{}/{}'.format(i, n))
129 i += 1
130 muon = False
131 for hit in event.strawtubesPoint:
132 if hit:
133 if not hit.GetEnergyLoss() > 0:
134 continue
135 trid = hit.GetTrackID()
136 assert trid > 0
137 weight = event.MCTrack[trid].GetWeight()
138 x = hit.GetX()
139 y = hit.GetY()
140 z = hit.GetZ()
141 px = hit.GetPx()
142 py = hit.GetPy()
143 pz = hit.GetPz()
144 pt = np.hypot(px, py)
145 P = np.hypot(pz, pt)
146 pid = hit.PdgCode()
147 assert pid not in [12, -12, 14, -14, 16, -16]
148 detector_ID = hit.GetDetectorID()
149 station = detector_ID // 10000000
150 h['T{}_all'.format(station)].Fill(x, y, weight)
151 if abs(pid) == 13:
152 muon = True
153 muonid = trid
154 h['T{}_mu'.format(station)].Fill(x, y, weight)
155 h['mu_p'].Fill(P, weight)
156 h['mu_pt'].Fill(pt, weight)
157 h['mu_ppt'].Fill(P, pt, weight)
158 for hit in event.EcalPoint:
159 if hit:
160 if not hit.GetEnergyLoss() > 0:
161 continue
162 trid = hit.GetTrackID()
163 assert trid > 0
164 weight = event.MCTrack[trid].GetWeight()
165 x = hit.GetX()
166 y = hit.GetY()
167 px = hit.GetPx()
168 py = hit.GetPy()
169 pz = hit.GetPz()
170 pt = np.hypot(px, py)
171 P = np.hypot(pz, pt)
172 pid = hit.PdgCode()
173 if pid in [12, -12, 14, -14, 16, -16]:
174 continue
175 h['ECAL_TP_all'].Fill(x, y, weight)
176 h['ECAL_Alt_all'].Fill(x, y, weight)
177 if abs(pid) == 13:
178 muon = True
179 muonid = trid
180 h['mu_p'].Fill(P, weight)
181 h['mu_pt'].Fill(pt, weight)
182 h['mu_ppt'].Fill(P, pt, weight)
183 h['ECAL_TP_mu'].Fill(x, y, weight)
184 h['ECAL_Alt_mu'].Fill(x, y, weight)
185 elif abs(pid) == 11:
186 Esq = px ** 2 + py ** 2 + pz ** 2 + 0.000511 ** 2
187 h['ECAL_e_E'].Fill(np.sqrt(Esq), weight)
188 h['ECAL_TP_e'].Fill(x, y, weight)
189 h['ECAL_Alt_e'].Fill(x, y, weight)
190 elif abs(pid) == 22:
191 Esq = px ** 2 + py ** 2 + pz ** 2
192 h['ECAL_gamma_E'].Fill(np.sqrt(Esq), weight)
193 h['ECAL_TP_gamma'].Fill(x, y, weight)
194 h['ECAL_Alt_gamma'].Fill(x, y, weight)
195 for hit in event.TTPoint:
196 if hit:
197 if not hit.GetEnergyLoss() > 0:
198 continue
199 trid = hit.GetTrackID()
200 assert trid > 0
201 detID = hit.GetDetectorID()
202 weight = event.MCTrack[trid].GetWeight()
203 x = hit.GetX()
204 y = hit.GetY()
205 z = hit.GetZ()
206 px = hit.GetPx()
207 py = hit.GetPy()
208 pz = hit.GetPz()
209 pt = np.hypot(px, py)
210 P = np.hypot(pz, pt)
211 pid = hit.PdgCode()
212 assert pid not in [12, -12, 14, -14, 16, -16]
213 if detID == 0:
214 h['TargetTracker_all'].Fill(x, y, weight)
215 h['TargetTracker_xvsz_all'].Fill(z, x, weight)
216 h['TargetTracker_yvsz_all'].Fill(z, y, weight)
217 if abs(pid) == 13:
218 muon = True
219 muonid = trid
220 h['mu_p'].Fill(P, weight)
221 h['mu_pt'].Fill(pt, weight)
222 h['mu_ppt'].Fill(P, pt, weight)
223 if detID == 0:
224 h['TargetTracker_mu'].Fill(x, y, weight)
225 h['TargetTracker_xvsz_mu'].Fill(z, x, weight)
226 h['TargetTracker_yvsz_mu'].Fill(z, y, weight)
227 for hit in event.ShipRpcPoint:
228 if hit:
229 if not hit.GetEnergyLoss() > 0:
230 continue
231 trid = hit.GetTrackID()
232 assert trid > 0
233 detID = hit.GetDetectorID()
234 nplane = detID - 10000
235 weight = event.MCTrack[trid].GetWeight()
236 x = hit.GetX()
237 y = hit.GetY()
238 z = hit.GetZ()
239 px = hit.GetPx()
240 py = hit.GetPy()
241 pz = hit.GetPz()
242 pt = np.hypot(px, py)
243 P = np.hypot(pz, pt)
244 pid = hit.PdgCode()
245 assert pid not in [12, -12, 14, -14, 16, -16]
246 h['NuTauMu_all'].Fill(x, y, weight)
247 if nplane >= 0:
248 h['NuTauMu_all_{}'.format(nplane)].Fill(x, y, weight)
249 h['NuTauMu_xvsz_all'].Fill(z, x, weight)
250 h['NuTauMu_yvsz_all'].Fill(z, y, weight)
251 if detID == 10000:
252 h['NuTauMu_all_p'].Fill(P, weight)
253 h['NuTauMu_all_pt'].Fill(pt, weight)
254 h['NuTauMu_all_ppt'].Fill(P, pt, weight)
255 if abs(pid) == 13:
256 muon = True
257 muonid = trid
258 h['mu_p'].Fill(P, weight)
259 h['mu_pt'].Fill(pt, weight)
260 h['mu_ppt'].Fill(P, pt, weight)
261 h['NuTauMu_mu'].Fill(x, y, weight)
262 if nplane >= 0:
263 # fill the histogram corresponding to nplane
264 h['NuTauMu_mu_{}'.format(nplane)].Fill(x, y, weight)
265 if detID == 10000:
266 h['NuTauMu_mu_p'].Fill(P, weight)
267 h['NuTauMu_mu_pt'].Fill(pt, weight)
268 h['NuTauMu_mu_ppt'].Fill(P, pt, weight)
269 h['NuTauMu_xvsz_mu'].Fill(z, x, weight)
270 h['NuTauMu_yvsz_mu'].Fill(z, y, weight)
271 for hit in event.TimeDetPoint:
272 if hit:
273 if not hit.GetEnergyLoss() > 0:
274 continue
275 trid = hit.GetTrackID()
276 assert trid > 0
277 weight = event.MCTrack[trid].GetWeight()
278 x = hit.GetX()
279 y = hit.GetY()
280 z = hit.GetZ()
281 px = hit.GetPx()
282 py = hit.GetPy()
283 pz = hit.GetPz()
284 pt = np.hypot(px, py)
285 P = np.hypot(pz, pt)
286 pid = hit.PdgCode()
287 assert pid not in [12, -12, 14, -14, 16, -16]
288 h['timing_all'].Fill(x, y, weight)
289 if abs(pid) == 13:
290 muon = True
291 muonid = trid
292 h['mu_p'].Fill(P, weight)
293 h['mu_pt'].Fill(pt, weight)
294 h['mu_ppt'].Fill(P, pt, weight)
295 h['timing_mu'].Fill(x, y, weight)
296 for hit in event.muonPoint:
297 if hit:
298 if not hit.GetEnergyLoss() > 0:
299 continue
300 trid = hit.GetTrackID()
301 assert trid > 0
302 weight = event.MCTrack[trid].GetWeight()
303 x = hit.GetX()
304 y = hit.GetY()
305 px = hit.GetPx()
306 py = hit.GetPy()
307 pz = hit.GetPz()
308 pt = np.hypot(px, py)
309 P = np.hypot(pz, pt)
310 pid = hit.PdgCode()
311 assert pid not in [12, -12, 14, -14, 16, -16]
312 h['muon_tiles_all'].Fill(x, y, weight)
313 h['muon_bars_x_all'].Fill(x, y, weight)
314 h['muon_bars_y_all'].Fill(x, y, weight)
315 if abs(pid) == 13:
316 muon = True
317 muonid = trid
318 h['mu_p'].Fill(P, weight)
319 h['mu_pt'].Fill(pt, weight)
320 h['mu_ppt'].Fill(P, pt, weight)
321 h['muon_tiles_mu'].Fill(x, y, weight)
322 h['muon_bars_y_mu'].Fill(x, y, weight)
323 h['muon_bars_x_mu'].Fill(x, y, weight)
324 for hit in event.vetoPoint:
325 if hit:
326 if not hit.GetEnergyLoss() > 0:
327 continue
328 trid = hit.GetTrackID()
329 assert trid > 0
330 weight = event.MCTrack[trid].GetWeight()
331 x = hit.GetX()
332 y = hit.GetY()
333 z = hit.GetZ()
334 px = hit.GetPx()
335 py = hit.GetPy()
336 pz = hit.GetPz()
337 pt = np.hypot(px, py)
338 P = np.hypot(pz, pt)
339 pid = hit.PdgCode()
340 detector_ID = hit.GetDetectorID()
341 assert pid not in [12, -12, 14, -14, 16, -16]
342 phi = r.TMath.ATan2(y, x)
343 if 99999 < detector_ID < 999999:
344 h['SBT_Liquid_all'].Fill(z, phi, weight)
345 if abs(pid) == 13:
346 muon = True
347 muonid = trid
348 h['mu_p'].Fill(P, weight)
349 h['mu_pt'].Fill(pt, weight)
350 h['mu_ppt'].Fill(P, pt, weight)
351 h['SBT_Liquid_mu'].Fill(z, phi, weight)
352 continue
353 elif detector_ID > 999999:
354 h['SBT_Plastic_all'].Fill(z, phi, weight)
355 if abs(pid) == 13:
356 muon = True
357 muonid = trid
358 h['mu_p'].Fill(P, weight)
359 h['mu_pt'].Fill(pt, weight)
360 h['mu_ppt'].Fill(P, pt, weight)
361 h['SBT_Plastic_mu'].Fill(z, phi, weight)
362 continue
363 log.warn('Unidentified vetoPoint.')
364 if muon:
365 original_muon = event.MCTrack[muonid]
366 weight = original_muon.GetWeight()
367 h['mu_p_original'].Fill(original_muon.GetP(), weight)
368 h['mu_pt_original'].Fill(original_muon.GetPt(), weight)
369 h['mu_ppt_original'].Fill(original_muon.GetP(),
370 original_muon.GetPt(), weight)
371 # NOTE: muons are counted several times if they create several hits
372 # But the original muon is only counted once.
373 log.info('Event loop done')
374 for key in h:
375 classname = h[key].Class().GetName()
376 if 'TH' in classname or 'TP' in classname:
377 h[key].Write()
378 f.Close()
379
380
381if __name__ == '__main__':
382 r.gErrorIgnoreLevel = r.kWarning
383 r.gROOT.SetBatch(True)
384 main()
flux_map.main
main()
Definition flux_map.py:13