Belle II Software  light-2212-foldex
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 ROOT
29 
30 tempdir = tempfile.mkdtemp()
31 os.chdir(tempdir)
32 
33 b2.conditions.append_testing_payloads('localdb/database.txt')
34 
35 fei.core.Teacher.MaximumNumberOfMVASamples = int(1e7)
36 fei.core.Teacher.MinimumNumberOfMVASamples = int(10)
37 
38 particles = fei.get_unittest_channels()
39 
40 # Construct path for production of mcParticlesCount.root at stage -1
41 path = b2.create_path()
42 
43 ma.inputMdst(filename=b2.find_file('mdst14.root', 'validation', False),
44  path=path)
45 
46 maxTracks = 12
47 empty_path = b2.create_path()
48 skimfilter = b2.register_module('VariableToReturnValue')
49 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
50 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
51 path.add_module(skimfilter)
52 
53 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=-1, training=True, monitor=False)
54 feistate = fei.get_path(particles, configuration)
55 path.add_path(feistate.path)
56 path.add_module('RootOutput')
57 
58 assert feistate.stage == 0 # corresponds to stage -1, since increased by 1 after creating path
59 path.add_module('Progress')
60 print(path)
61 b2.process(path, max_event=10000)
62 assert len(glob.glob('RootOutput.root')) == 1
63 assert len(glob.glob('mcParticlesCount.root')) == 1
64 
65 # Construct path for production of stage 0 training data
66 path = b2.create_path()
67 ma.inputMdstList(['./RootOutput.root'], path)
68 
69 maxTracks = 12
70 empty_path = b2.create_path()
71 skimfilter = b2.register_module('VariableToReturnValue')
72 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
73 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
74 path.add_module(skimfilter)
75 
76 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True, monitor=False)
77 feistate = fei.get_path(particles, configuration)
78 path.add_path(feistate.path)
79 path.add_module('RootOutput')
80 
81 assert feistate.stage == 1 # corresponds to stage 0, since increased by 1 after creating path
82 path.add_module('Progress')
83 print(path)
84 b2.process(path, max_event=10000)
85 
86 # Check availability of training input and cache file for stage 0
87 assert len(glob.glob('RootOutput.root')) == 1
88 assert len(glob.glob('training_input.root')) == 1
89 f = ROOT.TFile.Open("training_input.root", "read")
90 assert sum(['gamma' in key.GetName()[:5] for key in f.GetListOfKeys()]) == 2
91 assert sum(['mu+' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
92 assert sum(['pi+' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
93 assert sum(['K+' in key.GetName()[:3] for key in f.GetListOfKeys()]) == 1
94 f.Close()
95 
96 # Perform stage 0 training
97 fei.do_trainings(particles, configuration)
98 
99 # Moving training_input.root to training_input_stage0.root for later usage
100 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
101 
102 # Check availability of stage 0 *.xml training files
103 assert len(glob.glob('gamma*.xml')) == 2
104 assert len(glob.glob('mu+*.xml')) == 1
105 assert len(glob.glob('pi+*.xml')) == 1
106 assert len(glob.glob('K+*.xml')) == 1
107 
108 # Construct path for production of stage 1 training data
109 path = b2.create_path()
110 ma.inputMdstList(['./RootOutput.root'], path)
111 
112 maxTracks = 12
113 empty_path = b2.create_path()
114 skimfilter = b2.register_module('VariableToReturnValue')
115 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
116 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
117 path.add_module(skimfilter)
118 
119 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True, monitor=False)
120 feistate = fei.get_path(particles, configuration)
121 path.add_path(feistate.path)
122 path.add_module('RootOutput')
123 
124 assert feistate.stage == 2 # corresponds to stage 1, since increased by 1 after creating path
125 path.add_module('Progress')
126 print(path)
127 b2.process(path, max_event=10000)
128 
129 # Check availability of training input and cache file for stage 1
130 assert len(glob.glob('RootOutput.root')) == 1
131 assert len(glob.glob('training_input.root')) == 1
132 f = ROOT.TFile.Open("training_input.root", "read")
133 assert sum(['pi0' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
134 f.Close()
135 
136 # Perform stage 1 training
137 fei.do_trainings(particles, configuration)
138 
139 # Moving training_input.root to training_input_stage1.root for later usage
140 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
141 
142 # Check availability of stage 1 *.xml training files
143 assert len(glob.glob('pi0*.xml')) == 1
144 
145 # Construct path for production of stage 3 training data (stage 2 is skipped)
146 path = b2.create_path()
147 ma.inputMdstList(['./RootOutput.root'], path)
148 
149 maxTracks = 12
150 empty_path = b2.create_path()
151 skimfilter = b2.register_module('VariableToReturnValue')
152 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
153 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
154 path.add_module(skimfilter)
155 
156 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True, monitor=False)
157 feistate = fei.get_path(particles, configuration)
158 path.add_path(feistate.path)
159 path.add_module('RootOutput')
160 
161 assert feistate.stage == 4 # corresponds to stage 3, since increased by 1 after creating path
162 path.add_module('Progress')
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(filename=b2.find_file('mdst14.root', 'validation', False),
191  path=path)
192 
193 maxTracks = 12
194 empty_path = b2.create_path()
195 skimfilter = b2.register_module('VariableToReturnValue')
196 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
197 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
198 path.add_module(skimfilter)
199 
200 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, monitor=True)
201 feistate = fei.get_path(particles, configuration)
202 path.add_path(feistate.path)
203 
204 assert feistate.stage == 7 # corresponds to stage 6, since increased by 1 after creating path
205 path.add_module('Progress')
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)