perf: speed up time loop in processed var init sens (#5000)

* perf: speed up time loop in processed var init sens

* add changelog

Author: martinjrobins

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - Add a solver option to change `on_extrapolation` behavior to `"error"`, `"warn"`, or `"ignore"` on extrapolation events. ([#4993](https://github.com/pybamm-team/PyBaMM/pull/4993))
 - Improve reliability of `CasadiAlgebraicSolver` and added an option for the `step_tol` of the Newton iteration. ([#4985](https://github.com/pybamm-team/PyBaMM/pull/4985))
+- Speed up calculation of variable sensitivities in `ProcessedVariable` ([#5000](https://github.com/pybamm-team/PyBaMM/pull/5000))
 
 ## Bug fixes
 

src/pybamm/solvers/processed_variable.py

@@ -443,16 +443,18 @@ def initialise_sensitivity_explicit_forward(self):
             dvar_dp_func = casadi.Function(
                 "dvar_dp", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dp]
             )
-            for idx, t in enumerate(ts):
-                u = ys[:, idx]
-                next_dvar_dy_eval = dvar_dy_func(t, u, inputs_stacked)
-                next_dvar_dp_eval = dvar_dp_func(t, u, inputs_stacked)
-                if idx == 0:
-                    dvar_dy_eval = next_dvar_dy_eval
-                    dvar_dp_eval = next_dvar_dp_eval
-                else:
-                    dvar_dy_eval = casadi.diagcat(dvar_dy_eval, next_dvar_dy_eval)
-                    dvar_dp_eval = casadi.vertcat(dvar_dp_eval, next_dvar_dp_eval)
+            dvar_dy_eval = casadi.diagcat(
+                *[
+                    dvar_dy_func(t, ys[:, idx], inputs_stacked)
+                    for idx, t in enumerate(ts)
+                ]
+            )
+            dvar_dp_eval = casadi.vertcat(
+                *[
+                    dvar_dp_func(t, ys[:, idx], inputs_stacked)
+                    for idx, t in enumerate(ts)
+                ]
+            )
 
             # Compute sensitivity
             S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval