Belle II Software  release-05-01-25
eventSimulation.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
28 
29 
30 from basf2 import *
31 from beamparameters import add_beamparameters
32 from simulation import add_simulation
33 
34 import os
35 import sys
36 import random
37 
38 from pxd import add_pxd_reconstruction
39 from svd import add_svd_reconstruction
40 
41 
42 # ---------------------------------------------------------------------------------------
43 
44 # Set Random Seed for reproducable simulation. 0 means really random.
45 rndseed = 12345
46 # assume the first argument is the random seed
47 if(len(sys.argv) > 1):
48  rndseed = sys.argv[1]
49 
50 outputDir = './'
51 # assume second argument is the output directory
52 if(len(sys.argv) > 2):
53  outputDir = sys.argv[2]
54 
55 set_random_seed(rndseed)
56 
57 
58 # Set log level. Can be overridden with the "-l LEVEL" flag for basf2.
59 set_log_level(LogLevel.ERROR)
60 
61 # ---------------------------------------------------------------------------------------
62 main = create_path()
63 
64 eventinfosetter = register_module('EventInfoSetter')
65 # default phase3 geometry:
66 exp_number = 0
67 # if environment variable is set then phase2 (aka Beast2) geometry will be taken
68 if os.environ.get('USE_BEAST2_GEOMETRY'):
69  exp_number = 1002
70 eventinfosetter.param("expList", [exp_number])
71 main.add_module(eventinfosetter)
72 
73 eventinfoprinter = register_module('EventInfoPrinter')
74 main.add_module(eventinfoprinter)
75 
76 progress = register_module('Progress')
77 main.add_module(progress)
78 
79 # ---------------------------------------------------------------------------------------
80 # Simulation Settings:
81 
82 
83 # randomize the vertex position (flatly distributed) to make the sectormap more robust wrt. changing beam position
84 # minima and maxima of the beam position given in cm
85 random.seed(rndseed)
86 vertex_x_min = -0.1
87 vertex_x_max = 0.1
88 vertex_y_min = -0.1
89 vertex_y_max = 0.1
90 vertex_z_min = -0.5
91 vertex_z_max = 0.5
92 
93 vertex_x = random.uniform(vertex_x_min, vertex_x_max)
94 vertex_y = random.uniform(vertex_y_min, vertex_y_max)
95 vertex_z = random.uniform(vertex_z_min, vertex_z_max)
96 
97 print("WARNING: setting non-default beam vertex at x= " + str(vertex_x) + " y= " + str(vertex_y) + " z= " + str(vertex_z))
98 
99 # Particle Gun:
100 # One can add more particle gun modules if wanted.
101 
102 # additional flatly smear the muon vertex between +/- this value
103 vertex_delta = 0.005 # in cm
104 
105 particlegun = register_module('ParticleGun')
106 particlegun.logging.log_level = LogLevel.WARNING
107 param_pGun = {
108  'pdgCodes': [13, -13], # 13 = muon --> negatively charged!
109  'nTracks': 1,
110  'momentumGeneration': 'uniform',
111  'momentumParams': [0.1, 4],
112  'vertexGeneration': 'uniform',
113  'xVertexParams': [vertex_x - vertex_delta, vertex_x + vertex_delta], # in cm...
114  'yVertexParams': [vertex_y - vertex_delta, vertex_y + vertex_delta],
115  'zVertexParams': [vertex_z - vertex_delta, vertex_z + vertex_delta]
116 }
117 
118 particlegun.param(param_pGun)
119 main.add_module(particlegun)
120 
121 # EvtGen Simulation:
122 # TODO: There are newer convenience functions for this task -> Include them!
123 # Beam parameters
124 beamparameters = add_beamparameters(main, "Y4S")
125 beamparameters.param("vertex", [vertex_x, vertex_y, vertex_z])
126 # print_params(beamparameters)
127 
128 evtgenInput = register_module('EvtGenInput')
129 evtgenInput.logging.log_level = LogLevel.WARNING
130 main.add_module(evtgenInput)
131 
132 
133 # ---------------------------------------------------------------------------------------
134 
135 
136 # Detector Simulation:
137 add_simulation(path=main,
138  usePXDDataReduction=False, # for training one does not want the data reduction
139  components=None) # dont specify components because else not the whole geometry will be loaded!
140 
141 # needed for fitting
142 main.add_module('SetupGenfitExtrapolation')
143 
144 # this adds the clusters for PXD and SVD (was done by the usePXDDataReduction previously) which are needed in the next steps
145 add_pxd_reconstruction(path=main)
146 add_svd_reconstruction(path=main)
147 
148 # ---------------------------------------------------------------------------------------
149 # Setting up the MC based track finder.
150 mctrackfinder = register_module('TrackFinderMCTruthRecoTracks')
151 mctrackfinder.param('UseCDCHits', False)
152 mctrackfinder.param('UseSVDHits', True)
153 # Always use PXD hits! For SVD only, these will be filtered later when converting to SPTrackCand
154 # 15.02.2018: deactivated PXD again as we dont do pxd track finding anymore and to not bias the fit
155 mctrackfinder.param('UsePXDHits', False)
156 mctrackfinder.param('Smearing', False)
157 mctrackfinder.param('MinimalNDF', 6)
158 mctrackfinder.param('WhichParticles', ['primary'])
159 mctrackfinder.param('RecoTracksStoreArrayName', 'MCRecoTracks')
160 # set up the track finder to only use the first half loop of the track and discard all other hits
161 mctrackfinder.param('UseNLoops', 0.5)
162 mctrackfinder.param('discardAuxiliaryHits', True)
163 main.add_module(mctrackfinder)
164 
165 # include a track fit into the chain (sequence adopted from the tracking scripts)
166 # Correct time seed: Do I need it for VXD only tracks ????
167 main.add_module("IPTrackTimeEstimator", recoTracksStoreArrayName="MCRecoTracks", useFittedInformation=False)
168 # track fitting
169 daffitter = register_module("DAFRecoFitter")
170 daffitter.param('recoTracksStoreArrayName', "MCRecoTracks")
171 # daffitter.logging.log_level = LogLevel.DEBUG
172 main.add_module(daffitter)
173 # also used in the tracking sequence (multi hypothesis)
174 # may be overkill
175 main.add_module('TrackCreator', recoTrackColName="MCRecoTracks", pdgCodes=[211, 321, 2212])
176 
177 
178 # build the name of the output file
179 outputFileName = outputDir + './'
180 if os.environ.get('USE_BEAST2_GEOMETRY'):
181  outputFileName += "SimEvts_Beast2"
182 else:
183  outputFileName += "SimEvts_Belle2"
184 outputFileName += '_' + str(rndseed) + '.root'
185 
186 # Root output. Default filename can be overriden with '-o' basf2 option.
187 rootOutput = register_module('RootOutput')
188 rootOutput.param('outputFileName', outputFileName)
189 # to save some space exclude everything except stuff needed for tracking
190 rootOutput.param('excludeBranchNames', ["ARICHAeroHits",
191  "ARICHDigits",
192  "ARICHSimHits",
193  "BKLMDigits",
194  "BKLMSimHitPositions",
195  "BKLMSimHits",
196  "CDCHits",
197  "CDCHits4Trg",
198  "CDCSimHits",
199  "CDCSimHitsToCDCHits4Trg",
200  "ECLDigits",
201  "ECLDsps",
202  "ECLHits",
203  "ECLSimHits",
204  "ECLTrigs",
205  "ECLDiodeHits",
206  "EKLMDigits",
207  "EKLMSimHits",
208  "TOPBarHits",
209  "TOPDigits",
210  "TOPRawDigits",
211  "TOPSimHits"
212  ])
213 main.add_module(rootOutput)
214 
215 log_to_file('createSim.log', append=False)
216 
217 print_path(main)
218 
219 process(main)
220 print(statistics)
tracking.validation.run.TrackingValidationRun.create_path
def create_path(self)
Definition: run.py:114