Merge pull request #2883 from pybamm-team/i2789-short-discharge

#2789 add "all" to sensitivities for idaklu solver

Author: valentinsulzer

CHANGELOG.md

@@ -10,6 +10,7 @@
 ## Bug fixes
 
 - Fixed a bug in the discretisation of initial conditions of a scaled variable ([#2856](https://github.com/pybamm-team/PyBaMM/pull/2856))
+- Fixed keyerror on "all" when getting sensitivities from IDAKLU solver([#2883](https://github.com/pybamm-team/PyBaMM/pull/2883))
 
 # Breaking changes
 

pybamm/solvers/idaklu_solver.py

@@ -532,14 +532,19 @@ 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)
@@ -251,6 +252,10 @@ def test_ida_roberts_klu_sensitivities(self):
 
             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
         # example provided in sundials