def wrap(wrapper, f, g, dt, sde_type, var_type, wiener_type, show_code):
"""The base function to format a SRK method.
Parameters
----------
f : callable
The drift function of the SDE.
g : callable
The diffusion function of the SDE.
dt : float
The numerical precision.
sde_type : str
"utils.ITO_SDE" : Ito's Stochastic Calculus.
"utils.STRA_SDE" : Stratonovich's Stochastic Calculus.
wiener_type : str
var_type : str
"scalar" : with the shape of ().
"population" : with the shape of (N,) or (N1, N2) or (N1, N2, ...).
"system": with the shape of (d, ), (d, N), or (d, N1, N2).
show_code : bool
Whether show the formatted code.
Returns
-------
numerical_func : callable
The numerical function.
"""
var_type = constants.POPU_VAR if var_type is None else var_type
sde_type = constants.ITO_SDE if sde_type is None else sde_type
wiener_type = constants.SCALAR_WIENER if wiener_type is None else wiener_type
if var_type not in constants.SUPPORTED_VAR_TYPE:
raise errors.IntegratorError(f'Currently, BrainPy only supports variable types: '
f'{constants.SUPPORTED_VAR_TYPE}. But we got {var_type}.')
if sde_type != constants.ITO_SDE:
raise errors.IntegratorError(f'SRK method for SDEs with scalar noise only supports Ito SDE type, '
f'but we got {sde_type} integral.')
if wiener_type != constants.SCALAR_WIENER:
raise errors.IntegratorError(f'SRK method for SDEs with scalar noise only supports scalar '
f'Wiener Process, but we got "{wiener_type}" noise.')
show_code = False if show_code is None else show_code
dt = backend.get_dt() if dt is None else dt
if f is not None and g is not None:
return wrapper(f=f, g=g, dt=dt, show_code=show_code, sde_type=sde_type,
var_type=var_type, wiener_type=wiener_type)
elif f is not None:
return lambda g: wrapper(f=f, g=g, dt=dt, show_code=show_code, sde_type=sde_type,
var_type=var_type, wiener_type=wiener_type)
elif g is not None:
return lambda f: wrapper(f=f, g=g, dt=dt, show_code=show_code, sde_type=sde_type,
var_type=var_type, wiener_type=wiener_type)
else:
raise ValueError('Must provide "f" or "g".')
def __init__(self, diff_eq):
if not isinstance(diff_eq, ast_analysis.DiffEquation):
if diff_eq.__class__.__name__ != 'function':
raise errors.IntegratorError(
'"diff_eq" must be a function or an instance of DiffEquation .'
)
else:
diff_eq = ast_analysis.DiffEquation(func=diff_eq)
self.diff_eq = diff_eq
self._update_code = None
self._update_func = None
def __init__(self,
f,
var_type=None,
dt=None,
name=None,
adaptive=None,
tol=None,
show_code=False):
super(AdaptiveRKIntegrator, self).__init__(f=f,
var_type=var_type,
dt=dt,
name=name,
show_code=show_code)
# check parameters
self.adaptive = False if (adaptive is None) else adaptive
self.tol = 0.1 if tol is None else tol
self.var_type = C.POP_VAR if var_type is None else var_type
if self.var_type not in C.SUPPORTED_VAR_TYPE:
raise errors.IntegratorError(
f'"var_type" only supports {C.SUPPORTED_VAR_TYPE}, '
f'not {self.var_type}.')
# integrator keywords
keywords = {
C.F: 'the derivative function',
C.DT: 'the precision of numerical integration'
}
for v in self.variables:
keywords[f'{v}_new'] = 'the intermediate value'
for i in range(1, len(self.A) + 1):
keywords[f'd{v}_k{i}'] = 'the intermediate value'
for i in range(2, len(self.A) + 1):
keywords[f'k{i}_{v}_arg'] = 'the intermediate value'
keywords[f'k{i}_t_arg'] = 'the intermediate value'
if adaptive:
keywords['dt_new'] = 'the new numerical precision "dt"'
keywords['tol'] = 'the tolerance for the local truncation error'
keywords['error'] = 'the local truncation error'
for v in self.variables:
keywords[f'{v}_te'] = 'the local truncation error'
self.code_scope['tol'] = tol
self.code_scope['math'] = bm
utils.check_kws(self.arguments, keywords)
# build the integrator
self.build()
def heun(f=None,
g=None,
dt=None,
sde_type=None,
var_type=None,
wiener_type=None,
show_code=None):
if sde_type != constants.STRA_SDE:
raise errors.IntegratorError(
f'Heun method only supports Stranovich integral of SDEs, '
f'but we got {sde_type} integral.')
return Wrapper.wrap(Wrapper.euler_and_heun,
f=f,
g=g,
dt=dt,
sde_type=sde_type,
var_type=var_type,
wiener_type=wiener_type,
show_code=show_code)
def __init__(self,
f,
g,
dt=None,
name=None,
show_code=False,
var_type=None,
intg_type=None,
wiener_type=None):
if intg_type != constants.STRA_SDE:
raise errors.IntegratorError(
f'Heun method only supports Stranovich integral of SDEs, '
f'but we got {intg_type} integral.')
super(Heun, self).__init__(f=f,
g=g,
dt=dt,
show_code=show_code,
name=name,
var_type=var_type,
intg_type=intg_type,
wiener_type=wiener_type)
self.build()
def __init__(self,
f,
g,
dt=None,
name=None,
show_code=False,
var_type=None,
intg_type=None,
wiener_type=None):
super(SDEIntegrator, self).__init__(name=name)
# derivative functions
self.derivative = {constants.F: f, constants.G: g}
self.f = f
self.g = g
# integration function
self.integral = None
# essential parameters
self.dt = math.get_dt() if dt is None else dt
assert isinstance(
self.dt, (int, float)), f'"dt" must be a float, but got {self.dt}'
intg_type = constants.ITO_SDE if intg_type is None else intg_type
var_type = constants.SCALAR_VAR if var_type is None else var_type
wiener_type = constants.SCALAR_WIENER if wiener_type is None else wiener_type
if intg_type not in constants.SUPPORTED_INTG_TYPE:
raise errors.IntegratorError(
f'Currently, BrainPy only support SDE_INT types: '
f'{constants.SUPPORTED_INTG_TYPE}. But we got {intg_type}.')
if var_type not in constants.SUPPORTED_VAR_TYPE:
raise errors.IntegratorError(
f'Currently, BrainPy only supports variable types: '
f'{constants.SUPPORTED_VAR_TYPE}. But we got {var_type}.')
if wiener_type not in constants.SUPPORTED_WIENER_TYPE:
raise errors.IntegratorError(
f'Currently, BrainPy only supports Wiener '
f'Process types: {constants.SUPPORTED_WIENER_TYPE}. '
f'But we got {wiener_type}.')
self.var_type = var_type # variable type
self.intg_type = intg_type # integral type
self.wiener_type = wiener_type # wiener process type
# parse function arguments
variables, parameters, arguments = utils.get_args(f)
self.variables = variables # variable names, (before 't')
self.parameters = parameters # parameter names, (after 't')
self.arguments = list(arguments) + [f'{constants.DT}={self.dt}'
] # function arguments
# random seed
self.rng = math.random.RandomState()
# code scope
self.code_scope = {
constants.F: f,
constants.G: g,
'math': math,
'random': self.rng
}
# code lines
self.func_name = f_names(f)
self.code_lines = [
f'def {self.func_name}({", ".join(self.arguments)}):'
]
# others
self.show_code = show_code
def symbolic_build(self):
if self.var_type == constants.SYSTEM_VAR:
raise errors.IntegratorError(
f'Exponential Euler method do not support {self.var_type} variable type.'
)
if self.intg_type != constants.ITO_SDE:
raise errors.IntegratorError(
f'Exponential Euler method only supports Ito integral, but we got {self.intg_type}.'
)
if sympy is None or analysis_by_sympy is None:
raise errors.PackageMissingError('SymPy must be installed when '
'using exponential integrators.')
# check bound method
if hasattr(self.derivative[constants.F], '__self__'):
self.code_lines = [
f'def {self.func_name}({", ".join(["self"] + list(self.arguments))}):'
]
# 1. code scope
closure_vars = inspect.getclosurevars(self.derivative[constants.F])
self.code_scope.update(closure_vars.nonlocals)
self.code_scope.update(dict(closure_vars.globals))
self.code_scope['math'] = math
# 2. code lines
code_lines = self.code_lines
# code_lines = [f'def {self.func_name}({", ".join(self.arguments)}):']
code_lines.append(f' {constants.DT}_sqrt = {constants.DT} ** 0.5')
# 2.1 dg
# dg = g(x, t, *args)
all_dg = [f'{var}_dg' for var in self.variables]
code_lines.append(
f' {", ".join(all_dg)} = g({", ".join(self.variables + self.parameters)})'
)
code_lines.append(' ')
# 2.2 dW
noise_terms(code_lines, self.variables)
# 2.3 dgdW
# ----
# SCALAR_WIENER : dg * dW
# VECTOR_WIENER : math.sum(dg * dW, axis=-1)
if self.wiener_type == constants.SCALAR_WIENER:
for var in self.variables:
code_lines.append(f' {var}_dgdW = {var}_dg * {var}_dW')
else:
for var in self.variables:
code_lines.append(
f' {var}_dgdW = math.sum({var}_dg * {var}_dW, axis=-1)')
code_lines.append(' ')
# 2.4 new var
# ----
analysis = separate_variables(self.derivative[constants.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 = self.variables[vi]
diff_eq = analysis_by_sympy.SingleDiffEq(var_name=var_name,
variables=sd_variables,
expressions=expressions,
derivative_expr=key,
scope=self.code_scope,
func_name=self.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 = analysis_by_sympy.str2sympy(f_res.code).expr.expand()
s_df = sympy.Symbol(f"{f_res.var_name}")
code_lines.append(
f' {s_df.name} = {analysis_by_sympy.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} = {analysis_by_sympy.sympy2str(linear)}'
)
# linear exponential
code_lines.append(
f' {s_linear_exp.name} = math.exp({analysis_by_sympy.sympy2str(linear)} * {constants.DT})'
)
# df part
df_part = (s_linear_exp - 1) / s_linear * s_df
code_lines.append(
f' {s_df_part.name} = {analysis_by_sympy.sympy2str(df_part)}'
)
else:
# linear exponential
code_lines.append(
f' {s_linear_exp.name} = {constants.DT}_sqrt')
# df part
code_lines.append(
f' {s_df_part.name} = {s_df.name} * {constants.DT}')
# update expression
update = var + s_df_part
# The actual update step
code_lines.append(
f' {diff_eq.var_name}_new = {analysis_by_sympy.sympy2str(update)} + {var_name}_dgdW'
)
code_lines.append('')
# returns
new_vars = [f'{var}_new' for var in self.variables]
code_lines.append(f' return {", ".join(new_vars)}')
# return and compile
self.integral = utils.compile_code(
code_scope={k: v
for k, v in self.code_scope.items()},
code_lines=self.code_lines,
show_code=self.show_code,
func_name=self.func_name)
if hasattr(self.derivative[constants.F], '__self__'):
host = self.derivative[constants.F].__self__
self.integral = self.integral.__get__(host, host.__class__)
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 adaptive_rk_wrapper(f, dt, A, B1, B2, C, tol, adaptive, show_code,
var_type, im_return):
"""Adaptive Runge-Kutta numerical method for ordinary differential equations.
The embedded methods are designed to produce an estimate of the local
truncation error of a single Runge-Kutta step, and as result, allow to
control the error with adaptive stepsize. This is done by having two
methods in the tableau, one with order p and one with order :math:`p-1`.
The lower-order step is given by
.. math::
y^*_{n+1} = y_n + h\\sum_{i=1}^s b^*_i k_i,
where the :math:`k_{i}` are the same as for the higher order method. Then the error is
.. math::
e_{n+1} = y_{n+1} - y^*_{n+1} = h\\sum_{i=1}^s (b_i - b^*_i) k_i,
which is :math:`O(h^{p})`. The Butcher Tableau for this kind of method is extended to
give the values of :math:`b_{i}^{*}`
.. math::
\\begin{array}{c|cccc}
c_1 & a_{11} & a_{12}& \\dots & a_{1s}\\\\
c_2 & a_{21} & a_{22}& \\dots & a_{2s}\\\\
\\vdots & \\vdots & \\vdots& \\ddots& \\vdots\\\\
c_s & a_{s1} & a_{s2}& \\dots & a_{ss} \\\\
\\hline & b_1 & b_2 & \\dots & b_s\\\\
& b_1^* & b_2^* & \\dots & b_s^*\\\\
\\end{array}
Parameters
----------
f : callable
The derivative function.
show_code : bool
Whether show the formatted code.
dt : float
The numerical precision.
A : tuple, list
The A matrix in the Butcher tableau.
B1 : tuple, list
The B1 vector in the Butcher tableau.
B2 : tuple, list
The B2 vector in the Butcher tableau.
C : tuple, list
The C vector in the Butcher tableau.
adaptive : bool
tol : float
var_type : str
Returns
-------
integral_func : callable
The one-step numerical integration function.
"""
if var_type not in constants.SUPPORTED_VAR_TYPE:
raise errors.IntegratorError(
f'"var_type" only supports {constants.SUPPORTED_VAR_TYPE}, not {var_type}.'
)
class_kw, variables, parameters, arguments = utils.get_args(f)
dt_var = 'dt'
func_name = Tools.f_names(f)
if adaptive:
# code scope
code_scope = {'f': f, 'tol': tol}
arguments = list(arguments) + [f'dt={dt}']
else:
# code scope
code_scope = {'f': f, 'dt': dt}
# code lines
code_lines = [f'def {func_name}({", ".join(arguments)}):']
# stage steps
Tools.step(variables, dt_var, A, C, code_lines, parameters)
# variable update
return_args = Tools.update(variables, dt_var, B1, code_lines)
# error adaptive item
if adaptive:
errors_ = []
for v in variables:
result = []
for i, (b1, b2) in enumerate(zip(B1, B2)):
if isinstance(b1, str):
b1 = eval(b1)
if isinstance(b2, str):
b2 = eval(b2)
diff = b1 - b2
if diff != 0.:
result.append(f'd{v}_k{i + 1} * {dt_var} * {diff}')
if len(result) > 0:
if var_type == constants.SCALAR_VAR:
code_lines.append(f' {v}_te = abs({" + ".join(result)})')
else:
code_lines.append(
f' {v}_te = sum(abs({" + ".join(result)}))')
errors_.append(f'{v}_te')
if len(errors_) > 0:
code_lines.append(f' error = {" + ".join(errors_)}')
code_lines.append(f' if error > tol:')
code_lines.append(
f' {dt_var}_new = 0.9 * {dt_var} * (tol / error) ** 0.2')
code_lines.append(f' else:')
code_lines.append(f' {dt_var}_new = {dt_var}')
return_args.append(f'{dt_var}_new')
# returns
code_lines.append(f' return {", ".join(return_args)}')
# compilation
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)