# # This example shows how to write a basic calculator with variables. # from lark import Lark, Visitor, Transformer, v_args, Token try: input = raw_input # For Python2 compatibility except NameError: pass calc_grammar = """ ?varlist: "[" [NAME ("," NAME)*] "]" ?start: statement* ?statement: NAME "=" sum -> assign_var | avg "{" [NAME ("," NAME)*] "}" -> assign_avg_var | varlist ?sum: product | sum "+" product -> add | sum "-" product -> sub ?product: atom | product "*" atom -> mul | product "/" atom -> div ?atom: NUMBER -> number | "-" atom -> neg | NAME -> var | NAME "'" -> fluc | "$" NAME -> env | "(" sum ")" | inlinefunc "(" sum ")" -> icall | mathfunc "(" sum ")" -> fcall | derivative "(" NAME ")" -> dnx avg: "avg" [NAME] ?inlinefunc: "sqr" -> sqr | "pow3" -> pow3 ?mathfunc: "log" -> log | "exp" -> exp | "sqrt" -> sqrt | "abs" -> abs | "rxn_rate" -> rxn_rate ?derivative: "ddx" -> ddx | "d2dx" -> d2dx | "ddy" -> ddy | "d2dy" -> d2dy | "ddz" -> ddz | "d2dz" -> d2dz %import common.CNAME -> NAME %import common.NUMBER %import common.WS %ignore WS """ real_array_decl = "real*8, allocatable, dimension(:,:,:) :: {}" real_array_alloc = "allocate({0}(nxp,nyp,nzp), stat=ierr) ; {0} = 0." avg_array_decl = "real*8, allocatable, dimension(:) :: {}" avg_array_alloc = "allocate({0}(nxp), stat=ierr) ; {0} = 0." real_array_free = "deallocate({})" real_array_loop = """ do k = 1, nzp do j = 1, nyp do i = 1, nxp {0[0]}(i,j,k) = {0[1]} end do end do end do """ avg_array_sum = """ do k = 1, nzp do j = 1, nyp do i = 1, nxp {0}(i) = {0}(i) + {1} {2} end do end do end do """ avg_array_write = ''' integer :: i open (200, file="qEdge_X.dat") write (200,*) output_header do i=1,nxp write (200,'({0}e20.10)') real(i)*hxp, {1} end do close (200) ''' avg_array_divide = """ call MPI_ALLREDUCE(MPI_IN_PLACE, {0}, nxp, MPI_REAL8, MPI_SUM, MPI_COMM_TASK, mpi_err) {0} = {0} / denum {1} """ real_array_diff = "call {0[0]} ( {0[0]}_{0[1]}, {0[1]} )" class FortranCode: def __init__ (self, exp): self.exp = exp def __repr__ (self): return self.exp ''' def __add__ (self, other): self.exp = "( {} + {} )".format(self.exp, other.exp) return self def __sub__ (self, other): self.exp = "( {} - {} )".format(self.exp, other.exp) return self ''' @v_args(inline=True) # Affects the signatures of the methods class ToFortran(Transformer): def __init__(self, primary_set): self.primary = primary_set self.derivatives = {} self.dependency = {} self.codes = {} def number(self, numeral): return (FortranCode(str(float(numeral))), []) def env(self, name): return (FortranCode(name.value), []) def var(self, name): return (FortranCode(name + "(i,j,k)"), [name.value] if name.value not in self.primary else []) def fluc(self, name): fmt = "({0}(i,j,k) - {{0}}avg_{0}(i))" return (FortranCode(fmt.format(name)), [name.value] if name.value not in self.primary else []) def dnx (self, partial, b): signature = "{}_{}".format(partial.data, b) fcode = FortranCode(signature + "(i,j,k)") self.derivatives[signature] = (partial.data, b.value) self.dependency[signature] = [b.value] self.codes[signature] = { "decl" : real_array_decl.format(signature), "alloc" : real_array_alloc.format(signature), "free" : real_array_free.format(signature), "calc" : real_array_diff.format((partial.data, b)), "avg" : "", } return (fcode, [signature]) def icall (self, a, (b, dep)): fcode = "({0})".format(b) if a.data == "sqr": fcode = "(({0})*({0}))".format(b) elif a.data == "pow3": fcode = "(({0})*({0})*({0}))".format(b) return (fcode, dep) def fcall (self, a, (b, dep)): fcode = "( {} ( {} ) )".format(a, b) return (fcode, dep) def neg(self, (b, dep)): fcode = "( - {} )".format(b) return (fcode, dep) def add(self, (a, adep), (b, bdep)): fcode = "( {} + {} )".format(a, b) return (fcode, adep + bdep) def sub(self, (a, adep), (b, bdep)): fcode = "( {} - {} )".format(a, b) return (fcode, adep + bdep) def mul(self, (a, adep), (b, bdep)): fcode = "( {} * {} )".format(a, b) return (fcode, adep + bdep) def div(self, (a, adep), (b, bdep)): fcode = "( {} / {} )".format(a, b) return (fcode, adep + bdep) log = lambda self : "log" exp = lambda self : "exp" sqrt = lambda self : "sqrt" abs = lambda self : "dabs" rxn_rate = lambda self : "rxn_rate" class CheckPass(Visitor): def __init__(self): self.hasFluc = False @classmethod def check(cls, tree): self = cls() return self(tree) def __call__(self, tree): self.visit(tree) return self.hasFluc def fluc(self, tree): self.hasFluc = True @v_args(inline=True) # Affects the signatures of the methods class CalculateTree(Transformer): def __init__(self): self.primary = [] self.derived = {} self.averaged = {} self.averages = [] self.derivatives = {} self.dependency = {} self.fluctuation = {} self.definitions = {} self.exp_parser = ToFortran([]) self.codes = {} def varlist(self, *args): for arg in args: self.primary.append(arg.value) self.dependency[arg.value] = [] self.fluctuation[arg.value] = False return "" def assign_var(self, *args): # name, (value, dep)): vname, vdef = args self.definitions[vname.value] = vdef code, dep = self.exp_parser.transform(vdef) self.dependency[vname.value] = dep self.fluctuation[vname.value] = "{0}" in code self.codes[vname.value] = { "decl" : real_array_decl.format(vname.value), "alloc" : real_array_alloc.format(vname.value), "free" : real_array_free.format(vname.value), "calc" : real_array_loop.format((vname.value, code)), "avg" : "", } return "" def assign_avg_var(self, *args): # weight, *args): #name, (value, dep)): weight = args[0] vlist = args[1:] self.averaged[str(weight)] = map(str, vlist) w = str(weight) for v in vlist: avg_var = ( "" if w == str(None) else w + "_" ) + "avg_" + v self.averages.append(avg_var) self.fluctuation[avg_var] = False var = str(v) if self.has_fluc()[str(v)]: if w == str(None): var = var + "_" self.dependency[avg_var] = [str(v)+"_"] self.dependency[str(v)+"_"] = self.dependency[str(v)] self.fluctuation[str(v)+"_"] = True self.codes[var] = {} for k,val in self.codes[str(v)].items(): self.codes[var][k] = val.replace(str(v), var) self.codes[var]["calc"] = self.codes[var]["calc"].format("") else: var = var + "_" + w self.dependency[avg_var] = [str(v)+"_"+w, w] self.dependency[str(v)+"_"+w] = self.dependency[str(v)] self.fluctuation[str(v)+"_"+w] = True self.codes[var] = {} for k,val in self.codes[str(v)].items(): self.codes[var][k] = val.replace(str(v), var) self.codes[var]["calc"] = self.codes[var]["calc"].format(w+"_") else: self.dependency[avg_var] = [var] + ( [] if w == str(None) else [w] ) wfmt = "* {}(i,j,k)" pWeight = (wfmt.format(w) if w != str(None) else "") meanw = "/ avg_{}" dWeight = (meanw.format(w) if w != str(None) else "") self.codes[avg_var] = { "decl" : avg_array_decl.format(avg_var), "alloc" : avg_array_alloc.format(avg_var), "free" : real_array_free.format(avg_var), "calc" : avg_array_sum.format(avg_var, var+"(i,j,k)", pWeight), "avg" : avg_array_divide.format(avg_var, dWeight) } return "" def avg(self, *args): try: return args[0] except IndexError: return None def dep_graph (self): return dict( self.dependency.items() + self.exp_parser.dependency.items() ) def has_fluc (self): return dict( self.fluctuation.items() + [(k, False) for k, v in self.exp_parser.dependency.items()] ) import sys import pprint pp = pprint.PrettyPrinter() class FortranProgram: def __init__ (self, terms_input): self.tree_parser = Lark(calc_grammar, parser='lalr' ) self.parser = CalculateTree() tree = self.tree_parser.parse(terms_input) self.parser.transform(tree) dg = self.parser.dep_graph() fd = self.parser.has_fluc() def isFluc (a): for x in dg[a]: fd[a] = fd[a] or isFluc(x) return fd[a] average_names = self.parser.averages self.pass1avg = filter(lambda x: not isFluc(x), average_names) self.pass2avg = filter(isFluc, average_names) pass1var = filter(lambda x: not isFluc(x), average_names) pass2var = filter(isFluc, average_names) def dep_set (varset): c = set([]) for var in varset: c.update(dg[var]) return c def dep_closer (s): c = set(s) while len(dep_set(s)) > 0: s = dep_set(s) c.update(s) return c self.pass1set = dep_closer(set(pass1var)) - set(self.parser.primary) self.pass2set = dep_closer(set(pass2var)) - set(self.parser.primary) self.pass1list = self.sort_vars(dg, self.pass1set) self.pass2list = self.sort_vars(dg, self.pass2set) self.codes = dict(self.parser.codes.items() + self.parser.exp_parser.codes.items()) def sort_vars (self, dependency, group): order = [] remain = set(group) while len(remain) > 0: for v in remain: if len(set(dependency[v]) & remain) == 0: order.append(v) remain.remove(v) break return order def print_program (self): import StringIO output = StringIO.StringIO() output.write('First line.\n') print >>output, 'Second line.' decl = StringIO.StringIO() alloc = StringIO.StringIO() free = StringIO.StringIO() calc1 = StringIO.StringIO() avg1 = StringIO.StringIO() calc2 = StringIO.StringIO() avg2 = StringIO.StringIO() wres = StringIO.StringIO() hfmt = 'character (len = *), parameter :: output_header="{}"' print >>decl, hfmt.format(" ".join(["x"] + self.parser.averages)) avgarr = "{}(i)" print >>wres, avg_array_write.format( len(self.parser.averages)+1, ", ".join(map(avgarr.format, self.parser.averages)) ) for var in self.pass1set | self.pass2set: print >>decl, self.codes[var]["decl"] for var in self.pass1set | self.pass2set: print >>alloc, self.codes[var]["alloc"] for var in self.pass1set | self.pass2set: print >>free, self.codes[var]["free"] for var in self.pass1list : print >>calc1, self.codes[var]["calc"] for var in self.pass1avg : print >>avg1, self.codes[var]["avg"] for var in self.pass2list: print >>calc2, self.codes[var]["calc"] for var in self.pass2avg: print >>avg2, self.codes[var]["avg"] md = {} md["module_name"] = "terms" md["module_data"] = decl.getvalue() md["module_init"] = alloc.getvalue() md["module_finalize"] = free.getvalue() md["module_pass1"] = calc1.getvalue() md["module_pass1_avg"] = avg1.getvalue() md["module_pass2"] = calc2.getvalue() md["module_pass2_avg"] = avg2.getvalue() md["module_write_result"] = wres.getvalue() return md def main(): while True: try: s = input('> ') except EOFError: break print(calc(s)) def test(): with open("resources/m_template.f90") as template_file: mod_form = template_file.read() with open("terms.input") as inputfile: terms_raw = ((inputfile.read())) fp = FortranProgram(terms_raw) print mod_form.format( fp.print_program() ) if __name__ == '__main__': test()