Belle II Software development
caf_vxdcdc_alignment.py
1
8"""
9
10Full Simultaneous Global and Local VXD and CDC alignment with Millepede II
11
12The input collections are:
13- cosmics (hlt skim) - mandatorry
14- hadron - for "low" momentum tracks from IP
15- mumu - mumu_2trk or mumu_tight - for high momentum tracks from IP
16- offip - tracks from outside IP (beam background, beam-gas)
17
18"""
19
20import basf2
21import ROOT
22
23from prompt import CalibrationSettings, INPUT_DATA_FILTERS
24from prompt.calibrations.caf_cdc import settings as cdc_calibration
25from prompt.calibrations.caf_svd_time import settings as svd_time_calibration
26
27from prompt.utils import events_in_basf2_file
28from caf.utils import IoV
29from caf import strategies
30
31from softwaretrigger.constants import HLT_INPUT_OBJECTS
32import rawdata as raw
33import reconstruction as reco
34import modularAnalysis as ana
35import vertex as vtx
36
37from random import choice
38from random import seed
39
40import millepede_calibration as mpc
44
45collection_names = ["cosmic", "hadron", "mumu", "offip"]
46
47default_config = {
48 "only_prompt": False,
49 "cosmic.max_processed_events_per_file": 4000,
50 "hadron.max_processed_events_per_file": 1000,
51 "mumu.max_processed_events_per_file": 5000,
52 "offip.max_processed_events_per_file": 2000,
53 "beamspot.min_pxd_hits": 0,
54 "stage1.method": "decomposition",
55 "stage1.z_offset": True
56}
57
58quality_flags = [INPUT_DATA_FILTERS["Run Type"]["physics"],
59 INPUT_DATA_FILTERS["Data Quality Tag"]["Good Or Recoverable"],
60 INPUT_DATA_FILTERS["Magnet"]["On"]]
61
62
63settings = CalibrationSettings(name="Full VXD and CDC Alignment",
64 expert_username="tadeas.bilka",
65 subsystem="alignment",
66 description=__doc__,
67 input_data_formats=["raw"],
68 input_data_names=collection_names,
69 input_data_filters={
70 "cosmic": [INPUT_DATA_FILTERS["Data Tag"]["cosmic_calib"]] + quality_flags,
71 "mumu": [INPUT_DATA_FILTERS["Data Tag"]["mumu_tight_or_highm_calib"]] + quality_flags,
72 "hadron": [INPUT_DATA_FILTERS["Data Tag"]["hadron_calib"]] + quality_flags,
73 "offip": [INPUT_DATA_FILTERS["Data Tag"]["offip_calib"]] + quality_flags
74 },
75
76 expert_config=default_config,
77 depends_on=[cdc_calibration, svd_time_calibration],
78 produced_payloads=["VXDAlignment", "CDCAlignment"])
79
80
81def select_files(all_input_files, min_events, max_processed_events_per_file):
82 """
83 Select files randomly
84
85 Parameters
86 ----------
87 all_input_files : list(str)
88 List of all input file names
89 min_events : int
90 Minimum number of events to select from files
91 max_processed_events_per_file : int
92 Maximum number of events to consider per file
93 """
94 all_input_files = all_input_files[:]
95 # Let's iterate, taking a sample of files from the total (no repeats or replacement) until we get enough events
96 total_events = 0
97 chosen_files = []
98 while total_events < min_events:
99 # If the set is empty we must have used all available files. Here we break and continue. But you may want to
100 # raise an Error...
101 if not all_input_files:
102 break
103 # Randomly select a file
104 new_file_choice = choice(all_input_files)
105 # Remove it from the list so it can't be chosen again
106 all_input_files.remove(new_file_choice)
107 # Find the number of events in the file
108 total_events_in_file = events_in_basf2_file(new_file_choice)
109 if not total_events_in_file:
110 # Uh Oh! Zero event file, skip it
111 continue
112
113 events_contributed = min(total_events_in_file, max_processed_events_per_file)
114
115 chosen_files.append(new_file_choice)
116 total_events += events_contributed
117
118 basf2.B2INFO(f"Total chosen files = {len(chosen_files)}")
119 basf2.B2INFO(f"Total events in chosen files = {total_events}")
120 if total_events < min_events:
121 basf2.B2WARNING(
122 f"There weren't enough files events selected when max_processed_events_per_file={max_processed_events_per_file}")
123 return chosen_files
124
125
126def create_std_path():
127 """
128 Returns default path for collections with standard reconstruction
129 """
130 path = basf2.create_path()
131 path.add_module('Progress')
132 path.add_module('RootInput')
133 path.add_module('Gearbox')
134 path.add_module('Geometry')
135 raw.add_unpackers(path)
136 path.add_module('SetupGenfitExtrapolation')
137 reco.add_reconstruction(
138 path,
139 pruneTracks=False,
140 skipGeometryAdding=True,
141 )
142 path.add_module('DAFRecoFitter')
143 return path
144
145
146def create_cosmics_path():
147 """
148 Returns default path for cosmic collection
149 """
150 path = basf2.create_path()
151 path.add_module('Progress')
152 path.add_module('RootInput')
153 path.add_module('Gearbox')
154 path.add_module('Geometry')
155
156 raw.add_unpackers(path)
157 path.add_module('SetupGenfitExtrapolation')
158 reco.add_cosmics_reconstruction(
159 path,
160 pruneTracks=False,
161 skipGeometryAdding=True,
162 addClusterExpertModules=False,
163 merge_tracks=True
164 )
165
166 path.add_module('SetRecoTrackMomentum', automatic=True)
167 path.add_module('DAFRecoFitter', pdgCodesToUseForFitting=[13])
168
169 ana.fillParticleList(
170 'mu+:goodForVXDCDCAlignment',
171 '[z0 <= 57. or abs(d0) >= 26.5] and abs(dz) > 0.4 and nTracks == 1',
172 path=path)
173 path.add_module('SkimFilter', particleLists=['mu+:goodForVXDCDCAlignment']).if_false(basf2.create_path())
174
175 return path
176
177
178def make_mumu_collection(
179 name,
180 files=None,
181 muon_cut='p > 1.0 and abs(dz) < 2.0 and dr < 0.5 and nTracks==2',
182 dimuon_cut='9.5 < M and M < 11.',
183 prescale=1.):
184 """
185 Di-muons with vertex+beam constraint collection
186
187 Parameters
188 ----------
189 name : str
190 Collection name
191 files : list(str)
192 List of input data files
193 muon_cut : str
194 Cut string to select daughter muons
195 dimuon_cut : str
196 Cut string to apply for reconstructed di-muon decay
197 prescale : float
198 Process only 'prescale' fraction of events
199 """
200 path = basf2.create_path()
201 path.add_module('Progress')
202 path.add_module('RootInput')
203 if prescale != 1.:
204 path.add_module('Prescale', prescale=prescale).if_false(basf2.Path(), basf2.AfterConditionPath.END)
205
206 path.add_module('Gearbox')
207 path.add_module('Geometry')
208
209 raw.add_unpackers(path)
210
211 reco.add_reconstruction(path, pruneTracks=False)
212
213 path.add_module('DAFRecoFitter', pdgCodesToUseForFitting=[13])
214
215 ana.fillParticleList(f"mu+:{name}", muon_cut, path=path)
216 ana.reconstructDecay(f"Upsilon(4S):{name} -> mu+:{name} mu-:{name}", dimuon_cut, path=path)
217
218 vtx.raveFit(f"Upsilon(4S):{name}", 0.001, daughtersUpdate=True, silence_warning=True, path=path, constraint="ipprofile")
219
221 name,
222 files=files,
223 path=path,
224 primaryVertices=[f"Upsilon(4S):{name}"])
225
226
227def create_prompt(files, cfg):
228 """
229 Returns configured (original) prompt stage alignment
230
231 Parameters
232 ----------
233 files : list(str)
234 Dictionary with all input files by category (name)
235 cfg : dict
236 Expert config dictionary
237 """
238 mumu = select_files(files["mumu"], 0.2e6, cfg["mumu.max_processed_events_per_file"])
239 cosmic = select_files(files["cosmic"], 1e6, cfg["cosmic.max_processed_events_per_file"])
240 hadron = select_files(files["hadron"], 0.5e5, cfg["hadron.max_processed_events_per_file"])
241 offip = select_files(files["offip"], 0.2e6, cfg["offip.max_processed_events_per_file"])
242
243 cal = mpc.create(
244 name='VXDCDCalignment_prompt',
245 dbobjects=['VXDAlignment', 'CDCAlignment'],
246 collections=[
247 mpc.make_collection("cosmic", path=create_cosmics_path(), tracks=["RecoTracks"]),
248 mpc.make_collection("hadron", path=create_std_path(), tracks=["RecoTracks"]),
249 mpc.make_collection("mumu", path=create_std_path(), tracks=["RecoTracks"]),
250 mpc.make_collection("offip", path=create_std_path(), tracks=["RecoTracks"])
251 ],
252 tags=None,
253 files=dict(mumu=mumu, cosmic=cosmic, hadron=hadron, offip=offip),
254 timedep=None,
255 constraints=[
256 alignment.constraints.VXDHierarchyConstraints(type=2, pxd=True, svd=True),
257 alignment.constraints.CDCLayerConstraints(z_offset=True, z_scale=False, twist=True)
258 ],
259 fixed=alignment.parameters.vxd_sensors(rigid=False, surface2=False, surface3=False, surface4=False),
260 commands=[
261 "method diagonalization 3 0.1",
262 "scaleerrors 1. 1.",
263 "entries 1000"],
264 params=dict(minPValue=0.00001, externalIterations=0, granularity="run"),
265 min_entries=1000000)
266
267 cal.max_iterations = 5
268
269 return cal
270
271
272def create_beamspot(files, cfg):
273 """
274 Returns configured beamspot calibration
275
276 Parameters
277 ----------
278 files : list(str)
279 Dictionary with all input files by category (name)
280 cfg : dict
281 Expert config dictionary
282 """
283
284 mumu = select_files(files["mumu"], 10e6, cfg["mumu.max_processed_events_per_file"])
285
286
288
289 from ROOT.Belle2 import BeamSpotAlgorithm
290 from basf2 import create_path, register_module
291
292
294
295 from caf.framework import Calibration, Collection
296 from caf.strategies import SingleIOV
297
298 # module to be run prior the collector
299 path = create_path()
300
301 path.add_module('Progress')
302 path.add_module('RootInput')
303 path.add_module('Gearbox')
304 path.add_module('Geometry')
305 raw.add_unpackers(path)
306 path.add_module('SetupGenfitExtrapolation')
307 reco.add_reconstruction(path, skipGeometryAdding=True)
308
309 muSelection = '[p>1.0]'
310 muSelection += ' and abs(dz)<2.0 and abs(dr)<0.5'
311 muSelection += f' and nPXDHits >= {cfg["beamspot.min_pxd_hits"]} and nSVDHits >= 8 and nCDCHits >= 20'
312 ana.fillParticleList('mu+:BS', muSelection, path=path)
313 ana.reconstructDecay('Upsilon(4S):BS -> mu+:BS mu-:BS', '9.5<M<11.5', path=path)
314
315 collector_bs = register_module('BeamSpotCollector', Y4SPListName='Upsilon(4S):BS')
316 algorithm_bs = BeamSpotAlgorithm()
317 # algorithm_bs.setOuterLoss(cfg['outerLoss'])
318 # algorithm_bs.setInnerLoss(cfg['innerLoss'])
319
320 collection_bs = Collection(collector=collector_bs,
321 input_files=mumu,
322 pre_collector_path=path)
323
324 calibration_bs = Calibration('VXDCDCalignment_beamspot', algorithms=algorithm_bs)
325 calibration_bs.add_collection("mumu", collection_bs)
326
327 calibration_bs.strategies = SingleIOV
328
329 return calibration_bs
330
331
332def create_stage1(files, cfg):
333 """
334 Returns configured stage1 alignment (full constant alignment with wires, beamspot fixed)
335
336 Parameters
337 ----------
338 files : list(str)
339 Dictionary with all input files by category (name)
340 cfg : dict
341 Expert config dictionary
342 """
343
344 mumu = select_files(files["mumu"], 1.5e6, cfg["mumu.max_processed_events_per_file"])
345 cosmic = select_files(files["cosmic"], 0.7e6, cfg["cosmic.max_processed_events_per_file"])
346 hadron_and_offip = select_files(files["hadron"] + files["offip"], int(4.0e6 / 10.), cfg["hadron.max_processed_events_per_file"])
347
348 cal = mpc.create(
349 name='VXDCDCalignment_stage1',
350 dbobjects=['VXDAlignment', 'CDCAlignment', 'BeamSpot'],
351 collections=[
352 mpc.make_collection("cosmic", path=create_cosmics_path(), tracks=["RecoTracks"]),
353 mpc.make_collection("hadron", path=create_std_path(), tracks=["RecoTracks"]),
354 make_mumu_collection(name="mumu")
355 ],
356 tags=None,
357 files=dict(mumu=mumu, cosmic=cosmic, hadron=hadron_and_offip),
358 timedep=None,
359 constraints=[
360 alignment.constraints.VXDHierarchyConstraints(type=2, pxd=True, svd=True),
361 alignment.constraints.CDCLayerConstraints(z_offset=cfg["stage1.z_offset"], z_scale=False, twist=False),
362 alignment.constraints.CDCWireConstraints(layer_rigid=True, layer_radius=[53], cdc_radius=True, hemisphere=[55])
363 ],
364 fixed=alignment.parameters.vxd_sensors(rigid=False, surface2=False, surface3=False, surface4=False)
366 commands=[
367 f"method {cfg['stage1.method']} 6 0.001",
368 "entries 1000",
369 "threads 10 10"],
370 params=dict(minPValue=0.00001, externalIterations=0, granularity="run"),
371 min_entries=500000)
372
373 # Ignore results for BeamSpot
374 std_components = ROOT.vector('string')()
375 for component in ['VXDAlignment', 'CDCAlignment']:
376 std_components.push_back(component)
377 cal.algorithms[0].algorithm.setComponents(std_components)
378
379 cal.max_iterations = 0
380
381 return cal
382
383
384def create_stage2(files, cfg):
385 """
386 Returns configured stage2 alignment (run-dependent alignment)
387
388 Parameters
389 ----------
390 files : list(str)
391 Dictionary with all input files by category (name)
392 cfg : dict
393 Expert config dictionary
394 """
395 mumu = select_files(files["mumu"], 10e6, cfg["mumu.max_processed_events_per_file"])
396 cosmic = select_files(files["cosmic"], 2e6, cfg["cosmic.max_processed_events_per_file"])
397
398 cal = mpc.create(
399 name='VXDCDCalignment_stage2',
400 dbobjects=['VXDAlignment', 'CDCAlignment', 'BeamSpot'],
401 collections=[
402 mpc.make_collection("cosmic", path=create_cosmics_path(), tracks=["RecoTracks"]),
403 make_mumu_collection(name="mumu")],
404 tags=None,
405 files=dict(mumu=mumu, cosmic=cosmic),
406 timedep=None,
407 constraints=[
408 alignment.constraints.VXDHierarchyConstraints(type=2, pxd=True, svd=True),
409 alignment.constraints.CDCLayerConstraints(z_offset=True, z_scale=False, twist=False)],
410 fixed=alignment.parameters.vxd_sensors(layers=[3, 4, 5, 6]) +
411 alignment.parameters.vxd_sensors(layers=[1, 2], rigid=False, surface2=False, surface3=False, surface4=True),
412 commands=["method inversion 6 0.001", "entries 1000", "threads 10 10"],
413 params=dict(minPValue=0.00001, externalIterations=0, granularity="run"),
414 min_entries=80000)
415
416 # Ignore results for BeamSpot
417 std_components = ROOT.vector('string')()
418 for component in ['VXDAlignment', 'CDCAlignment']:
419 std_components.push_back(component)
420 cal.algorithms[0].algorithm.setComponents(std_components)
421
422 # Sequential run-by-run strategy
423 cal.max_iterations = 0
424 cal.algorithms[0].strategy = strategies.SequentialRunByRun
425
426 return cal
427
428
431
432
433def get_calibrations(input_data, **kwargs):
434 """
435 Required function called by b2caf-prompt-run.
436 Returns full configured 4-stage final alignment for prompt
437
438 """
439 seed(1234)
440
441 cfg = kwargs['expert_config']
442 files = dict()
443
444 for colname in collection_names:
445 file_to_iov = input_data[colname]
446 input_files = list(file_to_iov.keys())
447 files[colname] = input_files
448
449 prompt = create_prompt(files, cfg)
450 beamspot = create_beamspot(files, cfg)
451 stage1 = create_stage1(files, cfg)
452 stage2 = create_stage2(files, cfg)
453
454 requested_iov = kwargs.get("requested_iov", None)
455 output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
456
457 for cal in [prompt, beamspot, stage1, stage2]:
458 for colname in collection_names:
459 if colname not in cal.collections.keys():
460 continue
461 max_processed_events_per_file = cfg[f'{colname}.max_processed_events_per_file']
462 basf2.set_module_parameters(
463 cal.collections[colname].pre_collector_path,
464 'RootInput',
465 entrySequences=[f'0:{max_processed_events_per_file}'], branchNames=HLT_INPUT_OBJECTS)
466
467 for algorithm in cal.algorithms:
468 algorithm.params = {"apply_iov": output_iov}
469
470 # Bugfix for Condor:
471 for cal in [prompt, stage1, stage2]:
472 from alignment.prompt_utils import fix_mille_paths_for_algo
473 fix_mille_paths_for_algo(cal.algorithms[0])
474
475 beamspot.depends_on(prompt)
476 stage1.depends_on(beamspot)
477 stage2.depends_on(stage1)
478
479 # Do not save BeamSpot payloads in the final output database
480 # because alignment is changed -> BeamSpot payloads now invalid
481 beamspot.save_payloads = False
482
483 if cfg["only_prompt"]:
484 return [prompt]
485
486 return [prompt, beamspot, stage1, stage2]
487
488
489if __name__ == '__main__':
490 get_calibrations(dict(cosmic=dict(), mumu=dict(), hadron=dict(), offip=dict()),
491 requested_iov=IoV(0, 0, -1, -1), expert_config=default_config)
make_collection(name, files=None, path=None, **argk)
vxd_sensors(layers=None, rigid=True, surface=True, surface2=True, surface3=True, surface4=True, parameters=None)