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)