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