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