#1100 starting to add SDAEs (odes only for now)

valentinsulzer · valentinsulzer · commit a59abb58dc4f · 2020-07-08T18:28:18.000-04:00
diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py
@@ -205,6 +205,11 @@ def process_model(self, model, inplace=True, check_model=True):
         model_disc.rhs, model_disc.concatenated_rhs = rhs, concat_rhs
         model_disc.algebraic, model_disc.concatenated_algebraic = alg, concat_alg
 
+        # Save length of rhs and algebraic
+        model_disc.len_rhs = model_disc.concatenated_rhs.size
+        model_disc.len_alg = model_disc.concatenated_algebraic.size
+        model_disc.len_rhs_and_alg = model_disc.len_rhs + model_disc.len_alg
+
         # Process events
         processed_events = []
         pybamm.logger.info("Discretise events for {}".format(model.name))
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
@@ -30,6 +30,12 @@ class BaseSolver(object):
         specified by 'root_method' (e.g. "lm", "hybr", ...)
     root_tol : float, optional
         The tolerance for the initial-condition solver (default is 1e-6).
+    solve_sensitivity_equations : bool, optional
+        Whether to explicitly formulate the sensitivity equations for sensitivity
+        to input parameters. The formulation is as per "Park, S., Kato, D., Gima, Z.,
+        Klein, R., & Moura, S. (2018). Optimal experimental design for parameterization
+        of an electrochemical lithium-ion battery model. Journal of The Electrochemical
+        Society, 165(7), A1309.". See #1100 for details
     """
 
     def __init__(
@@ -40,6 +46,7 @@ def __init__(
         root_method=None,
         root_tol=1e-6,
         max_steps="deprecated",
+        solve_sensitivity_equations=False,
     ):
         self._method = method
         self._rtol = rtol
@@ -57,6 +64,7 @@ def __init__(
         self.name = "Base solver"
         self.ode_solver = False
         self.algebraic_solver = False
+        self.solve_sensitivity_equations = solve_sensitivity_equations
 
     @property
     def method(self):
@@ -191,17 +199,28 @@ def set_up(self, model, inputs=None):
             )
             model.convert_to_format = "casadi"
 
+        # Only allow solving sensitivity equations with the casadi format for now
+        if (
+            self.solve_sensitivity_equations is True
+            and model.convert_to_format != "casadi"
+        ):
+            raise NotImplementedError(
+                "model should be converted to casadi format in order to solve "
+                "sensitivity equations"
+            )
+
         if model.convert_to_format != "casadi":
             simp = pybamm.Simplification()
             # Create Jacobian from concatenated rhs and algebraic
-            y = pybamm.StateVector(slice(0, model.concatenated_initial_conditions.size))
+            y = pybamm.StateVector(slice(0, model.len_rhs_and_alg))
             # set up Jacobian object, for re-use of dict
             jacobian = pybamm.Jacobian()
         else:
             # Convert model attributes to casadi
             t_casadi = casadi.MX.sym("t")
-            y_diff = casadi.MX.sym("y_diff", model.concatenated_rhs.size)
-            y_alg = casadi.MX.sym("y_alg", model.concatenated_algebraic.size)
+            # Create the symbolic state vectors
+            y_diff = casadi.MX.sym("y_diff", model.len_rhs)
+            y_alg = casadi.MX.sym("y_alg", model.len_alg)
             y_casadi = casadi.vertcat(y_diff, y_alg)
             p_casadi = {}
             for name, value in inputs.items():
@@ -210,6 +229,13 @@ def set_up(self, model, inputs=None):
                 else:
                     p_casadi[name] = casadi.MX.sym(name, value.shape[0])
             p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()])
+            # sensitivity vectors
+            if self.solve_sensitivity_equations is True:
+                S_x = casadi.MX.sym("S_x", model.len_rhs * p_casadi_stacked.shape[0])
+                S_z = casadi.MX.sym("S_z", model.len_alg * p_casadi_stacked.shape[0])
+                y_and_S = casadi.vertcat(y_diff, S_x, y_alg, S_z)
+            else:
+                y_and_S = y_casadi
 
         def process(func, name, use_jacobian=None):
             def report(string):
@@ -258,16 +284,40 @@ def report(string):
                 # Process with CasADi
                 report(f"Converting {name} to CasADi")
                 func = func.to_casadi(t_casadi, y_casadi, inputs=p_casadi)
+                # Add sensitivity vectors to the rhs and algebraic equations
+                if self.solve_sensitivity_equations is True:
+                    if name == "rhs":
+                        report(f"Creating sensitivity equations for rhs using CasADi")
+                        df_dx = casadi.jacobian(func, y_diff)
+                        df_dp = casadi.jacobian(func, p_casadi_stacked)
+                        if model.len_alg == 0:
+                            S_rhs = df_dx @ S_x + df_dp
+                        else:
+                            df_dz = casadi.jacobian(func, y_alg)
+                            S_rhs = df_dx @ S_x + df_dz @ S_z + df_dp
+                        func = casadi.vertcat(func, S_rhs)
+                    elif name == "initial_conditions":
+                        if model.len_rhs == 0 or model.len_alg == 0:
+                            S_0 = casadi.jacobian(func, p_casadi_stacked).reshape(
+                                (-1, 1)
+                            )
+                            func = casadi.vertcat(func, S_0)
+                        else:
+                            x0 = func[: model.len_rhs]
+                            z0 = func[model.len_rhs :]
+                            Sx_0 = casadi.jacobian(x0, p_casadi_stacked)
+                            Sz_0 = casadi.jacobian(z0, p_casadi_stacked)
+                            func = casadi.vertcat(x0, Sx_0, z0, Sz_0)
                 if use_jacobian:
                     report(f"Calculating jacobian for {name} using CasADi")
-                    jac_casadi = casadi.jacobian(func, y_casadi)
+                    jac_casadi = casadi.jacobian(func, y_and_S)
                     jac = casadi.Function(
-                        name, [t_casadi, y_casadi, p_casadi_stacked], [jac_casadi]
+                        name, [t_casadi, y_and_S, p_casadi_stacked], [jac_casadi]
                     )
                 else:
                     jac = None
                 func = casadi.Function(
-                    name, [t_casadi, y_casadi, p_casadi_stacked], [func]
+                    name, [t_casadi, y_and_S, p_casadi_stacked], [func]
                 )
             if name == "residuals":
                 func_call = Residuals(func, name, model)
@@ -277,6 +327,7 @@ def report(string):
                 jac_call = SolverCallable(jac, name + "_jac", model)
             else:
                 jac_call = None
+
             return func, func_call, jac_call
 
         # Check for heaviside functions in rhs and algebraic and add discontinuity
@@ -324,8 +375,18 @@ def report(string):
         )[0]
         init_eval = InitialConditions(initial_conditions, model)
 
+        if self.solve_sensitivity_equations is True:
+            init_eval.y_dummy = np.zeros(
+                (
+                    model.len_rhs_and_alg * (np.vstack(list(inputs.values())).size + 1),
+                    1,
+                )
+            )
+        else:
+            init_eval.y_dummy = np.zeros((model.len_rhs_and_alg, 1))
+
         # Process rhs, algebraic and event expressions
-        rhs, rhs_eval, jac_rhs = process(model.concatenated_rhs, "RHS")
+        rhs, rhs_eval, jac_rhs = process(model.concatenated_rhs, "rhs")
         algebraic, algebraic_eval, jac_algebraic = process(
             model.concatenated_algebraic, "algebraic"
         )
@@ -423,7 +484,7 @@ def _set_initial_conditions(self, model, inputs, update_rhs):
                 y0_from_inputs = model.init_eval(inputs)
                 # Reuse old solution for algebraic equations
                 y0_from_model = model.y0
-                len_rhs = model.concatenated_rhs.size
+                len_rhs = model.len_rhs
                 # update model.y0, which is used for initialising the algebraic solver
                 if len_rhs == 0:
                     model.y0 = y0_from_model
@@ -861,7 +922,7 @@ def __init__(self, function, name, model):
 
     def __call__(self, t, y, inputs):
         y = y.reshape(-1, 1)
-        if self.name in ["RHS", "algebraic", "residuals"]:
+        if self.name in ["rhs", "algebraic", "residuals"]:
             pybamm.logger.debug(
                 "Evaluating {} for {} at t={}".format(
                     self.name, self.model.name, t * self.timescale
@@ -874,7 +935,7 @@ def __call__(self, t, y, inputs):
     def function(self, t, y, inputs):
         if self.form == "casadi":
             states_eval = self._function(t, y, inputs)
-            if self.name in ["RHS", "algebraic", "residuals", "event"]:
+            if self.name in ["rhs", "algebraic", "residuals", "event"]:
                 return states_eval.full()
             else:
                 # keep jacobians sparse
@@ -901,7 +962,6 @@ class InitialConditions(SolverCallable):
 
     def __init__(self, function, model):
         super().__init__(function, "initial conditions", model)
-        self.y_dummy = np.zeros(model.concatenated_initial_conditions.shape)
 
     def __call__(self, inputs):
         if self.form == "casadi":
diff --git a/pybamm/solvers/casadi_solver.py b/pybamm/solvers/casadi_solver.py
@@ -363,9 +363,8 @@ def get_integrator(self, model, inputs, t_eval=None):
 
     def _run_integrator(self, model, y0, inputs, t_eval):
         integrator, use_grid = self.integrators[model]
-        len_rhs = model.concatenated_rhs.size
-        y0_diff = y0[:len_rhs]
-        y0_alg = y0[len_rhs:]
+        y0_diff = y0[: model.len_rhs]
+        y0_alg = y0[model.len_rhs :]
         try:
             # Try solving
             if use_grid is True:
diff --git a/pybamm/solvers/processed_symbolic_variable.py b/pybamm/solvers/processed_symbolic_variable.py
@@ -110,9 +110,6 @@ def initialise_0D(self):
 
     def initialise_1D(self):
         "Create a 1D variable"
-        len_space = self.base_eval.shape[0]
-        entries = np.empty((len_space, len(self.t_sol)))
-
         # Evaluate the base_variable index-by-index
         for idx in range(len(self.t_sol)):
             t = self.t_sol[idx]
diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py
@@ -1,6 +1,7 @@
 #
 # Processed Variable class
 #
+import casadi
 import numbers
 import numpy as np
 import pybamm
@@ -59,6 +60,10 @@ def __init__(self, base_variable, solution, known_evals=None, warn=True):
         self.known_evals = known_evals
         self.warn = warn
 
+        # Sensitivity starts off uninitialized, only set when called
+        self._sensitivity = None
+        self.solution_sensitivity = solution.sensitivity
+
         # Set timescale
         self.timescale = solution.model.timescale.evaluate()
         self.t_pts = self.t_sol * self.timescale
@@ -553,6 +558,76 @@ def data(self):
         "Same as entries, but different name"
         return self.entries
 
+    @property
+    def sensitivity(self):
+        """
+        Returns a dictionary of sensitivity for each input parameter.
+        The keys are the input parameters, and the value is a matrix of size
+        (n_x * n_t, n_p), where n_x is the number of states, n_t is the number of time
+        points, and n_p is the size of the input parameter
+        """
+        # No sensitivity if there are no inputs
+        if len(self.inputs) == 0:
+            return {}
+        # Otherwise initialise and return sensitivity
+        if self._sensitivity is None:
+            self.initialise_sensitivity()
+        return self._sensitivity
+
+    def initialise_sensitivity(self):
+        "Set up the sensitivity dictionary"
+        inputs_stacked = casadi.vertcat(*[p for p in self.inputs.values()])
+
+        # Set up symbolic variables
+        t_casadi = casadi.MX.sym("t")
+        y_casadi = casadi.MX.sym("y", self.u_sol.shape[0])
+        p_casadi = {
+            name: casadi.MX.sym(name, value.shape[0])
+            for name, value in self.inputs.items()
+        }
+        p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()])
+
+        # Convert variable to casadi format for differentiating
+        var_casadi = self.base_variable.to_casadi(t_casadi, y_casadi, inputs=p_casadi)
+        dvar_dy = casadi.jacobian(var_casadi, y_casadi)
+        dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked)
+
+        # Convert to functions and evaluate index-by-index
+        dvar_dy_func = casadi.Function(
+            "dvar_dy", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dy]
+        )
+        dvar_dp_func = casadi.Function(
+            "dvar_dp", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dp]
+        )
+        for idx in range(len(self.t_sol)):
+            t = self.t_sol[idx]
+            u = self.u_sol[:, idx]
+            inp = inputs_stacked[:, idx]
+            next_dvar_dy_eval = dvar_dy_func(t, u, inp)
+            next_dvar_dp_eval = dvar_dp_func(t, u, inp)
+            if idx == 0:
+                dvar_dy_eval = next_dvar_dy_eval
+                dvar_dp_eval = next_dvar_dp_eval
+            else:
+                dvar_dy_eval = casadi.diagcat(dvar_dy_eval, next_dvar_dy_eval)
+                dvar_dp_eval = casadi.vertcat(dvar_dp_eval, next_dvar_dp_eval)
+
+        # Compute sensitivity
+        dy_dp = self.solution_sensitivity["all"]
+        S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval
+
+        sensitivity = {"all": S_var}
+
+        # Add the individual sensitivity
+        start = 0
+        for name, inp in self.inputs.items():
+            end = start + inp.shape[0]
+            sensitivity[name] = S_var[:, start:end]
+            start = end
+
+        # Save attribute
+        self._sensitivity = sensitivity
+
 
 def eval_dimension_name(name, x, r, y, z):
     if name == "x":
diff --git a/pybamm/solvers/scipy_solver.py b/pybamm/solvers/scipy_solver.py
@@ -23,10 +23,24 @@ class ScipySolver(pybamm.BaseSolver):
         Any options to pass to the solver.
         Please consult `SciPy documentation <https://tinyurl.com/yafgqg9y>`_ for
         details.
+    solve_sensitivity_equations : bool, optional
+        See :class:`pybamm.BaseSolver`
     """
 
