"""
Numerical integration functions.
"""
import operator
import string
from functools import reduce
import sympy
from pyparsing import (
Group,
Literal,
ParseException,
Suppress,
Word,
ZeroOrMore,
restOfLine,
)
from sympy.core.sympify import SympifyError
from brian2.equations.codestrings import is_constant_over_dt
from brian2.parsing.sympytools import str_to_sympy, sympy_to_str
from .base import (
StateUpdateMethod,
UnsupportedEquationsException,
extract_method_options,
)
__all__ = ["milstein", "heun", "euler", "rk2", "rk4", "ExplicitStateUpdater"]
# ===============================================================================
# Class for simple definition of explicit state updaters
# ===============================================================================
def _symbol(name, positive=None):
"""Shorthand for ``sympy.Symbol(name, real=True)``."""
return sympy.Symbol(name, real=True, positive=positive)
#: reserved standard symbols
SYMBOLS = {
"__x": _symbol("__x"),
"__t": _symbol("__t", positive=True),
"dt": _symbol("dt", positive=True),
"t": _symbol("t", positive=True),
"__f": sympy.Function("__f"),
"__g": sympy.Function("__g"),
"__dW": _symbol("__dW"),
}
[docs]
def split_expression(expr):
"""
Split an expression into a part containing the function ``f`` and another
one containing the function ``g``. Returns a tuple of the two expressions
(as sympy expressions).
Parameters
----------
expr : str
An expression containing references to functions ``f`` and ``g``.
Returns
-------
(non_stochastic, stochastic) : tuple of sympy expressions
A pair of expressions representing the non-stochastic (containing
function-independent terms and terms involving ``f``) and the
stochastic part of the expression (terms involving ``g`` and/or ``dW``).
Examples
--------
>>> split_expression('dt * __f(__x, __t)')
(dt*__f(__x, __t), None)
>>> split_expression('dt * __f(__x, __t) + __dW * __g(__x, __t)')
(dt*__f(__x, __t), __dW*__g(__x, __t))
>>> split_expression('1/(2*sqrt(dt))*(__g_support - __g(__x, __t))*(sqrt(__dW))')
(0, sqrt(__dW)*__g_support/(2*sqrt(dt)) - sqrt(__dW)*__g(__x, __t)/(2*sqrt(dt)))
"""
f = SYMBOLS["__f"]
g = SYMBOLS["__g"]
dW = SYMBOLS["__dW"]
# Arguments of the f and g functions
x_f = sympy.Wild("x_f", exclude=[f, g], real=True)
t_f = sympy.Wild("t_f", exclude=[f, g], real=True)
x_g = sympy.Wild("x_g", exclude=[f, g], real=True)
t_g = sympy.Wild("t_g", exclude=[f, g], real=True)
# Reorder the expression so that f(x,t) and g(x,t) are factored out
sympy_expr = sympy.sympify(expr, locals=SYMBOLS).expand()
sympy_expr = sympy.collect(sympy_expr, f(x_f, t_f))
sympy_expr = sympy.collect(sympy_expr, g(x_g, t_g))
# Constant part, contains neither f, g nor dW
independent = sympy.Wild("independent", exclude=[f, g, dW], real=True)
# The exponent of the random number
dW_exponent = sympy.Wild("dW_exponent", exclude=[f, g, dW, 0], real=True)
# The factor for the random number, not containing the g function
independent_dW = sympy.Wild("independent_dW", exclude=[f, g, dW], real=True)
# The factor for the f function
f_factor = sympy.Wild("f_factor", exclude=[f, g], real=True)
# The factor for the g function
g_factor = sympy.Wild("g_factor", exclude=[f, g], real=True)
match_expr = (
independent
+ f_factor * f(x_f, t_f)
+ independent_dW * dW**dW_exponent
+ g_factor * g(x_g, t_g)
)
matches = sympy_expr.match(match_expr)
if matches is None:
raise ValueError(
f'Expression "{sympy_expr}" in the state updater description could not be'
" parsed."
)
# Non-stochastic part
if x_f in matches:
# Includes the f function
non_stochastic = matches[independent] + (
matches[f_factor] * f(matches[x_f], matches[t_f])
)
else:
# Does not include f, might be 0
non_stochastic = matches[independent]
# Stochastic part
if independent_dW in matches and matches[independent_dW] != 0:
# includes a random variable term with a non-zero factor
stochastic = (
matches[g_factor] * g(matches[x_g], matches[t_g])
+ matches[independent_dW] * dW ** matches[dW_exponent]
)
elif x_g in matches:
# Does not include a random variable but the g function
stochastic = matches[g_factor] * g(matches[x_g], matches[t_g])
else:
# Contains neither random variable nor g function --> empty
stochastic = None
return (non_stochastic, stochastic)
[docs]
class ExplicitStateUpdater(StateUpdateMethod):
"""
An object that can be used for defining state updaters via a simple
description (see below). Resulting instances can be passed to the
``method`` argument of the `NeuronGroup` constructor. As other state
updater functions the `ExplicitStateUpdater` objects are callable,
returning abstract code when called with an `Equations` object.
A description of an explicit state updater consists of a (multi-line)
string, containing assignments to variables and a final "x_new = ...",
stating the integration result for a single timestep. The assignments
can be used to define an arbitrary number of intermediate results and
can refer to ``f(x, t)`` (the function being integrated, as a function of
``x``, the previous value of the state variable and ``t``, the time) and
``dt``, the size of the timestep.
For example, to define a Runge-Kutta 4 integrator (already provided as
`rk4`), use::
k1 = dt*f(x,t)
k2 = dt*f(x+k1/2,t+dt/2)
k3 = dt*f(x+k2/2,t+dt/2)
k4 = dt*f(x+k3,t+dt)
x_new = x+(k1+2*k2+2*k3+k4)/6
Note that for stochastic equations, the function `f` only corresponds to
the non-stochastic part of the equation. The additional function `g`
corresponds to the stochastic part that has to be multiplied with the
stochastic variable xi (a standard normal random variable -- if the
algorithm needs a random variable with a different variance/mean you have
to multiply/add it accordingly). Equations with more than one
stochastic variable do not have to be treated differently, the part
referring to ``g`` is repeated for all stochastic variables automatically.
Stochastic integrators can also make reference to ``dW`` (a normal
distributed random number with variance ``dt``) and ``g(x, t)``, the
stochastic part of an equation. A stochastic state updater could therefore
use a description like::
x_new = x + dt*f(x,t) + g(x, t) * dW
For simplicity, the same syntax is used for state updaters that only support
additive noise, even though ``g(x, t)`` does not depend on ``x`` or ``t``
in that case.
There a some restrictions on the complexity of the expressions (but most
can be worked around by using intermediate results as in the above Runge-
Kutta example): Every statement can only contain the functions ``f`` and
``g`` once; The expressions have to be linear in the functions, e.g. you
can use ``dt*f(x, t)`` but not ``f(x, t)**2``.
Parameters
----------
description : str
A state updater description (see above).
stochastic : {None, 'additive', 'multiplicative'}
What kind of stochastic equations this state updater supports: ``None``
means no support of stochastic equations, ``'additive'`` means only
equations with additive noise and ``'multiplicative'`` means
supporting arbitrary stochastic equations.
Raises
------
ValueError
If the parsing of the description failed.
Notes
-----
Since clocks are updated *after* the state update, the time ``t`` used
in the state update step is still at its previous value. Enumerating the
states and discrete times, ``x_new = x + dt*f(x, t)`` is therefore
understood as :math:`x_{i+1} = x_i + dt f(x_i, t_i)`, yielding the correct
forward Euler integration. If the integrator has to refer to the time at
the end of the timestep, simply use ``t + dt`` instead of ``t``.
See also
--------
euler, rk2, rk4, milstein
"""
# ===========================================================================
# Parsing definitions
# ===========================================================================
#: Legal names for temporary variables
TEMP_VAR = ~Literal("x_new") + Word(
f"{string.ascii_letters}_", f"{string.ascii_letters + string.digits}_"
).setResultsName("identifier")
#: A single expression
EXPRESSION = restOfLine.setResultsName("expression")
#: An assignment statement
STATEMENT = Group(TEMP_VAR + Suppress("=") + EXPRESSION).setResultsName("statement")
#: The last line of a state updater description
OUTPUT = Group(
Suppress(Literal("x_new")) + Suppress("=") + EXPRESSION
).setResultsName("output")
#: A complete state updater description
DESCRIPTION = ZeroOrMore(STATEMENT) + OUTPUT
def __init__(self, description, stochastic=None, custom_check=None):
self._description = description
self.stochastic = stochastic
self.custom_check = custom_check
try:
parsed = ExplicitStateUpdater.DESCRIPTION.parseString(
description, parseAll=True
)
except ParseException as p_exc:
ex = SyntaxError(f"Parsing failed: {str(p_exc.msg)}")
ex.text = str(p_exc.line)
ex.offset = p_exc.column
ex.lineno = p_exc.lineno
raise ex
self.statements = []
self.symbols = SYMBOLS.copy()
for element in parsed:
expression = str_to_sympy(element.expression)
# Replace all symbols used in state updater expressions by unique
# names that cannot clash with user-defined variables or functions
expression = expression.subs(sympy.Function("f"), self.symbols["__f"])
expression = expression.subs(sympy.Function("g"), self.symbols["__g"])
symbols = list(expression.atoms(sympy.Symbol))
unique_symbols = []
for symbol in symbols:
if symbol.name == "dt":
unique_symbols.append(symbol)
else:
unique_symbols.append(_symbol(f"__{symbol.name}"))
for symbol, unique_symbol in zip(symbols, unique_symbols):
expression = expression.subs(symbol, unique_symbol)
self.symbols.update({symbol.name: symbol for symbol in unique_symbols})
if element.getName() == "statement":
self.statements.append((f"__{element.identifier}", expression))
elif element.getName() == "output":
self.output = expression
else:
raise AssertionError(f"Unknown element name: {element.getName()}")
def __repr__(self):
# recreate a description string
description = "\n".join([f"{var} = {expr}" for var, expr in self.statements])
if len(description):
description += "\n"
description += f"x_new = {str(self.output)}"
classname = self.__class__.__name__
return f"{classname}('''{description}''', stochastic={self.stochastic!r})"
def __str__(self):
s = f"{self.__class__.__name__}\n"
if len(self.statements) > 0:
s += "Intermediate statements:\n"
s += "\n".join(
[f"{var} = {sympy_to_str(expr)}" for var, expr in self.statements]
)
s += "\n"
s += "Output:\n"
s += sympy_to_str(self.output)
return s
def _latex(self, *args):
from sympy import Symbol, latex
s = [r"\begin{equation}"]
for var, expr in self.statements:
expr = expr.subs(Symbol("x"), Symbol("x_t"))
s.append(f"{latex(Symbol(var))} = {latex(expr)}\\\\")
expr = self.output.subs(Symbol("x"), "x_t")
s.append(f"x_{{t+1}} = {latex(expr)}")
s.append(r"\end{equation}")
return "\n".join(s)
def _repr_latex_(self):
return self._latex()
[docs]
def replace_func(self, x, t, expr, temp_vars, eq_symbols, stochastic_variable=None):
"""
Used to replace a single occurance of ``f(x, t)`` or ``g(x, t)``:
`expr` is the non-stochastic (in the case of ``f``) or stochastic
part (``g``) of the expression defining the right-hand-side of the
differential equation describing `var`. It replaces the variable
`var` with the value given as `x` and `t` by the value given for
`t`. Intermediate variables will be replaced with the appropriate
replacements as well.
For example, in the `rk2` integrator, the second step involves the
calculation of ``f(k/2 + x, dt/2 + t)``. If `var` is ``v`` and
`expr` is ``-v / tau``, this will result in ``-(_k_v/2 + v)/tau``.
Note that this deals with only one state variable `var`, given as
an argument to the surrounding `_generate_RHS` function.
"""
try:
s_expr = str_to_sympy(str(expr))
except SympifyError as ex:
raise ValueError(f'Error parsing the expression "{expr}": {str(ex)}')
for var in eq_symbols:
# Generate specific temporary variables for the state variable,
# e.g. '_k_v' for the state variable 'v' and the temporary
# variable 'k'.
if stochastic_variable is None:
temp_var_replacements = {
self.symbols[temp_var]: _symbol(f"{temp_var}_{var}")
for temp_var in temp_vars
}
else:
temp_var_replacements = {
self.symbols[temp_var]: _symbol(
f"{temp_var}_{var}_{stochastic_variable}"
)
for temp_var in temp_vars
}
# In the expression given as 'x', replace 'x' by the variable
# 'var' and all the temporary variables by their
# variable-specific counterparts.
x_replacement = x.subs(self.symbols["__x"], eq_symbols[var])
x_replacement = x_replacement.subs(temp_var_replacements)
# Replace the variable `var` in the expression by the new `x`
# expression
s_expr = s_expr.subs(eq_symbols[var], x_replacement)
# If the expression given for t in the state updater description
# is not just "t" (or rather "__t"), then replace t in the
# equations by it, and replace "__t" by "t" afterwards.
if t != self.symbols["__t"]:
s_expr = s_expr.subs(SYMBOLS["t"], t)
s_expr = s_expr.replace(self.symbols["__t"], SYMBOLS["t"])
return s_expr
def _non_stochastic_part(
self,
eq_symbols,
non_stochastic,
non_stochastic_expr,
stochastic_variable,
temp_vars,
var,
):
non_stochastic_results = []
if stochastic_variable is None or len(stochastic_variable) == 0:
# Replace the f(x, t) part
replace_f = lambda x, t: self.replace_func(
x, t, non_stochastic, temp_vars, eq_symbols
)
non_stochastic_result = non_stochastic_expr.replace(
self.symbols["__f"], replace_f
)
# Replace x by the respective variable
non_stochastic_result = non_stochastic_result.subs(
self.symbols["__x"], eq_symbols[var]
)
# Replace intermediate variables
temp_var_replacements = {
self.symbols[temp_var]: _symbol(f"{temp_var}_{var}")
for temp_var in temp_vars
}
non_stochastic_result = non_stochastic_result.subs(temp_var_replacements)
non_stochastic_results.append(non_stochastic_result)
elif isinstance(stochastic_variable, str):
# Replace the f(x, t) part
replace_f = lambda x, t: self.replace_func(
x, t, non_stochastic, temp_vars, eq_symbols, stochastic_variable
)
non_stochastic_result = non_stochastic_expr.replace(
self.symbols["__f"], replace_f
)
# Replace x by the respective variable
non_stochastic_result = non_stochastic_result.subs(
self.symbols["__x"], eq_symbols[var]
)
# Replace intermediate variables
temp_var_replacements = {
self.symbols[temp_var]: _symbol(
f"{temp_var}_{var}_{stochastic_variable}"
)
for temp_var in temp_vars
}
non_stochastic_result = non_stochastic_result.subs(temp_var_replacements)
non_stochastic_results.append(non_stochastic_result)
else:
# Replace the f(x, t) part
replace_f = lambda x, t: self.replace_func(
x, t, non_stochastic, temp_vars, eq_symbols
)
non_stochastic_result = non_stochastic_expr.replace(
self.symbols["__f"], replace_f
)
# Replace x by the respective variable
non_stochastic_result = non_stochastic_result.subs(
self.symbols["__x"], eq_symbols[var]
)
# Replace intermediate variables
temp_var_replacements = {
self.symbols[temp_var]: reduce(
operator.add,
[_symbol(f"{temp_var}_{var}_{xi}") for xi in stochastic_variable],
)
for temp_var in temp_vars
}
non_stochastic_result = non_stochastic_result.subs(temp_var_replacements)
non_stochastic_results.append(non_stochastic_result)
return non_stochastic_results
def _stochastic_part(
self,
eq_symbols,
stochastic,
stochastic_expr,
stochastic_variable,
temp_vars,
var,
):
stochastic_results = []
if isinstance(stochastic_variable, str):
# Replace the g(x, t) part
replace_f = lambda x, t: self.replace_func(
x,
t,
stochastic.get(stochastic_variable, 0),
temp_vars,
eq_symbols,
stochastic_variable,
)
stochastic_result = stochastic_expr.replace(self.symbols["__g"], replace_f)
# Replace x by the respective variable
stochastic_result = stochastic_result.subs(
self.symbols["__x"], eq_symbols[var]
)
# Replace dW by the respective variable
stochastic_result = stochastic_result.subs(
self.symbols["__dW"], stochastic_variable
)
# Replace intermediate variables
temp_var_replacements = {
self.symbols[temp_var]: _symbol(
f"{temp_var}_{var}_{stochastic_variable}"
)
for temp_var in temp_vars
}
stochastic_result = stochastic_result.subs(temp_var_replacements)
stochastic_results.append(stochastic_result)
else:
for xi in stochastic_variable:
# Replace the g(x, t) part
replace_f = lambda x, t: self.replace_func(
x, t, stochastic.get(xi, 0), temp_vars, eq_symbols, xi # noqa: B023
)
stochastic_result = stochastic_expr.replace(
self.symbols["__g"], replace_f
)
# Replace x by the respective variable
stochastic_result = stochastic_result.subs(
self.symbols["__x"], eq_symbols[var]
)
# Replace dW by the respective variable
stochastic_result = stochastic_result.subs(self.symbols["__dW"], xi)
# Replace intermediate variables
temp_var_replacements = {
self.symbols[temp_var]: _symbol(f"{temp_var}_{var}_{xi}")
for temp_var in temp_vars
}
stochastic_result = stochastic_result.subs(temp_var_replacements)
stochastic_results.append(stochastic_result)
return stochastic_results
def _generate_RHS(
self,
eqs,
var,
eq_symbols,
temp_vars,
expr,
non_stochastic_expr,
stochastic_expr,
stochastic_variable=(),
):
"""
Helper function used in `__call__`. Generates the right hand side of
an abstract code statement by appropriately replacing f, g and t.
For example, given a differential equation ``dv/dt = -(v + I) / tau``
(i.e. `var` is ``v` and `expr` is ``(-v + I) / tau``) together with
the `rk2` step ``return x + dt*f(x + k/2, t + dt/2)``
(i.e. `non_stochastic_expr` is
``x + dt*f(x + k/2, t + dt/2)`` and `stochastic_expr` is ``None``),
produces ``v + dt*(-v - _k_v/2 + I + _k_I/2)/tau``.
"""
# Note: in the following we are silently ignoring the case that a
# state updater does not care about either the non-stochastic or the
# stochastic part of an equation. We do trust state updaters to
# correctly specify their own abilities (i.e. they do not claim to
# support stochastic equations but actually just ignore the stochastic
# part). We can't really check the issue here, as we are only dealing
# with one line of the state updater description. It is perfectly valid
# to write the euler update as:
# non_stochastic = dt * f(x, t)
# stochastic = dt**.5 * g(x, t) * xi
# return x + non_stochastic + stochastic
#
# In the above case, we'll deal with lines which do not define either
# the stochastic or the non-stochastic part.
non_stochastic, stochastic = expr.split_stochastic()
if non_stochastic_expr is not None:
# We do have a non-stochastic part in the state updater description
non_stochastic_results = self._non_stochastic_part(
eq_symbols,
non_stochastic,
non_stochastic_expr,
stochastic_variable,
temp_vars,
var,
)
else:
non_stochastic_results = []
if not (stochastic is None or stochastic_expr is None):
# We do have a stochastic part in the state
# updater description
stochastic_results = self._stochastic_part(
eq_symbols,
stochastic,
stochastic_expr,
stochastic_variable,
temp_vars,
var,
)
else:
stochastic_results = []
RHS = sympy.Number(0)
# All the parts (one non-stochastic and potentially more than one
# stochastic part) are combined with addition
for non_stochastic_result in non_stochastic_results:
RHS += non_stochastic_result
for stochastic_result in stochastic_results:
RHS += stochastic_result
return sympy_to_str(RHS)
[docs]
def __call__(self, eqs, variables=None, method_options=None):
"""
Apply a state updater description to model equations.
Parameters
----------
eqs : `Equations`
The equations describing the model
variables: dict-like, optional
The `Variable` objects for the model. Ignored by the explicit
state updater.
method_options : dict, optional
Additional options to the state updater (not used at the moment
for the explicit state updaters).
Examples
--------
>>> from brian2 import *
>>> eqs = Equations('dv/dt = -v / tau : volt')
>>> print(euler(eqs))
_v = -dt*v/tau + v
v = _v
>>> print(rk4(eqs))
__k_1_v = -dt*v/tau
__k_2_v = -dt*(__k_1_v/2 + v)/tau
__k_3_v = -dt*(__k_2_v/2 + v)/tau
__k_4_v = -dt*(__k_3_v + v)/tau
_v = __k_1_v/6 + __k_2_v/3 + __k_3_v/3 + __k_4_v/6 + v
v = _v
"""
extract_method_options(method_options, {})
# Non-stochastic numerical integrators should work for all equations,
# except for stochastic equations
if eqs.is_stochastic and self.stochastic is None:
raise UnsupportedEquationsException(
"Cannot integrate stochastic equations with this state updater."
)
if self.custom_check:
self.custom_check(eqs, variables)
# The final list of statements
statements = []
stochastic_variables = eqs.stochastic_variables
# The variables for the intermediate results in the state updater
# description, e.g. the variable k in rk2
intermediate_vars = [var for var, expr in self.statements]
# A dictionary mapping all the variables in the equations to their
# sympy representations
eq_variables = {var: _symbol(var) for var in eqs.eq_names}
# Generate the random numbers for the stochastic variables
for stochastic_variable in stochastic_variables:
statements.append(f"{stochastic_variable} = dt**.5 * randn()")
substituted_expressions = eqs.get_substituted_expressions(variables)
# Process the intermediate statements in the stateupdater description
for intermediate_var, intermediate_expr in self.statements:
# Split the expression into a non-stochastic and a stochastic part
non_stochastic_expr, stochastic_expr = split_expression(intermediate_expr)
# Execute the statement by appropriately replacing the functions f
# and g and the variable x for every equation in the model.
# We use the model equations where the subexpressions have
# already been substituted into the model equations.
for var, expr in substituted_expressions:
for xi in stochastic_variables:
RHS = self._generate_RHS(
eqs,
var,
eq_variables,
intermediate_vars,
expr,
non_stochastic_expr,
stochastic_expr,
xi,
)
statements.append(f"{intermediate_var}_{var}_{xi} = {RHS}")
if not stochastic_variables: # no stochastic variables
RHS = self._generate_RHS(
eqs,
var,
eq_variables,
intermediate_vars,
expr,
non_stochastic_expr,
stochastic_expr,
)
statements.append(f"{intermediate_var}_{var} = {RHS}")
# Process the "return" line of the stateupdater description
non_stochastic_expr, stochastic_expr = split_expression(self.output)
if eqs.is_stochastic and (
self.stochastic != "multiplicative"
and eqs.stochastic_type == "multiplicative"
):
# The equations are marked as having multiplicative noise and the
# current state updater does not support such equations. However,
# it is possible that the equations do not use multiplicative noise
# at all. They could depend on time via a function that is constant
# over a single time step (most likely, a TimedArray). In that case
# we can integrate the equations
dt_value = variables["dt"].get_value()[0] if "dt" in variables else None
for _, expr in substituted_expressions:
_, stoch = expr.split_stochastic()
if stoch is None:
continue
# There could be more than one stochastic variable (e.g. xi_1, xi_2)
for _, stoch_expr in stoch.items():
sympy_expr = str_to_sympy(stoch_expr.code)
# The equation really has multiplicative noise, if it depends
# on time (and not only via a function that is constant
# over dt), or if it depends on another variable defined
# via differential equations.
if not is_constant_over_dt(sympy_expr, variables, dt_value) or len(
stoch_expr.identifiers & eqs.diff_eq_names
):
raise UnsupportedEquationsException(
"Cannot integrate "
"equations with "
"multiplicative noise with "
"this state updater."
)
# Assign a value to all the model variables described by differential
# equations
for var, expr in substituted_expressions:
RHS = self._generate_RHS(
eqs,
var,
eq_variables,
intermediate_vars,
expr,
non_stochastic_expr,
stochastic_expr,
stochastic_variables,
)
statements.append(f"_{var} = {RHS}")
# Assign everything to the final variables
for var, _ in substituted_expressions:
statements.append(f"{var} = _{var}")
return "\n".join(statements)
# ===============================================================================
# Excplicit state updaters
# ===============================================================================
# these objects can be used like functions because they are callable
#: Forward Euler state updater
euler = ExplicitStateUpdater(
"x_new = x + dt * f(x,t) + g(x,t) * dW", stochastic="additive"
)
#: Second order Runge-Kutta method (midpoint method)
rk2 = ExplicitStateUpdater(
"""
k = dt * f(x,t)
x_new = x + dt*f(x + k/2, t + dt/2)"""
)
#: Classical Runge-Kutta method (RK4)
rk4 = ExplicitStateUpdater(
"""
k_1 = dt*f(x,t)
k_2 = dt*f(x+k_1/2,t+dt/2)
k_3 = dt*f(x+k_2/2,t+dt/2)
k_4 = dt*f(x+k_3,t+dt)
x_new = x+(k_1+2*k_2+2*k_3+k_4)/6
"""
)
[docs]
def diagonal_noise(equations, variables):
"""
Checks whether we deal with diagonal noise, i.e. one independent noise
variable per variable.
Raises
------
UnsupportedEquationsException
If the noise is not diagonal.
"""
if not equations.is_stochastic:
return
stochastic_vars = []
for _, expr in equations.get_substituted_expressions(variables):
expr_stochastic_vars = expr.stochastic_variables
if len(expr_stochastic_vars) > 1:
# More than one stochastic variable --> no diagonal noise
raise UnsupportedEquationsException(
"Cannot integrate stochastic "
"equations with non-diagonal "
"noise with this state "
"updater."
)
stochastic_vars.extend(expr_stochastic_vars)
# If there's no stochastic variable is used in more than one equation, we
# have diagonal noise
if len(stochastic_vars) != len(set(stochastic_vars)):
raise UnsupportedEquationsException(
"Cannot integrate stochastic "
"equations with non-diagonal "
"noise with this state "
"updater."
)
#: Derivative-free Milstein method
milstein = ExplicitStateUpdater(
"""
x_support = x + dt*f(x, t) + dt**.5 * g(x, t)
g_support = g(x_support, t)
k = 1/(2*dt**.5)*(g_support - g(x, t))*(dW**2)
x_new = x + dt*f(x,t) + g(x, t) * dW + k
""",
stochastic="multiplicative",
custom_check=diagonal_noise,
)
#: Stochastic Heun method (for multiplicative Stratonovic SDEs with non-diagonal
#: diffusion matrix)
heun = ExplicitStateUpdater(
"""
x_support = x + g(x,t) * dW
g_support = g(x_support,t+dt)
x_new = x + dt*f(x,t) + .5*dW*(g(x,t)+g_support)
""",
stochastic="multiplicative",
)