conditional propagators

Author: C.A.P. Linssen

odetoolbox/__init__.py

@@ -182,6 +182,32 @@ def _get_all_first_order_variables(indict) -> Iterable[str]:
     return variable_names
 
 
+def symbol_appears_in_any_expr(param_name, solver_json) -> bool:
+    if "update_expressions" in solver_json.keys():
+        for sym, expr in solver_json["update_expressions"].items():
+            if param_name in [str(sym) for sym in list(expr.atoms())]:
+                return True
+
+    if "propagators" in solver_json.keys():
+        for sym, expr in solver_json["propagators"].items():
+            if param_name in [str(sym) for sym in list(expr.atoms())]:
+                return True
+
+    if "conditions" in solver_json.keys():
+        for conditional_solver_json in solver_json["conditions"].values():
+            if "update_expressions" in conditional_solver_json.keys():
+                for sym, expr in conditional_solver_json["update_expressions"].items():
+                    if param_name in [str(sym) for sym in list(expr.atoms())]:
+                        return True
+
+            if "propagators" in conditional_solver_json.keys():
+                for sym, expr in solver_json["propagators"].items():
+                    if param_name in [str(sym) for sym in list(expr.atoms())]:
+                        return True
+
+    return False
+
+
 def _analysis(indict, disable_stiffness_check: bool = False, disable_analytic_solver: bool = False, disable_singularity_detection: bool = False, preserve_expressions: Union[bool, Iterable[str]] = False, log_level: Union[str, int] = logging.WARNING) -> Tuple[List[Dict], SystemOfShapes, List[Shape]]:
     r"""
     Like analysis(), but additionally returns ``shape_sys`` and ``shapes``.
@@ -320,20 +346,7 @@ def _analysis(indict, disable_stiffness_check: bool = False, disable_analytic_so
             solver_json["parameters"] = {}
             for param_name, param_expr in indict["parameters"].items():
                 # only make parameters appear in a solver if they are actually used there
-                symbol_appears_in_any_expr = False
-                if "update_expressions" in solver_json.keys():
-                    for sym, expr in solver_json["update_expressions"].items():
-                        if param_name in [str(sym) for sym in list(expr.atoms())]:
-                            symbol_appears_in_any_expr = True
-                            break
-
-                if "propagators" in solver_json.keys():
-                    for sym, expr in solver_json["propagators"].items():
-                        if param_name in [str(sym) for sym in list(expr.atoms())]:
-                            symbol_appears_in_any_expr = True
-                            break
-
-                if symbol_appears_in_any_expr:
+                if symbol_appears_in_any_expr(sym, solver_json):
                     sympy_expr = sympy.parsing.sympy_parser.parse_expr(param_expr, global_dict=Shape._sympy_globals)
 
                     # validate output for numerical problems
@@ -388,6 +401,26 @@ def _analysis(indict, disable_stiffness_check: bool = False, disable_analytic_so
             for sym, expr in solver_json["propagators"].items():
                 solver_json["propagators"][sym] = str(expr)
 
+        if "conditions" in solver_json.keys():
+            for cond, cond_solver in solver_json["conditions"].items():
+                if "update_expressions" in cond_solver:
+                    for sym, expr in cond_solver["update_expressions"].items():
+                        cond_solver["update_expressions"][sym] = str(expr)
+
+                        if preserve_expressions and sym in preserve_expressions:
+                            if "analytic" in solver_json["solver"]:
+                                logging.warning("Not preserving expression for variable \"" + sym + "\" as it is solved by propagator solver")
+                                continue
+
+                            logging.info("Preserving expression for variable \"" + sym + "\"")
+                            var_def_str = _find_variable_definition(indict, sym, order=1)
+                            assert var_def_str is not None
+                            cond_solver["update_expressions"][sym] = var_def_str.replace("'", Config().differential_order_symbol)
+
+                if "propagators" in cond_solver:
+                    for sym, expr in cond_solver["propagators"].items():
+                        cond_solver["propagators"][sym] = str(expr)
+
     logging.info("In ode-toolbox: returning outdict = ")
     logging.info(json.dumps(solvers_json, indent=4, sort_keys=True))
 

odetoolbox/analytic_integrator.py

@@ -19,6 +19,7 @@
 # along with NEST.  If not, see <http://www.gnu.org/licenses/>.
 #
 
+import logging
 from typing import Dict, List, Optional
 
 import sympy
@@ -59,7 +60,6 @@ def __init__(self, solver_dict, spike_times: Optional[Dict[str, List[float]]] =
         #
         #   define the necessary numerical state variables
         #
-
         self.dim = len(self.all_variable_symbols)
         self.initial_values = self.solver_dict["initial_values"].copy()
         self.set_initial_values(self.initial_values)
@@ -72,11 +72,16 @@ def __init__(self, solver_dict, spike_times: Optional[Dict[str, List[float]]] =
                     subs_dict[k_] = v_
             self.shape_starting_values[k] = float(expr.evalf(subs=subs_dict))
 
-        self.update_expressions = self.solver_dict["update_expressions"].copy()
-        for k, v in self.update_expressions.items():
-            if type(self.update_expressions[k]) is str:
-                self.update_expressions[k] = sympy.parsing.sympy_parser.parse_expr(self.update_expressions[k], global_dict=Shape._sympy_globals)
 
+        #
+        #   initialise update expressions depending on whether conditional solver or not
+        #
+
+        if "update_expressions" in self.solver_dict.keys():
+            self._pick_unconditional_solver()
+        else:
+            assert "conditions" in self.solver_dict.keys()
+            self._pick_solver_based_on_condition()
 
         #
         #  reset the system to t = 0
@@ -85,24 +90,65 @@ def __init__(self, solver_dict, spike_times: Optional[Dict[str, List[float]]] =
         self.reset()
 
 
+    def _condition_holds(self, condition_string) -> bool:
+        parts = condition_string.strip('()').split('==')
+        lhs_str = parts[0].strip()
+        rhs_str = parts[1].strip()
+
+        # Sympify each side individually and create the Eq
+        equation = sympy.Eq(sympy.sympify(lhs_str), sympy.sympify(rhs_str))
+
+        return equation.subs(self.solver_dict["parameters"])
+
+
+    def _pick_unconditional_solver(self):
+        self.update_expressions = self.solver_dict["update_expressions"].copy()
+        self.propagators = self.solver_dict["propagators"].copy()
+        self._process_update_expressions_from_solver_dict()
+
+    def _pick_solver_based_on_condition(self):
+        r"""In case of a conditional propagator solver: pick a solver depending on the conditions that hold (depending on parameter values)"""
+        self.update_expressions = self.solver_dict["conditions"]["default"]["update_expressions"]
+        self.propagators = self.solver_dict["conditions"]["default"]["propagators"]
+
+        for condition, conditional_solver in self.solver_dict["conditions"].items():
+            print("Checking condition " + str(condition) + ", params = " + str(self.solver_dict["parameters"]))
+            if condition != "default" and self._condition_holds(condition):
+                self.update_expressions = conditional_solver["update_expressions"]
+                self.propagators = conditional_solver["propagators"]
+                print("Picking solver based on condition: " + str(condition))
+
+                break
+
+        self._process_update_expressions_from_solver_dict()
+
+
+    def _process_update_expressions_from_solver_dict(self):
         #
-        #   in the update expression, replace symbolic variables with their numerical values
+        #   create substitution dictionary to replace symbolic variables with their numerical values
         #
 
-        self.subs_dict = {}
-        for prop_symbol, prop_expr in self.solver_dict["propagators"].items():
-            self.subs_dict[prop_symbol] = prop_expr
+        subs_dict = {}
+        for prop_symbol, prop_expr in self.propagators.items():
+            subs_dict[prop_symbol] = prop_expr
         if "parameters" in self.solver_dict.keys():
             for param_symbol, param_expr in self.solver_dict["parameters"].items():
-                self.subs_dict[param_symbol] = param_expr
+                subs_dict[param_symbol] = param_expr
 
+        #
+        #   parse the expressions from JSON if necessary
+        #
+
+        for k, v in self.update_expressions.items():
+            if type(self.update_expressions[k]) is str:
+                self.update_expressions[k] = sympy.parsing.sympy_parser.parse_expr(self.update_expressions[k], global_dict=Shape._sympy_globals)
 
         #
         #   perform substitution in update expressions ahead of time to save time later
         #
 
         for k, v in self.update_expressions.items():
-            self.update_expressions[k] = self.update_expressions[k].subs(self.subs_dict).subs(self.subs_dict)
+            self.update_expressions[k] = self.update_expressions[k].subs(subs_dict).subs(subs_dict)
 
         #
         #    autowrap

odetoolbox/system_of_shapes.py

@@ -253,13 +253,17 @@ def generate_propagator_solver(self, disable_singularity_detection: bool = False
                     solver_dict = {"solver": "analytical",
                                    "state_variables": default_solver["state_variables"],
                                    "initial_values": default_solver["initial_values"],
-                                   "conditions": {"default" : {"propagators": default_solver["propagators"],
-                                                               "update_expressions": default_solver["update_expressions"]}}}
+                                   "conditions": {"default": {"propagators": default_solver["propagators"],
+                                                              "update_expressions": default_solver["update_expressions"]}}}
 
                     # XXX: TODO: generate/loop over all combinations of conditions!!!
 
                     for cond_set in conditions:
-                        condition_str: str = " && ".join(["(" + str(eq.lhs) + " == " + str(eq.rhs) + ")" for eq in cond_set])
+                        if len(cond_set) == 1:
+                            eq = list(cond_set)[0]
+                            condition_str: str = str(eq.lhs) + " == " + str(eq.rhs)
+                        else:
+                            condition_str: str = " && ".join(["(" + str(eq.lhs) + " == " + str(eq.rhs) + ")" for eq in cond_set])
                         conditional_A = self.A_.copy()
                         conditional_b = self.b_.copy()
                         conditional_c = self.c_.copy()
@@ -297,7 +301,7 @@ def generate_propagator_solver(self, disable_singularity_detection: bool = False
             for col in range(P.shape[1]):
                 if not _is_zero(P[row, col]):
                     sym_str = Config().propagators_prefix + "__{}__{}".format(str(self.x_[row]), str(self.x_[col]))
-                    P_expr[sym_str] = str(P[row, col])
+                    P_expr[sym_str] = P[row, col]
                     if row != col and not _is_zero(self.b_[col]):
                         # the ODE for x_[row] depends on the inhomogeneous ODE of x_[col]. We can't solve this analytically in the general case (even though some specific cases might admit a solution)
                         raise PropagatorGenerationException("the ODE for " + str(self.x_[row]) + " depends on the inhomogeneous ODE of " + str(self.x_[col]) + ". We can't solve this analytically in the general case (even though some specific cases might admit a solution)")

tests/test_double_exponential.py

@@ -20,6 +20,7 @@
 #
 
 import numpy as np
+import pytest
 from scipy.integrate import odeint
 
 import odetoolbox
@@ -29,7 +30,7 @@
 
 try:
     import matplotlib as mpl
-    mpl.use('Agg')
+    mpl.use("Agg")
     import matplotlib.pyplot as plt
     INTEGRATION_TEST_DEBUG_PLOTS = True
 except ImportError:
@@ -39,23 +40,32 @@
 class TestDoubleExponential:
     r"""Test propagators generation for double exponential"""
 
-    def test_double_exponential(self):
-        r"""Test propagators generation for double exponential"""
+    @pytest.mark.parametrize("tau_1, tau_2", [(10., 2.), (10., 10.)])
+    def test_double_exponential(self, tau_1, tau_2):
+        r"""Test propagators generation for double exponential
+
+        Test for a case where tau_1 != tau_2 and where tau_1 == tau_2; this tests handling of numerical singularities.
+
+        tau_1: decay time constant (ms)
+        tau_2: rise time constant (ms)
+        """
 
         def time_to_max(tau_1, tau_2):
             r"""
             Time of maximum.
             """
-            tmax = (np.log(tau_1) - np.log(tau_2)) / (1. / tau_2 - 1. / tau_1)
-            return tmax
+            if tau_1 == tau_2:
+                return tau_1
+
+            return (np.log(tau_1) - np.log(tau_2)) / (1. / tau_2 - 1. / tau_1)
 
         def unit_amplitude(tau_1, tau_2):
             r"""
             Scaling factor ensuring that amplitude of solution is one.
             """
             tmax = time_to_max(tau_1, tau_2)
-            alpha = 1. / (np.exp(-tmax / tau_1) - np.exp(-tmax / tau_2))
-            return alpha
+
+            return 1. / (np.exp(-tmax / tau_1) - np.exp(-tmax / tau_2))
 
         def flow(y, t, tau_1, tau_2, alpha, dt):
             r"""
