Belle II Software development
printPIDLikelihoods.py
1#!/usr/bin/env python3
2
3
10
11# --------------------------------------------------------------------------
12# Prints PIDLikelihoods from a mdst file
13# File name must be specified using -i switch:
14# basf2 mdst/examples/printPIDLikelihoods.py -i <fileName.root>
15# --------------------------------------------------------------------------
16
17import basf2 as b2
18from ROOT import Belle2
19
20
21class printPIDLikelihoods(b2.Module):
22
23 '''
24 Prints PID log likelihoods + basic track info in a well formatted way.
25 User is prompted to continue or quit at each event
26 '''
27
28 def event(self):
29 '''
30 Print log likelihoods and wait for user respond.
31 '''
32
33 evtMetaData = Belle2.PyStoreObj('EventMetaData')
34 exp = evtMetaData.obj().getExperiment()
35 run = evtMetaData.obj().getRun()
36 evt = evtMetaData.obj().getEvent()
37 print()
38 print('Experiment ' + str(exp) + ' Run ' + str(run) + ' Event ' + str(evt) + ':')
39 print()
40 tracks = Belle2.PyStoreArray('Tracks')
41 for track in tracks:
42 fitresult = track.getTrackFitResult(Belle2.Const.pion)
43 if fitresult:
44 p_mag = fitresult.getMomentum().Mag()
45 pmom = '{:5.3f}'.format(p_mag) + ' GeV/c'
46 else:
47 pmom = '?'
48 mcpart = track.getRelated('MCParticles')
49 if mcpart:
50 pdg = str(mcpart.getPDG())
51 else:
52 pdg = '?'
53 index = track.getArrayIndex()
54 print('Track ' + str(index) + ': p = ' + pmom + ' MCtruth = ' + pdg)
55 likelihood = track.getRelated('PIDLikelihoods')
56 try:
57 print('logLikelihoods:')
58 likelihood.printArray()
59 except BaseException:
60 print('--> No relation to PIDLikelihood')
61 print()
62
63 # wait for user respond
64 response = input("Type Enter to continue or Q to quit.\n").lower().strip()
65 if response == "q":
66 evtMetaData.obj().setEndOfData()
67
68
69b2.set_log_level(b2.LogLevel.ERROR)
70
71# Create path
72main = b2.create_path()
73
74# Root input
75roinput = b2.register_module('RootInput')
76main.add_module(roinput)
77
78# print array of log likelihoods
79main.add_module(printPIDLikelihoods())
80
81# Process events
82b2.process(main)
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:67