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