@@ -66,26 +76,27 @@ def flow(y, t, tau_1, tau_2, alpha, dt):
 
             return np.array([dy1dt, dy2dt])
 
+        if tau_1 == tau_2:
+            alpha = 1.
+        else:
+            alpha = unit_amplitude(tau_1=tau_1, tau_2=tau_2)
+
         indict = {"dynamics": [{"expression": "I_aux' = -I_aux / tau_1",
                                 "initial_values": {"I_aux": "0."}},
                                {"expression": "I' = I_aux - I / tau_2",
                                 "initial_values": {"I": "0"}}],
                   "options": {"output_timestep_symbol": "__h"},
-                  "parameters": {"tau_1": "10",
-                                 "tau_2": "2",
+                  "parameters": {"tau_1": str(tau_1),
+                                 "tau_2": str(tau_2),
                                  "w": "3.14",
-                                 "alpha": str(unit_amplitude(tau_1=10., tau_2=2.)),
+                                 "alpha": str(alpha),
                                  "weighted_input_spikes": "0."}}
 
         w = 3.14                              # weight (amplitude; pA)
-        tau_1 = 10.                           # decay time constant (ms)
-        tau_2 = 2.                            # rise time constant (ms)
         dt = .125                             # time resolution (ms)
         T = 500.                              # simulation time (ms)
         input_spike_times = np.array([100., 300.])  # array of input spike times (ms)
 
-        alpha = unit_amplitude(tau_1, tau_2)
-
         stimuli = [{"type": "list",
                     "list": " ".join([str(el) for el in input_spike_times]),
                     "variables": ["I_aux"]}]
@@ -103,7 +114,7 @@ def flow(y, t, tau_1, tau_2, alpha, dt):
         N = int(np.ceil(T / dt) + 1)
         timevec = np.linspace(0., T, N)
         analytic_integrator = AnalyticIntegrator(solver_dict, spike_times)
-        analytic_integrator.shape_starting_values["I_aux"] = w * alpha * (1. / tau_2 - 1. / tau_1)
+        analytic_integrator.shape_starting_values["I_aux"] = w * alpha
         analytic_integrator.set_initial_values(ODE_INITIAL_VALUES)
         analytic_integrator.reset()
         state = {"timevec": [], "I": [], "I_aux": []}
@@ -119,24 +130,24 @@ def flow(y, t, tau_1, tau_2, alpha, dt):
         ts2 = np.arange(input_spike_times[1], T + dt, dt)
 
         y_ = odeint(flow, [0., 0.], ts0, args=(tau_1, tau_2, alpha, dt))
-        y_ = np.vstack([y_, odeint(flow, [y_[-1, 0] + w * alpha * (1. / tau_2 - 1. / tau_1), y_[-1, 1]], ts1, args=(tau_1, tau_2, alpha, dt))])
-        y_ = np.vstack([y_, odeint(flow, [y_[-1, 0] + w * alpha * (1. / tau_2 - 1. / tau_1), y_[-1, 1]], ts2, args=(tau_1, tau_2, alpha, dt))])
+        y_ = np.vstack([y_, odeint(flow, [y_[-1, 0] + w * alpha, y_[-1, 1]], ts1, args=(tau_1, tau_2, alpha, dt))])
+        y_ = np.vstack([y_, odeint(flow, [y_[-1, 0] + w * alpha, y_[-1, 1]], ts2, args=(tau_1, tau_2, alpha, dt))])
 
-        rec_I_interp = np.interp(np.hstack([ts0, ts1, ts2]), timevec, state['I'])
-        rec_I_aux_interp = np.interp(np.hstack([ts0, ts1, ts2]), timevec, state['I_aux'])
+        rec_I_interp = np.interp(np.hstack([ts0, ts1, ts2]), timevec, state["I"])
+        rec_I_aux_interp = np.interp(np.hstack([ts0, ts1, ts2]), timevec, state["I_aux"])
 
         if INTEGRATION_TEST_DEBUG_PLOTS:
             tmax = time_to_max(tau_1, tau_2)
-            mpl.rcParams['text.usetex'] = True
+            mpl.rcParams["text.usetex"] = True
 
             fig, ax = plt.subplots(nrows=2, figsize=(5, 4), dpi=300)
-            ax[0].plot(timevec, state['I_aux'], '--', lw=3, color='k', label=r'$I_\mathsf{aux}(t)$ (NEST)')
-            ax[0].plot(timevec, state['I'], '-', lw=3, color='k', label=r'$I(t)$ (NEST)')
-            ax[0].plot(np.hstack([ts0, ts1, ts2]), y_[:, 0], '--', lw=2, color='r', label=r'$I_\mathsf{aux}(t)$ (odeint)')
-            ax[0].plot(np.hstack([ts0, ts1, ts2]), y_[:, 1], '-', lw=2, color='r', label=r'$I(t)$ (odeint)')
+            ax[0].plot(timevec, state["I_aux"], "--", lw=3, color="k", label=r"$I_\mathsf{aux}(t)$ (ODEtb)")
+            ax[0].plot(timevec, state["I"], "-", lw=3, color="k", label=r"$I(t)$ (ODEtb)")
+            ax[0].plot(np.hstack([ts0, ts1, ts2]), y_[:, 0], "--", lw=2, color="r", label=r"$I_\mathsf{aux}(t)$ (odeint)")
+            ax[0].plot(np.hstack([ts0, ts1, ts2]), y_[:, 1], "-", lw=2, color="r", label=r"$I(t)$ (odeint)")
 
             for tin in input_spike_times:
-                ax[0].vlines(tin + tmax, ax[0].get_ylim()[0], ax[0].get_ylim()[1], colors='k', linestyles=':')
+                ax[0].vlines(tin + tmax, ax[0].get_ylim()[0], ax[0].get_ylim()[1], colors="k", linestyles=":")
 
             ax[1].semilogy(np.hstack([ts0, ts1, ts2]), np.abs(y_[:, 1] - rec_I_interp), label="I")
             ax[1].semilogy(np.hstack([ts0, ts1, ts2]), np.abs(y_[:, 0] - rec_I_aux_interp), linestyle="--", label="I_aux")
@@ -146,9 +157,9 @@ def flow(y, t, tau_1, tau_2, alpha, dt):
                 _ax.set_xlim(0., T + dt)
                 _ax.legend()
 
-            ax[-1].set_xlabel(r'time (ms)')
+            ax[-1].set_xlabel(r"time (ms)")
 
-            fig.savefig('double_exp_test.png')
+            fig.savefig("double_exp_test_[tau_1=" + str(tau_1) + "]_[tau_2=" + str(tau_2) + "].png")
 
         np.testing.assert_allclose(y_[:, 1], rec_I_interp, atol=1E-7)
 

tests/test_propagator_solver_homogeneous.py

@@ -32,5 +32,8 @@ def test_propagator_solver_homogeneous(self):
         assert len(solver_dict) == 1
         solver_dict = solver_dict[0]
         assert solver_dict["solver"] == "analytical"
-        assert float(solver_dict["propagators"]["__P__refr_t__refr_t"]) == 1.
-        assert solver_dict["propagators"]["__P__V_m__V_m"] == "1.0*exp(-__h/tau_m)"
+
+        for cond_solver_dict in solver_dict["conditions"].values():
+            assert float(cond_solver_dict["propagators"]["__P__refr_t__refr_t"]) == 1.
+
+        assert solver_dict["conditions"]["default"]["propagators"]["__P__V_m__V_m"] == "1.0*exp(-__h/tau_m)"