Belle II Software  release-08-01-10
klm_alignment.py
1 
8 
9 # KLM alignment validation
10 
11 from prompt import ValidationSettings
12 import sys
13 import os
14 import basf2
15 from ROOT.Belle2 import KLMCalibrationChecker
16 import uproot
17 import numpy as np
18 import matplotlib.pyplot as plt
19 from matplotlib.backends.backend_pdf import PdfPages
20 
21 
22 settings = ValidationSettings(name='KLM alignment',
23  description=__doc__,
24  download_files=['stdout'],
25  expert_config=None)
26 
27 
28 def 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 
51 def 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 
175 def 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 
241 def 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 
308 def 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 
340 if __name__ == "__main__":
341  run_validation(*sys.argv[1:])