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