Belle II Software  light-2205-abys
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 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 print(path)
82 b2.process(path, max_event=10000)
83 
84 # Check availability of training input and cache file for stage 0
85 assert len(glob.glob('RootOutput.root')) == 1
86 assert len(glob.glob('training_input.root')) == 1
87 f = ROOT.TFile.Open("training_input.root", "read")
88 assert sum(['gamma' in key.GetName()[:5] for key in f.GetListOfKeys()]) == 2
89 assert sum(['mu+' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
90 assert sum(['pi+' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
91 assert sum(['K+' in key.GetName()[:3] for key in f.GetListOfKeys()]) == 1
92 f.Close()
93 
94 # Perform stage 0 training
95 fei.do_trainings(particles, configuration)
96 
97 # Moving training_input.root to training_input_stage0.root for later usage
98 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
99 
100 # Check availability of stage 0 *.xml training files
101 assert len(glob.glob('gamma*.xml')) == 2
102 assert len(glob.glob('mu+*.xml')) == 1
103 assert len(glob.glob('pi+*.xml')) == 1
104 assert len(glob.glob('K+*.xml')) == 1
105 
106 # Construct path for production of stage 1 training data
107 path = b2.create_path()
108 ma.inputMdstList(['./RootOutput.root'], path)
109 
110 maxTracks = 12
111 empty_path = b2.create_path()
112 skimfilter = b2.register_module('VariableToReturnValue')
113 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
114 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
115 path.add_module(skimfilter)
116 
117 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True, monitor=False)
118 feistate = fei.get_path(particles, configuration)
119 path.add_path(feistate.path)
120 path.add_module('RootOutput')
121 
122 assert feistate.stage == 2 # corresponds to stage 1, since increased by 1 after creating path
123 print(path)
124 b2.process(path, max_event=10000)
125 
126 # Check availability of training input and cache file for stage 1
127 assert len(glob.glob('RootOutput.root')) == 1
128 assert len(glob.glob('training_input.root')) == 1
129 f = ROOT.TFile.Open("training_input.root", "read")
130 assert sum(['pi0' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
131 f.Close()
132 
133 # Perform stage 1 training
134 fei.do_trainings(particles, configuration)
135 
136 # Moving training_input.root to training_input_stage1.root for later usage
137 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
138 
139 # Check availability of stage 1 *.xml training files
140 assert len(glob.glob('pi0*.xml')) == 1
141 
142 # Construct path for production of stage 3 training data (stage 2 is skipped)
143 path = b2.create_path()
144 ma.inputMdstList(['./RootOutput.root'], path)
145 
146 maxTracks = 12
147 empty_path = b2.create_path()
148 skimfilter = b2.register_module('VariableToReturnValue')
149 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
150 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
151 path.add_module(skimfilter)
152 
153 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=feistate.stage, training=True, monitor=False)
154 feistate = fei.get_path(particles, configuration)
155 path.add_path(feistate.path)
156 path.add_module('RootOutput')
157 
158 assert feistate.stage == 4 # corresponds to stage 3, since increased by 1 after creating path
159 print(path)
160 b2.process(path, max_event=10000)
161 
162 # Check availability of training input and cache file for stage 3
163 assert len(glob.glob('RootOutput.root')) == 1
164 assert len(glob.glob('training_input.root')) == 1
165 f = ROOT.TFile.Open("training_input.root", "read")
166 assert sum(['D' in key.GetName()[:2] for key in f.GetListOfKeys()]) == 5
167 f.Close()
168 
169 # Perform stage 3 training
170 fei.do_trainings(particles, configuration)
171 
172 # Moving training_input.root to training_input_stage3.root for later usage
173 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
174 
175 # Check availability of stage 3 *.xml training files
176 assert len(glob.glob('D*.xml')) == 5
177 
178 # Merge training input files for validation
179 subprocess.call(["analysis-fei-mergefiles", "training_input.root"] + glob.glob("training_input_stage*.root"))
180 
181 # Expect 4 different training_input*.root files now
182 assert len(glob.glob('training_input*.root')) == 4
183 
184 # Construct path for stage 6 preparing evaluation (stages 4 and 5 skipped, input evaluated from stage 0 on)
185 path = b2.create_path()
186 ma.inputMdst(filename=b2.find_file('mdst14.root', 'validation', False),
187  path=path)
188 
189 maxTracks = 12
190 empty_path = b2.create_path()
191 skimfilter = b2.register_module('VariableToReturnValue')
192 skimfilter.param('variable', 'nCleanedTracks(dr < 2 and abs(dz) < 4)')
193 skimfilter.if_value('>{}'.format(maxTracks), empty_path, b2.AfterConditionPath.END)
194 path.add_module(skimfilter)
195 
196 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, monitor=True)
197 feistate = fei.get_path(particles, configuration)
198 path.add_path(feistate.path)
199 
200 assert feistate.stage == 7 # corresponds to stage 6, since increased by 1 after creating path
201 print(path)
202 b2.process(path, max_event=10000)
203 
204 # Check availability of monitoring files constructed at stage 6 for validation
205 assert len(glob.glob('Monitor_FSPLoader.root')) == 1
206 
207 assert len(glob.glob('Monitor_ModuleStatistics.root')) == 1
208 
209 assert len(glob.glob('Monitor_Final.root')) == 1
210 f = ROOT.TFile.Open("Monitor_Final.root", "read")
211 assert f.GetListOfKeys().GetEntries() == 7
212 f.Close()
213 
214 prereconstruction_files = [
215  'Monitor_PreReconstruction_BeforeRanking.root',
216  'Monitor_PreReconstruction_AfterRanking.root',
217  'Monitor_PreReconstruction_AfterVertex.root',
218  'Monitor_PostReconstruction_AfterMVA.root'
219 ]
220 
221 for fname in prereconstruction_files:
222  assert len(glob.glob(fname)) == 1
223  f = ROOT.TFile.Open(fname, "read")
224  assert f.GetListOfKeys().GetEntries() == 11
225  f.Close()
226 
227 postreconstruction_files = [
228  'Monitor_PostReconstruction_BeforePostCut.root',
229  'Monitor_PostReconstruction_BeforeRanking.root',
230  'Monitor_PostReconstruction_AfterRanking.root'
231 ]
232 
233 for fname in postreconstruction_files:
234  assert len(glob.glob(fname)) == 1
235  f = ROOT.TFile.Open(fname, "read")
236  assert f.GetListOfKeys().GetEntries() == 7
237  f.Close()
238 
239 shutil.rmtree(tempdir)