def test_time_shift():
from dolang.symbolic import parse_string
e = parse_string('sin(a(1) + b + a(0) + f(-(1)) + f(4) + a(1))')
to_source(e)
enes = stringify(e, variables=['a', 'f'])
print(to_source(enes))
assert (
to_source(enes) == "sin(a__1_ + b_ + a__0_ + f_m1_ + f__4_ + a__1_)")
def make_method_from_factory(fff: FlatFunctionFactory,
vectorize=True,
use_file=False,
debug=False):
mod = compile_factory(fff)
if debug:
print(to_source(mod))
if vectorize:
coredims = [len(v) for k, v in fff.arguments.items()]
signature = str.join(',', ['(n_{})'.format(d) for d in coredims])
n_out = len(fff.content)
if n_out in coredims:
signature += '->(n_{})'.format(n_out)
# ftylist = float64[:](*([float64[:]] * len(coredims)))
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature += ',(n_{})'.format(n_out)
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature = None
fun = eval_ast(mod)
from numba import jit, guvectorize
jitted = jit(fun, nopython=True)
if vectorize:
gufun = guvectorize([fty], signature, target='parallel',
nopython=True)(fun)
return jitted, gufun
else:
return jitted
def make_method_from_factory(fff: FlatFunctionFactory, vectorize=True, use_file=False, debug=False):
mod = compile_factory(fff)
if debug:
print(to_source(mod))
if vectorize:
coredims = [len(v) for k, v in fff.arguments.items()]
signature = str.join(',', ['(n_{})'.format(d) for d in coredims])
n_out = len(fff.content)
if n_out in coredims:
signature += '->(n_{})'.format(n_out)
# ftylist = float64[:](*([float64[:]] * len(coredims)))
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature += ',(n_{})'.format(n_out)
fty = "void(*[float64[:]]*{})".format(len(coredims) + 1)
else:
signature = None
fun = eval_ast(mod)
from numba import jit, guvectorize
jitted = jit(fun, nopython=True)
if vectorize:
gufun = guvectorize(
[fty], signature, target='parallel', nopython=True)(fun)
return jitted, gufun
else:
return jitted
def wrapper(*args, **kwds):
if not isinstance(args[0], str):
return f(*args, **kwds)
else:
a = parse_string(args[0])
nargs = tuple([a]) + args[1:]
res = f(*nargs, **kwds)
return to_source(res)
def wrapper(*args, **kwds):
if not isinstance(args[0], str):
return f(*args, **kwds)
else:
a = parse_string(args[0])
nargs = tuple([a]) + args[1:]
res = f(*nargs, **kwds)
return to_source(res)
def test_make_method_list_of_lists():
from dolang.parser import parse_string
from dolang.codegen import to_source
from dolang.function_compiler import make_method
#style 1: list of lists
eqs = ['x + lam*y']
arguments = [[('x', 1)], [('x', 0), ('y', 0)]]
constants = ['lam']
fun = make_method(eqs, arguments, constants)
print(to_source(fun))
def substitute_preamble(ff: FlatFunctionFactory):
import dolang
import copy
pr = copy.copy(ff.preamble)
for k in pr.keys():
pr[k] = dolang.parse_string(pr[k]).value
st = SubsTransformer(pr)
dd = copy.copy(ff.content)
for k in dd.keys():
eq = dolang.parse_string(dd[k]).value
dd[k] = to_source(st.visit(eq))
return FlatFunctionFactory({}, dd, ff.arguments, ff.funname)
def test_make_method_dictionary():
from dolang.parser import parse_string
from dolang.codegen import to_source
from dolang.function_compiler import make_method
# style 2: OrderedDict
eqs = ['x + lam*y']
from collections import OrderedDict
arguments = OrderedDict()
arguments['group_1'] = [('x', 1)]
arguments['group_2'] = [('x', 0), ('y', 0)]
constants = ['lam']
fun2 = make_method(eqs, arguments, constants)
print(to_source(fun2))
def substitute_preamble(ff: FlatFunctionFactory):
import dolang
import copy
from .symbolic import NameSubstituter
pr = copy.copy(ff.preamble)
for k in pr.keys():
pr[k] = dolang.parse_string(pr[k])
st = NameSubstituter(pr)
dd = copy.copy(ff.content)
for k in dd.keys():
dd[k] = to_source(st.transform(eq))
return FlatFunctionFactory({}, dd, ff.arguments, ff.funname)
def test_make_method_elaborate_definitions():
from dolang.parser import parse_string
from dolang.codegen import to_source
from dolang.function_compiler import make_method
# with elaborate definitions
from collections import OrderedDict
eqs = ['x + lam*y + z + w(1)']
arguments = OrderedDict()
arguments['group_1'] = [('x', 1)]
arguments['group_2'] = [('x', 0), ('y', 0)]
definitions = {'z': 'x', 'w': 'z+x(-1)+y(-1)'}
constants = ['lam']
fun = make_method(eqs, arguments, constants, definitions=definitions)
print(to_source(fun))
def read_equations(lines):
lines = [l.strip() for l in lines if len(l.strip()) >= 1]
conditions = []
equations = []
variables = []
complementarities = []
regex = re.compile("((.*):|)(.*)⟂(.*)")
for l in lines:
m = regex.match(l)
cond, eq, comp = m.groups()[1:]
if cond is not None:
cond = cond.strip()
conditions.append(cond)
comp = comp.strip()
d1 = match("_x <= _y", comp)
if d1:
lb = d1["_x"]
var = d1["_y"]
lb = to_source(lb)
else:
d2 = match("_y", comp)
lb = "-inf"
var = d2["_y"]
var = to_source(var)
equations.append(eq.strip())
variables.append(var)
complementarities.append(lb)
return [conditions, equations, variables, complementarities]
def test_sanitize():
from dolang.symbolic import sanitize, parse_string
from dolang.codegen import to_source
s = 'sin(a(1)+b+a+f(-1)+f(+4)+a(1))'
expected = "sin(a(1) + b + a(0) + f(-(1)) + f(4) + a(1))"
e = parse_string('sin(a(1)+b+a+f(-1)+f(+4)+a(1))')
enes = sanitize(e, variables=['a', 'f'])
assert (to_source(enes) == expected)
# it also works with the string directly
assert (sanitize(s, variables=['a', 'f']) == expected)
# we also deal with = signs, and convert to python exponents
assert (sanitize("a(1) = a^3 + b") == "a(1) == (a) ** (3) + b")
def eval_formula(expr: str, dataframe=None, context=None):
'''
expr: string
Symbolic expression to evaluate.
Example: `k(1)-delta*k(0)-i`
table: (optional) pandas dataframe
Each column is a time series, which can be indexed with dolo notations.
context: dict or CalibrationDict
'''
print("Evaluating: {}".format(expr))
if context is None:
dd = {} # context dictionary
elif isinstance(context, CalibrationDict):
dd = context.flat.copy()
else:
dd = context.copy()
# compat since normalize form for parameters doesn't match calib dict.
for k in [*dd.keys()]:
dd[stringify_symbol(k)] = dd[k]
expr_ast = parse_string(expr).value
variables = list_variables(expr_ast)
nexpr = stringify(expr_ast)
print(expr)
print(variables)
dd['log'] = log
dd['exp'] = exp
if dataframe is not None:
import pandas as pd
for (k, t) in variables:
dd[stringify_symbol((k, t))] = dataframe[k].shift(t)
dd['t'] = pd.Series(dataframe.index, index=dataframe.index)
expr = to_source(nexpr)
print(expr)
print(dd.keys())
res = eval(expr, dd)
return res
equations = ['{} - ({})'.format(*str.split(eq,'==')) for eq in equations]
equations = [parse_string(e) for e in equations]
with timeit("stringify equations"):
all_variables = [(v, 1) for v in model['symbols']['variables']] + \
[(v, 0) for v in model['symbols']['variables']] + \
[(v, -1) for v in model['symbols']['variables']] + \
[(v, 0) for v in model['symbols']['shocks']]
all_vnames = [e[0] for e in all_variables]
all_constants = model['symbols']['parameters']
# here comes the costly step
equations_stringified = [stringify(e, variables=all_vnames) for e in equations]
equations_stringified_strings = [to_source(e) for e in equations_stringified]
variables_stringified_strings = [stringify_variable(e) for e in all_variables]
stringify_variable( all_variables[-3] )
equations_stringified_strings[-1]
with timeit("Sympify equations"):
equations_stringified_sympy = [sympy.sympify(e) for e in equations_stringified_strings]
with timeit("Compute jacobian (sympy)"):
jac = []
for eq in equations_stringified_strings:
line = []
eqs = sympy.sympify(eq)
with timeit("stringify equations"):
all_variables = [(v, 1) for v in model['symbols']['variables']] + \
[(v, 0) for v in model['symbols']['variables']] + \
[(v, -1) for v in model['symbols']['variables']] + \
[(v, 0) for v in model['symbols']['shocks']]
all_vnames = [e[0] for e in all_variables]
all_constants = model['symbols']['parameters']
# here comes the costly step
equations_stringified = [
stringify(e, variables=all_vnames) for e in equations
]
equations_stringified_strings = [
to_source(e) for e in equations_stringified
]
variables_stringified_strings = [
stringify_variable(e) for e in all_variables
]
stringify_variable(all_variables[-3])
equations_stringified_strings[-1]
with timeit("Sympify equations"):
equations_stringified_sympy = [
sympy.sympify(e) for e in equations_stringified_strings
]
with timeit("Compute jacobian (sympy)"):
jac = []
equations = [parse_string(e) for e in equations]
with timeit("Normalize equations"):
all_variables = [(v, 1) for v in model['symbols']['variables']] + \
[(v, 0) for v in model['symbols']['variables']] + \
[(v, -1) for v in model['symbols']['variables']] + \
[(v, 0) for v in model['symbols']['shocks']]
all_vnames = [e[0] for e in all_variables]
all_constants = model['symbols']['parameters']
# here comes the costly step
equations_normalized = [
normalize(e, variables=all_vnames) for e in equations
]
equations_normalized_strings = [to_source(e) for e in equations_normalized]
variables_normalized_strings = [
stringify_variable(e) for e in all_variables
]
stringify_variable(all_variables[-3])
equations_normalized_strings[-1]
with timeit("Sympify equations"):
equations_normalized_sympy = [
sympy.sympify(e) for e in equations_normalized_strings
]
with timeit("Compute jacobian (sympy)"):
jac = []
for eq in equations_normalized_strings:
def test_parse_string():
from dolang.symbolic import parse_string
e = parse_string('sin(a(1)+b+f(1)+f(+4)+a(t+1))')
assert isinstance(e, ast.Expr)
s = to_source(e)
assert (s == "sin(a(1) + b + f(1) + f(+(4)) + a(t + 1))")