Belle II Software development
caf_cdc.py
1# disable doxygen check for this file
2# @cond
3
4
11
12"""Full CDC tracking calibration."""
13from prompt import CalibrationSettings, INPUT_DATA_FILTERS
14from prompt.utils import events_in_basf2_file, ExpRun
15import basf2
16from ROOT import Belle2
17from random import choice
18from caf.framework import Calibration
19from caf import strategies
20
21
22
23settings = CalibrationSettings(name="CDC Tracking",
24 expert_username="guanyu",
25 subsystem="cdc",
26 description=__doc__,
27 input_data_formats=["raw"],
28 input_data_names=["mumu_tight_or_highm_calib", "hadron_calib", "cosmic_calib"],
29 input_data_filters={"mumu_tight_or_highm_calib":
30 [INPUT_DATA_FILTERS["Data Tag"]["mumu_tight_or_highm_calib"],
31 INPUT_DATA_FILTERS["Data Quality Tag"]["Good"],
32 INPUT_DATA_FILTERS["Magnet"]["On"]],
33 "hadron_calib":
34 [INPUT_DATA_FILTERS["Data Tag"]["hadron_calib"],
35 INPUT_DATA_FILTERS["Data Quality Tag"]["Good"],
36 INPUT_DATA_FILTERS["Magnet"]["On"]],
37 "cosmic_calib":
38 [INPUT_DATA_FILTERS["Data Tag"]["cosmic_calib"],
39 INPUT_DATA_FILTERS["Data Quality Tag"]["Good"],
40 INPUT_DATA_FILTERS["Magnet"]["On"]]},
41 depends_on=[],
42 expert_config={
43 "min_events_per_file": 500,
44 "max_events_per_file": 20000,
45 "max_events_per_file_hadron_for_tz_tw": 5000,
46 "max_events_per_file_hadron_for_xt_sr": 12000,
47 "min_events_for_tz_tw_calibration": 500000,
48 "max_events_for_tz_tw_calibration": 1000000,
49 "min_events_for_xt_sr_calibration": 1000000, # 1M
50 "max_events_for_xt_sr_calibration": 10000000, # 10M
51 "fractions_for_each_type": [0.5, 1, 0.5], # [mumu, hadron, cosmic]
52 "max_job_for_each_type": [400, 700, 400],
53 "calib_mode": "quick", # manual or predefined: quick, full
54 "calibration_procedure": {"tz0": 1, "xt0": 0, "sr_tz0": 0, "tz2": 0},
55 "payload_boundaries": [],
56 "backend_args": {"request_memory": "4 GB"}},
57 produced_payloads=["CDCTimeZeros", "CDCTimeWalk", "CDCXtRelations", "CDCSpaceResols"])
58
59
60def select_files(all_input_files, min_events, max_events, max_processed_events_per_file, max_job=800, min_events_per_file=500):
61 basf2.B2INFO(f"Minimum number of events: {min_events}")
62 basf2.B2INFO(f"Maximum number of events: {max_events}")
63 basf2.B2INFO(f"Conditions: ({min_events_per_file} < #Event/file < {max_processed_events_per_file}) and max_job = {max_job}")
64 # Let's iterate, taking a sample of files from the total (no repeats or replacement) until we get enough events
65 total_events = 0
66 chosen_files = []
67 njob = 0
68 while total_events < max_events and njob < max_job:
69 # If the set is empty we must have used all available files. Here we break and continue. But you may want to
70 # raise an Error...
71 if not all_input_files:
72 break
73 # Randomly select a file
74 new_file_choice = choice(all_input_files)
75 # Remove it from the list so it can't be chosen again
76 all_input_files.remove(new_file_choice)
77 # Find the number of events in the file
78 total_events_in_file = events_in_basf2_file(new_file_choice)
79 if not total_events_in_file or total_events_in_file < min_events_per_file:
80 # Uh Oh! Zero event file, skip it
81 continue
82 events_contributed = 0
83 if total_events_in_file < max_processed_events_per_file:
84 # The file contains less than the max amount we have set (entrySequences)
85 events_contributed = total_events_in_file
86 else:
87 events_contributed = max_processed_events_per_file
88 chosen_files.append(new_file_choice)
89 total_events += events_contributed
90 njob += 1
91
92 basf2.B2INFO(f"Total chosen files = {len(chosen_files)}")
93 basf2.B2INFO(f"Total events in chosen files = {total_events}")
94 if total_events < min_events:
95 basf2.B2FATAL(
96 f"There is not enough required events with setup max_processed_events_per_file = {max_processed_events_per_file}")
97 return chosen_files
98
99
100
103
104
105def get_calibrations(input_data, **kwargs):
106 import basf2
107 # from prompt.utils import filter_by_max_files_per_run
108 # read expert_config values
109 expert_config = kwargs.get("expert_config")
110 calib_mode = expert_config["calib_mode"]
111 # max_files_per_run = expert_config["max_files_per_run"]
112 min_events_per_file = expert_config["min_events_per_file"]
113 max_events_per_file = expert_config["max_events_per_file"]
114
115 min_events_for_tz_tw = expert_config["min_events_for_tz_tw_calibration"] # for t0, tw calib.
116 max_events_for_tz_tw = expert_config["max_events_for_tz_tw_calibration"] # for t0, tw calib.
117 max_events_per_file_hadron_for_tz_tw = expert_config["max_events_per_file_hadron_for_tz_tw"]
118
119 min_events_for_xt_sr = expert_config["min_events_for_xt_sr_calibration"] # for xt, sr calib.
120 max_events_for_xt_sr = expert_config["max_events_for_xt_sr_calibration"] # for xt, sr calib.
121 max_events_per_file_hadron_for_xt_sr = expert_config["max_events_per_file_hadron_for_xt_sr"]
122
123 Max_events_per_file_for_tz_tw = [max_events_per_file, max_events_per_file_hadron_for_tz_tw, max_events_per_file]
124 Max_events_per_file_for_xt_sr = [max_events_per_file, max_events_per_file_hadron_for_xt_sr, max_events_per_file]
125
126 fraction_of_event_for_types = expert_config["fractions_for_each_type"]
127 max_jobs = expert_config["max_job_for_each_type"]
128 basf2.B2INFO(f"Number of job for each type are limited at [di-muon, hadron, cosmic]: {max_jobs}")
129 basf2.B2INFO(f"Fraction for [di-muon, hadron, cosmic]: {fraction_of_event_for_types}")
130 if len(fraction_of_event_for_types) != 3:
131 basf2.B2FATAL("fraction of event must be an array with the size of 3, with order [mumu, hadron, cosmic]")
132
133 payload_boundaries = []
134 payload_boundaries.extend([ExpRun(*boundary) for boundary in expert_config["payload_boundaries"]])
135 basf2.B2INFO(f"Payload boundaries from expert_config: {payload_boundaries}")
136
137 files_for_xt_sr_dict = {}
138 files_for_tz_tw_dict = {}
139
140 # select data file for t0 and tw calibration
141 if fraction_of_event_for_types[1] > 0:
142 basf2.B2INFO("*********************** Select Hadron data for calibration ****************")
143 min_hadron_events_for_tz_tw = fraction_of_event_for_types[1] * min_events_for_tz_tw
144 max_hadron_events_for_tz_tw = fraction_of_event_for_types[1] * max_events_for_tz_tw
145 min_hadron_events_for_xt_sr = fraction_of_event_for_types[1] * min_events_for_xt_sr
146 max_hadron_events_for_xt_sr = fraction_of_event_for_types[1] * max_events_for_xt_sr
147
148 file_to_iov_hadron = input_data["hadron_calib"]
149 # select data file for tw, t0
150 basf2.B2INFO("----> For T0 and Time walk correction")
151 chosen_files_hadron_for_tz_tw = select_files(list(file_to_iov_hadron.keys()),
152 min_hadron_events_for_tz_tw,
153 max_hadron_events_for_tz_tw,
154 max_events_per_file_hadron_for_tz_tw,
155 max_jobs[1],
156 min_events_per_file)
157 # select data file for xt, sigma
158 basf2.B2INFO("----> For XT, space resolution calib")
159 chosen_files_hadron_for_xt_sr = select_files(list(file_to_iov_hadron.keys()),
160 min_hadron_events_for_xt_sr,
161 max_hadron_events_for_xt_sr,
162 max_events_per_file_hadron_for_xt_sr,
163 max_jobs[1],
164 min_events_per_file)
165
166 files_for_xt_sr_dict["hadron_calib"] = chosen_files_hadron_for_xt_sr
167 files_for_tz_tw_dict["hadron_calib"] = chosen_files_hadron_for_tz_tw
168
169 if fraction_of_event_for_types[0] > 0:
170 basf2.B2INFO("***********************Select di-muon data for calibration ***************")
171 min_mumu_events_for_xt_sr = fraction_of_event_for_types[0] * min_events_for_xt_sr
172 max_mumu_events_for_xt_sr = fraction_of_event_for_types[0] * max_events_for_xt_sr
173 min_mumu_events_for_tz_tw = fraction_of_event_for_types[0] * min_events_for_tz_tw
174 max_mumu_events_for_tz_tw = fraction_of_event_for_types[0] * max_events_for_tz_tw
175 file_to_iov_mumu = input_data["mumu_tight_or_highm_calib"]
176 basf2.B2INFO("----> For T0 and Time walk correction")
177 chosen_files_mumu_for_tz_tw = select_files(list(file_to_iov_mumu.keys()),
178 min_mumu_events_for_tz_tw,
179 max_mumu_events_for_tz_tw,
180 max_events_per_file,
181 max_jobs[0],
182 min_events_per_file)
183
184 # select data file for xt, sigma calibration
185 basf2.B2INFO("----> For XT, space resolution calib")
186
187 chosen_files_mumu_for_xt_sr = select_files(list(file_to_iov_mumu.keys()), # input_files_mumu[:],
188 min_mumu_events_for_xt_sr,
189 max_mumu_events_for_xt_sr,
190 max_events_per_file,
191 max_jobs[0],
192 min_events_per_file)
193
194 files_for_xt_sr_dict["mumu_tight_or_highm_calib"] = chosen_files_mumu_for_xt_sr
195 files_for_tz_tw_dict["mumu_tight_or_highm_calib"] = chosen_files_mumu_for_tz_tw
196
197 ''' For cosmic data '''
198 if fraction_of_event_for_types[2] > 0:
199 basf2.B2INFO("********************* Select cosmic data for calibration *******************")
200 min_cosmic_events_for_tz_tw = fraction_of_event_for_types[2] * min_events_for_tz_tw
201 max_cosmic_events_for_tz_tw = fraction_of_event_for_types[2] * max_events_for_tz_tw
202 min_cosmic_events_for_xt_sr = fraction_of_event_for_types[2] * min_events_for_xt_sr
203 max_cosmic_events_for_xt_sr = fraction_of_event_for_types[2] * max_events_for_xt_sr
204
205 file_to_iov_cosmic = input_data["cosmic_calib"]
206
207 # Select cosmic data for tw and t0 calibration
208 basf2.B2INFO("---->For T0 and Time walk correction")
209 chosen_files_cosmic_for_tz_tw = select_files(list(file_to_iov_cosmic.keys()),
210 min_cosmic_events_for_tz_tw,
211 max_cosmic_events_for_tz_tw,
212 max_events_per_file,
213 max_jobs[2],
214 min_events_per_file)
215
216 # select cosmics data for xt and sigma calibration
217 basf2.B2INFO("----> For T0 and Time walk correction")
218 chosen_files_cosmic_for_xt_sr = select_files(list(file_to_iov_cosmic.keys()),
219 min_cosmic_events_for_xt_sr,
220 max_cosmic_events_for_xt_sr,
221 max_events_per_file,
222 max_jobs[2],
223 min_events_per_file)
224 files_for_xt_sr_dict["cosmic_calib"] = chosen_files_cosmic_for_xt_sr
225 files_for_tz_tw_dict["cosmic_calib"] = chosen_files_cosmic_for_tz_tw
226
227 basf2.B2INFO("Complete input data selection.")
228
229 # Get the overall IoV we want to cover, including the end values
230 requested_iov = kwargs.get("requested_iov", None)
231
232 from caf.utils import IoV
233 # The actual IoV we want for any prompt request is open-ended
234 output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
235
236 # for SingleIOV stratrgy, it's better to set the granularity to 'all' so that the collector jobs will run faster
237 collector_granularity = 'all'
238 if payload_boundaries:
239 basf2.B2INFO('Found payload_boundaries: set collector granularity to run')
240 collector_granularity = 'run'
241 if calib_mode == "full":
242 calibration_procedure = {
243 "tz0": 4,
244 "xt0": 0,
245 "sr_tz0": 0,
246 "xt1": 0,
247 "sr_tz1": 0,
248 "tw0": 1,
249 "tz1": 1,
250 "xt2": 0,
251 "sr_tz2": 0,
252 "tz2": 0
253 }
254 elif calib_mode == "quick":
255 calibration_procedure = {
256 "tz0": 3,
257 "xt0": 0,
258 "sr_tz0": 0,
259 "tw0": 1,
260 "tz1": 1,
261 "xt1": 0,
262 "sr_tz1": 0,
263 "tz2": 0
264 }
265 elif calib_mode == "manual":
266 calibration_procedure = expert_config["calibration_procedure"]
267 else:
268 basf2.B2FATAL(f"Calibration mode is not defined {calib_mode}, should be quick, full, or manual")
269 # t0
270 calib_keys = list(calibration_procedure)
271 cals = [None]*len(calib_keys)
272 basf2.B2INFO(f"Run calibration mode = {calib_mode}:")
273 print(calib_keys)
274 for i in range(len(cals)):
275 max_iter = calibration_procedure[calib_keys[i]]
276 alg = None
277 data_files = None
278 cal_name = ''.join([i for i in calib_keys[i] if not i.isdigit()])
279 if cal_name == "tz":
280 alg = [tz_algo()]
281 elif cal_name == "tw":
282 alg = [tw_algo()]
283 elif cal_name == "xt":
284 alg = [xt_algo()]
285 elif cal_name == "sr_tz":
286 alg = [sr_algo(), tz_algo()]
287 else:
288 basf2.B2FATAL(f"The calibration is not defined, check spelling: calib {i}: {calib_keys[i]}")
289
290 if cal_name == "xt" or cal_name == "sr_tz":
291 max_event = Max_events_per_file_for_xt_sr
292 data_files = files_for_xt_sr_dict
293 else:
294 max_event = Max_events_per_file_for_tz_tw
295 data_files = files_for_tz_tw_dict
296 basf2.B2INFO(f"calibration for {calib_keys[i]} with number of iteration={max_iter}")
297 cals[i] = CDCCalibration(name=calib_keys[i],
298 algorithms=alg,
299 input_file_dict=data_files,
300 max_iterations=max_iter,
301 max_events=max_event,
302 use_badWires=True if calib_keys[i] == "tz" else False,
303 collector_granularity=collector_granularity,
304 backend_args=expert_config["backend_args"],
305 dependencies=[cals[i-1]] if i > 0 else None
306 )
307 if payload_boundaries:
308 basf2.B2INFO("Found payload_boundaries: calibration strategies set to SequentialBoundaries.")
309 cals[0].strategies = strategies.SequentialBoundaries
310 for alg in cals:
311 for algorithm in alg.algorithms:
312 algorithm.params = {"iov_coverage": output_iov, "payload_boundaries": payload_boundaries}
313
314 else:
315 for alg in cals:
316 for algorithm in alg.algorithms:
317 algorithm.params = {"apply_iov": output_iov}
318
319 return cals
320
321
322
323def pre_collector(max_events=None, is_cosmic=False, use_badWires=False):
324 """
325 Define pre collection (reconstruction in our purpose).
326 Probably, we need only CDC and ECL data.
327 Parameters:
328 max_events [int] : number of events to be processed.
329 All events by Default.
330 Returns:
331 path : path for pre collection
332 """
333 from basf2 import create_path, register_module
334 from softwaretrigger.constants import HLT_INPUT_OBJECTS
335 reco_path = create_path()
336 if max_events is None:
337 root_input = register_module(
338 'RootInput',
339 branchNames=HLT_INPUT_OBJECTS
340 )
341 else:
342 root_input = register_module(
343 'RootInput',
344 branchNames=HLT_INPUT_OBJECTS,
345 entrySequences=[
346 f'0:{max_events}'])
347 reco_path.add_module(root_input)
348
349 gearbox = register_module('Gearbox')
350 reco_path.add_module(gearbox)
351 reco_path.add_module('Geometry', useDB=True)
352 Components = ['CDC', 'ECL']
353
354 from rawdata import add_unpackers
355 add_unpackers(reco_path, components=Components)
356
357 if is_cosmic:
358 from reconstruction import add_cosmics_reconstruction
359 # Add cdc tracking reconstruction modules
360 add_cosmics_reconstruction(path=reco_path,
361 components=Components,
362 pruneTracks=False,
363 merge_tracks=False,
364 posttracking=False)
365 else:
366 from reconstruction import default_event_abort, add_prefilter_pretracking_reconstruction
367 from tracking import add_prefilter_tracking_reconstruction
368
369 # Do not even attempt at reconstructing events w/ abnormally large occupancy.
370 doom = reco_path.add_module("EventsOfDoomBuster")
371 default_event_abort(doom, ">=1", Belle2.EventMetaData.c_ReconstructionAbort)
372 reco_path.add_module('StatisticsSummary').set_name('Sum_EventsofDoomBuster')
373
374 # Add modules that have to be run BEFORE track reconstruction
375 add_prefilter_pretracking_reconstruction(reco_path, components=Components)
376
377 # Add tracking reconstruction modules
378 add_prefilter_tracking_reconstruction(path=reco_path,
379 components=Components,
380 trackFitHypotheses=[211],
381 prune_temporary_tracks=False,
382 fit_tracks=True,
383 append_full_grid_cdc_eventt0=True,
384 skip_full_grid_cdc_eventt0_if_svd_time_present=False)
385 reco_path.add_module('StatisticsSummary').set_name('Sum_Tracking')
386
387 reco_path.add_module('Progress')
388
389 for module in reco_path.modules():
390 if module.name() == "TFCDC_WireHitPreparer":
391 module.param({"useBadWires": use_badWires})
392
393 return reco_path
394
395
396def collector(is_cosmic=False, granularity='all'):
397 """
398 Create a cdc calibration collector
399 Parameters:
400 bField [bool] : True if B field is on, else False
401 isCosmic [bool] : True if cosmic events,
402 else (collision) False.
403 Returns:
404 collector : collector module
405 """
406 from basf2 import register_module
407 col = register_module('CDCCalibrationCollector',
408 granularity=granularity,
409 calExpectedDriftTime=True,
410 eventT0Extraction=True,
411 isCosmic=is_cosmic
412 )
413 return col
414
415
416def tz_algo(max_rmsDt=0.25, max_badChannel=50):
417 """
418 Create a T0 calibration algorithm.
419 Returns:
420 algo : T0 algorithm
421 """
422 from ROOT import Belle2
424 algo.storeHisto(True)
425 algo.setMaxMeanDt(0.2)
426 algo.setMaxRMSDt(max_rmsDt)
427 algo.setMaxBadChannel(max_badChannel)
428 algo.setMinimumNDF(25)
429
430 return algo
431
432
433def tw_algo():
434 """
435 Create a time walk calibration algorithm.
436 Returns:
437 algo : TW algorithm
438 """
439 from ROOT import Belle2
441 algo.setStoreHisto(True)
442 algo.setMode(1)
443 return algo
444
445
446def xt_algo():
447 """
448 Create a XT calibration algorithm.
449 Parameters:
450 prefix : prefixed name for algorithm,
451 which should be consistent with one of collector..
452 Returns:
453 algo : XT algorithm
454 """
455 from ROOT import Belle2
457 algo.setStoreHisto(True)
458 algo.setLRSeparate(True)
459 algo.setThreshold(0.1)
460 return algo
461
462
463def sr_algo():
464 """
465 Create a Spacial resolution calibration algorithm.
466 Parameters:
467 prefix : prefixed name for algorithm,
468 which should be consistent with one of collector..
469 Returns:
470 algo : Spacial algorithm
471 """
472 from ROOT import Belle2
474 algo.setStoreHisto(True)
475 algo.setThreshold(0.1)
476 return algo
477
478
479class CDCCalibration(Calibration):
480 '''
481 CDCCalibration is a specialized calibration class for cdc.
482 Since collector is same in all elements, no need to specify it.
483 '''
484
485 def __init__(self,
486 name,
487 algorithms,
488 input_file_dict,
489 max_iterations=5,
490 dependencies=None,
491 max_events=[20000, 10000, 20000],
492 use_badWires=False,
493 collector_granularity='All',
494 backend_args=None):
495 for algo in algorithms:
496 algo.setHistFileName(name)
497
498 super().__init__(name=name,
499 algorithms=algorithms
500 )
501
502 from caf.framework import Collection
503
504 for skim_type, file_list in input_file_dict.items():
505 if skim_type == "cosmic_calib":
506 collection = Collection(collector=collector(is_cosmic=True,
507 granularity=collector_granularity),
508 input_files=file_list,
509 pre_collector_path=pre_collector(max_events=max_events[2],
510 is_cosmic=True,
511 use_badWires=use_badWires),
512 backend_args=backend_args)
513 elif skim_type == "hadron_calib":
514 collection = Collection(collector=collector(granularity=collector_granularity),
515 input_files=file_list,
516 pre_collector_path=pre_collector(max_events=max_events[1],
517 use_badWires=use_badWires),
518 backend_args=backend_args)
519 else:
520 collection = Collection(collector=collector(granularity=collector_granularity),
521 input_files=file_list,
522 pre_collector_path=pre_collector(max_events=max_events[0],
523 use_badWires=use_badWires),
524 backend_args=backend_args)
525
526 self.add_collection(name=skim_type, collection=collection)
527
528 self.max_iterations = max_iterations
529
530 if dependencies is not None:
531 for dep in dependencies:
532 self.depends_on(dep)
533
534# @endcond
Class to perform xt calibration for drift chamber.