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