DevelopmentPage/BeyondProjects: WriteHelasTensor.py

File WriteHelasTensor.py, 11.6 KB (added by Will Link, 14 years ago)

Other version broken, this should work.

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','T':'double complex T'}
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 elif elem[0] =='T':
114 CallList.append('%s%d' %(elem[0],index+1))
115 DeclareList.append('%s%d(18)' %(DeclareDict[elem[0]],index+1))
116 else:
117 CallList.append('%s%d' %(elem[0],index+1))
118 DeclareList.append('%s%d(6)' %(DeclareDict[elem[0]],index+1))
119 else:
120 OnShell = 0
121 OffShellParticle = index
122 DeclareList.append('%s%d(6)' %(DeclareDict[elem[0]],index+1))
123 if (elem[0]=='V' or elem[0]=='S' or elem[0]=='T'):
124 MomentumConserve.append('-%s%d' %(elem[0],index+1))
125 elif elem[0] =='F' and Counter ==0%2:
126 MomentumConserve.append('-F%d'%(index+1))
127 Counter +=1
128 else:
129 MomentumConserve.append('+F%d' %(index+1))
130 Counter+=1
131 return {'CallList':CallList,'OnShell':OnShell,'DeclareList':DeclareList,'OffShell':OffShellParticle,'Momentum':MomentumConserve}
132 def define_header(self):
133 CollectedVariables = self.collect_variables()
134 print CollectedVariables
135 Momenta = CollectedVariables['momenta']
136 Width = CollectedVariables['width']
137 Mass = CollectedVariables['mass']
138 OverM = CollectedVariables['om']
139 listout = self.make_call_lists()
140 CallList = listout['CallList']
141 OnShell = listout['OnShell']
142 DeclareList = listout['DeclareList']
143 OffShellParticle = listout['OffShell']
144 MomentumConserve = listout['Momentum']
145 if OnShell:
146 String='subroutine '+self.namestring+'(C,'+','.join(CallList+Mass+Width)+',vertex)\n'
147 DeclareList.append('double complex vertex')
148 else:
149 DeclareList.append('double complex denom')
150 String = 'subroutine '+self.namestring+'(C,'+','.join(CallList+Mass+Width) \
151 +','+self.particles[OffShellParticle][0]+'%d)\n' %(OffShellParticle+1)
152 String = String + 'implicit none \n'
153 for elem in DeclareList:
154 String = String + elem +'\n'
155 if len(Mass+Width) > 0:
156 String = String+'double precision ' + ','.join(Mass+Width)+'\n'
157 if len(OverM) >0:
158 String = String +'double complex ' + ','.join(OverM)+'\n'
159 if len(Momenta) >0:
160 String = String+'double precision '+ '(0:3),'.join(Momenta)+'(0:3)\n'
161 if not OnShell:
162 NegSign = MomentumConserve.pop(OffShellParticle)
163 NegSignBool = re.match('-F',NegSign)
164 if NegSignBool:
165 NegString = '('
166 else:
167 NegString = '-('
168 # Implement better routine here!!
169 if re.match('S',self.particles[OffShellParticle][0]):
170 MomString = '%s%d(2)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
171 elif re.match('T',self.particles[OffShellParticle][0]):
172 MomString = '%s%d(17)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
173 else:
174 MomString = '%s%d(5)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
175 for elem in MomentumConserve:
176 if re.match('-S',elem):
177 MomString = MomString +elem+'(2)'
178 elif re.match('-T',elem):
179 MomString = MomString +elem+'(17)'
180 else:
181 MomString = MomString +elem+'(5)'
182 MomString = MomString +')\n'
183 if re.match('S',self.particles[OffShellParticle][0]):
184 MomString = MomString+'%s%d(3)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
185 elif re.match('T',self.particles[OffShellParticle][0]):
186 MomString = MomString+'%s%d(18)='%(self.particles[OffShellParticle][0],OffShellParticle+1)+NegString
187 else:
188 MomString = MomString+'%s%d(6)='%(self.particles[OffShellParticle][0],OffShellParticle+1) +NegString
189 for elem in MomentumConserve:
190 if re.match('-S',elem):
191 MomString = MomString +elem+'(3)'
192 elif re.match('-T',elem):
193 MomString = MomString +elem+'(18)'
194 else:
195 MomString = MomString +elem+'(6)'
196 MomString = MomString + ')\n'
197 String = String+MomString
198 for mom in Momenta:
199 index = int(mom[-1])
200 if re.match('S',self.particles[index-1][0]):
201 String = String+'%s(0) = dble(%s%d(2))\n'%(mom,self.particles[index-1][0],index)
202 String = String+'%s(1) = dble(%s%d(3))\n'%(mom,self.particles[index-1][0],index)
203 String = String+'%s(2) = dimag(%s%d(3))\n'%(mom,self.particles[index-1][0],index)
204 String = String+'%s(3) = dimag(%s%d(2))\n'%(mom,self.particles[index-1][0],index)
205 elif re.match('T',self.particles[index-1][0]):
206 String = String+'%s(0) = dble(%s%d(17))\n'%(mom,self.particles[index-1][0],index)
207 String = String+'%s(1) = dble(%s%d(18))\n'%(mom,self.particles[index-1][0],index)
208 String = String+'%s(2) = dimag(%s%d(18))\n'%(mom,self.particles[index-1][0],index)
209 String = String+'%s(3) = dimag(%s%d(17))\n'%(mom,self.particles[index-1][0],index)
210 elif (index-1)==OffShellParticle and re.match('V',self.particles[index-1][0]):
211 String = String+'%s(0) = -dble(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
212 String = String+'%s(1) = -dble(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
213 String = String+'%s(2) = -dimag(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
214 String = String+'%s(3) = -dimag(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
215 elif (index-1)==OffShellParticle and re.match('T',self.particles[index-1][0]):
216 String = String+'%s(0) = -dble(%s%d(17))\n'%(mom,self.particles[index-1][0],index)
217 String = String+'%s(1) = -dble(%s%d(18))\n'%(mom,self.particles[index-1][0],index)
218 String = String+'%s(2) = -dimag(%s%d(18))\n'%(mom,self.particles[index-1][0],index)
219 String = String+'%s(3) = -dimag(%s%d(17))\n'%(mom,self.particles[index-1][0],index)
220 else:
221 String = String+'%s(0) = dble(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
222 String = String+'%s(1) = dble(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
223 String = String+'%s(2) = dimag(%s%d(6))\n'%(mom,self.particles[index-1][0],index)
224 String = String+'%s(3) = dimag(%s%d(5))\n'%(mom,self.particles[index-1][0],index)
225 for elem in OverM:
226 index = int(elem[-1])
227 String = String+'om%d = 0d0\n' %(index)
228 String = String + 'if (m%d .ne. 0d0) om%d'%(index,index)+'=1d0/dcmplx(m%d**2,-w%d*m%d)\n'%(index,index,index)
229 return String
230
231 def define_expression(self):
232 OnShell = 1
233 OutString =''
234 for index,elem in enumerate(self.particles):
235 if not elem[1]:
236 OnShell = 0
237 OffShellParticle = '%s%d'%(elem[0],index+1)
238 if OnShell:
239 for ind in self.obj.listindices():
240 string = 'Vertex = C*'+ self.write_obj(self.obj.get_rep(ind))
241 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
242 string = re.sub('chi','f',string)
243 string = re.sub('\+-','-',string)
244 string = re.sub('\((?P<num>[+-]*[0-9])(?P<num2>[+-][0-9])[Jj]\)\.','(\g<num>d0,\g<num2>d0)',string)
245 string = re.sub('(?P<num>[0-9])[Jj]\.','\g<num>.*(0d0,1d0)',string)
246 OutString = OutString + string + '\n'
247 else:
248 numerator = self.obj.numerator
249 denominator = self.obj.denominator
250 for ind in denominator.listindices():
251 denom = self.write_obj(denominator.get_rep(ind))
252 string = 'denom ='+'1d0/('+denom+')'
253 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
254 string = re.sub('\+-','-',string)
255 string = re.sub('\((?P<num>[+-]*[0-9])\+(?P<num2>[+-][0-9])[Jj]\)\.','(\g<num>d0,\g<num2>d0)',string)
256 string = re.sub('(?P<num>[0-9])[Jj]\.','\g<num>*(0d0,1d0)',string)
257 OutString = OutString + string + '\n'
258 counter = 1
259 for ind in numerator.listindices():
260 string = '%s(%d)= C*denom*' %(OffShellParticle,counter)
261 string += self.write_obj(numerator.get_rep(ind))
262 string = re.sub('\+-','-',string)
263 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
264 string = re.sub('\((?P<num>[+-]*[0-9])(?P<num2>[+-][0-9])[Jj]\)\.','(\g<num>d0,\g<num2>d0)',string)
265 string = re.sub('(?P<num>[0-9])[Jj]\.','\g<num>*(0d0,1d0)',string)
266 OutString = OutString +string+'\n'
267 counter +=1
268 return OutString
269
270 def define_foot(self):
271 return 'end'
272
273 def write(self):
274 writer = FortranWriter.FortranWriter()
275 writer.downcase = False
276 commentstring ='c This File Automatically generated by MadGraph 5/FeynRules HELAS writer \n'
277 commentstring +='c The process calculated in this file is: \n'
278 commentstring = commentstring + 'c '+self.comment+'\n'
279 head = self.define_header()
280 body = self.define_expression()
281 foot = self.define_foot()
282 out = commentstring+head+body+foot
283 for lines in out.split('\n'):
284 writer.write_fortran_line(self.out,lines)
285
286if __name__ == '__main__':
287 # Input as coming from FR!!!
288 obj = Helas.Mass(1,1)*Helas.P(1,1)
289# test = Mass(1)+Mass(2)
290# test = test.simplify().expand()
291 # Analysis
292# obj = obj.simplify() # -> Simplify sum
293 # expand
294# print obj.simplify().expand()
295 List = [('F',1),('V',1),('F',1)]
296 Write = HelasWriterForFortran(obj.expand(),List,'test')
297 Write.write()
298# for ind in test.listindices():
299# print Write.write_obj(test.get_rep(ind))
300# print test.get_rep(ind).prefactor
301# print Write.write_obj(test)
302 print 'all is done'
303# print low_level
304 #
305