Belle II Software  release-08-01-10
SVDValidationTTreeTrueHit.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 truehits, saving in a ttree in a ROOT file.
17  </description>
18  <noexecute>SVD validation helper class</noexecute>
19 </header>
20 """
21 import basf2 as b2
22 
23 # Some ROOT tools
24 import ROOT
25 from ROOT import Belle2 # make Belle2 namespace available
26 from ROOT import gROOT, addressof
27 
28 # Define a ROOT struct to hold output data in the TTree
29 gROOT.ProcessLine('struct EventDataTrueHit {\
30  int sensor_id;\
31  int layer;\
32  int ladder;\
33  int sensor;\
34  int sensor_type;\
35  int strip_dir;\
36  int reconstructed;\
37  };')
38 
39 from ROOT import EventDataTrueHit # noqa
40 
41 
42 class SVDValidationTTreeTrueHit(b2.Module):
43  '''class to create the true hit ttree'''
44 
45  def __init__(self):
46  """Initialize the module"""
47 
48  super(SVDValidationTTreeTrueHit, self).__init__()
49 
50  self.filefile = ROOT.TFile('../SVDValidationTTreeTrueHit.root', 'recreate')
51 
52  self.treetree = ROOT.TTree('tree', 'Event data of SVD validation events')
53 
54  self.datadata = EventDataTrueHit()
55 
56  # Declare tree branches
57  for key in EventDataTrueHit.__dict__:
58  if '__' not in key:
59  formstring = '/F'
60  if isinstance(self.datadata.__getattribute__(key), int):
61  formstring = '/I'
62  self.treetree.Branch(key, addressof(self.datadata, key), key + formstring)
63 
64  def event(self):
65  """ Start with truehits and use the relation to get the corresponding clusters """
66  svdtruehits = Belle2.PyStoreArray('SVDTrueHits')
67  for truehit in svdtruehits:
68  clusters = truehit.getRelationsFrom('SVDClusters')
69 
70  if len(clusters) == 0:
71  self.datadata.reconstructed = 0
72  # Sensor identification
73  sensorID = truehit.getSensorID()
74  self.datadata.sensor_id = int(sensorID)
75  sensorNum = sensorID.getSensorNumber()
76  self.datadata.sensor = sensorNum
77  layerNum = sensorID.getLayerNumber()
78  self.datadata.layer = layerNum
79  if (layerNum == 3):
80  sensorType = 1
81  else:
82  if (sensorNum == 1):
83  sensorType = 0
84  else:
85  sensorType = 1
86  self.datadata.sensor_type = sensorType
87  ladderNum = sensorID.getLadderNumber()
88  self.datadata.ladder = ladderNum
89  self.datadata.strip_dir = -1
90  # Fill tree
91  self.filefile.cd()
92  self.treetree.Fill()
93  else:
94  for cluster in clusters:
95  self.datadata.reconstructed = 1
96  # Sensor identification
97  sensorID = truehit.getSensorID()
98  self.datadata.sensor_id = int(sensorID)
99  sensorNum = sensorID.getSensorNumber()
100  self.datadata.sensor = sensorNum
101  layerNum = sensorID.getLayerNumber()
102  self.datadata.layer = layerNum
103  if (layerNum == 3):
104  sensorType = 1
105  else:
106  if (sensorNum == 1):
107  sensorType = 0
108  else:
109  sensorType = 1
110  self.datadata.sensor_type = sensorType
111  ladderNum = sensorID.getLadderNumber()
112  self.datadata.ladder = ladderNum
113  if cluster.isUCluster():
114  self.datadata.strip_dir = 0
115  else:
116  self.datadata.strip_dir = 1
117  # Fill tree
118  self.filefile.cd()
119  self.treetree.Fill()
120 
121  def terminate(self):
122  """Close the output file. """
123  self.filefile.cd()
124  self.filefile.Write()
125  self.filefile.Close()
Belle2::PyStoreArray
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72