Belle II Software release-09-00-11
klm_alignment.py
1
8
9# KLM alignment validation
10
11from prompt import ValidationSettings
12import os
13import basf2
14from ROOT.Belle2 import KLMCalibrationChecker
15import uproot
16import numpy as np
17import matplotlib.pyplot as plt
18from matplotlib.backends.backend_pdf import PdfPages
19import glob
20
21
22settings = ValidationSettings(name='KLM alignment',
23 description=__doc__,
24 download_files=['stdout'],
25 expert_config=None)
26
27# Avoid looking for the release globaltag
28basf2.conditions.override_globaltags()
29
30
31def get_result(job_path, tmp_dir):
32 database_file = f'{job_path}/outputdb/database.txt'
33 exp_run_list = []
34 with open(database_file) as f:
35 for line in f:
36 fields = line.split(' ')
37 if (fields[0] == 'dbstore/BKLMAlignment'):
38 iov = fields[2].split(',')
39 exp_run_list.append([int(iov[0]), int(iov[1])])
40
41 for exp_run in exp_run_list:
42 exp = exp_run[0]
43 run = exp_run[1]
44 checker = KLMCalibrationChecker()
45 checker.setExperimentRun(exp, run)
46 checker.setTestingPayload(database_file)
47 basf2.B2INFO(f'Creating alignment results tree for experiment {exp}, run {run}.')
48 checker.setAlignmentResultsFile(tmp_dir+f'/alignment_{exp}_{run}.root')
49 checker.checkAlignment()
50
51 return exp_run_list
52
53
54def get_residuals(data_path, data_path_prev):
55 # convert the alignment results to numpy.ndarray and calculate the residuals
56 data_prev = uproot.open(data_path_prev)
57 data = uproot.open(data_path)
58 BKLMModule_prev = data_prev[data_prev.keys()[0]]
59 EKLMModule_prev = data_prev[data_prev.keys()[1]]
60 BKLMModule = data[data.keys()[0]]
61 EKLMModule = data[data.keys()[1]]
62 EKLM = [EKLMModule_prev.arrays(library='pd'), EKLMModule.arrays(library='pd')]
63 BKLM = [BKLMModule_prev.arrays(library='pd'), BKLMModule.arrays(library='pd')]
64
65 # Fill numpy.ndarray with the alignment results
66 EKLM_values = np.zeros((2, len(EKLM[0]['section'].unique()),
67 len(EKLM[0]['sector'].unique()),
68 3, len(EKLM[0]['layer'].unique())))
69 EKLM_errors = np.zeros((2, len(EKLM[0]['section'].unique()),
70 len(EKLM[0]['sector'].unique()),
71 3, len(EKLM[0]['layer'].unique())))
72
73 BKLM_values = np.zeros((2, len(BKLM[0]['section'].unique()),
74 len(BKLM[0]['sector'].unique()),
75 3, len(BKLM[0]['layer'].unique())))
76 BKLM_errors = np.zeros((2, len(BKLM[0]['section'].unique()),
77 len(BKLM[0]['sector'].unique()),
78 3, len(BKLM[0]['layer'].unique())))
79
80 pars = {1: 1, 2: 2, 3: 6}
81
82 for i in [0, 1]:
83 for section in EKLM[i]['section'].unique():
84 for sector in EKLM[i]['sector'].unique():
85 for param in range(0, len(pars)):
86 req = (
87 (EKLM[i]['section'] == section) & (
88 EKLM[i]['sector'] == sector) & (
89 EKLM[i]['param'] == list(
90 pars.values())[param]))
91 if (section == 1):
92 EKLM_values[i][section-1][sector-1][list(pars.keys())[param] -
93 1] = np.append(np.array(EKLM[i][req]['value']), [0, 0])
94 EKLM_errors[i][section-1][sector-1][list(pars.keys())[param] -
95 1] = np.append(np.array(EKLM[i][req]['error']), [0, 0])
96 else:
97 EKLM_values[i][section-1][sector-1][list(pars.keys())[param]-1] = np.array(EKLM[i][req]['value'])
98 EKLM_errors[i][section-1][sector-1][list(pars.keys())[param]-1] = np.array(EKLM[i][req]['error'])
99
100 for i in [0, 1]:
101 for section in BKLM[i]['section'].unique():
102 for sector in BKLM[i]['sector'].unique():
103 for param in range(0, len(pars)):
104 req = (
105 (BKLM[i]['section'] == section) & (
106 BKLM[i]['sector'] == sector) & (
107 BKLM[i]['param'] == list(
108 pars.values())[param]))
109 BKLM_values[i][section][sector-1][list(pars.keys())[param]-1] = np.array(BKLM[i][req]['value'])
110 BKLM_errors[i][section][sector-1][list(pars.keys())[param]-1] = np.array(BKLM[i][req]['error'])
111
112 # Calculate the residuals
113 EKLM_res_values = np.zeros((len(EKLM[0]['section'].unique()),
114 len(EKLM[0]['sector'].unique()),
115 3, len(EKLM[0]['layer'].unique())))
116 EKLM_res_errors = np.zeros((len(EKLM[0]['section'].unique()),
117 len(EKLM[0]['sector'].unique()),
118 3, len(EKLM[0]['layer'].unique())))
119
120 BKLM_res_values = np.zeros((len(BKLM[0]['section'].unique()),
121 len(BKLM[0]['sector'].unique()),
122 3, len(BKLM[0]['layer'].unique())))
123 BKLM_res_errors = np.zeros((len(BKLM[0]['section'].unique()),
124 len(BKLM[0]['sector'].unique()),
125 3, len(BKLM[0]['layer'].unique())))
126
127 for section in range(0, EKLM_values[0].shape[0]):
128 for sector in range(0, EKLM_values[0].shape[1]):
129 for param in range(0, EKLM_values[0].shape[2]):
130 EKLM_res_values[section][sector][param] = EKLM_values[1][section][sector][param] - \
131 EKLM_values[0][section][sector][param]
132 EKLM_res_errors[section][sector][param] = np.sqrt(
133 EKLM_errors[1][section][sector][param]**2 +
134 EKLM_errors[0][section][sector][param]**2)
135
136 for section in range(0, BKLM_values[0].shape[0]):
137 for sector in range(0, BKLM_values[0].shape[1]):
138 for param in range(0, BKLM_values[0].shape[2]):
139 BKLM_res_values[section][sector][param] = BKLM_values[1][section][sector][param] - \
140 BKLM_values[0][section][sector][param]
141 BKLM_res_errors[section][sector][param] = np.sqrt(
142 BKLM_errors[1][section][sector][param]**2 +
143 BKLM_errors[0][section][sector][param]**2)
144
145 EKLM_chi2 = np.zeros((len(EKLM[0]['section'].unique()),
146 len(EKLM[0]['sector'].unique()),
147 3, len(EKLM[0]['layer'].unique())))
148
149 BKLM_chi2 = np.zeros((len(BKLM[0]['section'].unique()),
150 len(BKLM[0]['sector'].unique()),
151 3, len(BKLM[0]['layer'].unique())))
152
153 for section in range(0, EKLM_res_values.shape[0]):
154 for sector in range(0, EKLM_res_values.shape[1]):
155 for param in range(0, EKLM_res_values.shape[2]):
156 for layer in range(0, EKLM_res_values.shape[3]):
157 if ((EKLM_res_values[section][sector][param][layer] == 0) |
158 (EKLM_res_errors[section][sector][param][layer] == 0)):
159 EKLM_chi2[section][sector][param][layer] = 0
160 else:
161 EKLM_chi2[section][sector][param][layer] = (
162 EKLM_res_values[section][sector][param][layer]**2)/(EKLM_res_errors[section][sector][param][layer]**2)
163
164 for section in range(0, BKLM_res_values.shape[0]):
165 for sector in range(0, BKLM_res_values.shape[1]):
166 for param in range(0, BKLM_res_values.shape[2]):
167 for layer in range(0, BKLM_res_values.shape[3]):
168 if ((BKLM_res_values[section][sector][param][layer] == 0) |
169 (BKLM_res_errors[section][sector][param][layer] == 0)):
170 BKLM_chi2[section][sector][param][layer] = 0
171 else:
172 BKLM_chi2[section][sector][param][layer] = (
173 BKLM_res_values[section][sector][param][layer]**2)/(BKLM_res_errors[section][sector][param][layer]**2)
174
175 return [EKLM_res_values, EKLM_res_errors, EKLM_chi2, BKLM_res_values, BKLM_res_errors, BKLM_chi2]
176
177
178def draw_EKLM_pics(EKLM_values, EKLM_errors, EKLM_chi2, pdfPages):
179 # Draw the EKLM alignment residuals and add them to a .pdf file
180 plt.rcParams.update({
181 'font.size': 20,
182 'figure.figsize': (11, 10),
183 'axes.grid': True,
184 'grid.linestyle': '-',
185 'grid.alpha': 0.2,
186 'lines.markersize': 5.0,
187 'xtick.minor.visible': True,
188 'xtick.direction': 'in',
189 'xtick.major.size': 20.0,
190 'xtick.minor.size': 10.0,
191 'xtick.top': True,
192 'ytick.minor.visible': True,
193 'ytick.direction': 'in',
194 'ytick.major.size': 20.0,
195 'ytick.minor.size': 10.0,
196 'ytick.right': True,
197 'errorbar.capsize': 0.0,
198 })
199 param_meaning = {0: 'x', 1: 'y', 2: r'$\alpha$'}
200 section_meaning = {0: 'b', 1: 'f'}
201 layers = {'EKLM': np.arange(1, 15, 1), 'BKLM': np.arange(1, 16, 1)}
202 layers_err = {'EKLM': np.full(14, 0.5), 'BKLM': np.full(15, 0.5)}
203 for section in [0, 1]:
204 fig, axs = plt.subplots(4, 3, figsize=(20, 20))
205 for i in range(0, 12):
206 sector = i//3
207 param = i % 3
208 plt.sca(axs[sector][param])
209 plt.errorbar(
210 x=layers['EKLM'],
211 xerr=layers_err['EKLM'],
212 y=EKLM_values[section][sector][param],
213 yerr=EKLM_errors[section][sector][param],
214 ls='',
215 fmt='o',
216 ds='steps-mid',
217 color='black',
218 label='EKLM ' +
219 section_meaning[section] +
220 ' ' +
221 str(sector) +
222 r' $\chi^{2}$=' +
223 str(
224 np.around(
225 np.sum(
226 EKLM_chi2,
227 axis=3)[section][sector][param],
228 1)))
229 plt.hlines(0, 0, 14, color='red')
230 if (param == 2):
231 plt.ylim(-0.02, 0.02)
232 plt.ylabel(r'$\Delta$'+param_meaning[param]+' rad')
233 else:
234 plt.ylim(-2, 2)
235 plt.ylabel(r'$\Delta$'+param_meaning[param]+' cm')
236 plt.xlabel('Layer')
237 axs[sector][param].yaxis.set_label_coords(-0.1, 0.5)
238 plt.legend()
239 fig.tight_layout()
240 plt.savefig(pdfPages, format='pdf')
241 plt.close('all')
242
243
244def draw_BKLM_pics(BKLM_values, BKLM_errors, BKLM_chi2, pdfPages):
245 # Draw the BKLM alignment residuals and add them to a .pdf file
246 plt.rcParams.update({
247 'font.size': 20,
248 'figure.figsize': (11, 10),
249 'axes.grid': True,
250 'grid.linestyle': '-',
251 'grid.alpha': 0.2,
252 'lines.markersize': 5.0,
253 'xtick.minor.visible': True,
254 'xtick.direction': 'in',
255 'xtick.major.size': 20.0,
256 'xtick.minor.size': 10.0,
257 'xtick.top': True,
258 'ytick.minor.visible': True,
259 'ytick.direction': 'in',
260 'ytick.major.size': 20.0,
261 'ytick.minor.size': 10.0,
262 'ytick.right': True,
263 'errorbar.capsize': 0.0,
264 })
265 param_meaning = {0: 'x', 1: 'y', 2: r'$\alpha$'}
266 section_meaning = {0: 'b', 1: 'f'}
267 layers = {'EKLM': np.arange(1, 15, 1), 'BKLM': np.arange(1, 16, 1)}
268 layers_err = {'EKLM': np.full(14, 0.5), 'BKLM': np.full(15, 0.5)}
269 for section in [0, 1]:
270 for sector_shift in [0, 4]:
271 fig, axs = plt.subplots(4, 3, figsize=(20, 20))
272 for i in range(0, 12):
273 sector = i//3+sector_shift
274 param = i % 3
275 plt.sca(axs[sector-sector_shift][param])
276 plt.errorbar(
277 x=layers['BKLM'],
278 xerr=layers_err['BKLM'],
279 y=BKLM_values[section][sector][param],
280 yerr=BKLM_errors[section][sector][param],
281 ls='',
282 fmt='o',
283 ds='steps-mid',
284 color='black',
285 label='BKLM ' +
286 section_meaning[section] +
287 ' ' +
288 str(sector) +
289 r' $\chi^{2}$=' +
290 str(
291 np.around(
292 np.sum(
293 BKLM_chi2,
294 axis=3)[section][sector][param],
295 1)))
296 plt.hlines(0, 0, 15, color='red')
297 if (param == 2):
298 plt.ylim(-0.02, 0.02)
299 plt.ylabel(r'$\Delta$'+param_meaning[param]+' rad')
300 else:
301 plt.ylim(-2, 2)
302 plt.ylabel(r'$\Delta$'+param_meaning[param]+' cm')
303 plt.xlabel('Layer')
304 axs[sector-sector_shift][param].yaxis.set_label_coords(-0.1, 0.5)
305 plt.legend()
306 fig.tight_layout()
307 plt.savefig(pdfPages, format='pdf')
308 plt.close('all')
309
310
311def run_validation(calibration_results_dir, input_data_path=None, **kwargs):
312 '''
313 Run the validation.
314 The script compares the most recent alignment result with the previous results by calculating the residuals.
315 '''
316 tmp_work_dir = os.path.join(os.getcwd(), 'tmp_work')
317 tmp_plot_dir = os.path.join(os.getcwd(), 'tmp_plot')
318 if not os.path.exists(tmp_work_dir):
319 os.makedirs(tmp_work_dir)
320 if not os.path.exists(tmp_plot_dir):
321 os.makedirs(tmp_plot_dir)
322
323 # Find the latest iterations' directories
324 iterations = [d for d in glob.glob(f'{calibration_results_dir}/KLMAlignment/?')]
325 iterations = sorted(iterations, key=lambda x: int(x.split('/')[-1]), reverse=True)[:2]
326 if len(iterations) < 2:
327 raise ValueError("Not enough KLMAlignment iterations found.")
328
329 job_path = f'{iterations[0]}/algorithm_output'
330 job_path_prev = f'{iterations[1]}/algorithm_output'
331
332 # Create alignment results tree the recent and previous calibration and get IoVs
333 exp_run_list = get_result(job_path, tmp_work_dir)
334 exp_run_list_prev = get_result(job_path_prev, tmp_work_dir)
335 # Sort IoV from earliest to latest
336 sorted_exp_run_list = sorted(exp_run_list + exp_run_list_prev)
337 # Calculate the residuals for each adjacent pair of IoVs and saves the results of the comparison in a .pdf file
338 for i in range(0, len(sorted_exp_run_list)-1):
339 exp_prev = sorted_exp_run_list[i][0]
340 run_prev = sorted_exp_run_list[i][1]
341 exp = sorted_exp_run_list[i+1][0]
342 run = sorted_exp_run_list[i+1][1]
343 data_path = tmp_work_dir+f'/alignment_{exp_prev}_{run_prev}.root'
344 data_path_prev = tmp_work_dir+f'/alignment_{exp}_{run}.root'
345 EKLM_values, EKLM_errors, EKLM_chi2, BKLM_values, BKLM_errors, BKLM_chi2 = get_residuals(data_path, data_path_prev)
346 pdfPages = PdfPages(tmp_plot_dir+'/e'+str(exp_prev)+'r'+str(run_prev)+'_e'+str(exp)+'r'+str(run)+'.pdf')
347 draw_EKLM_pics(EKLM_values, EKLM_errors, EKLM_chi2, pdfPages)
348 draw_BKLM_pics(BKLM_values, BKLM_errors, BKLM_chi2, pdfPages)
349 pdfPages.close()
350
351
352if __name__ == "__main__":
353
354 import argparse
355 parser = argparse.ArgumentParser(description=__doc__,
356 formatter_class=argparse.RawTextHelpFormatter)
357
358 # b2val-prompt-run wants to pass to the script also input_data_path
359 # and requested_iov. As they are not required by this validation I just accept
360 # them together with calibration_results_dir and then ignore them
361 parser.add_argument('calibration_results_dir',
362 help='The directory that contains the collector outputs',
363 nargs='+')
364
365 parser.add_argument('-o', '--output_dir',
366 help='The directory where all the output will be saved',
367 default='KLMAlignmentValidation_output')
368 args = parser.parse_args()
369
370 run_validation(args.calibration_results_dir[0])
run_validation
Definition: run_validation.py:1