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