#1477 evaluating sensitivities ok for all convert_tos

Author: martinjrobins

pybamm/expression_tree/operations/evaluate.py

@@ -628,7 +628,7 @@ def get_sensitivities(self):
         self._sens_evaluate = jax.jit(jacobian_evaluate,
                                      static_argnums=self._static_argnums)
 
-        return EvaluatorJaxJacobian(self._jac_evaluate, self._constants)
+        return EvaluatorJaxSensitivities(self._sens_evaluate, self._constants)
 
 
     def debug(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
@@ -687,3 +687,25 @@ def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
             return result, known_evals
         else:
             return result
+
+class EvaluatorJaxSensitivities:
+    def __init__(self, jac_evaluate, constants):
+        self._jac_evaluate = jac_evaluate
+        self._constants = constants
+
+    def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None):
+        """
+        Acts as a drop-in replacement for :func:`pybamm.Symbol.evaluate`
+        """
+        # generated code assumes y is a column vector
+        if y is not None and y.ndim == 1:
+            y = y.reshape(-1, 1)
+
+        # execute code
+        result = self._jac_evaluate(*self._constants, t, y, y_dot, inputs, known_evals)
+
+        # don't need known_evals, but need to reproduce Symbol.evaluate signature
+        if known_evals is not None:
+            return result, known_evals
+        else:
+            return result

pybamm/solvers/base_solver.py

@@ -269,6 +269,7 @@ def report(string):
                         f"to parameters {calculate_sensitivites} using jax"
                     ))
                     jacp_dict = func.get_sensitivities()
+                    jacp_dict = jacp_dict.evaluate
                 else:
                     jacp_dict = None
                 if use_jacobian:
@@ -283,13 +284,11 @@ def report(string):
             elif model.convert_to_format != "casadi":
                 # Process with pybamm functions, optionally converting
                 # to python evaluator
-                print('calculate_sensitivites = ', calculate_sensitivites)
                 if calculate_sensitivites:
                     report((
                         f"Calculating sensitivities for {name} with respect "
                         f"to parameters {calculate_sensitivites}"
                     ))
-                    print(type(func))
                     jacp_dict = {
                         p: func.diff(pybamm.InputParameter(p))
                         for p in calculate_sensitivites
@@ -360,10 +359,17 @@ def report(string):
             else:
                 jac_call = None
             if jacp_dict is not None:
-                jacp_call = {
-                    k: SolverCallable(v, name + "_sensitivity_wrt_" + k, model)
-                    for k, v in jacp_dict.items()
-                }
+                if model.convert_to_format == "jax":
+                    jacp_call = SolverCallable(
+                        jacp_dict, name + "_sensitivity_wrt_inputs", model
+                    )
+                else:
+                    jacp_call = {
+                        k: SolverCallable(v, name + "_sensitivity_wrt_" + k, model)
+                        for k, v in jacp_dict.items()
+                    }
+            else:
+                jacp_call = None
             return func, func_call, jac_call, jacp_call
 
         # Check for heaviside and modulo functions in rhs and algebraic and add
@@ -520,6 +526,8 @@ def report(string):
                     "rhs", [t_casadi, y_casadi, p_casadi_stacked], [explicit_rhs]
                 )
             model.casadi_algebraic = algebraic
+            model.casadi_sensitivities_rhs = jacp_rhs
+            model.casadi_sensitivities_algebraic = jacp_algebraic
         if len(model.rhs) == 0:
             # No rhs equations: residuals is algebraic only
             model.residuals_eval = Residuals(algebraic, "residuals", model)

tests/unit/test_solvers/test_base_solver.py

@@ -323,15 +323,14 @@ def test_extrapolation_warnings(self):
             solver.solve(model, t_eval=[0, 1])
 
     def test_sensitivities(self):
-        pybamm.set_logging_level('DEBUG')
 
         def exact_diff_a(v, a, b):
-            return v**2 + 2 * a
+            return np.array([v**2 + 2 * a])
 
         def exact_diff_b(v, a, b):
-            return v
+            return np.array([v])
 
-        for f in ['', 'python', 'casadi']:
+        for f in ['', 'python', 'casadi', 'jax']:
             model = pybamm.BaseModel()
             v = pybamm.Variable("v")
             a = pybamm.InputParameter("a")
@@ -342,25 +341,45 @@ def exact_diff_b(v, a, b):
             solver = pybamm.ScipySolver()
             solver.set_up(model, calculate_sensitivites=True,
                           inputs={'a': 0, 'b': 0})
+            all_inputs = []
             for v_value in [0.1, -0.2, 1.5, 8.4]:
                 for a_value in [0.12, 1.5]:
                     for b_value in [0.82, 1.9]:
                         y = np.array([v_value])
                         t = 0
                         inputs = {'a': a_value, 'b': b_value}
+                        all_inputs.append((t, y, inputs))
+            for t, y, inputs in all_inputs:
+                if f == 'casadi':
+                    use_inputs = casadi.vertcat(*[x for x in inputs.values()])
+                else:
+                    use_inputs = inputs
+                if f == 'jax':
+                    sens = model.sensitivities_eval(
+                            t, y, use_inputs
+                    )
+                    np.testing.assert_array_equal(
+                        sens['a'],
+                        exact_diff_a(y, inputs['a'], inputs['b'])
+                    )
+                    np.testing.assert_array_equal(
+                        sens['b'],
+                        exact_diff_b(y, inputs['a'], inputs['b'])
+                    )
+                else:
+                    np.testing.assert_array_equal(
+                        model.sensitivities_eval['a'](
+                            t, y, use_inputs
+                        ),
+                        exact_diff_a(y, inputs['a'], inputs['b'])
+                    )
+                    np.testing.assert_array_equal(
+                        model.sensitivities_eval['b'](
+                            t, y, use_inputs
+                        ),
+                        exact_diff_b(y, inputs['a'], inputs['b'])
+                    )
 
-                        self.assertAlmostEqual(
-                            model.sensitivities_eval['a'](
-                                t=0, y=y, inputs=inputs
-                            ),
-                            exact_diff_a(v_value, a_value, b_value)
-                        )
-                        self.assertAlmostEqual(
-                            model.sensitivities_eval['b'](
-                                t=0, y=y, inputs=inputs
-                            ),
-                            exact_diff_b(v_value, a_value, b_value)
-                        )
 
 if __name__ == "__main__":
     print("Add -v for more debug output")