Belle II Software development
SVDValidationTTreeRecoTrack.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 clusters related to SVDRecoTracks, saving
16 in a ttree in a ROOT file.
17 </description>
18 <noexecute>SVD validation helper class</noexecute>
19</header>
20"""
21from cmath import nan
22from collections import OrderedDict
23
24import basf2 as b2
25
26# Some ROOT tools
27import ROOT
28from ROOT import Belle2
29from ROOT import gROOT, addressof
30
31# Define a ROOT struct to hold output data in the TTree
32gROOT.ProcessLine('struct EventDataRecoTrack {\
33 int sensor_id;\
34 int layer;\
35 int ladder;\
36 int sensor_type;\
37 int strip_dir;\
38 int matched;\
39 int cluster_truehits_number;\
40 float cluster_UVTimeDiff;\
41 float cluster_UUTimeDiff;\
42 float cluster_VVTimeDiff;\
43 int clusters_number;\
44};')
45from ROOT import EventDataRecoTrack # noqa
46
47
48class SVDValidationTTreeRecoTrack(b2.Module):
49 '''class to create the reco track ttree'''
50
51 def __init__(self):
52 """Initialize the module"""
53
54 super().__init__()
55
56 self.file = ROOT.TFile('../SVDValidationTTreeRecoTrack.root', 'recreate')
57
58 self.tree = ROOT.TTree('tree', 'Event data of SVD validation events')
59
60 self.data = EventDataRecoTrack()
61
62 # Declare tree branches
63 for key in EventDataRecoTrack.__dict__:
64 if '__' not in key:
65 formstring = '/F'
66 if isinstance(self.data.__getattribute__(key), int):
67 formstring = '/I'
68 self.tree.Branch(key, addressof(self.data, key), key + formstring)
69
70 def event(self):
71 """Take clusters from SVDRecoTracks with at least one truehit and save needed information"""
72 tracks = Belle2.PyStoreArray('SVDRecoTracks')
73 for track in tracks:
74 clusters = track.getRelationsWith('SVDClusters')
75 clusters_number = len(clusters)
76 dict_cluster = OrderedDict({'L3': [], 'L4': [], 'L5': [], 'L6': []}) # keys: layers, values: list of dicts
77 # clusters are ordered acording layers and sides, so we can read two at once
78 for i in range(0, len(clusters), 2):
79 c_U, c_V = clusters[i], clusters[i + 1] # read at once two clusters
80 cluster_U_truehits = c_U.getRelationsTo('SVDTrueHits') # SVDClustersToSVDTrueHits
81 U_id = c_U.getSensorID()
82 sensorNum_U = U_id.getSensorNumber()
83 layer_U = U_id.getLayerNumber()
84 if (layer_U == 3):
85 sensor_type_U = 1
86 else:
87 if (sensorNum_U == 1):
88 sensor_type_U = 0
89 else:
90 sensor_type_U = 1
91 cluster_U = {'sensor_id': int(U_id),
92 'layer': U_id.getLayerNumber(),
93 'ladder': U_id.getLadderNumber(),
94 'sensor_type': sensor_type_U,
95 'strip_dir': 0,
96 'matched': 1 if len(cluster_U_truehits) > 0 else 0,
97 'cluster_truehits_number': len(cluster_U_truehits),
98 'cluster_clsTime': c_U.getClsTime() if len(cluster_U_truehits) > 0 else nan,
99 'cluster_UVTimeDiff': nan,
100 'cluster_UUTimeDiff': nan,
101 'cluster_VVTimeDiff': nan,
102 'clusters_number': clusters_number}
103 cluster_V_truehits = c_V.getRelationsTo('SVDTrueHits') # SVDClustersToSVDTrueHits
104 V_id = c_V.getSensorID()
105 sensorNum_V = V_id.getSensorNumber()
106 layer_V = V_id.getLayerNumber()
107 if (layer_V == 3):
108 sensor_type_V = 1
109 else:
110 if (sensorNum_V == 1):
111 sensor_type_V = 0
112 else:
113 sensor_type_V = 1
114 cluster_V = {'sensor_id': int(V_id),
115 'layer': V_id.getLayerNumber(),
116 'ladder': V_id.getLadderNumber(),
117 'sensor_type': sensor_type_V,
118 'strip_dir': 1,
119 'matched': 1 if len(cluster_V_truehits) > 0 else 0,
120 'cluster_truehits_number': len(cluster_V_truehits),
121 'cluster_clsTime': c_V.getClsTime() if len(cluster_V_truehits) > 0 else nan,
122 'cluster_UVTimeDiff': c_U.getClsTime() - c_V.getClsTime()
123 if (len(cluster_U_truehits) > 0) and (len(cluster_V_truehits) > 0) else nan,
124 'cluster_UUTimeDiff': nan,
125 'cluster_VVTimeDiff': nan,
126 'clusters_number': clusters_number}
127 #
128 if U_id.getLayerNumber() == 3:
129 dict_cluster['L3'].append(cluster_U)
130 dict_cluster['L3'].append(cluster_V)
131 elif U_id.getLayerNumber() == 4:
132 dict_cluster['L4'].append(cluster_U)
133 dict_cluster['L4'].append(cluster_V)
134 elif U_id.getLayerNumber() == 5:
135 dict_cluster['L5'].append(cluster_U)
136 dict_cluster['L5'].append(cluster_V)
137 elif U_id.getLayerNumber() == 6:
138 dict_cluster['L6'].append(cluster_U)
139 dict_cluster['L6'].append(cluster_V)
140 else:
141 raise Exception('Incorrect number of layer')
142
143 for (i, layer) in enumerate(dict_cluster.items()):
144 if i < len(dict_cluster) - 1:
145 next_layer = list(dict_cluster.items())[i + 1]
146 # print(layer[0], next_layer[0], end='')
147 if (layer[0] == "L3" and next_layer[0] == "L4" and len(layer[1]) != 0 and len(next_layer[1]) != 0) or \
148 (layer[0] == "L4" and next_layer[0] == "L5" and len(layer[1]) != 0 and len(next_layer[1]) != 0) or \
149 (layer[0] == "L5" and next_layer[0] == "L6" and len(layer[1]) != 0 and len(next_layer[1]) != 0):
150 cluster_UUTimeDiff = \
151 layer[1][0]['cluster_clsTime'] - next_layer[1][0]['cluster_clsTime']
152 cluster_VVTimeDiff = \
153 layer[1][1]['cluster_clsTime'] - next_layer[1][1]['cluster_clsTime']
154 layer[1][0].update({'cluster_UUTimeDiff': cluster_UUTimeDiff})
155 layer[1][0].update({'cluster_VVTimeDiff': nan})
156 layer[1][1].update({'cluster_UUTimeDiff': nan})
157 layer[1][1].update({'cluster_VVTimeDiff': cluster_VVTimeDiff})
158 else:
159 # print("continue")
160 continue
161 # save all clusters
162 for layer in dict_cluster.items():
163 for c in layer[1]:
164 if len(c) != 0:
165 self.data.sensor_id = c['sensor_id']
166 self.data.layer = c['layer']
167 self.data.ladder = c['ladder']
168 self.data.sensor_type = c['sensor_type']
169 self.data.strip_dir = c['strip_dir']
170 self.data.matched = c['matched']
171 self.data.cluster_truehits_number = c['cluster_truehits_number']
172 self.data.cluster_UVTimeDiff = c['cluster_UVTimeDiff']
173 self.data.cluster_UUTimeDiff = c['cluster_UUTimeDiff']
174 self.data.cluster_VVTimeDiff = c['cluster_VVTimeDiff']
175 self.data.clusters_number = c['clusters_number']
176 else:
177 print(layer[0], ": empty list")
178 # Fill tree
179 self.file.cd()
180 self.tree.Fill()
181
182 def terminate(self):
183 """Close the output file. """
184 self.file.cd()
185 self.file.Write()
186 self.file.Close()
187
Belle2::PyStoreArray
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72