def test_separate_variables1(): eq_code = ''' def integral(V, m, h, n, t, Iext): alpha = 0.1 * (V + 40) / (1 - np.exp(-(V + 40) / 10)) beta = 4.0 * np.exp(-(V + 65) / 18) dmdt = alpha * (1 - m) - beta * m alpha = 0.07 * np.exp(-(V + 65) / 20.) beta = 1 / (1 + np.exp(-(V + 35) / 10)) dhdt = alpha * (1 - h) - beta * h return dmdt, dhdt ''' analyser = DiffEqReader() analyser.visit(ast.parse(eq_code)) print('returns: ') pprint(analyser.returns) print("code_lines: ") pprint(analyser.code_lines) print("rights: ") pprint(analyser.rights) print("variables: ") pprint(analyser.variables) r = separate_variables(returns=analyser.returns, variables=analyser.variables, right_exprs=analyser.rights, code_lines=analyser.code_lines) pprint(r)
def test_separate_variables2(): code = '''def derivative(V, m, h, n, t, C, gNa, ENa, gK, EK, gL, EL, Iext): alpha = 0.1 * (V + 40) / (1 - bp.backend.exp(-(V + 40) / 10)) beta = 4.0 * bp.backend.exp(-(V + 65) / 18) dmdt = alpha * (1 - m) - beta * m alpha = 0.07 * bp.backend.exp(-(V + 65) / 20.) beta = 1 / (1 + bp.backend.exp(-(V + 35) / 10)) dhdt = alpha * (1 - h) - beta * h alpha = 0.01 * (V + 55) / (1 - bp.backend.exp(-(V + 55) / 10)) beta = 0.125 * bp.backend.exp(-(V + 65) / 80) dndt = alpha * (1 - n) - beta * n I_Na = (gNa * m ** 3.0 * h) * (V - ENa) I_K = (gK * n ** 4.0) * (V - EK) I_leak = gL * (V - EL) dVdt = (- I_Na - I_K - I_leak + Iext) / C return dVdt, dmdt, dhdt, dndt ''' from pprint import pprint pprint(separate_variables(code))
def transform_integrals_to_model(integrals): from brainpy.integrators import sympy_analysis if callable(integrals): integrals = [integrals] all_scope = dict() all_variables = set() all_parameters = set() analyzers = [] for integral in integrals: # integral function if Dispatcher is not None and isinstance(integral, Dispatcher): integral = integral.py_func else: integral = integral # original function f = integral.origin_f if Dispatcher is not None and isinstance(f, Dispatcher): f = f.py_func func_name = f.__name__ # code scope closure_vars = inspect.getclosurevars(f) code_scope = dict(closure_vars.nonlocals) code_scope.update(dict(closure_vars.globals)) # separate variables analysis = ast_analysis.separate_variables(f) variables_for_returns = analysis['variables_for_returns'] expressions_for_returns = analysis['expressions_for_returns'] for vi, (key, vars) in enumerate(variables_for_returns.items()): variables = [] for v in vars: if len(v) > 1: raise ValueError( 'Cannot analyze multi-assignment code line.') variables.append(v[0]) expressions = expressions_for_returns[key] var_name = integral.variables[vi] DE = sympy_analysis.SingleDiffEq(var_name=var_name, variables=variables, expressions=expressions, derivative_expr=key, scope=code_scope, func_name=func_name) analyzers.append(DE) # others for var in integral.variables: if var in all_variables: raise errors.ModelDefError( f'Variable {var} has been defined before. Cannot group ' f'this integral as a dynamic system.') all_variables.add(var) all_parameters.update(integral.parameters) all_scope.update(code_scope) return DynamicModel(integrals=integrals, analyzers=analyzers, variables=list(all_variables), parameters=list(all_parameters), scopes=all_scope)
def exp_euler_wrapper(f, show_code, dt, var_type, im_return): try: import sympy from brainpy.integrators import sympy_analysis except ModuleNotFoundError: raise errors.PackageMissingError( 'SymPy must be installed when using exponential euler methods.') if var_type == constants.SYSTEM_VAR: raise errors.IntegratorError( f'Exponential Euler method do not support {var_type} variable type.' ) dt_var = 'dt' class_kw, variables, parameters, arguments = utils.get_args(f) func_name = Tools.f_names(f) code_lines = [f'def {func_name}({", ".join(arguments)}):'] # code scope closure_vars = inspect.getclosurevars(f) code_scope = dict(closure_vars.nonlocals) code_scope.update(dict(closure_vars.globals)) code_scope[dt_var] = dt code_scope['f'] = f code_scope['exp'] = ops.exp analysis = separate_variables(f) variables_for_returns = analysis['variables_for_returns'] expressions_for_returns = analysis['expressions_for_returns'] for vi, (key, vars) in enumerate(variables_for_returns.items()): # separate variables sd_variables = [] for v in vars: if len(v) > 1: raise ValueError('Cannot analyze multi-assignment code line.') sd_variables.append(v[0]) expressions = expressions_for_returns[key] var_name = variables[vi] diff_eq = sympy_analysis.SingleDiffEq(var_name=var_name, variables=sd_variables, expressions=expressions, derivative_expr=key, scope=code_scope, func_name=func_name) f_expressions = diff_eq.get_f_expressions( substitute_vars=diff_eq.var_name) # code lines code_lines.extend([f" {str(expr)}" for expr in f_expressions[:-1]]) # get the linear system using sympy f_res = f_expressions[-1] df_expr = sympy_analysis.str2sympy(f_res.code).expr.expand() s_df = sympy.Symbol(f"{f_res.var_name}") code_lines.append( f' {s_df.name} = {sympy_analysis.sympy2str(df_expr)}') var = sympy.Symbol(diff_eq.var_name, real=True) # get df part s_linear = sympy.Symbol(f'_{diff_eq.var_name}_linear') s_linear_exp = sympy.Symbol(f'_{diff_eq.var_name}_linear_exp') s_df_part = sympy.Symbol(f'_{diff_eq.var_name}_df_part') if df_expr.has(var): # linear linear = sympy.collect(df_expr, var, evaluate=False)[var] code_lines.append( f' {s_linear.name} = {sympy_analysis.sympy2str(linear)}') # linear exponential linear_exp = sympy.exp(linear * dt) code_lines.append( f' {s_linear_exp.name} = {sympy_analysis.sympy2str(linear_exp)}' ) # df part df_part = (s_linear_exp - 1) / s_linear * s_df code_lines.append( f' {s_df_part.name} = {sympy_analysis.sympy2str(df_part)}') else: # linear exponential code_lines.append(f' {s_linear_exp.name} = sqrt({dt})') # df part code_lines.append( f' {s_df_part.name} = {sympy_analysis.sympy2str(dt * s_df)}') # update expression update = var + s_df_part # The actual update step code_lines.append( f' {diff_eq.var_name}_new = {sympy_analysis.sympy2str(update)}') code_lines.append('') code_lines.append(f' return {", ".join([f"{v}_new" for v in variables])}') return Tools.compile_and_assign_attrs(code_lines=code_lines, code_scope=code_scope, show_code=show_code, func_name=func_name, variables=variables, parameters=parameters, dt=dt, var_type=var_type)
def exp_euler(f, g, dt, sde_type, var_type, wiener_type, show_code): try: import sympy from brainpy.integrators import sympy_analysis except ModuleNotFoundError: raise errors.PackageMissingError( 'SymPy must be installed when using exponential euler methods.' ) if var_type == constants.SYSTEM_VAR: raise errors.IntegratorError( f'Exponential Euler method do not support {var_type} variable type.' ) if sde_type != constants.ITO_SDE: raise errors.IntegratorError( f'Exponential Euler method only supports Ito integral, but we got {sde_type}.' ) vdt, variables, parameters, arguments, func_name = common.basic_info( f=f, g=g) # 1. code scope closure_vars = inspect.getclosurevars(f) code_scope = dict(closure_vars.nonlocals) code_scope.update(dict(closure_vars.globals)) code_scope['f'] = f code_scope['g'] = g code_scope[vdt] = dt code_scope[f'{vdt}_sqrt'] = dt**0.5 code_scope['ops'] = ops code_scope['exp'] = ops.exp # 2. code lines code_lines = [f'def {func_name}({", ".join(arguments)}):'] # 2.1 dg # dg = g(x, t, *args) all_dg = [f'{var}_dg' for var in variables] code_lines.append( f' {", ".join(all_dg)} = g({", ".join(variables + parameters)})') code_lines.append(' ') # 2.2 dW Tools.noise_terms(code_lines, variables) # 2.3 dgdW # ---- # SCALAR_WIENER : dg * dW # VECTOR_WIENER : ops.sum(dg * dW, axis=-1) if wiener_type == constants.SCALAR_WIENER: for var in variables: code_lines.append(f' {var}_dgdW = {var}_dg * {var}_dW') else: for var in variables: code_lines.append( f' {var}_dgdW = ops.sum({var}_dg * {var}_dW, axis=-1)') code_lines.append(' ') # 2.4 new var # ---- analysis = separate_variables(f) variables_for_returns = analysis['variables_for_returns'] expressions_for_returns = analysis['expressions_for_returns'] for vi, (key, vars) in enumerate(variables_for_returns.items()): # separate variables sd_variables = [] for v in vars: if len(v) > 1: raise ValueError( 'Cannot analyze multi-assignment code line.') sd_variables.append(v[0]) expressions = expressions_for_returns[key] var_name = variables[vi] diff_eq = sympy_analysis.SingleDiffEq(var_name=var_name, variables=sd_variables, expressions=expressions, derivative_expr=key, scope=code_scope, func_name=func_name) f_expressions = diff_eq.get_f_expressions( substitute_vars=diff_eq.var_name) # code lines code_lines.extend( [f" {str(expr)}" for expr in f_expressions[:-1]]) # get the linear system using sympy f_res = f_expressions[-1] df_expr = sympy_analysis.str2sympy(f_res.code).expr.expand() s_df = sympy.Symbol(f"{f_res.var_name}") code_lines.append( f' {s_df.name} = {sympy_analysis.sympy2str(df_expr)}') var = sympy.Symbol(diff_eq.var_name, real=True) # get df part s_linear = sympy.Symbol(f'_{diff_eq.var_name}_linear') s_linear_exp = sympy.Symbol(f'_{diff_eq.var_name}_linear_exp') s_df_part = sympy.Symbol(f'_{diff_eq.var_name}_df_part') if df_expr.has(var): # linear linear = sympy.collect(df_expr, var, evaluate=False)[var] code_lines.append( f' {s_linear.name} = {sympy_analysis.sympy2str(linear)}') # linear exponential linear_exp = sympy.exp(linear * dt) code_lines.append( f' {s_linear_exp.name} = {sympy_analysis.sympy2str(linear_exp)}' ) # df part df_part = (s_linear_exp - 1) / s_linear * s_df code_lines.append( f' {s_df_part.name} = {sympy_analysis.sympy2str(df_part)}' ) else: # linear exponential code_lines.append(f' {s_linear_exp.name} = sqrt({dt})') # df part code_lines.append( f' {s_df_part.name} = {sympy_analysis.sympy2str(dt * s_df)}' ) # update expression update = var + s_df_part # The actual update step code_lines.append( f' {diff_eq.var_name}_new = {sympy_analysis.sympy2str(update)} + {var_name}_dgdW' ) code_lines.append('') # returns new_vars = [f'{var}_new' for var in variables] code_lines.append(f' return {", ".join(new_vars)}') # return and compile return common.compile_and_assign_attrs(code_lines=code_lines, code_scope=code_scope, show_code=show_code, variables=variables, parameters=parameters, func_name=func_name, sde_type=sde_type, var_type=var_type, wiener_type=wiener_type, dt=dt)