def ast_to_sympy(expr):
'''Converts an AST expression to a sympy expression (STUPID)'''
from dolang import to_source
s = to_source(expr)
not_to_be_treated_as_functions = ['alpha', 'beta', 'gamma', 'zeta', 'Chi']
d = {v: sympy.Symbol(v) for v in not_to_be_treated_as_functions}
return sympy.sympify(s, locals=d)
def ast_to_sympy(expr):
'''Converts an AST expression to a sympy expression (STUPID)'''
from dolang import to_source
s = to_source(expr)
not_to_be_treated_as_functions = ['alpha','beta', 'gamma','zeta', 'Chi']
d = {v: sympy.Symbol(v) for v in not_to_be_treated_as_functions}
return sympy.sympify(s, locals=d)
def eval_formula(expr, dataframe=None, context=None):
'''
expr: string
Symbolic expression to evaluate.
Example: `k(1)-delta*k-i`
table: (optional) pandas dataframe
Each column is a time series, which can be indexed with dolo notations.
context: dict or CalibrationDict
'''
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(k)] = dd[k]
from numpy import log, exp
dd['log'] = log
dd['exp'] = exp
if dataframe is not None:
import pandas as pd
tvariables = dataframe.columns
for k in tvariables:
if k in dd:
dd[k + '_ss'] = dd[k] # steady-state value
dd[stringify((k, 0))] = dataframe[k]
for h in range(1, 3): # maximum number of lags
dd[stringify((k, -h))] = dataframe[k].shift(h)
dd[stringify((k, h))] = dataframe[k].shift(-h)
dd['t'] = pd.Series(dataframe.index, index=dataframe.index)
import ast
expr_ast = ast.parse(expr).body[0].value
# nexpr = StandardizeDatesSimple(tvariables).visit(expr_ast)
print(tvariables)
nexpr = normalize(expr_ast, variables=tvariables)
expr = to_source(nexpr)
res = eval(expr, dd)
return res
def get_factory(model, eq_type: str, tshift: int = 0):
from dolo.compiler.model import decode_complementarity
from dolo.compiler.recipes import recipes
from dolang.symbolic import stringify, stringify_symbol
equations = model.equations
if eq_type == "auxiliary":
eqs = [('{}({})'.format(s, 0)) for s in model.symbols['auxiliaries']]
specs = {
'eqs': [['exogenous', 0, 'm'], ['states', 0, 's'],
['controls', 0, 'x'], ['parameters', 0, 'p']]
}
else:
eqs = equations[eq_type]
if eq_type in ('controls_lb', 'controls_ub'):
specs = {
'eqs':
recipes['dtcc']['specs']['arbitrage']['complementarities'][
'left-right']
}
else:
specs = recipes['dtcc']['specs'][eq_type]
specs = shift_spec(specs, tshift=tshift)
preamble_tshift = set([s[1] for s in specs['eqs'] if s[0] == 'states'])
preamble_tshift = preamble_tshift.intersection(
set([s[1] for s in specs['eqs'] if s[0] == 'controls']))
args = []
for sg in specs['eqs']:
if sg[0] == 'parameters':
args.append([s for s in model.symbols["parameters"]])
else:
args.append([(s, sg[1]) for s in model.symbols[sg[0]]])
args = [[stringify_symbol(e) for e in vg] for vg in args]
arguments = dict(zip([sg[2] for sg in specs['eqs']], args))
# temp
eqs = [eq.replace("==","=").replace("=","==") for eq in eqs]
if 'target' in specs:
sg = specs['target']
targets = [(s, sg[1]) for s in model.symbols[sg[0]]]
eqs = [eq.split('==')[1] for eq in eqs]
else:
eqs = [("({1})-({0})".format(*eq.split('==')) if '==' in eq else eq)
for eq in eqs]
targets = [('out{}'.format(i), 0) for i in range(len(eqs))]
eqs = [str.strip(eq) for eq in eqs]
eqs = [dolang.parse_string(eq) for eq in eqs]
es = ExpressionSanitizer(model.variables)
eqs = [es.visit(eq) for eq in eqs]
eqs = [time_shift(eq, tshift) for eq in eqs]
eqs = [stringify(eq) for eq in eqs]
eqs = [dolang.to_source(eq) for eq in eqs]
targets = [stringify_symbol(e) for e in targets]
# sanitize defs ( should be )
defs = dict()
for k in model.definitions:
if '(' not in k:
s = "{}(0)".format(k)
val = model.definitions[k]
val = es.visit(dolang.parse_string(val))
for t in preamble_tshift:
s = stringify_symbol((k, t))
vv = stringify(time_shift(val, t))
defs[s] = dolang.to_source(vv)
preamble = reorder_preamble(defs)
eqs = dict(zip(targets, eqs))
ff = FlatFunctionFactory(preamble, eqs, arguments, eq_type)
return ff
def print_matlab(sexpr):
ss = (to_source(sexpr))
ss = ss.replace(' ** ', '.^')
ss = ss.replace(' * ', '.*')
ss = ss.replace(' / ', './')
return ss
def print_matlab(sexpr):
ss = (to_source(sexpr))
ss = ss.replace(' ** ', '.^')
ss = ss.replace(' * ', '.*')
ss = ss.replace(' / ', './')
return ss