refactor: read solution variables through per-model observer caches

CHANGELOG.md

@@ -8,6 +8,7 @@
 ## Deprecated
 
 - Setting `Array.entries_string` or `Interpolant.entries_string`, and calling `StateVector.set_evaluation_array`, are in-place symbol updates and now emit `SymbolMutationDeprecationWarning`. Construct a new symbol instead. ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825))
+- `ProcessedVariable`'s `base_variables_casadi` argument and attribute are deprecated; pass the CasADi functions positionally or as `observer=`, and call the variable to evaluate it. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786))
 
 ## Features
 
@@ -46,6 +47,8 @@
 - `Solution` stores its attributes in `__slots__`, so an experiment, which keeps four solutions per step, retains about 40% less memory per step on Python 3.13. Solutions pickled by earlier versions still load. ([#5827](https://github.com/pybamm-team/PyBaMM/pull/5827))
 - Symbol mutation is no longer guarded by intercepting every attribute write, and `children` and `domain` return the stored immutable sequences instead of building a view on each access. Writes to a symbol's slots after construction are instead rejected by a static check of PyBaMM's code, tests, examples and documentation. Under the test suite's mutation guard, constructing an expression is twice as fast and parameterising it is 25% faster, which recovers most of the `test_parameterise` benchmark regression from [#5779](https://github.com/pybamm-team/PyBaMM/pull/5779). ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825))
 - Faster reads of `output_variables` results: an off-grid `solution["Voltage [V]"](t)` on an SPM drops from about 420 µs to 1 µs. ([#5785](https://github.com/pybamm-team/PyBaMM/pull/5785))
+- Repeat evaluations of a `ProcessedVariable` are faster: a repeat `solution["Voltage [V]"](t)` on an SPMe discharge takes about 40% less time. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786))
+- Time integrals are converted to CasADi once per model instead of once per solution segment. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786))
 
 # [v26.9.0.0](https://github.com/pybamm-team/PyBaMM/tree/pybamm-v26.9.0.0) - 2026-09-28
 

packages/pybamm/src/pybamm/models/base_model.py

@@ -126,7 +126,8 @@ def __init__(self, name="Unnamed model"):
         self._fixed_input_parameters = {}
         self._parameter_info = None
         self._is_standard_form_dae = None
-        self._variables_casadi = {}
+        # Built on first read by pybamm.solvers.observation.ObserverCache
+        self._observer_cache = None
         self._geometry = pybamm.Geometry({})
         self._symbol_processor = SymbolProcessor()
         self._solution_observable = ModelSolutionObservability.ENABLED
@@ -226,6 +227,8 @@ def __setstate__(self, state: dict) -> None:
         if "convert_to_format" in state:
             state = dict(state)
             state["_convert_to_format"] = state.pop("convert_to_format")
+        # Older pickles carry no observer cache
+        self._observer_cache = None
         self.__dict__.update(state)
 
     @property
@@ -1192,7 +1195,8 @@ def new_copy(self):
         new_model._variables = self.variables.copy()
         new_model._variables_processed = self._variables_processed.copy()
         new_model._events = self.events.copy()
-        new_model._variables_casadi = self._variables_casadi.copy()
+        if self._observer_cache is not None:
+            new_model._observer_cache = self._observer_cache.copy()
         new_model._symbol_processor = self.symbol_processor.copy()
         new_model._solution_observable = self._solution_observable
         new_model.switching_control_variables = self.switching_control_variables.copy()

packages/pybamm/src/pybamm/solvers/idaklu_solver.py

@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 import dataclasses
 import logging
-import math
 import numbers
 import warnings
 from enum import IntEnum
@@ -15,6 +14,7 @@
 import pybamm
 from pybamm.codegen.compilation import aot_compile
 from pybamm.solvers.base_solver import flatten_inputs, stack_inputs
+from pybamm.solvers.observation import OutputAssembly
 
 _UNSET = object()
 
@@ -675,6 +675,11 @@ def to_idaklu(fn):
             "number_of_sensitivity_parameters": number_of_sensitivity_parameters,
             "standard_form_dae": model.is_standard_form_dae,
             "output_variables": self.output_variables,
