Belle II Software  release-05-01-25
ParticleGunFull.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 
15 
16 from basf2 import set_log_level, register_module, process, LogLevel, \
17  set_random_seed, print_params, create_path, statistics
18 from simulation import add_simulation
19 from reconstruction import add_reconstruction
20 
21 # suppress messages and warnings during processing:
22 set_log_level(LogLevel.WARNING)
23 
24 # to run the framework the used modules need to be registered
25 particlegun = register_module('ParticleGun')
26 
27 # ============================================================================
28 # Setting the random seed for particle generation
29 set_random_seed(123)
30 
31 # ============================================================================
32 # Setting the list of particle codes (PDG codes) for the generated particles
33 # the codes are given in an array, if only one code is used, the brackets
34 # should be kept:
35 # particlegun.param('pdgCodes', [11])
36 
37 # there is no limit on how many codes can be given the particle gun will select
38 # randomly amongst the PDGcodes using a uniform distribution if nTracks>0,
39 # otherwise there will be one particle for each code in the list default is
40 # [-11, 11]
41 particlegun.param('pdgCodes', [-11, 11])
42 
43 # ============================================================================
44 # Setting the number of tracks to be generated per event: this number can be
45 # any int>=0 default is 1
46 particlegun.param('nTracks', 10)
47 
48 # a value o 0 means that a track should be created for each entry in the
49 # pdgCodes list, e.g. the following two lines would create two electrons and
50 # one pion per event.
51 # particlegun.param('pdgCodes', [11,11,211])
52 # particlegun.param('nTracks', 0)
53 
54 # ============================================================================
55 # Varying the number of tracks per event can be achieved by setting varyNTacks
56 # to True. If so, the number of tracks will be randomized using possion
57 # distribution around the value of nTracks. Only valid if nTracks>0, default
58 # is False
59 particlegun.param('varyNTracks', False)
60 
61 # ============================================================================
62 # Each particle has a total momentum, a direction given as theta and phi and a
63 # vertex position in x, y and z. For each variable we need to specify the
64 # distribution and the parameters for the distribution. The number of
65 # parameters depends on the chosen distribution. All available distributions
66 # are listed here with required parameters are given in []. The distributions
67 # with suffix "Pt" can only be used for momentum generation and the ones with
68 # suffix "Cos" can only used for the generation of theta and phi.
69 #
70 # - fixed: always use the exact same value [value]
71 # - uniform: uniform distribution between min and max [min, max]
72 # - uniformPt: uniform distribution of transverse momentum between a given
73 # minimal and maximum value [min, max]
74 # - uniformCos: uniform distribution of the cosine of the angle between min
75 # and max of the absolute value [min, max]
76 # - normal: normal (gaussian) distribution around mean with width of sigma
77 # [mean, sigma]
78 # - normalPt: normal distribution of transverse momentum [mean, sigma]
79 # - normalCos: normal distribution of the cosine of the angle, not the
80 # absolute value [mean, sigma]
81 # - inversePt: generate momentum to obtain a flat distribution of the track
82 # curvature (inverse transverse momentum) between a minimal and
83 # maximal transverse momentum, [min_pt, max_pt]
84 # - polyline: create the momentum to follow an arbitrary distribution given
85 # as a list of x and y coordinates. All y coordinates must be
86 # non-negative and at leas one y coordinate must be positive
87 # [x1, x2, ..., xn, y1, y2, ..., yn]
88 # - polylinePt: same as polyline but for the transverse momentum, not the
89 # total momentum.
90 # [x1, x2, ..., xn, y1, y2, ..., yn]
91 # - polylineCos: same as polyline but for the cosine of the angle, not the
92 # absolute value
93 # [x1, x2, ..., xn, y1, y2, ..., yn]
94 # - discrete: a discrete spectrum consisting of possible values xi and their
95 # weights wi (useful e.g. for radioactive sources)
96 # [x1, x2, ..., xn, w1, w2, ..., wn]
97 
98 # ============================================================================
99 # Momentum generation
100 #
101 # The default is a uniform momentum distribution between 0.05 and 3 GeV
102 particlegun.param('momentumGeneration', 'uniform')
103 particlegun.param('momentumParams', [0.05, 3])
104 
105 # we could also generate a fixed momentum of 1 GeV
106 # particlegun.param('momentumGeneration', "fixed")
107 # particlegun.param('momentumParams', [1.0])
108 
109 # or we could generate a normal distributed transverse momentum around 2 GeV
110 # with a width of 0.5 GeV
111 # particlegun.param('momentumGeneration', "normalPt")
112 # particlegun.param('momentumParams', [2.0, 0.5])
113 
114 # to generate the momentum according to a cadmium 109
115 # source we could use
116 # particlegun.param('momentumGeneration', 'discreteSpectrum'),
117 # particlegun.param('momentumParams', [
118 # 22.1e-6, 25.0e-6, 88.0e-6, # photon energies
119 # 82.6, 14.7, 3.65, # weights
120 # ])
121 
122 # ============================================================================
123 # polar angle, theta
124 # The default is a uniform theta distribution between 17 and 150 degree
125 
126 particlegun.param('thetaGeneration', 'uniform')
127 particlegun.param('thetaParams', [17, 150])
128 
129 # We could also generate between 17 and 150 degrees with a flat distribution in
130 # cos(theta)
131 # particlegun.param('thetaGeneration', 'uniformCos')
132 # particlegun.param('thetaParams', [17, 150])
133 
134 # or we could create a theta angle between 17 and 150 degree where the
135 # cos(theta) distribution is flat
136 # particlegun.param('thetaGeneration', "normal")
137 # particlegun.param('thetaParams', [90,5])
138 
139 # Finally, we could use numpy to create events where the cos(theta)
140 # distribution follows a parabolic shape:
141 # import numpy as np
142 # # Create a set of 1000 x values spread uniformly over -1, 1
143 # x = np.linspace(-1,1,1000)
144 # # Create the corresponding distribution
145 # y = 1-x**2
146 # # Set the parameters by concatenating the x and y positions
147 # particlegun.param('thetaGeneration', "polylineCos")
148 # particlegun.param('thetaParams', list(x) + list(y))
149 
150 # ============================================================================
151 # azimuth angle, phi
152 # The default is a uniform theta distribution between 0 and 360 degree
153 
154 particlegun.param('phiGeneration', 'uniform')
155 particlegun.param('phiParams', [0, 360])
156 
157 # or we could create a normal distributed phi angle around 90 degrees with a
158 # width of 5 degrees
159 # particlegun.param('phiGeneration', "normal")
160 # particlegun.param('phiParams', [90,5])
161 
162 # ============================================================================
163 # Vertex generation
164 # The default is a fixed vertex at (0,0,0) for all tracks
165 particlegun.param('vertexGeneration', 'fixed')
166 particlegun.param('xVertexParams', [0])
167 particlegun.param('yVertexParams', [0])
168 particlegun.param('zVertexParams', [0])
169 
170 # We could also generate a normal distributed vertex with mean at (0,0,0) and
171 # width of 10µm, 60nm and 190 µm in x, y and z respectively
172 # particlegun.param('vertexGeneration', 'fixed')
173 # particlegun.param('xVertexParams', [0, 10e-4])
174 # particlegun.param('yVertexParams', [0, 60e-7])
175 # particlegun.param('zVertexParams', [0, 190e-4])
176 
177 # We can also specify a different distribution for any of the three axes, e.g.
178 # to make the z-vertex uniform between -10 and 10 cm and the xy vertex normal
179 # distributed around zero with a width of 0.1 cm we could use
180 # particlegun.param('vertexGeneration', 'normal')
181 # particlegun.param('xVertexParams', [0, 0.1])
182 # particlegun.param('yVertexParams', [0, 0.1])
183 # particlegun.param('zVertexGeneration', 'uniform')
184 # particlegun.param('zVertexParams', [-10, 10])
185 
186 # ============================================================================
187 # Setting independent vertices for each particle The default is to create one
188 # event vertex for all particles per event. By setting independentVertices to
189 # True, a new vertex will be created for each particle default is False
190 particlegun.param('independentVertices', False)
191 
192 # ============================================================================
193 # Print the parameters of the particle gun
194 print_params(particlegun)
195 
196 # ============================================================================
197 # Now lets create the necessary modules to perform a simulation
198 #
199 # Create Event information
200 eventinfosetter = register_module('EventInfoSetter')
201 # Show progress of processing
202 progress = register_module('Progress')
203 # Load parameters
204 gearbox = register_module('Gearbox')
205 # Create geometry
206 geometry = register_module('Geometry')
207 # Save output of simulation
208 output = register_module('RootOutput')
209 
210 # Setting the option for all non particle gun modules: want to process 100 MC
211 # events
212 eventinfosetter.param({'evtNumList': [100], 'runList': [1]})
213 
214 # Set output filename
215 output.param('outputFileName', 'ParticleGunOutput.root')
216 
217 # ============================================================================
218 main = create_path()
219 main.add_module(eventinfosetter)
220 main.add_module(progress)
221 main.add_module(particlegun)
222 
223 main.add_module(output)
224 
225 # Process events
226 process(main)
227 
228 # Print call statistics
229 print(statistics)