Fork me on GitHub

source: git/validation/flatGunLHEventProducer.py@ c749957

ImprovedOutputFile Timing llp
Last change on this file since c749957 was 6038170, checked in by Michele Selvaggi (michele.selvaggi@…>, 7 years ago

updated validation scripts

  • Property mode set to 100644
File size: 9.7 KB
Line 
1#!/usr/bin/env python
2#################################################################################################
3#
4# Requires python v2.7
5#
6# Simple script for generation of a LHE file containing events distributed with a flat (log) e, pt and eta spectrum.
7#
8# For help type:
9# python flatGunLHEventProducer.py -h
10#
11#################################################################################################
12import os, random, math, argparse, sys
13
14#__________________________________________________________
15def print_params(args):
16
17 print 'sqrt(s) : ', args.ecm
18 print 'pdg codes : ', args.pdg
19 print 'gun mode : ', args.guntype
20 print 'pmax : ', args.pmax
21 print 'pmin : ', args.pmin
22 print 'etamin : ', args.etamin
23 print 'etamax : ', args.etamax
24 print 'nevts : ', args.nevts
25 print 'Seed : ', args.seed
26 print 'flat log : ', args.log
27 print 'output : ', args.output
28#__________________________________________________________
29def write_init(args):
30
31 ebeam = args.ecm/2.
32
33 out = open(args.output, "a")
34 out.write('<LesHouchesEvents version="3.0">\n')
35 out.write('<init>\n')
36 out.write('2212 2212 {0} {1} 0 0 1 1 -4 1\n'.format(ebeam, ebeam))
37 out.write('1000. 1. 1000. 1\n')
38 out.write('</init>\n')
39 out.close()
40#__________________________________________________________
41def write_event(args, pt, eta, phi):
42
43 # define particle masses
44 mass = dict()
45 mass[1] = 0.
46 mass[2] = 0.
47 mass[3] = 0.
48 mass[4] = 1.29
49 mass[5] = 4.7
50 mass[6] = 173.
51 mass[11] = 0.
52 mass[12] = 0.
53 mass[13] = 0.
54 mass[14] = 0.
55 mass[15] = 0.
56 mass[16] = 0.
57 mass[21] = 0.
58 mass[22] = 0.
59 mass[23] = 9.118800e+01
60 mass[24] = 80.419002
61 mass[25] = 1.250000e+02
62
63 def alpha_s(q):
64 return 12*math.pi/((33.-6.)*math.log(q**2/0.3**2))
65
66 # randomly pick a pdg code from supplied list
67 pdg = random.choice(args.pdg)
68
69 # compute particles 4-vectors: px, py, pz, e
70
71 px = pt*math.cos(phi)
72 py = pt*math.sin(phi)
73 pz = pt*math.sinh(eta)
74 e = math.sqrt(mass[pdg]**2 + px**2 + py**2 + pz**2)
75
76 p1 = [0., 0., e, e]
77 p2 = [0., 0., -e, e]
78 p3 = [px, py, pz, e]
79 p4 = [-px, -py, -pz, e]
80
81 cf_list = []
82
83 # initialize stuff
84 color = [501, 501, 0, 0]
85
86 out = open(args.output, "a")
87 out.write('<event>\n')
88 out.write(' 4 1 +1000. {:.8e} 7.54677100e-03 {:.8e}\n'.format(2*e, alpha_s(e)))
89
90
91 ########## g g -> g g ###################
92
93 if pdg == 21:
94
95 cf_list.append([503, 501, 504, 502, 503, 502, 504, 501])
96 cf_list.append([504, 501, 503, 502, 503, 501, 504, 502])
97
98 color = cf_list[random.randint(0, 1)]
99
100 out.write(' 21 -1 0 0 {} {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 1.0000e+00\n'.format(color[0], color[1], p1[0], p1[1], p1[2], p1[3], 0.))
101 out.write(' 21 -1 0 0 {} {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 1.0000e+00\n'.format(color[2], color[3], p2[0], p2[1], p2[2], p2[3], 0.))
102 out.write(' 21 1 1 2 {} {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 1.0000e+00\n'.format(color[4], color[5], p3[0], p3[1], p3[2], p3[3], 0.))
103 out.write(' 21 1 1 2 {} {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 1.0000e+00\n'.format(color[6], color[7], p4[0], p4[1], p4[2], p4[3], 0.))
104
105 ########## q q -> q' q' ###############
106
107 elif pdg > 0 and pdg < 7:
108
109 cf_list.append([501, 501, 502, 502])
110 cf_list.append([501, 502, 501, 502])
111
112 color = cf_list[random.randint(0, 1)]
113
114 out.write(' 1 -1 0 0 {} 0 {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 1.0000e+00\n'.format(color[0], p1[0], p1[1], p1[2], p1[3], 0.))
115 out.write(' -1 -1 0 0 0 {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 -1.0000e+00\n'.format(color[1], p2[0], p2[1], p2[2], p2[3], 0.))
116 out.write(' {} 1 1 2 {} 0 {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} {:.8e} 1.0000e+00\n'.format(pdg, color[2], p3[0], p3[1], p3[2], p3[3], mass[pdg], 0.))
117 out.write(' -{} 1 1 2 0 {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} {:.8e} -1.0000e+00\n'.format(pdg, color[3], p4[0], p4[1], p4[2], p4[3], mass[pdg], 0.))
118
119 ########## q q -> l l (l = e, mu, tau, ve, vm , vtau) ###############
120
121
122 elif pdg > 10 and pdg < 17:
123
124 out.write(' 1 -1 0 0 {} 0 {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 1.0000e+00\n'.format(color[0], p1[0], p1[1], p1[2], p1[3], 0.))
125 out.write(' -1 -1 0 0 0 {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 -1.0000e+00\n'.format(color[1], p2[0], p2[1], p2[2], p2[3], 0.))
126 out.write(' {} 1 1 2 {} 0 {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} {:.8e} -1.0000e+00\n'.format(pdg, color[2], p3[0], p3[1], p3[2], p3[3], mass[pdg], 0.))
127 out.write(' -{} 1 1 2 0 {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} {:.8e} 1.0000e+00\n'.format(pdg, color[3], p4[0], p4[1], p4[2], p4[3], mass[pdg], 0.))
128
129 ########## q q -> B B (B = gamma, Z, W, H) ###############
130
131 elif pdg > 21 and pdg < 26:
132
133 out.write(' 1 -1 0 0 {} 0 {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 1.0000e+00\n'.format(color[0], p1[0], p1[1], p1[2], p1[3], 0.))
134 out.write(' -1 -1 0 0 0 {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} 0.0000e+00 -1.0000e+00\n'.format(color[1], p2[0], p2[1], p2[2], p2[3], 0.))
135 out.write(' {} 1 1 2 {} 0 {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} {:.8e} 0.0000e+00\n'.format(pdg, color[2], p3[0], p3[1], p3[2], p3[3], mass[pdg], 0.))
136 if pdg == 24:
137 out.write(' -{} 1 1 2 0 {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} {:.8e} 0.0000e+00\n'.format(pdg, color[3], p4[0], p4[1], p4[2], p4[3], mass[pdg], 0.))
138 else:
139 out.write(' {} 1 1 2 0 {} {:+.8e} {:+.8e} {:+.8e} {:.8e} {:.8e} {:.8e} 0.0000e+00\n'.format(pdg, color[3], p4[0], p4[1], p4[2], p4[3], mass[pdg], 0.))
140
141 out.write('</event>\n')
142
143
144
145
146#__________________________________________________________
147
148if __name__=="__main__":
149
150 parser = argparse.ArgumentParser()
151 parser.add_argument("--pdg", nargs='+', type=int)
152 parser.add_argument('--guntype', dest='guntype', help='pt or e gun. The parameters pmin(max) are then interpreted as ptmin(max) or emin(max) depending on the specified gunmode', default='pt')
153 parser.add_argument("--pmin", type=float, help="minimum pt/e [GeV] (default: 1.)", default=1.)
154 parser.add_argument("--pmax", type=float, help="maximum pt/e [GeV] (default: 50000.)", default=50000.)
155 parser.add_argument("--etamin", type=float, help="minimum eta (default: -2.5)", default=-6.)
156 parser.add_argument("--etamax", type=float, help="maximum eta (default: 2.5)", default=6.)
157 parser.add_argument("--ecm,", dest='ecm', type=float, help="center of mass energy (default: 13000)", default=13000)
158 parser.add_argument("--nevts", type=int, help="number of events to generate (default: 1000)", default=1000)
159 parser.add_argument("--seed", type=int, help="random seed (default uses cpu time)", default=None)
160 parser.add_argument('--log', dest='log', help="flat in log pt (default: yes)", action='store_true')
161 parser.add_argument('--nolog', dest='log', help="flat in pt (default: false)", action='store_false')
162 parser.set_defaults(log=True)
163 parser.add_argument("--output", help="output LHE file (default: events.lhe)", default='events.lhe')
164
165 args = parser.parse_args()
166
167 # check if provided pdgCode is allowed
168 allowed_pdgCodes = [1,2,3,4,5,6,11,12,13,14,15,16,21,22,23,24,25]
169 for p1 in args.pdg:
170 if p1 not in allowed_pdgCodes:
171 print args.pdg ,'Please provide a list of supported pdgCodes : 1, 2, 3, 4, 5 or 21'
172 sys.exit(0)
173
174 # print user-defined parameters
175 print_params(args)
176 print ''
177
178 # intialize file and write LHE file header
179 out = open(args.output, "w+")
180 out.close()
181 write_init(args)
182
183 # initialize random seed
184 random.seed(args.seed)
185
186 print 'Start event generation ...'
187
188 ebeam = args.ecm/2.
189 count = 0
190
191 if args.guntype not in ['pt', 'e']:
192 print 'Please specify either gun mode "pt" or "e".'
193 sys.exit(0)
194
195 # start event loop
196 nfail = 0
197 while count < args.nevts:
198
199 phi = random.uniform(0., math.pi)
200 eta = random.uniform(args.etamin, args.etamax)
201
202 # flat in e/pt or in log e/pt
203 if args.log:
204 gun = math.pow(10, random.uniform(math.log10(args.pmin), math.log10(args.pmax)))
205 else:
206 gun = random.uniform(args.pmin, args.pmax)
207
208 # generating "balanced" collision, i.e x1 = x2 = 2*energy/sqrt(s)
209 if args.guntype == 'pt':
210 pt = gun
211 e = pt*math.cosh(eta)
212 elif args.guntype == 'e':
213 e = gun
214 pt = e/math.cosh(eta)
215
216 # write event corresponding to required process
217 if e <= ebeam:
218 write_event(args, pt, eta, phi)
219 count += 1
220 if (count+1)%500 == 0:
221 print ' ... processed {} events ...'.format(count+1)
222 else:
223 nfail += 1
224
225 if nfail > 10*args.nevts:
226 print 'Too many events fail phase space requirements. Usually means ptmax is too high given eta cuts.'
227 print 'To fix this either increase center of mass/ change cuts.'
228 out.close()
229 os.system('rm {}'.format(args.output))
230 sys.exit(0)
231
232 out = open(args.output, "a")
233 out.write('</LesHouchesEvents>\n')
234 out.close()
235
236 os.system('gzip {}'.format(args.output))
237
238 print ''
239 print 'Event generation completed.'
240 print 'Output file:'
241 print '{}'.format(os.path.abspath(args.output +'.gz'))
242
Note: See TracBrowser for help on using the repository browser.