Belle II Software development
neurotrigger.py
1
8
9import basf2
10import nntd
11
12
15
16# standard name for hardware tracks coming from the unpacker
17hwneurotracks = 'CDCTriggerNeuroTracks'
18hwneuroinput2dfindertracks = 'CDCTriggerNNInput2DFinderTracks'
19hwneuroinputsegmenthits = 'CDCTriggerNNInputSegmentHits'
20
21
22# software neurotracks, that were calculated from hardware 2d-inputtracks
23# and hardware stereo track segments coming over all 48cc within one event
24# OR
25# standard name for tracks resulting from the software neurotrigger
26# module when it is rerun on the hardware tracks
27hwsimneurotracks = 'TSimNeuroTracks'
28
29# software neurotracks, which were all simulated in software starting only
30# from real CDCHITs. Also simulated track segments and 2dtracks are needed therefore.
31simneurotracks_swtssw2d = 'TRGCDCNeuroTracks'
32simsegmenthits = 'SimSegmentHits'
33sim2dtracks_swts = 'TRGCDC2DFinderTracks'
34
35
38
39
40def filterTRG(path):
41 nullpath = basf2.create_path()
42 mfilter = basf2.register_module("CDCTriggerUnpacker", unpackNeuro=True)
43 path.add_module(mfilter)
44 mfilter.if_value('<1', nullpath)
45
46
47class casefilter(basf2.Module):
48 def initialize(self):
49 self.filterarrayname = "RecoTracks"
50 self.lowerthan = 0
51 self.higherthan = -1
52
53 def param(self, params):
54 for key, value in params.items():
55 setattr(self, key, value)
56
57 def event(self):
58 from ROOT import Belle2 # noqa
59 lower = self.lowerthan == 0 or bool(len(Belle2.PyStoreArray(self.filterarrayname)) < self.lowerthan)
60 higher = bool(len(Belle2.PyStoreArray(self.filterarrayname)) > self.higherthan)
61 self.return_value(lower and higher)
62
63
64class nnt_eventfilter(basf2.Module):
65 def initialize(self):
66 self.tracksegmentsname = hwneuroinputsegmenthits,
67 self.twodtracksname = hwneuroinput2dfindertracks,
68 self.neurotracksname = hwneurotracks,
69 self.recotracksname = "RecoTracks"
70 self.nullpath = basf2.create_path()
71
72 def event(self):
73 self.return_value(bool(self.neurotrack_allgoodquality()))
74 self.if_false(self.nullpath)
75
76 def hastrginfo(self):
77 from ROOT import Belle2 # noqa
78 return bool(Belle2.PyStoreArray(self.twodtracksname).getEntries() > 0)
79
80 def neurotrack_allgoodquality(self):
81 from ROOT import Belle2 # noqa
82 isgoodquality = True
83 for tr in Belle2.PyStoreArray("CDCTriggerNeuroTracks"):
84 if tr.getQualityVector() > 0:
85 isgoodquality = False
86 break
87 return isgoodquality
88
89
90class randommaker(basf2.Module):
91 def initialize(self):
92 self.counter = 0
93
94 def event(self):
95 # print("counter is " + str(self.counter))
96 if self.counter % 100 == 0:
97 # print("case 0")
98 self.return_value(0)
99 elif self.counter % 3 == 0:
100 self.return_value(1)
101 elif self.counter % 3 == 1:
102 self.return_value(2)
103 elif self.counter % 3 == 2:
104 self.return_value(3)
105 else:
106 print("some kind of error")
107 self.counter += 1
108
109
110def add_nnta_gzip_test_output(
111 path,
112 baseAnaFileName,
113 configFileName,
114 baseGzipFileName,
115 baseOutputFileName,
116 excludeBranchNames=[],
117 branchNames=[]):
118 # create 4 output paths:
119 outpaths = []
120 for x in range(4):
121 outpaths.append([".random_"+str(x), basf2.create_path()])
122 # add the randommaker module:
123 rm = basf2.register_module(randommaker())
124 path.add_module(rm)
125
126 # also add nnta module:
127
128 rm.if_value('==0', outpaths[0][1])
129 rm.if_value('==1', outpaths[1][1])
130 rm.if_value('==2', outpaths[2][1])
131 rm.if_value('==3', outpaths[3][1])
132 for p in outpaths:
133 nnta = basf2.register_module(nntd.nntd())
134 nnta.param({"filename": baseAnaFileName+p[0]+".pkl", "netname": "hardware"})
135 p[1].add_module(nnta)
136 p[1].add_module('CDCTriggerNeuroData',
137 # input and target arrays
138 NeuroTrackInputMode=False,
139 inputCollectionName=hwneuroinput2dfindertracks, # the hardware input tracks
140 # inputCollectionName=sim2dtracks_swts, # the software 2dtracks from real hits
141 # inputCollectionName='TRGCDC2DFinderTracks', # the mcparticle based 2d tracks
142 hitCollectionName=hwneuroinputsegmenthits, # simsegmenthits, #'SimSegmentHits',
143 EventTimeName="CDCTriggerNeuroETFT0",
144 targetCollectionName='RecoTracks',
145 trainOnRecoTracks=True,
146 gzipFilename=baseGzipFileName+p[0]+".gz",
147 configFileName=configFileName,
148 writeconfigFileName=baseGzipFileName.split("/gzip")[0]+"/"+configFileName.split('/')[-1],
149 )
150 if p[0] == ".random_3":
151 p[1].add_module("RootOutput", outputFileName=baseOutputFileName +
152 p[0]+".root", excludeBranchNames=excludeBranchNames, branchNames=branchNames)
153
154
155def add_nnta_gzip_output(path, baseAnaFileName, configFileName, baseGzipFileName):
156 # create 4 output paths:
157 outpaths = []
158 for x in range(4):
159 outpaths.append([".random_"+str(x), basf2.create_path()])
160 # add the randommaker module:
161 rm = basf2.register_module(randommaker())
162 path.add_module(rm)
163
164 # also add nnta module:
165
166 rm.if_value('==0', outpaths[0][1])
167 rm.if_value('==1', outpaths[1][1])
168 rm.if_value('==2', outpaths[2][1])
169 rm.if_value('==3', outpaths[3][1])
170 for p in outpaths:
171 nnta = basf2.register_module(nntd.nntd())
172 nnta.param({"filename": baseAnaFileName+p[0]+".pkl", "netname": "hardware"})
173 p[1].add_module(nnta)
174 p[1].add_module('CDCTriggerNeuroData',
175 # input and target arrays
176 NeuroTrackInputMode=False,
177 inputCollectionName=hwneuroinput2dfindertracks, # the hardware input tracks
178 # inputCollectionName=sim2dtracks_swts, # the software 2dtracks from real hits
179 # inputCollectionName='TRGCDC2DFinderTracks', # the mcparticle based 2d tracks
180 hitCollectionName=hwneuroinputsegmenthits, # simsegmenthits, #'SimSegmentHits',
181 EventTimeName="CDCTriggerNeuroETFT0",
182 targetCollectionName='RecoTracks',
183 trainOnRecoTracks=True,
184 gzipFilename=baseGzipFileName+p[0]+".gz",
185 configFileName=configFileName,
186 writeconfigFileName=baseGzipFileName.split("/gzip")[0]+"/"+configFileName.split('/')[-1],
187 )
188
189
190def add_all_output(
191 path,
192 baseOutputFileName,
193 baseAnaFileName,
194 configFileName,
195 baseGzipFileName,
196 excludeBranchNames=[],
197 branchNames=[]):
198 # create 4 output paths:
199 outpaths = []
200 for x in range(4):
201 outpaths.append([".random_"+str(x), basf2.create_path()])
202 # add the randommaker module:
203 rm = basf2.register_module(randommaker())
204 path.add_module(rm)
205
206 # also add nnta module:
207
208 rm.if_value('==0', outpaths[0][1])
209 rm.if_value('==1', outpaths[1][1])
210 rm.if_value('==2', outpaths[2][1])
211 rm.if_value('==3', outpaths[3][1])
212 for p in outpaths:
213 nnta = basf2.register_module(nntd.nntd())
214 nnta.param({"filename": baseAnaFileName+p[0]+".pkl", "netname": "hardware"})
215 p[1].add_module("RootOutput", outputFileName=baseOutputFileName +
216 p[0]+".root", excludeBranchNames=excludeBranchNames, branchNames=branchNames)
217 p[1].add_module(nnta)
218 p[1].add_module('CDCTriggerNeuroData',
219 # input and target arrays
220 NeuroTrackInputMode=False,
221 inputCollectionName=hwneuroinput2dfindertracks, # the hardware input tracks
222 # inputCollectionName=sim2dtracks_swts, # the software 2dtracks from real hits
223 # inputCollectionName='TRGCDC2DFinderTracks', # the mcparticle based 2d tracks
224 hitCollectionName=hwneuroinputsegmenthits, # simsegmenthits, #'SimSegmentHits',
225 EventTimeName="CDCTriggerNeuroETFT0",
226 targetCollectionName='RecoTracks',
227 trainOnRecoTracks=True,
228 gzipFilename=baseGzipFileName+p[0]+".gz",
229 configFileName=configFileName,
230 writeconfigFileName=baseGzipFileName.split("/gzip")[0]+"/"+configFileName.split('/')[-1],
231 )
232
233
234def add_nnta_root_output(path, baseOutputFileName, baseAnaFileName, excludeBranchNames=[], branchNames=[]):
235 # create 4 output paths:
236 outpaths = []
237 for x in range(4):
238 outpaths.append([".random_"+str(x), basf2.create_path()])
239 # add the randommaker module:
240 rm = basf2.register_module(randommaker())
241 path.add_module(rm)
242
243 # also add nnta module:
244
245 rm.if_value('==0', outpaths[0][1])
246 rm.if_value('==1', outpaths[1][1])
247 rm.if_value('==2', outpaths[2][1])
248 rm.if_value('==3', outpaths[3][1])
249 for p in outpaths:
250 nnta = basf2.register_module(nntd.nntd())
251 nnta.param({"filename": baseAnaFileName+p[0]+".pkl", "netname": "hardware"})
252 p[1].add_module("RootOutput", outputFileName=baseOutputFileName +
253 p[0]+".root", excludeBranchNames=excludeBranchNames, branchNames=branchNames)
254 p[1].add_module(nnta)
255
256
257def add_neuro_unpacker(path, debug_level=4, debugout=False, **kwargs):
258 #
259 unpacker = basf2.register_module('CDCTriggerUnpacker')
260 if debugout:
261 unpacker.logging.log_level = basf2.LogLevel.DEBUG
262 unpacker.logging.debug_level = debug_level
263 unpacker.logging.set_info(basf2.LogLevel.DEBUG, basf2.LogInfo.LEVEL | basf2.LogInfo.MESSAGE)
264 # size (number of words) of the Belle2Link header
265 unpacker.param('headerSize', 3)
266 # unpack the data from the 2D tracker and save its Bitstream
267 unpacker.param('unpackTracker2D', False)
268 # make CDCTriggerTrack and CDCTriggerSegmentHit objects from the 2D output
269 unpacker.param('decode2DFinderTrack', False)
270 # make CDCTriggerSegmentHit objects from the 2D input
271 unpacker.param('decode2DFinderInput', False)
272 unpacker.param('2DNodeId', [
273 [0x11000001, 0],
274 [0x11000001, 1],
275 [0x11000002, 0],
276 [0x11000002, 1]])
277 unpacker.param('NeuroNodeId', [
278 [0x11000005, 0],
279 [0x11000005, 1],
280 [0x11000006, 0],
281 [0x11000006, 1],
282 ])
283 if 'useDB' in kwargs:
284 unpacker.param('useDB', kwargs['useDB'])
285 else:
286 unpacker.param('useDB', True)
287
288 if 'sim13dt' in kwargs:
289 unpacker.param('sim13dt', kwargs['sim13dt'])
290 else:
291 unpacker.param('sim13dt', False)
292 unpacker.param('unpackNeuro', True)
293 unpacker.param('decodeNeuro', True)
294 path.add_module(unpacker)
295
296
297def add_neuro_2d_unpackers(path, debug_level=4, debugout=False, **kwargs):
298 #
299 unpacker = basf2.register_module('CDCTriggerUnpacker')
300 if debugout:
301 unpacker.logging.log_level = basf2.LogLevel.DEBUG
302 unpacker.logging.debug_level = debug_level
303 unpacker.logging.set_info(basf2.LogLevel.DEBUG, basf2.LogInfo.LEVEL | basf2.LogInfo.MESSAGE)
304 # size (number of words) of the Belle2Link header
305 unpacker.param('headerSize', 3)
306 # unpack the data from the 2D tracker and save its Bitstream
307 unpacker.param('unpackTracker2D', True)
308 # make CDCTriggerTrack and CDCTriggerSegmentHit objects from the 2D output
309 unpacker.param('decode2DFinderTrack', True)
310 # make CDCTriggerSegmentHit objects from the 2D input
311 unpacker.param('decode2DFinderInput', True)
312 unpacker.param('2DNodeId', [
313 [0x11000001, 0],
314 [0x11000001, 1],
315 [0x11000002, 0],
316 [0x11000002, 1],
317 ])
318 unpacker.param('NeuroNodeId', [
319 [0x11000005, 0],
320 [0x11000005, 1],
321 [0x11000006, 0],
322 [0x11000006, 1],
323 ])
324 if 'useDB' in kwargs:
325 unpacker.param('useDB', kwargs['useDB'])
326 else:
327 unpacker.param('useDB', True)
328
329 if 'sim13dt' in kwargs:
330 unpacker.param('sim13dt', kwargs['sim13dt'])
331 else:
332 unpacker.param('sim13dt', False)
333 unpacker.param('unpackNeuro', True)
334 unpacker.param('decodeNeuro', True)
335 path.add_module(unpacker)
336
337 if (('unpackDNN' in kwargs) and (kwargs['unpackDNN'])):
338 # DNN unpacker, using same module
339 DNN_unpacker = basf2.register_module('CDCTriggerUnpacker')
340 if debugout:
341 DNN_unpacker.logging.log_level = basf2.LogLevel.DEBUG
342 DNN_unpacker.logging.debug_level = debug_level
343 DNN_unpacker.logging.set_info(basf2.LogLevel.DEBUG, basf2.LogInfo.LEVEL | basf2.LogInfo.MESSAGE)
344
345 # size (number of words) of the Belle2Link header
346 DNN_unpacker.param('headerSize', 3)
347 # always use DB for DNNT
348 DNN_unpacker.param('useDB', True)
349 # sim 13 bit drift time if required
350 if 'sim13dt' in kwargs:
351 DNN_unpacker.param('sim13dt', kwargs['sim13dt'])
352 else:
353 DNN_unpacker.param('sim13dt', False)
354 # DNN board ID
355 DNN_unpacker.param('NeuroNodeId_pcie40', [
356 # UT4 3D
357 [0x10000001, 4],
358 [0x10000001, 5],
359 [0x10000001, 6],
360 [0x10000001, 7]])
361 # unpack DNN track, decodeing and related 2D track
362 DNN_unpacker.param('unpackNeuro', True)
363 DNN_unpacker.param('decodeNeuro', True)
364 DNN_unpacker.param('unpackMerger', True)
365 DNN_unpacker.param('unpackTracker2D', True)
366 DNN_unpacker.param('decode2DFinderTrack', True)
367 # DNN track default name
368 DNN_unpacker.param("NNTrackName", "CDCTriggerDNNTracks")
369 # DNN relate TS default name
370 DNN_unpacker.param("NNInStereoTS", "CDCTriggerDNNInputAllStereoSegmentHits")
371 # DNN select TS default name
372 DNN_unpacker.param("NNInSelectTS", "CDCTriggerDNNInputSegmentHits")
373 # DNN 2D input
374 DNN_unpacker.param("NNIn2DTrack", "CDCTriggerDNNInput2DFinderTracks")
375 # DNN ETF input
376 DNN_unpacker.param("NNInETFT0", "CDCTriggerDNNETFT0")
377 # DNN scaled input
378 DNN_unpacker.param("NNScaledInput", "CDCTriggerDNNTracksInput")
379 # DNNT bit map
380 DNN_unpacker.param("configName", "DNNT_Config")
381 # flag to use DNN unpacker
382 DNN_unpacker.param("isDNN", True)
383 path.add_module(DNN_unpacker)
384
385
386def add_neurotrigger_sim(path, nntweightfile=None, debug_level=4, debugout=False, **kwargs):
387 nnt = basf2.register_module('CDCTriggerNeuro')
388 if 'inputCollectionName' in kwargs:
389 nnt.param('inputCollectionName', kwargs['inputCollectionName'])
390 else:
391 nnt.param('inputCollectionName', hwneuroinput2dfindertracks)
392 if 'outputCollectionName' in kwargs:
393 nnt.param('outputCollectionName', kwargs['outputCollectionName'])
394 else:
395 nnt.param('outputCollectionName', hwsimneurotracks)
396 if 'hitCollectionName' in kwargs:
397 nnt.param('hitCollectionName', kwargs['hitCollectionName'])
398 else:
399 nnt.param('hitCollectionName', hwneuroinputsegmenthits)
400 if 'writeMLPinput' in kwargs:
401 nnt.param('writeMLPinput', kwargs['writeMLPinput'])
402 else:
403 nnt.param('writeMLPinput', True)
404 if 'fixedPoint' in kwargs:
405 nnt.param('fixedPoint', kwargs['fixedPoint'])
406 else:
407 nnt.param('fixedPoint', True)
408 if nntweightfile is not None:
409 nnt.param('filename', basf2.find_file(nntweightfile))
410
411 if 'et_option' in kwargs:
412 nnt.param('et_option', kwargs['et_option'])
413 if 'EventTimeName' in kwargs:
414 nnt.param('EventTimeName', kwargs['EventTimeName'])
415 if debugout:
416 nnt.logging.log_level = basf2.LogLevel.DEBUG
417 nnt.logging.debug_level = debug_level
418 path.add_module(nnt)
419
420
421def add_neurotrigger_hw(path, nntweightfile=None, debug_level=4, debugout=False, **kwargs):
422 nnt = basf2.register_module('CDCTriggerNeuro')
423 if 'inputCollectionName' in kwargs:
424 nnt.param('inputCollectionName', kwargs['inputCollectionName'])
425 else:
426 nnt.param('inputCollectionName', hwneurotracks)
427 if 'outputCollectionName' in kwargs:
428 nnt.param('outputCollectionName', kwargs['outputCollectionName'])
429 else:
430 nnt.param('outputCollectionName', hwsimneurotracks)
431 if 'hitCollectionName' in kwargs:
432 nnt.param('hitCollectionName', kwargs['hitCollectionName'])
433 else:
434 nnt.param('hitCollectionName', hwneuroinputsegmenthits)
435 if 'writeMLPinput' in kwargs:
436 nnt.param('writeMLPinput', kwargs['writeMLPinput'])
437 else:
438 nnt.param('writeMLPinput', True)
439 if 'fixedPoint' in kwargs:
440 nnt.param('fixedPoint', kwargs['fixedPoint'])
441 else:
442 nnt.param('fixedPoint', True)
443 if 'realinputCollectonName' in kwargs:
444 nnt.param('realinputCollectionName', kwargs['realinputCollectionName'])
445 else:
446 nnt.param('realinputCollectionName', hwneuroinput2dfindertracks)
447
448 if nntweightfile is not None:
449 nnt.param('filename', basf2.find_file(nntweightfile))
450 nnt.param('NeuroHWTrackInputMode', True)
451 if 'et_option' in kwargs:
452 nnt.param('et_option', kwargs['et_option'])
453 else:
454 nnt.param('et_option', 'etfhwin')
455
456 if 'EventTimeName' in kwargs:
457 nnt.param('EventTimeName', kwargs['EventTimeName'])
458 else:
459 nnt.param('EventTimeName', 'CDCTriggerNeuroETFT0')
460 if debugout:
461 nnt.logging.log_level = basf2.LogLevel.DEBUG
462 nnt.logging.debug_level = debug_level
463 path.add_module(nnt)
464
465
466def add_neuro_simulation(path, nntweightfile=None, **kwargs):
467 nnt = basf2.register_module('CDCTriggerNeuro')
468 tsf = basf2.register_module('CDCTriggerTSF')
469 if "InnerTSLUTFile" in kwargs:
470 tsf.param("InnerTSLUTFile", kwargs["InnerTSLUTFile"])
471 else:
472 tsf.param("InnerTSLUTFile", basf2.find_file("data/trg/cdc/innerLUT_Bkg_p0.70_b0.80.coe"))
473 if "OuterTSLUTFile" in kwargs:
474 tsf.param("OuterTSLUTFile", kwargs["OuterTSLUTFile"])
475 else:
476 tsf.param("OuterTSLUTFile", basf2.find_file("data/trg/cdc/outerLUT_Bkg_p0.70_b0.80.coe"))
477 tsf.param("TSHitCollectionName", simsegmenthits)
478 tsf.param("CDCHitCollectionName", "CDCHits")
479 if "relateAllHits" in kwargs:
480 tsf.param("relateAllHits", kwargs["relateAllHits"])
481 if "makeRecoLRTable" in kwargs:
482 tsf.param("makeRecoLRTable", kwargs["makeRecoLRTable"])
483 if "outerRecoLRTableFilename" in kwargs:
484 tsf.param("outerRecoLRTableFilename", kwargs["outerRecoLRTableFilename"])
485 if "innerRecoLRTableFilename" in kwargs:
486 tsf.param("innerRecoLRTableFilename", kwargs["innerRecoLRTableFilename"])
487 if "makeTrueLRTable" in kwargs:
488 tsf.param("makeTrueLRTable", kwargs["makeTrueLRTable"])
489 path.add_module(tsf)
490 path.add_module('CDCTrigger2DFinder',
491 minHits=4, minHitsShort=4, minPt=0.3,
492 hitCollectionName=simsegmenthits,
493 outputCollectionName=sim2dtracks_swts)
494 if nntweightfile is not None:
495 nnt.param('filename', basf2.find_file(nntweightfile))
496 if 'et_option' in kwargs:
497 nnt.param('et_option', kwargs['et_option'])
498 nnt.param('inputCollectionName', sim2dtracks_swts)
499 nnt.param('outputCollectionName', simneurotracks_swtssw2d)
500 nnt.param('hitCollectionName', simsegmenthits)
501 nnt.param('writeMLPinput', True)
502 nnt.param('fixedPoint', False)
503 path.add_module(nnt)
Belle2::PyStoreArray
A (simplified) python wrapper for StoreArray.
Definition PyStoreArray.h:73
neurotrigger.nnt_eventfilter
Definition neurotrigger.py:64
neurotrigger.nnt_eventfilter.neurotracksname
neurotracksname
Definition neurotrigger.py:68
neurotrigger.nnt_eventfilter.nullpath
nullpath
Definition neurotrigger.py:70
neurotrigger.nnt_eventfilter.tracksegmentsname
tracksegmentsname
Definition neurotrigger.py:66
neurotrigger.nnt_eventfilter.neurotrack_allgoodquality
neurotrack_allgoodquality(self)
Definition neurotrigger.py:80
neurotrigger.nnt_eventfilter.recotracksname
str recotracksname
Definition neurotrigger.py:69
neurotrigger.nnt_eventfilter.twodtracksname
twodtracksname
Definition neurotrigger.py:67