Source code for brian2.equations.codestrings

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 import Hashable

import sympy

from brian2.utils.logger import get_logger
from brian2.utils.stringtools import get_identifiers
from brian2.parsing.sympytools import str_to_sympy, sympy_to_str

__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(Expression, self).__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(f"Error when separating expression " f"'{self.code}' into stochastic and non-" f"stochastic term: non-stochastic " f"part was determined to be '{s_expr}' but " f"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 " f"stochastic and non-stochastic " f"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