4 airflow script for PXD gain calibration.
5 Author: qingyuan.liu@desy.de
10 from prompt.utils import filter_by_max_files_per_run, filter_by_max_events_per_run
11 from prompt
import CalibrationSettings
12 from caf.utils
import vector_from_runs, ExpRun, IoV
13 from itertools
import groupby
14 from itertools
import chain
15 from math
import ceil, inf
16 from prompt.calibrations.caf_beamspot
import settings
as beamspot_calibration
19 settings = CalibrationSettings(name=
"PXD gain calibration",
20 expert_username=
"qyliu",
22 input_data_formats=[
"cdst"],
23 input_data_names=[
"physics"],
33 "gain_method":
"analytic",
34 "min_files_per_chunk": 10,
35 "min_events_per_file": 1000,
36 "max_events_per_run": 4000000,
37 "max_files_per_run": 20,
38 "payload_boundaries": []
40 depends_on=[beamspot_calibration])
43 def get_calibrations(input_data, **kwargs):
46 input_data (dict): Should contain every name from the 'input_data_names' variable as a key.
47 Each value is a dictionary with {"/path/to/file_e1_r5.root": IoV(1,5,1,5), ...}. Useful for
48 assigning to calibration.files_to_iov
50 **kwargs: Configuration options to be sent in. Since this may change we use kwargs as a way to help prevent
51 backwards compatibility problems. But you could use the correct arguments in b2caf-prompt-run for this
52 release explicitly if you want to.
54 Currently only kwargs["output_iov"] is used. This is the output IoV range that your payloads should
55 correspond to. Generally your highest ExpRun payload should be open ended e.g. IoV(3,4,-1,-1)
58 list(caf.framework.Calibration): All of the calibration objects we want to assign to the CAF process
62 requested_iov = kwargs.get(
"requested_iov",
None)
63 output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
65 expert_config = kwargs.get(
"expert_config")
66 gain_method = expert_config[
"gain_method"]
67 debug = expert_config[
"debug"]
68 total_jobs = expert_config[
"total_jobs"]
69 max_events_per_run = expert_config[
"max_events_per_run"]
70 max_files_per_run = expert_config[
"max_files_per_run"]
71 min_files_per_chunk = expert_config[
"min_files_per_chunk"]
72 min_events_per_file = expert_config[
"min_events_per_file"]
73 cal_kwargs = expert_config.get(
"kwargs", {})
76 basf2.B2INFO(f
"Requested iov: {requested_iov} ")
77 basf2.B2INFO(f
"Expert config: {expert_config} ")
81 file_to_iov_physics = input_data[
"physics"]
85 if max_events_per_run < 0:
86 basf2.B2INFO(
"No file reduction applied.")
87 reduced_file_to_iov_physics = file_to_iov_physics
88 elif max_events_per_run == 0:
89 basf2.B2INFO(f
"Reducing to a maximum of {max_files_per_run} files per run.")
90 reduced_file_to_iov_physics = filter_by_max_files_per_run(file_to_iov_physics,
91 max_files_per_run, min_events_per_file)
93 basf2.B2INFO(f
"Reducing to a maximum of {max_events_per_run} events per run.")
94 reduced_file_to_iov_physics = filter_by_max_events_per_run(file_to_iov_physics,
95 max_events_per_run, random_select=
True)
98 input_iov_set_physics = set(reduced_file_to_iov_physics.values())
99 exp_set = set([iov.exp_low
for iov
in input_iov_set_physics])
102 payload_boundaries = [ExpRun(output_iov.exp_low, output_iov.run_low)]
103 payload_boundaries.extend([ExpRun(*boundary)
for boundary
in expert_config[
"payload_boundaries"]])
105 payload_boundaries.extend([ExpRun(exp, 0)
for exp
in sorted(exp_set)[1:]])
106 basf2.B2INFO(f
"Final Boundaries: {payload_boundaries}")
109 chunks_head = payload_boundaries
110 chunks_tail = payload_boundaries[1:] + [ExpRun(inf, inf)]
111 iov_chunks = [list(g)
for k, g
in groupby(sorted(input_iov_set_physics),
112 lambda x: [i
for i, j
in zip(chunks_head, chunks_tail)
if i <= x < j])]
115 input_file_to_iov = reduced_file_to_iov_physics
117 for ichunk, chunk
in enumerate(iov_chunks):
118 first_iov = IoV(chunk[0].exp_low, chunk[0].run_low, -1, -1)
119 last_iov = IoV(chunk[-1].exp_low, chunk[-1].run_low, -1, -1)
120 if last_iov < output_iov:
123 input_files = list(chain.from_iterable([list(g)
for k, g
in groupby(
124 input_file_to_iov,
lambda x: input_file_to_iov[x]
in chunk)
if k]))
126 if len(input_files) < min_files_per_chunk:
127 basf2.B2WARNING(f
"No enough file in sub run chunk [{chunk[0]},{chunk[-1]}]: {len(input_files)},\
128 but {min_files_per_chunk} required!")
131 specific_iov = first_iov
if iCal > 0
else output_iov
132 basf2.B2INFO(f
"Total number of files actually used as input = {len(input_files)} for the output {specific_iov}")
133 cal_name = f
"{ichunk+1}_PXDAnalyticGainCalibration"
135 cal = gain_calibration(
137 gain_method=gain_method,
139 input_files=input_files,
141 for alg
in cal.algorithms:
142 alg.params[
"iov_coverage"] = specific_iov
145 basf2.B2INFO(f
"Dry run on Calibration(name={cal_name})")
154 total_input_files = len(reduced_file_to_iov_physics)
157 fraction_of_input_files = len(cal.input_files) / total_input_files
159 cal.max_collector_jobs = ceil(fraction_of_input_files * total_jobs)
160 basf2.B2INFO(f
"{cal.name} will submit a maximum of {cal.max_collector_jobs} batch jobs")