DevelopmentPage/BeyondProjects: WriteHelas.py

File WriteHelas.py, 9.9 KB (added by Olivier Mattelaer, 14 years ago)
Line 
1################################################################################
2#
3# Copyright (c) 2010 The MadGraph Development team and Contributors
4#
5# This file is a part of the MadGraph 5 project, an application which
6# automatically generates Feynman diagrams and matrix elements for arbitrary
7# high-energy processes in the Standard Model and beyond.
8#
9# It is subject to the MadGraph license which should accompany this
10# distribution.
11#
12# For more information, please visit: http://madgraph.phys.ucl.ac.be
13#
14"""This authorizes to write the Helas Routine in whatever languages.
15
16"""
17import madgraph.iolibs.export_v4 as FortranWriter
18import helasamp_object as Helas
19import helas_amp_lib as Helas_Lib
20import re
21
22class WriteHelas:
23 """ generic function if needed """
24
25 def __init__(self, object, particlelist, out_path):
26 """ store basic information """
27 self.obj = object
28 self.out = open(out_path+'.f','w')
29 self.particles = particlelist
30 self.namestring = out_path
31
32 def collect_variables(self):
33 """ Collects Momenta,Mass,Widths for Declares """
34 MomentaList = []
35 MassList = []
36 WidthList = []
37 for elem in self.obj.tag:
38 if elem.startswith('p'):
39 MomentaList.append(elem)
40 elif elem.startswith('m'):
41 MassList.append(elem)
42 elif elem.startswith ('w'):
43 WidthList.append(elem)
44 return {'momenta':MomentaList,'width':WidthList,'mass':MassList}
45
46
47 def write(self):
48 """ launch the writing routine """
49 self.define_header()
50 self.define_content()
51 self.define_footer()
52
53 def define_header(self):
54 """How to define the header of the file """
55 #perhaps something as variable_list = self.collect_variable()
56 #This should probably done language in language specific subclass?
57 def define_content(self):
58 """basic way to write the output"""
59# for ind in self.obj.listindices():
60
61 #be carefull for scalar expression -> no ind are outputting ?
62 #perhpas need to separate lorentz/spin indices?
63
64 def define_footer(self):
65 """ How to define the footer of the file """
66 # Also language specific?
67 def write_indices_part(self, indices, obj):
68 """How to write this part"""
69
70 #Example
71 text = 'output(%s) =' % indices
72 text += self.write_obj(obj)
73 print text
74 def write_obj(self,obj):
75 """Schedulling which routine returns the correct string representation"""
76 if isinstance(obj, Helas_Lib.MultVariable):
77 return self.write_obj_Mult(obj)
78 elif isinstance(obj,Helas_Lib.AddVariable):
79 return self.write_obj_Add(obj)
80 elif isinstance(obj,Helas_Lib.Variable):
81 return self.write_obj_Var(obj)
82# elif isinstance(obj,FracVariable):
83# return self.write_obj_Frac(obj)
84 else:
85 return str(obj)
86# def write_obj_Frac(self,obj):
87# denominator = self.write_obj(obj.denominator)
88# denominator = '1/'+denom
89# numerator = self.write_obj(obj.numerator)
90# return (denominator,numerator)
91
92 def write_obj_Mult(self,obj):
93 """ basic example """
94 mult_list = [self.write_obj(factor) for factor in obj]
95 print mult_list
96 text='('
97 if obj.prefactor != 1:
98 text += str(obj.prefactor)+ '*'
99 return text + '*'.join(mult_list) + ')'
100
101
102 def write_obj_Add(self,obj):
103 """ basic example """
104 mult_list = [self.write_obj(factor) for factor in obj]
105 text='('
106 if obj.prefactor != 1:
107 text = str(obj.prefactor)+'*'+text
108 return text + '+'.join(mult_list)+')'
109
110 def write_obj_Var(self,obj):
111 text =''
112 if obj.prefactor != 1:
113 text += str(obj.prefactor)+'*'
114 text += obj.variable
115 return text
116
117
118
119class HelasWriterForFortran(WriteHelas):
120
121 def define_header(self):
122 CollectedVariables = self.collect_variables()
123 writer = FortranWriter.FortranWriter()
124 # Set lowercase/uppercase Fortran code
125 writer.downcase = False
126
127 DeclareDict = {'F':'double complex f','V':'double complex V'}
128 # Note -> talk to Olivier about producing a list of masses, widths,momenta involved.
129 # Set up the lists needed for information.
130 Momenta = CollectedVariables['momenta']
131 Width = CollectedVariables['width']
132 Mass = CollectedVariables['mass']
133
134 CallList = []
135 DeclareList =['double precision C']
136 MomentumConserve = []
137 OnShell = 1
138 Counter = 0
139 for index,elem in enumerate(self.particles):
140 if elem[1]:
141 CallList.append('%s%d' %(elem[0],index+1))
142 DeclareList.append('%s%d(6)' %(DeclareDict[elem[0]],index+1))
143 else:
144 OnShell = 0
145 OffShellParticle = index
146 DeclareList.append('%s%d(6)' %(DeclareDict[elem[0]],index+1))
147 if elem[0] == 'V':
148 MomentumConserve.append('V%d' %(index+1))
149 # because fermions come in twos, we know which is out, which is in.
150 elif elem[0] == 'F' and Counter == 0%2:
151 MomentumConserve.append('F%d' %(index +1))
152 Counter +=1
153 else:
154 MomentumConserve.append('-F%d' %(index+1))
155 Counter += 1
156 # Now that we've set up lists, we can begin writing out pieces
157 if OnShell:
158 FirstString = 'subroutine '+self.namestring+'(C,'+','.join(CallList+Mass+Width)+',vertex)'
159 writer.write_fortran_line(self.out,FirstString)
160 writer.write_fortran_line(self.out,'implicit none')
161 DeclareList.append('double complex vertex')
162 else:
163 FirstString = 'subroutine '+self.namestring+'(C,'+','.join(CallList+Mass+Width)+','+self.particles[OffShellParticle][0]+'%d)' %(OffShellParticle+1)
164 writer.write_fortran_line(self.out,FirstString)
165 writer.write_fortran_line(self.out,'implicit none')
166 DeclareList.append('double complex %s%d(6)'%(self.particles[OffShellParticle][0],OffShellParticle+1))
167 for elem in DeclareList:
168 writer.write_fortran_line(self.out,elem)
169 # Declare widths,mass,momenta
170 if len(Mass+Width) > 0:
171 string = 'double precision '+','.join(Mass+Width)
172 writer.write_fortran_line(self.out,string)
173 if len(Momenta) >0:
174 string = 'double precision ' +'(0:3),'.join(Momenta)+'(0:3)'
175 writer.write_fortran_line(self.out,string)
176 if not OnShell:
177 NegSign = MomentumConserve.pop(OffShellParticle)
178 NegSign = re.match('-',NegSign)
179 if NegSign:
180 Front = '-('
181 else:
182 Front = '('
183 # Broken for Scalars, might have to rethink
184 String = '%s%d(5) ='%(self.particles[OffShellParticle][0],OffShellParticle+1) +Front + '(5)'.join(MomentumConserve)+'(5))'
185 writer.write_fortran_line(self.out,String)
186 String = '%s%d(6) = '%(self.particles[OffShellParticle][0],OffShellParticle+1)+Front + '(6)'.join(MomentumConserve)+'(6))'
187 writer.write_fortran_line(self.out,String)
188 for mom in Momenta:
189 index = int(mom[-1])
190 string = '%s(0) = dbl(%s%d(5))'%(mom,self.particles[index-1][0],index)
191 writer.write_fortran_line(self.out,string)
192 string = '%s(1) = dbl(%s%d(6))'%(mom,self.particles[index-1][0],index)
193 writer.write_fortran_line(self.out,string)
194 string = '%s(2) = dimag(%s%d(6))'%(mom,self.particles[index-1][0],index)
195 writer.write_fortran_line(self.out,string)
196 string = '%s(3) = dimag(%s%d(5))'%(mom,self.particles[index-1][0],index)
197 writer.write_fortran_line(self.out,string)
198
199
200
201
202 def define_expression(self):
203 writer = FortranWriter.FortranWriter()
204 # Set lowercase/uppercase Fortran code
205 writer.downcase = False
206 OnShell = 1
207 for index,elem in enumerate(self.particles):
208 if not elem[1]:
209 OnShell = 0
210 OffShellParticle = '%s%d'%(elem[0],index+1)
211 break
212 if OnShell:
213 for ind in self.obj.listindices():
214 string = 'Out = C*'+ self.write_obj(self.obj.get_rep(ind))
215 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
216 string = re.sub('chi','f',string)
217 string = re.sub('(?P<num>[0-9])[Jj]','\g<num>*(0,1)',string)
218 writer.write_fortran_line(self.out,string)
219 else:
220 counter = 1
221 numerator = self.obj.numerator
222 denominator = self.obj.denominator
223 for ind in denominator.listindices():
224 denom = self.write_obj(denominator.get_rep(ind))
225 for ind in numerator.listindices():
226 if counter == 1:
227 string = 'denom ='+'1/('+denom+')'
228 writer.write_fortran_line(self.out,string)
229 string = '%s(%d)= C*denom*' %(OffShellParticle,counter)
230 string += self.write_obj(numerator.get_rep(ind))
231 string = re.sub('\_(?P<num>[0-9])','(\g<num>)',string)
232 string = re.sub('chi','f',string)
233 string = re.sub('(?P<num>[0-9])[Jj]','\g<num>*(0,1)',string)
234 writer.write_fortran_line(self.out,string)
235 counter +=1
236
237 def define_foot(self):
238 writer = FortranWriter.FortranWriter()
239 writer.downcase = False
240 string = 'end'
241 writer.write_fortran_line(self.out,string)
242
243 def write(self):
244 self.define_header()
245 self.define_expression()
246 self.define_foot()
247
248if __name__ == '__main__':
249 # Input as coming from FR!!!
250 obj = Helas.Mass(1)/Helas.Width(1)
251# test = Mass(1)+Mass(2)
252# test = test.simplify().expand()
253 # Analysis
254# obj = obj.simplify() # -> Simplify sum
255 # expand
256# print obj.simplify().expand()
257 List = [('F',0),('V',1),('F',1)]
258 Write = HelasWriterForFortran(obj.expand(),List,'Test')
259 Write.write()
260# for ind in test.listindices():
261# print Write.write_obj(test.get_rep(ind))
262# print test.get_rep(ind).prefactor
263# print Write.write_obj(test)
264 print 'all is done'
265# print low_level
266 #
267