Belle II Software  release-05-01-25
event_generation.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 import os
5 
6 import basf2
7 import simulation
8 import generators
9 import beamparameters
10 
11 from . import utilities
12 from .minimal import MinimalRun
13 import tracking.adjustments as adjustments
14 
15 import argparse
16 import logging
17 
18 
19 def get_logger():
20  return logging.getLogger(__name__)
21 
22 
23 # Standard run for generating or reading event #
24 
25 
27  """Read generated events or generate new events"""
28 
29 
30  description = "Simulate events using various generator and detector setups from command line."
31 
32  # Declarative section which can be redefined in a subclass
33 
34 
35  generator_module = None
36 
37  detector_setup = "Default"
38 
39  bkg_files = []
40 
41  components = None
42 
43  disable_deltas = False
44 
45  simulation_output = None
46 
47  def create_argument_parser(self, **kwds):
48  """Convert command-line arguments to basf2 argument list"""
49  argument_parser = super().create_argument_parser(**kwds)
50 
51  setup_argument_group = argument_parser.add_argument_group("Detector setup arguments")
52  setup_argument_group.add_argument(
53  '-d',
54  '--detector',
55  dest='detector_setup',
56  default=argparse.SUPPRESS,
57  metavar='DETECTOR_SETUP_NAME',
58  choices=utilities.NonstrictChoices(detector_setups_by_short_name.keys()),
59  help=('Name of the detector setup to be used')
60  )
61 
62  setup_argument_group.add_argument(
63  '-c',
64  '--component',
65  dest='components',
66  nargs='+',
67  default=argparse.SUPPRESS,
68  metavar='COMPONENTS',
69  action='store',
70  help=('Overrides the components of the detector setup')
71  )
72 
73  generator_argument_group = argument_parser.add_argument_group("Generator arguments")
74  generator_argument_group.add_argument(
75  '-g',
76  '--generator',
77  dest='generator_module',
78  default=argparse.SUPPRESS,
79  metavar='GENERATOR_NAME',
80  choices=utilities.NonstrictChoices(valid_generator_short_names),
81  help='Name module or short name of the generator to be used.',
82  )
83 
84  simulation_argument_group = argument_parser.add_argument_group("Simulation arguments")
85  simulation_argument_group.add_argument(
86  '-b',
87  '--bkg-file',
88  dest='bkg_files',
89  nargs='+',
90  default=self.bkg_files,
91  metavar='BACKGROUND_DIRECTORY',
92  help='Path to folder of files or to a file containing the background to be used. ' +
93  'Can be given multiple times.',
94  )
95 
96  simulation_argument_group.add_argument(
97  '--disable-deltas',
98  action='store_true',
99  help='Disable the generation of delta rays in the simulation'
100  )
101 
102  simulation_argument_group.add_argument(
103  '-so',
104  '--simulation-output',
105  nargs='?',
106  default=self.simulation_output,
107  const=self.root_input_file,
108  dest='simulation_output',
109  help='Only generate and simulate the events and write them to the given output file. Skip rest of the path.'
110  )
111 
112  return argument_parser
113 
114  def configure(self, arguments):
115  """Configure for basf2 job; disable ROOT input if simulating events"""
116  super().configure(arguments)
117  if self.simulation_output:
118  get_logger().info("Requested to simulation run. Deactivate input file")
119 
120  self.root_input_file = None
121 
122  def execute(self):
123  """Run the basf2 job"""
124  if not self.simulation_output:
125  super().execute()
126  return
127 
128  # Run only simulation
129  path = ReadOrGenerateEventsRun.create_path(self)
130  self.run(path)
131 
132  def create_path(self):
133  """Create and configure the basf2 path"""
134  path = super().create_path()
135 
136  # Gearbox & Geometry must always be registered
137  path.add_module("Gearbox")
138  path.add_module("Geometry", useDB=True)
139  if self.detector_setup:
140  detector_setup = self.detector_setup
141  detector_setup_function = detector_setups_by_short_name[detector_setup]
142  components = detector_setup_function(path)
143 
144  if self.components:
145  components = self.components
146 
147  # Only generate events if no input file has been provided
148  if self.root_input_file is None:
149  # Check if generator means a decay file
150  if isinstance(self.generator_module, str) and utilities.find_file(self.generator_module):
151  dec_file_path = utilities.find_file(self.generator_module)
152  add_evtgen_generator(path, dec_file_path)
153  else:
154  # All other possibilities
155  utilities.extend_path(path,
156  self.generator_module,
157  generators_by_short_name,
158  allow_function_import=True)
159 
160  # Only simulate if generator is setup
161  if self.root_input_file is None:
162  bkg_file_paths = get_bkg_file_paths(self.bkg_files)
163 
165  components=components,
166  bkgfiles=bkg_file_paths)
167 
168  if self.disable_deltas:
169  adjustments.disable_deltas(path)
170 
171  # Catch if no generator is added, no background should be simulated and events
172  # are not read from a file.
173  if not bkg_file_paths and self.generator_module is None:
174  raise RuntimeError('Need at least one of root_input_file,'
175  ' generator_module or bkg_files specified.')
176  else:
177  if not os.path.exists(self.root_input_file):
178  raise RuntimeError("Could not find file " + str(self.root_input_file) + ". Generate it with -- -so?")
179 
180  # early write out if simulation output was requested
181  if self.simulation_output:
182  root_output_module = path.add_module('RootOutput',
183  outputFileName=self.simulation_output)
184 
185  return path
186 
187 
189  """Generate events using the EvtGen generator"""
190 
191  generator_module = "EvtGenInput"
192 
193 
194 # Default settings and shorthand names for generator with specific settings #
195 
196 
197 # PDG code of an electorn
198 electron_pdg_code = 11
199 
200 # PDG code of a muon
201 muon_pdg_code = 13
202 
203 # PDG code of a tau
204 tau_pdg_code = 15
205 
206 # PDG code of a pion
207 pion_pdg_code = 211
208 
209 # PDG code of a kaon
210 kaon_pdg_code = 321
211 
212 # PDG code of a protons
213 proton_pdg_code = 2212
214 
215 
216 def add_single_gun_generator(path):
217  """Add ParticleGun with a single muon"""
218  path.add_module("ParticleGun",
219  pdgCodes=[muon_pdg_code, -muon_pdg_code],
220  nTracks=1,
221  varyNTracks=False,
222  momentumGeneration='inversePt',
223  momentumParams=[0.6, 1.4],
224  thetaGeneration='uniform',
225  thetaParams=[17., 150.])
226 
227 
228 def add_transverse_gun_generator(path):
229  """Add ParticleGun to illuminate a region of the phase space with low efficiency"""
230  path.add_module("ParticleGun",
231  pdgCodes=[muon_pdg_code, -muon_pdg_code],
232  nTracks=1,
233  varyNTracks=False,
234  momentumGeneration='inversePt',
235  momentumParams=[0.275, 0.276],
236  thetaGeneration='uniform',
237  thetaParams=[89., 91.])
238 
239 
240 def add_simple_gun_generator(path):
241  """Add ParticleGun firing 10 muons at medium energy"""
242  path.add_module("ParticleGun",
243  pdgCodes=[muon_pdg_code, -muon_pdg_code],
244  nTracks=10,
245  varyNTracks=False,
246  momentumGeneration='inversePt',
247  momentumParams=[0.6, 1.4],
248  thetaGeneration='uniform')
249 
250 
251 def add_low_gun_generator(path):
252  """Add ParticleGun firing 10 muons at low energy"""
253  path.add_module("ParticleGun",
254  pdgCodes=[muon_pdg_code, -muon_pdg_code],
255  nTracks=10,
256  varyNTracks=False,
257  momentumGeneration='inversePt',
258  momentumParams=[0.4, 0.8],
259  thetaGeneration='uniform')
260 
261 
262 def add_gun_generator(path):
263  """Add ParticleGun firing 10 muons with wide energy range"""
264  path.add_module("ParticleGun",
265  pdgCodes=[muon_pdg_code, -muon_pdg_code],
266  nTracks=10,
267  varyNTracks=False,
268  momentumGeneration='inversePt',
269  thetaGeneration='uniform',
270  thetaParams=[17., 150.])
271 
272 
273 def add_eloss_gun_generator(path):
274  """Add ParticleGun particle gun for energy loss estimations"""
275  path.add_module("ParticleGun",
276  pdgCodes=[
277  # muon_pdg_code,
278  # -muon_pdg_code,
279  # electron_pdg_code,
280  # -electron_pdg_code,
281  pion_pdg_code,
282  -pion_pdg_code,
283  # kaon_pdg_code,
284  # -kaon_pdg_code,
285  # proton_pdg_code,
286  # -proton_pdg_code,
287  ],
288  momentumParams=[0.3, 2],
289  nTracks=10,
290  varyNTracks=False,
291  thetaGeneration='uniform',
292  # thetaParams=[17., 150.],
293  thetaParams=[89., 91],
294  )
295 
296 
297 def add_forward_gun_generator(path):
298  """Add ParticleGun with one muon in rather forward direction"""
299  path.add_module("ParticleGun",
300  pdgCodes=[muon_pdg_code, -muon_pdg_code],
301  nTracks=1,
302  varyNTracks=False,
303  momentumGeneration='inversePt',
304  thetaGeneration='uniform',
305  thetaParams=[30., 31.])
306 
307 
308 def add_evtgen_generator(path, dec_file_path=None):
309  """Add Y4S generator"""
311  if dec_file_path:
312  path.add_module("EvtGenInput", userDECFile=dec_file_path)
313  else:
314  path.add_module("EvtGenInput")
315 
316 
317 def add_cosmics_generator(path):
318  """Add simple cosmics generator"""
320 
321 
322 def add_sector_tb_generator(path, sector=1):
323  phiBounds = (240 + 60.0 * sector % 360.0, 300 + 60.0 * sector % 360.0)
324  path.add_module("ParticleGun",
325  pdgCodes=[muon_pdg_code, -muon_pdg_code],
326  nTracks=1,
327  varyNTracks=False,
328  momentumGeneration='uniform',
329  momentumParams=[2.0, 4.0],
330  phiGeneration='uniform',
331  phiParams=phiBounds,
332  thetaGeneration='uniform',
333  thetaParams=[70., 110.])
334 
335 
336 def add_cosmics_tb_generator(path):
337  """Add simple cosmics generator resembling the test beam setup"""
338  path.add_module("Cosmics",
339  ipRequirement=1,
340  ipdr=5,
341  ipdz=15,
342  )
343 
344  # Use point trigger
345  tof_mode = 1
346  # Do not add time of propagation in the scintilator
347  top_mode = False
348 
349  cosmics_selector = path.add_module('CDCCosmicSelector',
350  xOfCounter=0.0,
351  yOfCounter=0.0,
352  zOfCounter=0.0,
353  TOF=tof_mode,
354  TOP=top_mode,
355  cryGenerator=False,
356  )
357 
358  cosmics_selector.if_false(basf2.create_path())
359 
360 
361 def add_cry_tb_generator(path):
362  """Add cry generator resembling the test beam setup"""
363  generators.add_cosmics_generator(path, accept_box=[0.7, 0.3, 0.3],
364  keep_box=[0.7, 0.3, 0.3],
365  pre_general_run_setup="normal")
366 
367 
368 def add_no_generator(path):
369  """Add no generator for e.g. background only studies"""
370  # Nothing to do here since the background files are included in the add_simulation
371  pass
372 
373 # Generator module names hashed by shorthand menomics. Includes
374 # None as a special value for background only simulation
375 generators_by_short_name = {
376  'single_gun': add_single_gun_generator,
377  'transverse_gun': add_transverse_gun_generator,
378  'simple_gun': add_simple_gun_generator,
379  'low_gun': add_low_gun_generator,
380  'forward_gun': add_forward_gun_generator,
381  'gun': add_gun_generator,
382  'eloss_gun': add_eloss_gun_generator,
383  'generic': add_evtgen_generator,
384  "EvtGenInput": add_evtgen_generator, # <- requires beam parameters
385  'cosmics': add_cosmics_generator,
386  'cosmics_tb': add_cosmics_tb_generator,
387  'cry_tb': add_cry_tb_generator,
388  'sector_tb': add_sector_tb_generator,
389  'bkg': add_no_generator,
390  'none': add_no_generator,
391 }
392 
393 # Names of module names and short names of the generators usable in this script.
394 valid_generator_short_names = list(generators_by_short_name.keys())
395 
396 
397 # Memorandum of geometry setups #
398 # ############################# #
399 def setup_default_detector(path):
400  pass
401 
402 
403 def setup_tracking_detector(path):
404  components = ["BeamPipe", "PXD", "SVD", "CDC", "MagneticField"]
405  override = [
406  ("/DetectorComponent[@name='MagneticField']//Component[@type='3d'][2]/ExcludeRadiusMax",
407  "4.20", "m", ) # Remove the second compontent which is the magnetic field outside the tracking volumn.
408  ]
409 
410  adjustments.adjust_module(path, "Gearbox", override=override)
411  adjustments.adjust_module(path, "Geometry", components=components)
412  return components
413 
414 
415 def setup_tracking_detector_constant_b(path):
416  components = ["BeamPipe", "PXD", "SVD", "CDC", "MagneticFieldConstant4LimitedRCDC"]
417  adjustments.adjust_module(path, "Geometry", components=components)
418  return components
419 
420 
421 def setup_cdc_cr_test(path):
422  components = ["CDC"]
423  override = [
424  # Reset the top volume to accomodate cosmics that can hit all parts of the detector
425  ("/Global/length", "8.", "m"),
426  ("/Global/width", "8.", "m"),
427  ("/Global/height", "1.5", "m"),
428 
429  # Adjustments of the CDC setup
430  ("/DetectorComponent[@name='CDC']//t0FileName", "t0.dat", ""),
431  ("/DetectorComponent[@name='CDC']//xtFileName", "xt_noB_v1.dat", ""),
432  ("/DetectorComponent[@name='CDC']//GlobalPhiRotation", "1.875", "deg"),
433  # ("/DetectorComponent[@name='CDC']//bwFileName", "badwire_CDCTOP.dat", ""),
434  ]
435 
436  adjustments.adjust_module(path, "Gearbox", override=override)
437  adjustments.adjust_module(path, "Geometry", components=components)
438  return components
439 
440 
441 def setup_cdc_top_test(path):
442  components = ["CDC"]
443  override = [
444  # Reset the top volume: must be larger than the generated surface and higher than the detector
445  # It is the users responsibility to ensure a full angular coverage
446  ("/Global/length", "8.", "m"),
447  ("/Global/width", "8.", "m"),
448  ("/Global/height", "1.5", "m"),
449 
450  # Adjustments of the CDC setup
451  ("/DetectorComponent[@name='CDC']//t0FileName", "t0.dat", ""),
452  ("/DetectorComponent[@name='CDC']//xtFileName", "xt_noB_v1.dat", ""),
453  # ("/DetectorComponent[@name='CDC']//bwFileName", "badwire_CDCTOP.dat", ""),
454  ("/DetectorComponent[@name='CDC']//GlobalPhiRotation", "1.875", "deg"),
455  ("/DetectorComponent[@name='MagneticField']//Component/Z", "0", ""),
456  ]
457 
458  adjustments.adjust_module(path, "Gearbox",
459  override=override,
460  fileName="geometry/CDCcosmicTests.xml" # <- does something mysterious to the reconstruction...
461  )
462 
463  adjustments.adjust_module(path, "Geometry", components=components)
464  return components
465 
466 detector_setups_by_short_name = {
467  "Default": setup_default_detector,
468  'TrackingDetector': setup_tracking_detector,
469  'TrackingDetectorConstB': setup_tracking_detector_constant_b,
470  'CDCCRTest': setup_cdc_cr_test,
471  'CDCTOPTest': setup_cdc_top_test,
472 }
473 
474 
475 # Heuristic to find background files #
476 # ################################## #
477 
478 def is_bkg_file(bkg_file_path):
479  """Test if a file path points to a file containing background mixins.
480 
481  Returns
482  -------
483  bool
484 
485  Note
486  ----
487  Simple test only checks if file exists and ends with ".root"
488  """
489  return os.path.isfile(bkg_file_path) and bkg_file_path.endswith('.root')
490 
491 
492 def get_bkg_file_paths(bkg_dir_or_file_paths):
493  """Unpacks the content of a single or a list of directories and/or files filtering for
494  files containing background mixins.
495 
496  Parameters
497  ----------
498  bkg_dir_or_file_paths : string or iterable of strings
499  Single file or single directory in which background files are located or
500  a list of files and/or directories.
501 
502  Returns
503  -------
504  list(string)
505  A list of paths to individual background files.
506  """
507 
508  if isinstance(bkg_dir_or_file_paths, str):
509  bkg_dir_or_file_path = bkg_dir_or_file_paths
510  bkg_dir_or_file_paths = [bkg_dir_or_file_path]
511 
512  result = []
513  for bkg_dir_or_file_path in bkg_dir_or_file_paths:
514  if is_bkg_file(bkg_dir_or_file_path):
515  bkg_file_path = bkg_dir_or_file_path
516  result.append(bkg_file_path)
517  elif os.path.isdir(bkg_dir_or_file_path):
518 
519  bkg_dir_path = bkg_dir_or_file_path
520  bkg_dir_contents = os.listdir(bkg_dir_path)
521  for dir_or_file_name in bkg_dir_contents:
522  dir_or_file_path = os.path.join(bkg_dir_path, dir_or_file_name)
523  if is_bkg_file(dir_or_file_path):
524  bkg_file_path = dir_or_file_path
525  result.append(bkg_file_path)
526 
527  return result
tracking.run.event_generation.ReadOrGenerateEventsRun.simulation_output
simulation_output
By default, do no store the simulation output.
Definition: event_generation.py:45
tracking.run.minimal.EmptyRun.run
def run(self, path)
Definition: minimal.py:52
tracking.run.event_generation.ReadOrGenerateEventsRun.bkg_files
list bkg_files
By default, no background overlay.
Definition: event_generation.py:39
tracking.run.event_generation.ReadOrGenerateEventsRun.generator_module
generator_module
By default, do not generate events.
Definition: event_generation.py:35
simulation.add_simulation
def add_simulation(path, components=None, bkgfiles=None, bkgOverlay=True, forceSetPXDDataReduction=False, usePXDDataReduction=True, cleanupPXDDataReduction=True, generate_2nd_cdc_hits=False, simulateT0jitter=False, usePXDGatedMode=False)
Definition: simulation.py:114
tracking.adjustments
Definition: adjustments.py:1
tracking.run.event_generation.ReadOrGenerateEventsRun.components
components
By default, do specific components.
Definition: event_generation.py:41
beamparameters.add_beamparameters
def add_beamparameters(path, name, E_cms=None, **argk)
Definition: beamparameters.py:189
tracking.run.event_generation.StandardEventGenerationRun
Definition: event_generation.py:188
tracking.run.event_generation.ReadOrGenerateEventsRun
Definition: event_generation.py:26
tracking.run.event_generation.ReadOrGenerateEventsRun.disable_deltas
bool disable_deltas
By default, do not disable delta-ray generation.
Definition: event_generation.py:43
tracking.run.event_generation.ReadOrGenerateEventsRun.execute
def execute(self)
Definition: event_generation.py:122
tracking.run.minimal.MinimalRun
Definition: minimal.py:95
tracking.run.event_generation.ReadOrGenerateEventsRun.root_input_file
root_input_file
generating events, so there is no ROOT input file
Definition: event_generation.py:120
tracking.run.event_generation.ReadOrGenerateEventsRun.configure
def configure(self, arguments)
Definition: event_generation.py:114
tracking.run.event_generation.ReadOrGenerateEventsRun.create_argument_parser
def create_argument_parser(self, **kwds)
Definition: event_generation.py:47
tracking.run.event_generation.ReadOrGenerateEventsRun.create_path
def create_path(self)
Definition: event_generation.py:132
tracking.run.event_generation.ReadOrGenerateEventsRun.detector_setup
string detector_setup
By default, use the default detector setup.
Definition: event_generation.py:37
tracking.run.utilities.NonstrictChoices
Definition: utilities.py:24
generators.add_cosmics_generator
def add_cosmics_generator(path, components=None, global_box_size=None, accept_box=None, keep_box=None, geometry_xml_file='geometry/Beast2_phase2.xml', cosmics_data_dir='data/generators/modules/cryinput/', setup_file='generators/scripts/cry.setup', data_taking_period='gcr2017', top_in_counter=False)
Definition: generators.py:677