incomp-flame-post/code/code_gen/code_gen.py
ignis e9524d645b
All checks were successful
CI Test Suite / run-tests (push) Successful in 1m12s
Deploy Documentation / build-and-deploy (push) Successful in 52s
Consolidate build info module generation into code_gen.py and delete build_info_gen.py
2026-06-02 23:51:59 +00:00

257 lines
7.6 KiB
Python

import os
import sys
import argparse
import datetime
import subprocess as sp
from lark import Lark, Token
from jinja2 import Template
from post import Stage1, Stage2, Stage3, Stage4
from resources.m_template import mod_form
parser = argparse.ArgumentParser()
parser.add_argument("term_file", help="name of file containing postprocessing term spec")
parser.add_argument("-b", "--build-info", action="store_true", help="print build_info instead of code")
parser.add_argument("-x", "--print-latex", action="store_true", help="print latex equation instead of code")
parser.add_argument("-p", "--print-report", action="store_true", help="print code generation report instead of code")
parser.add_argument("-m", "--make-build-module", action="store_true", help="print build_info Fortran module instead of code")
args = parser.parse_args()
calc_grammar = """
?varlist: "[" [NAME ("," NAME)*] "]"
?start: statement*
?statement: NAME [ attr_list ] "=" sum -> assign_var
| avg "{" [NAME ("," NAME)*] "}" -> assign_avg_var
| varlist
attr_list: "(" [attr_pair ("," attr_pair)*] ")"
?attr_pair: NAME "=" BOOL
| NAME "=" INT
| NAME "=" ESCAPED_STRING
?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 ")" -> paren
| inlinefunc "(" sum ")" -> icall
| mathfunc "(" sum ("," 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
| "$" NAME -> udf
?derivative: "ddx" -> ddx
| "d2dx" -> d2dx
| "ddy" -> ddy
| "d2dy" -> d2dy
| "ddz" -> ddz
| "d2dz" -> d2dz
%import common.CNAME -> NAME
%import common.NUMBER
%import common.ESCAPED_STRING
%import common.INT
%import common.WS
BOOL: "true" | "false"
COMMENT: /#.*/
%ignore COMMENT
%ignore WS
"""
def tok_to_bool(tok):
"Convert the value of `tok` from string to bool, while maintaining line number & column."
# tok.type == 'BOOL'
return Token.new_borrow_pos(tok.type, tok.value == "true", tok)
def tok_to_int(tok):
"Convert the value of `tok` from string to int, while maintaining line number & column."
# tok.type == 'INT'
return Token.new_borrow_pos(tok.type, int(tok), tok)
def tok_to_str(tok):
"Convert the value of `tok` from string to string, while maintaining line number & column."
# tok.type == 'ESCAPED_STRING'
return Token.new_borrow_pos(tok.type, tok.value.strip('"'), tok)
class VersionInfo(object):
def __init__ (self, input_raw):
bd = 'build date: {}'
bb = 'build base: {}'
self.build_date = bd.format(datetime.datetime.today().ctime())
try:
self.build_base = bb.format(sp.check_output(["git", "rev-parse", "HEAD"]).strip())
except sp.CalledProcessError:
self.build_base = "None"
try:
sp.check_output("git diff --exit-code".split())
sp.check_output("git diff --cached --exit-code".split())
self.off_base = ""
except sp.CalledProcessError:
self.off_base = "built with changes not commited"
tinfo = '''
================================================================================
************************************ terms *************************************
--------------------------------------------------------------------------------
{}
================================================================================
'''
self.term_info = tinfo.format(input_raw)
def fparams (self):
return "\n".join([self.build_date, self.build_base, self.off_base, self.term_info])
def test(terms_raw, report=False, latex=False):
"""Compiles the post-processing term specifications from DSL into targets.
This compiles the Lark AST through all 4 pipeline stages (Stage 1 to 4)
and outputs the resulting Fortran source code via Jinja2, a LaTeX equation
dictionary, or a code generation IR report.
Args:
terms_raw (str): The raw string contents of the DSL terms specification input file.
report (bool): If True, generates and saves the intermediate code generation report.
latex (bool): If True, prints LaTeX equations of the averaged terms to stdout.
"""
parser = Lark(calc_grammar,
parser='lalr',
lexer_callbacks =
{
'ESCAPED_STRING': tok_to_str,
'INT': tok_to_int,
'BOOL': tok_to_bool
}
)
tree = parser.parse(terms_raw)
ir1 = Stage1(tree)
ir2 = Stage2(ir1)
ir3 = Stage3(ir2)
ir4 = Stage4(ir3)
cdict = ir4.print_program()
if report:
ir4.save_ir()
elif latex:
print("{")
for avg in ir4.averaged.values():
print('"{}" : '.format(avg.name))
# print r"\begin{equation}"
print('r"${}$" ,'.format(avg.latex))
# print r"\end{equation}"
print("}")
else:
print(Template(mod_form).render(**cdict))
def escape_fortran_string(s):
"""Escapes string content for standard Fortran source formatting."""
escaped = s.replace('"', '""')
lines = escaped.splitlines()
if not lines:
return '""'
formatted_lines = [f'"{line}"' for line in lines]
return ' // char(10) // &\n '.join(formatted_lines)
def generate_build_info_module(terms_raw):
"""Generates a complete standard Fortran module containing build info and LaTeX equations."""
vinfo = VersionInfo(terms_raw)
build_info_str = vinfo.fparams()
parser = Lark(calc_grammar,
parser='lalr',
lexer_callbacks={
'ESCAPED_STRING': tok_to_str,
'INT': tok_to_int,
'BOOL': tok_to_bool
})
tree = parser.parse(terms_raw)
ir1 = Stage1(tree)
ir2 = Stage2(ir1)
ir3 = Stage3(ir2)
ir4 = Stage4(ir3)
latex_lines = ["{"]
for avg in ir4.averaged.values():
latex_lines.append(f' "{avg.name}" : r"${avg.latex}$",')
latex_lines.append("}")
latex_str = "\n".join(latex_lines)
fortran_build_info = escape_fortran_string(build_info_str)
fortran_latex = escape_fortran_string(latex_str)
fortran_code = f"""module m_build_info
implicit none
character(len=*), parameter :: build_info_str = &
{fortran_build_info}
character(len=*), parameter :: latex_equations_str = &
{fortran_latex}
end module m_build_info
"""
print(fortran_code)
def build_info(terms_raw):
"""Generates and prints build information for the current code generation run.
Args:
terms_raw (str): The raw string contents of the DSL terms specification input file.
"""
vinfo = VersionInfo(terms_raw)
print(vinfo.fparams())
if __name__ == '__main__':
with open(args.term_file) as inputfile:
terms_raw = ((inputfile.read()))
if args.make_build_module:
generate_build_info_module(terms_raw)
elif args.build_info:
build_info(terms_raw)
else:
test(terms_raw, report=args.print_report, latex=args.print_latex)