Belle II Software  release-06-01-15
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 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(specific=True)
42 
43 # Construct path for reconstruction of signal side B-meson
44 sig_path = b2.create_path()
45 ma.inputMdst(environmentType='default',
46  filename=b2.find_file('mdst14.root', 'validation', False),
47  path=sig_path)
48 ma.fillParticleList('mu+:sig', 'muonID > 0.5 and dr < 1 and abs(dz) < 2', writeOut=True, path=sig_path)
49 ma.reconstructDecay('tau+:sig -> mu+:sig', '', 1, writeOut=True, path=sig_path)
50 ma.matchMCTruth('tau+:sig', path=sig_path)
51 ma.reconstructDecay('B+:sig -> tau+:sig', '', writeOut=True, path=sig_path)
52 ma.matchMCTruth('B+:sig', path=sig_path)
53 ma.buildRestOfEvent('B+:sig', path=sig_path)
54 
55 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=-1, training=True, monitor=False)
56 feistate = fei.get_path(particles, configuration)
57 
58 # Construct path for production of mcParticlesCount.root at stage -1
59 path = b2.create_path()
60 roe_path = b2.create_path()
61 cond_module = b2.register_module('SignalSideParticleFilter')
62 cond_module.param('particleLists', ['B+:sig'])
63 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
64 roe_path.add_module(cond_module)
65 path.add_path(sig_path)
66 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
67 
68 assert feistate.stage == 0 # corresponds to stage -1, since increased by 1 after creating path
69 print(path)
70 b2.process(path, max_event=10000)
71 assert len(glob.glob('mcParticlesCount.root')) == 1
72 
73 # Construct path for production of stage 0 training data
74 path = b2.create_path()
75 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, training=True, monitor=False)
76 feistate = fei.get_path(particles, configuration)
77 roe_path = b2.create_path()
78 cond_module = b2.register_module('SignalSideParticleFilter')
79 cond_module.param('particleLists', ['B+:sig'])
80 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
81 roe_path.add_module(cond_module)
82 path.add_path(sig_path)
83 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
84 
85 assert feistate.stage == 1 # corresponds to stage 0, since increased by 1 after creating path
86 print(path)
87 b2.process(path, max_event=10000)
88 
89 # Check availability of training input and cache file for stage 0
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 feistate = fei.get_path(particles, configuration)
113 roe_path = b2.create_path()
114 cond_module = b2.register_module('SignalSideParticleFilter')
115 cond_module.param('particleLists', ['B+:sig'])
116 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
117 roe_path.add_module(cond_module)
118 path.add_path(sig_path)
119 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
120 
121 assert feistate.stage == 2 # corresponds to stage 1, since increased by 1 after creating path
122 print(path)
123 b2.process(path, max_event=10000)
124 
125 # Check availability of training input and cache file for stage 1
126 assert len(glob.glob('training_input.root')) == 1
127 f = ROOT.TFile.Open("training_input.root", "read")
128 assert sum(['pi0' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
129 f.Close()
130 
131 # Perform stage 1 training
132 fei.do_trainings(particles, configuration)
133 
134 # Moving training_input.root to training_input_stage1.root for later usage
135 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
136 
137 # Check availability of stage 1 *.xml training files
138 assert len(glob.glob('pi0*.xml')) == 1
139 
140 # Construct path for production of stage 3 training data (stage 2 is skipped)
141 path = b2.create_path()
142 feistate = fei.get_path(particles, configuration)
143 roe_path = b2.create_path()
144 cond_module = b2.register_module('SignalSideParticleFilter')
145 cond_module.param('particleLists', ['B+:sig'])
146 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
147 roe_path.add_module(cond_module)
148 path.add_path(sig_path)
149 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
150 
151 assert feistate.stage == 4 # corresponds to stage 3, since increased by 1 after creating path
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 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)