DevelopmentPage/BeyondProjects: NewWriteHelas.py

File NewWriteHelas.py, 10.1 KB (added by Will Link, 14 years ago)
Line 
1import FortranWriter as FortranWriter
2import helasamp_object as Helas
3import helasamp_lib as Helas_Lib
4import re
5
6class WriteHelas:
7 """ Generic writing functions """
8
9 def __init__(self,object,particlelist,out_path,comment):
10 self.obj = object
11 self.out = open(out_path+'.f','w')
12 self.particles = particlelist
13 self.namestring = out_path
14 self.comment = comment
15 def collect_variables(self):
16 """Collects Momenta,Mass,Width into lists"""
17 MomentaList = set()
18 MassList = set()
19 WidthList = set()
20 OverMList = set()
21 for elem in self.obj.tag:
22 if elem.startswith('P'):
23 MomentaList.add(elem)
24 elif elem.startswith('M'):
25 MassList.add(elem)
26 elif elem.startswith('W'):
27 WidthList.add(elem)
28 elif elem.startswith('O'):
29 OverMList.add(elem)
30
31 MomentaList = list(MomentaList)
32 MassList = list(MassList)
33 WidthList = list(WidthList)
34 OverMList = list(OverMList)
35 return {'momenta':MomentaList,'width':WidthList,'mass':MassList,'om':OverMList}
36
37 def define_header(self):
38 """ Prototype for language specific header"""
39 pass
40 def define_content(self):
41 """Prototype for language specific body"""
42 pass
43 def define_foote (self):
44 """Prototype for language specific footer"""
45 pass
46
47 def write_indices_part(self,indices,obj):
48 """Routine for making a string out of indice objects"""
49 text = 'output(%s)' %indices
50 return text
51
52 def write_obj(self,obj):
53 """Calls the appropriate writing routine"""
54 if isinstance(obj,Helas_Lib.MultVariable):
55 return self.write_obj_Mult(obj)
56 elif isinstance(obj,Helas_Lib.AddVariable):
57 return self.write_obj_Add(obj)
58 elif isinstance(obj,Helas_Lib.Variable):
59 return self.write_obj_Var(obj)
60 else:
61 return str(obj)
62
63 def write_obj_Mult(self,obj):
64 """Turn a multvariable into a string"""
65 mult_list = [self.write_obj(factor) for factor in obj]
66 text = '('
67 if obj.prefactor != 1:
68 if obj.prefactor != -1:
69 text = str(obj.prefactor)+'.'+'*'+text
70 else:
71 text = '-'+text
72 return text+'*'.join(mult_list)+')'
73
74 def write_obj_Add(self,obj):
75 """Turns add variable into a string"""
76 mult_list = [self.write_obj(factor) for factor in obj]
77 text = '('
78 if obj.prefactor != 1:
79 if obj.prefactor != -1:
80 text = str(obj.prefactor)+'.'+'*'+text
81 else:
82 text = '-'+text
83 return text+ '+'.join(mult_list)+')'
84
85 def write_obj_Var(self,obj):
86 text = ''
87 if obj.prefactor !=1:
88 if obj.prefactor != -1:
89 text = str(obj.prefactor)+'.'+'*'+text
90 else:
91 text = '-'+text
92 text +=obj.variable
93 if obj.power !=1:
94 text = text+'**'+str(obj.power)
95 return text
96
97class HelasWriterForFortran(WriteHelas):
98 """routines for writing out Fortran"""
99
100 def make_call_lists(self):
101 CallList = []
102 DeclareList = ['double precision C']
103 MomentumConserve = []
104 DeclareDict = {'F':'double complex f','V':'double complex V','S':'double complex s'}
105 OnShell = 1
106 Counter = 0
107 OffShellParticle = 999
108 for index,elem in enumerate(self.particles):
109 if elem[1]:
110 if elem[0] =='S':
111 CallList.append('%s%d' %(elem[0],index+1))
112 DeclareList.append('%s%d(3)' %(DeclareDict[elem[0]],index+1))
113 else:
114 CallList.append('%s%d' %(elem[0],index+1))
115 DeclareList.append('%s%d(6)' %(DeclareDict[elem[0]],index+1))
116 else:
117 OnShell = 0
118 OffShellParticle = index
119 DeclareList.append('%s%d(6)' %(DeclareDict[elem[0]],index+1))
120 if (elem[0]=='V' or elem[0]=='S'):
121 MomentumConserve.append('-%s%d' %(elem[0],index+1))
122 elif elem[0] =='F' and Counter ==0%2:
123 MomentumConserve.append('-F%d'%(index+1))
124 Counter +=1
125 else:
126 MomentumConserve.append('+F%d' %(index+1))
127 Counter+=1
128 return {'CallList':CallList,'OnShell':OnShell,'DeclareList':DeclareList,'OffShell':OffShellParticle,'Momentum':MomentumConserve}
129 def define_header(self):
130 CollectedVariables = self.collect_variables()
131 print CollectedVariables
132 Momenta = CollectedVariables['momenta']
133 Width = CollectedVariables['width']
134 Mass = CollectedVariables['mass']
135 OverM = CollectedVariables['om']
136 listout = self.make_call_lists()
137 CallList = listout['CallList']
138 OnShell = listout['OnShell']
139 DeclareList = listout['DeclareList']
140 OffShellParticle = listout['OffShell']
141 MomentumConserve = listout['Momentum']
142 if OnShell:
143 String='subroutine '+self.namestring+'(C,'+','.join(CallList+Mass+Width)+',vertex)\n'
144 DeclareList.append('double complex vertex')
145 else:
146 DeclareList.append('double complex denom')
147 String = 'subroutine '+self.namestring+'(C,'+','.join(CallList+Mass+Width) \
148 +','+self.particles[OffShellParticle][0]+'%d)\n' %(OffShellParticle+1)
149 String = String + 'implicit none \n'
150 for elem in DeclareList:
151 String = String + elem +'\n'
152 if len(Mass+Width) > 0:
153 String = String+'double precision ' + ','.join(Mass+Width)+'\n'
154 if len(OverM) >0:
155 String = String +'double complex ' + ','.join(OverM)+'\n'
156 if len(Momenta) >0:
157 String = String+'double precision '+ '(0:3),'.join(Momenta)+'(0:3)\n'
158 if not OnShell:
159 NegSign = MomentumConserve.pop(OffShellParticle)
160 NegSignBool = re.match('-F',NegSign)
161 if NegSignBool:
162 NegString = '('
163 else:
164 NegString = '-('
165 if re.match('S',self.particles[OffShellParticle][0]):
166 MomString = '%s%d(2)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
167 else:
168 MomString = '%s%d(5)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
169 for elem in MomentumConserve:
170 if re.match('-S',elem):
171 MomString = MomString +elem+'(2)'
172 else:
173 MomString = MomString +elem+'(5)'
174 MomString = MomString +')\n'
175 if re.match('S',self.particles[OffShellParticle][0]):
176 MomString = MomString+'%s%d(3)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
177 else:
178 MomString = MomString+'%s%d(6)='%(self.particles[OffShellParticle][0],OffShellParticle+1) +NegString
179 for elem in MomentumConserve:
180 if re.match('-S',elem):
181 MomString = MomString +elem+'(3)'
182 else:
183 MomString = MomString +elem+'(6)'
184 MomString = MomString + ')\n'
185 String = String+MomString
186 for mom in Momenta:
187 index = int(mom[-1])
188 if re.match('S',self.particles[index-1][0]):
189 String = String+'%s(0) = dble(%s%d(2))\n'%(mom,self.particles[index-1][0],index)
190 String = String+'%s(1) = dble(%s%d(3))\n'%(mom,self.particles[index-1][0],index)
191 String = String+'%s(2) = dimag(%s%d(3))\n'%(mom,self.particles[index-1][0],index)
192 String = String+'%s(3) = dimag(%s%d(2))\n'%(mom,self.particles[index-1][0],index)
193 elif (index-1)==OffShellParticle and re.match('V',self.particles[index-1][0]):
194 String = String+'%s(0) = -dble(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
195 String = String+'%s(1) = -dble(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
196 String = String+'%s(2) = -dimag(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
197 String = String+'%s(3) = -dimag(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
198 else:
199 String = String+'%s(0) = dble(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
200 String = String+'%s(1) = dble(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
201 String = String+'%s(2) = dimag(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
202 String = String+'%s(3) = dimag(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
203 for elem in OverM:
204 index = int(elem[-1])
205 String = String+'om%d = 0d0\n' %(index)
206 String = String + 'if (m%d .ne. 0d0) om%d'%(index,index)+'=1d0/dcmplx(m%d**2,-w%d*m%d)\n'%(index,index,index)
207 return String
208
209 def define_expression(self):
210 OnShell = 1
211 OutString =''
212 for index,elem in enumerate(self.particles):
213 if not elem[1]:
214 OnShell = 0
215 OffShellParticle = '%s%d'%(elem[0],index+1)
216 if OnShell:
217 for ind in self.obj.listindices():
218 string = 'Vertex = C*'+ self.write_obj(self.obj.get_rep(ind))
219 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
220 string = re.sub('chi','f',string)
221 string = re.sub('\+-','-',string)
222 string = re.sub('\((?P<num>[+-]*[0-9])(?P<num2>[+-][0-9])[Jj]\)\.','(\g<num>d0,\g<num2>d0)',string)
223 string = re.sub('(?P<num>[0-9])[Jj]\.','\g<num>.*(0d0,1d0)',string)
224 OutString = OutString + string + '\n'
225 else:
226 numerator = self.obj.numerator
227 denominator = self.obj.denominator
228 for ind in denominator.listindices():
229 denom = self.write_obj(denominator.get_rep(ind))
230 string = 'denom ='+'1d0/('+denom+')'
231 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
232 string = re.sub('\+-','-',string)
233 string = re.sub('\((?P<num>[+-]*[0-9])\+(?P<num2>[+-][0-9])[Jj]\)\.','(\g<num>d0,\g<num2>d0)',string)
234 string = re.sub('(?P<num>[0-9])[Jj]\.','\g<num>*(0d0,1d0)',string)
235 OutString = OutString + string + '\n'
236 counter = 1
237 for ind in numerator.listindices():
238 string = '%s(%d)= C*denom*' %(OffShellParticle,counter)
239 string += self.write_obj(numerator.get_rep(ind))
240 string = re.sub('\+-','-',string)
241 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
242 string = re.sub('\((?P<num>[+-]*[0-9])(?P<num2>[+-][0-9])[Jj]\)\.','(\g<num>d0,\g<num2>d0)',string)
243 string = re.sub('(?P<num>[0-9])[Jj]\.','\g<num>*(0d0,1d0)',string)
244 OutString = OutString +string+'\n'
245 counter +=1
246 return OutString
247
248 def define_foot(self):
249 return 'end'
250
251 def write(self):
252 writer = FortranWriter.FortranWriter()
253 writer.downcase = False
254 commentstring ='c This File Automatically generated by MadGraph 5/FeynRules HELAS writer \n'
255 commentstring +='c The process calculated in this file is: \n'
256 commentstring = commentstring + 'c '+self.comment+'\n'
257 head = self.define_header()
258 body = self.define_expression()
259 foot = self.define_foot()
260 out = commentstring+head+body+foot
261 for lines in out.split('\n'):
262 writer.write_fortran_line(self.out,lines)
263
264if __name__ == '__main__':
265 # Input as coming from FR!!!
266 obj = Helas.Mass(1,1)*Helas.P(1,1)
267# test = Mass(1)+Mass(2)
268# test = test.simplify().expand()
269 # Analysis
270# obj = obj.simplify() # -> Simplify sum
271 # expand
272# print obj.simplify().expand()
273 List = [('F',1),('V',1),('F',1)]
274 Write = HelasWriterForFortran(obj.expand(),List,'test')
275 Write.write()
276# for ind in test.listindices():
277# print Write.write_obj(test.get_rep(ind))
278# print test.get_rep(ind).prefactor
279# print Write.write_obj(test)
280 print 'all is done'
281# print low_level
282 #
283