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