+            "output_assembly": OutputAssembly(
+                self.output_variables,
+                self.computed_var_fcns,
+                time_integrals=self._time_integral_vars,
+            ),
             "var_fcns": self.computed_var_fcns,
             "var_idaklu_fcns": [],
             "dvar_dy_idaklu_fcns": [],
@@ -968,85 +973,17 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict, t_ev
                 sol.yS_term, sensitivity_names
             )
 
-        # Populate variables and sensitivities dictionaries directly
-        number_of_samples = sol.y.shape[0] // number_of_timesteps
-        sol.y = sol.y.reshape((number_of_timesteps, number_of_samples))
-        sensitivity_params = (
-            model.calculate_sensitivities if model.calculate_sensitivities else []
+        # sol.y holds the outputs, and sol.yS their sensitivities (n_t, n_rows, n_p)
+        self._setup["output_assembly"].attach(
+            newsol,
+            np.asarray(sol.y).reshape(number_of_timesteps, -1),
+            sensitivities=(
+                np.asarray(sol.yS) if number_of_sensitivity_parameters else None
+            ),
+            sensitivity_names=sensitivity_names,
         )
-
-        start_idx = 0
-        for var in self.output_variables:
-            var_nnz, var_shape, base_variables = self._get_variable_info(model, var)
-            end_idx = start_idx + var_nnz
-            data = sol.y[:, start_idx:end_idx]
-            if var_nnz != math.prod(var_shape):
-                # Scatter the returned structural nonzeros to their flat indices
-                indices = base_variables[0].sparsity_out(0).find()
-                dense = np.zeros((number_of_timesteps, math.prod(var_shape)))
-                dense[:, indices] = data
-                data = dense
-            time_integral = self._time_integral_vars.get(var)
-            values = data
-
-            # handle any time integral variables
-            if time_integral is not None:
-                # time integral variables should all be 1D
-                values = time_integral.postfix(data.reshape(-1), sol.t, inputs_dict)
-
-            newsol._variables[var] = pybamm.ProcessedVariableComputed(
-                [model.get_processed_variable_or_event(var)],
-                base_variables,
-                [values],
-                newsol,
-                time_integral=time_integral,
-            )
-
-            # Add sensitivities
-            newsol[var]._sensitivities = {}
-            if sensitivity_params:
-                if var_nnz != math.prod(var_shape):
-                    raise pybamm.SolverError(
-                        f"Sensitivity of sparse variables not supported. {var} is a sparse variable with number of non-zeros {var_nnz} and shape {var_shape}"
-                    )
-                sens_data = sol.yS[:, start_idx:end_idx, :]
-                sens_data = sens_data.reshape(
-                    number_of_timesteps * (end_idx - start_idx),
-                    number_of_sensitivity_parameters,
-                )
-                if time_integral is not None:
-                    sens_data = time_integral.postfix_sensitivities(
-                        var,
-                        data.reshape(-1),
-                        sol.t,
-                        inputs_dict,
-                        sensitivity_names,
-                        sens_data,
-                    )
-                newsol[var]._sensitivities["all"] = sens_data
-
-                # Add the individual sensitivity
-                for i, name in enumerate(sensitivity_names):
-                    sens = newsol[var]._sensitivities["all"][:, i : i + 1].reshape(-1)
-                    newsol[var]._sensitivities[name] = sens
-
-            start_idx += var_nnz
         return newsol
 
-    def _get_variable_info(self, model, var) -> tuple:
-        """Get variable length and base variables based on model format."""
-        if model.convert_to_format == "casadi":
-            base_var = self._setup["var_fcns"][var]
-            sparsity = base_var.sparsity_out(0)
-            var_nnz = sparsity.nnz()
-            var_shape = sparsity.shape
-            return var_nnz, var_shape, [base_var]
-        else:  # pragma: no cover
-            raise pybamm.SolverError(
-                f"Unsupported evaluation engine for convert_to_format="
-                f"{model.convert_to_format}"
-            )
-
     def _set_consistent_initialization(self, model, time, inputs_list):
         """
         Initialize y0 and ydot0 for the solver. In addition to calculating