"""
Module defining `CodeString`, a class for a string of code together with
information about its namespace. Only serves as a parent class, its subclasses
`Expression` and `Statements` are the ones that are actually used.
"""
from collections.abc import Hashable
import sympy
from brian2.parsing.sympytools import str_to_sympy, sympy_to_str
from brian2.utils.logger import get_logger
from brian2.utils.stringtools import get_identifiers
__all__ = ["Expression", "Statements"]
logger = get_logger(__name__)
[docs]
class CodeString(Hashable):
"""
A class for representing "code strings", i.e. a single Python expression
or a sequence of Python statements.
Parameters
----------
code : str
The code string, may be an expression or a statement(s) (possibly
multi-line).
"""
def __init__(self, code):
self._code = code.strip()
# : Set of identifiers in the code string
self.identifiers = get_identifiers(code)
code = property(lambda self: self._code, doc="The code string")
def __str__(self):
return self.code
def __repr__(self):
return f"{self.__class__.__name__}({self.code!r})"
def __eq__(self, other):
if not isinstance(other, CodeString):
return NotImplemented
return self.code == other.code
def __ne__(self, other):
return not self == other
def __hash__(self):
return hash(self.code)
[docs]
class Statements(CodeString):
"""
Class for representing statements.
Parameters
----------
code : str
The statement or statements. Several statements can be given as a
multi-line string or separated by semicolons.
Notes
-----
Currently, the implementation of this class does not add anything to
`~brian2.equations.codestrings.CodeString`, but it should be used instead
of that class for clarity and to allow for future functionality that is
only relevant to statements and not to expressions.
"""
pass
[docs]
class Expression(CodeString):
"""
Class for representing an expression.
Parameters
----------
code : str, optional
The expression. Note that the expression has to be written in a form
that is parseable by sympy. Alternatively, a sympy expression can be
provided (in the ``sympy_expression`` argument).
sympy_expression : sympy expression, optional
A sympy expression. Alternatively, a plain string expression can be
provided (in the ``code`` argument).
"""
def __init__(self, code=None, sympy_expression=None):
if code is None and sympy_expression is None:
raise TypeError("Have to provide either a string or a sympy expression")
if code is not None and sympy_expression is not None:
raise TypeError(
"Provide a string expression or a sympy expression, not both"
)
if code is None:
code = sympy_to_str(sympy_expression)
else:
# Just try to convert it to a sympy expression to get syntax errors
# for incorrect expressions
str_to_sympy(code)
super().__init__(code=code)
stochastic_variables = property(
lambda self: {
variable
for variable in self.identifiers
if variable == "xi" or variable.startswith("xi_")
},
doc="Stochastic variables in this expression",
)
[docs]
def split_stochastic(self):
"""
Split the expression into a stochastic and non-stochastic part.
Splits the expression into a tuple of one `Expression` objects f (the
non-stochastic part) and a dictionary mapping stochastic variables
to `Expression` objects. For example, an expression of the form
``f + g * xi_1 + h * xi_2`` would be returned as:
``(f, {'xi_1': g, 'xi_2': h})``
Note that the `Expression` objects for the stochastic parts do not
include the stochastic variable itself.
Returns
-------
(f, d) : (`Expression`, dict)
A tuple of an `Expression` object and a dictionary, the first
expression being the non-stochastic part of the equation and
the dictionary mapping stochastic variables (``xi`` or starting
with ``xi_``) to `Expression` objects. If no stochastic variable
is present in the code string, a tuple ``(self, None)`` will be
returned with the unchanged `Expression` object.
"""
stochastic_variables = []
for identifier in self.identifiers:
if identifier == "xi" or identifier.startswith("xi_"):
stochastic_variables.append(identifier)
# No stochastic variable
if not len(stochastic_variables):
return (self, None)
stochastic_symbols = [
sympy.Symbol(variable, real=True) for variable in stochastic_variables
]
# Note that collect only works properly if the expression is expanded
collected = (
str_to_sympy(self.code).expand().collect(stochastic_symbols, evaluate=False)
)
f_expr = None
stochastic_expressions = {}
for var, s_expr in collected.items():
expr = Expression(sympy_expression=s_expr)
if var == 1:
if any(s_expr.has(s) for s in stochastic_symbols):
raise AssertionError(
"Error when separating expression "
f"'{self.code}' into stochastic and non-"
"stochastic term: non-stochastic "
f"part was determined to be '{s_expr}' but "
"contains a stochastic symbol."
)
f_expr = expr
elif var in stochastic_symbols:
stochastic_expressions[str(var)] = expr
else:
raise ValueError(
f"Expression '{self.code}' cannot be separated into "
"stochastic and non-stochastic "
"term"
)
if f_expr is None:
f_expr = Expression("0.0")
return f_expr, stochastic_expressions
def _repr_pretty_(self, p, cycle):
"""
Pretty printing for ipython.
"""
if cycle:
raise AssertionError("Cyclical call of 'CodeString._repr_pretty'")
# Make use of sympy's pretty printing
p.pretty(str_to_sympy(self.code))
def __eq__(self, other):
if not isinstance(other, Expression):
return NotImplemented
return self.code == other.code
def __ne__(self, other):
return not self.__eq__(other)
def __hash__(self):
return hash(self.code)
[docs]
def is_constant_over_dt(expression, variables, dt_value):
"""
Check whether an expression can be considered as constant over a time step.
This is *not* the case when the expression either:
1. contains the variable ``t`` (except as the argument of a function that
can be considered as constant over a time step, e.g. a `TimedArray` with
a dt equal to or greater than the dt used to evaluate this expression)
2. refers to a stateful function such as ``rand()``.
Parameters
----------
expression : `sympy.Expr`
The (sympy) expression to analyze
variables : dict
The variables dictionary.
dt_value : float or None
The length of a timestep (without units), can be ``None`` if the
time step is not yet known.
Returns
-------
is_constant : bool
Whether the expression can be considered to be constant over a time
step.
"""
t_symbol = sympy.Symbol("t", real=True, positive=True)
if expression == t_symbol:
return False # The full expression is simply "t"
func_name = str(expression.func)
func_variable = variables.get(func_name, None)
if func_variable is not None and not func_variable.stateless:
return False
for arg in expression.args:
if arg == t_symbol and dt_value is not None:
# We found "t" -- if it is not the only argument of a locally
# constant function we bail out
if not (
func_variable is not None
and func_variable.is_locally_constant(dt_value)
):
return False
else:
if not is_constant_over_dt(arg, variables, dt_value):
return False
return True