Belle II Software development
SVDValidationTTreeTrueHit.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 truehits, saving in a ttree in a ROOT file.
16 </description>
17 <noexecute>SVD validation helper class</noexecute>
18</header>
19"""
20import basf2 as b2
21
22# Some ROOT tools
23import ROOT
24from ROOT import Belle2 # noqa: make Belle2 namespace available
25from ROOT import gROOT, addressof
26
27# Define a ROOT struct to hold output data in the TTree
28gROOT.ProcessLine('struct EventDataTrueHit {\
29 int sensor_id;\
30 int layer;\
31 int ladder;\
32 int sensor;\
33 int sensor_type;\
34 int strip_dir;\
35 int reconstructed;\
36 };')
37
38from ROOT import EventDataTrueHit # noqa
39
40
41class SVDValidationTTreeTrueHit(b2.Module):
42 '''class to create the true hit ttree'''
43
44 def __init__(self):
45 """Initialize the module"""
46
47 super().__init__()
48
49 self.file = ROOT.TFile('../SVDValidationTTreeTrueHit.root', 'recreate')
50
51 self.tree = ROOT.TTree('tree', 'Event data of SVD validation events')
52
53 self.data = EventDataTrueHit()
54
55 # Declare tree branches
56 for key in EventDataTrueHit.__dict__:
57 if '__' not in key:
58 formstring = '/F'
59 if isinstance(self.data.__getattribute__(key), int):
60 formstring = '/I'
61 self.tree.Branch(key, addressof(self.data, key), key + formstring)
62
63 def event(self):
64 """ Start with truehits and use the relation to get the corresponding clusters """
65 svdtruehits = Belle2.PyStoreArray('SVDTrueHits')
66 for truehit in svdtruehits:
67 clusters = truehit.getRelationsFrom('SVDClusters')
68
69 if len(clusters) == 0:
70 self.data.reconstructed = 0
71 # Sensor identification
72 sensorID = truehit.getSensorID()
73 self.data.sensor_id = int(sensorID)
74 sensorNum = sensorID.getSensorNumber()
75 self.data.sensor = sensorNum
76 layerNum = sensorID.getLayerNumber()
77 self.data.layer = layerNum
78 if (layerNum == 3):
79 sensorType = 1
80 else:
81 if (sensorNum == 1):
82 sensorType = 0
83 else:
84 sensorType = 1
85 self.data.sensor_type = sensorType
86 ladderNum = sensorID.getLadderNumber()
87 self.data.ladder = ladderNum
88 self.data.strip_dir = -1
89 # Fill tree
90 self.file.cd()
91 self.tree.Fill()
92 else:
93 for cluster in clusters:
94 self.data.reconstructed = 1
95 # Sensor identification
96 sensorID = truehit.getSensorID()
97 self.data.sensor_id = int(sensorID)
98 sensorNum = sensorID.getSensorNumber()
99 self.data.sensor = sensorNum
100 layerNum = sensorID.getLayerNumber()
101 self.data.layer = layerNum
102 if (layerNum == 3):
103 sensorType = 1
104 else:
105 if (sensorNum == 1):
106 sensorType = 0
107 else:
108 sensorType = 1
109 self.data.sensor_type = sensorType
110 ladderNum = sensorID.getLadderNumber()
111 self.data.ladder = ladderNum
112 if cluster.isUCluster():
113 self.data.strip_dir = 0
114 else:
115 self.data.strip_dir = 1
116 # Fill tree
117 self.file.cd()
118 self.tree.Fill()
119
120 def terminate(self):
121 """Close the output file. """
122 self.file.cd()
123 self.file.Write()
124 self.file.Close()
Belle2::PyStoreArray
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72