Belle II Software light-2406-ragdoll
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', treeName='pi+:generic ==> pi+:FSP variables',
328 ignoreCommandLineOverride=True,
329 variables=['p', 'dr', 'isPrimarySignal'],
330 particleList='pi+:generic_0', sampling=('isPrimarySignal', {}))
331 path.add_module('VariablesToNtuple', fileName='training_input.root', treeName='K+:generic ==> K+:FSP variables',
332 ignoreCommandLineOverride=True,
333 variables=['p', 'dr', 'isPrimarySignal'],
334 particleList='K+:generic_0', sampling=('isPrimarySignal', {}))
335 path.add_module('VariablesToNtuple', fileName='training_input.root',
336 ignoreCommandLineOverride=True,
337 treeName='D0:generic ==> K-:generic pi+:generic variables',
338 variables=['M', 'p', 'isSignal'],
339 particleList='D0:generic_0', sampling=('isSignal', {}))
340 path.add_module('VariablesToNtuple', fileName='training_input.root',
341 ignoreCommandLineOverride=True,
342 treeName='D0:generic ==> pi-:generic pi+:generic variables',
343 variables=['M', 'p', 'isSignal'],
344 particleList='D0:generic_1', sampling=('isSignal', {}))
345 print_path(path, x.reconstruct())
346 self.assertEqual(x.reconstruct(), path)
347
348 def test_with_monitoring(self):
349 particles = get_small_unittest_channels()
350 config = fei.config.FeiConfiguration(monitor=True)
351 x = fei.core.TrainingData(particles, config, self.mc_counts)
352
353 path = basf2.create_path()
354 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
355 ignoreCommandLineOverride=True,
356 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'p', 'dr', 'isPrimarySignal']),
357 variables_2d=fei.config.variables2binnings_2d([('p', 'isPrimarySignal'), ('dr', 'isPrimarySignal')]),
358 fileName='Monitor_TrainingData.root', directory='pi+:generic ==> pi+:FSP')
359 path.add_module('VariablesToNtuple', fileName='training_input.root', treeName='pi+:generic ==> pi+:FSP variables',
360 ignoreCommandLineOverride=True,
361 variables=['p', 'dr', 'isPrimarySignal'],
362 particleList='pi+:generic_0', sampling=('isPrimarySignal', {}))
363 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
364 ignoreCommandLineOverride=True,
365 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'p', 'dr', 'isPrimarySignal']),
366 variables_2d=fei.config.variables2binnings_2d([('p', 'isPrimarySignal'), ('dr', 'isPrimarySignal')]),
367 fileName='Monitor_TrainingData.root', directory='K+:generic ==> K+:FSP')
368 path.add_module('VariablesToNtuple', fileName='training_input.root', treeName='K+:generic ==> K+:FSP variables',
369 ignoreCommandLineOverride=True,
370 variables=['p', 'dr', 'isPrimarySignal'],
371 particleList='K+:generic_0', sampling=('isPrimarySignal', {}))
372 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
373 ignoreCommandLineOverride=True,
374 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'M', 'p', 'isSignal']),
375 variables_2d=fei.config.variables2binnings_2d([('M', 'isSignal'), ('p', 'isSignal')]),
376 fileName='Monitor_TrainingData.root', directory='D0:generic ==> K-:generic pi+:generic')
377 path.add_module('VariablesToNtuple', fileName='training_input.root',
378 ignoreCommandLineOverride=True,
379 treeName='D0:generic ==> K-:generic pi+:generic variables',
380 variables=['M', 'p', 'isSignal'],
381 particleList='D0:generic_0', sampling=('isSignal', {}))
382 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
383 ignoreCommandLineOverride=True,
384 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'M', 'p', 'isSignal']),
385 variables_2d=fei.config.variables2binnings_2d([('M', 'isSignal'), ('p', 'isSignal')]),
386 fileName='Monitor_TrainingData.root', directory='D0:generic ==> pi-:generic pi+:generic')
387 path.add_module('VariablesToNtuple', fileName='training_input.root',
388 ignoreCommandLineOverride=True,
389 treeName='D0:generic ==> pi-:generic pi+:generic variables',
390 variables=['M', 'p', 'isSignal'],
391 particleList='D0:generic_1', sampling=('isSignal', {}))
392 print_path(path, x.reconstruct())
393 self.assertEqual(x.reconstruct(), path)
394
395
396class TestPreReconstruction(unittest.TestCase):
397
398 def test_without_monitoring(self):
399 particles = get_small_unittest_channels()
400 config = fei.config.FeiConfiguration(monitor=False)
401 x = fei.core.PreReconstruction(particles, config)
402
403 path = basf2.create_path()
404 path.add_module('ParticleListManipulator', inputListNames=['pi+:FSP'], outputListName='pi+:generic_0',
405 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
406 path.add_module('VariablesToExtraInfo', particleList='pi+:generic_0', variables={'constant(0)': 'decayModeID'})
407 path.add_module('BestCandidateSelection', particleList='pi+:generic_0', variable='piid', selectLowest=False,
408 numBest=20, outputVariable='preCut_rank')
409
410 path.add_module('ParticleListManipulator', inputListNames=['K+:FSP'], outputListName='K+:generic_0',
411 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
412 path.add_module('VariablesToExtraInfo', particleList='K+:generic_0', variables={'constant(0)': 'decayModeID'})
413 path.add_module('BestCandidateSelection', particleList='K+:generic_0', variable='Kid', selectLowest=False,
414 numBest=20, outputVariable='preCut_rank')
415
416 path.add_module('ParticleCombiner', decayString='D0:generic_0 -> K-:generic pi+:generic', writeOut=True,
417 decayMode=0, cut='1.7 < M < 1.95')
418 path.add_module('BestCandidateSelection', particleList='D0:generic_0',
419 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
420 path.add_module('ParticleVertexFitter', listName='D0:generic_0', confidenceLevel=-2.0,
421 vertexFitter='KFit', fitType='vertex')
422
423 path.add_module('ParticleCombiner', decayString='D0:generic_1 -> pi-:generic pi+:generic', writeOut=True,
424 decayMode=1, cut='1.7 < M < 1.95')
425 path.add_module('BestCandidateSelection', particleList='D0:generic_1',
426 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
427 path.add_module('ParticleVertexFitter', listName='D0:generic_1', confidenceLevel=-2.0,
428 vertexFitter='KFit', fitType='vertex')
429
430 print_path(path, x.reconstruct())
431 self.assertEqual(x.reconstruct(), path)
432
433 def test_with_monitoring(self):
434 particles = get_small_unittest_channels()
435 config = fei.config.FeiConfiguration(monitor=True)
436 x = fei.core.PreReconstruction(particles, config)
437
438 path = basf2.create_path()
439 path.add_module('ParticleListManipulator', inputListNames=['pi+:FSP'], outputListName='pi+:generic_0',
440 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
441 path.add_module('VariablesToExtraInfo', particleList='pi+:generic_0', variables={'constant(0)': 'decayModeID'})
442 path.add_module('MCMatcherParticles', listName='pi+:generic_0')
443 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
444 variables=fei.config.variables2binnings(['piid', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal']),
445 variables_2d=fei.config.variables2binnings_2d([('piid', 'isPrimarySignal'),
446 ('piid', 'mcErrors'),
447 ('piid', 'mcParticleStatus')]),
448 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='pi+:generic ==> pi+:FSP',
449 ignoreCommandLineOverride=True)
450 path.add_module('BestCandidateSelection', particleList='pi+:generic_0', variable='piid', selectLowest=False,
451 numBest=20, outputVariable='preCut_rank')
452 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
453 variables=fei.config.variables2binnings(['piid', 'mcErrors', 'mcParticleStatus',
454 'isPrimarySignal', 'extraInfo(preCut_rank)']),
455 variables_2d=fei.config.variables2binnings_2d([('piid', 'isPrimarySignal'),
456 ('piid', 'mcErrors'),
457 ('piid', 'mcParticleStatus'),
458 ('extraInfo(preCut_rank)', 'isPrimarySignal'),
459 ('extraInfo(preCut_rank)', 'mcErrors'),
460 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
461 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='pi+:generic ==> pi+:FSP',
462 ignoreCommandLineOverride=True)
463 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
464 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
465 'isPrimarySignal']),
466 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isPrimarySignal'),
467 ('chiProb', 'mcErrors'),
468 ('chiProb', 'mcParticleStatus')]),
469 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='pi+:generic ==> pi+:FSP',
470 ignoreCommandLineOverride=True)
471
472 path.add_module('ParticleListManipulator', inputListNames=['K+:FSP'], outputListName='K+:generic_0',
473 cut='[dr < 2] and [abs(dz) < 4]', writeOut=True)
474 path.add_module('VariablesToExtraInfo', particleList='K+:generic_0', variables={'constant(0)': 'decayModeID'})
475
476 path.add_module('MCMatcherParticles', listName='K+:generic_0')
477 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
478 variables=fei.config.variables2binnings(['Kid', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal']),
479 variables_2d=fei.config.variables2binnings_2d([('Kid', 'isPrimarySignal'),
480 ('Kid', 'mcErrors'),
481 ('Kid', 'mcParticleStatus')]),
482 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='K+:generic ==> K+:FSP',
483 ignoreCommandLineOverride=True)
484 path.add_module('BestCandidateSelection', particleList='K+:generic_0', variable='Kid', selectLowest=False,
485 numBest=20, outputVariable='preCut_rank')
486 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
487 variables=fei.config.variables2binnings(['Kid', 'mcErrors', 'mcParticleStatus',
488 'isPrimarySignal', 'extraInfo(preCut_rank)']),
489 variables_2d=fei.config.variables2binnings_2d([('Kid', 'isPrimarySignal'),
490 ('Kid', 'mcErrors'),
491 ('Kid', 'mcParticleStatus'),
492 ('extraInfo(preCut_rank)', 'isPrimarySignal'),
493 ('extraInfo(preCut_rank)', 'mcErrors'),
494 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
495 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='K+:generic ==> K+:FSP',
496 ignoreCommandLineOverride=True)
497 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
498 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
499 'isPrimarySignal']),
500 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isPrimarySignal'),
501 ('chiProb', 'mcErrors'),
502 ('chiProb', 'mcParticleStatus')]),
503 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='K+:generic ==> K+:FSP',
504 ignoreCommandLineOverride=True)
505
506 path.add_module('ParticleCombiner', decayString='D0:generic_0 -> K-:generic pi+:generic', writeOut=True,
507 decayMode=0, cut='1.7 < M < 1.95')
508 path.add_module('MCMatcherParticles', listName='D0:generic_0')
509 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
510 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus', 'isSignal']),
511 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
512 ('abs(dM)', 'mcErrors'),
513 ('abs(dM)', 'mcParticleStatus')]),
514 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='D0:generic ==> K-:generic pi+:generic',
515 ignoreCommandLineOverride=True)
516 path.add_module('BestCandidateSelection', particleList='D0:generic_0',
517 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
518 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
519 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus',
520 'isSignal', 'extraInfo(preCut_rank)']),
521 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
522 ('abs(dM)', 'mcErrors'),
523 ('abs(dM)', 'mcParticleStatus'),
524 ('extraInfo(preCut_rank)', 'isSignal'),
525 ('extraInfo(preCut_rank)', 'mcErrors'),
526 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
527 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='D0:generic ==> K-:generic pi+:generic',
528 ignoreCommandLineOverride=True)
529 path.add_module('ParticleVertexFitter', listName='D0:generic_0', confidenceLevel=-2.0,
530 vertexFitter='KFit', fitType='vertex')
531 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
532 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
533 'isSignal']),
534 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isSignal'),
535 ('chiProb', 'mcErrors'),
536 ('chiProb', 'mcParticleStatus')]),
537 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='D0:generic ==> K-:generic pi+:generic',
538 ignoreCommandLineOverride=True)
539
540 path.add_module('ParticleCombiner', decayString='D0:generic_1 -> pi-:generic pi+:generic', writeOut=True,
541 decayMode=1, cut='1.7 < M < 1.95')
542 path.add_module('MCMatcherParticles', listName='D0:generic_1')
543 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
544 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus', 'isSignal']),
545 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
546 ('abs(dM)', 'mcErrors'),
547 ('abs(dM)', 'mcParticleStatus')]),
548 fileName='Monitor_PreReconstruction_BeforeRanking.root', directory='D0:generic ==> pi-:generic pi+:generic',
549 ignoreCommandLineOverride=True)
550 path.add_module('BestCandidateSelection', particleList='D0:generic_1',
551 variable='abs(dM)', selectLowest=True, numBest=20, outputVariable='preCut_rank')
552 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
553 variables=fei.config.variables2binnings(['abs(dM)', 'mcErrors', 'mcParticleStatus',
554 'isSignal', 'extraInfo(preCut_rank)']),
555 variables_2d=fei.config.variables2binnings_2d([('abs(dM)', 'isSignal'),
556 ('abs(dM)', 'mcErrors'),
557 ('abs(dM)', 'mcParticleStatus'),
558 ('extraInfo(preCut_rank)', 'isSignal'),
559 ('extraInfo(preCut_rank)', 'mcErrors'),
560 ('extraInfo(preCut_rank)', 'mcParticleStatus')]),
561 fileName='Monitor_PreReconstruction_AfterRanking.root', directory='D0:generic ==> pi-:generic pi+:generic',
562 ignoreCommandLineOverride=True)
563 path.add_module('ParticleVertexFitter', listName='D0:generic_1', confidenceLevel=-2.0,
564 vertexFitter='KFit', fitType='vertex')
565 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
566 variables=fei.config.variables2binnings(['chiProb', 'mcErrors', 'mcParticleStatus',
567 'isSignal']),
568 variables_2d=fei.config.variables2binnings_2d([('chiProb', 'isSignal'),
569 ('chiProb', 'mcErrors'),
570 ('chiProb', 'mcParticleStatus')]),
571 fileName='Monitor_PreReconstruction_AfterVertex.root', directory='D0:generic ==> pi-:generic pi+:generic',
572 ignoreCommandLineOverride=True)
573
574 print_path(path, x.reconstruct())
575 self.assertEqual(x.reconstruct(), path)
576
577
578class TestPostReconstruction(unittest.TestCase):
579
580 def test_get_missing_channels(self):
581 pion = Particle('pi+:unittest', fei.config.MVAConfiguration(variables=['p', 'dr'], target='isPrimarySignal'))
582 pion.addChannel(['pi+:FSP'])
583 D0 = Particle('D0:unittest', fei.config.MVAConfiguration(variables=['M', 'p'], target='isSignal'))
584 D0.addChannel(['K-:unittest', 'pi+:unittest'])
585 D0.addChannel(['pi-:unittest', 'pi+:unittest'])
586 config = fei.config.FeiConfiguration(monitor=False, prefix="UNITTEST")
587 x = fei.core.PostReconstruction([pion, D0], config)
588
589 self.assertEqual(x.get_missing_channels(), ['pi+:unittest ==> pi+:FSP', 'D0:unittest ==> K-:unittest pi+:unittest',
590 'D0:unittest ==> pi-:unittest pi+:unittest'])
591 self.assertEqual(x.available(), False)
592
593 content = """
594 <?xml version="1.0" encoding="utf-8"?>
595 <method>Trivial</method>
596 <weightfile>{channel}.xml</weightfile>
597 <treename>tree</treename>
598 <target_variable>isSignal</target_variable>
599 <weight_variable>__weight__</weight_variable>
600 <signal_class>1</signal_class>
601 <max_events>0</max_events>
602 <number_feature_variables>1</number_feature_variables>
603 <variable0>M</variable0>
604 <number_spectator_variables>0</number_spectator_variables>
605 <number_data_files>1</number_data_files>
606 <datafile0>train.root</datafile0>
607 <Trivial_version>1</Trivial_version>
608 <Trivial_output>0</Trivial_output>
609 <signal_fraction>0.066082567</signal_fraction>
610 """
611
612 channel = 'D0:unittest ==> K-:unittest pi+:unittest'
613 with open(channel + ".xml", "w") as f:
614 f.write(content.format(channel=channel))
615 basf2_mva.upload(channel + ".xml", 'UNITTEST_' + channel)
616
617 self.assertEqual(x.get_missing_channels(), ['pi+:unittest ==> pi+:FSP',
618 'D0:unittest ==> pi-:unittest pi+:unittest'])
619 self.assertEqual(x.available(), False)
620
621 channel = 'D0:unittest ==> pi-:unittest pi+:unittest'
622 with open(channel + ".xml", "w") as f:
623 f.write(content.format(channel=channel))
624 basf2_mva.upload(channel + ".xml", 'UNITTEST_' + channel)
625
626 self.assertEqual(x.get_missing_channels(), ['pi+:unittest ==> pi+:FSP'])
627 self.assertEqual(x.available(), False)
628
629 channel = 'pi+:unittest ==> pi+:FSP'
630 with open(channel + ".xml", "w") as f:
631 f.write(content.format(channel=channel))
632 basf2_mva.upload(channel + ".xml", 'UNITTEST_' + channel)
633
634 self.assertEqual(x.get_missing_channels(), [])
635 self.assertEqual(x.available(), True)
636
637 def tearDown(self):
638 if os.path.isfile('pi+:unittest ==> pi+:FSP.xml'):
639 os.remove('pi+:unittest ==> pi+:FSP.xml')
640 if os.path.isfile('D0:unittest ==> pi-:unittest pi+:unittest.xml'):
641 os.remove('D0:unittest ==> pi-:unittest pi+:unittest.xml')
642 if os.path.isfile('D0:unittest ==> K-:unittest pi+:unittest.xml'):
643 os.remove('D0:unittest ==> K-:unittest pi+:unittest.xml')
644
645 def test_without_monitoring(self):
646 particles = get_small_unittest_channels()
647 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST')
648 x = fei.core.PostReconstruction(particles, config)
649
650 path = basf2.create_path()
651
652 path.add_module('MVAExpert', identifier='UNITTEST_pi+:generic ==> pi+:FSP', extraInfoName='SignalProbability',
653 listNames=['pi+:generic_0'])
654 path.add_module('TagUniqueSignal', particleList='pi+:generic_0', target='isPrimarySignal',
655 extraInfoName='uniqueSignal')
656 path.add_module('ParticleListManipulator', outputListName='pi+:generic', inputListNames=['pi+:generic_0'],
657 variable='particleSource', writeOut=True)
658 path.add_module('ParticleSelector', decayString='pi+:generic', cut='0.01 < extraInfo(SignalProbability)')
659 path.add_module('BestCandidateSelection', particleList='pi+:generic', variable='extraInfo(SignalProbability)',
660 selectLowest=False, numBest=10, outputVariable='postCut_rank')
661
662 path.add_module('MVAExpert', identifier='UNITTEST_K+:generic ==> K+:FSP', extraInfoName='SignalProbability',
663 listNames=['K+:generic_0'])
664 path.add_module('TagUniqueSignal', particleList='K+:generic_0', target='isPrimarySignal',
665 extraInfoName='uniqueSignal')
666 path.add_module('ParticleListManipulator', outputListName='K+:generic', inputListNames=['K+:generic_0'],
667 variable='particleSource', writeOut=True)
668 path.add_module('ParticleSelector', decayString='K+:generic', cut='0.01 < extraInfo(SignalProbability)')
669 path.add_module('BestCandidateSelection', particleList='K+:generic', variable='extraInfo(SignalProbability)',
670 selectLowest=False, numBest=10, outputVariable='postCut_rank')
671
672 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> K-:generic pi+:generic',
673 extraInfoName='SignalProbability', listNames=['D0:generic_0'])
674 path.add_module('TagUniqueSignal', particleList='D0:generic_0', target='isSignal',
675 extraInfoName='uniqueSignal')
676
677 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> pi-:generic pi+:generic',
678 extraInfoName='SignalProbability', listNames=['D0:generic_1'])
679 path.add_module('TagUniqueSignal', particleList='D0:generic_1', target='isSignal',
680 extraInfoName='uniqueSignal')
681
682 path.add_module('ParticleListManipulator', outputListName='D0:generic',
683 inputListNames=['D0:generic_0', 'D0:generic_1'], variable='particleSource',
684 writeOut=True)
685 path.add_module('ParticleSelector', decayString='D0:generic', cut='0.001 < extraInfo(SignalProbability)')
686 path.add_module('BestCandidateSelection', particleList='D0:generic', variable='extraInfo(SignalProbability)',
687 selectLowest=False, numBest=10, outputVariable='postCut_rank')
688
689 print_path(path, x.reconstruct())
690 self.assertEqual(x.reconstruct(), path)
691
692 def test_without_monitoring_training_mode(self):
693 particles = get_small_unittest_channels()
694 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST', training=True)
695 x = fei.core.PostReconstruction(particles, config)
696
697 path = basf2.create_path()
698
699 path.add_module('MVAExpert', identifier='pi+:generic ==> pi+:FSP.xml', extraInfoName='SignalProbability',
700 listNames=['pi+:generic_0'])
701 path.add_module('TagUniqueSignal', particleList='pi+:generic_0', target='isPrimarySignal',
702 extraInfoName='uniqueSignal')
703 path.add_module('ParticleListManipulator', outputListName='pi+:generic', inputListNames=['pi+:generic_0'],
704 variable='particleSource', writeOut=True)
705 path.add_module('ParticleSelector', decayString='pi+:generic', cut='0.01 < extraInfo(SignalProbability)')
706 path.add_module('BestCandidateSelection', particleList='pi+:generic', variable='extraInfo(SignalProbability)',
707 selectLowest=False, numBest=10, outputVariable='postCut_rank')
708
709 path.add_module('MVAExpert', identifier='K+:generic ==> K+:FSP.xml', extraInfoName='SignalProbability',
710 listNames=['K+:generic_0'])
711 path.add_module('TagUniqueSignal', particleList='K+:generic_0', target='isPrimarySignal',
712 extraInfoName='uniqueSignal')
713 path.add_module('ParticleListManipulator', outputListName='K+:generic', inputListNames=['K+:generic_0'],
714 variable='particleSource', writeOut=True)
715 path.add_module('ParticleSelector', decayString='K+:generic', cut='0.01 < extraInfo(SignalProbability)')
716 path.add_module('BestCandidateSelection', particleList='K+:generic', variable='extraInfo(SignalProbability)',
717 selectLowest=False, numBest=10, outputVariable='postCut_rank')
718
719 path.add_module('MVAExpert', identifier='D0:generic ==> K-:generic pi+:generic.xml',
720 extraInfoName='SignalProbability', listNames=['D0:generic_0'])
721 path.add_module('TagUniqueSignal', particleList='D0:generic_0', target='isSignal',
722 extraInfoName='uniqueSignal')
723
724 path.add_module('MVAExpert', identifier='D0:generic ==> pi-:generic pi+:generic.xml',
725 extraInfoName='SignalProbability', listNames=['D0:generic_1'])
726 path.add_module('TagUniqueSignal', particleList='D0:generic_1', target='isSignal',
727 extraInfoName='uniqueSignal')
728
729 path.add_module('ParticleListManipulator', outputListName='D0:generic',
730 inputListNames=['D0:generic_0', 'D0:generic_1'], variable='particleSource',
731 writeOut=True)
732 path.add_module('ParticleSelector', decayString='D0:generic', cut='0.001 < extraInfo(SignalProbability)')
733 path.add_module('BestCandidateSelection', particleList='D0:generic', variable='extraInfo(SignalProbability)',
734 selectLowest=False, numBest=10, outputVariable='postCut_rank')
735
736 print_path(path, x.reconstruct())
737 self.assertEqual(x.reconstruct(), path)
738
739 def test_with_monitoring(self):
740 particles = get_small_unittest_channels()
741 config = fei.config.FeiConfiguration(monitor=True, prefix='UNITTEST')
742 x = fei.core.PostReconstruction(particles, config)
743
744 path = basf2.create_path()
745
746 path.add_module('MVAExpert', identifier='UNITTEST_pi+:generic ==> pi+:FSP', extraInfoName='SignalProbability',
747 listNames=['pi+:generic_0'])
748 path.add_module('TagUniqueSignal', particleList='pi+:generic_0', target='isPrimarySignal',
749 extraInfoName='uniqueSignal')
750
751 path.add_module('VariablesToHistogram', particleList='pi+:generic_0',
752 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
753 'extraInfo(SignalProbability)',
754 'isPrimarySignal', 'extraInfo(decayModeID)']),
755 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isPrimarySignal'),
756 ('extraInfo(SignalProbability)', 'mcErrors'),
757 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
758 ('extraInfo(decayModeID)', 'isPrimarySignal'),
759 ('extraInfo(decayModeID)', 'mcErrors'),
760 ('extraInfo(decayModeID)', 'extraInfo(uniqueSignal)'),
761 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
762 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='pi+:generic ==> pi+:FSP',
763 ignoreCommandLineOverride=True)
764 path.add_module('ParticleListManipulator', outputListName='pi+:generic', inputListNames=['pi+:generic_0'],
765 variable='particleSource', writeOut=True)
766 path.add_module('VariablesToHistogram', particleList='pi+:generic',
767 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
768 'extraInfo(SignalProbability)',
769 'isPrimarySignal', 'extraInfo(decayModeID)']),
770 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
771 ('extraInfo(decayModeID)', 'mcErrors'),
772 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
773 fileName='Monitor_PostReconstruction_BeforePostCut.root', directory='pi+:generic',
774 ignoreCommandLineOverride=True)
775 path.add_module('ParticleSelector', decayString='pi+:generic', cut='0.01 < extraInfo(SignalProbability)')
776 path.add_module('VariablesToHistogram', particleList='pi+:generic',
777 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
778 'extraInfo(SignalProbability)',
779 'isPrimarySignal', 'extraInfo(decayModeID)']),
780 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
781 ('extraInfo(decayModeID)', 'mcErrors'),
782 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
783 fileName='Monitor_PostReconstruction_BeforeRanking.root', directory='pi+:generic',
784 ignoreCommandLineOverride=True)
785 path.add_module('BestCandidateSelection', particleList='pi+:generic', variable='extraInfo(SignalProbability)',
786 selectLowest=False, numBest=10, outputVariable='postCut_rank')
787 path.add_module('VariablesToHistogram', particleList='pi+:generic',
788 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
789 'extraInfo(SignalProbability)', 'isPrimarySignal',
790 'extraInfo(decayModeID)', 'extraInfo(postCut_rank)']),
791 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
792 ('extraInfo(decayModeID)', 'mcErrors'),
793 ('extraInfo(decayModeID)', 'mcParticleStatus'),
794 ('extraInfo(decayModeID)', 'extraInfo(postCut_rank)'),
795 ('isPrimarySignal', 'extraInfo(postCut_rank)'),
796 ('mcErrors', 'extraInfo(postCut_rank)'),
797 ('mcParticleStatus', 'extraInfo(postCut_rank)')]),
798 fileName='Monitor_PostReconstruction_AfterRanking.root', directory='pi+:generic',
799 ignoreCommandLineOverride=True)
800 path.add_module('VariablesToNtuple', fileName='Monitor_Final.root', treeName='pi+:generic variables',
801 ignoreCommandLineOverride=True,
802 variables=['extraInfo(SignalProbability)', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal',
803 'extraInfo(uniqueSignal)', 'extraInfo(decayModeID)'],
804 particleList='pi+:generic')
805
806 path.add_module('MVAExpert', identifier='UNITTEST_K+:generic ==> K+:FSP', extraInfoName='SignalProbability',
807 listNames=['K+:generic_0'])
808 path.add_module('TagUniqueSignal', particleList='K+:generic_0', target='isPrimarySignal',
809 extraInfoName='uniqueSignal')
810 path.add_module('VariablesToHistogram', particleList='K+:generic_0',
811 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
812 'extraInfo(SignalProbability)',
813 'isPrimarySignal', 'extraInfo(decayModeID)']),
814 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isPrimarySignal'),
815 ('extraInfo(SignalProbability)', 'mcErrors'),
816 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
817 ('extraInfo(decayModeID)', 'isPrimarySignal'),
818 ('extraInfo(decayModeID)', 'mcErrors'),
819 ('extraInfo(decayModeID)', 'extraInfo(uniqueSignal)'),
820 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
821 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='K+:generic ==> K+:FSP',
822 ignoreCommandLineOverride=True)
823 path.add_module('ParticleListManipulator', outputListName='K+:generic', inputListNames=['K+:generic_0'],
824 variable='particleSource', writeOut=True)
825 path.add_module('VariablesToHistogram', particleList='K+:generic',
826 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
827 'extraInfo(SignalProbability)',
828 'isPrimarySignal', 'extraInfo(decayModeID)']),
829 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
830 ('extraInfo(decayModeID)', 'mcErrors'),
831 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
832 fileName='Monitor_PostReconstruction_BeforePostCut.root', directory='K+:generic',
833 ignoreCommandLineOverride=True)
834 path.add_module('ParticleSelector', decayString='K+:generic', cut='0.01 < extraInfo(SignalProbability)')
835 path.add_module('VariablesToHistogram', particleList='K+:generic',
836 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
837 'extraInfo(SignalProbability)',
838 'isPrimarySignal', 'extraInfo(decayModeID)']),
839 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
840 ('extraInfo(decayModeID)', 'mcErrors'),
841 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
842 fileName='Monitor_PostReconstruction_BeforeRanking.root', directory='K+:generic',
843 ignoreCommandLineOverride=True)
844 path.add_module('BestCandidateSelection', particleList='K+:generic', variable='extraInfo(SignalProbability)',
845 selectLowest=False, numBest=10, outputVariable='postCut_rank')
846 path.add_module('VariablesToHistogram', particleList='K+:generic',
847 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
848 'extraInfo(SignalProbability)', 'isPrimarySignal',
849 'extraInfo(decayModeID)', 'extraInfo(postCut_rank)']),
850 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isPrimarySignal'),
851 ('extraInfo(decayModeID)', 'mcErrors'),
852 ('extraInfo(decayModeID)', 'mcParticleStatus'),
853 ('extraInfo(decayModeID)', 'extraInfo(postCut_rank)'),
854 ('isPrimarySignal', 'extraInfo(postCut_rank)'),
855 ('mcErrors', 'extraInfo(postCut_rank)'),
856 ('mcParticleStatus', 'extraInfo(postCut_rank)')]),
857 fileName='Monitor_PostReconstruction_AfterRanking.root', directory='K+:generic',
858 ignoreCommandLineOverride=True)
859 path.add_module('VariablesToNtuple', fileName='Monitor_Final.root', treeName='K+:generic variables',
860 ignoreCommandLineOverride=True,
861 variables=['extraInfo(SignalProbability)', 'mcErrors', 'mcParticleStatus', 'isPrimarySignal',
862 'extraInfo(uniqueSignal)', 'extraInfo(decayModeID)'],
863 particleList='K+:generic')
864
865 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> K-:generic pi+:generic',
866 extraInfoName='SignalProbability', listNames=['D0:generic_0'])
867 path.add_module('TagUniqueSignal', particleList='D0:generic_0', target='isSignal',
868 extraInfoName='uniqueSignal')
869 path.add_module('VariablesToHistogram', particleList='D0:generic_0',
870 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
871 'extraInfo(SignalProbability)',
872 'isSignal', 'extraInfo(decayModeID)']),
873 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isSignal'),
874 ('extraInfo(SignalProbability)', 'mcErrors'),
875 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
876 ('extraInfo(decayModeID)', 'isSignal'),
877 ('extraInfo(decayModeID)', 'mcErrors'),
878 ('extraInfo(decayModeID)', 'extraInfo(uniqueSignal)'),
879 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
880 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='D0:generic ==> K-:generic pi+:generic',
881 ignoreCommandLineOverride=True)
882
883 path.add_module('MVAExpert', identifier='UNITTEST_D0:generic ==> pi-:generic pi+:generic',
884 extraInfoName='SignalProbability', listNames=['D0:generic_1'])
885 path.add_module('TagUniqueSignal', particleList='D0:generic_1', target='isSignal',
886 extraInfoName='uniqueSignal')
887 path.add_module('VariablesToHistogram', particleList='D0:generic_1',
888 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
889 'extraInfo(SignalProbability)',
890 'isSignal', 'extraInfo(decayModeID)']),
891 variables_2d=fei.config.variables2binnings_2d([('extraInfo(SignalProbability)', 'isSignal'),
892 ('extraInfo(SignalProbability)', 'mcErrors'),
893 ('extraInfo(SignalProbability)', 'mcParticleStatus'),
894 ('extraInfo(decayModeID)', 'isSignal'),
895 ('extraInfo(decayModeID)', 'mcErrors'),
896 ('extraInfo(decayModeID)', 'extraInfo(uniqueSignal)'),
897 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
898 fileName='Monitor_PostReconstruction_AfterMVA.root', directory='D0:generic ==> pi-:generic pi+:generic',
899 ignoreCommandLineOverride=True)
900
901 path.add_module('ParticleListManipulator', outputListName='D0:generic',
902 inputListNames=['D0:generic_0', 'D0:generic_1'], variable='particleSource',
903 writeOut=True)
904 path.add_module('VariablesToHistogram', particleList='D0:generic',
905 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
906 'extraInfo(SignalProbability)',
907 'isSignal', 'extraInfo(decayModeID)']),
908 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isSignal'),
909 ('extraInfo(decayModeID)', 'mcErrors'),
910 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
911 fileName='Monitor_PostReconstruction_BeforePostCut.root', directory='D0:generic',
912 ignoreCommandLineOverride=True)
913 path.add_module('ParticleSelector', decayString='D0:generic', cut='0.001 < extraInfo(SignalProbability)')
914 path.add_module('VariablesToHistogram', particleList='D0:generic',
915 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
916 'extraInfo(SignalProbability)',
917 'isSignal', 'extraInfo(decayModeID)']),
918 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isSignal'),
919 ('extraInfo(decayModeID)', 'mcErrors'),
920 ('extraInfo(decayModeID)', 'mcParticleStatus')]),
921 fileName='Monitor_PostReconstruction_BeforeRanking.root', directory='D0:generic',
922 ignoreCommandLineOverride=True)
923 path.add_module('BestCandidateSelection', particleList='D0:generic', variable='extraInfo(SignalProbability)',
924 selectLowest=False, numBest=10, outputVariable='postCut_rank')
925 path.add_module('VariablesToHistogram', particleList='D0:generic',
926 variables=fei.config.variables2binnings(['mcErrors', 'mcParticleStatus', 'extraInfo(uniqueSignal)',
927 'extraInfo(SignalProbability)', 'isSignal',
928 'extraInfo(decayModeID)', 'extraInfo(postCut_rank)']),
929 variables_2d=fei.config.variables2binnings_2d([('extraInfo(decayModeID)', 'isSignal'),
930 ('extraInfo(decayModeID)', 'mcErrors'),
931 ('extraInfo(decayModeID)', 'mcParticleStatus'),
932 ('extraInfo(decayModeID)', 'extraInfo(postCut_rank)'),
933 ('isSignal', 'extraInfo(postCut_rank)'),
934 ('mcErrors', 'extraInfo(postCut_rank)'),
935 ('mcParticleStatus', 'extraInfo(postCut_rank)')]),
936 fileName='Monitor_PostReconstruction_AfterRanking.root', directory='D0:generic',
937 ignoreCommandLineOverride=True)
938 path.add_module('VariablesToNtuple', fileName='Monitor_Final.root', treeName='D0:generic variables',
939 ignoreCommandLineOverride=True,
940 variables=['extraInfo(SignalProbability)', 'mcErrors', 'mcParticleStatus', 'isSignal',
941 'extraInfo(uniqueSignal)', 'extraInfo(decayModeID)'],
942 particleList='D0:generic')
943
944 print_path(path, x.reconstruct())
945 self.assertEqual(x.reconstruct(), path)
946
947
948class TestTeacher(unittest.TestCase):
949 def setUp(self):
950 fei.core.Teacher.MaximumNumberOfMVASamples = int(1e7)
951 fei.core.Teacher.MinimumNumberOfMVASamples = int(10)
952
953 f = ROOT.TFile('training_input.root', 'RECREATE')
954 f.cd()
955 tree = ROOT.TTree('pi+:generic ==> pi+:FSP variables', 'pi+:generic ==> pi+:FSP variables')
956 isSignal = np.zeros(1, dtype=float)
957 p = np.zeros(1, dtype=float)
958 pt = np.zeros(1, dtype=float)
959 tree.Branch('isPrimarySignal', isSignal, 'isPrimarySignal/D')
960 tree.Branch('p', p, 'p/D')
961 tree.Branch('dr', pt, 'dr/D')
962 for i in range(1000):
963 isSignal[0] = i % 2
964 p[0] = i
965 pt[0] = i * 2
966 tree.Fill()
967 tree.Write("", ROOT.TObject.kOverwrite)
968
969 # Test case, where tree not existent for K+:generic ==> K+:FSP, so skip creating it
970
971 tree = ROOT.TTree('D0:generic ==> K-:generic pi+:generic variables', 'D0:generic ==> K-:generic pi+:generic variables')
972 isSignal = np.zeros(1, dtype=float)
973 p = np.zeros(1, dtype=float)
974 pt = np.zeros(1, dtype=float)
975 tree.Branch('isSignal', isSignal, 'isSignal/D')
976 tree.Branch('M', p, 'M/D')
977 tree.Branch('p', pt, 'p/D')
978 # Too few signal events here!
979 for i in range(10):
980 isSignal[0] = i % 2
981 p[0] = i
982 pt[0] = i * 2
983 tree.Fill()
984 tree.Write("", ROOT.TObject.kOverwrite)
985
986 tree = ROOT.TTree('D0:generic ==> pi-:generic pi+:generic variables', 'D0:generic ==> pi-:generic pi+:generic variables')
987 isSignal = np.zeros(1, dtype=float)
988 p = np.zeros(1, dtype=float)
989 pt = np.zeros(1, dtype=float)
990 tree.Branch('isSignal', isSignal, 'isSignal/D')
991 tree.Branch('M', p, 'M/D')
992 tree.Branch('p', pt, 'p/D')
993 for i in range(1000):
994 isSignal[0] = i % 2
995 p[0] = i
996 pt[0] = i * 2
997 tree.Fill()
998 tree.Write("", ROOT.TObject.kOverwrite)
999
1000 def tearDown(self):
1001 if os.path.isfile('UNITTEST_TEACHER.xml'):
1002 os.remove('UNITTEST_TEACHER.xml')
1003 if os.path.isfile('UNITTEST_pi+:generic ==> pi+:FSP.xml'):
1004 os.remove('UNITTEST_pi+:generic ==> pi+:FSP.xml')
1005 if os.path.isfile('UNITTEST_K+:generic ==> K+:FSP.xml'):
1006 os.remove('UNITTEST_K+:generic ==> K+:FSP.xml')
1007 if os.path.isfile('UNITTEST_D0:generic ==> K-:generic pi+:generic.xml'):
1008 os.remove('UNITTEST_D0:generic ==> K-:generic pi+:generic.xml')
1009 if os.path.isfile('UNITTEST_D0:generic ==> pi-:generic pi+:generic.xml'):
1010 os.remove('UNITTEST_D0:generic ==> pi-:generic pi+:generic.xml')
1011 if os.path.isfile('training_input.root'):
1012 os.remove('training_input.root')
1013
1014 def test_create_fake_weightfile(self):
1015 self.assertEqual(os.path.isfile('UNITTEST_pi+:generic ==> pi+:FSP.xml'), False)
1016 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP.xml'), False)
1017 fei.core.Teacher.create_fake_weightfile('UNITTEST_pi+:generic ==> pi+:FSP')
1018 self.assertEqual(os.path.isfile('UNITTEST_pi+:generic ==> pi+:FSP.xml'), True)
1019 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP.xml'), True)
1020
1021 def test_upload(self):
1022 particles = get_small_unittest_channels()
1023 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST', externTeacher='basf2_mva_teacher')
1024 x = fei.core.Teacher(particles, config)
1025 fei.core.Teacher.create_fake_weightfile('TEACHER')
1026 self.assertEqual(basf2_mva.available('UNITTEST_TEACHER'), False)
1027 r = x.upload('TEACHER')
1028 self.assertEqual(basf2_mva.available('UNITTEST_TEACHER'), True)
1029 self.assertEqual(r, ('TEACHER.xml', 'UNITTEST_TEACHER'))
1030
1031 def test_without_monitoring(self):
1032 particles = get_small_unittest_channels()
1033 config = fei.config.FeiConfiguration(monitor=False, prefix='UNITTEST', externTeacher='basf2_mva_teacher')
1034 x = fei.core.Teacher(particles, config)
1035
1036 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP'), False)
1037 self.assertEqual(basf2_mva.available('UNITTEST_K+:generic ==> K+:FSP'), False)
1038 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> K-:generic pi+:generic'), False)
1039 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> pi-:generic pi+:generic'), False)
1040
1041 x.do_all_trainings()
1042
1043 self.assertEqual(basf2_mva.available('UNITTEST_pi+:generic ==> pi+:FSP'), True)
1044 self.assertEqual(basf2_mva.available('UNITTEST_K+:generic ==> K+:FSP'), False)
1045 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> K-:generic pi+:generic'), True)
1046 self.assertEqual(basf2_mva.available('UNITTEST_D0:generic ==> pi-:generic pi+:generic'), True)
1047
1048
1049"""
1050class TestConvertLegacyTraining(unittest.TestCase):
1051 pass
1052"""
1053
1054
1055class TestGetPath(unittest.TestCase):
1056
1057 def setUp(self):
1058 particles = fei.get_unittest_channels()
1059
1060 f = ROOT.TFile('mcParticlesCount.root', 'RECREATE')
1061 f.cd()
1062
1063 for pdgnumber in {abs(pdg.from_name(particle.name)) for particle in particles}:
1064 hist = ROOT.TH1F(f"NumberOfMCParticlesInEvent__bo{pdgnumber}__bc",
1065 f"NumberOfMCParticlesInEvent__bo{pdgnumber}__bc", 11, -0.5, 10.5)
1066 for i in range(10):
1067 hist.Fill(5)
1068 f.Write(f"NumberOfMCParticlesInEvent__bo{pdgnumber}__bc")
1069
1070 def tearDown(self):
1071 if os.path.isfile('mcParticlesCount.root'):
1072 os.remove('mcParticlesCount.root')
1073 if os.path.isfile('Summary.pickle'):
1074 os.remove('Summary.pickle')
1075
1076 def test_get_path_default_cache(self):
1077 particles = fei.get_unittest_channels()
1078 config = fei.config.FeiConfiguration(training=True)
1079 x = fei.core.Teacher(particles, config) # noqa
1080
1081 # Should try to create mcParticlesCount
1082 # -> Returns at stage 0
1083 feistate = fei.core.get_path(particles, config)
1084 self.assertEqual(feistate.stage, 0)
1085
1086 # Should try to create training data for FSPs
1087 # -> Returns at stage 1
1088 config = fei.config.FeiConfiguration(training=True)
1089 feistate = fei.core.get_path(particles, config)
1090 self.assertEqual(feistate.stage, 1)
1091
1092 # No weightfiles were created, hence the algorithm
1093 # cannot go forward and tries to create the training data again
1094 config = fei.config.FeiConfiguration(training=True)
1095 feistate = fei.core.get_path(particles, config)
1096 self.assertEqual(feistate.stage, 1)
1097
1098 # We create the FSP weightfiles by hand
1099 fei.core.Teacher.create_fake_weightfile('pi+:generic ==> pi+:FSP')
1100 fei.core.Teacher.create_fake_weightfile('K+:generic ==> K+:FSP')
1101 fei.core.Teacher.create_fake_weightfile('mu+:generic ==> mu+:FSP')
1102 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:FSP')
1103 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:V0')
1104
1105 # Should try to create training data for pi0
1106 # -> Returns at stage 2
1107 config = fei.config.FeiConfiguration(training=True)
1108 feistate = fei.core.get_path(particles, config)
1109 self.assertEqual(feistate.stage, 2)
1110
1111 # No weightfiles were created, hence the algorithm
1112 # cannot go forward and tries to create the training data again
1113 config = fei.config.FeiConfiguration(training=True)
1114 feistate = fei.core.get_path(particles, config)
1115 self.assertEqual(feistate.stage, 2)
1116
1117 # We create the pi0 weightfile by hand
1118 fei.core.Teacher.create_fake_weightfile('pi0:generic ==> gamma:generic gamma:generic')
1119
1120 # Should try to create training data for D0
1121 # -> Returns stage 4 (stage 3 is skipped because it only contains K_S0)
1122 config = fei.config.FeiConfiguration(training=True)
1123 feistate = fei.core.get_path(particles, config)
1124 self.assertEqual(feistate.stage, 4)
1125
1126 # We create the D0 weightfiles by hand
1127 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic')
1128 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic pi0:generic')
1129 fei.core.Teacher.create_fake_weightfile('D0:generic ==> pi-:generic pi+:generic')
1130 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic mu+:generic')
1131 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic pi0:generic mu+:generic')
1132
1133 # Unittest channels do not contain anymore stages,
1134 # -> Returns last stage + 1
1135 config = fei.config.FeiConfiguration(training=True)
1136 feistate = fei.core.get_path(particles, config)
1137 self.assertEqual(feistate.stage, 7)
1138
1139 # If we start at 0, we should be get the whole path
1140 # -> Returns last stage + 1
1141 config = fei.config.FeiConfiguration(cache=0, training=True)
1142 feistate = fei.core.get_path(particles, config)
1143 self.assertEqual(feistate.stage, 7)
1144
1145 # If training is False we should get the whole path without
1146 # checking if weightfiles exist.
1147 os.remove('pi+:generic ==> pi+:FSP.xml')
1148 os.remove('K+:generic ==> K+:FSP.xml')
1149 os.remove('mu+:generic ==> mu+:FSP.xml')
1150 os.remove('gamma:generic ==> gamma:FSP.xml')
1151 os.remove('gamma:generic ==> gamma:V0.xml')
1152 os.remove('pi0:generic ==> gamma:generic gamma:generic.xml')
1153 os.remove('D0:generic ==> K-:generic pi+:generic.xml')
1154 os.remove('D0:generic ==> K-:generic pi+:generic pi0:generic.xml')
1155 os.remove('D0:generic ==> pi-:generic pi+:generic.xml')
1156 os.remove('D0:semileptonic ==> K-:generic mu+:generic.xml')
1157 os.remove('D0:semileptonic ==> K-:generic pi0:generic mu+:generic.xml')
1158
1159 config = fei.config.FeiConfiguration(cache=0, training=False)
1160 feistate = fei.core.get_path(particles, config)
1161 self.assertEqual(feistate.stage, 7)
1162
1163 def test_get_path(self):
1164 particles = fei.get_unittest_channels()
1165 config = fei.config.FeiConfiguration(cache=-1, training=True)
1166 x = fei.core.Teacher(particles, config) # noqa
1167
1168 # Should try to create mcParticlesCount
1169 # -> Returns at stage 0
1170 feistate = fei.core.get_path(particles, config)
1171 self.assertEqual(feistate.stage, 0)
1172
1173 # Should try to create training data for FSPs
1174 # -> Returns at stage 1
1175 config = fei.config.FeiConfiguration(cache=0, training=True)
1176 feistate = fei.core.get_path(particles, config)
1177 self.assertEqual(feistate.stage, 1)
1178
1179 # No weightfiles were created, hence the algorithm
1180 # cannot go forward and tries to create the training data again
1181 config = fei.config.FeiConfiguration(cache=1, training=True)
1182 feistate = fei.core.get_path(particles, config)
1183 self.assertEqual(feistate.stage, 1)
1184
1185 # We create the FSP weightfiles by hand
1186 fei.core.Teacher.create_fake_weightfile('pi+:generic ==> pi+:FSP')
1187 fei.core.Teacher.create_fake_weightfile('K+:generic ==> K+:FSP')
1188 fei.core.Teacher.create_fake_weightfile('mu+:generic ==> mu+:FSP')
1189 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:FSP')
1190 fei.core.Teacher.create_fake_weightfile('gamma:generic ==> gamma:V0')
1191
1192 # Should try to create training data for pi0
1193 # -> Returns at stage 2
1194 config = fei.config.FeiConfiguration(cache=1, training=True)
1195 feistate = fei.core.get_path(particles, config)
1196 self.assertEqual(feistate.stage, 2)
1197
1198 # No weightfiles were created, hence the algorithm
1199 # cannot go forward and tries to create the training data again
1200 config = fei.config.FeiConfiguration(cache=2, training=True)
1201 feistate = fei.core.get_path(particles, config)
1202 self.assertEqual(feistate.stage, 2)
1203
1204 # No weightfiles were created, hence the algorithm
1205 # cannot go forward and tries to create the training data again
1206 # This applies as well if the stage is even bigger
1207 config = fei.config.FeiConfiguration(cache=4, training=True)
1208 feistate = fei.core.get_path(particles, config)
1209 self.assertEqual(feistate.stage, 2)
1210
1211 # We create the pi0 weightfile by hand
1212 fei.core.Teacher.create_fake_weightfile('pi0:generic ==> gamma:generic gamma:generic')
1213
1214 # Should try to create training data for D0
1215 # -> Returns stage 4 (stage 3 is skipped because it only contains K_S0)
1216 config = fei.config.FeiConfiguration(cache=2, training=True)
1217 feistate = fei.core.get_path(particles, config)
1218 self.assertEqual(feistate.stage, 4)
1219
1220 # We create the D0 weightfiles by hand
1221 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic')
1222 fei.core.Teacher.create_fake_weightfile('D0:generic ==> K-:generic pi+:generic pi0:generic')
1223 fei.core.Teacher.create_fake_weightfile('D0:generic ==> pi-:generic pi+:generic')
1224 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic mu+:generic')
1225 fei.core.Teacher.create_fake_weightfile('D0:semileptonic ==> K-:generic pi0:generic mu+:generic')
1226
1227 # Unittest channels do not contain anymore stages,
1228 # -> Returns last stage + 1
1229 config = fei.config.FeiConfiguration(cache=4, training=True)
1230 feistate = fei.core.get_path(particles, config)
1231 self.assertEqual(feistate.stage, 7)
1232
1233 # If we start at 0, we should be get the whole path
1234 # -> Returns last stage + 1
1235 config = fei.config.FeiConfiguration(cache=0, training=True)
1236 feistate = fei.core.get_path(particles, config)
1237 self.assertEqual(feistate.stage, 7)
1238
1239 # If training is False we should get the whole path without
1240 # checking if weightfiles exist.
1241 os.remove('pi+:generic ==> pi+:FSP.xml')
1242 os.remove('K+:generic ==> K+:FSP.xml')
1243 os.remove('mu+:generic ==> mu+:FSP.xml')
1244 os.remove('gamma:generic ==> gamma:FSP.xml')
1245 os.remove('gamma:generic ==> gamma:V0.xml')
1246 os.remove('pi0:generic ==> gamma:generic gamma:generic.xml')
1247 os.remove('D0:generic ==> K-:generic pi+:generic.xml')
1248 os.remove('D0:generic ==> K-:generic pi+:generic pi0:generic.xml')
1249 os.remove('D0:generic ==> pi-:generic pi+:generic.xml')
1250 os.remove('D0:semileptonic ==> K-:generic mu+:generic.xml')
1251 os.remove('D0:semileptonic ==> K-:generic pi0:generic mu+:generic.xml')
1252
1253 config = fei.config.FeiConfiguration(cache=0, training=False)
1254 feistate = fei.core.get_path(particles, config)
1255 self.assertEqual(feistate.stage, 7)
1256
1257
1258if __name__ == '__main__':
1259 with b2test_utils.clean_working_directory():
1260 basf2.conditions.testing_payloads = ['localdb/database.txt']
1261 unittest.main()
1262
1263# @endcond
b2bii.setB2BII
def setB2BII()
Definition: b2bii.py:21
b2bii.unsetB2BII
def unsetB2BII()
Definition: b2bii.py:25
b2test_utils.clean_working_directory
def clean_working_directory()
Definition: __init__.py:189
pdg.from_name
def from_name(name)
Definition: pdg.py:63