Belle II Software development
caf_cdcdedx_electron.py
1
8
9"""
10Airflow script for automatic CDC dEdx calibration. It is the electron based
11calibration, where at present only RunGain, injection time, Cosine, WireGain and 1D are implemented.
12The remaining two 2D will be implemented in the near future.
13"""
14
15import ROOT
16from ROOT import gSystem
17from ROOT.Belle2 import CDCDedxRunGainAlgorithm, CDCDedxCosineAlgorithm, CDCDedxWireGainAlgorithm
18from ROOT.Belle2 import CDCDedxCosEdgeAlgorithm, CDCDedxBadWireAlgorithm, CDCDedxInjectTimeAlgorithm
19from ROOT.Belle2 import CDCDedx1DCellAlgorithm, CDCDedxValidationAlgorithm
20
21from caf.framework import Calibration
22from caf.strategies import SequentialRunByRun, SequentialBoundaries
23from prompt import CalibrationSettings, INPUT_DATA_FILTERS
24import reconstruction as recon
25from random import seed
26import basf2
27
28gSystem.Load('libreconstruction.so')
29ROOT.gROOT.SetBatch(True)
30
31settings = CalibrationSettings(
32 name="CDC dedx",
33 expert_username="renu92garg",
34 subsystem="cdc",
35 description=__doc__,
36 input_data_formats=["cdst"],
37 input_data_names=["bhabha_all_calib"],
38 expert_config={
39 "payload_boundaries": [],
40 "calib_datamode": False,
41 "maxevt_rg": 75000,
42 "maxevt_cc": 18e6,
43 "maxevt_wg": 18e6,
44 "adjustment": 1.00798,
45 "calib_mode": "full", # manual or predefined: full or quick
46 "calibration_procedure": {"rungain0": 0, "rungain1": 0, "rungain2": 0}
47 },
48 input_data_filters={
49 "bhabha_all_calib": [
50 INPUT_DATA_FILTERS['Run Type']['physics'],
51 INPUT_DATA_FILTERS['Data Tag']['bhabha_all_calib'],
52 INPUT_DATA_FILTERS['Data Quality Tag']['Good Or Recoverable'],
53 INPUT_DATA_FILTERS['Magnet']['On'],
54 INPUT_DATA_FILTERS['Beam Energy']['4S'],
55 INPUT_DATA_FILTERS['Beam Energy']['Continuum'],
56 INPUT_DATA_FILTERS['Beam Energy']['Scan']]},
57 depends_on=[],
58 produced_payloads=["CDCDedxInjectionTime", "CDCDedxCosineEdge", "CDCDedxBadWires",
59 "CDCDedxWireGain", "CDCDedx1DCell", "CDCDedxCosineCor", "CDCDedxRunGain"])
60
61
62def get_calibrations(input_data, **kwargs):
63 """ REQUIRED FUNCTION used by b2caf-prompt-run tool
64 This function return a list of Calibration
65 objects we assign to the CAF process
66 """
67
68 import basf2
69 file_to_iov_physics = input_data["bhabha_all_calib"]
70
71 expert_config = kwargs.get("expert_config")
72 calib_mode = expert_config["calib_mode"]
73
74 # extracting parameters
75 fulldataMode = expert_config["calib_datamode"]
76 adjustment = expert_config["adjustment"]
77
78 if fulldataMode:
79 input_files_rungain = list(file_to_iov_physics.keys())
80 input_files_coscorr = list(file_to_iov_physics.keys())
81 input_files_wiregain = list(file_to_iov_physics.keys())
82 else:
83 seed(271492)
84
85 maxevt_rg = expert_config["maxevt_rg"]
86 maxevt_cc = expert_config["maxevt_cc"]
87 maxevt_wg = expert_config["maxevt_wg"]
88
89 from prompt.utils import filter_by_max_events_per_run, filter_by_select_max_events_from_files
90
91 # collection for rungains
92 max_files_for_maxevents = maxevt_rg # allevents to accp bhabha event ratio = 0.60
93 reduced_file_to_iov_rungain = filter_by_max_events_per_run(file_to_iov_physics, max_files_for_maxevents, True)
94 input_files_rungain = list(reduced_file_to_iov_rungain.keys())
95 basf2.B2INFO(f"Total number of files used for rungains = {len(input_files_rungain)}")
96
97 # collection for cosinecorr
98 input_files_coscorr = filter_by_select_max_events_from_files(list(file_to_iov_physics.keys()), maxevt_cc)
99 basf2.B2INFO(f"Total number of files used for cosine = {len(input_files_coscorr)}")
100 if not input_files_coscorr:
101 raise ValueError(
102 f"Cosine: all requested ({maxevt_cc}) events not found")
103
104 # collection for wiregain
105 if maxevt_wg == maxevt_cc:
106 input_files_wiregain = input_files_coscorr
107 else:
108 input_files_wiregain = filter_by_select_max_events_from_files(list(file_to_iov_physics.keys()), maxevt_wg)
109
110 basf2.B2INFO(f"Total number of files used for wiregains = {len(input_files_wiregain)}")
111 if not input_files_wiregain:
112 raise ValueError(
113 f"WireGain: all requested ({maxevt_wg}) events not found")
114
115 requested_iov = kwargs.get("requested_iov", None)
116 from caf.utils import ExpRun, IoV
117 output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
118
119 payload_boundaries = [ExpRun(output_iov.exp_low, output_iov.run_low)]
120 payload_boundaries.extend([ExpRun(*boundary) for boundary in expert_config["payload_boundaries"]])
121 basf2.B2INFO(f"Expert set payload boundaries are: {expert_config['payload_boundaries']}")
122
123 collector_granularity = 'all'
124 if expert_config["payload_boundaries"] is not None:
125 basf2.B2INFO('Found payload_boundaries: set collector granularity to run')
126 collector_granularity = 'run'
127
128 if calib_mode == "full":
129 calibration_procedure = {
130 "rungain0": 0, # Run Gain trail (No Payload saving and take of effect of previous rungains)
131 "timegain0": 0, # Injection time gain Pre (No payload saving)
132 "timegain1": 0, # Injection time gain
133 "rungain1": 0, # Run Gain Pre (No Payload saving)
134 "coscorr0": 0, # Cosine Corr Gain Pre (No Payload saving)
135 "cosedge0": 0, # Cosine edge Corr Gain
136 "badwire0": 0, # Bad wire
137 "wiregain0": 0, # WireGain Gain
138 "onedcell0": 0, # OneD cell correction
139 "coscorr1": 0, # Cosine Corr Gain
140 "rungain2": 0, # Final Run Gain to take Wire and Cosine correction in effect
141 "validation0": 0 # get data for validation
142 }
143 elif calib_mode == "quick":
144 calibration_procedure = {
145 "rungain0": 0,
146 "timegain0": 0,
147 "rungain1": 0,
148 "coscorr0": 0,
149 "cosedge0": 0,
150 "badwire0": 0,
151 "wiregain0": 0,
152 "rungain2": 0,
153 "validation0": 0
154 }
155 elif calib_mode == "manual":
156 calibration_procedure = expert_config["calibration_procedure"]
157 else:
158 basf2.B2FATAL(f"Calibration mode is not defined {calib_mode}, should be full, quick, or manual")
159
160 calib_keys = list(calibration_procedure)
161 cals = [None]*len(calib_keys)
162 basf2.B2INFO(f"Run calibration mode = {calib_mode}:")
163
164 for i in range(len(cals)):
165 max_iter = calibration_procedure[calib_keys[i]]
166 alg = None
167 data_files = [input_files_rungain, input_files_coscorr, input_files_wiregain]
168 cal_name = ''.join([i for i in calib_keys[i] if not i.isdigit()])
169 if cal_name == "rungain":
170 alg = [rungain_algo(calib_keys[i], adjustment)]
171 elif cal_name == "coscorr":
172 alg = [cos_algo()]
173 elif cal_name == "cosedge":
174 alg = [cosedge_algo()]
175 elif cal_name == "timegain":
176 alg = [injection_time_algo()]
177 elif cal_name == "badwire":
178 alg = [badwire_algo()]
179 elif cal_name == "wiregain":
180 alg = [wiregain_algo()]
181 elif cal_name == "onedcell":
182 alg = [onedcell_algo()]
183 elif cal_name == "validation":
184 alg = [validation_algo()]
185 else:
186 basf2.B2FATAL(f"The calibration is not defined, check spelling: calib {i}: {calib_keys[i]}")
187
188 basf2.B2INFO(f"calibration for {calib_keys[i]} with number of iteration={max_iter}")
189
190 cals[i] = CDCDedxCalibration(name=calib_keys[i],
191 algorithms=alg,
192 input_file_dict=data_files,
193 max_iterations=max_iter,
194 collector_granularity=collector_granularity,
195 dependencies=[cals[i-1]] if i > 0 else None
196 )
197 if payload_boundaries:
198 basf2.B2INFO("Found payload_boundaries: calibration strategies set to SequentialBoundaries.")
199 if cal_name == "rungain" or cal_name == "timegain":
200 cals[i].strategies = SequentialRunByRun
201 for algorithm in cals[i].algorithms:
202 algorithm.params = {"iov_coverage": output_iov}
203 if calib_keys[i] == "rungain0" or calib_keys[i] == "rungain1" or calib_keys[i] == "timegain0":
204 cals[i].save_payloads = False
205 else:
206 cals[i].strategies = SequentialBoundaries
207 for algorithm in cals[i].algorithms:
208 algorithm.params = {"iov_coverage": output_iov, "payload_boundaries": payload_boundaries}
209 if calib_keys[i] == "coscorr0":
210 cals[i].save_payloads = False
211
212 else:
213 for algorithm in cals[i].algorithms:
214 algorithm.params = {"apply_iov": output_iov}
215
216 return cals
217
218
219# Precollector path
220def pre_collector(name='rg'):
221 """
222 Define pre collection.
223 Parameters:
224 name : name of the calibration
225 rungain rungain0 by Default.
226 Returns:
227 path : path for pre collection
228 """
229
230 reco_path = basf2.create_path()
231 recon.prepare_cdst_analysis(path=reco_path)
232 if (name == "validation"):
233 basf2.B2INFO("no trigger skim")
234 elif (name == "timegain" or name == "onedcell"):
235 trg_bhabhaskim = reco_path.add_module("TriggerSkim", triggerLines=["software_trigger_cut&skim&accept_radee"])
236 trg_bhabhaskim.if_value("==0", basf2.Path(), basf2.AfterConditionPath.END)
237 ps_bhabhaskim = reco_path.add_module("Prescale", prescale=0.80)
238 ps_bhabhaskim.if_value("==0", basf2.Path(), basf2.AfterConditionPath.END)
239
240 elif (name == "cosedge"):
241 trg_bhabhaskim = reco_path.add_module(
242 "TriggerSkim",
243 triggerLines=[
244 "software_trigger_cut&skim&accept_bhabha",
245 "software_trigger_cut&filter&ee_flat_90_180",
246 "software_trigger_cut&filter&ee_flat_0_19"])
247 trg_bhabhaskim.if_value("==0", basf2.Path(), basf2.AfterConditionPath.END)
248 else:
249 trg_bhabhaskim = reco_path.add_module("TriggerSkim", triggerLines=["software_trigger_cut&skim&accept_bhabha"])
250 trg_bhabhaskim.if_value("==0", basf2.Path(), basf2.AfterConditionPath.END)
251
252 reco_path.add_module(
253 'CDCDedxCorrection',
254 relativeCorrections=False,
255 scaleCor=True,
256 runGain=True,
257 timeGain=True,
258 cosineCor=True,
259 wireGain=True,
260 twoDCell=True,
261 oneDCell=True)
262 return reco_path
263
264# Collector setup
265
266
267def collector(granularity='all', name=''):
268 """
269 Create a cdcdedx calibration collector
270 Parameters:
271 name : name of calibration
272 granularity : granularity : all or run
273 Returns:
274 collector : collector module
275 """
276
277 from basf2 import register_module
278 if name == "validation":
279 col = register_module('ElectronValCollector', cleanupCuts=True)
280 return col
281
282 else:
283 col = register_module('CDCDedxElectronCollector', cleanupCuts=True)
284 if name == "timegain":
285 CollParam = {'isRun': True, 'isInjTime': True, 'granularity': 'run'}
286
287 elif name == "coscorr" or name == "cosedge":
288 CollParam = {'isCharge': True, 'isCosth': True, 'granularity': granularity}
289
290 elif name == "badwire":
291 isHit = True
292 CollParam = {'isWire': True, 'isDedxhit': isHit, 'isADCcorr': not isHit, 'granularity': granularity}
293
294 elif name == "wiregain":
295 CollParam = {'isWire': True, 'isDedxhit': True, 'granularity': granularity}
296
297 elif name == "onedcell":
298 CollParam = {
299 'isPt': True,
300 'isCosth': True,
301 'isLayer': True,
302 'isDedxhit': True,
303 'isEntaRS': True,
304 'granularity': granularity}
305
306 else:
307 CollParam = {'isRun': True, 'granularity': 'run'}
308
309 col.param(CollParam)
310 return col
311
312# Rungain Algorithm setup
313
314
315def rungain_algo(name, adjustment):
316 """
317 Create a rungain calibration algorithm.
318 Returns:
319 algo : rungain algorithm
320 """
321 algo = CDCDedxRunGainAlgorithm()
322 algo.setMonitoringPlots(True)
323 if name == "rungain2":
324 algo.setAdjustment(adjustment)
325 return algo
326
327# Injection Algorithm setup
328
329
330def injection_time_algo():
331 """
332 Create a injection time calibration algorithm.
333 Returns:
334 algo : injection time algorithm
335 """
336 algo = CDCDedxInjectTimeAlgorithm()
337 algo.setMonitoringPlots(True)
338 return algo
339
340# Cosine Algorithm setup
341
342
343def cos_algo():
344 """
345 Create a cosine calibration algorithm.
346 Returns:
347 algo : cosine algorithm
348 """
349 algo = CDCDedxCosineAlgorithm()
350 algo.setMonitoringPlots(True)
351 return algo
352
353# CosineEdge Algorithm setup
354
355
356def cosedge_algo():
357 """
358 Create a cosine edge calibration algorithm.
359 Returns:
360 algo : cosine edge algorithm
361 """
362 algo = CDCDedxCosEdgeAlgorithm()
363 algo.setMonitoringPlots(True)
364 return algo
365
366# Badwire Algorithm setup
367
368
369def badwire_algo():
370 """
371 Create a badwire calibration algorithm.
372 Returns:
373 algo : badwire algorithm
374 """
375 algo = CDCDedxBadWireAlgorithm()
376 # threshold (mean and rms) pars for dedx
377 algo.setHighFracThres(0.2)
378 algo.setMeanThres(0.4)
379 algo.setRMSThres(0.4)
380 algo.setHistPars(150, 0, 5)
381 algo.setMonitoringPlots(True)
382 return algo
383
384# WireGain Algorithm setup
385
386
387def wiregain_algo():
388 """
389 Create a wire gain calibration algorithm.
390 Returns:
391 algo : wiregain algorithm
392 """
393 algo = CDCDedxWireGainAlgorithm()
394 algo.enableExtraPlots(True)
395 return algo
396
397
398def onedcell_algo():
399 """
400 Create oned cell calibration algorithm.
401 Returns:
402 algo : oned cell correction algorithm
403 """
404 algo = CDCDedx1DCellAlgorithm()
405 algo.enableExtraPlots(True)
406 algo.setMergePayload(True)
407 return algo
408
409
410def validation_algo():
411 """
412 Create validation algorithm
413 Returns:
414 algo : validation algorithm
415 """
416 algo = CDCDedxValidationAlgorithm()
417 return algo
418
419
420class CDCDedxCalibration(Calibration):
421 '''
422 CDCDedxCalibration is a specialized calibration for cdcdedx.
423 '''
424
425 def __init__(self,
426 name,
427 algorithms,
428 input_file_dict,
429 max_iterations=5,
430 dependencies=None,
431 collector_granularity='All'):
432 '''
433 parameters:
434 name: name of calibration
435 algorithms: algorithm of calibration
436 input_file_dict: input files list
437 max_iterations: maximum number of iterations
438 dependencies: depends on the previous calibration
439 collector_granularity: granularity : all or run
440 '''
441 super().__init__(name=name,
442 algorithms=algorithms
443 )
444
445 from caf.framework import Collection
446 cal_name = ''.join([i for i in name if not i.isdigit()])
447 if cal_name == "badwire" or cal_name == "wiregain":
448 collection = Collection(collector=collector(granularity=collector_granularity, name=cal_name),
449 input_files=input_file_dict[2],
450 pre_collector_path=pre_collector(cal_name)
451 )
452 elif cal_name == "coscorr" or cal_name == "cosedge" or cal_name == "onedcell":
453 collection = Collection(collector=collector(granularity=collector_granularity, name=cal_name),
454 input_files=input_file_dict[1],
455 pre_collector_path=pre_collector(cal_name)
456 )
457 else:
458 collection = Collection(collector=collector(granularity=collector_granularity, name=cal_name),
459 input_files=input_file_dict[0],
460 pre_collector_path=pre_collector(cal_name)
461 )
462 self.add_collection(name=cal_name, collection=collection)
463
464
465 self.max_iterations = max_iterations
466
467 if dependencies is not None:
468 for dep in dependencies:
469 self.depends_on(dep)
__init__(self, name, algorithms, input_file_dict, max_iterations=5, dependencies=None, collector_granularity='All')