Belle II Software  release-05-01-25
test_fei_generic.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 """
5 <header>
6  <contact>wsut@uni-bonn.de</contact>
7 </header>
8 """
9 
10 import fei
11 from fei.default_channels import get_unittest_channels
12 import os
13 import sys
14 
15 import tempfile
16 import shutil
17 import glob
18 import sys
19 
20 import basf2 as b2
21 import modularAnalysis as ma
22 from ROOT import Belle2
23 import basf2_mva
24 import pdg
25 from basf2 import conditions
26 
27 basf2_mva.loadRootDictionary()
28 
29 tempdir = tempfile.mkdtemp()
30 os.chdir(tempdir)
31 
32 b2.conditions.append_testing_payloads('localdb/database.txt')
33 
34 fei.core.Teacher.MaximumNumberOfMVASamples = int(1e7)
35 fei.core.Teacher.MinimumNumberOfMVASamples = int(10)
36 
37 particles = fei.get_unittest_channels()
38 
39 path = b2.create_path()
40 
41 ma.inputMdst(environmentType='default',
42  filename=b2.find_file('mdst12.root', 'validation', False),
43  path=path)
44 
45 maxTracks = 12
46 empty_path = b2.create_path()
47 skimfilter = ma.register_module('VariableToReturnValue')
48 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
49 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
50 path.add_module(skimfilter)
51 
52 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=-1, training=True)
53 feistate = fei.get_path(particles, configuration)
54 path.add_path(feistate.path)
55 path.add_module('RootOutput')
56 
57 assert feistate.stage == 0
58 print(path)
59 b2.process(path, max_event=10000)
60 assert len(glob.glob('RootOutput.root')) == 1
61 assert len(glob.glob('mcParticlesCount.root')) == 1
62 
63 fei.do_trainings(particles, configuration)
64 
65 path = b2.create_path()
66 ma.inputMdstList('default', ['./RootOutput.root'], path)
67 
68 maxTracks = 12
69 empty_path = b2.create_path()
70 skimfilter = ma.register_module('VariableToReturnValue')
71 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
72 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
73 path.add_module(skimfilter)
74 
75 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True)
76 feistate = fei.get_path(particles, configuration)
77 path.add_path(feistate.path)
78 path.add_module('RootOutput')
79 
80 assert feistate.stage == 1
81 print(path)
82 b2.process(path, max_event=10000)
83 assert len(glob.glob('RootOutput.root')) == 1
84 assert len(glob.glob('gamma*')) == 2
85 assert len(glob.glob('mu+*')) == 1
86 assert len(glob.glob('pi+*')) == 1
87 assert len(glob.glob('K+*')) == 1
88 
89 fei.do_trainings(particles, configuration)
90 
91 assert len(glob.glob('gamma*')) == 6
92 assert len(glob.glob('mu+*')) == 3
93 assert len(glob.glob('pi+*')) == 3
94 assert len(glob.glob('K+*')) == 3
95 
96 path = b2.create_path()
97 ma.inputMdstList('default', ['./RootOutput.root'], path)
98 
99 maxTracks = 12
100 empty_path = b2.create_path()
101 skimfilter = ma.register_module('VariableToReturnValue')
102 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
103 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
104 path.add_module(skimfilter)
105 
106 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True)
107 feistate = fei.get_path(particles, configuration)
108 path.add_path(feistate.path)
109 path.add_module('RootOutput')
110 
111 assert feistate.stage == 2
112 print(path)
113 b2.process(path, max_event=10000)
114 assert len(glob.glob('pi0*')) == 1
115 
116 fei.do_trainings(particles, configuration)
117 assert len(glob.glob('pi0*')) == 3
118 
119 path = b2.create_path()
120 ma.inputMdstList('default', ['./RootOutput.root'], path)
121 
122 maxTracks = 12
123 empty_path = b2.create_path()
124 skimfilter = ma.register_module('VariableToReturnValue')
125 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
126 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
127 path.add_module(skimfilter)
128 
129 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True)
130 feistate = fei.get_path(particles, configuration)
131 path.add_path(feistate.path)
132 path.add_module('RootOutput')
133 
134 assert feistate.stage == 4
135 print(path)
136 b2.process(path, max_event=10000)
137 assert len(glob.glob('RootOutput.root')) == 1
138 assert len(glob.glob('D*')) == 5
139 
140 # One training will fail D -> pi pi due to low statistic
141 fei.do_trainings(particles, configuration)
142 print(len(glob.glob('D*')))
143 assert len(glob.glob('D*')) == 15
144 
145 path = b2.create_path()
146 ma.inputMdstList('default', ['./RootOutput.root'], path)
147 
148 maxTracks = 12
149 empty_path = b2.create_path()
150 skimfilter = ma.register_module('VariableToReturnValue')
151 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
152 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
153 path.add_module(skimfilter)
154 
155 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True)
156 feistate = fei.get_path(particles, configuration)
157 path.add_path(feistate.path)
158 path.add_module('RootOutput')
159 
160 assert feistate.stage == 7
161 
162 path = b2.create_path()
163 ma.inputMdst(environmentType='default',
164  filename=b2.find_file('mdst12.root', 'validation', False),
165  path=path)
166 
167 maxTracks = 12
168 empty_path = b2.create_path()
169 skimfilter = ma.register_module('VariableToReturnValue')
170 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
171 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
172 path.add_module(skimfilter)
173 
174 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, monitor=True)
175 feistate = fei.get_path(particles, configuration)
176 path.add_path(feistate.path)
177 path.add_module('RootOutput')
178 
179 assert feistate.stage == 7
180 print(path)
181 b2.process(path, max_event=10000)
182 assert len(glob.glob('Monitor_FSPLoader.root')) == 1
183 assert len(glob.glob('Monitor_TrainingData_*')) == 11
184 assert len(glob.glob('Monitor_PreReconstruction_BeforeRanking_*')) == 11
185 assert len(glob.glob('Monitor_PreReconstruction_AfterRanking_*')) == 11
186 assert len(glob.glob('Monitor_PreReconstruction_AfterVertex_*')) == 11
187 assert len(glob.glob('Monitor_PostReconstruction_AfterMVA_*')) == 11
188 assert len(glob.glob('Monitor_PostReconstruction_BeforePostCut_*')) == 7
189 assert len(glob.glob('Monitor_PostReconstruction_BeforeRanking_*')) == 7
190 assert len(glob.glob('Monitor_PostReconstruction_AfterRanking_*')) == 7
191 assert len(glob.glob('Monitor_Final_*')) == 7
192 assert len(glob.glob('Monitor_ModuleStatistics.root')) == 1
193 
194 shutil.rmtree(tempdir)
fei.default_channels
Definition: default_channels.py:1