#1100 get SODEs working for casadi solver

valentinsulzer · valentinsulzer · commit 54d45f94944d · 2020-07-09T16:38:18.000-04:00
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
@@ -8,6 +8,7 @@
 import numpy as np
 import sys
 import itertools
+from scipy.linalg import block_diag
 
 
 class BaseSolver(object):
@@ -426,12 +427,21 @@ def report(string):
         ):
             # can use DAE solver to solve model with algebraic equations only
             if len(model.rhs) > 0:
-                mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
+                if self.solve_sensitivity_equations is True:
+                    # Copy mass matrix blocks diagonally
+                    single_mass_matrix_inv = model.mass_matrix_inv.entries.toarray()
+                    n_inputs = p_casadi_stacked.shape[0]
+                    block_mass_matrix = block_diag(
+                        *[single_mass_matrix_inv] * (n_inputs + 1)
+                    )
+                    mass_matrix_inv = casadi.MX(block_mass_matrix)
+                else:
+                    mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
                 explicit_rhs = mass_matrix_inv @ rhs(
-                    t_casadi, y_casadi, p_casadi_stacked
+                    t_casadi, y_and_S, p_casadi_stacked
                 )
                 model.casadi_rhs = casadi.Function(
-                    "rhs", [t_casadi, y_casadi, p_casadi_stacked], [explicit_rhs]
+                    "rhs", [t_casadi, y_and_S, p_casadi_stacked], [explicit_rhs]
                 )
             model.casadi_algebraic = algebraic
         if len(model.rhs) == 0:
@@ -703,10 +713,6 @@ def solve(self, model, t_eval=None, external_variables=None, inputs=None):
         solution.set_up_time = set_up_time
         solution.solve_time = timer.time()
 
-        # Add model and inputs to solution
-        solution.model = model
-        solution.inputs = ext_and_inputs
-
         # Identify the event that caused termination
         termination = self.get_termination_reason(solution, model.events)
 
diff --git a/pybamm/solvers/casadi_solver.py b/pybamm/solvers/casadi_solver.py
@@ -55,6 +55,9 @@ class CasadiSolver(pybamm.BaseSolver):
         Any options to pass to the CasADi integrator when calling the integrator.
         Please consult `CasADi documentation <https://tinyurl.com/y5rk76os>`_ for
         details.
+    solve_sensitivity_equations : bool, optional
+        Whether to explicitly formulate and solve the forward sensitivity equations.
+        See :class:`pybamm.BaseSolver`
 
     """
 
@@ -69,8 +72,16 @@ def __init__(
         dt_max=None,
         extra_options_setup=None,
         extra_options_call=None,
+        solve_sensitivity_equations=False,
     ):
-        super().__init__("problem dependent", rtol, atol, root_method, root_tol)
+        super().__init__(
+            "problem dependent",
+            rtol,
+            atol,
+            root_method,
+            root_tol,
+            solve_sensitivity_equations=solve_sensitivity_equations,
+        )
         if mode in ["safe", "fast"]:
             self.mode = mode
         else:
@@ -106,24 +117,26 @@ def _integrate(self, model, t_eval, inputs=None):
             Any external variables or input parameters to pass to the model when solving
         """
         # Record whether there are any symbolic inputs
-        inputs = inputs or {}
-        has_symbolic_inputs = any(isinstance(v, casadi.MX) for v in inputs.values())
+        inputs_dict = inputs or {}
+        has_symbolic_inputs = any(
+            isinstance(v, casadi.MX) for v in inputs_dict.values()
+        )
 
         # convert inputs to casadi format
-        inputs = casadi.vertcat(*[x for x in inputs.values()])
+        inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
 
         if has_symbolic_inputs:
             # Create integrax`tor without grid to avoid having to create several times
             self.get_integrator(model, inputs)
-            solution = self._run_integrator(model, model.y0, inputs, t_eval)
+            solution = self._run_integrator(model, model.y0, inputs_dict, t_eval)
             solution.termination = "final time"
             return solution
         elif self.mode == "fast" or not model.events:
             if not model.events:
                 pybamm.logger.info("No events found, running fast mode")
             # Create an integrator with the grid (we just need to do this once)
             self.get_integrator(model, inputs, t_eval)
-            solution = self._run_integrator(model, model.y0, inputs, t_eval)
+            solution = self._run_integrator(model, model.y0, inputs_dict, t_eval)
             solution.termination = "final time"
             return solution
         elif self.mode == "safe":
@@ -143,7 +156,9 @@ def _integrate(self, model, t_eval, inputs=None):
             pybamm.logger.info("Start solving {} with {}".format(model.name, self.name))
 
             # Initialize solution
