8 from caf.utils
import ExpRun, IoV
9 from prompt
import CalibrationSettings
20 from prompt.calibrations.vxdcdc_alignment
import settings
as vxdcdc_alignment
30 settings = CalibrationSettings(name=
"KLM alignmnent",
31 expert_username=
"oskin",
33 input_data_formats=[
"raw"],
34 input_data_names=[
"raw_physics",
"raw_cosmic"],
35 depends_on=[vxdcdc_alignment],
37 "required_events": 5000000,
38 "required_events_experiment": 500000,
39 "events_per_file": 1000,
40 "millepede_entries": 1000000,
41 "millepede_entries_exp7": 500000
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,
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
61 events_per_file (int): Number of events per file. If non-positive, then
62 the number is not limited.
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)
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
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]:
95 if files_per_run >= len(files):
97 input_file = files[files_per_run]
98 events = events_in_basf2_file(input_file)
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]:
112 if files_per_run >= len(files):
114 input_file = files[files_per_run]
115 events = events_in_basf2_file(input_file)
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
134 all_experiments_selected =
False
135 if all_experiments_selected:
137 basf2.B2INFO(f
'The total number of collected events is {collected_events}.')
150 def get_calibrations(input_data, **kwargs):
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
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.
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)
165 list(caf.framework.Calibration): All of the calibration objects we want to assign to the CAF process
170 expert_config = kwargs.get(
"expert_config")
174 file_to_iov_physics = input_data[
"raw_physics"]
175 file_to_iov_cosmic = input_data[
"raw_cosmic"]
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"])
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)}")
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)}")
192 if not input_files_physics
and not input_files_cosmic:
193 raise Exception(
"No valid input files found!")
197 requested_iov = kwargs[
"requested_iov"]
199 from caf.utils
import IoV
201 output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
206 from ROOT
import Belle2
207 from ROOT.Belle2
import KLMChannelIndex, KLMElementNumbers
208 from alignment
import MillepedeCalibration
211 millepede = MillepedeCalibration([
'BKLMAlignment',
'EKLMAlignment',
'EKLMSegmentAlignment'])
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]:
221 if ipar
in [1, 2, 6]:
226 for ipar
in [1, 2, 6]:
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()
241 millepede.fixGlobalParam(
246 cal_klm = millepede.create(
'KLMAlignment', [])
248 millepede.algo.ignoreUndeterminedParams(
True)
249 millepede.algo.invertSign()
254 from caf.framework
import Collection
259 from klm_calibration_utils
import get_alignment_pre_collector_path_physics, get_alignment_pre_collector_path_cosmic
262 if expert_config[
"events_per_file"] > 0:
263 last_entry = expert_config[
"events_per_file"]
264 entries = f
'0:{last_entry}'
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)
270 collection_physics =
Collection(collector=coll_physics,
271 input_files=input_files_physics,
272 pre_collector_path=rec_path_physics)
274 cal_klm.add_collection(name=
"physics", collection=collection_physics)
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)
280 collection_cosmic =
Collection(collector=coll_cosmic,
281 input_files=input_files_cosmic,
282 pre_collector_path=rec_path_cosmic)
284 cal_klm.add_collection(name=
"cosmic", collection=collection_cosmic)
289 from klm_alignment
import KLMAlignment
291 cal_klm.algorithms = [millepede.algo]
295 for algorithm
in cal_klm.algorithms:
296 fix_mille_paths_for_algo(algorithm)
298 for algorithm
in cal_klm.algorithms:
299 algorithm.strategy = KLMAlignment
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"]
313 def get_collector(input_data_name):
315 Return the correct MillepedeCollector module setup for each data type.
316 Placed here so it can be different for prompt compared to standard.
319 if input_data_name ==
"raw_physics":
320 return basf2.register_module(
321 'MillepedeCollector',
322 components=[
'BKLMAlignment',
'EKLMAlignment',
323 'EKLMSegmentAlignment'],
326 elif input_data_name ==
"raw_cosmic":
327 return basf2.register_module(
328 'MillepedeCollector',
329 components=[
'BKLMAlignment',
'EKLMAlignment',
330 'EKLMSegmentAlignment'],
334 raise Exception(
"Unknown input data name used when setting up collector")