Belle II Software development
caf_klm_alignment.py
1
8
9"""
10Calibration of KLM alignment. It provides calibration constants for the BKLMAlignment
11and EKLMAlignment database objects.
12"""
13
14import collections
15
16import basf2
17from caf.utils import ExpRun, IoV
18from prompt import CalibrationSettings, INPUT_DATA_FILTERS
19from prompt.utils import events_in_basf2_file
20
21
28
29from prompt.calibrations.caf_vxdcdc_alignment import settings as caf_vxdcdc_alignment
30from prompt.calibrations.caf_cdc import settings as caf_cdc
31from prompt.calibrations.caf_klm_channel_status import settings as caf_klm_channel_status
32
33
41settings = CalibrationSettings(
42 name="KLM alignmnent",
43 expert_username="nbrenny",
44 subsystem="klm",
45 description=__doc__,
46 input_data_formats=["raw"],
47 input_data_names=[
48 "raw_physics",
49 "raw_cosmic"],
50 input_data_filters={
51 'raw_physics': [
52 INPUT_DATA_FILTERS['Run Type']['physics'],
53 INPUT_DATA_FILTERS['Data Tag']['mumu_tight_or_highm_calib'],
54 INPUT_DATA_FILTERS['Data Quality Tag']['Good Or Recoverable']],
55 'raw_cosmic': [
56 INPUT_DATA_FILTERS['Run Type']['physics'],
57 INPUT_DATA_FILTERS['Data Tag']['cosmic_calib'],
58 INPUT_DATA_FILTERS['Data Quality Tag']['Good Or Recoverable']]},
59 depends_on=[
60 caf_vxdcdc_alignment,
61 caf_cdc,
62 caf_klm_channel_status],
63 expert_config={
64 "required_events": 5000000,
65 "required_events_experiment": 500000,
66 "events_per_file": 1000,
67 "millepede_entries": 1000000,
68 "millepede_entries_exp7": 500000},
69 produced_payloads=["EKLMAlignment_CORRECTIONS",
70 "BKLMAlignment_CORRECTIONS",
71 "EKLMAlignment",
72 "BKLMAlignment",
73 "EKLMAlignment_ERRORS",
74 "BKLMAlignment_ERRORS",
75 "EKLMSegmentAlignment_ERRORS",
76 "EKLMSegmentAlignment_CORRECTIONS"])
77
78
79def select_input_files(file_to_iov_physics, file_to_iov_cosmic,
80 reduced_file_to_iov_physics, reduced_file_to_iov_cosmic,
81 required_events, required_events_experiment,
82 events_per_file):
83 """
84 Parameters:
85 files_to_iov_physics (dict): Dictionary {run : IOV} for physics data.
86 files_to_iov_cosmic (dict): Dictionary {run : IOV} for cosmic data.
87 reduced_file_to_iov_physics (dict): Selected physics data.
88 reduced_file_to_iov_cosmic (dict): Selected cosmic data.
89 required_events (int): Required total number of events.
90 required_events_experiment (int): Required number of events
91 per experiment.
92 events_per_file (int): Number of events per file. If non-positive, then
93 the number is not limited.
94 """
95 experiment_numbers = set()
96 run_to_files_physics = collections.defaultdict(list)
97 for input_file, file_iov in file_to_iov_physics.items():
98 run = ExpRun(exp=file_iov.exp_low, run=file_iov.run_low)
99 run_to_files_physics[run].append(input_file)
100 experiment_numbers.add(file_iov.exp_low)
101 run_to_files_cosmic = collections.defaultdict(list)
102 for input_file, file_iov in file_to_iov_cosmic.items():
103 run = ExpRun(exp=file_iov.exp_low, run=file_iov.run_low)
104 run_to_files_cosmic[run].append(input_file)
105 experiment_numbers.add(file_iov.exp_low)
106 max_files_per_run = 0
107 for files in run_to_files_physics.values():
108 files_per_run = len(files)
109 if files_per_run > max_files_per_run:
110 max_files_per_run = files_per_run
111 for files in run_to_files_cosmic.values():
112 files_per_run = len(files)
113 if files_per_run > max_files_per_run:
114 max_files_per_run = files_per_run
115 files_per_run = 0
116 collected_events = 0
117 collected_events_experiment = collections.defaultdict(int)
118 select_events_experiment = collections.defaultdict(bool)
119 for exp in experiment_numbers:
120 collected_events_experiment[exp] = 0
121 select_events_experiment[exp] = True
122 while files_per_run < max_files_per_run:
123 for run, files in run_to_files_physics.items():
124 if not select_events_experiment[run.exp]:
125 continue
126 if files_per_run >= len(files):
127 continue
128 input_file = files[files_per_run]
129 events = events_in_basf2_file(input_file)
130 # Reject files without events.
131 if events == 0:
132 continue
133 if events_per_file > 0 and events > events_per_file:
134 events = events_per_file
135 collected_events = collected_events + events
136 collected_events_experiment[run.exp] = \
137 collected_events_experiment[run.exp] + events
138 reduced_file_to_iov_physics[input_file] = IoV(*run, *run)
139 basf2.B2INFO(f'File {input_file} with {events} events is selected.')
140 for run, files in run_to_files_cosmic.items():
141 if not select_events_experiment[run.exp]:
142 continue
143 if files_per_run >= len(files):
144 continue
145 input_file = files[files_per_run]
146 events = events_in_basf2_file(input_file)
147 # Reject files without events.
148 if events == 0:
149 continue
150 if events_per_file > 0 and events > events_per_file:
151 events = events_per_file
152 collected_events = collected_events + events
153 collected_events_experiment[run.exp] = \
154 collected_events_experiment[run.exp] + events
155 reduced_file_to_iov_cosmic[input_file] = IoV(*run, *run)
156 basf2.B2INFO(f'File {input_file} with {events} events is selected.')
157 files_per_run = files_per_run + 1
158 if collected_events >= required_events:
159 all_experiments_selected = True
160 for exp in experiment_numbers:
161 if collected_events_experiment[exp] >= \
162 required_events_experiment:
163 select_events_experiment[exp] = False
164 else:
165 all_experiments_selected = False
166 if all_experiments_selected:
167 break
168 basf2.B2INFO(f'The total number of collected events is {collected_events}.')
169
170
179
180
181def get_calibrations(input_data, **kwargs):
182 """
183 Parameters:
184 input_data (dict): Should contain every name from the 'input_data_names' variable as a key.
185 Each value is a dictionary with {"/path/to/file_e1_r5.root": IoV(1,5,1,5), ...}. Useful for
186 assigning to calibration.files_to_iov
187
188 **kwargs: Configuration options to be sent in. Since this may change we use kwargs as a way to help prevent
189 backwards compatibility problems. But you could use the correct arguments in b2caf-prompt-run for this
190 release explicitly if you want to.
191
192 Currently only kwargs["requested_iov"] is used. This is the output IoV range that your payloads should
193 correspond to. Generally your highest ExpRun payload should be open ended e.g. IoV(3,4,-1,-1)
194
195 Returns:
196 list(caf.framework.Calibration): All of the calibration objects we want to assign to the CAF process
197 """
198 # Set up config options
199
200 # Expert configuration.
201 expert_config = kwargs.get("expert_config")
202
203 # In this script we want to use one sources of input data.
204 # Get the input files from the input_data variable
205 file_to_iov_physics = input_data["raw_physics"]
206 file_to_iov_cosmic = input_data["raw_cosmic"]
207
208 # Select input files.
209 reduced_file_to_iov_physics = collections.OrderedDict()
210 reduced_file_to_iov_cosmic = collections.OrderedDict()
211 select_input_files(file_to_iov_physics, file_to_iov_cosmic,
212 reduced_file_to_iov_physics, reduced_file_to_iov_cosmic,
213 expert_config["required_events"],
214 expert_config["required_events_experiment"],
215 expert_config["events_per_file"])
216
217 input_files_physics = sorted(list(reduced_file_to_iov_physics.keys()))
218 basf2.B2INFO(f"Total number of 'physics' files actually used as input = {len(input_files_physics)}")
219
220 input_files_cosmic = sorted(list(reduced_file_to_iov_cosmic.keys()))
221 basf2.B2INFO(f"Total number of 'cosmic' files actually used as input = {len(input_files_cosmic)}")
222
223 if not input_files_physics and not input_files_cosmic:
224 raise Exception("No valid input files found!")
225
226 # Get the overall IoV we our process should cover. Includes the end values that we may want to ignore since our output
227 # IoV should be open ended. We could also use this as part of the input data selection in some way.
228 requested_iov = kwargs["requested_iov"]
229
230 from caf.utils import IoV
231 # The actual value our output IoV payload should have. Notice that we've set it open ended.
232 output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
233
234
236
237 from ROOT import Belle2
238 from ROOT.Belle2 import KLMChannelIndex, KLMElementNumbers
239 from alignment import MillepedeCalibration
240
241 # Create the algorithm.
242 millepede = MillepedeCalibration(['BKLMAlignment', 'EKLMAlignment', 'EKLMSegmentAlignment'])
243
244 # Fix module parameters.
245 index = KLMChannelIndex(KLMChannelIndex.c_IndexLevelLayer)
246 index2 = KLMChannelIndex(KLMChannelIndex.c_IndexLevelLayer)
247 while (index != index2.end()):
248 module = index.getKLMModuleNumber()
249 if (index.getSubdetector() == KLMElementNumbers.c_BKLM):
250 for ipar in [1, 2, 3, 4, 5, 6]:
251 # Free parameters are idU, dV, dGamma.
252 if ipar in [1, 2, 6]:
253 continue
254 millepede.fixGlobalParam(Belle2.BKLMAlignment.getGlobalUniqueID(),
255 module, ipar)
256 else:
257 for ipar in [1, 2, 6]:
258 # No parameters are fixed; if necessary, uncomment the following:
259 # millepede.fixGlobalParam(Belle2.EKLMAlignment.getGlobalUniqueID(),
260 # module, ipar)
261 continue
262 index.increment()
263
264 # Fix EKLM segment parameters.
265 index.setIndexLevel(KLMChannelIndex.c_IndexLevelStrip)
266 index2.setIndexLevel(KLMChannelIndex.c_IndexLevelStrip)
267 index = index2.beginEKLM()
268 index.useEKLMSegments()
269 while (index != index2.endEKLM()):
270 segment = index.getEKLMSegmentNumber()
271 for ipar in [2, 6]:
272 millepede.fixGlobalParam(
274 segment, ipar)
275 index.increment()
276
277 cal_klm = millepede.create('KLMAlignment', [])
278 # Number of entries is set in algorithm strategy now.
279 millepede.algo.ignoreUndeterminedParams(True)
280 millepede.algo.invertSign()
281
282
284
285 from caf.framework import Collection
286
287
289
290 from klm_calibration_utils import get_alignment_pre_collector_path_physics, get_alignment_pre_collector_path_cosmic
291
292 entries = ''
293 if expert_config["events_per_file"] > 0:
294 last_entry = expert_config["events_per_file"]
295 entries = f'0:{last_entry}'
296
297 if input_files_physics:
298 coll_physics = get_collector("raw_physics")
299 rec_path_physics = get_alignment_pre_collector_path_physics(entry_sequence=entries)
300 # remove cdcdedxpid module
301 tmp = basf2.create_path()
302 for m in rec_path_physics.modules():
303 if m.name() not in ["CDCDedxPID", "TOPBunchFinder", "VXDDedxPID"]:
304 tmp.add_module(m)
305 elif m.name() == "CDCDedxPID":
306 basf2.B2INFO('removed CDCDedxPID')
307 elif m.name() == "TOPBunchFinder":
308 basf2.B2INFO('removed TOPBunchFinder')
309 rec_path_physics = tmp
310
311 collection_physics = Collection(collector=coll_physics,
312 input_files=input_files_physics,
313 pre_collector_path=rec_path_physics)
314
315 cal_klm.add_collection(name="physics", collection=collection_physics)
316
317 if input_files_cosmic:
318 coll_cosmic = get_collector("raw_cosmic")
319 rec_path_cosmic = get_alignment_pre_collector_path_cosmic(entry_sequence=entries)
320 # remove cdcdedxpid module
321 tmp = basf2.create_path()
322 for m in rec_path_cosmic.modules():
323 if m.name() not in ["CDCDedxPID", "TOPBunchFinder", "VXDDedxPID"]:
324 tmp.add_module(m)
325 elif m.name() == "CDCDedxPID":
326 basf2.B2INFO('removed CDCDedxPID')
327 elif m.name() == "TOPBunchFinder":
328 basf2.B2INFO('removed TOPBunchFinder')
329 rec_path_cosmic = tmp
330
331 collection_cosmic = Collection(collector=coll_cosmic,
332 input_files=input_files_cosmic,
333 pre_collector_path=rec_path_cosmic)
334
335 cal_klm.add_collection(name="cosmic", collection=collection_cosmic)
336
337
339
340 from klm_alignment import KLMAlignment
341
342 cal_klm.algorithms = [millepede.algo]
343
344 # Bugfix for Condor:
345 from alignment.prompt_utils import fix_mille_paths_for_algo
346 for algorithm in cal_klm.algorithms:
347 fix_mille_paths_for_algo(algorithm)
348
349 for algorithm in cal_klm.algorithms:
350 algorithm.strategy = KLMAlignment
351 algorithm.params = {
352 "has_experiment_settings": True,
353 "iov_coverage": output_iov,
354 "millepede_entries": expert_config["millepede_entries"],
355 "millepede_entries_exp7": expert_config["millepede_entries_exp7"]
356 }
357
358 # You must return all calibrations you want to run in the prompt process, even if it's only one
359 return [cal_klm]
360
361
362
363
364def get_collector(input_data_name):
365 """
366 Return the correct MillepedeCollector module setup for each data type.
367 Placed here so it can be different for prompt compared to standard.
368 """
369
370 if input_data_name == "raw_physics":
371 return basf2.register_module(
372 'MillepedeCollector',
373 components=['BKLMAlignment', 'EKLMAlignment',
374 'EKLMSegmentAlignment'],
375 useGblTree=True,
376 minPValue=1e-5)
377 elif input_data_name == "raw_cosmic":
378 return basf2.register_module(
379 'MillepedeCollector',
380 components=['BKLMAlignment', 'EKLMAlignment',
381 'EKLMSegmentAlignment'],
382 useGblTree=True,
383 minPValue=1e-5)
384
385 raise Exception("Unknown input data name used when setting up collector")
static unsigned short getGlobalUniqueID()
Get global unique identifier.
static unsigned short getGlobalUniqueID()
Get global unique identifier.