MadWeightTool: convert_lhe_lhco.py.txt

File convert_lhe_lhco.py.txt, 14.6 KB (added by (none), 12 years ago)
Line 
1#!/usr/bin/env python
2##########################################################################
3## ##
4## MadWeight ##
5## --------- ##
6##########################################################################
7## ##
8## author: Mattelaer Olivier (CP3) ##
9## email: olivier.mattelaer@uclouvain.be ##
10## ##
11##########################################################################
12## ##
13## license: GNU ##
14## last-modif:08/07/08 ##
15## version 1.1 ##
16## ##
17##########################################################################
18## ##
19## README: ##
20## input: a regular lhe file ##
21## output: a lhco file ##
22## ##
23## how to launch the program: ##
24## $> python ./convert_lhe_lhco input_file output_file ##
25## (output file we overwritted) ##
26## or ##
27## $> python ./convert_lhe_lhc and follows instruction ##
28##########################################################################
29## ##
30## Content ##
31## ------- ##
32## ##
33## + Particle ##
34## | + init ##
35## | + def_pid ##
36## | + cal_mass ##
37## | + cal_pt ##
38## | + cal_phi ##
39## | + cal_eta ##
40## | | + check_def ##
41## | + lhco_line ##
42## + part_quadvec -> Particle ##
43## | + init ##
44## | + add ##
45## + Pass_from_lhe_to_lhco ##
46## | + init ##
47## | + find_begin_blok ##
48## | + find_end_blok ##
49## | + put_in_output ##
50## | | + define_particle ##
51## | | + write_output ##
52## ##
53## ##
54##########################################################################
55import sys
56import math
57import re
58##########################################################################
59## START CODE
60##########################################################################
61
62
63#1 ########################################################################
64class Particle:
65 """particle information"""
66
67 #2 ########################################################################
68 def __init__(self):
69 """start a particle"""
70
71
72
73 self.pid=''
74 self.E=''
75 self.px=''
76 self.py=''
77 self.pz=''
78 self.pt=''
79 self.phi=''
80 self.y=''
81 self.mass=''
82
83 #2 ########################################################################
84 def def_pid(self,pid):
85 """ define pid of the particle """
86 self.pid=int(pid)
87
88 #2 ########################################################################
89 def cal_mass(self):
90 """ define mass from quadri impulsion """
91
92 if not self.check_def(['E','px','py','pz']):
93 sys.exit('Particle error: Quadri impulsion not define (error for mass routine)')
94
95
96
97 if self.E**2-self.px**2-self.py**2-self.pz**2>1e-7: #precision problem
98 self.mass=math.sqrt(self.E**2-self.px**2-self.py**2-self.pz**2)
99 else:
100 self.mass=0
101
102 #2 ########################################################################
103 def def_mass(self,mass):
104 """ put the value for the mass """
105
106 self.mass=float(mass)
107
108
109 #2 ########################################################################
110 def cal_pt(self):
111 """define pt from quadri impulsion"""
112
113 if not self.check_def(['E','px','py','pz']):
114 sys.exit('Particle error: Quadri impulsion not define (error for pt routine)')
115
116 self.pt =math.sqrt(self.px**2+self.py**2)
117
118 #2 ########################################################################
119 def cal_phi(self):
120 """define phi from quadri impulsion"""
121
122 if not self.check_def(['E','px','py','pz']):
123 sys.exit('Particle error: Quadri impulsion not define (error for phi routine)')
124
125 if(self.px>0):
126 self.phi=math.atan(self.py/self.px)
127 elif(self.px<0):
128 self.phi=math.atan(self.py/self.px)+math.pi
129 elif(self.py>0): #remind that p(1)=0
130 self.phi=math.pi/2.0
131 elif(self.py<0): # remind that p(1)=0
132 self.phi=-math.pi/2.0
133 else:
134 print "Warning self.phi not properly defined put value to 0"
135 self.phi=0
136
137 if(self.phi<0):
138 self.phi=self.phi+2*math.pi
139
140 return self.phi
141 #2 ########################################################################
142 def cal_eta(self):
143 """define y from quadri impulsion"""
144
145 if not self.check_def(['E','px','py','pz']):
146 sys.exit('Particle error: Quadri impulsion not define (error for eta routine)')
147
148 theta=math.acos(self.pz/math.sqrt(self.px**2+self.py**2+self.pz**2))
149 self.eta=-math.log(math.tan(theta/2.0))
150
151 #3 ########################################################################
152 def check_def(self,data_list):
153 """ check if everything is defined """
154
155 if type(data_list)!=list:
156 data_list=[data_list]
157
158 for value in data_list:
159 if type(eval('self.'+value))==str:
160 print "failed for", value
161 return 0
162 return 1
163
164 #2 ########################################################################
165 def lhco_line(self):
166 """ define the line for the lhco line (without the initial tag)
167 this is a first simple, trully naive, routine.
168 """
169 if not self.check_def(['eta','phi','pt','mass','pid']):
170 sys.exit('Particle error: some attribute not defined')
171
172 jet=[1,2,3,4,5,6,21]
173 inv_list=[12,14,16,18,1000022,1000023,1000024,1000025,1000035]
174
175 #define pid-> type
176 pid_to_type={11:1,-11:1,13:2,-13:2,15:3,-15:3,22:0}
177 for data in jet:
178 pid_to_type[data]=4
179 pid_to_type[-data]=4
180 for data in inv_list:
181 pid_to_type[data]=6
182 pid_to_type[-data]=6
183
184
185
186 type=''
187 for key in pid_to_type.keys():
188 if self.pid==key:
189 type=pid_to_type[key]
190 break
191
192 if type=='':
193 print 'Warning unknown type'
194 return ''
195
196 text =' '+str(type) #type LHCO
197 text+=' '+str(self.eta) #ETA
198 text+=' '+str(self.phi) #PHI
199 text+=' '+str(self.pt) #PT
200 text+=' '+str(self.mass) #JMASS
201 if self.pid in [11,13]: #NTRK
202 text+=' -1'
203 else:
204 text+=' 1'
205 if self.pid in [-5,5]: #BTAG
206 text+=' 2'
207 else:
208 text+=' 0'
209 text+=' 0' #HAD/EM
210 text+=' 0' #DUMMY 1
211 text+=' 0' #DUMMY 2
212
213 return text
214
215#1 ########################################################################
216class part_quadvec(Particle):
217 """ particle define from quadri impulsion """
218
219 #2 ########################################################################
220 def __init__(self,E,px,py,pz):
221 """ define from a quadri vector """
222 Particle.__init__(self)
223 self.E=float(E)
224 self.px=float(px)
225 self.py=float(py)
226 self.pz=float(pz)
227 self.cal_pt()
228 self.cal_phi()
229 self.cal_eta()
230 #self.cal_mass()
231 print self.E,self.px,self.py,self.pz
232 print self.pt,self.phi,self.eta
233
234 #2 ########################################################################
235 def add(self,particle):
236 """ add two quadri vector """
237
238 if not self.check_def(['E','px','py','pz']):
239 sys.exit('Particle error: Quadri impulsion not define')
240 if not particle.check_def(['E','px','py','pz']):
241 sys.exit('Particle error: Quadri impulsion not define')
242
243 neut=part_quadvec(self.E+particle.E,self.px+particle.px,self.py+particle.py,self.pz+particle.pz)
244 neut.cal_mass()
245 return neut
246
247
248#1 ########################################################################
249class Pass_from_lhe_to_lhco:
250
251
252 #2 ########################################################################
253 def __init__(self,input,output):
254
255 self.still_to_read=1
256 self.nb_data=0
257 self.input=open(input,'r')
258 self.output=open(output,'w')
259 while 1:
260 self.find_begin_blok()
261 text=self.find_end_blok()
262 if text:
263 self.put_in_output(text)
264 else:
265 break
266
267 #2 ########################################################################
268 def find_begin_blok(self):
269
270 tag=re.compile('''<event>''',re.I)
271 while 1:
272 line=self.input.readline()
273 if line=='':
274 self.still_to_read=0
275 return
276 if tag.search(line):
277 return
278
279 #2 ########################################################################
280 def find_end_blok(self):
281
282 tag=re.compile(r'''</event>''')
283 text=[]
284
285 while 1:
286 line=self.input.readline()
287 if line=='':
288 self.still_to_read=0
289 return
290 text.append(line)
291 if tag.search(line):
292 return text
293
294 #2 ########################################################################
295 def put_in_output(self,text):
296 """ analyse the event and write the lhco """
297
298 inv_list=[12,14,16,18,1000022,1000023,1000024,1000025,1000035]
299 inv_list+=[-12,-14,-16,-18,-1000022,-1000023,-1000024,-1000025,-1000035]
300 particle_content=[]
301 neut=0
302 for line in text:
303 particle=self.define_particle(line)
304 if particle==0:
305 continue
306
307 if particle.pid in inv_list:
308 if neut==0:
309 neut=particle
310 else:
311 neut=neut.add(particle)
312 else:
313 particle_content.append(particle)
314
315 if neut:
316 neut.pid=12
317 particle_content.append(neut)
318
319 self.write_output(particle_content)
320
321 #3 ########################################################################
322 def define_particle(self,line):
323 """ define the particle from the lhe line """
324
325 pattern=re.compile(r'''^\s*
326 (?P<pid>-?\d+)\s+ #PID
327 (?P<status>1)\s+ #status (1 for output particle)
328 (?P<mother>-?\d+)\s+ #mother
329 (?P=mother)\s+ #mother
330 (?P<color1>[+-e.\d]*)\s+ #color1
331 (?P<color2>[+-e.\d]*)\s+ #color2
332 (?P<px>[+-e.\d]*)\s+ #px
333 (?P<py>[+-e.\d]*)\s+ #py
334 (?P<pz>[+-e.\d]*)\s+ #pz
335 (?P<E>[+-e.\d]*)\s+ #E
336 (?P<mass>[+-e.\d]*)\s+ #mass
337 (?P<dum1>[+-e.\d]*)\s+ #dummy1
338 (?P<dum2>[+-e.\d]*)\s* #dummy2
339 $ #end of string
340 ''',66) #verbose+ignore case
341
342 if pattern.search(line):
343 obj=pattern.search(line)
344 E=obj.group('E')
345 px=obj.group('px')
346 py=obj.group('py')
347 pz=obj.group('pz')
348 particle=part_quadvec(E,px,py,pz)
349 particle.def_mass(obj.group('mass'))
350 particle.def_pid(obj.group('pid'))
351 return particle
352 else:
353 return 0
354
355
356
357 #3 ########################################################################
358 def write_output(self,content):
359 """ write the input from the content info """
360 text="""# typ eta phi pt jmass ntrk btag had/em dummy dummy\n"""
361 self.output.writelines(text)
362 text="0 "+str(self.nb_data)+" "+str(len(content))+"\n"
363 self.output.writelines(text)
364
365 i=1
366 for particle in content:
367 text=str(i)+' '+particle.lhco_line()+'\n'
368 self.output.writelines(text)
369 i+=1
370
371
372
373##########################################################################
374## END CODE
375##########################################################################
376if __name__=='__main__':
377 opt=sys.argv
378 if len(opt)!=3:
379 input=raw_input('enter the position of the input lhe file')
380 output=raw_input('enter the position of the output lhco file')
381 else:
382 input=opt[1]
383 output=opt[2]
384
385 Pass_from_lhe_to_lhco(input,output)