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