Belle II Software development
klm_strip_efficiency.py
1
8
9'''
10Validation of KLM strip efficiency calibration.
11'''
12
13
14import basf2
15from prompt import ValidationSettings
16import ROOT
17from ROOT.Belle2 import BKLMElementNumbers, KLMCalibrationChecker, KLMElementNumbers
18import sys
19import subprocess
20import math
21import os
22import json
23
24
25settings = ValidationSettings(name='KLM strip efficiency',
26 description=__doc__,
27 download_files=['stdout'],
28 expert_config={
29 "chunk_size": 100
30 })
31
32
33def save_graph_to_root(graph_name):
34 '''
35 Save a TGraph in a ROOT file.
36 '''
37 graph = ROOT.gPad.GetPrimitive('Graph')
38 assert isinstance(graph, ROOT.TGraph) == 1
39 graph.SetName(graph_name)
40 graph.Write()
41
42
43def save_graph_to_pdf(canvas, root_file, graph_name, exp, chunk):
44 '''
45 Save a drawn TGraph in a PDF file.
46 '''
47 graph = root_file.Get(graph_name)
48 assert isinstance(graph, ROOT.TGraph) == 1
49 graph.SetMarkerStyle(ROOT.EMarkerStyle.kFullDotSmall)
50 graph.SetMarkerColor(ROOT.EColor.kAzure + 10)
51 graph.GetXaxis().SetTitle(f'Exp. {exp} -- Run number')
52 graph.GetYaxis().SetTitle('Plane efficiency')
53 graph.SetMinimum(0.)
54 graph.SetMaximum(1.)
55 graph.Draw('AP')
56 ROOT.gPad.SetGridy()
57 canvas.SaveAs(f'efficiency_exp{exp}_chunk{chunk}_{graph_name}.pdf')
58
59
60def run_validation(job_path, input_data_path, requested_iov, expert_config):
61 '''
62 Run the validation.
63 Nota bene:
64 - job_path will be replaced with path/to/calibration_results
65 - input_data_path will be replaced with path/to/data_path used for calibration, e.g. /group/belle2/dataprod/Data/PromptSkim/
66 '''
67
68 # Avoid looking for the release globaltag
69 basf2.conditions.override_globaltags()
70
71 # Grab the expert configurations.
72 expert_config = json.loads(expert_config)
73 chunk_size = expert_config['chunk_size']
74
75 # Ignore the ROOT command line options.
76 ROOT.PyConfig.IgnoreCommandLineOptions = True # noqa
77 # Run ROOT in batch mode.
78 ROOT.gROOT.SetBatch(True)
79 # Set the Belle II style.
80 ROOT.gROOT.SetStyle("BELLE2")
81 # And unset the stat box.
82 ROOT.gStyle.SetOptStat(0)
83
84 # Path to the database.txt file.
85 database_file = f'{job_path}/KLMStripEfficiency/outputdb/database.txt'
86
87 # Dictionary with the definition of the validation plots.
88 bklm = KLMElementNumbers.c_BKLM
89 eklm = KLMElementNumbers.c_EKLM
90 first_rpc = BKLMElementNumbers.c_FirstRPCLayer
91 graph_dictionary = {'barrel_rpcs': f'subdetector=={bklm} && layer>={first_rpc}',
92 'barrel_scintillators': f'subdetector=={bklm} && layer<{first_rpc}',
93 'endcap_scintillators': f'subdetector=={eklm}'}
94
95 # Check the list of runs from the file database.txt.
96 exp_run_dict = {}
97 previous_exp = -666
98 with open(database_file) as f:
99 for line in f:
100 fields = line.split(' ')
101 if (fields[0] == 'dbstore/KLMStripEfficiency'):
102 iov = fields[2].split(',')
103 exp = int(iov[0])
104 run = int(iov[1])
105 if (exp != previous_exp):
106 exp_run_dict[exp] = [run]
107 previous_exp = exp
108 else:
109 exp_run_dict[exp].append(run)
110
111 # Tweak the IoV range if the first run is 0.
112 # This is needed for display purposes.
113 for exp, run_list in exp_run_dict.items():
114 run_list.sort()
115 if len(run_list) > 1:
116 if run_list[0] == 0 and run_list[1] > 5:
117 run_list[0] = run_list[1] - 5
118
119 # Run the KLMCalibrationChecker class.
120 for exp, run_list in exp_run_dict.items():
121 for run in run_list:
122 checker = KLMCalibrationChecker()
123 checker.setExperimentRun(exp, run)
124 checker.setTestingPayload(database_file)
125 basf2.B2INFO(f'Creating strip efficiency results tree for experiment {exp}, run {run}.')
126 checker.setStripEfficiencyResultsFile(f'strip_efficiency_exp{exp}_run{run}.root')
127 checker.checkStripEfficiency()
128
129 # Run the validation.
130 for exp, run_list in exp_run_dict.items():
131 # For each experiment, merge the files in chunks of some runs.
132 chunks = math.ceil(len(run_list) / chunk_size)
133 for chunk in range(chunks):
134 file_name = f'strip_efficiency_exp{exp}_chunk{chunk}.root'
135 run_files = [
136 f'strip_efficiency_exp{exp}_run{run}.root' for run in run_list[chunk * chunk_size:(chunk + 1) * chunk_size]]
137 subprocess.run(['hadd', '-f', file_name] + run_files, check=True)
138 input_file = ROOT.TFile(f'{file_name}')
139 output_file = ROOT.TFile(f'histograms_{file_name}', 'recreate')
140 output_file.cd()
141 tree = input_file.Get('efficiency')
142 assert isinstance(tree, ROOT.TTree) == 1
143 canvas = ROOT.TCanvas(f'canvas_exp{exp}_chunk{chunk}', 'canvas', 800, 500)
144 canvas.cd()
145 for graph_name, graph_cut in graph_dictionary.items():
146 tree.Draw('efficiency:run', graph_cut)
147 save_graph_to_root(graph_name)
148 save_graph_to_pdf(canvas, output_file, graph_name, exp, chunk)
149 input_file.Close()
150 output_file.Close()
151 # Let's delete the files for single IoVs.
152 for run_file in run_files:
153 try:
154 os.remove(run_file)
155 except OSError as e:
156 basf2.B2ERROR(f'The file {run_file} can not be removed: {e.strerror}')
157
158
159if __name__ == "__main__":
160 run_validation(*sys.argv[1:])