-            solution = pybamm.Solution(np.array([t]), y0[:, np.newaxis])
+            solution = pybamm.Solution(
+                np.array([t]), y0[:, np.newaxis], model=model, inputs=inputs_dict
+            )
             solution.solve_time = 0
 
             # Try to integrate in global steps of size dt_max. Note: dt_max must
@@ -178,7 +193,7 @@ def _integrate(self, model, t_eval, inputs=None):
                     # halve the step size and try again.
                     try:
                         current_step_sol = self._run_integrator(
-                            model, y0, inputs, t_window
+                            model, y0, inputs_dict, t_window
                         )
                         solved = True
                     except pybamm.SolverError:
@@ -257,7 +272,9 @@ def event_fun(t):
                         t_window = np.array([t, t_event])
 
                     # integrator = self.get_integrator(model, t_window, inputs)
-                    current_step_sol = self._run_integrator(model, y0, inputs, t_window)
+                    current_step_sol = self._run_integrator(
+                        model, y0, inputs_dict, t_window
+                    )
 
                     # assign temporary solve time
                     current_step_sol.solve_time = np.nan
@@ -361,10 +378,18 @@ def get_integrator(self, model, inputs, t_eval=None):
             self.integrators[model] = (integrator, use_grid)
             return integrator
 
-    def _run_integrator(self, model, y0, inputs, t_eval):
+    def _run_integrator(self, model, y0, inputs_dict, t_eval):
+        inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
         integrator, use_grid = self.integrators[model]
-        y0_diff = y0[: model.len_rhs]
-        y0_alg = y0[model.len_rhs :]
+        # Split up initial conditions into differential and algebraic
+        # Check y0 to see if it includes sensitivities
+        if model.len_rhs_and_alg == y0.shape[0]:
+            len_rhs = model.len_rhs
+        else:
+            len_rhs = model.len_rhs * (inputs.shape[0] + 1)
+        y0_diff = y0[:len_rhs]
+        y0_alg = y0[len_rhs:]
+        # Solve
         try:
             # Try solving
             if use_grid is True:
@@ -379,7 +404,7 @@ def _run_integrator(self, model, y0, inputs, t_eval):
                     **self.extra_options_call
                 )
                 y_sol = np.concatenate([sol["xf"].full(), sol["zf"].full()])
-                return pybamm.Solution(t_eval, y_sol)
+                return pybamm.Solution(t_eval, y_sol, model=model, inputs=inputs_dict)
             else:
                 # Repeated calls to the integrator
                 x = y0_diff
@@ -399,10 +424,14 @@ def _run_integrator(self, model, y0, inputs, t_eval):
                     if not z.is_empty():
                         y_alg = casadi.horzcat(y_alg, z)
                 if z.is_empty():
-                    return pybamm.Solution(t_eval, y_diff)
+                    return pybamm.Solution(
+                        t_eval, y_diff, model=model, inputs=inputs_dict
+                    )
                 else:
                     y_sol = casadi.vertcat(y_diff, y_alg)
-                    return pybamm.Solution(t_eval, y_sol)
+                    return pybamm.Solution(
+                        t_eval, y_sol, model=model, inputs=inputs_dict
+                    )
         except RuntimeError as e:
             # If it doesn't work raise error
             raise pybamm.SolverError(e.args[0])
diff --git a/pybamm/solvers/scipy_solver.py b/pybamm/solvers/scipy_solver.py
@@ -24,6 +24,7 @@ class ScipySolver(pybamm.BaseSolver):
         Please consult `SciPy documentation <https://tinyurl.com/yafgqg9y>`_ for
         details.
     solve_sensitivity_equations : bool, optional
+        Whether to explicitly formulate and solve the forward sensitivity equations.
         See :class:`pybamm.BaseSolver`
     """
 
diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py
@@ -66,6 +66,7 @@ def __init__(
         # y only has the shape or the rhs and alg solution)
         if model is None or model.len_rhs_and_alg == y.shape[0]:
             self._y = y
+            self.sensitivity = {}
         else:
             n_states = model.len_rhs_and_alg
             n_t = len(t)
@@ -133,11 +134,9 @@ def __init__(
         if copy_this is None:
             self.set_up_time = None
             self.solve_time = None
-            self.has_symbolic_inputs = False
         else:
             self.set_up_time = copy_this.set_up_time
             self.solve_time = copy_this.solve_time
-            self.has_symbolic_inputs = copy_this.has_symbolic_inputs
 
         # initiaize empty variables and data
         self._variables = pybamm.FuzzyDict()
diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py
diff --git a/tests/unit/test_solvers/test_scipy_solver.py b/tests/unit/test_solvers/test_scipy_solver.py
