26 from ROOT
import Belle2
27 from ROOT
import gROOT, AddressOf
30 gROOT.ProcessLine(
'struct TreeStruct {\
41 from ROOT
import TreeStruct
49 This module writes its output to a ROOT tree.
50 Adapted from pxd/validation/PXDValidationTTreeSimHit.py
54 """Initialize the module"""
56 super(TreeWriterModule, self).
__init__()
59 self.
file = ROOT.TFile(
'PID_TTree.root',
'recreate')
61 self.
tree = ROOT.TTree(
'tree',
'')
64 """ Declare tree branches """
65 for key
in TreeStruct.__dict__.keys():
68 if isinstance(self.
data.__getattribute__(key), int):
70 self.
tree.Branch(key, AddressOf(self.
data, key), key + formstring)
73 """Store TOP and dE/dx info in tree"""
77 track = pid.getRelatedFrom(
'Tracks')
79 mcpart = track.getRelatedTo(
'MCParticles')
81 pdg = abs(mcpart.getPDG())
82 if pdg != 211
and pdg != 321:
84 momentumVec = mcpart.getMomentum()
85 momentum = momentumVec.Mag()
88 costheta = momentumVec.CosTheta()
89 phi = momentumVec.Phi()
90 fitresult = track.getTrackFitResultWithClosestMass(Belle2.Const.pion)
92 trackmomentum = fitresult.getMomentum().Mag()
97 selectedpart = Belle2.Const.kaon
100 logl_sel = pid.getLogL(selectedpart, pid_dedx)
101 logl_pi = pid.getLogL(Belle2.Const.pion, pid_dedx)
102 dedx_DLL = logl_pi - logl_sel
104 logl_sel = pid.getLogL(selectedpart, pid_top)
105 logl_pi = pid.getLogL(Belle2.Const.pion, pid_top)
106 top_DLL = logl_pi - logl_sel
108 self.
data.lld_dedx = dedx_DLL
109 self.
data.lld_top = top_DLL
110 self.
data.p = momentum
112 self.
data.costheta = costheta
113 self.
data.trackmomentum = trackmomentum
114 self.
data.iskaon = pdg == 321
122 B2WARNING(
'problems with track <-> mcparticle relations')
124 print(
'event: %d, track: %d' % (event, track.getArrayIndex()))
127 """ Close the output file."""
137 main.add_module(register_module(
'RootInput'))