#1477 draft out a test for idaklu and changes to base solver for sensitivities

Author: martinjrobins

pybamm/expression_tree/operations/evaluate.py

@@ -557,7 +557,7 @@ def __init__(self, symbol):
                 constants[symbol_id] = jax.device_put(constants[symbol_id])
 
         # get a list of constant arguments to input to the function
-        arg_list = [
+        self._arg_list = [
             id_to_python_variable(symbol_id, True) for symbol_id in constants.keys()
         ]
 
@@ -578,8 +578,8 @@ def __init__(self, symbol):
 
         # add function def to first line
         args = "t=None, y=None, y_dot=None, inputs=None, known_evals=None"
-        if arg_list:
-            args = ",".join(arg_list) + ", " + args
+        if self._arg_list:
+            args = ",".join(self._arg_list) + ", " + args
         python_str = "def evaluate_jax({}):\n".format(args) + python_str
 
         # calculate the final variable that will output the result of calling `evaluate`
@@ -604,17 +604,33 @@ def __init__(self, symbol):
         compiled_function = compile(python_str, result_var, "exec")
         exec(compiled_function)
 
-        n = len(arg_list)
-        static_argnums = tuple(static_argnums)
-        self._jit_evaluate = jax.jit(self._evaluate_jax, static_argnums=static_argnums)
+        self._static_argnums = tuple(static_argnums)
+        self._jit_evaluate = jax.jit(self._evaluate_jax,
+                                     static_argnums=self._static_argnums)
 
-        # store a jit version of evaluate_jax's jacobian
+    def get_jacobian(self):
+        n = len(self._arg_list)
+
+        # forward mode autodiff  wrt y, which is argument 1 after arg_list
         jacobian_evaluate = jax.jacfwd(self._evaluate_jax, argnums=1 + n)
-        self._jac_evaluate = jax.jit(jacobian_evaluate, static_argnums=static_argnums)
 
-    def get_jacobian(self):
+        self._jac_evaluate = jax.jit(jacobian_evaluate,
+                                     static_argnums=self._static_argnums)
+
+        return EvaluatorJaxJacobian(self._jac_evaluate, self._constants)
+
+    def get_sensitivities(self):
+        n = len(self._arg_list)
+
+        # forward mode autodiff wrt inputs, which is argument 3 after arg_list
+        jacobian_evaluate = jax.jacfwd(self._evaluate_jax, argnums=3 + n)
+
+        self._sens_evaluate = jax.jit(jacobian_evaluate,
+                                     static_argnums=self._static_argnums)
+
         return EvaluatorJaxJacobian(self._jac_evaluate, self._constants)
 
+
     def debug(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
         # generated code assumes y is a column vector
         if y is not None and y.ndim == 1:

pybamm/solvers/base_solver.py

@@ -34,6 +34,15 @@ class BaseSolver(object):
         The tolerance for the initial-condition solver (default is 1e-6).
     extrap_tol : float, optional
         The tolerance to assert whether extrapolation occurs or not. Default is 0.
+    sensitivity : str, optional
+        Whether (and how) to calculate sensitivities when solving. Options are:
+        - "explicit forward": explicitly formulate the sensitivity equations. \
+        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 \
+        - see individual solvers for other options
+
     """
 
     def __init__(
@@ -45,6 +54,7 @@ def __init__(
         root_tol=1e-6,
         extrap_tol=0,
         max_steps="deprecated",
+        sensitivity=None
     ):
         self._method = method
         self._rtol = rtol
@@ -63,6 +73,7 @@ def __init__(
         self.name = "Base solver"
         self.ode_solver = False
         self.algebraic_solver = False
+        self.sensitivity = sensitivity
 
     @property
     def method(self):
@@ -203,6 +214,10 @@ def set_up(self, model, inputs=None, t_eval=None):
             y = pybamm.StateVector(slice(0, model.concatenated_initial_conditions.size))
             # set up Jacobian object, for re-use of dict
             jacobian = pybamm.Jacobian()
+            jacobian_parameters = {
+                p: pybamm.Jacobian() for p in inputs.keys()
+            }
+
         else:
             # Convert model attributes to casadi
             t_casadi = casadi.MX.sym("t")
@@ -225,32 +240,56 @@ def report(string):
 
             if use_jacobian is None:
                 use_jacobian = model.use_jacobian
-            if model.convert_to_format != "casadi":
-                # Process with pybamm functions
 
-                if model.convert_to_format == "jax":
-                    report(f"Converting {name} to jax")
-                    jax_func = pybamm.EvaluatorJax(func)
+            if model.convert_to_format == "jax":
+                report(f"Converting {name} to jax")
+                func = pybamm.EvaluatorJax(func)
+                if self.sensitivity:
+                    report(f"Calculating sensitivities for {name} using jax")
+                    jacp_dict = func.get_sensitivities()
+                else:
+                    jacp_dict = None
+                if use_jacobian:
+                    report(f"Calculating jacobian for {name} using jax")
+                    jac = func.get_jacobian()
+                    jac = jac.evaluate
+                else:
+                    jac = None
+
+                func = func.evaluate
+
+            elif model.convert_to_format != "casadi":
+                # Process with pybamm functions, optionally converting
+                # to python evaluator
+                if self.sensitivity:
+                    report(f"Calculating sensitivities for {name}")
+                    jacp_dict = {
+                        p: jwrtp.jac(func, pybamm.InputParameter(p))
+                        for jwrtp, p in
+                        zip(jacobian_parameters, inputs.keys())
+                    }
+                    if model.convert_to_format == "python":
+                        report(f"Converting sensitivities for {name} to python")
+                        jacp_dict = {
+                            p: pybamm.EvaluatorPython(jacp)
+                            for p, jacp in jacp_dict.items()
+                        }
+                else:
+                    jacp_dict = None
 
                 if use_jacobian:
                     report(f"Calculating jacobian for {name}")
                     jac = jacobian.jac(func, y)
                     if model.convert_to_format == "python":
                         report(f"Converting jacobian for {name} to python")
                         jac = pybamm.EvaluatorPython(jac)
-                    elif model.convert_to_format == "jax":
-                        report(f"Converting jacobian for {name} to jax")
-                        jac = jax_func.get_jacobian()
                     jac = jac.evaluate
                 else:
                     jac = None
 
                 if model.convert_to_format == "python":
                     report(f"Converting {name} to python")
                     func = pybamm.EvaluatorPython(func)
-                if model.convert_to_format == "jax":
-                    report(f"Converting {name} to jax")
-                    func = jax_func
 
                 func = func.evaluate
 
@@ -266,6 +305,16 @@ def report(string):
                     )
                 else:
                     jac = None
+
+                if self.sensitivity:
+                    report(f"Calculating sensitivities for {name} using CasADi")
+                    jacp_dict = {
+                        name: casadi.jacobian(func, p)
+                        for name, p in p_casadi.items()
+                    }
+                else:
+                    jacp_dict = None
+
                 func = casadi.Function(
                     name, [t_casadi, y_casadi, p_casadi_stacked], [func]
                 )

pybamm/solvers/idaklu_solver.py

@@ -38,6 +38,15 @@ class IDAKLUSolver(pybamm.BaseSolver):
         The tolerance for the initial-condition solver (default is 1e-6).
     extrap_tol : float, optional
         The tolerance to assert whether extrapolation occurs or not (default is 0).
+    sensitivity : str, optional
+        Whether (and how) to calculate sensitivities when solving. Options are:
+        - "explicit forward": explicitly formulate the sensitivity equations. \
+        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 \
+        - "idas": use Sundials IDAS to compute forward sensitivities
+
     """
 
     def __init__(
@@ -48,13 +57,21 @@ def __init__(
         root_tol=1e-6,
         extrap_tol=0,
         max_steps="deprecated",
+        sensitivity="idas"
     ):
 
         if idaklu_spec is None:
             raise ImportError("KLU is not installed")
 
         super().__init__(
-            "ida", rtol, atol, root_method, root_tol, extrap_tol, max_steps
+            "ida",
+            rtol,
+            atol,
+            root_method,
+            root_tol,
+            extrap_tol,
+            max_steps,
+            sensitivity=sensitivity,
         )
         self.name = "IDA KLU solver"
 

tests/unit/test_solvers/test_idaklu_solver.py

@@ -46,6 +46,59 @@ def test_ida_roberts_klu(self):
             true_solution = 0.1 * solution.t
             np.testing.assert_array_almost_equal(solution.y[0, :], true_solution)
 
+    def test_ida_roberts_klu_sensitivities(self):
+        # this test implements a python version of the ida Roberts
+        # example provided in sundials
+        # see sundials ida examples pdf
+        for form in ["python", "casadi"]:
+            model = pybamm.BaseModel()
+            model.convert_to_format = form
+            u = pybamm.Variable("u")
+            v = pybamm.Variable("v")
+            a = pybamm.InputParameter("a")
+            model.rhs = {u: a * v}
+            model.algebraic = {v: 1 - v}
+            model.initial_conditions = {u: 0, v: 1}
+            model.events = [pybamm.Event("1", u - 0.2), pybamm.Event("2", v)]
+
+            disc = pybamm.Discretisation()
+            disc.process_model(model)
+
+            solver = pybamm.IDAKLUSolver(root_method="lm")
+
+            t_eval = np.linspace(0, 3, 100)
+            a_value = 0.1
+            sol = solver.solve(model, t_eval, inputs={"a": a_value})
+
+            # test that final time is time of event
+            # y = 0.1 t + y0 so y=0.2 when t=2
+            np.testing.assert_array_almost_equal(sol.t[-1], 2.0)
+
+            # test that final value is the event value
+            np.testing.assert_array_almost_equal(sol.y[0, -1], 0.2)
+
+            # test that y[1] remains constant
+            np.testing.assert_array_almost_equal(
+                sol.y[1, :], np.ones(sol.t.shape)
+            )
+
+            # test that y[0] = to true solution
+            true_solution = 0.1 * sol.t
+            np.testing.assert_array_almost_equal(sol.y[0, :], true_solution)
+
+            # evaluate the sensitivities using idas
+            dyda_ida = sol.sensitivities["a"]
+
+            # evaluate the sensitivities using finite difference
+            h = 1e-6
+            sol_plus = solver.solve(model, t_eval, inputs={"a": a_value + 0.5 * h})
+            sol_neg = solver.solve(model, t_eval, inputs={"a": a_value - 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_atol(self):
         model = pybamm.lithium_ion.DFN()
         geometry = model.default_geometry