30def run_track_pattern_recognition(input_file, geo_file, output_file, method):
31 """
32 Runs all steps of track pattern recognition.
33 Parameters
34 ----------
35 input_file : string
36 Path to an input .root file with events.
37 geo_file : string
38 Path to a file with SHiP geometry.
39 output_file : string
40 Path to an output .root file with quality plots.
41 method : string
42 Name of a track pattern recognition method.
43 """
44
45
46
47
48
49 try:
50 fgeo = ROOT.TFile(geo_file)
51 except:
52 print("An error with opening the ship geo file.")
53 raise
54
55 sGeo = fgeo.FAIRGeom
56
57
58 if not fgeo.FindKey('ShipGeo'):
59
60 if sGeo.GetVolume('EcalModule3') :
61 ecalGeoFile = "ecal_ellipse6x12m2.geo"
62 else:
63 ecalGeoFile = "ecal_ellipse5x10m2.geo"
64
65 if dy:
66 ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", Yheight = dy, EcalGeoFile = ecalGeoFile)
67 else:
68 ShipGeo = ConfigRegistry.loadpy("$FAIRSHIP/geometry/geometry_config.py", EcalGeoFile = ecalGeoFile)
69
70 else:
71 upkl = Unpickler(fgeo)
72 ShipGeo = upkl.load('ShipGeo')
73
74
75 global_variables.ShipGeo = ShipGeo
76
77
78
79 import shipDet_conf
80 run = ROOT.FairRunSim()
81 run.SetName("TGeant4")
82 run.SetOutputFile("dummy")
83 run.SetUserConfig("g4Config_basic.C")
84 rtdb = run.GetRuntimeDb()
85
87 run.Init()
88
89
90
91
92
93
94
95
96 import geomGeant4
97 if hasattr(ShipGeo.Bfield,"fieldMap"):
99 else:
100 print("no fieldmap given, geofile too old, not anymore support")
101 exit(-1)
102 sGeo = fgeo.FAIRGeom
103 geoMat = ROOT.genfit.TGeoMaterialInterface()
104 ROOT.genfit.MaterialEffects.getInstance().init(geoMat)
105 bfield = ROOT.genfit.FairShipFields()
106 bfield.setField(fieldMaker.getGlobalField())
107 fM = ROOT.genfit.FieldManager.getInstance()
108 fM.init(bfield)
109
110
111
112
113
114 try:
115 fn = ROOT.TFile(input_file,'update')
116 except:
117 print("An error with opening the input data file.")
118 raise
119
120 sTree = fn.cbmsim
121 sTree.Write()
122
123
124
125 h = init_book_hist()
126
127
128 import shipPatRec
129
130
131 metrics = {'n_hits': [],
132 'reconstructible': 0,
133 'passed_y12': 0, 'passed_stereo12': 0, 'passed_12': 0,
134 'passed_y34': 0, 'passed_stereo34': 0, 'passed_34': 0,
135 'passed_combined': 0, 'reco_passed': 0, 'reco_passed_no_clones': 0,
136 'frac_y12': [], 'frac_stereo12': [], 'frac_12': [],
137 'frac_y34': [], 'frac_stereo34': [], 'frac_34': [],
138 'reco_frac_tot': [],
139 'reco_mc_p': [], 'reco_mc_theta': [],
140 'fitted_p': [], 'fitted_pval': [], 'fitted_chi': [],
141 'fitted_x': [], 'fitted_y': [], 'fitted_z': [], 'fitted_mass': []}
142
143
144 nEvents = sTree.GetEntries()
145
146
147 for iEvent in range(nEvents):
148
149 if iEvent%1000 == 0:
150 print('Event ', iEvent)
151
152
153
154 rc = sTree.GetEvent(iEvent)
155
156
157
158 reconstructible_tracks = getReconstructibleTracks(iEvent, sTree, sGeo, ShipGeo)
159
160 metrics['reconstructible'] += len(reconstructible_tracks)
161 for i_reco in reconstructible_tracks:
162 h['TracksPassed'].Fill("Reconstructible tracks", 1)
163 h['TracksPassedU'].Fill("Reconstructible tracks", 1)
164
165 in_y12 = []
166 in_stereo12 = []
167 in_12 = []
168 in_y34 = []
169 in_stereo34 = []
170 in_34 = []
171 in_combo = []
172
173 found_track_ids = []
174 n_tracks = len(reconstructible_tracks)
175 n_recognized = 0
176 n_clones = 0
177 n_ghosts = 0
178 n_others = 0
179
180 min_eff = 0.
181
182
183
184 nTracklets = sTree.Tracklets.GetEntriesFast()
185
186 for i_track in range(nTracklets):
187
188 atracklet = sTree.Tracklets[i_track]
189
190 if atracklet.getType() != 1:
191 continue
192
193 atrack = atracklet.getList()
194
195 if atrack.size() == 0:
196 continue
197
198 hits = {'X': [], 'Y': [], 'Z': [],
199 'DetID': [], 'TrackID': [],
200 'Pz': [], 'Px': [], 'Py': [],
201 'dist2Wire': [], 'Pdg': []}
202
203 for ihit in atrack:
204 ahit = sTree.strawtubesPoint[ihit]
205 hits['X'] += [ahit.GetX()]
206 hits['Y'] += [ahit.GetY()]
207 hits['Z'] += [ahit.GetZ()]
208 hits['DetID'] += [ahit.GetDetectorID()]
209 hits['TrackID'] += [ahit.GetTrackID()]
210 hits['Pz'] += [ahit.GetPz()]
211 hits['Px'] += [ahit.GetPx()]
212 hits['Py'] += [ahit.GetPy()]
213 hits['dist2Wire'] += [ahit.dist2Wire()]
214 hits['Pdg'] += [ahit.PdgCode()]
215
216
217 for key in hits.keys():
218 hits[key] = numpy.array(hits[key])
219
220
221 statnb, vnb, pnb, lnb, snb = decodeDetectorID(hits['DetID'])
222 is_stereo = ((vnb == 1) + (vnb == 2))
223 is_y = ((vnb == 0) + (vnb == 3))
224 is_before = ((statnb == 1) + (statnb == 2))
225 is_after = ((statnb == 3) + (statnb == 4))
226
227
228
229 metrics['n_hits'] += [get_n_hits(hits['TrackID'])]
230
231
232 frac_y12, tmax_y12 = fracMCsame(hits['TrackID'][is_before * is_y])
233 n_hits_y12 = get_n_hits(hits['TrackID'][is_before * is_y])
234 frac_stereo12, tmax_stereo12 = fracMCsame(hits['TrackID'][is_before * is_stereo])
235 n_hits_stereo12 = get_n_hits(hits['TrackID'][is_before * is_stereo])
236 frac_12, tmax_12 = fracMCsame(hits['TrackID'][is_before])
237 n_hits_12 = get_n_hits(hits['TrackID'][is_before])
238
239 frac_y34, tmax_y34 = fracMCsame(hits['TrackID'][is_after * is_y])
240 n_hits_y34 = get_n_hits(hits['TrackID'][is_after * is_y])
241 frac_stereo34, tmax_stereo34 = fracMCsame(hits['TrackID'][is_after * is_stereo])
242 n_hits_stereo34 = get_n_hits(hits['TrackID'][is_after * is_stereo])
243 frac_34, tmax_34 = fracMCsame(hits['TrackID'][is_after])
244 n_hits_34 = get_n_hits(hits['TrackID'][is_after])
245 frac_tot, tmax_tot = fracMCsame(hits['TrackID'])
246 n_hits_tot = get_n_hits(hits['TrackID'])
247
248 if tmax_y12 == tmax_stereo12 and tmax_y12 == tmax_y34 and tmax_y12 == tmax_stereo34:
249 if frac_y12 >= min_eff and frac_stereo12 >= min_eff and frac_y34 >= min_eff and frac_stereo34 >= min_eff:
250
251 if tmax_y12 in reconstructible_tracks and tmax_y12 not in found_track_ids:
252 n_recognized += 1
253 found_track_ids.append(tmax_y12)
254 elif tmax_y12 in reconstructible_tracks and tmax_y12 in found_track_ids:
255 n_clones += 1
256 elif tmax_y12 not in reconstructible_tracks:
257 n_others += 1
258
259 else:
260 n_ghosts += 1
261 else:
262 n_ghosts += 1
263
264 is_reconstructed = 0
265 is_reconstructed_no_clones = 0
266
267
268 if tmax_y12 in reconstructible_tracks:
269 metrics['passed_y12'] += 1
270 metrics['frac_y12'] += [frac_y12]
271 h['TracksPassed'].Fill("Y view station 1&2", 1)
272 if tmax_y12 not in in_y12:
273 h['TracksPassedU'].Fill("Y view station 1&2", 1)
274 in_y12.append(tmax_y12)
275
276 if tmax_stereo12 == tmax_y12:
277 metrics['passed_stereo12'] += 1
278 metrics['frac_stereo12'] += [frac_stereo12]
279 h['TracksPassed'].Fill("Stereo station 1&2", 1)
280 if tmax_stereo12 not in in_stereo12:
281 h['TracksPassedU'].Fill("Stereo station 1&2", 1)
282 in_stereo12.append(tmax_stereo12)
283
284 if tmax_12 == tmax_y12:
285 metrics['passed_12'] += 1
286 metrics['frac_12'] += [frac_12]
287 h['TracksPassed'].Fill("station 1&2", 1)
288 if tmax_12 not in in_12:
289 h['TracksPassedU'].Fill("station 1&2", 1)
290 in_12.append(tmax_12)
291
292 if tmax_y34 in reconstructible_tracks:
293 metrics['passed_y34'] += 1
294 metrics['frac_y34'] += [frac_y34]
295 h['TracksPassed'].Fill("Y view station 3&4", 1)
296 if tmax_y34 not in in_y34:
297 h['TracksPassedU'].Fill("Y view station 3&4", 1)
298 in_y34.append(tmax_y34)
299
300 if tmax_stereo34 == tmax_y34:
301 metrics['passed_stereo34'] += 1
302 metrics['frac_stereo34'] += [frac_stereo34]
303 h['TracksPassed'].Fill("Stereo station 3&4", 1)
304 if tmax_stereo34 not in in_stereo34:
305 h['TracksPassedU'].Fill("Stereo station 3&4", 1)
306 in_stereo34.append(tmax_stereo34)
307
308 if tmax_34 == tmax_y34:
309 metrics['passed_34'] += 1
310 metrics['frac_34'] += [frac_34]
311 h['TracksPassed'].Fill("station 3&4", 1)
312 if tmax_34 not in in_34:
313 h['TracksPassedU'].Fill("station 3&4", 1)
314 in_34.append(tmax_34)
315
316 if tmax_12 == tmax_34:
317 metrics['passed_combined'] += 1
318 h['TracksPassed'].Fill("Combined stations 1&2/3&4", 1)
319 metrics['reco_passed'] += 1
320 is_reconstructed = 1
321 if tmax_34 not in in_combo:
322 h['TracksPassedU'].Fill("Combined stations 1&2/3&4", 1)
323 metrics['reco_passed_no_clones'] += 1
324 in_combo.append(tmax_34)
325 is_reconstructed_no_clones = 1
326
327
328 if is_reconstructed == 0:
329 continue
330
331 metrics['reco_frac_tot'] += [frac_tot]
332
333
334 Pz = hits['Pz']
335 Px = hits['Px']
336 Py = hits['Py']
337
338 p, px, py, pz = getPtruthFirst(sTree, tmax_tot)
339 pt = math.sqrt(px**2 + py**2)
340
341 Z_true = []
342 X_true = []
343 Y_true = []
344 for ahit in sTree.strawtubesPoint:
345 if ahit.GetTrackID() == tmax_tot:
346 az, ax, ay = ahit.GetZ(),ahit.GetX(),ahit.GetY()
347 Z_true.append(az)
348 X_true.append(ax)
349 Y_true.append(ay)
350
351
352 metrics['reco_mc_p'] += [p]
353 h['TracksPassed_p'].Fill(p, 1)
354
355
356 Z = hits['Z'][(hits['TrackID'] == tmax_tot) * is_before]
357 X = hits['X'][(hits['TrackID'] == tmax_tot) * is_before]
358 Y = hits['Y'][(hits['TrackID'] == tmax_tot) * is_before]
359 Z = Z - Z[0]
360 X = X - X[0]
361 Y = Y - Y[0]
362 R = numpy.sqrt(X**2 + Y**2 + Z**2)
363 Theta = numpy.arccos(Z[1:] / R[1:])
364 theta = numpy.mean(Theta)
365
366 metrics['reco_mc_theta'] += [theta]
367
368 h['n_hits_reco_y12'].Fill(n_hits_y12)
369 h['n_hits_reco_stereo12'].Fill(n_hits_stereo12)
370 h['n_hits_reco_12'].Fill(n_hits_12)
371 h['n_hits_reco_y34'].Fill(n_hits_y34)
372 h['n_hits_reco_stereo34'].Fill(n_hits_stereo34)
373 h['n_hits_reco_34'].Fill(n_hits_34)
374 h['n_hits_reco'].Fill(n_hits_tot)
375
376 h['n_hits_y12'].Fill(p, n_hits_y12)
377 h['n_hits_stereo12'].Fill(p, n_hits_stereo12)
378 h['n_hits_12'].Fill(p, n_hits_12)
379 h['n_hits_y34'].Fill(p, n_hits_y34)
380 h['n_hits_stereo34'].Fill(p, n_hits_stereo34)
381 h['n_hits_34'].Fill(p, n_hits_34)
382 h['n_hits_total'].Fill(p, n_hits_tot)
383
384 h['frac_y12'].Fill(p, frac_y12)
385 h['frac_stereo12'].Fill(p, frac_stereo12)
386 h['frac_12'].Fill(p, frac_12)
387 h['frac_y34'].Fill(p, frac_y34)
388 h['frac_stereo34'].Fill(p, frac_stereo34)
389 h['frac_34'].Fill(p, frac_34)
390 h['frac_total'].Fill(p, frac_tot)
391
392 h['frac_y12_dist'].Fill(frac_y12)
393 h['frac_stereo12_dist'].Fill(frac_stereo12)
394 h['frac_12_dist'].Fill(frac_12)
395 h['frac_y34_dist'].Fill(frac_y34)
396 h['frac_stereo34_dist'].Fill(frac_stereo34)
397 h['frac_34_dist'].Fill(frac_34)
398 h['frac_total_dist'].Fill(frac_tot)
399
400
401
402 thetrack = sTree.FitTracks[i_track]
403
404 fitStatus = thetrack.getFitStatus()
405 thetrack.prune("CFL")
406
407 nmeas = fitStatus.getNdf()
408 pval = fitStatus.getPVal()
409 chi2 = fitStatus.getChi2() / nmeas
410
411 metrics['fitted_pval'] += [pval]
412 metrics['fitted_chi'] += [chi2]
413
414 h['chi2fittedtracks'].Fill(chi2)
415 h['pvalfittedtracks'].Fill(pval)
416
417
418 try:
419
420 fittedState = thetrack.getFittedState()
421 fittedMom = fittedState.getMomMag()
422 fittedMom = fittedMom
423
424 px_fit,py_fit,pz_fit = fittedState.getMom().x(),fittedState.getMom().y(),fittedState.getMom().z()
425 p_fit = fittedMom
426 pt_fit = math.sqrt(px_fit**2 + py_fit**2)
427
428
429 metrics['fitted_p'] += [p_fit]
430 perr = (p - p_fit) / p
431 h['ptrue-p/ptrue'].Fill(perr)
432 h['perr'].Fill(p, perr)
433 h['perr_direction'].Fill(numpy.rad2deg(theta), perr)
434
435 pterr = (pt - pt_fit) / pt
436 h['pttrue-pt/pttrue'].Fill(pterr)
437
438 pxerr = (px - px_fit) / px
439 h['pxtrue-px/pxtrue'].Fill(pxerr)
440
441 pyerr = (py - py_fit) / py
442 h['pytrue-py/pytrue'].Fill(pyerr)
443
444 pzerr = (pz - pz_fit) / pz
445 h['pztrue-pz/pztrue'].Fill(pzerr)
446
447 if math.fabs(p) > 0.0 :
448 h['pvspfitted'].Fill(p, fittedMom)
449 fittedtrackDir = fittedState.getDir()
450 fittedx=math.degrees(math.acos(fittedtrackDir[0]))
451 fittedy=math.degrees(math.acos(fittedtrackDir[1]))
452 fittedz=math.degrees(math.acos(fittedtrackDir[2]))
453 fittedmass = fittedState.getMass()
454 h['momentumfittedtracks'].Fill(fittedMom)
455 h['xdirectionfittedtracks'].Fill(fittedx)
456 h['ydirectionfittedtracks'].Fill(fittedy)
457 h['zdirectionfittedtracks'].Fill(fittedz)
458 h['massfittedtracks'].Fill(fittedmass)
459
460 metrics['fitted_x'] += [fittedx]
461 metrics['fitted_y'] += [fittedy]
462 metrics['fitted_z'] += [fittedz]
463 metrics['fitted_mass'] += [fittedmass]
464
465
466 Z_fit = []
467 X_fit = []
468 Y_fit = []
469 for az in Z_true:
470 rc,pos,mom = extrapolateToPlane(thetrack, az)
471 Z_fit.append(pos.Z())
472 X_fit.append(pos.X())
473 Y_fit.append(pos.Y())
474
475 for i in range(len(Z_true)):
476 xerr = abs(X_fit[i] - X_true[i])
477 yerr = abs(Y_fit[i] - Y_true[i])
478 h['abs(x - x-true)'].Fill(xerr)
479 h['abs(y - y-true)'].Fill(yerr)
480
481 rmse_x = numpy.sqrt(numpy.mean((numpy.array(X_fit) - numpy.array(X_true))**2))
482 rmse_y = numpy.sqrt(numpy.mean((numpy.array(Y_fit) - numpy.array(Y_true))**2))
483 h['rmse_x'].Fill(rmse_x)
484 h['rmse_y'].Fill(rmse_y)
485
486
487
488 except:
489 print("Problem with fitted state.")
490
491
492 h['Reco_tracks'].Fill("N total", n_tracks)
493 h['Reco_tracks'].Fill("N recognized tracks", n_recognized)
494 h['Reco_tracks'].Fill("N clones", n_clones)
495 h['Reco_tracks'].Fill("N ghosts", n_ghosts)
496 h['Reco_tracks'].Fill("N others", n_others)
497
498
499
500 save_hists(h, output_file)
501
502 return metrics
503
504
505
506
addVMCFields(shipGeo, controlFile='', verbose=False, withVirtualMC=True)