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