[WIP] Add prefix/postfix method to Equations

brian2/equations/equations.py

@@ -7,7 +7,7 @@
 import string
 
 import sympy
-from brian2.utils.stringtools import get_identifiers
+from brian2.utils.stringtools import get_identifiers, word_substitute
 from pyparsing import (Group, ZeroOrMore, OneOrMore, Optional, Word, CharsNotIn,
                        Combine, Suppress, restOfLine, LineEnd, ParseException)
 
@@ -536,10 +536,12 @@ class Equations(collections.Hashable, collections.Mapping):
 
     Parameters
     ----------
-    eqs : `str` or list of `SingleEquation` objects
-        A multiline string of equations (see above) -- for internal purposes
-        also a list of `SingleEquation` objects can be given. This is done for
-        example when adding new equations to implement the refractory
+    eqs : `str`, list of `SingleEquation`, or `Equations`
+        A multiline string of equations (see above). To make a copy of an
+        existing `Equations` object (with potential changes in the variable
+        names, see below), an existing `Equations` object can be provided as
+        well. Finally, a list of `SingleEquation` objects can be given. This is
+        done for example when adding new equations to implement the refractory
         mechanism. Note that in this case the variable names are not checked
         to allow for "internal names", starting with an underscore.
     kwds: keyword arguments
