MadWeightTool: convert_lhe_lhco.py.txt

File convert_lhe_lhco.py.txt, 14.6 KB (added by anonymous, 7 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)