#1100 working on casadi solver

valentinsulzer · valentinsulzer · commit b8875155570f · 2020-07-22T14:02:49.000-04:00
diff --git a/pybamm/solvers/casadi_algebraic_solver.py b/pybamm/solvers/casadi_algebraic_solver.py
@@ -40,6 +40,9 @@ def __init__(self, tol=1e-6, extra_options=None, sensitivity=None):
         self.extra_options = extra_options or {}
         pybamm.citations.register("Andersson2019")
 
+        self.rootfinders = {}
+        self.y_sols = {}
+
     @property
     def tol(self):
         return self._tol
@@ -62,24 +65,12 @@ def _integrate(self, model, t_eval, inputs=None):
             Any input parameters to pass to the model when solving.
         """
         # Record whether there are any symbolic inputs
-        inputs = inputs or {}
+        inputs_dict = inputs or {}
         # Create casadi objects for the root-finder
-        inputs_dict = inputs
         inputs = casadi.vertcat(*[v for v in inputs.values()])
 
-        if self.sensitivity == "casadi" and inputs_dict != {}:
-            # Create symbolic inputs for sensitivity analysis
-            symbolic_inputs_list = []
-            for name, value in inputs_dict.items():
-                if isinstance(value, numbers.Number):
-                    symbolic_inputs_list.append(casadi.MX.sym(name))
-                else:
-                    symbolic_inputs_list.append(casadi.MX.sym(name, value.shape[0]))
-            symbolic_inputs = casadi.vertcat(*[p for p in symbolic_inputs_list])
-            inputs_for_alg = symbolic_inputs
-        else:
-            symbolic_inputs = casadi.DM()
-            inputs_for_alg = inputs
+        # Create symbolic inputs
+        symbolic_inputs = casadi.MX.sym("inputs", inputs.shape[0])
 
         y0 = model.y0
         # The casadi algebraic solver can read rhs equations, but leaves them unchanged
@@ -101,34 +92,50 @@ def _integrate(self, model, t_eval, inputs=None):
 
         y_alg = None
 
-        # Set up
-        t_sym = casadi.MX.sym("t")
-        y_alg_sym = casadi.MX.sym("y_alg", y0_alg.shape[0])
-        y_sym = casadi.vertcat(y0_diff, y_alg_sym)
-
-        t_and_inputs_sym = casadi.vertcat(t_sym, symbolic_inputs)
-        alg = model.casadi_algebraic(t_sym, y_sym, inputs_for_alg)
-
-        # Set constraints vector in the casadi format
-        # Constrain the unknowns. 0 (default): no constraint on ui, 1: ui >= 0.0,
-        # -1: ui <= 0.0, 2: ui > 0.0, -2: ui < 0.0.
-        constraints = np.zeros_like(model.bounds[0], dtype=int)
-        # If the lower bound is positive then the variable must always be positive
-        constraints[model.bounds[0] >= 0] = 1
-        # If the upper bound is negative then the variable must always be negative
-        constraints[model.bounds[1] <= 0] = -1
-
-        # Set up rootfinder
-        roots = casadi.rootfinder(
-            "roots",
-            "newton",
-            dict(x=y_alg_sym, p=t_and_inputs_sym, g=alg),
-            {
-                **self.extra_options,
-                "abstol": self.tol,
-                "constraints": list(constraints[len_rhs:]),
-            },
-        )
+        if model in self.rootfinders:
+            if self.sensitivity == "casadi":
+                # Reuse (symbolic) solution with new inputs
+                y_sol = self.y_sols[model]
+                return pybamm.Solution(
+                    t_eval,
+                    y_sol,
+                    termination="success",
+                    model=model,
+                    inputs=inputs_dict,
+                )
+            roots = self.rootfinders[model]
+        else:
+            # Set up
+            t_sym = casadi.MX.sym("t")
+            y_alg_sym = casadi.MX.sym("y_alg", y0_alg.shape[0])
+            y_sym = casadi.vertcat(y0_diff, y_alg_sym)
+
+            t_and_inputs_sym = casadi.vertcat(t_sym, symbolic_inputs)
+            alg = model.casadi_algebraic(t_sym, y_sym, symbolic_inputs)
+
+            # Set constraints vector in the casadi format
+            # Constrain the unknowns. 0 (default): no constraint on ui, 1: ui >= 0.0,
+            # -1: ui <= 0.0, 2: ui > 0.0, -2: ui < 0.0.
+            constraints = np.zeros_like(model.bounds[0], dtype=int)
+            # If the lower bound is positive then the variable must always be positive
+            constraints[model.bounds[0] >= 0] = 1
+            # If the upper bound is negative then the variable must always be negative
+            constraints[model.bounds[1] <= 0] = -1
+
+            # Set up rootfinder
+            roots = casadi.rootfinder(
+                "roots",
+                "newton",
+                dict(x=y_alg_sym, p=t_and_inputs_sym, g=alg),
+                {
+                    **self.extra_options,
+                    "abstol": self.tol,
+                    "constraints": list(constraints[len_rhs:]),
+                },
+            )
+
+            self.rootfinders[model] = roots
+
         for idx, t in enumerate(t_eval):
             # Evaluate algebraic with new t and previous y0, if it's already close
             # enough then keep it
@@ -145,10 +152,15 @@ def _integrate(self, model, t_eval, inputs=None):
                     y_alg = casadi.horzcat(y_alg, y0_alg)
             # Otherwise calculate new y_sol
             else:
-                t_eval_inputs_sym = casadi.vertcat(t, symbolic_inputs)
+                # 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)
+                else:
+                    t_eval_and_inputs = casadi.vertcat(t, inputs)
                 # Solve
                 try:
-                    y_alg_sol = roots(y0_alg, t_eval_inputs_sym)
+                    y_alg_sol = roots(y0_alg, t_eval_and_inputs)
                     success = True
                     message = None
                     # Check final output
@@ -193,6 +205,8 @@ def _integrate(self, model, t_eval, inputs=None):
         # If doing sensitivity, return the solution as a function of the inputs
         if self.sensitivity == "casadi":
             y_sol = casadi.Function("y_sol", [symbolic_inputs], [y_sol])
+            # Save the solution, can just reuse and change the inputs
+            self.y_sols[model] = y_sol
         # Return solution object (no events, so pass None to t_event, y_event)
         return pybamm.Solution(
             t_eval, y_sol, termination="success", model=model, inputs=inputs_dict
diff --git a/pybamm/solvers/casadi_solver.py b/pybamm/solvers/casadi_solver.py
@@ -59,9 +59,12 @@ class CasadiSolver(pybamm.BaseSolver):
         Please consult `CasADi documentation <https://tinyurl.com/y5rk76os>`_ for
         details.
     sensitivity : bool, optional
-        Whether to explicitly formulate and solve the forward sensitivity equations.
-        See :class:`pybamm.BaseSolver`
+        Whether (and how) to calculate sensitivities when solving. Options are:
 
+        - None: no sensitivities
+        - "explicit forward": explicitly formulate the sensitivity equations.
+        See :class:`pybamm.BaseSolver`
+        - "casadi": use casadi to differentiate through the integrator
     """
 
     def __init__(
@@ -104,6 +107,7 @@ def __init__(
         # Initialize
         self.integrators = {}
         self.integrator_specs = {}
+        self.y_sols = {}
 
         pybamm.citations.register("Andersson2019")
 
@@ -122,24 +126,29 @@ def _integrate(self, model, t_eval, inputs=None):
         """
         # Record whether there are any symbolic inputs
         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_dict.values()])
 
-        if has_symbolic_inputs:
-            # Create integrator without grid to avoid having to create several times
-            self.create_integrator(model, inputs)
-            solution = self._run_integrator(model, model.y0, inputs_dict, t_eval)
+        if self.sensitivity == "casadi" and inputs_dict != {}:
+            # If the solution has already been created, we can reuse it
+            if model in self.y_sols:
+                y_sol = self.y_sols[model]
+                solution = pybamm.Solution(
+                    t_eval, y_sol, model=model, inputs=inputs_dict
+                )
+            else:
+                # Create integrator without grid, which will be called repeatedly
+                # This is necessary for casadi to compute sensitivities
+                self.create_integrator(model, inputs_dict)
+                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.create_integrator(model, inputs, t_eval)
+            self.create_integrator(model, inputs_dict, t_eval)
             solution = self._run_integrator(model, model.y0, inputs_dict, t_eval)
             solution.termination = "final time"
             return solution
@@ -161,7 +170,7 @@ def _integrate(self, model, t_eval, inputs=None):
                 # in "safe without grid" mode,
                 # create integrator once, without grid,
                 # to avoid having to create several times
-                self.create_integrator(model, inputs)
+                self.create_integrator(model, inputs_dict)
                 # Initialize solution
                 solution = pybamm.Solution(
                     np.array([t]), y0[:, np.newaxis], model=model, inputs=inputs_dict
@@ -314,12 +323,15 @@ def event_fun(t):
                     y0 = solution.y[:, -1]
             return solution
 
-    def create_integrator(self, model, inputs, t_eval=None):
+    def create_integrator(self, model, inputs_dict, t_eval=None):
         """
         Method to create a casadi integrator object.
         If t_eval is provided, the integrator uses t_eval to make the grid.
         Otherwise, the integrator has grid [0,1].
         """
+        # convert inputs to casadi format
+        inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
+
         # Use grid if t_eval is given
         use_grid = not (t_eval is None)
         # Only set up problem once
@@ -400,6 +412,13 @@ def create_integrator(self, model, inputs, t_eval=None):
 
     def _run_integrator(self, model, y0, inputs_dict, t_eval):
         inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
+        symbolic_inputs = casadi.MX.sym("inputs", inputs.shape[0])
+        # If doing sensitivity with casadi, evaluate with symbolic inputs
+        # Otherwise, evaluate with actual inputs
+        if self.sensitivity == "casadi":
+            inputs_eval = symbolic_inputs
+        else:
+            inputs_eval = inputs
         integrator, use_grid = self.integrators[model]
         # Split up initial conditions into differential and algebraic
         # Check y0 to see if it includes sensitivities
@@ -415,10 +434,9 @@ def _run_integrator(self, model, y0, inputs_dict, t_eval):
             if use_grid is True:
                 # Call the integrator once, with the grid
                 sol = integrator(
-                    x0=y0_diff, z0=y0_alg, p=inputs, **self.extra_options_call
+                    x0=y0_diff, z0=y0_alg, p=inputs_eval, **self.extra_options_call
                 )
                 y_sol = np.concatenate([sol["xf"].full(), sol["zf"].full()])
-                return pybamm.Solution(t_eval, y_sol, model=model, inputs=inputs_dict)
             else:
                 # Repeated calls to the integrator
                 x = y0_diff
@@ -428,7 +446,7 @@ def _run_integrator(self, model, y0, inputs_dict, t_eval):
                 for i in range(len(t_eval) - 1):
                     t_min = t_eval[i]
                     t_max = t_eval[i + 1]
-                    inputs_with_tlims = casadi.vertcat(inputs, t_min, t_max)
+                    inputs_with_tlims = casadi.vertcat(inputs_eval, t_min, t_max)
                     sol = integrator(
                         x0=x, z0=z, p=inputs_with_tlims, **self.extra_options_call
                     )
@@ -438,14 +456,15 @@ def _run_integrator(self, model, y0, inputs_dict, 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, model=model, inputs=inputs_dict
-                    )
+                    y_sol = y_diff
                 else:
                     y_sol = casadi.vertcat(y_diff, y_alg)
-                    return pybamm.Solution(
-                        t_eval, y_sol, model=model, inputs=inputs_dict
-                    )
+            # If doing sensitivity, return the solution as a function of the inputs
+            if self.sensitivity == "casadi":
+                y_sol = casadi.Function("y_sol", [symbolic_inputs], [y_sol])
+                # Save the solution, can just reuse and change the inputs
+                self.y_sols[model] = 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/processed_variable.py b/pybamm/solvers/processed_variable.py
@@ -722,11 +722,7 @@ def initialise_1D_symbolic():
                     + "implemented)"
                 )
 
-        # Make entries a function and compute jacobian
-        casadi_entries_fn = casadi.Function(
-            "variable", [self.symbolic_inputs], [entries_MX]
-        )
-
+        # Compute jacobian
         sens_MX = casadi.jacobian(entries_MX, self.symbolic_inputs)
         casadi_sens_fn = casadi.Function("variable", [self.symbolic_inputs], [sens_MX])
 
diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py
