Belle II Software  release-05-01-25
tracking_efficiency_helpers.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
13 
14 """
15 <header>
16 <contact>Michael Ziegler, michael.ziegler2@kit.edu</contact>
17 <description> To be filled by the contact person</description>
18 </header>
19 """
20 
21 from simulation import add_simulation
22 from reconstruction import add_reconstruction, add_mc_reconstruction
23 from basf2 import *
24 import glob
25 import os
26 import sys
27 
28 
29 def get_generated_pdg_code():
30  return 13
31 
32 
33 def get_simulation_components():
34  return None
35 
36 
37 def get_reconstruction_components():
38  # Could be different from get_simulation_components, if only testing subdetectors.
39  return get_simulation_components()
40 
41 
42 def get_generated_pt_value(index):
43  """
44  List with generated pt values is created in this function. With the
45  parameter index, one is able to access single pt values. If index == -1,
46  the number of pt values in the list is returned.
47 
48  @param index
49 
50  @return double
51  """
52 
53  # pt_values = [0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 2.5, 3.0]
54  pt_values = [
55  0.05,
56  0.1,
57  0.25,
58  0.4,
59  0.6,
60  1.,
61  1.5,
62  2.,
63  3.,
64  4.,
65  ]
66 
67  if index == -1:
68  return len(pt_values)
69  try:
70  return pt_values[index]
71  except IndexError:
72  print('ERROR in %s. Index is out of range. Only %d elements in list.'
73  % (get_generated_pt_value.__name__, len(pt_values)))
74  sys.exit(1)
75 
76 
77 def get_generated_pt_values():
78  """
79  Collects all pt values with help of the function get_generated_pt_value
80  and stores them in a list, that is returned
81 
82  @return list of pt values
83  """
84 
85  list = []
86  for index in range(get_generated_pt_value(-1)):
87  list.append(get_generated_pt_value(index))
88 
89  return list
90 
91 
92 def additional_options(path):
93  """
94  This sets the specifics of simulation and reconstruction in order
95  to use consistent settings across the various involved modules.
96  """
97 
98  realisticCDCGeo = 1
99  useInWirePropagation = 1
100  useTime = 1
101  useWireSag = 1
102  for m in path.modules():
103  if realisticCDCGeo:
104  if m.type() == 'CDCDigitizer':
105  m.param('AddInWirePropagationDelay', useInWirePropagation)
106  m.param('AddTimeOfFlight', useTime)
107  m.param('CorrectForWireSag', useWireSag)
108  else:
109  if m.type() == 'CDCDigitizer':
110  m.param('AddInWirePropagationDelay', 0)
111  m.param('AddTimeOfFlight', 0)
112  m.param('CorrectForWireSag', 0)
113 
114  if m.type() == 'DAFRecoFitter':
115  m.param('pdgCodesToUseForFitting', [get_generated_pdg_code()])
116 
117  if m.type() == "TrackCreator":
118  m.param('pdgCodes', [get_generated_pdg_code(), 211])
119 
120 
121 def run_simulation(path, pt_value, output_filename=''):
122  """Add needed modules to the path and set parameters and start it"""
123 
124  # evt meta
125  eventinfosetter = register_module('EventInfoSetter')
126 
127  # generate one event
128  eventinfosetter.param('expList', [0])
129  eventinfosetter.param('runList', [1])
130  eventinfosetter.param('evtNumList', [200])
131 
132  path.add_module(eventinfosetter)
133 
134  progress = register_module('Progress')
135  path.add_module(progress)
136 
137  # ParticleGun
138  pgun = register_module('ParticleGun')
139  # choose the particles you want to simulate
140  param_pgun = {
141  'pdgCodes': [get_generated_pdg_code(), -get_generated_pdg_code()],
142  'nTracks': 10,
143  'varyNTracks': 0,
144  'momentumGeneration': 'uniformPt',
145  'momentumParams': [pt_value, pt_value],
146  'vertexGeneration': 'fixed',
147  'xVertexParams': [0.0],
148  'yVertexParams': [0.0],
149  'zVertexParams': [0.0],
150  'thetaGeneration': 'uniformCos',
151  }
152 
153  pgun.param(param_pgun)
154 
155  path.add_module(pgun)
156 
157  background_files = []
158  if 'BELLE2_BACKGROUND_DIR' in os.environ:
159  background_files += glob.glob(os.environ['BELLE2_BACKGROUND_DIR'] + '/*.root')
160 
161  print('Number of used background files (%d): ' % len(background_files))
162 
163  if len(background_files) == 0:
164  background_files = None
165 
166  # add simulation modules to the path
167  add_simulation(path, get_simulation_components(), background_files)
168 
169  additional_options(path)
170  # write output root file
171  if output_filename != '':
172  root_output = register_module('RootOutput')
173  root_output.param('outputFileName', output_filename)
174  path.add_module(root_output)
175 
176 
177 def run_reconstruction(path, output_file_name, input_file_name=''):
178  if input_file_name != '':
179  root_input = register_module('RootInput')
180  root_input.param('inputFileNames', input_file_name)
181  path.add_module(root_input)
182 
183  gearbox = register_module('Gearbox')
184  path.add_module(gearbox)
185 
186  geometry = register_module('Geometry')
187  geometry.param("useDB", True)
188  path.add_module(geometry)
189 
190  add_reconstruction(path, get_reconstruction_components(), pruneTracks=0)
191  # add_mc_reconstruction(path, get_reconstruction_components(), pruneTracks=0)
192 
193  tracking_efficiency = register_module('StandardTrackingPerformance')
194  # tracking_efficiency.logging.log_level = LogLevel.DEBUG
195  tracking_efficiency.param('outputFileName', output_file_name)
196  path.add_module(tracking_efficiency)
197 
198  additional_options(path)