Belle II Software  light-2205-abys
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 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 print(path)
83 b2.process(path, max_event=10000)
84 
85 # Check availability of training input and cache file for stage 0
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 feistate = fei.get_path(particles, configuration)
109 roe_path = b2.create_path()
110 cond_module = b2.register_module('SignalSideParticleFilter')
111 cond_module.param('particleLists', ['B+:sig'])
112 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
113 roe_path.add_module(cond_module)
114 path.add_path(sig_path)
115 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
116 
117 assert feistate.stage == 2 # corresponds to stage 1, since increased by 1 after creating path
118 print(path)
119 b2.process(path, max_event=10000)
120 
121 # Check availability of training input and cache file for stage 1
122 assert len(glob.glob('training_input.root')) == 1
123 f = ROOT.TFile.Open("training_input.root", "read")
124 assert sum(['pi0' in key.GetName()[:4] for key in f.GetListOfKeys()]) == 1
125 f.Close()
126 
127 # Perform stage 1 training
128 fei.do_trainings(particles, configuration)
129 
130 # Moving training_input.root to training_input_stage1.root for later usage
131 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
132 
133 # Check availability of stage 1 *.xml training files
134 assert len(glob.glob('pi0*.xml')) == 1
135 
136 # Construct path for production of stage 3 training data (stage 2 is skipped)
137 path = b2.create_path()
138 feistate = fei.get_path(particles, configuration)
139 roe_path = b2.create_path()
140 cond_module = b2.register_module('SignalSideParticleFilter')
141 cond_module.param('particleLists', ['B+:sig'])
142 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
143 roe_path.add_module(cond_module)
144 path.add_path(sig_path)
145 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
146 
147 assert feistate.stage == 4 # corresponds to stage 3, since increased by 1 after creating path
148 print(path)
149 b2.process(path, max_event=10000)
150 
151 # Check availability of training input and cache file for stage 3
152 assert len(glob.glob('training_input.root')) == 1
153 f = ROOT.TFile.Open("training_input.root", "read")
154 assert sum(['D' in key.GetName()[:2] for key in f.GetListOfKeys()]) == 5
155 f.Close()
156 
157 # Perform stage 3 training
158 fei.do_trainings(particles, configuration)
159 
160 # Moving training_input.root to training_input_stage3.root for later usage
161 shutil.move("training_input.root", f"training_input_stage{feistate.stage-1}.root")
162 
163 # Check availability of stage 3 *.xml training files
164 assert len(glob.glob('D*.xml')) == 5
165 
166 # Merge training input files for validation
167 subprocess.call(["analysis-fei-mergefiles", "training_input.root"] + glob.glob("training_input_stage*.root"))
168 
169 # Expect 4 different training_input*.root files now
170 assert len(glob.glob('training_input*.root')) == 4
171 
172 # Construct path for stage 6 preparing evaluation (stages 4 and 5 skipped, input evaluated from stage 0 on)
173 configuration = fei.config.FeiConfiguration(prefix='FEI_VALIDATION', cache=0, monitor=True)
174 path = b2.create_path()
175 feistate = fei.get_path(particles, configuration)
176 roe_path = b2.create_path()
177 cond_module = b2.register_module('SignalSideParticleFilter')
178 cond_module.param('particleLists', ['B+:sig'])
179 cond_module.if_true(feistate.path, b2.AfterConditionPath.END)
180 roe_path.add_module(cond_module)
181 path.add_path(sig_path)
182 path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
183 
184 assert feistate.stage == 7 # corresponds to stage 6, since increased by 1 after creating path
185 print(path)
186 b2.process(path, max_event=10000)
187 
188 # Check availability of monitoring files constructed at stage 6 for validation
189 assert len(glob.glob('Monitor_FSPLoader.root')) == 1
190 
191 assert len(glob.glob('Monitor_ModuleStatistics.root')) == 1
192 
193 assert len(glob.glob('Monitor_Final.root')) == 1
194 f = ROOT.TFile.Open("Monitor_Final.root", "read")
195 assert f.GetListOfKeys().GetEntries() == 7
196 f.Close()
197 
198 prereconstruction_files = [
199  'Monitor_PreReconstruction_BeforeRanking.root',
200  'Monitor_PreReconstruction_AfterRanking.root',
201  'Monitor_PreReconstruction_AfterVertex.root',
202  'Monitor_PostReconstruction_AfterMVA.root'
203 ]
204 
205 for fname in prereconstruction_files:
206  assert len(glob.glob(fname)) == 1
207  f = ROOT.TFile.Open(fname, "read")
208  assert f.GetListOfKeys().GetEntries() == 11
209  f.Close()
210 
211 postreconstruction_files = [
212  'Monitor_PostReconstruction_BeforePostCut.root',
213  'Monitor_PostReconstruction_BeforeRanking.root',
214  'Monitor_PostReconstruction_AfterRanking.root'
215 ]
216 
217 for fname in postreconstruction_files:
218  assert len(glob.glob(fname)) == 1
219  f = ROOT.TFile.Open(fname, "read")
220  assert f.GetListOfKeys().GetEntries() == 7
221  f.Close()
222 
223 shutil.rmtree(tempdir)