conditional propagators

Author: C.A.P. Linssen

odetoolbox/__init__.py

@@ -181,6 +181,30 @@ 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"""
@@ -320,20 +344,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 +399,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

@@ -72,11 +72,17 @@ 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.update_expressions = self.solver_dict["update_expressions"].copy()
+            self._process_update_expressions_from_solver_dict()
+        else:
+            assert "conditions" in self.solver_dict.keys()
+            self._pick_solver_based_on_condition()
 
         #
         #  reset the system to t = 0
@@ -85,24 +91,57 @@ 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_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():
+            if condition != "default" and self._condition_holds(condition):
+                self.update_expressions = conditional_solver["update_expressions"]
+                self.propagators = conditional_solver["update_expressions"]
+                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

@@ -259,7 +259,11 @@ def generate_propagator_solver(self, disable_singularity_detection: bool = False
                     # 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()