Belle II Software  release-05-01-25
caf_svd_time.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 '''
4 Script to perform the svd time calibration with the CoG6, CoG3 and ELS3 algorithms
5 '''
6 
7 import sys
8 import datetime
9 import random
10 
11 from ROOT.Belle2 import SVDCoGTimeCalibrationAlgorithm
12 from ROOT.Belle2 import SVD3SampleCoGTimeCalibrationAlgorithm
13 from ROOT.Belle2 import SVD3SampleELSTimeCalibrationAlgorithm
14 from ROOT.Belle2 import SVDTimeValidationAlgorithm
15 
16 import basf2 as b2
17 
18 import rawdata as raw
19 from softwaretrigger.constants import HLT_INPUT_OBJECTS
20 from tracking import add_tracking_reconstruction
21 
22 from caf.framework import Calibration
23 from caf import strategies
24 from caf.utils import IoV
25 from prompt import CalibrationSettings
26 from prompt.utils import filter_by_max_events_per_run
27 
28 b2.set_log_level(b2.LogLevel.INFO)
29 
30 random.seed(42)
31 
32 now = datetime.datetime.now()
33 
34 settings = CalibrationSettings(name="caf_svd_time",
35  expert_username="gdujany",
36  description=__doc__,
37  input_data_formats=["raw"],
38  input_data_names=["hlt_hadron"],
39  depends_on=[],
40  expert_config={
41  "max_events_per_run": 10000,
42  "isMC": False,
43  })
44 
45 
47 
48 
49 def remove_module(path, name):
50 
51  new_path = b2.create_path()
52  for m in path.modules():
53  if name != m.name():
54  new_path.add_module(m)
55  return new_path
56 
57 
58 
59 NEW_RECO_DIGITS_NAME = "SVDRecoDigitsFromTracks"
60 NEW_SHAPER_DIGITS_NAME = "SVDShaperDigitsFromTracks"
61 
62 
63 def create_collector(name="SVDTimeCalibrationCollector",
64  clusters="SVDClustersFromTracks",
65  event_info="SVDEventInfo",
66  event_t0="EventT0",
67  granularity="run"):
68  """
69  Simply creates a SVDTimeCalibrationCollector module with some options.
70 
71  Returns:
72  pybasf2.Module
73  """
74 
75  collector = b2.register_module("SVDTimeCalibrationCollector")
76  collector.set_name(name)
77  collector.param("SVDClustersFromTracksName", clusters)
78  collector.param("SVDEventInfoName", event_info)
79  collector.param("EventT0Name", event_t0)
80  collector.param("granularity", granularity)
81 
82  return collector
83 
84 
85 def create_validation_collector(name="SVDTimeValidationCollector",
86  clusters="SVDClusters",
87  clusters_onTracks="SVDClustersOnTracks",
88  event_info="SVDEventInfo",
89  event_t0="EventT0",
90  granularity="run"):
91  """
92  Simply creates a SVDTimeCalibrationCollector module with some options.
93 
94  Returns:
95  pybasf2.Module
96  """
97 
98  collector = b2.register_module("SVDTimeValidationCollector")
99  collector.set_name(name)
100  collector.param("SVDClustersName", clusters)
101  collector.param("SVDClustersOnTracksName", clusters_onTracks)
102  collector.param("SVDEventInfoName", event_info)
103  collector.param("EventT0Name", event_t0)
104  collector.param("granularity", granularity)
105 
106  return collector
107 
108 
109 def create_algorithm(unique_id, prefix="", min_entries=10000):
110  """
111  Simply creates a SVDCoGTimeCalibrationAlgorithm class with some options.
112 
113  Returns:
114  ROOT.Belle2.SVDCoGTimeCalibrationAlgorithm
115  """
116  if "CoG6" in prefix:
117  algorithm = SVDCoGTimeCalibrationAlgorithm(unique_id)
118  if "CoG3" in prefix:
119  algorithm = SVD3SampleCoGTimeCalibrationAlgorithm(unique_id)
120  if "ELS3" in prefix:
121  algorithm = SVD3SampleELSTimeCalibrationAlgorithm(unique_id)
122  if prefix:
123  algorithm.setPrefix(prefix)
124  algorithm.setMinEntries(min_entries)
125 
126  return algorithm
127 
128 
129 def create_validation_algorithm(prefix="", min_entries=10000):
130  """
131  Simply creates a SVDCoGTimeValidationAlgorithm class with some options.
132 
133  Returns:
134  ROOT.Belle2.SVDCoGTimeValidationAlgorithm
135  """
136  algorithm = SVDTimeValidationAlgorithm()
137  if prefix:
138  algorithm.setPrefix(prefix)
139  algorithm.setMinEntries(min_entries)
140 
141  return algorithm
142 
143 
144 def create_svd_clusterizer(name="ClusterReconstruction",
145  clusters="SVDClustersFromTracks",
146  reco_digits=None,
147  shaper_digits=None,
148  time_algorithm=0,
149  get_3sample_raw_time=False):
150  """
151  Simply creates a SVDSimpleClusterizer module with some options.
152 
153  Returns:
154  pybasf2.Module
155  """
156 
157  cluster = b2.register_module("SVDSimpleClusterizer")
158  cluster.set_name(name)
159  cluster.param("Clusters", clusters)
160  if reco_digits is not None:
161  cluster.param("RecoDigits", reco_digits)
162  if shaper_digits is not None:
163  cluster.param("ShaperDigits", shaper_digits)
164  cluster.param("timeAlgorithm", time_algorithm)
165  if get_3sample_raw_time:
166  cluster.param("Calibrate3SampleWithEventT0", False)
167  return cluster
168 
169 
170 def create_pre_collector_path(clusterizers, isMC=False, is_validation=False):
171  """
172  Create a basf2 path that runs a common reconstruction path and also runs several SVDSimpleClusterizer
173  modules with different configurations. This way they re-use the same reconstructed objects.
174 
175  Parameters:
176  clusterizers (list[pybasf2.Module]): All the differently configured SVDSimpleClusterizer modules.
177  They should output to different datastore objects.
178 
179  returns:
180  pybasf2.Path
181  """
182  # Set-up re-processing path
183  path = b2.create_path()
184 
185  # Read from file only what is needed
186  path.add_module("RootInput", branchNames=HLT_INPUT_OBJECTS)
187 
188  # unpack raw data to do the tracking
189  if not isMC:
190  raw.add_unpackers(path, components=['PXD', 'SVD', 'CDC'])
191  else:
192  path.add_module("Gearbox")
193  path.add_module("Geometry")
194 
195  # proceed only if we acquired 6-sample strips
196  skim6SampleEvents = b2.register_module("SVD6SampleEventSkim")
197  path.add_module(skim6SampleEvents)
198  emptypath = b2.create_path()
199  skim6SampleEvents.if_false(emptypath)
200 
201  # run tracking reconstruction
202  add_tracking_reconstruction(path)
203  path = remove_module(path, "V0Finder")
204  if not is_validation:
205  b2.set_module_parameters(path, 'SVDCoGTimeEstimator', CalibrationWithEventT0=False)
206 
207  # repeat svd reconstruction using only SVDShaperDigitsFromTracks
208  path.add_module("SVDShaperDigitsFromTracks")
209  cog = b2.register_module("SVDCoGTimeEstimator")
210  cog.set_name("CoGReconstruction")
211  path.add_module(cog)
212  b2.set_module_parameters(path, 'CoGReconstruction',
213  ShaperDigits=NEW_SHAPER_DIGITS_NAME,
214  RecoDigits=NEW_RECO_DIGITS_NAME,
215  CalibrationWithEventT0=is_validation)
216 
217  for cluster in clusterizers:
218  path.add_module(cluster)
219 
220  path = remove_module(path, "SVDMissingAPVsClusterCreator")
221 
222  return path
223 
224 
225 def get_calibrations(input_data, **kwargs):
226 
227  file_to_iov_physics = input_data["hlt_hadron"]
228  expert_config = kwargs.get("expert_config")
229  max_events_per_run = expert_config["max_events_per_run"] # Maximum number of events selected per each run
230  isMC = expert_config["isMC"]
231 
232  reduced_file_to_iov_physics = filter_by_max_events_per_run(file_to_iov_physics,
233  max_events_per_run, random_select=True)
234  good_input_files = list(reduced_file_to_iov_physics.keys())
235 
236  b2.B2INFO(f"Total number of files used as input = {len(good_input_files)}")
237 
238  exps = [i.exp_low for i in reduced_file_to_iov_physics.values()]
239  runs = sorted([i.run_low for i in reduced_file_to_iov_physics.values()])
240 
241  firstRun = runs[0]
242  lastRun = runs[-1]
243  expNum = exps[0]
244 
245  if not len(good_input_files):
246  print("No good input files found! Check that the input files have entries != 0!")
247  sys.exit(1)
248 
249  cog6_suffix = "_CoG6"
250  cog3_suffix = "_CoG3"
251  els3_suffix = "_ELS3"
252 
253  calType = "Prompt"
254  if isMC:
255  calType = "MC"
256  unique_id_cog6 = f"SVDCoGTimeCalibrations_{calType}_{now.isoformat()}_INFO:_3rdOrderPol_TBindep_" \
257  f"Exp{expNum}_runsFrom{firstRun}to{lastRun}"
258  print(f"\nUniqueID_CoG6:\n{unique_id_cog6}")
259 
260  unique_id_cog3 = f"SVD3SampleCoGTimeCalibrations_{calType}_{now.isoformat()}_INFO:_3rdOrderPol_TBindep_" \
261  f"Exp{expNum}_runsFrom{firstRun}to{lastRun}"
262  print(f"\nUniqueID_CoG3:\n{unique_id_cog3}")
263 
264  unique_id_els3 = f"SVD3SampleELSTimeCalibrations_{calType}_{now.isoformat()}_INFO:_TBindep_" \
265  f"Exp{expNum}_runsFrom{firstRun}to{lastRun}"
266  print(f"\nUniqueID_ELS3:\n{unique_id_els3}")
267 
268  requested_iov = kwargs.get("requested_iov", None)
269  output_iov = IoV(requested_iov.exp_low, requested_iov.run_low, -1, -1)
270 
271 
274 
275  cog6 = create_svd_clusterizer(
276  name=f"ClusterReconstruction{cog6_suffix}",
277  clusters=f"SVDClustersFromTracks{cog6_suffix}",
278  reco_digits=NEW_RECO_DIGITS_NAME,
279  shaper_digits=NEW_SHAPER_DIGITS_NAME,
280  time_algorithm=0,
281  get_3sample_raw_time=True)
282 
283  cog3 = create_svd_clusterizer(
284  name=f"ClusterReconstruction{cog3_suffix}",
285  clusters=f"SVDClustersFromTracks{cog3_suffix}",
286  reco_digits=NEW_RECO_DIGITS_NAME,
287  shaper_digits=NEW_SHAPER_DIGITS_NAME,
288  time_algorithm=1,
289  get_3sample_raw_time=True)
290 
291  els3 = create_svd_clusterizer(
292  name=f"ClusterReconstruction{els3_suffix}",
293  clusters=f"SVDClustersFromTracks{els3_suffix}",
294  reco_digits=NEW_RECO_DIGITS_NAME,
295  shaper_digits=NEW_SHAPER_DIGITS_NAME,
296  time_algorithm=2,
297  get_3sample_raw_time=True)
298 
299 
302  eventInfo = "SVDEventInfo"
303  if isMC:
304  eventInfo = "SVDEventInfoSim"
305 
306  coll_cog6 = create_collector(
307  name=f"SVDTimeCalibrationCollector{cog6_suffix}",
308  clusters=f"SVDClustersFromTracks{cog6_suffix}",
309  event_info=eventInfo,
310  event_t0="EventT0")
311 
312  algo_cog6 = create_algorithm(
313  unique_id_cog6,
314  prefix=coll_cog6.name(),
315  min_entries=10000)
316 
317  coll_cog3 = create_collector(
318  name=f"SVDTimeCalibrationCollector{cog3_suffix}",
319  clusters=f"SVDClustersFromTracks{cog3_suffix}",
320  event_info=eventInfo,
321  event_t0="EventT0")
322 
323  algo_cog3 = create_algorithm(
324  unique_id_cog3,
325  prefix=coll_cog3.name(),
326  min_entries=10000)
327 
328  coll_els3 = create_collector(
329  name=f"SVDTimeCalibrationCollector{els3_suffix}",
330  clusters=f"SVDClustersFromTracks{els3_suffix}",
331  event_info=eventInfo,
332  event_t0="EventT0")
333 
334  algo_els3 = create_algorithm(
335  unique_id_els3,
336  prefix=coll_els3.name(),
337  min_entries=10000)
338 
339 
343 
344  pre_collector_path = create_pre_collector_path(
345  clusterizers=[cog6, cog3, els3],
346  isMC=isMC)
347  pre_collector_path.add_module(coll_cog6)
348  pre_collector_path.add_module(coll_cog3)
349  # We leave the coll_els3 to be the one "managed" by the CAF
350 
351  # calibration setup
352  calibration = Calibration("SVDTime",
353  collector=coll_els3, # The other collectors are in the pre_collector_path itself
354  algorithms=[algo_cog6, algo_cog3, algo_els3],
355  input_files=good_input_files,
356  pre_collector_path=pre_collector_path)
357 
358  calibration.strategies = strategies.SequentialBoundaries
359 
360  for algorithm in calibration.algorithms:
361  algorithm.params = {"iov_coverage": output_iov}
362 
363 
366 
367  val_cog6 = create_svd_clusterizer(
368  name=f"ClusterReconstruction{cog6_suffix}",
369  clusters=f"SVDClusters{cog6_suffix}",
370  time_algorithm=0)
371 
372  val_cog6_onTracks = create_svd_clusterizer(
373  name=f"ClusterReconstruction{cog6_suffix}_onTracks",
374  clusters=f"SVDClusters{cog6_suffix}_onTracks",
375  reco_digits=NEW_RECO_DIGITS_NAME,
376  shaper_digits=NEW_SHAPER_DIGITS_NAME,
377  time_algorithm=0)
378 
379  val_cog3 = create_svd_clusterizer(
380  name=f"ClusterReconstruction{cog3_suffix}",
381  clusters=f"SVDClusters{cog3_suffix}",
382  time_algorithm=1)
383 
384  val_cog3_onTracks = create_svd_clusterizer(
385  name=f"ClusterReconstruction{cog3_suffix}_onTracks",
386  clusters=f"SVDClusters{cog3_suffix}_onTracks",
387  reco_digits=NEW_RECO_DIGITS_NAME,
388  shaper_digits=NEW_SHAPER_DIGITS_NAME,
389  time_algorithm=1)
390 
391  val_els3 = create_svd_clusterizer(
392  name=f"ClusterReconstruction{els3_suffix}",
393  clusters=f"SVDClusters{els3_suffix}",
394  time_algorithm=2)
395 
396  val_els3_onTracks = create_svd_clusterizer(
397  name=f"ClusterReconstruction{els3_suffix}_onTracks",
398  clusters=f"SVDClusters{els3_suffix}_onTracks",
399  reco_digits=NEW_RECO_DIGITS_NAME,
400  shaper_digits=NEW_SHAPER_DIGITS_NAME,
401  time_algorithm=2)
402 
403  val_coll_cog6 = create_validation_collector(
404  name=f"SVDTimeValidationCollector{cog6_suffix}",
405  clusters=f"SVDClusters{cog6_suffix}",
406  clusters_onTracks=f"SVDClusters{cog6_suffix}_onTracks",
407  event_info=eventInfo,
408  event_t0="EventT0")
409 
410  val_algo_cog6 = create_validation_algorithm(
411  prefix=val_coll_cog6.name(),
412  min_entries=10000)
413 
414  val_coll_cog3 = create_validation_collector(
415  name=f"SVDTimeValidationCollector{cog3_suffix}",
416  clusters=f"SVDClusters{cog3_suffix}",
417  clusters_onTracks=f"SVDClusters{cog3_suffix}_onTracks",
418  event_info=eventInfo,
419  event_t0="EventT0")
420 
421  val_algo_cog3 = create_validation_algorithm(
422  prefix=val_coll_cog3.name(),
423  min_entries=10000)
424 
425  val_coll_els3 = create_validation_collector(
426  name=f"SVDTimeValidationCollector{els3_suffix}",
427  clusters=f"SVDClusters{els3_suffix}",
428  clusters_onTracks=f"SVDClusters{els3_suffix}_onTracks",
429  event_info=eventInfo,
430  event_t0="EventT0")
431 
432  val_algo_els3 = create_validation_algorithm(
433  prefix=val_coll_els3.name(),
434  min_entries=10000)
435 
436  val_pre_collector_path = create_pre_collector_path(
437  clusterizers=[val_cog6, val_cog6_onTracks,
438  val_cog3, val_cog3_onTracks,
439  val_els3, val_els3_onTracks],
440  isMC=isMC, is_validation=True)
441  val_pre_collector_path.add_module(val_coll_cog6)
442  val_pre_collector_path.add_module(val_coll_cog3)
443 
444  val_calibration = Calibration("SVDTimeValidation",
445  collector=val_coll_els3,
446  algorithms=[val_algo_cog6, val_algo_cog3,
447  val_algo_els3],
448  input_files=good_input_files,
449  pre_collector_path=val_pre_collector_path)
450 
451  val_calibration.strategies = strategies.SequentialRunByRun
452 
453  for algorithm in val_calibration.algorithms:
454  algorithm.params = {"iov_coverage": output_iov}
455 
456  val_calibration.depends_on(calibration)
457 
458  return [calibration, val_calibration]
prompt.utils
Definition: utils.py:1
strategies.SequentialRunByRun
Definition: strategies.py:251
softwaretrigger.constants
Definition: constants.py:1
strategies.SequentialBoundaries
Definition: strategies.py:648
Calibration
Definition: Calibration.py:1