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):
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}.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 [EKLM_res_values, EKLM_res_errors, EKLM_chi2, BKLM_res_values, BKLM_res_errors, BKLM_chi2]
178
179
180def draw_EKLM_pics(EKLM_values, EKLM_errors, EKLM_chi2, pdfPages):
181 # Draw the EKLM alignment residuals and add them to a .pdf file
182 plt.rcParams.update({
183 'font.size': 20,
184 'figure.figsize': (11, 10),
185 'axes.grid': True,
186 'grid.linestyle': '-',
187 'grid.alpha': 0.2,
188 'lines.markersize': 5.0,
189 'xtick.minor.visible': True,
190 'xtick.direction': 'in',
191 'xtick.major.size': 20.0,
192 'xtick.minor.size': 10.0,
193 'xtick.top': True,
194 'ytick.minor.visible': True,
195 'ytick.direction': 'in',
196 'ytick.major.size': 20.0,
197 'ytick.minor.size': 10.0,
198 'ytick.right': True,
199 'errorbar.capsize': 0.0,
200 })
201 param_meaning = {0: 'x', 1: 'y', 2: r'$\alpha$'}
202 section_meaning = {0: 'b', 1: 'f'}
203 layers = {'EKLM': np.arange(1, 15, 1), 'BKLM': np.arange(1, 16, 1)}
204 layers_err = {'EKLM': np.full(14, 0.5), 'BKLM': np.full(15, 0.5)}
205 for section in [0, 1]:
206 fig, axs = plt.subplots(4, 3, figsize=(20, 20))
207 for i in range(0, 12):
208 sector = i//3
209 param = i % 3
210 plt.sca(axs[sector][param])
211 plt.errorbar(
212 x=layers['EKLM'],
213 xerr=layers_err['EKLM'],
214 y=EKLM_values[section][sector][param],
215 yerr=EKLM_errors[section][sector][param],
216 ls='',
217 fmt='o',
218 ds='steps-mid',
219 color='black',
220 label='EKLM ' +
221 section_meaning[section] +
222 ' ' +
223 str(sector) +
224 r' $\chi^{2}$=' +
225 str(
226 np.around(
227 np.sum(
228 EKLM_chi2,
229 axis=3)[section][sector][param],
230 1)))
231 plt.hlines(0, 0, 14, color='red')
232 if (param == 2):
233 plt.ylim(-0.02, 0.02)
234 plt.ylabel(r'$\Delta$'+param_meaning[param]+' rad')
235 else:
236 plt.ylim(-2, 2)
237 plt.ylabel(r'$\Delta$'+param_meaning[param]+' cm')
238 plt.xlabel('Layer')
239 axs[sector][param].yaxis.set_label_coords(-0.1, 0.5)
240 plt.legend()
241 fig.tight_layout()
242 plt.savefig(pdfPages, format='pdf')
243 plt.close('all')
244
245
246def draw_BKLM_pics(BKLM_values, BKLM_errors, BKLM_chi2, pdfPages):
247 # Draw the BKLM alignment residuals and add them to a .pdf file
248 plt.rcParams.update({
249 'font.size': 20,
250 'figure.figsize': (11, 10),
251 'axes.grid': True,
252 'grid.linestyle': '-',
253 'grid.alpha': 0.2,
254 'lines.markersize': 5.0,
255 'xtick.minor.visible': True,
256 'xtick.direction': 'in',
257 'xtick.major.size': 20.0,
258 'xtick.minor.size': 10.0,
259 'xtick.top': True,
260 'ytick.minor.visible': True,
261 'ytick.direction': 'in',
262 'ytick.major.size': 20.0,
263 'ytick.minor.size': 10.0,
264 'ytick.right': True,
265 'errorbar.capsize': 0.0,
266 })
267 param_meaning = {0: 'x', 1: 'y', 2: r'$\alpha$'}
268 section_meaning = {0: 'b', 1: 'f'}
269 layers = {'EKLM': np.arange(1, 15, 1), 'BKLM': np.arange(1, 16, 1)}
270 layers_err = {'EKLM': np.full(14, 0.5), 'BKLM': np.full(15, 0.5)}
271 for section in [0, 1]:
272 for sector_shift in [0, 4]:
273 fig, axs = plt.subplots(4, 3, figsize=(20, 20))
274 for i in range(0, 12):
275 sector = i//3+sector_shift
276 param = i % 3
277 plt.sca(axs[sector-sector_shift][param])
278 plt.errorbar(
279 x=layers['BKLM'],
280 xerr=layers_err['BKLM'],
281 y=BKLM_values[section][sector][param],
282 yerr=BKLM_errors[section][sector][param],
283 ls='',
284 fmt='o',
285 ds='steps-mid',
286 color='black',
287 label='BKLM ' +
288 section_meaning[section] +
289 ' ' +
290 str(sector) +
291 r' $\chi^{2}$=' +
292 str(
293 np.around(
294 np.sum(
295 BKLM_chi2,
296 axis=3)[section][sector][param],
297 1)))
298 plt.hlines(0, 0, 15, color='red')
299 if (param == 2):
300 plt.ylim(-0.02, 0.02)
301 plt.ylabel(r'$\Delta$'+param_meaning[param]+' rad')
302 else:
303 plt.ylim(-2, 2)
304 plt.ylabel(r'$\Delta$'+param_meaning[param]+' cm')
305 plt.xlabel('Layer')
306 axs[sector-sector_shift][param].yaxis.set_label_coords(-0.1, 0.5)
307 plt.legend()
308 fig.tight_layout()
309 plt.savefig(pdfPages, format='pdf')
310 plt.close('all')
311
312
313def run_validation(calibration_results_dir, input_data_path=None, **kwargs):
314 '''
315 Run the validation.
316 The script compares the most recent alignment result with the previous results by calculating the residuals.
317 '''
318 tmp_work_dir = os.path.join(os.getcwd(), 'tmp_work')
319 tmp_plot_dir = os.path.join(os.getcwd(), 'tmp_plot')
320 if not os.path.exists(tmp_work_dir):
321 os.makedirs(tmp_work_dir)
322 if not os.path.exists(tmp_plot_dir):
323 os.makedirs(tmp_plot_dir)
324
325 # Find the latest iterations' directories
326 iterations = [d for d in glob.glob(f'{calibration_results_dir}/KLMAlignment/?')]
327 iterations = sorted(iterations, key=lambda x: int(x.split('/')[-1]), reverse=True)[:2]
328 if len(iterations) < 2:
329 raise ValueError("Not enough KLMAlignment iterations found.")
330
331 job_path = f'{iterations[0]}/algorithm_output'
332 job_path_prev = f'{iterations[1]}/algorithm_output'
333
334 # Create alignment results tree the recent and previous calibration and get IoVs
335 exp_run_list = get_result(job_path, tmp_work_dir)
336 exp_run_list_prev = get_result(job_path_prev, tmp_work_dir)
337 # Sort IoV from earliest to latest
338 sorted_exp_run_list = sorted(exp_run_list + exp_run_list_prev)
339 # Calculate the residuals for each adjacent pair of IoVs and saves the results of the comparison in a .pdf file
340 for i in range(0, len(sorted_exp_run_list)//2):
341 exp_prev = sorted_exp_run_list[2*i][0]
342 run_prev = sorted_exp_run_list[2*i][1]
343 exp = sorted_exp_run_list[2*i+1][0]
344 run = sorted_exp_run_list[2*i+1][1]
345 data_path = tmp_work_dir+f'/alignment_{exp_prev}_{run_prev}.root'
346 data_path_prev = tmp_work_dir+f'/alignment_{exp}_{run}.root'
347 EKLM_values, EKLM_errors, EKLM_chi2, BKLM_values, BKLM_errors, BKLM_chi2 = get_residuals(data_path, data_path_prev)
348 pdfPages = PdfPages(tmp_plot_dir+'/e'+str(exp_prev)+'r'+str(run_prev)+'_e'+str(exp)+'r'+str(run)+'.pdf')
349 draw_EKLM_pics(EKLM_values, EKLM_errors, EKLM_chi2, pdfPages)
350 draw_BKLM_pics(BKLM_values, BKLM_errors, BKLM_chi2, pdfPages)
351 pdfPages.close()
352
353
354if __name__ == "__main__":
355
356 import argparse
357 parser = argparse.ArgumentParser(description=__doc__,
358 formatter_class=argparse.RawTextHelpFormatter)
359
360 # b2val-prompt-run wants to pass to the script also input_data_path
361 # and requested_iov. As they are not required by this validation I just accept
362 # them together with calibration_results_dir and then ignore them
363 parser.add_argument('calibration_results_dir',
364 help='The directory that contains the collector outputs',
365 nargs='+')
366
367 parser.add_argument('-o', '--output_dir',
368 help='The directory where all the output will be saved',
369 default='KLMAlignmentValidation_output')
370 args = parser.parse_args()
371
372 run_validation(args.calibration_results_dir[0])
run_validation
Definition run_validation.py:1