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