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};')
41
42from ROOT import EventDataRecoDigit # noqa
43
44
46 '''class to create reco digits ttree'''
47
48 def __init__(self):
49 """Initialize the module"""
50
51 super().__init__()
52
53 self.file = ROOT.TFile('../SVDValidationTTreeRecoDigit.root', 'recreate')
54
55 self.tree = ROOT.TTree('tree', 'Event data of SVD validation events')
56
57 self.data = EventDataRecoDigit()
58
59 # Declare tree branches
60 for key in EventDataRecoDigit.__dict__:
61 if '__' not in key:
62 formstring = '/F'
63 if isinstance(self.data.__getattribute__(key), int):
64 formstring = '/I'
65 self.tree.Branch(key, addressof(self.data, key), key + formstring)
66
67 def event(self):
68 """Take digits from SVDRecoDigits with at least one truehit and save needed information"""
69 digits = Belle2.PyStoreArray('SVDRecoDigits')
70 shaperDigits = Belle2.PyStoreArray('SVDShaperDigits')
71 for digit in digits:
72 # get the true hit from the related SVDShaperDigit
73 # it works because there is a 1-to-1 correspondence between
74 # ShaperDigits and RecoDigits
75 digit_truehits = shaperDigits[digit.getArrayIndex()].getRelationsTo('SVDTrueHits')
76
77 if len(digit_truehits) == 0:
78 continue
79
80 # find the trueHit with highest energy deposit (the "best" match)
81 energy = 0
82 bestTrueHitIndex = 0
83
84 for i, trueHit in enumerate(digit_truehits):
85 if trueHit.getEnergyDep() > energy:
86 energy = trueHit.getEnergyDep()
87 bestTrueHitIndex = i
88 bestTrueHit = digit_truehits[bestTrueHitIndex]
89
90 # Sensor identification
91 sensorID = digit.getSensorID()
92 self.data.sensor_id = int(sensorID)
93 sensorNum = sensorID.getSensorNumber()
94 self.data.sensor = sensorNum
95 layerNum = sensorID.getLayerNumber()
96 self.data.layer = layerNum
97 if (layerNum == 3):
98 sensorType = 1 # Barrel
99 else:
100 if (sensorNum == 1):
101 sensorType = 0
102 else:
103 sensorType = 1
104 self.data.sensor_type = sensorType
105 ladderNum = sensorID.getLadderNumber()
106 self.data.ladder = ladderNum
107 if digit.isUStrip():
108 self.data.strip_dir = 0
109 else:
110 self.data.strip_dir = 1
111 self.data.recodigit_charge = digit.getCharge()
112 self.data.recodigit_time = digit.getTime()
113 self.data.truehit_time = bestTrueHit.getGlobalTime()
114 # Fill tree
115 self.file.cd()
116 self.tree.Fill()
117
118 def terminate(self):
119 """Close the output file. """
120 self.file.cd()
121 self.file.Write()
122 self.file.Close()
A (simplified) python wrapper for StoreArray.