@@ -552,6 +554,8 @@ class Equations(collections.Hashable, collections.Mapping):
     """
 
     def __init__(self, eqns, **kwds):
+        if isinstance(eqns, Equations):
+            eqns = eqns._equations.values()
         if isinstance(eqns, basestring):
             self._equations = parse_string_equations(eqns)
             # Do a basic check for the identifiers
@@ -780,6 +784,177 @@ def get_substituted_expressions(self, variables=None,
             return [(name, expr) for name, expr in self._substituted_expressions
                     if self[name].type == DIFFERENTIAL_EQUATION]
 
+    def _prefix_postfix(self, use_prefix, prepostfix, internal=True,
+                        external=False, exclude=None):
+        '''
+        Helper method to implement `~Equations.prefix` and `~Equations.postfix`
+        without duplicating the code. See the respective documentation.
+        '''
+        if internal is True:
+            internal = self.names
+        elif internal is False:
+            internal = []
+        if exclude is None:
+            exclude = []
+        if use_prefix:
+            substitutions = {n: prepostfix + n for n in internal
+                             if n not in exclude}
+        else:
+            substitutions = {n: n + prepostfix for n in internal
+                             if n not in exclude}
+        eq_substituted_vars = Equations(self, **substitutions)
+        if external is True:
+            external_names = {n for n in self.identifiers
+                              if n not in self.names}
+            external = external_names
+        elif external is False:
+            external = []
+        if use_prefix:
+            substitutions = {n: prepostfix + n for n in external
+                             if n not in exclude}
+        else:
+            substitutions = {n: n + prepostfix for n in external
+                             if n not in exclude}
+        single_equations = eq_substituted_vars._equations.values()
+        new_single_equations = []
+        for single_eq in single_equations:
+            if single_eq.type in (DIFFERENTIAL_EQUATION, SUBEXPRESSION):
+                expr = single_eq.expr.code
+                new_expr = Expression(word_substitute(expr, substitutions))
+                new_single_equation = SingleEquation(varname=single_eq.varname,
+                                                     type=single_eq.type,
+                                                     dimensions=single_eq.dim,
+                                                     var_type=single_eq.var_type,
+                                                     expr=new_expr,
+                                                     flags=single_eq.flags)
+            else:
+                new_single_equation = single_eq
+            new_single_equations.append(new_single_equation)
+
+        return Equations(new_single_equations)
+
+    def prefix(self, prefix, internal=True, external=False, exclude=None):
+        '''
+        Return a copy of the equations where variables and/or constants are
+        prefixed with a given string.
+
+        Parameters
+        ----------
+        prefix : str
+            The string to use as a prefix (no separator character is added,
+            i.e. to transform ``'v'`` into ``'soma_v'``, the prefix has to be
+            ``'soma_'``).
+        internal : bool or list of str, optional
+            A list of variable names (variables defined by differential
+            equations, neuron-/synapse-specific parameters, and subexpressions)
+            that should be prefixed. Alternatively, ``True`` can be used to
+            prefix all variables, or ``False`` to prefix none. By default, all
+            names are prefixed.
+        external : bool or list of str, optional
+            A list of constant or external variable names (i.e., all names that
+            are not defined as part of the equations) that should be prefixed.
+            Alternatively, ``True`` can be used to prefix all variables and
+            constants, or ``False`` to prefix none. By default, no external
+            names are prefixed.
+        exclude : list of str or None, optional
+            A list of names that should not be prefixed. Can be used together
+            with ``internal=True`` or ``external=True``, when it is easier to
+            specify which names should not be prefixed instead of listing the
+            names that should be. This argument applies both to internal and
+            external variables. Can be set to ``None`` (the default value) to
+            not exclude any names.
+
+        Returns
+        -------
+        new_equations : `Equations`
+            The prefixed equations
+
+        Examples
+        --------
+
+        >>> from brian2 import *
+        >>> eqs = Equations("""
+        ... I = g*m**3*h*(E - v) : amp
+        ... dm/dt=q10*(minf - m) / mtau : 1
+        ... dh/dt=q10*(hinf - h) / htau : 1""")
+        >>> print(eqs.prefix('Na_'))
+        Na_I = g*Na_m**3*Na_h*(E - v) : amp
+        dNa_h/dt = q10*(hinf - Na_h) / htau : 1
+        dNa_m/dt = q10*(minf - Na_m) / mtau : 1
+        >>> print(eqs.prefix('Na_', external=['E']))
+        Na_I = g*Na_m**3*Na_h*(Na_E - v) : amp
+        dNa_h/dt = q10*(hinf - Na_h) / htau : 1
+        dNa_m/dt = q10*(minf - Na_m) / mtau : 1
+        >>> print(eqs.prefix('Na_', internal=['I'], external=True))
+        Na_I = Na_g*m**3*h*(Na_E - Na_v) : amp
+        dh/dt = Na_q10*(Na_hinf - h) / Na_htau : 1
+        dm/dt = Na_q10*(Na_minf - m) / Na_mtau : 1
+        '''
+        return self._prefix_postfix(use_prefix=True, prepostfix=prefix,
+                                    internal=internal, external=external,
+                                    exclude=exclude)
+
+    def postfix(self, postfix, internal=True, external=False, exclude=None):
+        '''
+        Return a copy of the equations where variables and/or constants are
+        postfixed with a given string.
+
+        Parameters
+        ----------
+        postfix : str
+            The string to use as a postfix (no separator character is added,
+            i.e. to transform ``'v'`` into ``'v_soma'``, the postfix has to be
+            ``'_soma'``).
+        internal : bool or list of str, optional
+            A list of variable names (variables defined by differential
+            equations, neuron-/synapse-specific parameters, and subexpressions)
+            that should be postfixed. Alternatively, ``True`` can be used to
+            postfix all variables, or ``False`` to postfix none. By default, all
+            names are postfixed.
+        external : bool or list of str, optional
+            A list of constant or external variable names (i.e., all names that
+            are not defined as part of the equations) that should be postfixed.
+            Alternatively, ``True`` can be used to postfix all variables and
+            constants, or ``False`` to postfix none. By default, no external
+            names are postfixed.
+        exclude : list of str or None, optional
+            A list of names that should not be postfixed. Can be used together
+            with ``internal=True`` or ``external=True``, when it is easier to
+            specify which names should not be postfixed instead of listing the
+            names that should be. This argument applies both to internal and
+            external variables. Can be set to ``None`` (the default value) to
+            not exclude any names.
+
+        Returns
+        -------
+        new_equations : `Equations`
+            The postfixed equations
+
+        Examples
+        --------
+
+        >>> from brian2 import *
+        >>> eqs = Equations("""
+        ... I = g*m**3*h*(E - v) : amp
+        ... dm/dt=q10*(minf - m) / mtau : 1
+        ... dh/dt=q10*(hinf - h) / htau : 1""")
+        >>> print(eqs.postfix('_Na'))
+        I_Na = g*m_Na**3*h_Na*(E - v) : amp
+        dh_Na/dt = q10*(hinf - h_Na) / htau : 1
+        dm_Na/dt = q10*(minf - m_Na) / mtau : 1
+        >>> print(eqs.postfix('_Na', external=['E']))
+        I_Na = g*m_Na**3*h_Na*(E_Na - v) : amp
+        dh_Na/dt = q10*(hinf - h_Na) / htau : 1
+        dm_Na/dt = q10*(minf - m_Na) / mtau : 1
+        >>> print(eqs.postfix('_Na', internal=['I'], external=True))
+        I_Na = g_Na*m**3*h*(E_Na - v_Na) : amp
+        dh/dt = q10_Na*(hinf_Na - h) / htau_Na : 1
+        dm/dt = q10_Na*(minf_Na - m) / mtau_Na : 1
+        '''
+        return self._prefix_postfix(use_prefix=False, prepostfix=postfix,
+                                    internal=internal, external=external,
+                                    exclude=exclude)
+
     def _get_stochastic_type(self):
         '''
         Returns the type of stochastic differential equations (additivive or
@@ -822,7 +997,8 @@ def _get_stochastic_type(self):
 
     # Lists of equations or (variable, expression tuples)
     ordered = property(lambda self: sorted(self._equations.itervalues(),
-                                           key=lambda key: key.update_order),
+                                           key=lambda key: (key.update_order,
+                                                            key.varname)),
                                            doc='A list of all equations, sorted '
                                            'according to the order in which they should '
                                            'be updated')