Belle II Software  release-05-01-25
caf_pxd_gain.py
1 # -*- coding: utf-8 -*-
2 
3 """
4 airflow script for PXD gain calibration.
5 Author: qingyuan.liu@desy.de
6 """
7 
8 import basf2
9 from pxd.calibration import gain_calibration
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
17 
18 
19 settings = CalibrationSettings(name="PXD gain calibration",
20  expert_username="qyliu",
21  description=__doc__,
22  input_data_formats=["cdst"],
23  input_data_names=["physics"],
24  # input_data_filters={
25  # "physics": [
26  # "bhabha_all_calib",
27  # "4S", "Continuum", "Scan",
28  # "physics", "Good"]
29  # },
30  expert_config={
31  "debug": False,
32  "total_jobs": 1000,
33  "gain_method": "analytic",
34  "min_files_per_chunk": 10,
35  "min_events_per_file": 1000, # avoid empty files
36  "max_events_per_run": 4000000,
37  "max_files_per_run": 20, # only valid when max_events/run = 0
38  "payload_boundaries": []
39  },
40  depends_on=[beamspot_calibration])
41 
42 
43 def get_calibrations(input_data, **kwargs):
44  """
45  Parameters:
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
49 
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.
53 
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)
56 
57  Returns:
58  list(caf.framework.Calibration): All of the calibration objects we want to assign to the CAF process
59  """
60 
61  # Set up config options
62  requested_iov = kwargs.get("requested_iov", None)
63  output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
64  # expert config
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", {})
74 
75  # print all config
76  basf2.B2INFO(f"Requested iov: {requested_iov} ")
77  basf2.B2INFO(f"Expert config: {expert_config} ")
78  # basf2.B2INFO(f"Expert sets payload boundaries are: {expert_config['payload_boundaries']} ")
79 
80  # Read input_data
81  file_to_iov_physics = input_data["physics"]
82 
83  # Reduce data and create calibration instances for different data categories
84  cal_list = []
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)
92  else:
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)
96 
97  # input_files_physics = list(reduced_file_to_iov_physics.keys())
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])
100 
101  # boundaries setting for run chunks (At certain runs, gain was tuned)
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"]])
104  # We don't need run 0 for the first exp as it's handled by output_iov
105  payload_boundaries.extend([ExpRun(exp, 0) for exp in sorted(exp_set)[1:]])
106  basf2.B2INFO(f"Final Boundaries: {payload_boundaries}")
107 
108  # run chunk creation
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])]
113 
114  # Create calibrations from chunks
115  input_file_to_iov = reduced_file_to_iov_physics
116  iCal = 0
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: # All the chunk iovs are earlier than the requested
121  continue
122  else:
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]))
125  # Check the minimum number of files in the physics/beam run chunk
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!")
129  continue
130  # From the second chunk within the requested range, we have the iov defined by the first run
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"
134  if (not debug):
135  cal = gain_calibration(
136  cal_name=cal_name,
137  gain_method=gain_method,
138  # boundaries=vector_from_runs(payload_boundaries),
139  input_files=input_files,
140  **cal_kwargs)
141  for alg in cal.algorithms:
142  alg.params["iov_coverage"] = specific_iov
143  cal_list.append(cal)
144  else:
145  basf2.B2INFO(f"Dry run on Calibration(name={cal_name})")
146  iCal += 1
147 
148  # The number of calibrations depends on the 'chunking' above. We would like to make sure that the total number of
149  # batch jobs submitted is approximately constant and reasonable, no matter how many files and chunks are used.
150  # So we define 1000 total jobs and split this between the calibrations depending on the fraction of total input
151  # files in the calibrations.
152 
153  # total_jobs = expert_config["total_jobs"]
154  total_input_files = len(reduced_file_to_iov_physics)
155 
156  for cal in cal_list:
157  fraction_of_input_files = len(cal.input_files) / total_input_files
158  # Assign the max collector jobs to be roughly the same fraction of total jobs
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")
161 
162  return cal_list
pxd.calibration
Definition: calibration.py:1
prompt.utils
Definition: utils.py:1