def test_predicate(): from dolang.symbolic import parse_string, str_expression e = parse_string("a[t] <= (x[t]+b)") print(e.pretty()) print(str_expression(e)) e = parse_string("∀t, a[t] <= (x[t]+b)") print(e.pretty()) print(str_expression(e))
def check(self): from dolang.symbolic import remove_timing, parse_string, str_expression p = self.data.get("projection") if p is None: raise AggregateException("Missing 'projection section'.") else: exo = self.model.symbols["exogenous"] proj = [] eqs = parse_string(p, start="assignment_block") for eq in eqs.children: lhs, _ = eq.children lhs = remove_timing(lhs) lhs = str_expression(lhs) proj.append(lhs) exo_in_proj = [e for e in exo if e in proj] diff = set(proj).difference(set(exo_in_proj)) if diff: raise AggregateException( f"Some projected values were not defined as exogenous in the agent's program: {', '.join(diff)}." ) t = self.data.get("transition") states = self.symbols.get("states") if t is not None and states is None: raise AggregateException( f"Aggregate transition equations are defined, whereas no aggregate state is filled in." ) elif t is None and states is not None: raise AggregateException( f"Aggregate states are defined, whereas no transition equation is filled in." ) elif t is not None and states is not None: trans = [] eqs = parse_string(t, start="assignment_block") for eq in eqs.children: lhs, _ = eq.children lhs = remove_timing(lhs) lhs = str_expression(lhs) trans.append(lhs) states_in_trans = [e for e in states if e in trans] trans_in_states = [e for e in trans if e in states] diff = set(trans).difference(set(states_in_trans)) if diff: raise AggregateException( f"Some variables defined in transition equations are not filled in aggregate states: {', '.join(diff)}." ) diff = set(states).difference(set(trans_in_states)) if diff: raise AggregateException( f"Some aggregate states do not have transition equations filled in: {', '.join(diff)}." )
def test_subperiod(): from dolang.symbolic import parse_string, str_expression e = parse_string("a[t$1] = a[t+1]") print(e.pretty()) print(str_expression(e)) e = parse_string("a[t$consumption] = a[t+1]") print(e.pretty()) print(str_expression(e)) from lark.lark import Lark
def test_list_symbols(): from dolang.symbolic import parse_string e = parse_string('sin(a(1)+b+f(1)+f(+4)+a(1))+cos(0)') ll = list_symbols(e) # cos is recognized as a usual function assert (ll.variables == [('a', 1), ('f', 1), ('f', 4)]) assert (ll.parameters == ['b']) e = parse_string('sin(a(1)+b+f(1)+f(+4)+a(1))+cos(0)') ll = list_symbols(e, funs=['f']) # now we add a custom function assert (ll.variables == [('a', 1)]) assert (ll.parameters == ['b'])
def compile_factory(fff: FlatFunctionFactory): arguments = [*fff.arguments.keys()] funname = fff.funname unpacking = [] for i, (arg_group_name, arg_group) in enumerate(fff.arguments.items()): for pos, sym in enumerate(arg_group): rhs = Subscript(value=Name(id=arg_group_name, ctx=Load()), slice=Index(Num(pos)), ctx=Load()) val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs) unpacking.append(val) body = [] for (k, neq) in fff.preamble.items(): val = parse_string(neq).value line = Assign(targets=[Name(id=k, ctx=Store())], value=val) body.append(line) for n, (k, neq) in enumerate(fff.content.items()): # should the result of parse_string always of type Expr ? val = parse_string(neq).value line = Assign(targets=[Name(id=k, ctx=Store())], value=val) body.append(line) # for n, (lhs, neq) in enumerate(fff.content.items()): line = Assign(targets=[ Subscript(value=Name(id='out', ctx=Load()), slice=Index(Num(n)), ctx=Store()) ], value=Name(id=lhs, ctx=Load())) body.append(line) f = FunctionDef(name=funname, args=ast_arguments(args=[arg(arg=a) for a in arguments] + [arg(arg='out')], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]), body=unpacking + body, decorator_list=[]) mod = Module(body=[f]) mmod = ast.fix_missing_locations(mod) return mmod
def test_parsing_unicode(): from dolang.symbolic import parse_string s = "αα" e = parse_string(s) print(e)
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, Tree) s = to_source(e) assert s == "sin(a[t+1] + b + f[t+1] + f[t+4] + a[t+1])"
def ℰ(self): if self.__equilibrium__ is None: if self.features["with-aggregate-states"]: arguments_ = { "e": [(e, 0) for e in self.model.symbols["exogenous"]], "s": [(e, 0) for e in self.model.symbols["states"]], "x": [(e, 0) for e in self.model.symbols["controls"]], "m": [(e, 0) for e in self.symbols["exogenous"]], "S": [(e, 0) for e in self.symbols["states"]], "X": [(e, 0) for e in self.symbols["aggregate"]], "m_1": [(e, 1) for e in self.symbols["exogenous"]], "S_1": [(e, 1) for e in self.symbols["states"]], "X_1": [(e, 1) for e in self.symbols["aggregate"]], "p": self.symbols["parameters"], } else: arguments_ = { "e": [(e, 0) for e in self.model.symbols["exogenous"]], "s": [(e, 0) for e in self.model.symbols["states"]], "x": [(e, 0) for e in self.model.symbols["controls"]], "m": [(e, 0) for e in self.symbols["exogenous"]], "X": [(e, 0) for e in self.symbols["aggregate"]], "m_1": [(e, 1) for e in self.symbols["exogenous"]], "X_1": [(e, 1) for e in self.symbols["aggregate"]], "p": self.symbols["parameters"], } vars = sum([[e[0] for e in h] for h in [*arguments_.values()][:-1]], []) arguments = { k: [dolang.symbolic.stringify_symbol(e) for e in v] for k, v in arguments_.items() } preamble = {} # for now from dolang.symbolic import sanitize, stringify eqs = parse_string(self.data["equilibrium"], start="equation_block") eqs = sanitize(eqs, variables=vars) eqs = stringify(eqs) content = {} for i, eq in enumerate(eqs.children): lhs, rhs = eq.children content[f"eq_{i}"] = "({1})-({0})".format( str_expression(lhs), str_expression(rhs)) fff = FlatFunctionFactory(preamble, content, arguments, "equilibrium") _, gufun = dolang.function_compiler.make_method_from_factory( fff, debug=self.debug) from dolang.vectorize import standard_function self.__equilibrium__ = standard_function(gufun, len(content)) return self.__equilibrium__
def test_remove_timing(): from dolang.symbolic import parse_string, remove_timing e = parse_string("sin(a(1)+b+f(1)+f(4)+a[t+1])") assert isinstance(e, Tree) s = to_source(e) rr = remove_timing(s) assert rr == "sin(a + b + f + f + a)"
def test_parsing(): from dolang.grammar import str_expression, sanitize from dolang.symbolic import parse_string e = parse_string("s + a(0) + b[t-1] + b[t] + b[t+1]") print(e.pretty()) e = parse_string("chi*n^eta*c^sigma - w(1)") print(e.pretty()) e = parse_string("chi*n^eta*c^sigma - w(1) | 0.01 <= n <= 1.0") print(e.pretty()) e = parse_string("chi*n^eta*c^sigma - w(1) ⟂ 0.01 <= n <= 1.0") print(e.pretty()) e = parse_string("i = exp(z)*k^alpha*n^(1-alpha) - (m)^(-1/sigma)") print(e.pretty()) f = sanitize(e, variables=["m"]) print(str_expression(f)) s = "i = exp(z)*k^alpha*n^(1-alpha) - (m)^(-1/sigma)" e = parse_string(s) s = "i = exp(z)*k^alpha*n^(1-alpha) - (m)^(-1/sigma)" # vars = ['z', 'p', 'k', 'n', 'i', 'm', 'V', 'u', 'y', 'c', 'rk', 'w', 'y', 'c'] # print("HI") # v = sanitize(s, variables=vars) # print(v) e = parse_string(s) print(str_expression(e))
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 test_expectation(): from dolang.symbolic import parse_string s = "�[ (x[t+1] / x[t]) ]" e = parse_string(s) print(e.pretty()) from dolang.symbolic import str_expression print(str_expression(e))
def test_list_symbols_debug(): from dolang.symbolic import parse_string e = parse_string('sin(a(1)+b+f(1)+f(+4)+a(1)+a+a(t+1))+cos(0)') l = ListSymbols(known_functions=['sin', 'f']) l.visit(e) # note that cos is recognized as variable assert (l.variables == [(('a', 1), 4), (('a', 1), 22), (('cos', 0), 37)]) assert (l.constants == [('b', 9), ('a', 27)]) assert (l.functions == [('sin', 0), ('f', 11), ('f', 16)]) assert (l.problems == [['a', 0, 29, 'incorrect subscript']])
def test_list_symbols(): from dolang.symbolic import parse_string e = parse_string("sin(a(1)+b+f(1)+f(4)+a(1))+cos(0)") ll = list_symbols(e) print(ll.variables) print(ll.parameters) # cos is recognized as a usual function assert ll.variables == [("a", 1), ("f", 1), ("f", 4)] assert ll.parameters == ["b"]
def projection(self): # , m: 'n_e', y: "n_y", p: "n_p"): # TODO: # behaves in a very misleading way if wrong number of argument is supplied # if no aggregate states, projection(m,x) (instead of projection(m,x,p)) returns zeros if self.__projection__ is None: if self.features["with-aggregate-states"]: arguments_ = { "m": [(e, 0) for e in self.symbols["exogenous"]], "S": [(e, 0) for e in self.symbols["states"]], "X": [(e, 0) for e in self.symbols["aggregate"]], "p": self.symbols["parameters"], } else: arguments_ = { "m": [(e, 0) for e in self.symbols["exogenous"]], "X": [(e, 0) for e in self.symbols["aggregate"]], "p": self.symbols["parameters"], } vars = sum([[e[0] for e in h] for h in [*arguments_.values()][:-1]], []) arguments = { k: [dolang.symbolic.stringify_symbol(e) for e in v] for k, v in arguments_.items() } preamble = {} # for now from dolang.symbolic import sanitize, stringify eqs = parse_string(self.data["projection"], start="assignment_block") eqs = sanitize(eqs, variables=vars) eqs = stringify(eqs) content = {} for eq in eqs.children: lhs, rhs = eq.children content[str_expression(lhs)] = str_expression(rhs) fff = FlatFunctionFactory(preamble, content, arguments, "equilibrium") _, gufun = dolang.function_compiler.make_method_from_factory( fff, debug=self.debug) from dolang.vectorize import standard_function self.__projection__ = standard_function(gufun, len(content)) return self.__projection__
def compile_factory(fff: FlatFunctionFactory): arguments = [*fff.arguments.keys()] funname = fff.funname unpacking = [] for i, (arg_group_name, arg_group) in enumerate(fff.arguments.items()): for pos, sym in enumerate(arg_group): rhs = Subscript(value=Name(id=arg_group_name, ctx=Load()), slice=Index(Num(pos)), ctx=Load()) val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs) unpacking.append(val) body = [] for (k, neq) in fff.preamble.items(): val = parse_string(neq).value line = Assign(targets=[Name(id=k, ctx=Store())], value=val) body.append(line) for n, (k, neq) in enumerate(fff.content.items()): # should the result of parse_string always of type Expr ? val = parse_string(neq).value line = Assign(targets=[Name(id=k, ctx=Store())], value=val) body.append(line) # for n, (lhs, neq) in enumerate(fff.content.items()): line = Assign(targets=[Subscript(value=Name(id='out', ctx=Load()), slice=Index(Num(n)), ctx=Store())], value=Name(id=lhs, ctx=Load())) body.append(line) f = FunctionDef(name=funname, args=ast_arguments(args=[arg(arg=a) for a in arguments] + [arg(arg='out')], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[]), body=unpacking + body, decorator_list=[]) mod = Module(body=[f]) mmod = ast.fix_missing_locations(mod) return mmod
def 𝒢(self): if (self.__transition__ is None) and self.features["with-aggregate-states"]: arguments_ = { "m_m1": [(e, -1) for e in self.symbols["exogenous"]], "S_m1": [(e, -1) for e in self.symbols["states"]], "X_m1": [(e, -1) for e in self.symbols["aggregate"]], "m": [(e, 0) for e in self.symbols["exogenous"]], "p": self.symbols["parameters"], } vars = sum([[e[0] for e in h] for h in [*arguments_.values()][:-1]], []) arguments = { k: [dolang.symbolic.stringify_symbol(e) for e in v] for k, v in arguments_.items() } preamble = {} # for now from dolang.symbolic import ( sanitize, parse_string, str_expression, stringify, ) eqs = parse_string(self.data["transition"], start="assignment_block") eqs = sanitize(eqs, variables=vars) eqs = stringify(eqs) content = {} for i, eq in enumerate(eqs.children): lhs, rhs = eq.children content[str_expression(lhs)] = str_expression(rhs) from dolang.factory import FlatFunctionFactory fff = FlatFunctionFactory(preamble, content, arguments, "transition") _, gufun = dolang.function_compiler.make_method_from_factory( fff, debug=self.debug) from dolang.vectorize import standard_function self.__transition__ = standard_function(gufun, len(content)) return self.__transition__
def get_calibration(self): from dolang.symbolic import remove_timing import copy calibration = dict() for k, v in self.data.get("calibration", {}).items(): if v.tag == "tag:yaml.org,2002:str": expr = parse_string(v) expr = remove_timing(expr) expr = str_expression(expr) else: expr = float(v.value) kk = remove_timing(parse_string(k)) kk = str_expression(kk) calibration[kk] = expr from dolang.triangular_solver import solve_triangular_system return solve_triangular_system(calibration)
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
model['symbols'].keys() with timeit("Solve triangular system"): calibration = model['calibration'] sol = solve_triangular_system(calibration) with timeit("Parse equations"): equations = model['equations'] compat = lambda s: s.replace("^", "**").replace('==','=').replace('=','==') equations = [compat(eq) for eq in equations] 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]
def test_list_variables(): from dolang.symbolic import parse_string e = parse_string("sin(a(1)+b+f(1)+f(4)+sin(a)+k*cos(a(0)))") list_variables(e) assert list_variables(e) == [("a", 1), ("f", 1), ("f", 4), ("a", 0)]
def test_multiline(): from dolang.symbolic import parse_string e = parse_string("a ⟂ x <= y <= exp(z)\nb ⟂ x <= y <= z", start="complementarity_block") assert len(e.children) == 2
def compile_factory(fff: FlatFunctionFactory): arguments = [*fff.arguments.keys()] funname = fff.funname unpacking = [] for i, (arg_group_name, arg_group) in enumerate(fff.arguments.items()): for pos, sym in enumerate(arg_group): rhs = Subscript( value=Name(id=arg_group_name, ctx=Load()), slice=Index(Num(pos)), ctx=Load(), ) val = Assign(targets=[Name(id=sym, ctx=Store())], value=rhs) unpacking.append(val) body = [] for (k, neq) in fff.preamble.items(): tree = parse_string(neq) val = tree_to_ast(tree).value line = Assign(targets=[Name(id=k, ctx=Store())], value=val) body.append(line) for n, (k, neq) in enumerate(fff.content.items()): tree = parse_string(neq) val = tree_to_ast(tree).value line = Assign(targets=[Name(id=k, ctx=Store())], value=val) body.append(line) # for n, (lhs, neq) in enumerate(fff.content.items()): line = Assign( targets=[ Subscript(value=Name(id="out", ctx=Load()), slice=Index(Num(n)), ctx=Store()) ], value=Name(id=lhs, ctx=Load()), ) body.append(line) if sys.version_info >= (3, 8, 0): f = FunctionDef( name=funname, args=ast_arguments( posonlyargs=[], args=[arg(arg=a) for a in arguments] + [arg(arg="out")], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[], ), body=unpacking + body, decorator_list=[], ) mod = Module(body=[f], type_ignores=[]) else: f = FunctionDef( name=funname, args=ast_arguments( args=[arg(arg=a) for a in arguments] + [arg(arg="out")], vararg=None, kwarg=None, kwonlyargs=[], kw_defaults=[], defaults=[], ), body=unpacking + body, decorator_list=[], ) mod = Module(body=[f]) mmod = ast.fix_missing_locations(mod) return mmod
model = json.load(f) model['symbols'].keys() with timeit("Solve triangular system"): calibration = model['calibration'] sol = solve_triangular_system(calibration) with timeit("Parse equations"): equations = model['equations'] compat = lambda s: s.replace("^", "**").replace('==', '=').replace( '=', '==') equations = [compat(eq) for eq in equations] 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 = [
def equations(self): import yaml.nodes if self.__equations__ is None: vars = self.variables + [*self.definitions.keys()] d = dict() for g, v in self.data["equations"].items(): # new style if isinstance(v, yaml.nodes.ScalarNode): assert v.style == "|" if g in ("arbitrage", ): start = "complementarity_block" else: start = "assignment_block" eqs = parse_string(v, start=start) eqs = sanitize(eqs, variables=vars) eq_list = eqs.children # old style else: eq_list = [] for eq_string in v: start = "equation" # it should be assignment eq = parse_string(eq_string, start=start) eq = sanitize(eq, variables=vars) eq_list.append(eq) if g in ("arbitrage", ): ll = [] # List[str] ll_lb = [] # List[str] ll_ub = [] # List[str] with_complementarity = False for i, eq in enumerate(eq_list): if eq.data == "double_complementarity": v = eq.children[1].children[1].children[ 0].children[0].value t = int(eq.children[1].children[1].children[1]. children[0].value) expected = ( self.symbols["controls"][i], 0, ) # TODO raise nice error message if (v, t) != expected: raise Exception( f"Incorrect variable in complementarity: expected {expected}. Found {(v,t)}" ) ll_lb.append( str_expression(eq.children[1].children[0])) ll_ub.append( str_expression(eq.children[1].children[2])) eq = eq.children[0] with_complementarity = True else: ll_lb.append("-inf") ll_ub.append("inf") from dolang.symbolic import list_symbols # syms = list_symbols(eq) ll.append(str_expression(eq)) d[g] = ll if with_complementarity: d[g + "_lb"] = ll_lb d[g + "_ub"] = ll_ub else: # TODO: we should check here that equations are well specified d[g] = [str_expression(e) for e in eq_list] # if "controls_lb" not in d: # for ind, g in enumerate(("controls_lb", "controls_ub")): # eqs = [] # for i, eq in enumerate(d['arbitrage']): # if "⟂" not in eq: # if ind == 0: # eq = "-inf" # else: # eq = "inf" # else: # comp = eq.split("⟂")[1].strip() # v = self.symbols["controls"][i] # eq = decode_complementarity(comp, v+"[t]")[ind] # eqs.append(eq) # d[g] = eqs self.__equations__ = d return self.__equations__
def get_calibration(self): # if self.__calibration__ is None: from dolang.symbolic import remove_timing import copy symbols = self.symbols calibration = dict() for k, v in self.data.get("calibration", {}).items(): if v.tag == "tag:yaml.org,2002:str": expr = parse_string(v) expr = remove_timing(expr) expr = str_expression(expr) else: expr = float(v.value) kk = remove_timing(parse_string(k)) kk = str_expression(kk) calibration[kk] = expr definitions = self.definitions initial_values = { "exogenous": 0, "expectations": 0, "values": 0, "controls": float("nan"), "states": float("nan"), } # variables defined by a model equation default to using these definitions initialized_from_model = { "values": "value", "expectations": "expectation", "direct_responses": "direct_response", } for k, v in definitions.items(): kk = remove_timing(k) if kk not in calibration: if isinstance(v, str): vv = remove_timing(v) else: vv = v calibration[kk] = vv for symbol_group in symbols: if symbol_group not in initialized_from_model.keys(): if symbol_group in initial_values: default = initial_values[symbol_group] else: default = float("nan") for s in symbols[symbol_group]: if s not in calibration: calibration[s] = default from dolang.triangular_solver import solve_triangular_system return solve_triangular_system(calibration)
def definitions(self): from yaml import ScalarNode if self.__definitions__ is None: # at this stage, basic_symbols doesn't contain auxiliaries basic_symbols = self.symbols vars = sum( [ basic_symbols[k] for k in basic_symbols.keys() if k != "parameters" ], [], ) # # auxiliaries = [remove_timing(parse_string(k)) for k in self.data.get('definitions', {})] # # auxiliaries = [str_expression(e) for e in auxiliaries] # # symbols['auxiliaries'] = auxiliaries if "definitions" not in self.data: self.__definitions__ = {} # self.__symbols__['auxiliaries'] = [] elif isinstance(self.data["definitions"], ScalarNode): definitions = {} # new-style from lark import Token def_block_tree = parse_string(self.data["definitions"], start="assignment_block") def_block_tree = sanitize( def_block_tree ) # just to replace (v,) by (v,0) # TODO: remove auxiliaries = [] for eq_tree in def_block_tree.children: lhs, rhs = eq_tree.children tok_name: Token = lhs.children[0].children[0] tok_date: Token = lhs.children[1].children[0] name = tok_name.value date = int(tok_date.value) if name in vars: raise Exception( f"definitions:{tok_name.line}:{tok_name.column}: Auxiliary variable '{name}'' already defined." ) if date != 0: raise Exception( f"definitions:{tok_name.line}:{tok_name.column}: Auxiliary variable '{name}' must be defined at date 't'." ) # here we could check some stuff from dolang import list_symbols syms = list_symbols(rhs) for p in syms.parameters: if p in vars: raise Exception( f"definitions:{tok_name.line}: Symbol '{p}' is defined as a variable. Can't appear as a parameter." ) if p not in self.symbols["parameters"]: raise Exception( f"definitions:{tok_name.line}: Paremeter '{p}' must be defined as a model symbol." ) for v in syms.variables: if v[0] not in vars: raise Exception( f"definitions:{tok_name.line}: Variable '{v[0]}[t]' is not defined." ) auxiliaries.append(name) vars.append(name) definitions[str_expression(lhs)] = str_expression(rhs) self.__symbols__["auxiliaries"] = auxiliaries self.__definitions__ = definitions else: # old style from dolang.symbolic import remove_timing auxiliaries = [ remove_timing(parse_string(k)) for k in self.data.get("definitions", {}) ] auxiliaries = [str_expression(e) for e in auxiliaries] self.__symbols__["auxiliaries"] = auxiliaries vars = self.variables auxs = [] definitions = self.data["definitions"] d = dict() for i in range(len(definitions.value)): kk = definitions.value[i][0] if self.__compat__: k = parse_string(kk.value) if k.data == "symbol": # TODO: warn that definitions should be timed from dolang.grammar import create_variable k = create_variable(k.children[0].value, 0) else: k = parse_string(kk.value, start="variable") k = sanitize(k, variables=vars) assert k.children[1].children[0].value == "0" vv = definitions.value[i][1] v = parse_string(vv, start="formula") v = sanitize(v, variables=vars) v = str_expression(v) key = str_expression(k) vars.append(key) d[key] = v auxs.append(remove_timing(key)) self.__symbols__["auxiliaries"] = auxs self.__definitions__ = d return self.__definitions__
def test_list_variables(): from dolang.symbolic import parse_string e = parse_string('sin(a(1)+b+f(1)+f(+4)+sin(a)+k*cos(a(0)))') list_variables(e) assert (list_variables(e) == [('a', 1), ('f', 1), ('f', 4), ('a', 0)]) assert (list_variables(e, funs=['f']) == [('a', 1), ('a', 0)])
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))")