Belle II Software  release-05-01-25
pid_ttree_writer.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
19 
20 import sys
21 import math
22 from basf2 import *
23 
24 # Some ROOT tools
25 import ROOT
26 from ROOT import Belle2
27 from ROOT import gROOT, AddressOf
28 
29 # Define a ROOT struct to hold output data in the TTree.
30 gROOT.ProcessLine('struct TreeStruct {\
31  float lld_dedx;\
32  float lld_top;\
33  float p;\
34  float phi;\
35  float costheta;\
36  float trackmomentum;\
37  int iskaon;\
38 };'
39  )
40 
41 from ROOT import TreeStruct
42 
43 # define the python module to save the PID information
44 
45 
46 class TreeWriterModule(Module):
47 
48  """
49  This module writes its output to a ROOT tree.
50  Adapted from pxd/validation/PXDValidationTTreeSimHit.py
51  """
52 
53  def __init__(self):
54  """Initialize the module"""
55 
56  super(TreeWriterModule, self).__init__()
57 
58 
59  self.file = ROOT.TFile('PID_TTree.root', 'recreate')
60 
61  self.tree = ROOT.TTree('tree', '')
62 
63  self.data = TreeStruct()
64  """ Declare tree branches """
65  for key in TreeStruct.__dict__.keys():
66  if '__' not in key:
67  formstring = '/F'
68  if isinstance(self.data.__getattribute__(key), int):
69  formstring = '/I'
70  self.tree.Branch(key, AddressOf(self.data, key), key + formstring)
71 
72  def event(self):
73  """Store TOP and dE/dx info in tree"""
74 
75  pids = Belle2.PyStoreArray('PIDLikelihoods')
76  for pid in pids:
77  track = pid.getRelatedFrom('Tracks')
78  if track:
79  mcpart = track.getRelatedTo('MCParticles')
80  try:
81  pdg = abs(mcpart.getPDG())
82  if pdg != 211 and pdg != 321:
83  continue
84  momentumVec = mcpart.getMomentum()
85  momentum = momentumVec.Mag()
86  if momentum > 3.5: # cut off
87  continue
88  costheta = momentumVec.CosTheta()
89  phi = momentumVec.Phi()
90  fitresult = track.getTrackFitResultWithClosestMass(Belle2.Const.pion)
91  if fitresult:
92  trackmomentum = fitresult.getMomentum().Mag()
93  else:
94  trackmomentum = 0.0
95 
96  # particle to compare with pions
97  selectedpart = Belle2.Const.kaon
98  pid_dedx = Belle2.Const.PIDDetectorSet(Belle2.Const.CDC)
99  pid_top = Belle2.Const.PIDDetectorSet(Belle2.Const.TOP)
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
103 
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
107 
108  self.data.lld_dedx = dedx_DLL
109  self.data.lld_top = top_DLL
110  self.data.p = momentum
111  self.data.phi = phi
112  self.data.costheta = costheta
113  self.data.trackmomentum = trackmomentum
114  self.data.iskaon = pdg == 321
115 
116  # Fill tree
117  self.file.cd()
118  self.tree.Fill()
119  except:
120 
121  # some tracks don't have an mcparticle (fixed now)
122  B2WARNING('problems with track <-> mcparticle relations')
123  event = Belle2.PyStoreObj('EventMetaData').obj().getEvent()
124  print('event: %d, track: %d' % (event, track.getArrayIndex()))
125 
126  def terminate(self):
127  """ Close the output file."""
128 
129  self.file.cd()
130  self.file.Write()
131  self.file.Close()
132 
133 # create path
134 main = create_path()
135 
136 # use the input file defined via command line
137 main.add_module(register_module('RootInput'))
138 
139 # add the python module defined above
140 main.add_module(TreeWriterModule())
141 
142 # process events and print call statistics
143 process(main)
144 print(statistics)
pid_ttree_writer.TreeWriterModule.file
file
Output ROOT file.
Definition: pid_ttree_writer.py:59
pid_ttree_writer.TreeWriterModule.terminate
def terminate(self)
Definition: pid_ttree_writer.py:126
pid_ttree_writer.TreeWriterModule.tree
tree
TTree for output data.
Definition: pid_ttree_writer.py:61
Belle2::PyStoreObj
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:69
pid_ttree_writer.TreeWriterModule.event
def event(self)
Definition: pid_ttree_writer.py:72
pid_ttree_writer.TreeWriterModule
Definition: pid_ttree_writer.py:46
Belle2::Const::RestrictedDetectorSet
A class for sets of detector IDs whose content is limited to restricted set of valid detector IDs.
Definition: Const.h:172
pid_ttree_writer.TreeWriterModule.data
data
Instance of EventData class.
Definition: pid_ttree_writer.py:63
Belle2::PyStoreArray
a (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:58
pid_ttree_writer.TreeWriterModule.__init__
def __init__(self)
Definition: pid_ttree_writer.py:53