#1100 working on casadi solver sensitivities

Author: valentinsulzer

pybamm/solvers/base_solver.py

@@ -510,20 +510,27 @@ def _set_initial_conditions(self, model, inputs, update_rhs):
             Whether to update the rhs. True for 'solve', False for 'step'.
 
         """
+        # Make inputs symbolic if calculating sensitivities with casadi
+        if self.sensitivity == "casadi":
+            symbolic_inputs = casadi.MX.sym(
+                "inputs", casadi.vertcat(*inputs.values()).shape[0]
+            )
+        else:
+            symbolic_inputs = inputs
         if self.algebraic_solver is True:
             # Don't update model.y0
             return None
         elif len(model.algebraic) == 0:
             if update_rhs is True:
                 # Recalculate initial conditions for the rhs equations
-                model.y0 = model.init_eval(inputs)
+                y0 = model.init_eval(symbolic_inputs)
             else:
                 # Don't update model.y0
                 return None
         else:
             if update_rhs is True:
                 # Recalculate initial conditions for the rhs equations
-                y0_from_inputs = model.init_eval(inputs)
+                y0_from_inputs = model.init_eval(symbolic_inputs)
                 # Reuse old solution for algebraic equations
                 y0_from_model = model.y0
                 len_rhs = model.len_rhs
@@ -534,7 +541,12 @@ def _set_initial_conditions(self, model, inputs, update_rhs):
                     model.y0 = casadi.vertcat(
                         y0_from_inputs[:len_rhs], y0_from_model[len_rhs:]
                     )
-            model.y0 = self.calculate_consistent_state(model, 0, inputs)
+            y0 = self.calculate_consistent_state(model, 0, inputs)
+        # Make y0 a function of inputs if doing symbolic with casadi
+        if self.sensitivity == "casadi":
+            model.y0 = casadi.Function("y0", [symbolic_inputs], [y0])
+        else:
+            model.y0 = y0
 
     def calculate_consistent_state(self, model, time=0, inputs=None):
         """

pybamm/solvers/casadi_algebraic_solver.py

@@ -107,10 +107,11 @@ def _integrate(self, model, t_eval, inputs=None):
         else:
             # Set up
             t_sym = casadi.MX.sym("t")
+            y0_diff_sym = casadi.MX.sym("y0_diff", y0_diff.shape[0])
             y_alg_sym = casadi.MX.sym("y_alg", y0_alg.shape[0])
-            y_sym = casadi.vertcat(y0_diff, y_alg_sym)
+            y_sym = casadi.vertcat(y0_diff_sym, y_alg_sym)
 
-            t_and_inputs_sym = casadi.vertcat(t_sym, symbolic_inputs)
+            t_y0diff_inputs_sym = casadi.vertcat(t_sym, y0_diff_sym, symbolic_inputs)
             alg = model.casadi_algebraic(t_sym, y_sym, symbolic_inputs)
 
             # Set constraints vector in the casadi format
@@ -126,7 +127,7 @@ def _integrate(self, model, t_eval, inputs=None):
             roots = casadi.rootfinder(
                 "roots",
                 "newton",
-                dict(x=y_alg_sym, p=t_and_inputs_sym, g=alg),
+                dict(x=y_alg_sym, p=t_y0diff_inputs_sym, g=alg),
                 {
                     **self.extra_options,
                     "abstol": self.tol,
@@ -155,12 +156,12 @@ def _integrate(self, model, t_eval, inputs=None):
                 # If doing sensitivity with casadi, evaluate with symbolic inputs
                 # Otherwise, evaluate with actual inputs
                 if self.sensitivity == "casadi":
-                    t_eval_and_inputs = casadi.vertcat(t, symbolic_inputs)
+                    t_y0_diff_inputs = casadi.vertcat(t, y0_diff, symbolic_inputs)
                 else:
-                    t_eval_and_inputs = casadi.vertcat(t, inputs)
+                    t_y0_diff_inputs = casadi.vertcat(t, y0_diff, inputs)
                 # Solve
                 try:
-                    y_alg_sol = roots(y0_alg, t_eval_and_inputs)
+                    y_alg_sol = roots(y0_alg, t_y0_diff_inputs)
                     success = True
                     message = None
                     # Check final output

pybamm/solvers/casadi_solver.py

@@ -208,7 +208,7 @@ def _integrate(self, model, t_eval, inputs=None):
 
                     if self.mode == "safe":
                         # update integrator with the grid
-                        self.create_integrator(model, inputs, t_window)
+                        self.create_integrator(model, inputs_dict, t_window)
                     # Try to solve with the current global step, if it fails then
                     # halve the step size and try again.
                     try:
@@ -347,7 +347,6 @@ def create_integrator(self, model, inputs_dict, t_eval=None):
                 self.integrators[model] = (integrator, use_grid)
                 return integrator
         else:
-            y0 = model.y0
             rhs = model.casadi_rhs
             algebraic = model.casadi_algebraic
 
@@ -370,6 +369,12 @@ def create_integrator(self, model, inputs_dict, t_eval=None):
             # set up and solve
             t = casadi.MX.sym("t")
             p = casadi.MX.sym("p", inputs.shape[0])
+            # If the initial conditions depend on inputs, evaluate the function
+            if isinstance(model.y0, casadi.Function):
+                y0 = model.y0(p)
+            else:
+                y0 = model.y0
+
             y_diff = casadi.MX.sym("y_diff", rhs(0, y0, p).shape[0])
 
             if use_grid is False:
@@ -420,6 +425,13 @@ def _run_integrator(self, model, y0, inputs_dict, t_eval):
         else:
             inputs_eval = inputs
         integrator, use_grid = self.integrators[model]
+
+        # If the initial conditions depend on inputs, evaluate the function
+        if isinstance(y0, casadi.Function):
+            y0 = y0(symbolic_inputs)
+        else:
+            y0 = y0
+
         # 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]: