Belle II Software development
cosmicsExtrapolationPlots.py
1#!/usr/bin/env python3
2
3
10
11"""
12<header>
13 <contact>software-tracking@belle2.org</contact>
14 <input>CosmicsExtrapolation.root</input>
15 <output>CosmicsExtrapolationPlots.root</output>
16 <description>Validation of cosmic track extrapolation.</description>
17</header>
18"""
19
20
21import basf2
22import ROOT
23from ROOT import Belle2, TNamed
24
25ACTIVE = True
26
27contact = 'Kirill Chilikin (chilikin@lebedev.ru)'
28
29bklm_numbers = Belle2.BKLMElementNumbers()
30eklm_numbers = Belle2.EKLMElementNumbers.Instance()
31
32
33def run():
34 """
35 Create the cosmics extrapolation validation plots.
36 """
37
38 class CosmicsExtapolationPlotModule(basf2.Module):
39 """ Class for creation of cosmics extrapolation plot module. """
40
41 def set_options_coordinate(self, histogram, description, shifter):
42 """ Set options for coordinate plot. """
43 histogram.SetXTitle('cm')
44 histogram.SetYTitle('Events')
45 function_list = histogram.GetListOfFunctions()
46 function_list.Add(TNamed('Description', description))
47 function_list.Add(TNamed('Check', ' No bias, no large background, resolution ~ 2 cm.'))
48 function_list.Add(TNamed('Contact', contact))
49 if shifter:
50 basf2.B2WARNING("temporarily removed all shifter plots for the CosmicsExtrapolationPlots")
51 # function_list.Add(TNamed('MetaOptions', 'shifter'))
52
53 def set_options_momentum(self, histogram, description):
54 """ Set options for momentum plot. """
55 histogram.SetXTitle('p [GeV]')
56 histogram.SetYTitle('Events')
57 function_list = histogram.GetListOfFunctions()
58 function_list.Add(TNamed('Description', description))
59 function_list.Add(TNamed('Check', ' No bias, no large background, resolution ~ 2 cm.'))
60 function_list.Add(TNamed('Contact', contact))
61
62 def __init__(self):
63 """Initialization."""
64 super().__init__()
65
66
67 self.output_file = ROOT.TFile('CosmicsExtrapolationPlots.root',
68 'recreate')
69
70
71 self.hist_xres_forward_mum = \
72 ROOT.TH1F('xres_forward_mum',
73 'X resolution (#mu^{-}, forward)',
74 100, -20, 20)
75 self.set_options_coordinate(
76 self.hist_xres_forward_mum,
77 'Extrapolated hit X resolution (mu-, forward propagation).', True)
78
79
80 self.hist_xres_forward_mup = \
81 ROOT.TH1F('xres_forward_mup',
82 'X resolution (#mu^{+}, forward)',
83 100, -20, 20)
84 self.set_options_coordinate(
85 self.hist_xres_forward_mup,
86 'Extrapolated hit X resolution (mu+, forward propagation).', True)
87
88
89 self.hist_yres_forward_mum = \
90 ROOT.TH1F('yres_forward_mum',
91 'Y resolution (#mu^{-}, forward)',
92 100, -20, 20)
93 self.set_options_coordinate(
94 self.hist_yres_forward_mum,
95 'Extrapolated hit Y resolution (mu-, forward propagation).', False)
96
97
98 self.hist_yres_forward_mup = \
99 ROOT.TH1F('yres_forward_mup',
100 'Y resolution (#mu^{+}, forward)',
101 100, -20, 20)
102 self.set_options_coordinate(
103 self.hist_yres_forward_mup,
104 'Extrapolated hit Y resolution (mu+, forward propagation).', False)
105
106
107 self.hist_zres_forward_mum = \
108 ROOT.TH1F('zres_forward_mum',
109 'Z resolution (#mu^{-}, forward)',
110 100, -20, 20)
111 self.set_options_coordinate(
112 self.hist_zres_forward_mum,
113 'Extrapolated hit Z resolution (mu-, forward propagation).', False)
114
115
116 self.hist_zres_forward_mup = \
117 ROOT.TH1F('zres_forward_mup',
118 'Z resolution (#mu^{+}, forward)',
119 100, -20, 20)
120 self.set_options_coordinate(
121 self.hist_zres_forward_mup,
122 'Extrapolated hit Z resolution (mu+, forward propagation).', False)
123
124
125 self.hist_xres_backward_mum = \
126 ROOT.TH1F('xres_backward_mum',
127 'X resolution (#mu^{-}, backward)',
128 100, -20, 20)
129 self.set_options_coordinate(
130 self.hist_xres_backward_mum,
131 'Extrapolated hit X resolution (mu-, backward propagation).', True)
132
133
134 self.hist_xres_backward_mup = \
135 ROOT.TH1F('xres_backward_mup',
136 'X resolution (#mu^{+}, backward)',
137 100, -20, 20)
138 self.set_options_coordinate(
139 self.hist_xres_backward_mup,
140 'Extrapolated hit X resolution (mu+, backward propagation).', True)
141
142
143 self.hist_yres_backward_mum = \
144 ROOT.TH1F('yres_backward_mum',
145 'Y resolution (#mu^{-}, backward)',
146 100, -20, 20)
147 self.set_options_coordinate(
148 self.hist_yres_backward_mum,
149 'Extrapolated hit Y resolution (mu-, backward propagation).', False)
150
151
152 self.hist_yres_backward_mup = \
153 ROOT.TH1F('yres_backward_mup',
154 'Y resolution (#mu^{+}, backward)',
155 100, -20, 20)
156 self.set_options_coordinate(
157 self.hist_yres_backward_mup,
158 'Extrapolated hit Y resolution (mu+, backward propagation).', False)
159
160
161 self.hist_zres_backward_mum = \
162 ROOT.TH1F('zres_backward_mum',
163 'Z resolution (#mu^{-}, backward)',
164 100, -20, 20)
165 self.set_options_coordinate(
166 self.hist_zres_backward_mum,
167 'Extrapolated hit Z resolution (mu-, backward propagation).', False)
168
169
170 self.hist_zres_backward_mup = \
171 ROOT.TH1F('zres_backward_mup',
172 'Z resolution (#mu^{+}, backward)',
173 100, -20, 20)
174 self.set_options_coordinate(
175 self.hist_zres_backward_mup,
176 'Extrapolated hit Z resolution (mu+, backward propagation).', False)
177
178
179 self.hist_pres_forward_mum = \
180 ROOT.TH1F('pres_forward_mum',
181 'Momentum resolution (#mu^{-}, forward)',
182 100, -10, 10)
183 self.set_options_momentum(
184 self.hist_pres_forward_mum,
185 'Momentum resolution (mu-, forward propagation).')
186
187
188 self.hist_pres_forward_mup = \
189 ROOT.TH1F('pres_forward_mup',
190 'Momentum resolution (#mu^{+}, forward)',
191 100, -10, 10)
192 self.set_options_momentum(
193 self.hist_pres_forward_mup,
194 'Momentum resolution (mu+, forward propagation).')
195
196
197 self.hist_pres_backward_mum = \
198 ROOT.TH1F('pres_backward_mum',
199 'Momentum resolution (#mu^{-}, backward)',
200 100, -10, 10)
201 self.set_options_momentum(
202 self.hist_pres_backward_mum,
203 'Momentum resolution (mu-, backward propagation).')
204
205
206 self.hist_pres_backward_mup = \
207 ROOT.TH1F('pres_backward_mup',
208 'Momentum resolution (#mu^{+}, backward)',
209 100, -10, 10)
210 self.set_options_momentum(
211 self.hist_pres_backward_mup,
212 'Momentum resolution (mu+, backward propagation).')
213
214 def fill_histograms_exthit(self, exthit, klmhit2d):
215 """ Fill histograms with ExtHit data. """
216 ext_position = exthit.getPosition()
217 if exthit.isBackwardPropagated():
218 if exthit.getPdgCode() == 13:
219 self.hist_xres_backward_mum.Fill(
220 ext_position.X() - klmhit2d.getPositionX())
221 self.hist_yres_backward_mum.Fill(
222 ext_position.Y() - klmhit2d.getPositionY())
223 self.hist_zres_backward_mum.Fill(
224 ext_position.Z() - klmhit2d.getPositionZ())
225 elif exthit.getPdgCode() == -13:
226 self.hist_xres_backward_mup.Fill(
227 ext_position.X() - klmhit2d.getPositionX())
228 self.hist_yres_backward_mup.Fill(
229 ext_position.Y() - klmhit2d.getPositionY())
230 self.hist_zres_backward_mup.Fill(
231 ext_position.Z() - klmhit2d.getPositionZ())
232 else:
233 if exthit.getPdgCode() == 13:
234 self.hist_xres_forward_mum.Fill(
235 ext_position.X() - klmhit2d.getPositionX())
236 self.hist_yres_forward_mum.Fill(
237 ext_position.Y() - klmhit2d.getPositionY())
238 self.hist_zres_forward_mum.Fill(
239 ext_position.Z() - klmhit2d.getPositionZ())
240 elif exthit.getPdgCode() == -13:
241 self.hist_xres_forward_mup.Fill(
242 ext_position.X() - klmhit2d.getPositionX())
243 self.hist_yres_forward_mup.Fill(
244 ext_position.Y() - klmhit2d.getPositionY())
245 self.hist_zres_forward_mup.Fill(
246 ext_position.Z() - klmhit2d.getPositionZ())
247
248 def event(self):
249 """ Event function. """
250 klmhit2ds = Belle2.PyStoreArray('KLMHit2ds')
251 exthits = Belle2.PyStoreArray('ExtHits')
252 tracks = Belle2.PyStoreArray('Tracks')
253 mcparticles = Belle2.PyStoreArray('MCParticles')
254
255 for klmhit2d in klmhit2ds:
256 subdetector = klmhit2d.getSubdetector()
257 section = klmhit2d.getSection()
258 sector = klmhit2d.getSector()
259 layer = klmhit2d.getLayer()
260 if (subdetector == Belle2.KLMElementNumbers.c_BKLM):
261 for exthit in exthits:
262 if exthit.getDetectorID() != Belle2.Const.BKLM:
263 continue
264 module = exthit.getCopyID()
265 section_ext = bklm_numbers.getSectionByModule(module)
266 sector_ext = bklm_numbers.getSectorByModule(module)
267 layer_ext = bklm_numbers.getLayerByModule(module)
268 if (section_ext != section or sector_ext != sector or
269 layer_ext != layer):
270 continue
271 self.fill_histograms_exthit(exthit, klmhit2d)
272 else:
273 section = klmhit2d.getSection()
274 sector = klmhit2d.getSector()
275 layer = klmhit2d.getLayer()
276 for exthit in exthits:
277 if exthit.getDetectorID() != Belle2.Const.EKLM:
278 continue
279 strip_global = exthit.getCopyID()
280 section_ext = eklm_numbers.getSectionByGlobalStrip(strip_global)
281 sector_ext = eklm_numbers.getSectorByGlobalStrip(strip_global)
282 layer_ext = eklm_numbers.getLayerByGlobalStrip(strip_global)
283 if (section_ext != section or sector_ext != sector or
284 layer_ext != layer):
285 continue
286 self.fill_histograms_exthit(exthit, klmhit2d)
287
288 for track in tracks:
289 track_fit_result = track.getTrackFitResult(Belle2.Const.muon)
290 momentum = track_fit_result.getMomentum()
291 track_exthits = track.getRelationsTo('ExtHits')
292 if len(track_exthits) == 0:
293 continue
294 mc_momentum = mcparticles[0].getMomentum()
295 p_diff = momentum.R() - mc_momentum.R()
296 muon_found = False
297 for i in range(len(track_exthits)):
298 if abs(track_exthits[i].getPdgCode()) == 13:
299 exthit = track_exthits[i]
300 muon_found = True
301 if not muon_found:
302 continue
303 if exthit.isBackwardPropagated():
304 if exthit.getPdgCode() == 13:
305 self.hist_pres_backward_mum.Fill(p_diff)
306 elif exthit.getPdgCode() == -13:
307 self.hist_pres_backward_mup.Fill(p_diff)
308 else:
309 if exthit.getPdgCode() == 13:
310 self.hist_pres_forward_mum.Fill(p_diff)
311 elif exthit.getPdgCode() == -13:
312 self.hist_pres_forward_mup.Fill(p_diff)
313
314 def terminate(self):
315 """ Termination function. """
316 self.output_file.cd()
317 self.hist_xres_forward_mum.Write()
318 self.hist_xres_forward_mup.Write()
319 self.hist_yres_forward_mum.Write()
320 self.hist_yres_forward_mup.Write()
321 self.hist_zres_forward_mum.Write()
322 self.hist_zres_forward_mup.Write()
323 self.hist_xres_backward_mum.Write()
324 self.hist_xres_backward_mup.Write()
325 self.hist_yres_backward_mum.Write()
326 self.hist_yres_backward_mup.Write()
327 self.hist_zres_backward_mum.Write()
328 self.hist_zres_backward_mup.Write()
329 self.hist_pres_forward_mum.Write()
330 self.hist_pres_forward_mup.Write()
331 self.hist_pres_backward_mum.Write()
332 self.hist_pres_backward_mup.Write()
333 self.output_file.Close()
334
335 # Input.
336 root_input = basf2.register_module('RootInput')
337 root_input.param('inputFileName', '../CosmicsExtrapolation.root')
338
339 # Plotting.
340 plot = CosmicsExtapolationPlotModule()
341
342 # Create main path.
343 main = basf2.create_path()
344
345 # Add modules to main path
346 main.add_module(root_input)
347 main.add_module(plot)
348
349 main.add_module('Progress')
350
351 # Run.
352 basf2.process(main)
353
354 print(basf2.statistics)
355
356
357if __name__ == '__main__':
358 if ACTIVE:
359 run()
360 else:
361 print("This validation deactivated and thus basf2 is not executed.\n"
362 "If you want to run this validation, please set the 'ACTIVE' flag above to 'True'.\n"
363 "Exiting.")
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72