Belle II Software  release-08-01-10
test_fei_generic.py
1 #!/usr/bin/env python3
2 
3 
10 
11 """
12 <header>
13  <contact>wsut@uni-bonn.de</contact>
14 </header>
15 """
16 
17 import os
18 import subprocess
19 
20 import tempfile
21 import shutil
22 import glob
23 
24 import fei
25 import basf2 as b2
26 import modularAnalysis as ma
27 import ROOT
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 # Construct path for production of mcParticlesCount.root at stage -1
40 path = b2.create_path()
41 
42 ma.inputMdst(filename=b2.find_file('mdst14.root', 'validation', False),
43  path=path)
44 
45 maxTracks = 12
46 empty_path = b2.create_path()
47 skimfilter = b2.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, monitor=False)
53 feistate = fei.get_path(particles, configuration)
54 path.add_path(feistate.path)
55 path.add_module('RootOutput')
56 
57 assert feistate.stage == 0 # corresponds to stage -1, since increased by 1 after creating path
58 path.add_module('Progress')
59 print(path)
60 b2.process(path, max_event=10000)
61 assert len(glob.glob('RootOutput.root')) == 1
62 assert len(glob.glob('mcParticlesCount.root')) == 1
63 
64 # Construct path for production of stage 0 training data
65 path = b2.create_path()
66 ma.inputMdstList(['./RootOutput.root'], path)
67 
68 maxTracks = 12
69 empty_path = b2.create_path()
70 skimfilter = b2.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, monitor=False)
76 feistate = fei.get_path(particles, configuration)
77 path.add_path(feistate.path)
78 path.add_module('RootOutput')
79 
80 assert feistate.stage == 1 # corresponds to stage 0, since increased by 1 after creating path
81 path.add_module('Progress')
82 print(path)
83 b2.process(path, max_event=10000)
84 
85 # Check availability of training input and cache file for stage 0
86 assert len(glob.glob('RootOutput.root')) == 1
87 assert len(glob.glob('training_input.root')) == 1
88 f = ROOT.TFile.Open("training_input.root", "read")
89 assert sum(['gamma' in key.GetName()[:5] for key in f.GetListOfKeys()]) == 2
90 assert sum(['mu+' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
91 assert sum(['pi+' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
92 assert sum(['K+' in key.GetName()[:3] for key in f.GetListOfKeys()]) == 1
93 f.Close()
94 
95 # Perform stage 0 training
96 fei.do_trainings(particles, configuration)
97 
98 # Moving training_input.root to training_input_stage0.root for later usage
99 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
100 
101 # Check availability of stage 0 *.xml training files
102 assert len(glob.glob('gamma*.xml')) == 2
103 assert len(glob.glob('mu+*.xml')) == 1
104 assert len(glob.glob('pi+*.xml')) == 1
105 assert len(glob.glob('K+*.xml')) == 1
106 
107 # Construct path for production of stage 1 training data
108 path = b2.create_path()
109 ma.inputMdstList(['./RootOutput.root'], path)
110 
111 maxTracks = 12
112 empty_path = b2.create_path()
113 skimfilter = b2.register_module('VariableToReturnValue')
114 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
115 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
116 path.add_module(skimfilter)
117 
118 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True, monitor=False)
119 feistate = fei.get_path(particles, configuration)
120 path.add_path(feistate.path)
121 path.add_module('RootOutput')
122 
123 assert feistate.stage == 2 # corresponds to stage 1, since increased by 1 after creating path
124 path.add_module('Progress')
125 print(path)
126 b2.process(path, max_event=10000)
127 
128 # Check availability of training input and cache file for stage 1
129 assert len(glob.glob('RootOutput.root')) == 1
130 assert len(glob.glob('training_input.root')) == 1
131 f = ROOT.TFile.Open("training_input.root", "read")
132 assert sum(['pi0' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
133 f.Close()
134 
135 # Perform stage 1 training
136 fei.do_trainings(particles, configuration)
137 
138 # Moving training_input.root to training_input_stage1.root for later usage
139 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
140 
141 # Check availability of stage 1 *.xml training files
142 assert len(glob.glob('pi0*.xml')) == 1
143 
144 # Construct path for production of stage 3 training data (stage 2 is skipped)
145 path = b2.create_path()
146 ma.inputMdstList(['./RootOutput.root'], path)
147 
148 maxTracks = 12
149 empty_path = b2.create_path()
150 skimfilter = b2.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, monitor=False)
156 feistate = fei.get_path(particles, configuration)
157 path.add_path(feistate.path)
158 path.add_module('RootOutput')
159 
160 assert feistate.stage == 4 # corresponds to stage 3, since increased by 1 after creating path
161 path.add_module('Progress')
162 print(path)
163 b2.process(path, max_event=10000)
164 
165 # Check availability of training input and cache file for stage 3
166 assert len(glob.glob('RootOutput.root')) == 1
167 assert len(glob.glob('training_input.root')) == 1
168 f = ROOT.TFile.Open("training_input.root", "read")
169 assert sum(['D' in key.GetName()[:2] for key in f.GetListOfKeys()]) == 5
170 f.Close()
171 
172 # Perform stage 3 training
173 fei.do_trainings(particles, configuration)
174 
175 # Moving training_input.root to training_input_stage3.root for later usage
176 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
177 
178 # Check availability of stage 3 *.xml training files
179 assert len(glob.glob('D*.xml')) == 5
180 
181 # Merge training input files for validation
182 subprocess.call(["analysis-fei-mergefiles", "training_input.root"] + glob.glob("training_input_stage*.root"))
183 
184 # Expect 4 different training_input*.root files now
185 assert len(glob.glob('training_input*.root')) == 4
186 
187 # Construct path for stage 6 preparing evaluation (stages 4 and 5 skipped, input evaluated from stage 0 on)
188 path = b2.create_path()
189 ma.inputMdst(filename=b2.find_file('mdst14.root', 'validation', False),
190  path=path)
191 
192 maxTracks = 12
193 empty_path = b2.create_path()
194 skimfilter = b2.register_module('VariableToReturnValue')
195 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
196 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
197 path.add_module(skimfilter)
198 
199 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, monitor=True)
200 feistate = fei.get_path(particles, configuration)
201 path.add_path(feistate.path)
202 
203 assert feistate.stage == 7 # corresponds to stage 6, since increased by 1 after creating path
204 path.add_module('Progress')
205 print(path)
206 b2.process(path, max_event=10000)
207 
208 # Check availability of monitoring files constructed at stage 6 for validation
209 assert len(glob.glob('Monitor_FSPLoader.root')) == 1
210 
211 assert len(glob.glob('Monitor_ModuleStatistics.root')) == 1
212 
213 assert len(glob.glob('Monitor_Final.root')) == 1
214 f = ROOT.TFile.Open("Monitor_Final.root", "read")
215 assert f.GetListOfKeys().GetEntries() == 7
216 f.Close()
217 
218 prereconstruction_files = [
219  'Monitor_PreReconstruction_BeforeRanking.root',
220  'Monitor_PreReconstruction_AfterRanking.root',
221  'Monitor_PreReconstruction_AfterVertex.root',
222  'Monitor_PostReconstruction_AfterMVA.root'
223 ]
224 
225 for fname in prereconstruction_files:
226  assert len(glob.glob(fname)) == 1
227  f = ROOT.TFile.Open(fname, "read")
228  assert f.GetListOfKeys().GetEntries() == 11
229  f.Close()
230 
231 postreconstruction_files = [
232  'Monitor_PostReconstruction_BeforePostCut.root',
233  'Monitor_PostReconstruction_BeforeRanking.root',
234  'Monitor_PostReconstruction_AfterRanking.root'
235 ]
236 
237 for fname in postreconstruction_files:
238  assert len(glob.glob(fname)) == 1
239  f = ROOT.TFile.Open(fname, "read")
240  assert f.GetListOfKeys().GetEntries() == 7
241  f.Close()
242 
243 shutil.rmtree(tempdir)