Belle II Software development
SVDValidationTTreeRecoDigit.py
1#!/usr/bin/env python3
2
3
10
11"""
12<header>
13 <description>
14 This module is used for the SVD validation.
15 It gets information about RecoDigits, saving
16 in a ttree in a ROOT file.
17 </description>
18 <noexecute>SVD validation helper class</noexecute>
19</header>
20"""
21
22import basf2 as b2
23
24# Some ROOT tools
25import ROOT
26from ROOT import Belle2
27from ROOT import gROOT, addressof
28
29# Define a ROOT struct to hold output data in the TTree
30gROOT.ProcessLine('struct EventDataRecoDigit {\
31 int sensor_id;\
32 int layer;\
33 int ladder;\
34 int sensor;\
35 int sensor_type;\
36 int strip_dir;\
37 float recodigit_charge;\
38 float recodigit_time;\
39 float truehit_time;\
40};')
41from ROOT import EventDataRecoDigit # noqa
42
43
44class SVDValidationTTreeRecoDigit(b2.Module):
45 '''class to create reco digitis ttree'''
46
47 def __init__(self):
48 """Initialize the module"""
49
50 super().__init__()
51
52 self.file = ROOT.TFile('../SVDValidationTTreeRecoDigit.root', 'recreate')
54 self.tree = ROOT.TTree('tree', 'Event data of SVD validation events')
56 self.data = EventDataRecoDigit()
58 # Declare tree branches
59 for key in EventDataRecoDigit.__dict__:
60 if '__' not in key:
61 formstring = '/F'
62 if isinstance(self.data.__getattribute__(key), int):
63 formstring = '/I'
64 self.tree.Branch(key, addressof(self.data, key), key + formstring)
65
66 def event(self):
67 """Take digits from SVDRecoDigits with at least one truehit and save needed information"""
68 digits = Belle2.PyStoreArray('SVDRecoDigits')
69 shaperDigits = Belle2.PyStoreArray('SVDShaperDigits')
70 for digit in digits:
71 # get the true hit from the related SVDShaperDigit
72 # it works because there is a 1-to-1 correspondence between
73 # ShaperDigits and RecoDigits
74 digit_truehits = shaperDigits[digit.getArrayIndex()].getRelationsTo('SVDTrueHits')
75
76 if len(digit_truehits) == 0:
77 continue
78
79 # find the trueHit with highest energy deposit (the "best" match)
80 energy = 0
81 bestTrueHitIndex = 0
82
83 for i, trueHit in enumerate(digit_truehits):
84 if trueHit.getEnergyDep() > energy:
85 energy = trueHit.getEnergyDep()
86 bestTrueHitIndex = i
87 bestTrueHit = digit_truehits[bestTrueHitIndex]
88
89 # Sensor identification
90 sensorID = digit.getSensorID()
91 self.data.sensor_id = int(sensorID)
92 sensorNum = sensorID.getSensorNumber()
93 self.data.sensor = sensorNum
94 layerNum = sensorID.getLayerNumber()
95 self.data.layer = layerNum
96 if (layerNum == 3):
97 sensorType = 1 # Barrel
98 else:
99 if (sensorNum == 1):
100 sensorType = 0
101 else:
102 sensorType = 1
103 self.data.sensor_type = sensorType
104 ladderNum = sensorID.getLadderNumber()
105 self.data.ladder = ladderNum
106 if digit.isUStrip():
107 self.data.strip_dir = 0
108 else:
109 self.data.strip_dir = 1
110 self.data.recodigit_charge = digit.getCharge()
111 self.data.recodigit_time = digit.getTime()
112 self.data.truehit_time = bestTrueHit.getGlobalTime()
113 # Fill tree
114 self.file.cd()
115 self.tree.Fill()
116
117 def terminate(self):
118 """Close the output file. """
119 self.file.cd()
120 self.file.Write()
121 self.file.Close()
122
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72