1 | import FortranWriter as FortranWriter
|
---|
2 | import helasamp_object as Helas
|
---|
3 | import helasamp_lib as Helas_Lib
|
---|
4 | import re
|
---|
5 |
|
---|
6 | class 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 |
|
---|
97 | class 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 |
|
---|
286 | if __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 |
|
---|