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.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