-    def __init__(self, method="BDF", rtol=1e-6, atol=1e-6, extra_options=None):
-        super().__init__(method, rtol, atol)
+    def __init__(
+        self,
+        method="BDF",
+        rtol=1e-6,
+        atol=1e-6,
+        extra_options=None,
+        solve_sensitivity_equations=False,
+    ):
+        super().__init__(
+            method=method,
+            rtol=rtol,
+            atol=atol,
+            solve_sensitivity_equations=solve_sensitivity_equations,
+        )
         self.ode_solver = True
         self.extra_options = extra_options or {}
         self.name = "Scipy solver ({})".format(method)
@@ -52,8 +66,10 @@ def _integrate(self, model, t_eval, inputs=None):
             various diagnostic messages.
 
         """
+        # Save inputs dictionary, and if necessary convert inputs to a casadi vector
+        inputs_dict = inputs
         if model.convert_to_format == "casadi":
-            inputs = casadi.vertcat(*[x for x in inputs.values()])
+            inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
 
         extra_options = {**self.extra_options, "rtol": self.rtol, "atol": self.atol}
 
@@ -107,6 +123,14 @@ def event_fn(t, y):
                 termination = "final time"
                 t_event = None
                 y_event = np.array(None)
-            return pybamm.Solution(sol.t, sol.y, t_event, y_event, termination)
+            return pybamm.Solution(
+                sol.t,
+                sol.y,
+                t_event,
+                y_event,
+                termination,
+                model=model,
+                inputs=inputs_dict,
+            )
         else:
             raise pybamm.SolverError(sol.message)
diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py
diff --git a/tests/unit/test_solvers/test_scipy_solver.py b/tests/unit/test_solvers/test_scipy_solver.py
