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