Belle II Software  release-05-01-25
readout.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
7 
8 
9 # we need to modify this to get data from different coppers
10 
11 # 4 2Ds
12 # hslb = (('11000001', 'b'),
13 # ('11000001', 'a'),
14 # ('11000002', 'a'),
15 # ('11000002', 'b'))
16 
17 # 2D0 and NN
18 
19 hslb = (('11000001', 'b'),
20  ('11000003', 'b'),
21  )
22 
23 
24 integrity_check = False
25 
26 from basf2 import *
27 import interactive
28 
29 from ROOT import Belle2
30 import numpy as np
31 import sys
32 # from testpkg.bitstring import BitArray
33 from bitstring import BitArray
34 import pickle
35 import re
36 
37 [steering, srootFile] = sys.argv[:2]
38 if len(sys.argv) >= 3:
39  pickleSigFile = sys.argv[2]
40  isPickleFile = pickleSigFile[pickleSigFile.rfind('.'):] == '.p'
41 else:
42  pickleSigFile = None
43  isPickleFile = False
44 
45 pickleIt = len(sys.argv) == 2 or isPickleFile
46 
47 
48 def finesse(x): return ord(x) - ord('a')
49 
50 
51 def copper(x): return int(x, 16)
52 
53 
54 hslb = [(copper(ele[0]), finesse(ele[1])) for ele in hslb]
55 
56 
57 def join(ary):
58  return BitArray([]).join([BitArray(uint=i, length=32) for i in ary])
59 
60 
61 data = []
62 meta = []
63 
64 
65 class MinModule(Module):
66 
67  """
68  Example module to drop into ipython and create some objects to look at.
69  If you just want to start IPython and create PyStoreArray etc.
70  interactively in your own steering file, the 'Interactive' module
71  might be of interest.
72  """
73 
74  def event(self):
75  """
76  reimplement Module::event()
77  """
78  self.return_value(0)
79  # function namespace caching
80  frombuffer = np.frombuffer
81  GetNodeID = Belle2.RawTRG.GetNodeID
82  GetDetectorNwords = Belle2.RawTRG.GetDetectorNwords
83  GetDetectorBuffer = Belle2.RawTRG.GetDetectorBuffer
84  trgary = Belle2.PyStoreArray("RawTRGs")
85 
86  evtmeta = Belle2.PyStoreObj("EventMetaData")
87  trgs = []
88 
89  # integrity check
90  if integrity_check:
91  print(len(trgary), 'copper(s)')
92  for evt in trgary:
93  for entry in range(evt.GetNumEntries()): # flattened. Usually only 1 entry
94  print('{:0x}'.format(GetNodeID(evt, entry)))
95  entrylist = []
96  for bid in range(4):
97  if (GetNodeID(evt, entry), bid) in hslb:
98  count = GetDetectorNwords(evt, entry, bid)
99  print(bid, "exist, ", count, "words.")
100  return
101  dataList = []
102  for evt in trgary:
103  # assuming smaller Copper ID comes first
104  for entry in range(evt.GetNumEntries()): # flattened. Usually only 1 entry
105  # print('{:0x}'.format(GetNodeID(evt, entry)))
106  entrylist = []
107  for bid in range(4):
108  if (GetNodeID(evt, entry), bid) not in hslb:
109  continue
110  count = GetDetectorNwords(evt, entry, bid)
111  if count == 0:
112  # continue
113  pass
114  bf = GetDetectorBuffer(evt, entry, bid)
115  ary = frombuffer(bf, np.uintc, count)
116  # obsolete: skipping dummy buffers
117  # if '{:0x}'.format(ary[0])[:4] == 'dddd':
118  # self.return_value(1)
119  # else:
120  # return
121  self.return_value(1)
122  dataList.append(join(ary))
123 
124  event = evtmeta.getEvent()
125  run = evtmeta.getRun()
126  subrun = evtmeta.getSubrun()
127  meta.append((event, run, subrun))
128  data.append(dataList)
129 
130 
131 # Set the log level to show only error and fatal messages
132 # set_log_level(LogLevel.ERROR)
133 set_log_level(LogLevel.INFO)
134 
135 # Create main path
136 main = create_path()
137 
138 # input
139 if srootFile[-5:] == 'sroot':
140  root_input = register_module('SeqRootInput')
141 else:
142  root_input = register_module('RootInput')
143 root_input.param('inputFileName', srootFile)
144 
145 prog = register_module('Progress')
146 
147 # Add modules to main path
148 main.add_module(root_input)
149 main.add_module(prog)
150 
151 readout = MinModule()
152 main.add_module(readout)
153 
154 emptypath = create_path()
155 readout.if_false(emptypath)
156 
157 # check signal file before processing
158 if not pickleIt:
159  import b2vcd_48
160  vcdFile = sys.argv[3] if len(sys.argv) >= 4 else re.sub(r'.+/', '', re.sub(r'sroot', 'vcd', srootFile))
161  with open(pickleSigFile) as fin:
162  evtsize = [int(width) for width in fin.readline().split()]
163  B2INFO('Interpreting B2L data format with dimension ' + str(evtsize))
164  atlas = b2vcd_48.makeAtlas(fin.read(), evtsize)
165 
166 # Process all events
167 process(main)
168 
169 if pickleIt:
170  pica = pickleSigFile if isPickleFile else re.sub(r'.+/', 'ana/', re.sub(r'sroot', 'p', srootFile))
171  wfp = open(pica, 'wb')
172  pickle.dump(data, wfp, protocol=2)
173  pickle.dump(meta, wfp, protocol=2)
174  wfp.close()
175 
176 print(statistics)
177 
178 if pickleIt:
179  B2INFO('Output pickle file ' + pica + ' saved.')
180 else:
181  b2vcd_48.writeVCD(meta, data, atlas, vcdFile, evtsize)
scripts.readout.MinModule
Definition: readout.py:65
Belle2::PyStoreObj
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:69
scripts.readout.MinModule.event
def event(self)
Definition: readout.py:74
Belle2::PyStoreArray
a (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:58