Belle II Software  release-08-01-10
test_fei_specific.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(specific=True)
38 
39 # Construct path for reconstruction of signal side B-meson
40 sig_path = b2.create_path()
41 ma.inputMdst(filename=b2.find_file('mdst14.root', 'validation', False),
42  path=sig_path)
43 ma.fillParticleList('mu+:sig', 'muonID > 0.5 and dr < 1 and abs(dz) < 2', writeOut=True, path=sig_path)
44 ma.reconstructDecay('tau+:sig -> mu+:sig', '', 1, writeOut=True, path=sig_path)
45 ma.matchMCTruth('tau+:sig', path=sig_path)
46 ma.reconstructDecay('B+:sig -> tau+:sig', '', writeOut=True, path=sig_path)
47 ma.matchMCTruth('B+:sig', path=sig_path)
48 ma.buildRestOfEvent('B+:sig', path=sig_path)
49 
50 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=-1, training=True, monitor=False)
51 feistate = fei.get_path(particles, configuration)
52 
53 # Construct path for production of mcParticlesCount.root at stage -1
54 path = b2.create_path()
55 roe_path = b2.create_path()
56 cond_module = b2.register_module('SignalSideParticleFilter')
57 cond_module.param('particleLists', ['B+:sig'])
58 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
59 roe_path.add_module(cond_module)
60 path.add_path(sig_path)
61 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
62 
63 assert feistate.stage == 0 # corresponds to stage -1, since increased by 1 after creating path
64 path.add_module('Progress')
65 print(path)
66 b2.process(path, max_event=10000)
67 assert len(glob.glob('mcParticlesCount.root')) == 1
68 
69 # Construct path for production of stage 0 training data
70 path = b2.create_path()
71 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, training=True, monitor=False)
72 feistate = fei.get_path(particles, configuration)
73 roe_path = b2.create_path()
74 cond_module = b2.register_module('SignalSideParticleFilter')
75 cond_module.param('particleLists', ['B+:sig'])
76 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
77 roe_path.add_module(cond_module)
78 path.add_path(sig_path)
79 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
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('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 feistate = fei.get_path(particles, configuration)
110 roe_path = b2.create_path()
111 cond_module = b2.register_module('SignalSideParticleFilter')
112 cond_module.param('particleLists', ['B+:sig'])
113 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
114 roe_path.add_module(cond_module)
115 path.add_path(sig_path)
116 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
117 
118 assert feistate.stage == 2 # corresponds to stage 1, since increased by 1 after creating path
119 path.add_module('Progress')
120 print(path)
121 b2.process(path, max_event=10000)
122 
123 # Check availability of training input and cache file for stage 1
124 assert len(glob.glob('training_input.root')) == 1
125 f = ROOT.TFile.Open("training_input.root", "read")
126 assert sum(['pi0' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
127 f.Close()
128 
129 # Perform stage 1 training
130 fei.do_trainings(particles, configuration)
131 
132 # Moving training_input.root to training_input_stage1.root for later usage
133 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
134 
135 # Check availability of stage 1 *.xml training files
136 assert len(glob.glob('pi0*.xml')) == 1
137 
138 # Construct path for production of stage 3 training data (stage 2 is skipped)
139 path = b2.create_path()
140 feistate = fei.get_path(particles, configuration)
141 roe_path = b2.create_path()
142 cond_module = b2.register_module('SignalSideParticleFilter')
143 cond_module.param('particleLists', ['B+:sig'])
144 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
145 roe_path.add_module(cond_module)
146 path.add_path(sig_path)
147 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
148 
149 assert feistate.stage == 4 # corresponds to stage 3, since increased by 1 after creating path
150 path.add_module('Progress')
151 print(path)
152 b2.process(path, max_event=10000)
153 
154 # Check availability of training input and cache file for stage 3
155 assert len(glob.glob('training_input.root')) == 1
156 f = ROOT.TFile.Open("training_input.root", "read")
157 assert sum(['D' in key.GetName()[:2] for key in f.GetListOfKeys()]) == 5
158 f.Close()
159 
160 # Perform stage 3 training
161 fei.do_trainings(particles, configuration)
162 
163 # Moving training_input.root to training_input_stage3.root for later usage
164 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
165 
166 # Check availability of stage 3 *.xml training files
167 assert len(glob.glob('D*.xml')) == 5
168 
169 # Merge training input files for validation
170 subprocess.call(["analysis-fei-mergefiles", "training_input.root"] + glob.glob("training_input_stage*.root"))
171 
172 # Expect 4 different training_input*.root files now
173 assert len(glob.glob('training_input*.root')) == 4
174 
175 # Construct path for stage 6 preparing evaluation (stages 4 and 5 skipped, input evaluated from stage 0 on)
176 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, monitor=True)
177 path = b2.create_path()
178 feistate = fei.get_path(particles, configuration)
179 roe_path = b2.create_path()
180 cond_module = b2.register_module('SignalSideParticleFilter')
181 cond_module.param('particleLists', ['B+:sig'])
182 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
183 roe_path.add_module(cond_module)
184 path.add_path(sig_path)
185 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
186 
187 assert feistate.stage == 7 # corresponds to stage 6, since increased by 1 after creating path
188 path.add_module('Progress')
189 print(path)
190 b2.process(path, max_event=10000)
191 
192 # Check availability of monitoring files constructed at stage 6 for validation
193 assert len(glob.glob('Monitor_FSPLoader.root')) == 1
194 
195 assert len(glob.glob('Monitor_ModuleStatistics.root')) == 1
196 
197 assert len(glob.glob('Monitor_Final.root')) == 1
198 f = ROOT.TFile.Open("Monitor_Final.root", "read")
199 assert f.GetListOfKeys().GetEntries() == 7
200 f.Close()
201 
202 prereconstruction_files = [
203  'Monitor_PreReconstruction_BeforeRanking.root',
204  'Monitor_PreReconstruction_AfterRanking.root',
205  'Monitor_PreReconstruction_AfterVertex.root',
206  'Monitor_PostReconstruction_AfterMVA.root'
207 ]
208 
209 for fname in prereconstruction_files:
210  assert len(glob.glob(fname)) == 1
211  f = ROOT.TFile.Open(fname, "read")
212  assert f.GetListOfKeys().GetEntries() == 11
213  f.Close()
214 
215 postreconstruction_files = [
216  'Monitor_PostReconstruction_BeforePostCut.root',
217  'Monitor_PostReconstruction_BeforeRanking.root',
218  'Monitor_PostReconstruction_AfterRanking.root'
219 ]
220 
221 for fname in postreconstruction_files:
222  assert len(glob.glob(fname)) == 1
223  f = ROOT.TFile.Open(fname, "read")
224  assert f.GetListOfKeys().GetEntries() == 7
225  f.Close()
226 
227 shutil.rmtree(tempdir)