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