#2789 add "all" to sensitivities for idaklu solver

Author: martinjrobins

pybamm/solvers/idaklu_solver.py

@@ -532,14 +532,18 @@ def _integrate(self, model, t_eval, inputs_dict=None):
         y_out = sol.y.reshape((number_of_timesteps, number_of_states))
 
         # return sensitivity solution, we need to flatten yS to
-        # (#timesteps * #states,) to match format used by Solution
+        # (#timesteps * #states (where t is changing the quickest),) to match format used by Solution
+        # note that yS is (n_p, n_t, n_y)  
         if number_of_sensitivity_parameters != 0:
             yS_out = {
                 name: sol.yS[i].reshape(-1, 1)
                 for i, name in enumerate(sensitivity_names)
             }
+            # add "all" stacked sensitivities ((#timesteps * #states,#sens_params))
+            yS_out["all"] = np.hstack([yS_out[name] for name in sensitivity_names])
         else:
             yS_out = False
+            
         if sol.flag in [0, 2]:
             # 0 = solved for all t_eval
             if sol.flag == 0:

tests/unit/test_solvers/test_idaklu_solver.py

@@ -200,6 +200,7 @@ def test_ida_roberts_klu_sensitivities(self):
             model.rhs = {u: a * v}
             model.algebraic = {v: 1 - v}
             model.initial_conditions = {u: 0, v: 1}
+            model.variables = {"2u": 2 * u}
 
             disc = pybamm.Discretisation()
             disc.process_model(model)
@@ -250,6 +251,11 @@ def test_ida_roberts_klu_sensitivities(self):
             dyda_fd = dyda_fd.transpose().reshape(-1, 1)
 
             np.testing.assert_array_almost_equal(dyda_ida, dyda_fd)
+            
+            # get the sensitivities for the variable
+            d2uda = sol["2u"].sensitivities["a"]
+            np.testing.assert_array_almost_equal(2 * dyda_ida[0:200:2], d2uda)
+            
 
     def test_sensitivities_with_events(self):
         # this test implements a python version of the ida Roberts