Belle II Software development
test_fei.py
1#!/usr/bin/env python3
2
3
10
11import basf2
12import basf2_mva
13import unittest
14import unittest.mock
15import b2test_utils
16import os
17import ROOT
18
19import b2bii
20
21import fei
22from fei.config import Particle
23
24import numpy as np
25
26import pdg
27
28# @cond
29
30# Define equality operators for a bunch of pybasf2 classes
31import pybasf2
32pybasf2.Module.__eq__ = lambda a, b: a.type() == b.type() and\
33 all(x == y for x, y in zip(a.available_params(), b.available_params()))
34pybasf2.ModuleParamInfo.__eq__ = lambda a, b: a.name == b.name and a.values == b.values
35pybasf2.Path.__eq__ = lambda a, b: all(x == y for x, y in zip(a.modules(), b.modules()))
36
37
38def print_path(a, b):
39 """
40 Print the parts of the paths which are different
41 """
42 for x, y in zip(a.modules(), b.modules()):
43 if x.type() != y.type():
44 print(x.type(), y.type())
45 for n, m in zip(x.available_params(), y.available_params()):
46 if n.name != m.name:
47 print(n.name, m.name)
48 if n.values != m.values:
49 print(n.name, n.values, m.values)
50
51
52def get_small_unittest_channels():
53 pion = Particle('pi+',
54 fei.config.MVAConfiguration(variables=['p', 'dr'],
55 target='isPrimarySignal'),
56 fei.config.PreCutConfiguration(userCut='[dr < 2] and [abs(dz) < 4]',
57 bestCandidateMode='highest',
58 bestCandidateVariable='piid',
59 bestCandidateCut=20),
60 fei.config.PostCutConfiguration(bestCandidateCut=10, value=0.01))
61 pion.addChannel(['pi+:FSP'])
62
63 kaon = Particle('K+',
64 fei.config.MVAConfiguration(variables=['p', 'dr'],
65 target='isPrimarySignal'),
66 fei.config.PreCutConfiguration(userCut='[dr < 2] and [abs(dz) < 4]',
67 bestCandidateMode='highest',
68 bestCandidateVariable='Kid',
69 bestCandidateCut=20),
70 fei.config.PostCutConfiguration(bestCandidateCut=10, value=0.01))
71 kaon.addChannel(['K+:FSP'])
72
73 D0 = Particle('D0',
74 fei.config.MVAConfiguration(variables=['M', 'p'],
75 target='isSignal'),
76 fei.config.PreCutConfiguration(userCut='1.7 < M < 1.95',
77 bestCandidateMode='lowest',
78 bestCandidateVariable='abs(dM)',
79 bestCandidateCut=20),
80 fei.config.PostCutConfiguration(bestCandidateCut=10, value=0.001))
81 D0.addChannel(['K-', 'pi+'])
82 D0.addChannel(['pi-', 'pi+'])
83
84 particles = [pion, kaon, D0]
85 return particles
86
87
88class TestGetStagesFromParticles(unittest.TestCase):
89 def test_get_stages_from_particles(self):
90 particles = fei.get_unittest_channels()
91 stages = fei.core.get_stages_from_particles(particles)
92 self.assertEqual(len(stages), 7)
93 self.assertEqual(len(stages[0]), 4)
94 self.assertEqual(stages[0][0].identifier, 'gamma:generic')
95 self.assertEqual(stages[0][1].identifier, 'mu+:generic')
96 self.assertEqual(stages[0][2].identifier, 'pi+:generic')
97 self.assertEqual(stages[0][3].identifier, 'K+:generic')
98 self.assertEqual(len(stages[1]), 1)
99 self.assertEqual(stages[1][0].identifier, 'pi0:generic')
100 self.assertEqual(len(stages[2]), 0)
101 self.assertEqual(len(stages[3]), 2)
102 self.assertEqual(stages[3][0].identifier, 'D0:generic')
103 self.assertEqual(stages[3][1].identifier, 'D0:semileptonic')
104 self.assertEqual(len(stages[4]), 0)
105 self.assertEqual(len(stages[5]), 0)
106 self.assertEqual(len(stages[6]), 0)
107
108
109class TestTrainingDataInformation(unittest.TestCase):
110
111 def tearDown(self):
112 if os.path.isfile('mcParticlesCount.root'):
113 os.remove('mcParticlesCount.root')
114
115 def test_reconstruct(self):
116 particles = fei.get_unittest_channels()
117 x = fei.core.TrainingDataInformation(particles)
118
119 path = basf2.create_path()
120 path.add_module('VariablesToHistogram', fileName='mcParticlesCount.root',
121 ignoreCommandLineOverride=True,
122 variables=[
123 ('NumberOfMCParticlesInEvent(321)', 100, -0.5, 99.5),
124 ('NumberOfMCParticlesInEvent(421)', 100, -0.5, 99.5),
125 ('NumberOfMCParticlesInEvent(13)', 100, -0.5, 99.5),
126 ('NumberOfMCParticlesInEvent(111)', 100, -0.5, 99.5),
127 ('NumberOfMCParticlesInEvent(211)', 100, -0.5, 99.5),
128 ('NumberOfMCParticlesInEvent(22)', 100, -0.5, 99.5)]
129 )
130 print_path(path, x.reconstruct())
131 self.assertEqual(x.reconstruct(), path)
132
133 def test_available(self):
134 particles = fei.get_unittest_channels()
135 x = fei.core.TrainingDataInformation(particles)
136 self.assertEqual(x.available(), False)
137 f = ROOT.TFile('mcParticlesCount.root', 'RECREATE') # noqa
138 self.assertEqual(x.available(), True)
139
140 def test_get_mc_counts(self):
141 f = ROOT.TFile('mcParticlesCount.root', 'RECREATE')
142 f.cd()
143 hist = ROOT.TH1F("NumberOfMCParticlesInEvent__bo211__bc", "NumberOfMCParticlesInEvent__bo211__bc", 11, -0.5, 10.5)
144 for i in range(10):
145 hist.Fill(5)
146 for i in range(5):
147 hist.Fill(4)
148 for i in range(3):
149 hist.Fill(3)
150 f.Write("NumberOfMCParticlesInEvent__bo211__bc")
151
152 hist = ROOT.TH1F("NumberOfMCParticlesInEvent__bo321__bc", "NumberOfMCParticlesInEvent__bo321__bc", 11, -0.5, 10.5)
153 for i in range(8):
154 hist.Fill(4)
155 for i in range(5):
156 hist.Fill(2)
157 for i in range(5):
158 hist.Fill(7)
159 f.Write("NumberOfMCParticlesInEvent__bo321__bc")
160
161 hist = ROOT.TH1F("NumberOfMCParticlesInEvent__bo13__bc", "NumberOfMCParticlesInEvent__bo13__bc", 11, -0.5, 10.5)
162 for i in range(18):
163 hist.Fill(5)
164 f.Write("NumberOfMCParticlesInEvent__bo13__bc")
165
166 hist = ROOT.TH1F("NumberOfMCParticlesInEvent__bo22__bc", "NumberOfMCParticlesInEvent__bo222__bc", 11, -0.5, 10.5)
167 for i in range(18):
168 hist.Fill(0)
169 f.Write("NumberOfMCParticlesInEvent__bo22__bc")
170
171 hist = ROOT.TH1F("NumberOfMCParticlesInEvent__bo111__bc", "NumberOfMCParticlesInEvent__bo111__bc", 11, -0.5, 10.5)
172 for i in range(5):
173 hist.Fill(5)
174 for i in range(10):
175 hist.Fill(4)
176 for i in range(3):
177 hist.Fill(3)
178 f.Write("NumberOfMCParticlesInEvent__bo111__bc")
179
180 hist = ROOT.TH1F("NumberOfMCParticlesInEvent__bo421__bc", "NumberOfMCParticlesInEvent__bo421__bc", 11, -0.5, 10.5)
181 for i in range(10):
182 hist.Fill(9)
183 for i in range(5):
184 hist.Fill(0)
185 for i in range(3):
186 hist.Fill(1)
187 f.Write("NumberOfMCParticlesInEvent__bo421__bc")
188
189 particles = fei.get_unittest_channels()
190 x = fei.core.TrainingDataInformation(particles)
191
192 mcCounts = {
193 211: {'max': 5.0, 'min': 3.0, 'sum': 79.0, 'avg': 4.3888888888888893, 'std': 0.7556372504852998},
194 321: {'max': 7.0, 'min': 2.0, 'sum': 77.0, 'avg': 4.2777777777777777, 'std': 1.8798804795209585},
195 13: {'max': 5.0, 'min': 5.0, 'sum': 90.0, 'avg': 5.0, 'std': 0.0},
196 22: {'max': 0.0, 'min': 0.0, 'sum': 0.0, 'avg': 0.0, 'std': 0.0},
197 111: {'max': 5.0, 'min': 3.0, 'sum': 74.0, 'avg': 4.1111111111111107, 'std': 0.6573421981221816},
198 421: {'max': 9.0, 'min': 0.0, 'sum': 93.0, 'avg': 5.166666666666667, 'std': 4.2979323194092087},
199 0: {'sum': 18.0}}
200 self.assertDictEqual(x.get_mc_counts(), mcCounts)
201
202
203class TestFSPLoader(unittest.TestCase):
204
205 def test_belle2_without_monitoring(self):
206 particles = get_small_unittest_channels()
207 config = fei.config.FeiConfiguration(monitor=False)
208 x = fei.core.FSPLoader(particles, config)
209
210 path = basf2.create_path()
211 fsps = ['K+:FSP', 'pi+:FSP', 'e+:FSP', 'mu+:FSP', 'gamma:FSP', 'p+:FSP', 'K_L0:FSP']
212 path.add_module('ParticleLoader', decayStrings=fsps, writeOut=True)
213 for fsp in fsps:
214 path.add_module('ParticleListManipulator', outputListName=fsp,
215 inputListNames=[fsp.split(':')[0] + ':all'], writeOut=True)
216 if 'gamma' in fsp:
217 path.add_module('ParticleSelector', decayString='gamma:FSP', cut='isFromECL')
218 path.add_module('ParticleLoader', decayStrings=['K_S0:V0 -> pi+ pi-'], writeOut=True)
219 path.add_module('ParticleLoader', decayStrings=['Lambda0:V0 -> p+ pi-'], writeOut=True)
220 path.add_module('ParticleLoader', decayStrings=['gamma:V0 -> e+ e-'], addDaughters=True, writeOut=True)
221 print_path(path, x.reconstruct())
222 self.assertEqual(x.reconstruct(), path)
223
224 def test_belle2_with_monitoring(self):
225 particles = get_small_unittest_channels()
226 config = fei.config.FeiConfiguration(monitor=True)
227 x = fei.core.FSPLoader(particles, config)
228
229 path = basf2.create_path()
230 fsps = ['K+:FSP', 'pi+:FSP', 'e+:FSP', 'mu+:FSP', 'gamma:FSP', 'p+:FSP', 'K_L0:FSP']
231 path.add_module('ParticleLoader', decayStrings=fsps, writeOut=True)
232 for fsp in fsps:
233 path.add_module('ParticleListManipulator', outputListName=fsp,
234 inputListNames=[fsp.split(':')[0] + ':all'], writeOut=True)
235 if 'gamma' in fsp:
236 path.add_module('ParticleSelector', decayString='gamma:FSP', cut='isFromECL')
237 path.add_module('ParticleLoader', decayStrings=['K_S0:V0 -> pi+ pi-'], writeOut=True)
238 path.add_module('ParticleLoader', decayStrings=['Lambda0:V0 -> p+ pi-'], writeOut=True)
239 path.add_module('ParticleLoader', decayStrings=['gamma:V0 -> e+ e-'], addDaughters=True, writeOut=True)
240 hist_variables = [(f'NumberOfMCParticlesInEvent({pdgcode})', 100, -0.5, 99.5)
241 for pdgcode in {11, 321, 211, 13, 22, 310, 2212, 130, 3122, 111}]
242 path.add_module('VariablesToHistogram', particleList='',
243 ignoreCommandLineOverride=True,
244 variables=hist_variables,
245 fileName='Monitor_FSPLoader.root')
246 print_path(path, x.reconstruct())
247 self.assertEqual(x.reconstruct(), path)
248
249 def test_belle1_without_monitoring(self):
250 particles = get_small_unittest_channels()
251 b2bii.setB2BII()
252 config = fei.config.FeiConfiguration(monitor=False)
253 x = fei.core.FSPLoader(particles, config)
254
255 path = basf2.create_path()
256 fsps = ['K+:FSP', 'pi+:FSP', 'e+:FSP', 'mu+:FSP', 'p+:FSP']
257 path.add_module('ParticleLoader', decayStrings=fsps, writeOut=True)
258 for fsp in fsps:
259 path.add_module('ParticleListManipulator', outputListName=fsp,
260 inputListNames=[fsp.split(':')[0] + ':all'], writeOut=True)
261 path.add_module('ParticleListManipulator', outputListName='gamma:FSP', inputListNames=['gamma:mdst'], writeOut=True)
262 path.add_module('ParticleCopier', inputListNames=['gamma:FSP'])
263 path.add_module('ParticleListManipulator', outputListName='K_S0:V0', inputListNames=['K_S0:mdst'], writeOut=True)
264 path.add_module('ParticleCopier', inputListNames=['K_S0:V0'])
265 path.add_module('ParticleListManipulator', outputListName='Lambda0:V0', inputListNames=['Lambda0:mdst'], writeOut=True)
266 path.add_module('ParticleCopier', inputListNames=['Lambda0:V0'])
267 path.add_module('ParticleListManipulator', outputListName='K_L0:FSP', inputListNames=['K_L0:mdst'], writeOut=True)
268 path.add_module('ParticleCopier', inputListNames=['K_L0:FSP'])
269 path.add_module('ParticleListManipulator', outputListName='pi0:FSP', inputListNames=['pi0:mdst'], writeOut=True)
270 path.add_module('ParticleCopier', inputListNames=['pi0:FSP'])
271 path.add_module('ParticleListManipulator', outputListName='gamma:V0', inputListNames=['gamma:v0mdst'], writeOut=True)
272 path.add_module('ParticleCopier', inputListNames=['gamma:V0'])
273 print_path(path, x.reconstruct())
274 self.assertEqual(x.reconstruct(), path)
275 b2bii.unsetB2BII()
276
277 def test_belle1_with_monitoring(self):
278 particles = get_small_unittest_channels()
279 b2bii.setB2BII()
280 config = fei.config.FeiConfiguration(monitor=True)
281 x = fei.core.FSPLoader(particles, config)
282
283 path = basf2.create_path()
284 fsps = ['K+:FSP', 'pi+:FSP', 'e+:FSP', 'mu+:FSP', 'p+:FSP']
285 path.add_module('ParticleLoader', decayStrings=fsps, writeOut=True)
286 for fsp in fsps:
287 path.add_module('ParticleListManipulator', outputListName=fsp,
288 inputListNames=[fsp.split(':')[0] + ':all'], writeOut=True)
289 path.add_module('ParticleListManipulator', outputListName='gamma:FSP', inputListNames=['gamma:mdst'], writeOut=True)
290 path.add_module('ParticleCopier', inputListNames=['gamma:FSP'])
291 path.add_module('ParticleListManipulator', outputListName='K_S0:V0', inputListNames=['K_S0:mdst'], writeOut=True)
292 path.add_module('ParticleCopier', inputListNames=['K_S0:V0'])
293 path.add_module('ParticleListManipulator', outputListName='Lambda0:V0', inputListNames=['Lambda0:mdst'], writeOut=True)
294 path.add_module('ParticleCopier', inputListNames=['Lambda0:V0'])
295 path.add_module('ParticleListManipulator', outputListName='K_L0:FSP', inputListNames=['K_L0:mdst'], writeOut=True)
296 path.add_module('ParticleCopier', inputListNames=['K_L0:FSP'])
297 path.add_module('ParticleListManipulator', outputListName='pi0:FSP', inputListNames=['pi0:mdst'], writeOut=True)
298 path.add_module('ParticleCopier', inputListNames=['pi0:FSP'])
299 path.add_module('ParticleListManipulator', outputListName='gamma:V0', inputListNames=['gamma:v0mdst'], writeOut=True)
300 path.add_module('ParticleCopier', inputListNames=['gamma:V0'])
301 hist_variables = [(f'NumberOfMCParticlesInEvent({pdgcode})', 100, -0.5, 99.5)
302 for pdgcode in {11, 321, 211, 13, 22, 310, 2212, 130, 3122, 111}]
303 path.add_module('VariablesToHistogram', particleList='',
304 ignoreCommandLineOverride=True,
305 variables=hist_variables,
306 fileName='Monitor_FSPLoader.root')
307 print_path(path, x.reconstruct())
308 self.assertEqual(x.reconstruct(), path)
309 b2bii.unsetB2BII()
310
311
312class TestTrainingData(unittest.TestCase):
313
314 def setUp(self):
315 self.mc_counts = {
316 211: {'sum': 79, 'avg': 4.3888888888888893, 'max': 5, 'min': 3, 'std': 0.75563725048530228},
317 321: {'sum': 77, 'avg': 4.2777777777777777, 'max': 7, 'min': 2, 'std': 1.8798804795209592},
318 421: {'sum': 93, 'avg': 5.166666666666667, 'max': 9, 'min': 0, 'std': 4.2979323194092087},
319 0: {'sum': 18}}
320
321 def test_without_monitoring(self):
322 particles = get_small_unittest_channels()
323 config = fei.config.FeiConfiguration(monitor=False)
324 x = fei.core.TrainingData(particles, config, self.mc_counts)
325
326 path = basf2.create_path()
327 path.add_module('VariablesToNtuple', fileName='training_input.root',
328 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('pi+:generic ==> pi+:FSP variables'),
329 ignoreCommandLineOverride=True,
330 variables=['p', 'dr', 'isPrimarySignal'],
331 particleList='pi+:generic_0', sampling=('isPrimarySignal', {}))
332 path.add_module('VariablesToNtuple', fileName='training_input.root',
333 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('K+:generic ==> K+:FSP variables'),
334 ignoreCommandLineOverride=True,
335 variables=['p', 'dr', 'isPrimarySignal'],
336 particleList='K+:generic_0', sampling=('isPrimarySignal', {}))
337 path.add_module(
338 'VariablesToNtuple',
339 fileName='training_input.root',
340 ignoreCommandLineOverride=True,
341 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> K-:generic pi+:generic variables'),
342 variables=[
343 'M',
344 'p',
345 'isSignal'],
346 particleList='D0:generic_0',
347 sampling=(
348 'isSignal',
349 {}))
350 path.add_module(
351 'VariablesToNtuple',
352 fileName='training_input.root',
353 ignoreCommandLineOverride=True,
354 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> pi-:generic pi+:generic variables'),
355 variables=[
356 'M',
357 'p',
358 'isSignal'],
359 particleList='D0:generic_1',
360 sampling=(
361 'isSignal',
362 {}))
363 print_path(path, x.reconstruct())
364 self.assertEqual(x.reconstruct(), path)
365
366 def test_with_monitoring(self):
367 particles = get_small_unittest_channels()
368 config = fei.config.FeiConfiguration(monitor=True)
369 x = fei.core.TrainingData(particles, config, self.mc_counts)
370
371 path = basf2.create_path()
372 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
373 ignoreCommandLineOverride=True,
374 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'p', 'dr', 'isPrimarySignal']),
375 variables_2d=fei.config.variables2binnings_2d([('p', 'isPrimarySignal'), ('dr', 'isPrimarySignal')]),
376 fileName='Monitor_TrainingData.root', directory='pi+:generic ==> pi+:FSP')
377 path.add_module('VariablesToNtuple', fileName='training_input.root',
378 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('pi+:generic ==> pi+:FSP variables'),
379 ignoreCommandLineOverride=True,
380 variables=['p', 'dr', 'isPrimarySignal'],
381 particleList='pi+:generic_0', sampling=('isPrimarySignal', {}))
382 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
383 ignoreCommandLineOverride=True,
384 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'p', 'dr', 'isPrimarySignal']),
385 variables_2d=fei.config.variables2binnings_2d([('p', 'isPrimarySignal'), ('dr', 'isPrimarySignal')]),
386 fileName='Monitor_TrainingData.root', directory='K+:generic ==> K+:FSP')
387 path.add_module('VariablesToNtuple', fileName='training_input.root',
388 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('K+:generic ==> K+:FSP variables'),
389 ignoreCommandLineOverride=True,
390 variables=['p', 'dr', 'isPrimarySignal'],
391 particleList='K+:generic_0', sampling=('isPrimarySignal', {}))
392 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
393 ignoreCommandLineOverride=True,
394 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'M', 'p', 'isSignal']),
395 variables_2d=fei.config.variables2binnings_2d([('M', 'isSignal'), ('p', 'isSignal')]),
396 fileName='Monitor_TrainingData.root', directory='D0:generic ==> K-:generic pi+:generic')
397 path.add_module(
398 'VariablesToNtuple',
399 fileName='training_input.root',
400 ignoreCommandLineOverride=True,
401 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> K-:generic pi+:generic variables'),
402 variables=[
403 'M',
404 'p',
405 'isSignal'],
406 particleList='D0:generic_0',
407 sampling=(
408 'isSignal',
409 {}))
410 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
411 ignoreCommandLineOverride=True,
412 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'M', 'p', 'isSignal']),
413 variables_2d=fei.config.variables2binnings_2d([('M', 'isSignal'), ('p', 'isSignal')]),
414 fileName='Monitor_TrainingData.root', directory='D0:generic ==> pi-:generic pi+:generic')
415 path.add_module(
416 'VariablesToNtuple',
417 fileName='training_input.root',
418 ignoreCommandLineOverride=True,
419 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> pi-:generic pi+:generic variables'),
420 variables=[
421 'M',
422 'p',
423 'isSignal'],
424 particleList='D0:generic_1',
425 sampling=(
426 'isSignal',
427 {}))
428 print_path(path, x.reconstruct())
429 self.assertEqual(x.reconstruct(), path)
430
431
432class TestPreReconstruction(unittest.TestCase):
433
434 def test_without_monitoring(self):
435 particles = get_small_unittest_channels()
436 config = fei.config.FeiConfiguration(monitor=False)
437 x = fei.core.PreReconstruction(particles, config)
438
439 path = basf2.create_path()
440 path.add_module('ParticleListManipulator', inputListNames=['pi+:FSP'], outputListName='pi+:generic_0',
441 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
442 path.add_module('VariablesToExtraInfo', particleList='pi+:generic_0', variables={'constant(0)': 'decayModeID'})
443 path.add_module('BestCandidateSelection', particleList='pi+:generic_0', variable='piid', selectLowest=False,
444 numBest=20, outputVariable='preCut_rank')
445
446 path.add_module('ParticleListManipulator', inputListNames=['K+:FSP'], outputListName='K+:generic_0',
447 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
448 path.add_module('VariablesToExtraInfo', particleList='K+:generic_0', variables={'constant(0)': 'decayModeID'})
449 path.add_module('BestCandidateSelection', particleList='K+:generic_0', variable='Kid', selectLowest=False,
450 numBest=20, outputVariable='preCut_rank')
451
452 path.add_module('ParticleCombiner', decayString='D0:generic_0 -> K-:generic pi+:generic', writeOut=True,
453 decayMode=0, cut='1.7 < M < 1.95')
454 path.add_module('BestCandidateSelection', particleList='D0:generic_0',
455 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
456 path.add_module('ParticleVertexFitter', listName='D0:generic_0', confidenceLevel=-2.0,
457 vertexFitter='KFit', fitType='vertex')
458
459 path.add_module('ParticleCombiner', decayString='D0:generic_1 -> pi-:generic pi+:generic', writeOut=True,
460 decayMode=1, cut='1.7 < M < 1.95')
461 path.add_module('BestCandidateSelection', particleList='D0:generic_1',
462 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
463 path.add_module('ParticleVertexFitter', listName='D0:generic_1', confidenceLevel=-2.0,
464 vertexFitter='KFit', fitType='vertex')
465
466 print_path(path, x.reconstruct())
467 self.assertEqual(x.reconstruct(), path)
468
469 def test_with_monitoring(self):
470 particles = get_small_unittest_channels()
471 config = fei.config.FeiConfiguration(monitor=True)
472 x = fei.core.PreReconstruction(particles, config)
473
474 path = basf2.create_path()
475 path.add_module('ParticleListManipulator', inputListNames=['pi+:FSP'], outputListName='pi+:generic_0',
476 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
477 path.add_module('VariablesToExtraInfo', particleList='pi+:generic_0', variables={'constant(0)': 'decayModeID'})
478 path.add_module('MCMatcherParticles', listName='pi+:generic_0')
479 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
480 variables=fei.config.variables2binnings(['piid', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal']),
481 variables_2d=fei.config.variables2binnings_2d([('piid', 'isPrimarySignal'),
482 ('piid', 'mcErrors'),
483 ('piid', 'mcParticleStatus')]),
484 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='pi+:generic ==> pi+:FSP',
485 ignoreCommandLineOverride=True)
486 path.add_module('BestCandidateSelection', particleList='pi+:generic_0', variable='piid', selectLowest=False,
487 numBest=20, outputVariable='preCut_rank')
488 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
489 variables=fei.config.variables2binnings(['piid', 'mcErrors', 'mcParticleStatus',
490 'isPrimarySignal', 'extraInfo(preCut_rank)']),
491 variables_2d=fei.config.variables2binnings_2d([('piid', 'isPrimarySignal'),
492 ('piid', 'mcErrors'),
493 ('piid', 'mcParticleStatus'),
494 ('extraInfo(preCut_rank)', 'isPrimarySignal'),
495 ('extraInfo(preCut_rank)', 'mcErrors'),
496 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
497 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='pi+:generic ==> pi+:FSP',
498 ignoreCommandLineOverride=True)
499 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
500 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
501 'isPrimarySignal']),
502 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isPrimarySignal'),
503 ('chiProb', 'mcErrors'),
504 ('chiProb', 'mcParticleStatus')]),
505 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='pi+:generic ==> pi+:FSP',
506 ignoreCommandLineOverride=True)
507
508 path.add_module('ParticleListManipulator', inputListNames=['K+:FSP'], outputListName='K+:generic_0',
509 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
510 path.add_module('VariablesToExtraInfo', particleList='K+:generic_0', variables={'constant(0)': 'decayModeID'})
511
512 path.add_module('MCMatcherParticles', listName='K+:generic_0')
513 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
514 variables=fei.config.variables2binnings(['Kid', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal']),
515 variables_2d=fei.config.variables2binnings_2d([('Kid', 'isPrimarySignal'),
516 ('Kid', 'mcErrors'),
517 ('Kid', 'mcParticleStatus')]),
518 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='K+:generic ==> K+:FSP',
519 ignoreCommandLineOverride=True)
520 path.add_module('BestCandidateSelection', particleList='K+:generic_0', variable='Kid', selectLowest=False,
521 numBest=20, outputVariable='preCut_rank')
522 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
523 variables=fei.config.variables2binnings(['Kid', 'mcErrors', 'mcParticleStatus',
524 'isPrimarySignal', 'extraInfo(preCut_rank)']),
525 variables_2d=fei.config.variables2binnings_2d([('Kid', 'isPrimarySignal'),
526 ('Kid', 'mcErrors'),
527 ('Kid', 'mcParticleStatus'),
528 ('extraInfo(preCut_rank)', 'isPrimarySignal'),
529 ('extraInfo(preCut_rank)', 'mcErrors'),
530 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
531 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='K+:generic ==> K+:FSP',
532 ignoreCommandLineOverride=True)
533 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
534 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
535 'isPrimarySignal']),
536 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isPrimarySignal'),
537 ('chiProb', 'mcErrors'),
538 ('chiProb', 'mcParticleStatus')]),
539 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='K+:generic ==> K+:FSP',
540 ignoreCommandLineOverride=True)
541
542 path.add_module('ParticleCombiner', decayString='D0:generic_0 -> K-:generic pi+:generic', writeOut=True,
543 decayMode=0, cut='1.7 < M < 1.95')
544 path.add_module('MCMatcherParticles', listName='D0:generic_0')
545 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
546 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus', 'isSignal']),
547 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
548 ('abs(dM)', 'mcErrors'),
549 ('abs(dM)', 'mcParticleStatus')]),
550 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='D0:generic ==> K-:generic pi+:generic',
551 ignoreCommandLineOverride=True)
552 path.add_module('BestCandidateSelection', particleList='D0:generic_0',
553 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
554 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
555 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus',
556 'isSignal', 'extraInfo(preCut_rank)']),
557 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
558 ('abs(dM)', 'mcErrors'),
559 ('abs(dM)', 'mcParticleStatus'),
560 ('extraInfo(preCut_rank)', 'isSignal'),
561 ('extraInfo(preCut_rank)', 'mcErrors'),
562 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
563 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='D0:generic ==> K-:generic pi+:generic',
564 ignoreCommandLineOverride=True)
565 path.add_module('ParticleVertexFitter', listName='D0:generic_0', confidenceLevel=-2.0,
566 vertexFitter='KFit', fitType='vertex')
567 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
568 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
569 'isSignal']),
570 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isSignal'),
571 ('chiProb', 'mcErrors'),
572 ('chiProb', 'mcParticleStatus')]),
573 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='D0:generic ==> K-:generic pi+:generic',
574 ignoreCommandLineOverride=True)
575
576 path.add_module('ParticleCombiner', decayString='D0:generic_1 -> pi-:generic pi+:generic', writeOut=True,
577 decayMode=1, cut='1.7 < M < 1.95')
578 path.add_module('MCMatcherParticles', listName='D0:generic_1')
579 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
580 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus', 'isSignal']),
581 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
582 ('abs(dM)', 'mcErrors'),
583 ('abs(dM)', 'mcParticleStatus')]),
584 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='D0:generic ==> pi-:generic pi+:generic',
585 ignoreCommandLineOverride=True)
586 path.add_module('BestCandidateSelection', particleList='D0:generic_1',
587 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
588 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
589 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus',
590 'isSignal', 'extraInfo(preCut_rank)']),
591 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
592 ('abs(dM)', 'mcErrors'),
593 ('abs(dM)', 'mcParticleStatus'),
594 ('extraInfo(preCut_rank)', 'isSignal'),
595 ('extraInfo(preCut_rank)', 'mcErrors'),
596 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
597 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='D0:generic ==> pi-:generic pi+:generic',
598 ignoreCommandLineOverride=True)
599 path.add_module('ParticleVertexFitter', listName='D0:generic_1', confidenceLevel=-2.0,
600 vertexFitter='KFit', fitType='vertex')
601 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
602 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
603 'isSignal']),
604 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isSignal'),
605 ('chiProb', 'mcErrors'),
606 ('chiProb', 'mcParticleStatus')]),
607 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='D0:generic ==> pi-:generic pi+:generic',
608 ignoreCommandLineOverride=True)
609
610 print_path(path, x.reconstruct())
611 self.assertEqual(x.reconstruct(), path)
612
613
614class TestPostReconstruction(unittest.TestCase):
615
616 def test_get_missing_channels(self):
617 pion = Particle('pi+:unittest', fei.config.MVAConfiguration(variables=['p', 'dr'], target='isPrimarySignal'))
618 pion.addChannel(['pi+:FSP'])
619 D0 = Particle('D0:unittest', fei.config.MVAConfiguration(variables=['M', 'p'], target='isSignal'))
620 D0.addChannel(['K-:unittest', 'pi+:unittest'])
621 D0.addChannel(['pi-:unittest', 'pi+:unittest'])
622 config = fei.config.FeiConfiguration(monitor=False, prefix="UNITTEST")
623 x = fei.core.PostReconstruction([pion, D0], config)
624
625 self.assertEqual(x.get_missing_channels(), ['pi+:unittest ==> pi+:FSP', 'D0:unittest ==> K-:unittest pi+:unittest',
626 'D0:unittest ==> pi-:unittest pi+:unittest'])
627 self.assertEqual(x.available(), False)
628
629 content = """
630 <?xml version="1.0" encoding="utf-8"?>
631 <method>Trivial</method>
632 <weightfile>{channel}.xml</weightfile>
633 <treename>tree</treename>
634 <target_variable>isSignal</target_variable>
635 <weight_variable>__weight__</weight_variable>
636 <signal_class>1</signal_class>
637 <max_events>0</max_events>
638 <number_feature_variables>1</number_feature_variables>
639 <variable0>M</variable0>
640 <number_spectator_variables>0</number_spectator_variables>
641 <number_data_files>1</number_data_files>
642 <datafile0>train.root</datafile0>
643 <Trivial_version>1</Trivial_version>
644 <Trivial_output>0</Trivial_output>
645 <signal_fraction>0.066082567</signal_fraction>
646 """
647
648 channel = 'D0:unittest ==> K-:unittest pi+:unittest'
649 with open(channel + ".xml", "w") as f:
650 f.write(content.format(channel=channel))
651 basf2_mva.upload(channel + ".xml", 'UNITTEST_' + channel)
652
653 self.assertEqual(x.get_missing_channels(), ['pi+:unittest ==> pi+:FSP',
654 'D0:unittest ==> pi-:unittest pi+:unittest'])
655 self.assertEqual(x.available(), False)
656
657 channel = 'D0:unittest ==> pi-:unittest pi+:unittest'
658 with open(channel + ".xml", "w") as f:
659 f.write(content.format(channel=channel))
660 basf2_mva.upload(channel + ".xml", 'UNITTEST_' + channel)
661
662 self.assertEqual(x.get_missing_channels(), ['pi+:unittest ==> pi+:FSP'])
663 self.assertEqual(x.available(), False)
664
665 channel = 'pi+:unittest ==> pi+:FSP'
666 with open(channel + ".xml", "w") as f:
667 f.write(content.format(channel=channel))
668 basf2_mva.upload(channel + ".xml", 'UNITTEST_' + channel)
669
670 self.assertEqual(x.get_missing_channels(), [])
671 self.assertEqual(x.available(), True)
672
673 def tearDown(self):
674 if os.path.isfile('pi+:unittest ==> pi+:FSP.xml'):
675 os.remove('pi+:unittest ==> pi+:FSP.xml')
676 if os.path.isfile('D0:unittest ==> pi-:unittest pi+:unittest.xml'):
677 os.remove('D0:unittest ==> pi-:unittest pi+:unittest.xml')
678 if os.path.isfile('D0:unittest ==> K-:unittest pi+:unittest.xml'):
679 os.remove('D0:unittest ==> K-:unittest pi+:unittest.xml')
680
681 def test_without_monitoring(self):
682 particles = get_small_unittest_channels()
683 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST')
684 x = fei.core.PostReconstruction(particles, config)
685
686 path = basf2.create_path()
687
688 path.add_module('MVAExpert', identifier='UNITTEST_pi+:generic ==> pi+:FSP', extraInfoName='SignalProbability',
689 listNames=['pi+:generic_0'])
690 path.add_module('ParticleListManipulator', outputListName='pi+:generic', inputListNames=['pi+:generic_0'],
691 variable='particleSource', writeOut=True)
692 path.add_module('ParticleSelector', decayString='pi+:generic', cut='0.01 < extraInfo(SignalProbability)')
693 path.add_module('BestCandidateSelection', particleList='pi+:generic', variable='extraInfo(SignalProbability)',
694 selectLowest=False, numBest=10, outputVariable='postCut_rank')
695 path.add_module('TagUniqueSignal', particleList='pi+:generic', target='isPrimarySignal',
696 extraInfoName='uniqueSignal')
697
698 path.add_module('MVAExpert', identifier='UNITTEST_K+:generic ==> K+:FSP', extraInfoName='SignalProbability',
699 listNames=['K+:generic_0'])
700 path.add_module('ParticleListManipulator', outputListName='K+:generic', inputListNames=['K+:generic_0'],
701 variable='particleSource', writeOut=True)
702 path.add_module('ParticleSelector', decayString='K+:generic', cut='0.01 < extraInfo(SignalProbability)')
703 path.add_module('BestCandidateSelection', particleList='K+:generic', variable='extraInfo(SignalProbability)',
704 selectLowest=False, numBest=10, outputVariable='postCut_rank')
705 path.add_module('TagUniqueSignal', particleList='K+:generic', target='isPrimarySignal',
706 extraInfoName='uniqueSignal')
707
708 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> K-:generic pi+:generic',
709 extraInfoName='SignalProbability', listNames=['D0:generic_0'])
710
711 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> pi-:generic pi+:generic',
712 extraInfoName='SignalProbability', listNames=['D0:generic_1'])
713
714 path.add_module('ParticleListManipulator', outputListName='D0:generic',
715 inputListNames=['D0:generic_0', 'D0:generic_1'], variable='particleSource',
716 writeOut=True)
717 path.add_module('ParticleSelector', decayString='D0:generic', cut='0.001 < extraInfo(SignalProbability)')
718 path.add_module('BestCandidateSelection', particleList='D0:generic', variable='extraInfo(SignalProbability)',
719 selectLowest=False, numBest=10, outputVariable='postCut_rank')
720 path.add_module('TagUniqueSignal', particleList='D0:generic', target='isSignal',
721 extraInfoName='uniqueSignal')
722
723 print_path(path, x.reconstruct())
724 self.assertEqual(x.reconstruct(), path)
725
726 def test_without_monitoring_training_mode(self):
727 particles = get_small_unittest_channels()
728 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST', training=True)
729 x = fei.core.PostReconstruction(particles, config)
730
731 path = basf2.create_path()
732
733 path.add_module('MVAExpert', identifier='pi+:generic ==> pi+:FSP.xml', extraInfoName='SignalProbability',
734 listNames=['pi+:generic_0'])
735 path.add_module('ParticleListManipulator', outputListName='pi+:generic', inputListNames=['pi+:generic_0'],
736 variable='particleSource', writeOut=True)
737 path.add_module('ParticleSelector', decayString='pi+:generic', cut='0.01 < extraInfo(SignalProbability)')
738 path.add_module('BestCandidateSelection', particleList='pi+:generic', variable='extraInfo(SignalProbability)',
739 selectLowest=False, numBest=10, outputVariable='postCut_rank')
740 path.add_module('TagUniqueSignal', particleList='pi+:generic', target='isPrimarySignal',
741 extraInfoName='uniqueSignal')
742
743 path.add_module('MVAExpert', identifier='K+:generic ==> K+:FSP.xml', extraInfoName='SignalProbability',
744 listNames=['K+:generic_0'])
745 path.add_module('ParticleListManipulator', outputListName='K+:generic', inputListNames=['K+:generic_0'],
746 variable='particleSource', writeOut=True)
747 path.add_module('ParticleSelector', decayString='K+:generic', cut='0.01 < extraInfo(SignalProbability)')
748 path.add_module('BestCandidateSelection', particleList='K+:generic', variable='extraInfo(SignalProbability)',
749 selectLowest=False, numBest=10, outputVariable='postCut_rank')
750 path.add_module('TagUniqueSignal', particleList='K+:generic', target='isPrimarySignal',
751 extraInfoName='uniqueSignal')
752
753 path.add_module('MVAExpert', identifier='D0:generic ==> K-:generic pi+:generic.xml',
754 extraInfoName='SignalProbability', listNames=['D0:generic_0'])
755
756 path.add_module('MVAExpert', identifier='D0:generic ==> pi-:generic pi+:generic.xml',
757 extraInfoName='SignalProbability', listNames=['D0:generic_1'])
758
759 path.add_module('ParticleListManipulator', outputListName='D0:generic',
760 inputListNames=['D0:generic_0', 'D0:generic_1'], variable='particleSource',
761 writeOut=True)
762 path.add_module('ParticleSelector', decayString='D0:generic', cut='0.001 < extraInfo(SignalProbability)')
763 path.add_module('BestCandidateSelection', particleList='D0:generic', variable='extraInfo(SignalProbability)',
764 selectLowest=False, numBest=10, outputVariable='postCut_rank')
765 path.add_module('TagUniqueSignal', particleList='D0:generic', target='isSignal',
766 extraInfoName='uniqueSignal')
767
768 print_path(path, x.reconstruct())
769 self.assertEqual(x.reconstruct(), path)
770
771 def test_with_monitoring(self):
772 particles = get_small_unittest_channels()
773 config = fei.config.FeiConfiguration(monitor=True, prefix='UNITTEST')
774 x = fei.core.PostReconstruction(particles, config)
775
776 path = basf2.create_path()
777
778 path.add_module('MVAExpert', identifier='UNITTEST_pi+:generic ==> pi+:FSP', extraInfoName='SignalProbability',
779 listNames=['pi+:generic_0'])
780
781 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
782 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
783 'extraInfo(SignalProbability)',
784 'isPrimarySignal', 'extraInfo(decayModeID)']),
785 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isPrimarySignal'),
786 ('extraInfo(SignalProbability)', 'mcErrors'),
787 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
788 ('extraInfo(decayModeID)', 'isPrimarySignal'),
789 ('extraInfo(decayModeID)', 'mcErrors'),
790 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
791 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='pi+:generic ==> pi+:FSP',
792 ignoreCommandLineOverride=True)
793 path.add_module('ParticleListManipulator', outputListName='pi+:generic', inputListNames=['pi+:generic_0'],
794 variable='particleSource', writeOut=True)
795 path.add_module('VariablesToHistogram', particleList='pi+:generic',
796 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
797 'extraInfo(SignalProbability)',
798 'isPrimarySignal', 'extraInfo(decayModeID)']),
799 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
800 ('extraInfo(decayModeID)', 'mcErrors'),
801 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
802 fileName='Monitor_PostReconstruction_BeforePostCut.root', directory='pi+:generic',
803 ignoreCommandLineOverride=True)
804 path.add_module('ParticleSelector', decayString='pi+:generic', cut='0.01 < extraInfo(SignalProbability)')
805 path.add_module('VariablesToHistogram', particleList='pi+:generic',
806 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
807 'extraInfo(SignalProbability)',
808 'isPrimarySignal', 'extraInfo(decayModeID)']),
809 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
810 ('extraInfo(decayModeID)', 'mcErrors'),
811 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
812 fileName='Monitor_PostReconstruction_BeforeRanking.root', directory='pi+:generic',
813 ignoreCommandLineOverride=True)
814 path.add_module('BestCandidateSelection', particleList='pi+:generic', variable='extraInfo(SignalProbability)',
815 selectLowest=False, numBest=10, outputVariable='postCut_rank')
816 path.add_module('TagUniqueSignal', particleList='pi+:generic', target='isPrimarySignal',
817 extraInfoName='uniqueSignal')
818 path.add_module('VariablesToHistogram', particleList='pi+:generic',
819 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
820 'extraInfo(SignalProbability)', 'isPrimarySignal',
821 'extraInfo(decayModeID)', 'extraInfo(postCut_rank)']),
822 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
823 ('extraInfo(decayModeID)', 'mcErrors'),
824 ('extraInfo(decayModeID)', 'mcParticleStatus'),
825 ('extraInfo(decayModeID)', 'extraInfo(postCut_rank)'),
826 ('isPrimarySignal', 'extraInfo(postCut_rank)'),
827 ('mcErrors', 'extraInfo(postCut_rank)'),
828 ('mcParticleStatus', 'extraInfo(postCut_rank)')]),
829 fileName='Monitor_PostReconstruction_AfterRanking.root', directory='pi+:generic',
830 ignoreCommandLineOverride=True)
831 path.add_module('VariablesToNtuple', fileName='Monitor_Final.root',
832 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('pi+:generic variables'),
833 ignoreCommandLineOverride=True,
834 variables=['extraInfo(SignalProbability)', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal',
835 'extraInfo(uniqueSignal)', 'extraInfo(decayModeID)'],
836 particleList='pi+:generic')
837
838 path.add_module('MVAExpert', identifier='UNITTEST_K+:generic ==> K+:FSP', extraInfoName='SignalProbability',
839 listNames=['K+:generic_0'])
840 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
841 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
842 'extraInfo(SignalProbability)',
843 'isPrimarySignal', 'extraInfo(decayModeID)']),
844 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isPrimarySignal'),
845 ('extraInfo(SignalProbability)', 'mcErrors'),
846 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
847 ('extraInfo(decayModeID)', 'isPrimarySignal'),
848 ('extraInfo(decayModeID)', 'mcErrors'),
849 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
850 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='K+:generic ==> K+:FSP',
851 ignoreCommandLineOverride=True)
852 path.add_module('ParticleListManipulator', outputListName='K+:generic', inputListNames=['K+:generic_0'],
853 variable='particleSource', writeOut=True)
854 path.add_module('VariablesToHistogram', particleList='K+:generic',
855 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
856 'extraInfo(SignalProbability)',
857 'isPrimarySignal', 'extraInfo(decayModeID)']),
858 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
859 ('extraInfo(decayModeID)', 'mcErrors'),
860 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
861 fileName='Monitor_PostReconstruction_BeforePostCut.root', directory='K+:generic',
862 ignoreCommandLineOverride=True)
863 path.add_module('ParticleSelector', decayString='K+:generic', cut='0.01 < extraInfo(SignalProbability)')
864 path.add_module('VariablesToHistogram', particleList='K+:generic',
865 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
866 'extraInfo(SignalProbability)',
867 'isPrimarySignal', 'extraInfo(decayModeID)']),
868 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
869 ('extraInfo(decayModeID)', 'mcErrors'),
870 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
871 fileName='Monitor_PostReconstruction_BeforeRanking.root', directory='K+:generic',
872 ignoreCommandLineOverride=True)
873 path.add_module('BestCandidateSelection', particleList='K+:generic', variable='extraInfo(SignalProbability)',
874 selectLowest=False, numBest=10, outputVariable='postCut_rank')
875 path.add_module('TagUniqueSignal', particleList='K+:generic', target='isPrimarySignal',
876 extraInfoName='uniqueSignal')
877 path.add_module('VariablesToHistogram', particleList='K+:generic',
878 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
879 'extraInfo(SignalProbability)', 'isPrimarySignal',
880 'extraInfo(decayModeID)', 'extraInfo(postCut_rank)']),
881 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
882 ('extraInfo(decayModeID)', 'mcErrors'),
883 ('extraInfo(decayModeID)', 'mcParticleStatus'),
884 ('extraInfo(decayModeID)', 'extraInfo(postCut_rank)'),
885 ('isPrimarySignal', 'extraInfo(postCut_rank)'),
886 ('mcErrors', 'extraInfo(postCut_rank)'),
887 ('mcParticleStatus', 'extraInfo(postCut_rank)')]),
888 fileName='Monitor_PostReconstruction_AfterRanking.root', directory='K+:generic',
889 ignoreCommandLineOverride=True)
890 path.add_module('VariablesToNtuple', fileName='Monitor_Final.root',
891 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('K+:generic variables'),
892 ignoreCommandLineOverride=True,
893 variables=['extraInfo(SignalProbability)', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal',
894 'extraInfo(uniqueSignal)', 'extraInfo(decayModeID)'],
895 particleList='K+:generic')
896
897 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> K-:generic pi+:generic',
898 extraInfoName='SignalProbability', listNames=['D0:generic_0'])
899 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
900 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
901 'extraInfo(SignalProbability)',
902 'isSignal', 'extraInfo(decayModeID)']),
903 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isSignal'),
904 ('extraInfo(SignalProbability)', 'mcErrors'),
905 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
906 ('extraInfo(decayModeID)', 'isSignal'),
907 ('extraInfo(decayModeID)', 'mcErrors'),
908 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
909 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='D0:generic ==> K-:generic pi+:generic',
910 ignoreCommandLineOverride=True)
911
912 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> pi-:generic pi+:generic',
913 extraInfoName='SignalProbability', listNames=['D0:generic_1'])
914 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
915 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
916 'extraInfo(SignalProbability)',
917 'isSignal', 'extraInfo(decayModeID)']),
918 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isSignal'),
919 ('extraInfo(SignalProbability)', 'mcErrors'),
920 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
921 ('extraInfo(decayModeID)', 'isSignal'),
922 ('extraInfo(decayModeID)', 'mcErrors'),
923 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
924 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='D0:generic ==> pi-:generic pi+:generic',
925 ignoreCommandLineOverride=True)
926
927 path.add_module('ParticleListManipulator', outputListName='D0:generic',
928 inputListNames=['D0:generic_0', 'D0:generic_1'], variable='particleSource',
929 writeOut=True)
930 path.add_module('VariablesToHistogram', particleList='D0:generic',
931 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
932 'extraInfo(SignalProbability)',
933 'isSignal', 'extraInfo(decayModeID)']),
934 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isSignal'),
935 ('extraInfo(decayModeID)', 'mcErrors'),
936 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
937 fileName='Monitor_PostReconstruction_BeforePostCut.root', directory='D0:generic',
938 ignoreCommandLineOverride=True)
939 path.add_module('ParticleSelector', decayString='D0:generic', cut='0.001 < extraInfo(SignalProbability)')
940 path.add_module('VariablesToHistogram', particleList='D0:generic',
941 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
942 'extraInfo(SignalProbability)',
943 'isSignal', 'extraInfo(decayModeID)']),
944 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isSignal'),
945 ('extraInfo(decayModeID)', 'mcErrors'),
946 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
947 fileName='Monitor_PostReconstruction_BeforeRanking.root', directory='D0:generic',
948 ignoreCommandLineOverride=True)
949 path.add_module('BestCandidateSelection', particleList='D0:generic', variable='extraInfo(SignalProbability)',
950 selectLowest=False, numBest=10, outputVariable='postCut_rank')
951 path.add_module('TagUniqueSignal', particleList='D0:generic', target='isSignal',
952 extraInfoName='uniqueSignal')
953 path.add_module('VariablesToHistogram', particleList='D0:generic',
954 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus',
955 'extraInfo(SignalProbability)', 'isSignal',
956 'extraInfo(decayModeID)', 'extraInfo(postCut_rank)']),
957 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isSignal'),
958 ('extraInfo(decayModeID)', 'mcErrors'),
959 ('extraInfo(decayModeID)', 'mcParticleStatus'),
960 ('extraInfo(decayModeID)', 'extraInfo(postCut_rank)'),
961 ('isSignal', 'extraInfo(postCut_rank)'),
962 ('mcErrors', 'extraInfo(postCut_rank)'),
963 ('mcParticleStatus', 'extraInfo(postCut_rank)')]),
964 fileName='Monitor_PostReconstruction_AfterRanking.root', directory='D0:generic',
965 ignoreCommandLineOverride=True)
966 path.add_module('VariablesToNtuple', fileName='Monitor_Final.root',
967 treeName=ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic variables'),
968 ignoreCommandLineOverride=True,
969 variables=['extraInfo(SignalProbability)', 'mcErrors', 'mcParticleStatus', 'isSignal',
970 'extraInfo(uniqueSignal)', 'extraInfo(decayModeID)'],
971 particleList='D0:generic')
972
973 print_path(path, x.reconstruct())
974 self.assertEqual(x.reconstruct(), path)
975
976
977class TestTeacher(unittest.TestCase):
978 def setUp(self):
979 fei.core.Teacher.MaximumNumberOfMVASamples = int(1e7)
980 fei.core.Teacher.MinimumNumberOfMVASamples = int(10)
981
982 f = ROOT.TFile('training_input.root', 'RECREATE')
983 f.cd()
984
985 tree = ROOT.TTree(
986 ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('pi+:generic ==> pi+:FSP variables'),
987 ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('pi+:generic ==> pi+:FSP variables'))
988 isSignal = np.zeros(1, dtype=float)
989 p = np.zeros(1, dtype=float)
990 pt = np.zeros(1, dtype=float)
991 tree.Branch('isPrimarySignal', isSignal, 'isPrimarySignal/D')
992 tree.Branch('p', p, 'p/D')
993 tree.Branch('dr', pt, 'dr/D')
994 for i in range(1000):
995 isSignal[0] = i % 2
996 p[0] = i
997 pt[0] = i * 2
998 tree.Fill()
999 tree.Write("", ROOT.TObject.kOverwrite)
1000
1001 # Test case, where tree not existent for K+:generic ==> K+:FSP, so skip creating it
1002
1003 tree = ROOT.TTree(
1004 ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> K-:generic pi+:generic variables'),
1005 ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> K-:generic pi+:generic variables'))
1006 isSignal = np.zeros(1, dtype=float)
1007 p = np.zeros(1, dtype=float)
1008 pt = np.zeros(1, dtype=float)
1009 tree.Branch('isSignal', isSignal, 'isSignal/D')
1010 tree.Branch('M', p, 'M/D')
1011 tree.Branch('p', pt, 'p/D')
1012 # Too few signal events here!
1013 for i in range(10):
1014 isSignal[0] = i % 2
1015 p[0] = i
1016 pt[0] = i * 2
1017 tree.Fill()
1018 tree.Write("", ROOT.TObject.kOverwrite)
1019
1020 tree = ROOT.TTree(
1021 ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> pi-:generic pi+:generic variables'),
1022 ROOT.Belle2.MakeROOTCompatible.makeROOTCompatible('D0:generic ==> pi-:generic pi+:generic variables'))
1023 isSignal = np.zeros(1, dtype=float)
1024 p = np.zeros(1, dtype=float)
1025 pt = np.zeros(1, dtype=float)
1026 tree.Branch('isSignal', isSignal, 'isSignal/D')
1027 tree.Branch('M', p, 'M/D')
1028 tree.Branch('p', pt, 'p/D')
1029 for i in range(1000):
1030 isSignal[0] = i % 2
1031 p[0] = i
1032 pt[0] = i * 2
1033 tree.Fill()
1034 tree.Write("", ROOT.TObject.kOverwrite)
1035
1036 def tearDown(self):
1037 if os.path.isfile('UNITTEST_TEACHER.xml'):
1038 os.remove('UNITTEST_TEACHER.xml')
1039 if os.path.isfile('UNITTEST_pi+:generic ==> pi+:FSP.xml'):
1040 os.remove('UNITTEST_pi+:generic ==> pi+:FSP.xml')
1041 if os.path.isfile('UNITTEST_K+:generic ==> K+:FSP.xml'):
1042 os.remove('UNITTEST_K+:generic ==> K+:FSP.xml')
1043 if os.path.isfile('UNITTEST_D0:generic ==> K-:generic pi+:generic.xml'):
1044 os.remove('UNITTEST_D0:generic ==> K-:generic pi+:generic.xml')
1045 if os.path.isfile('UNITTEST_D0:generic ==> pi-:generic pi+:generic.xml'):
1046 os.remove('UNITTEST_D0:generic ==> pi-:generic pi+:generic.xml')
1047 if os.path.isfile('training_input.root'):
1048 os.remove('training_input.root')
1049 if os.path.isfile('pi+:generic ==> pi+:FSP.xml'):
1050 os.remove('pi+:generic ==> pi+:FSP.xml')
1051 if os.path.isfile('D0:generic ==> pi-:generic pi+:generic.xml'):
1052 os.remove('D0:generic ==> pi-:generic pi+:generic.xml')
1053 if os.path.isfile('D0:generic ==> K-:generic pi+:generic.xml'):
1054 os.remove('D0:generic ==> K-:generic pi+:generic.xml')
1055 if os.path.isfile('TEACHER.xml'):
1056 os.remove('TEACHER.xml')
1057
1058 def test_create_fake_weightfile(self):
1059 self.assertEqual(os.path.isfile('UNITTEST_pi+:generic ==> pi+:FSP.xml'), False)
1060 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP.xml'), False)
1061 fei.core.Teacher.create_fake_weightfile('UNITTEST_pi+:generic ==> pi+:FSP')
1062 self.assertEqual(os.path.isfile('UNITTEST_pi+:generic ==> pi+:FSP.xml'), True)
1063 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP.xml'), True)
1064
1065 def test_upload(self):
1066 particles = get_small_unittest_channels()
1067 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST', externTeacher='basf2_mva_teacher')
1068 x = fei.core.Teacher(particles, config)
1069 fei.core.Teacher.create_fake_weightfile('TEACHER')
1070 self.assertEqual(basf2_mva.available('UNITTEST_TEACHER'), False)
1071 r = x.upload('TEACHER')
1072 self.assertEqual(basf2_mva.available('UNITTEST_TEACHER'), True)
1073 self.assertEqual(r, ('TEACHER.xml', 'UNITTEST_TEACHER'))
1074
1075 def test_without_monitoring(self):
1076 particles = get_small_unittest_channels()
1077 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST', externTeacher='basf2_mva_teacher')
1078 x = fei.core.Teacher(particles, config)
1079
1080 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP'), False)
1081 self.assertEqual(basf2_mva.available('UNITTEST_K+:generic ==> K+:FSP'), False)
1082 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> K-:generic pi+:generic'), False)
1083 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> pi-:generic pi+:generic'), False)
1084
1085 x.do_all_trainings()
1086
1087 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP'), True)
1088 self.assertEqual(basf2_mva.available('UNITTEST_K+:generic ==> K+:FSP'), False)
1089 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> K-:generic pi+:generic'), True)
1090 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> pi-:generic pi+:generic'), True)
1091
1092
1093"""
1094class TestConvertLegacyTraining(unittest.TestCase):
1095 pass
1096"""
1097
1098
1099class TestGetPath(unittest.TestCase):
1100
1101 def setUp(self):
1102 particles = fei.get_unittest_channels()
1103
1104 f = ROOT.TFile('mcParticlesCount.root', 'RECREATE')
1105 f.cd()
1106
1107 for pdgnumber in {abs(pdg.from_name(particle.name)) for particle in particles}:
1108 hist = ROOT.TH1F(f"NumberOfMCParticlesInEvent__bo{pdgnumber}__bc",
1109 f"NumberOfMCParticlesInEvent__bo{pdgnumber}__bc", 11, -0.5, 10.5)
1110 for i in range(10):
1111 hist.Fill(5)
1112 f.Write(f"NumberOfMCParticlesInEvent__bo{pdgnumber}__bc")
1113
1114 def tearDown(self):
1115 if os.path.isfile('mcParticlesCount.root'):
1116 os.remove('mcParticlesCount.root')
1117 if os.path.isfile('Summary.pickle'):
1118 os.remove('Summary.pickle')
1119
1120 def test_get_path_default_cache(self):
1121 particles = fei.get_unittest_channels()
1122 config = fei.config.FeiConfiguration(training=True)
1123 x = fei.core.Teacher(particles, config) # noqa
1124
1125 # Should try to create mcParticlesCount
1126 # -> Returns at stage 0
1127 feistate = fei.core.get_path(particles, config)
1128 self.assertEqual(feistate.stage, 0)
1129
1130 # Should try to create training data for FSPs
1131 # -> Returns at stage 1
1132 config = fei.config.FeiConfiguration(training=True)
1133 feistate = fei.core.get_path(particles, config)
1134 self.assertEqual(feistate.stage, 1)
1135
1136 # No weightfiles were created, hence the algorithm
1137 # cannot go forward and tries to create the training data again
1138 with self.assertRaises(RuntimeError):
1139 config = fei.config.FeiConfiguration(training=True)
1140 feistate = fei.core.get_path(particles, config)
1141
1142 # We create the FSP weightfiles by hand
1143 fei.core.Teacher.create_fake_weightfile('pi+:generic ==> pi+:FSP')
1144 fei.core.Teacher.create_fake_weightfile('K+:generic ==> K+:FSP')
1145 fei.core.Teacher.create_fake_weightfile('mu+:generic ==> mu+:FSP')
1146 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:FSP')
1147 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:V0')
1148
1149 # Should try to create training data for pi0
1150 # -> Returns at stage 2
1151 config = fei.config.FeiConfiguration(training=True)
1152 feistate = fei.core.get_path(particles, config)
1153 self.assertEqual(feistate.stage, 2)
1154
1155 # No weightfiles were created, hence the algorithm
1156 # cannot go forward and tries to create the training data again
1157 with self.assertRaises(RuntimeError):
1158 config = fei.config.FeiConfiguration(training=True)
1159 feistate = fei.core.get_path(particles, config)
1160
1161 # We create the pi0 weightfile by hand
1162 fei.core.Teacher.create_fake_weightfile('pi0:generic ==> gamma:generic gamma:generic')
1163
1164 # Should try to create training data for D0
1165 # -> Returns stage 4 (stage 3 is skipped because it only contains K_S0)
1166 config = fei.config.FeiConfiguration(training=True)
1167 feistate = fei.core.get_path(particles, config)
1168 self.assertEqual(feistate.stage, 4)
1169
1170 # We create the D0 weightfiles by hand
1171 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic')
1172 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic pi0:generic')
1173 fei.core.Teacher.create_fake_weightfile('D0:generic ==> pi-:generic pi+:generic')
1174 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic mu+:generic')
1175 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic pi0:generic mu+:generic')
1176
1177 # Unittest channels do not contain anymore stages,
1178 # -> Returns last stage + 1
1179 config = fei.config.FeiConfiguration(training=True)
1180 feistate = fei.core.get_path(particles, config)
1181 self.assertEqual(feistate.stage, 7)
1182
1183 # If we start at 0, we should be get the whole path
1184 # -> Returns last stage + 1
1185 config = fei.config.FeiConfiguration(cache=0, training=True)
1186 feistate = fei.core.get_path(particles, config)
1187 self.assertEqual(feistate.stage, 7)
1188
1189 # If training is False we should get the whole path without
1190 # checking if weightfiles exist.
1191 os.remove('pi+:generic ==> pi+:FSP.xml')
1192 os.remove('K+:generic ==> K+:FSP.xml')
1193 os.remove('mu+:generic ==> mu+:FSP.xml')
1194 os.remove('gamma:generic ==> gamma:FSP.xml')
1195 os.remove('gamma:generic ==> gamma:V0.xml')
1196 os.remove('pi0:generic ==> gamma:generic gamma:generic.xml')
1197 os.remove('D0:generic ==> K-:generic pi+:generic.xml')
1198 os.remove('D0:generic ==> K-:generic pi+:generic pi0:generic.xml')
1199 os.remove('D0:generic ==> pi-:generic pi+:generic.xml')
1200 os.remove('D0:semileptonic ==> K-:generic mu+:generic.xml')
1201 os.remove('D0:semileptonic ==> K-:generic pi0:generic mu+:generic.xml')
1202
1203 config = fei.config.FeiConfiguration(cache=0, training=False)
1204 feistate = fei.core.get_path(particles, config)
1205 self.assertEqual(feistate.stage, 7)
1206
1207 def test_get_path(self):
1208 particles = fei.get_unittest_channels()
1209 config = fei.config.FeiConfiguration(cache=-1, training=True)
1210 x = fei.core.Teacher(particles, config) # noqa
1211
1212 # Should try to create mcParticlesCount
1213 # -> Returns at stage 0
1214 feistate = fei.core.get_path(particles, config)
1215 self.assertEqual(feistate.stage, 0)
1216
1217 # Should try to create training data for FSPs
1218 # -> Returns at stage 1
1219 config = fei.config.FeiConfiguration(cache=0, training=True)
1220 feistate = fei.core.get_path(particles, config)
1221 self.assertEqual(feistate.stage, 1)
1222
1223 # No weightfiles were created, hence the algorithm
1224 # cannot go forward and tries to create the training data again
1225 with self.assertRaises(RuntimeError):
1226 config = fei.config.FeiConfiguration(cache=1, training=True)
1227 feistate = fei.core.get_path(particles, config)
1228
1229 # We create the FSP weightfiles by hand
1230 fei.core.Teacher.create_fake_weightfile('pi+:generic ==> pi+:FSP')
1231 fei.core.Teacher.create_fake_weightfile('K+:generic ==> K+:FSP')
1232 fei.core.Teacher.create_fake_weightfile('mu+:generic ==> mu+:FSP')
1233 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:FSP')
1234 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:V0')
1235
1236 # Should try to create training data for pi0
1237 # -> Returns at stage 2
1238 config = fei.config.FeiConfiguration(cache=1, training=True)
1239 feistate = fei.core.get_path(particles, config)
1240 self.assertEqual(feistate.stage, 2)
1241
1242 # No weightfiles were created, hence the algorithm
1243 # cannot go forward and tries to create the training data again
1244 with self.assertRaises(RuntimeError):
1245 config = fei.config.FeiConfiguration(cache=2, training=True)
1246 feistate = fei.core.get_path(particles, config)
1247
1248 # No weightfiles were created, hence the algorithm
1249 # cannot go forward and tries to create the training data again
1250 # This applies as well if the stage is even bigger
1251 with self.assertRaises(RuntimeError):
1252 config = fei.config.FeiConfiguration(cache=4, training=True)
1253 feistate = fei.core.get_path(particles, config)
1254
1255 # We create the pi0 weightfile by hand
1256 fei.core.Teacher.create_fake_weightfile('pi0:generic ==> gamma:generic gamma:generic')
1257
1258 # Should try to create training data for D0
1259 # -> Returns stage 4 (stage 3 is skipped because it only contains K_S0)
1260 config = fei.config.FeiConfiguration(cache=2, training=True)
1261 feistate = fei.core.get_path(particles, config)
1262 self.assertEqual(feistate.stage, 4)
1263
1264 # We create the D0 weightfiles by hand
1265 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic')
1266 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic pi0:generic')
1267 fei.core.Teacher.create_fake_weightfile('D0:generic ==> pi-:generic pi+:generic')
1268 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic mu+:generic')
1269 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic pi0:generic mu+:generic')
1270
1271 # Unittest channels do not contain anymore stages,
1272 # -> Returns last stage + 1
1273 config = fei.config.FeiConfiguration(cache=4, training=True)
1274 feistate = fei.core.get_path(particles, config)
1275 self.assertEqual(feistate.stage, 7)
1276
1277 # If we start at 0, we should be get the whole path
1278 # -> Returns last stage + 1
1279 config = fei.config.FeiConfiguration(cache=0, training=True)
1280 feistate = fei.core.get_path(particles, config)
1281 self.assertEqual(feistate.stage, 7)
1282
1283 # If training is False we should get the whole path without
1284 # checking if weightfiles exist.
1285 os.remove('pi+:generic ==> pi+:FSP.xml')
1286 os.remove('K+:generic ==> K+:FSP.xml')
1287 os.remove('mu+:generic ==> mu+:FSP.xml')
1288 os.remove('gamma:generic ==> gamma:FSP.xml')
1289 os.remove('gamma:generic ==> gamma:V0.xml')
1290 os.remove('pi0:generic ==> gamma:generic gamma:generic.xml')
1291 os.remove('D0:generic ==> K-:generic pi+:generic.xml')
1292 os.remove('D0:generic ==> K-:generic pi+:generic pi0:generic.xml')
1293 os.remove('D0:generic ==> pi-:generic pi+:generic.xml')
1294 os.remove('D0:semileptonic ==> K-:generic mu+:generic.xml')
1295 os.remove('D0:semileptonic ==> K-:generic pi0:generic mu+:generic.xml')
1296
1297 config = fei.config.FeiConfiguration(cache=0, training=False)
1298 feistate = fei.core.get_path(particles, config)
1299 self.assertEqual(feistate.stage, 7)
1300
1301
1302if __name__ == '__main__':
1303 with b2test_utils.clean_working_directory():
1304 basf2.conditions.testing_payloads = ['localdb/database.txt']
1305 unittest.main()
1306
1307# @endcond
b2bii.setB2BII
setB2BII()
Definition b2bii.py:21
b2bii.unsetB2BII
unsetB2BII()
Definition b2bii.py:25
b2test_utils.clean_working_directory
clean_working_directory()
Definition __init__.py:198
pdg.from_name
from_name(name)
Definition pdg.py:63