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