#1477 going to take out sensitivity=casadi option

Author: martinjrobins

pybamm/solvers/base_solver.py

@@ -42,11 +42,7 @@ class BaseSolver(object):
         the solution instance returned. At the moment this is only implemented for the
         IDAKLU solver.\
         - "explicit forward": explicitly formulate the sensitivity equations for
-        the chosen 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. At the moment this is only
+        the chosen input parameters. . At the moment this is only
         implemented using convert_to_format = 'casadi'. \
         - see individual solvers for other options
     """
@@ -60,7 +56,6 @@ def __init__(
         root_tol=1e-6,
         extrap_tol=0,
         max_steps="deprecated",
-        sensitivity=None,
     ):
         self._method = method
         self._rtol = rtol
@@ -79,7 +74,6 @@ def __init__(
         self.name = "Base solver"
         self.ode_solver = False
         self.algebraic_solver = False
-        self.sensitivity = sensitivity
 
     @property
     def method(self):
@@ -140,8 +134,6 @@ def copy(self):
         new_solver.models_set_up = {}
         return new_solver
 
-
-
     def set_up(self, model, inputs=None, t_eval=None,
                calculate_sensitivites=False):
         """Unpack model, perform checks, and calculate jacobian.
@@ -238,21 +230,17 @@ def set_up(self, model, inputs=None, t_eval=None,
                 calculate_sensitivites = [p for p in inputs.keys()]
             else:
                 calculate_sensitivites = []
+
+        calculate_sensitivites_explicit = False
+        if calculate_sensitivites and not isinstance(self, pybamm.IDAKLUSolver):
+            calculate_sensitivites_explicit = True
+
         # save sensitivity parameters so we can identify them later on
         # (FYI: this is used in the Solution class)
         model.calculate_sensitivities = calculate_sensitivites
-        model.len_rhs_sens = model.len_rhs * len(calculate_sensitivites)
-        model.len_alg_sens = model.len_alg * len(calculate_sensitivites)
-
-        # Only allow solving explicit sensitivity equations with the casadi format for now
-        if (
-            self.sensitivity == "explicit forward"
-            and model.convert_to_format != "casadi"
-        ):
-            raise NotImplementedError(
-                "model should be converted to casadi format in order to solve "
-                "explicit sensitivity equations"
-            )
+        if calculate_sensitivites_explicit:
+            model.len_rhs_sens = model.len_rhs * len(calculate_sensitivites)
+            model.len_alg_sens = model.len_alg * len(calculate_sensitivites)
 
         if model.convert_to_format != "casadi":
             # Create Jacobian from concatenated rhs and algebraic
@@ -275,7 +263,7 @@ def set_up(self, model, inputs=None, t_eval=None,
                     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.sensitivity == "explicit forward":
+            if calculate_sensitivites_explicit:
                 pS_casadi_stacked = casadi.vertcat(
                     *[p_casadi[name] for name in calculate_sensitivites]
                 )
@@ -297,15 +285,19 @@ def report(string):
             if model.convert_to_format == "jax":
                 report(f"Converting {name} to jax")
                 func = pybamm.EvaluatorJax(func)
-                if calculate_sensitivites:
+                jacp = None
+                if calculate_sensitivites_explicit:
+                    raise NotImplementedError(
+                        "sensitivities using convert_to_format = 'jax' "
+                        "only implemented for IDAKLUSolver"
+                    )
+                elif calculate_sensitivites:
                     report((
                         f"Calculating sensitivities for {name} with respect "
                         f"to parameters {calculate_sensitivites} using jax"
                     ))
                     jacp = func.get_sensitivities()
                     jacp = jacp.evaluate
-                else:
-                    jacp = None
                 if use_jacobian:
                     report(f"Calculating jacobian for {name} using jax")
                     jac = func.get_jacobian()
@@ -319,6 +311,10 @@ def report(string):
                 # Process with pybamm functions, optionally converting
                 # to python evaluator
                 if calculate_sensitivites:
+                    raise NotImplementedError(
+                        "sensitivities only implemented with "
+                        "convert_to_format = 'casadi' or convert_to_format = 'jax'"
+                    )
                     report((
                         f"Calculating sensitivities for {name} with respect "
                         f"to parameters {calculate_sensitivites}"
@@ -362,9 +358,16 @@ def jacp(*args, **kwargs):
                 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.sensitivity == "explicit forward":
+                jacp = None
+                if calculate_sensitivites_explicit:
+                    # 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
                     if name == "rhs" and model.len_rhs > 0:
-                        report("Creating sensitivity equations for rhs using CasADi")
+                        report(
+                            "Creating explicit forward sensitivity equations for rhs using CasADi")
                         df_dx = casadi.jacobian(func, y_diff)
                         df_dp = casadi.jacobian(func, pS_casadi_stacked)
                         S_x_mat = S_x.reshape(
@@ -383,7 +386,7 @@ def jacp(*args, **kwargs):
                         func = casadi.vertcat(func, S_rhs)
                     if name == "algebraic" and model.len_alg > 0:
                         report(
-                            "Creating sensitivity equations for algebraic using CasADi"
+                            "Creating explicit forward sensitivity equations for algebraic using CasADi"
                         )
                         dg_dz = casadi.jacobian(func, y_alg)
                         dg_dp = casadi.jacobian(func, pS_casadi_stacked)
@@ -401,7 +404,12 @@ def jacp(*args, **kwargs):
                                 (-1, 1)
                             )
                         func = casadi.vertcat(func, S_alg)
-                    elif name == "initial_conditions":
+                    if name == "residuals":
+                        raise NotImplementedError(
+                            "explicit forward equations not implimented for residuals"
+                        )
+
+                    if name == "initial_conditions":
                         if model.len_rhs == 0 or model.len_alg == 0:
                             S_0 = casadi.jacobian(func, pS_casadi_stacked).reshape(
                                 (-1, 1)
@@ -417,16 +425,7 @@ def jacp(*args, **kwargs):
                                 (-1, 1)
                             )
                             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_and_S)
-                    jac = casadi.Function(
-                        name, [t_casadi, y_and_S, p_casadi_stacked], [jac_casadi]
-                    )
-                else:
-                    jac = None
-
-                if calculate_sensitivites and self.sensitivity != "explicit forward":
+                elif calculate_sensitivites:
                     report((
                         f"Calculating sensitivities for {name} with respect "
                         f"to parameters {calculate_sensitivites} using CasADi"
@@ -444,8 +443,14 @@ def jacp(*args, **kwargs):
                         return {k: v(*args, **kwargs)
                                 for k, v in jacp_dict.items()}
 
+                if use_jacobian:
+                    report(f"Calculating jacobian for {name} using CasADi")
+                    jac_casadi = casadi.jacobian(func, y_and_S)
+                    jac = casadi.Function(
+                        name, [t_casadi, y_and_S, p_casadi_stacked], [jac_casadi]
+                    )
                 else:
-                    jacp = None
+                    jac = None
 
                 func = casadi.Function(
                     name, [t_casadi, y_and_S, p_casadi_stacked], [func]
@@ -538,7 +543,7 @@ def jacp(*args, **kwargs):
         )[0]
         init_eval = InitialConditions(initial_conditions, model)
 
-        if self.sensitivity == "explicit forward":
+        if calculate_sensitivites_explicit:
             y0_total_size = (
                 model.len_rhs + model.len_rhs_sens
                 + model.len_alg + model.len_alg_sens
@@ -555,7 +560,6 @@ def jacp(*args, **kwargs):
 
         # Calculate initial conditions
         model.y0 = init_eval(inputs)
-        print('YYYYY', model.y0)
 
         casadi_terminate_events = []
         terminate_events_eval = []
@@ -726,7 +730,6 @@ def _set_initial_conditions(self, model, inputs, update_rhs):
             model.y0 = casadi.Function("y0", [symbolic_inputs], [y0])
         else:
             model.y0 = y0
-        print('ASDF', model.y0)
 
     def calculate_consistent_state(self, model, time=0, inputs=None):
         """

pybamm/solvers/casadi_solver.py

@@ -62,13 +62,6 @@ 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.
-    sensitivity : str, optional
-        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__(
@@ -83,7 +76,6 @@ def __init__(
         extrap_tol=0,
         extra_options_setup=None,
         extra_options_call=None,
-        sensitivity=None,
     ):
         super().__init__(
             "problem dependent",
@@ -92,7 +84,6 @@ def __init__(
             root_method,
             root_tol,
             extrap_tol,
-            sensitivity=sensitivity,
         )
         if mode in ["safe", "fast", "fast with events", "safe without grid"]:
             self.mode = mode
@@ -138,6 +129,9 @@ def _integrate(self, model, t_eval, inputs_dict=None):
 
         # Record whether there are any symbolic inputs
         inputs_dict = inputs_dict 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()])
@@ -176,7 +170,7 @@ def _integrate(self, model, t_eval, inputs_dict=None):
             else:
                 # Create integrator without grid, which will be called repeatedly
                 # This is necessary for casadi to compute sensitivities
-                self.create_integrator(model, inputs_dict)
+                self.create_integrator(model, inputs)
                 solution = self._run_integrator(
                     model, model.y0, inputs_dict, inputs, t_eval
                 )
@@ -216,7 +210,7 @@ def _integrate(self, model, t_eval, inputs_dict=None):
                 # in "safe without grid" mode,
                 # create integrator once, without grid,
                 # to avoid having to create several times
-                self.create_integrator(model, inputs_dict)
+                self.create_integrator(model, inputs)
                 # Initialize solution
                 solution = pybamm.Solution(
                     np.array([t]), y0, model, inputs_dict,
@@ -258,7 +252,7 @@ def _integrate(self, model, t_eval, inputs_dict=None):
 
                     if self.mode == "safe":
                         # update integrator with the grid
-                        self.create_integrator(model, inputs_dict, t_window)
+                        self.create_integrator(model, inputs, t_window)
                     # Try to solve with the current global step, if it fails then
                     # halve the step size and try again.
                     try:

pybamm/solvers/idaklu_solver.py

@@ -62,6 +62,11 @@ def __init__(
         if idaklu_spec is None:
             raise ImportError("KLU is not installed")
 
+        if sensitivity == "explicit forward":
+            raise NotImplementedError(
+                "Cannot use explicit forward equations with IDAKLUSolver"
+            )
+
         super().__init__(
             "ida",
             rtol,
@@ -188,12 +193,9 @@ def _integrate(self, model, t_eval, inputs_dict=None):
             atol = self._atol
 
         y0 = model.y0
-        print('idaklu, y0', y0)
         if isinstance(y0, casadi.DM):
             y0 = y0.full().flatten()
 
-        print('idaklu, y0', y0)
-
         rtol = self._rtol
         atol = self._check_atol_type(atol, y0.size)
 

pybamm/solvers/processed_variable.py

@@ -46,7 +46,7 @@ def __init__(self, base_variables, base_variables_casadi, solution, warn=True):
         self.auxiliary_domains = base_variables[0].auxiliary_domains
         self.warn = warn
 
-        self.symbolic_inputs = solution._symbolic_inputs
+        self.symbolic_inputs = solution.has_symbolic_inputs
 
         self.u_sol = solution.y
         self.y_sym = solution._y_sym

tests/integration/test_solvers/test_idaklu.py

@@ -19,6 +19,47 @@ def test_on_spme(self):
         solution = pybamm.IDAKLUSolver().solve(model, t_eval)
         np.testing.assert_array_less(1, solution.t.size)
 
+    def test_on_spme_sensitivities(self):
+        param_name = "Negative electrode conductivity [S.m-1]"
+        neg_electrode_cond = 100.0
+        model = pybamm.lithium_ion.SPMe()
+        geometry = model.default_geometry
+        param = model.default_parameter_values
+        param.update({param_name: "[input]"})
+        inputs = {param_name: neg_electrode_cond}
+        param.process_model(model)
+        param.process_geometry(geometry)
+        mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)
+        disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
+        disc.process_model(model)
+        t_eval = np.linspace(0, 3600, 100)
+        solver = pybamm.IDAKLUSolver()
+        solution = solver.solve(
+            model, t_eval,
+            inputs=inputs,
+            calculate_sensitivities=True,
+        )
+        np.testing.assert_array_less(1, solution.t.size)
+
+        # evaluate the sensitivities using idas
+        dyda_ida = solution.sensitivities[param_name]
+
+        # evaluate the sensitivities using finite difference
+        h = 1e-6
+        sol_plus = solver.solve(
+            model, t_eval,
+            inputs={param_name: neg_electrode_cond + 0.5 * h}
+        )
+        sol_neg = solver.solve(
+            model, t_eval,
+            inputs={param_name: neg_electrode_cond - 0.5 * h}
+        )
+        dyda_fd = (sol_plus.y - sol_neg.y) / h
+
+        np.testing.assert_array_almost_equal(
+            dyda_ida, dyda_fd
+        )
+
     def test_set_tol_by_variable(self):
         model = pybamm.lithium_ion.SPMe()
         geometry = model.default_geometry
@@ -68,6 +109,7 @@ def test_changing_grid(self):
 if __name__ == "__main__":
     print("Add -v for more debug output")
 
+    pybamm.set_logging_level('INFO')
     if "-v" in sys.argv:
         debug = True
     pybamm.settings.debug_mode = True