perf: replace all_inputs_casadi with all_inputs_stacked (#5413)

* perf: replace all_inputs_casadi with all_inputs_stacked

* Update solution.py

* add test

* Update CHANGELOG.md

* add cache tests and remove cached property

Author: MarcBerliner

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- Improved the performance of processed variables by replacing `casadi.vertcat` input stacking with numpy vectors. ([#5413](https://github.com/pybamm-team/PyBaMM/pull/5413))
 - Allow out of bounds initial state of charge to enable initialising a simulation at a voltage outside the voltage limits. ([#5386](https://github.com/pybamm-team/PyBaMM/pull/5386))
 - Added `cache_esoh` option to `Simulation` that caches the electrode SOH computation across repeated `solve` calls, avoiding redundant recalculation when eSOH-relevant parameters have not changed. The cached eSOH solver/simulation object is also reused on cache misses to skip expensive model rebuilding. ([#5408](https://github.com/pybamm-team/PyBaMM/pull/5408))
 - Eliminated the mass matrix inverse and temporary dense matrix objects when building the simulation. ([#5391](https://github.com/pybamm-team/PyBaMM/pull/5391))

src/pybamm/solvers/idaklu_solver.py

@@ -1090,6 +1090,7 @@ def reduce_solution(
         )
 
         # Propagate metadata from the original solution
+        new_sol._all_inputs_stacked = solution.all_inputs_stacked
         new_sol._all_inputs_casadi = solution.all_inputs_casadi
         new_sol.closest_event_idx = solution.closest_event_idx
 

src/pybamm/solvers/processed_variable.py

@@ -52,7 +52,7 @@ def __init__(
         self.all_ys = solution.all_ys
         self.all_yps = solution.all_yps
         self.all_inputs = solution.all_inputs
-        self.all_inputs_casadi = solution.all_inputs_casadi
+        self.all_inputs_stacked = solution.all_inputs_stacked
 
         self.mesh = base_variables[0].mesh
         self.domain = base_variables[0].domain
@@ -127,7 +127,7 @@ def _setup_inputs(self, t, full_range):
         ts = self.all_ts
         ys = self.all_ys
         yps = self.all_yps
-        inputs = self.all_inputs_casadi
+        inputs = self.all_inputs_stacked
 
         # Remove all empty ts
         idxs = np.where([ti.size > 0 for ti in ts])[0]
@@ -143,7 +143,7 @@ def _setup_inputs(self, t, full_range):
         ys = [ys[idx] for idx in idxs]
         if self.hermite_interpolation:
             yps = [yps[idx] for idx in idxs]
-        inputs = [self.all_inputs_casadi[idx] for idx in idxs]
+        inputs = [inputs[idx] for idx in idxs]
 
         is_f_contiguous = _is_f_contiguous(ys)
 
@@ -422,7 +422,7 @@ def initialise_sensitivity_explicit_forward(self):
         for ts, ys, inputs_stacked, inputs, base_variable, dy_dp in zip(
             self.all_ts,
             self.all_ys,
-            self.all_inputs_casadi,
+            self.all_inputs_stacked,
             self.all_inputs,
             self.base_variables,
             self.all_solution_sensitivities["all"],
@@ -508,19 +508,21 @@ def _stub_solution(self):
             """
 
             class StubSolution:
-                def __init__(self, ts, ys, inputs, inputs_casadi, sensitivities, t_pts):
+                def __init__(
+                    self, ts, ys, inputs, inputs_stacked, sensitivities, t_pts
+                ):
                     self.all_ts = ts
                     self.all_ys = ys
                     self.all_inputs = inputs
-                    self.all_inputs_casadi = inputs_casadi
+                    self.all_inputs_stacked = inputs_stacked
                     self.sensitivities = sensitivities
                     self.t = t_pts
 
             return StubSolution(
                 self.all_ts,
                 self.all_ys,
                 self.all_inputs,
-                self.all_inputs_casadi,
+                self.all_inputs_stacked,
                 self.sensitivities,
                 self.t_pts,
             )
@@ -568,7 +570,7 @@ def _observe_postfix(self, entries, t):
         if self.time_integral is None:
             return entries
         return self.time_integral.postfix(
-            entries, self.t_pts, self.all_inputs_casadi[0]
+            entries, self.t_pts, self.all_inputs_stacked[0]
         )
 
     def _interp_setup(self, entries, t):

src/pybamm/solvers/processed_variable_computed.py

@@ -63,7 +63,7 @@ def __init__(
         self.all_ts = solution.all_ts
         self.all_ys = solution.all_ys
         self.all_inputs = solution.all_inputs
-        self.all_inputs_casadi = solution.all_inputs_casadi
+        self.all_inputs_stacked = solution.all_inputs_stacked
 
         self.mesh = base_variables[0].mesh
         self.domain = base_variables[0].domain

src/pybamm/solvers/solution.py

@@ -155,6 +155,9 @@ def __init__(
         self.solve_time = None
         self.integration_time = None
 
+        self._all_inputs_stacked = None
+        self._all_inputs_casadi = None
+
         # initialize empty variable cache and data
         self._variables = {}
         self._data = pybamm.FuzzyDict()
@@ -337,9 +340,21 @@ def all_models(self):
         """Model(s) used for solution"""
         return self._all_models
 
-    @cached_property
-    def all_inputs_casadi(self):
-        return [casadi.vertcat(*inp.values()) for inp in self.all_inputs]
+    @property
+    def all_inputs_stacked(self) -> list[np.ndarray]:
+        if self._all_inputs_stacked is None:
+            self._all_inputs_stacked = [
+                np.asarray(list(inp.values())).reshape(-1) for inp in self.all_inputs
+            ]
+        return self._all_inputs_stacked
+
+    @property
+    def all_inputs_casadi(self) -> list[casadi.DM]:
+        if self._all_inputs_casadi is None:
+            self._all_inputs_casadi = [
+                casadi.vertcat(inp) for inp in self.all_inputs_stacked
+            ]
+        return self._all_inputs_casadi
 
     @property
     def all_yps(self) -> list[np.ndarray | casadi.DM | casadi.MX] | None:
@@ -421,6 +436,7 @@ def first_state(self):
             all_sensitivities=sensitivities,
             all_yps=all_yps,
         )
+        new_sol._all_inputs_stacked = self.all_inputs_stacked[:1]
         new_sol._all_inputs_casadi = self.all_inputs_casadi[:1]
         new_sol._sub_solutions = self.sub_solutions[:1]
 
@@ -464,6 +480,7 @@ def last_state(self):
             all_sensitivities=sensitivities,
             all_yps=all_yps,
         )
+        new_sol._all_inputs_stacked = self.all_inputs_stacked[-1:]
         new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:]
         new_sol._sub_solutions = self.sub_solutions[-1:]
         new_sol.solve_time = 0
@@ -621,20 +638,23 @@ def _convert_to_casadi(self, var_pybamm, inputs, ys_shape):
     def process_casadi_var(self, var_pybamm, inputs, ys_shape):
         t_MX = casadi.MX.sym("t")
         y_MX = casadi.MX.sym("y", ys_shape[0])
-        inputs_MX_dict = {
-            key: casadi.MX.sym("input", value.shape[0]) for key, value in inputs.items()
-        }
-        inputs_MX = casadi.vertcat(*[p for p in inputs_MX_dict.values()])
+        total_input_size = sum(v.size for v in inputs.values())
+        inputs_MX = casadi.MX.sym("p", total_input_size)
+        inputs_MX_dict = {}
+        offset = 0
+        for key, value in inputs.items():
+            n = value.size
+            inputs_MX_dict[key] = inputs_MX[offset : offset + n]
+            offset += n
         var_sym = var_pybamm.to_casadi(t_MX, y_MX, inputs=inputs_MX_dict)
 
         opts = {
             "cse": True,
             "inputs_check": False,
-            "is_diff_in": [False, False, False],
-            "is_diff_out": [False],
+            "is_diff_in": [False, True, False],
+            "is_diff_out": [True],
             "regularity_check": False,
             "error_on_fail": False,
-            "enable_jacobian": False,
         }
 
         # Casadi has a bug where it does not correctly handle arrays with
@@ -940,6 +960,7 @@ def __add__(self, other):
         )
 
         new_sol.closest_event_idx = other.closest_event_idx
+        new_sol._all_inputs_stacked = self.all_inputs_stacked + other.all_inputs_stacked
         new_sol._all_inputs_casadi = self.all_inputs_casadi + other.all_inputs_casadi
 
         # Add timers (if available)
@@ -978,6 +999,7 @@ def copy(self):
             all_t_evals=self._all_t_evals,
             variables_returned=self.variables_returned,
         )
+        new_sol._all_inputs_stacked = self.all_inputs_stacked
         new_sol._all_inputs_casadi = self.all_inputs_casadi
         new_sol._sub_solutions = self.sub_solutions
         new_sol.closest_event_idx = self.closest_event_idx
@@ -1105,6 +1127,7 @@ def make_cycle_solution(
     if sum_sols.variables_returned:
         cycle_solution._variables = sum_sols._variables
 
+    cycle_solution._all_inputs_stacked = sum_sols.all_inputs_stacked
     cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi
     cycle_solution._sub_solutions = sum_sols.sub_solutions
 

tests/unit/test_solvers/test_solution.py

@@ -280,6 +280,10 @@ def test_add_solutions(self):
             sol_sum.y, np.concatenate([y1, y2[:, 1:]], axis=1)
         )
         np.testing.assert_array_equal(sol_sum.all_inputs, [{"a": 1}, {"a": 2}])
+        assert sol_sum.all_inputs_stacked[0] is sol1.all_inputs_stacked[0]
+        assert sol_sum.all_inputs_stacked[1] is sol2.all_inputs_stacked[0]
+        assert sol_sum.all_inputs_casadi[0] is sol1.all_inputs_casadi[0]
+        assert sol_sum.all_inputs_casadi[1] is sol2.all_inputs_casadi[0]
 
         # Test sub-solutions
         assert len(sol_sum.sub_solutions) == 2
@@ -402,7 +406,8 @@ def test_copy(self):
         for ys_copy, ys1 in zip(sol_copy.all_ys, sol1.all_ys, strict=False):
             np.testing.assert_array_equal(ys_copy, ys1)
         assert sol_copy.all_inputs == sol1.all_inputs
-        assert sol_copy.all_inputs_casadi == sol1.all_inputs_casadi
+        assert sol_copy.all_inputs_stacked is sol1.all_inputs_stacked
+        assert sol_copy.all_inputs_casadi is sol1.all_inputs_casadi
         assert sol_copy.set_up_time == sol1.set_up_time
         assert sol_copy.solve_time == sol1.solve_time
         assert sol_copy.integration_time == sol1.integration_time
@@ -432,6 +437,26 @@ def test_copy_with_computed_variables(self):
         )
         assert sol2.variables_returned is True
 
+    def test_all_inputs(self):
+        t = [np.linspace(0, 1, 10), np.linspace(1, 2, 10)]
+        t[1][0] = np.nextafter(t[1][0], np.inf)
+        y = [np.tile(t[0], (5, 1)), np.tile(t[1], (5, 1))]
+        inputs = [{"a": 1.0, "b": 2.0, "c": 3.0}, {"a": 4.0, "b": 5.0, "c": 6.0}]
+        sol = pybamm.Solution(t, y, pybamm.BaseModel(), inputs)
+
+        stacked = sol.all_inputs_stacked
+        assert len(stacked) == 2
+        for s, inp in zip(stacked, inputs, strict=True):
+            assert isinstance(s, np.ndarray)
+            # check that it's a vector
+            assert s.shape == (len(inp),)
+            np.testing.assert_array_equal(s, np.array(list(inp.values())))
+
+        casadi_inputs = sol.all_inputs_casadi
+        assert len(casadi_inputs) == 2
+        for c, s in zip(casadi_inputs, stacked, strict=True):
+            np.testing.assert_array_equal(np.array(c).flatten(), s)
+
     def test_last_state(self):
         # Set up first solution
         t1 = [np.linspace(0, 1), np.linspace(1, 2, 5)]
@@ -447,6 +472,7 @@ def test_last_state(self):
         assert sol_last_state.all_ts[0] == 2
         np.testing.assert_array_equal(sol_last_state.all_ys[0], 2)
         assert sol_last_state.all_inputs == sol1.all_inputs[-1:]
+        assert sol_last_state.all_inputs_stacked == sol1.all_inputs_stacked[-1:]
         assert sol_last_state.all_inputs_casadi == sol1.all_inputs_casadi[-1:]
         assert sol_last_state.all_models == sol1.all_models[-1:]
         assert sol_last_state.set_up_time == 0