From d44963608e6b9e46c00031a7df273ad0cb880e49 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Wed, 23 Sep 2026 19:29:49 +0100 Subject: [PATCH 1/7] refactor: put solution observation behind an ObservationBackend seam A Solution reads variables through one ObservationBackend and a ProcessedVariable through one VariableObserver (CasADi only); OutputAssembly owns the output_variables layout and keys its sensitivities by sensitivity parameter. Solution carries its backend in an _observation slot, which a solution pickled before it existed loads as CASADI_OBSERVATION. --- CHANGELOG.md | 2 + .../src/pybamm/solvers/idaklu_solver.py | 92 +--- .../pybamm/src/pybamm/solvers/observation.py | 356 ++++++++++++++++ .../src/pybamm/solvers/processed_variable.py | 346 +++++---------- .../pybamm/src/pybamm/solvers/solution.py | 84 +--- .../src/pybamm/solvers/variable_observer.py | 400 ++++++++++++++++++ .../test_legacy_mutation.py | 1 - .../unit/test_solvers/test_idaklu_solver.py | 37 ++ .../test_solvers/test_observation_backend.py | 228 ++++++++++ .../test_solvers/test_observation_layout.py | 159 +++++++ .../test_solvers/test_processed_variable.py | 41 ++ .../tests/unit/test_solvers/test_solution.py | 110 +++++ 12 files changed, 1468 insertions(+), 388 deletions(-) create mode 100644 packages/pybamm/src/pybamm/solvers/observation.py create mode 100644 packages/pybamm/src/pybamm/solvers/variable_observer.py create mode 100644 packages/pybamm/tests/unit/test_solvers/test_observation_backend.py create mode 100644 packages/pybamm/tests/unit/test_solvers/test_observation_layout.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 1949337530..20836598b1 100644 --- a/CHANGELOG.md +++ b/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)) +- The `base_variables_casadi` argument and attribute of `ProcessedVariable` are deprecated. Pass the CasADi functions as the third positional argument or as `observer=`, and call the variable to evaluate it rather than reading the functions back. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) ## Features @@ -16,6 +17,7 @@ - A multi-input `IDAKLUSolver` solve now hands each thread its next input set as soon as it is free, instead of splitting the sets into equal static blocks plus a serial remainder solved on the calling thread, so a heterogeneous sweep (for example a current sweep whose high currents stop early on a voltage cut-off) no longer waits on its slowest block. A failing set no longer hides the others: the `SolverError` names every set that failed, whether it failed from the start or part-way through, each prefixed `input set N: `, and a single-set solve's message now starts with `input set 0: `. Each `on_failure="warn"` warning names its set the same way. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) - Added `Solution.sensitivity_names`, the sensitivity inputs in the column order of `sensitivities["all"]`. ([#5785](https://github.com/pybamm-team/PyBaMM/pull/5785)) - An outputs-only `pybammsolvers` solution also returns `y_init`, `yS_init` and `yS_term`: its initial states and its states' sensitivities at both ends. ([#5794](https://github.com/pybamm-team/PyBaMM/pull/5794)) +- Evaluating a `ProcessedVariable` again, for example at new times, no longer re-serialises its CasADi functions or rescans every sub-solution's time span: a repeat `solution["Voltage [V]"](t)` on an SPMe discharge is about 40% faster, and finding the sub-solutions that cover the requested times, across 60 of them, takes 1 µs rather than 10 µs. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) ## Bug fixes diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index 50994a3af8..f32a73c166 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/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,18 @@ 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 [] + # On this path `sol.y` carries the concatenated outputs rather than the + # states, and `sol.yS` their sensitivities as (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 diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py new file mode 100644 index 0000000000..e4b81da87d --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -0,0 +1,356 @@ +"""Observation backends: how a :class:`pybamm.Solution` reads its variables. + +An :class:`ObservationBackend` turns a variable name into a processed variable +over a run of a Solution's sub-solutions, its *segments*. A *leaf* is one +segment's compiled form of one variable, so a variable spanning ``n`` +sub-solutions has ``n`` leaves. + +:class:`OutputAssembly` is the ``output_variables`` counterpart: the solver has +already computed those variables, so it owns the payload's row layout and +populates the Solution eagerly. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from itertools import accumulate, pairwise + +import casadi +import numpy as np +import numpy.typing as npt + +import pybamm +from pybamm.solvers.base_processed_variable import BaseProcessedVariable +from pybamm.solvers.variable_observer import ( + CasadiObserver, + check_variable_in_solve, + pack_sensitivity_dict, +) + + +class ObservationBackend(ABC): + """How a Solution turns a variable name into something it can evaluate. + + A backend covers an ordered run of a Solution's sub-solutions, one + *segment* each. ``backend[key]`` indexes *segments*, matching the slice the + Solution applies to its sub-solutions; :func:`join_observations` combines + the backends of consecutive Solutions. + """ + + @abstractmethod + def __getitem__(self, key: slice) -> ObservationBackend: + """This backend restricted to a slice of the Solution's segments.""" + + @abstractmethod + def build_variable( + self, solution: pybamm.Solution, name: str + ) -> BaseProcessedVariable: + """The processed variable for ``name``, ready to evaluate. + + Parameters + ---------- + solution : :class:`pybamm.Solution` + The solution being observed; supplies the trajectories, models and + inputs. Its segments are 1:1 with this backend's. + name : str + Variable name, as registered on the models. + + Returns + ------- + :class:`pybamm.solvers.base_processed_variable.BaseProcessedVariable` + Evaluable over this backend's whole run of segments. + """ + + +class CasadiObservation(ObservationBackend): + """Variables converted to CasADi and evaluated by the IDAKLU kernels. + + Stateless: every per-segment artifact it needs is reached through the + Solution it is handed, so any two instances are interchangeable and one, + :data:`CASADI_OBSERVATION`, covers any run of segments. + """ + + def __getitem__(self, key): + return self + + def __eq__(self, other): + return type(other) is type(self) + + def __hash__(self): + return hash(type(self)) + + def build_variable(self, solution, name): + time_integral = None + pybamm.logger.debug(f"Post-processing {name}") + + # Iterate through all models, some may be in the list several times and + # therefore only get set up once + vars_pybamm = [ + model.get_processed_variable_or_event(name) for model in solution.all_models + ] + vars_casadi = [None] * len(solution.all_models) + for i, (model, ys, inputs) in enumerate( + zip(solution.all_models, solution.all_ys, solution.all_inputs, strict=True) + ): + _var_pybamm = vars_pybamm[i] + check_variable_in_solve(solution, name, _var_pybamm) + if isinstance(_var_pybamm, pybamm.VectorField): + comp_casadi = [] + for k, comp in enumerate(_var_pybamm.components): + cc, _, _ = self._model_leaf( + solution, + model, + comp, + inputs=inputs, + ys_shape=ys.shape, + time_integral=None, + cache_key=f"{name}[{k}]", + ) + comp_casadi.append(cc) + vars_casadi[i] = comp_casadi + else: + var_casadi, var_pybamm, time_integral = self._model_leaf( + solution, + model, + _var_pybamm, + inputs=inputs, + ys_shape=ys.shape, + time_integral=time_integral, + cache_key=name, + ) + vars_pybamm[i] = var_pybamm + vars_casadi[i] = var_casadi + return pybamm.process_variable( + name, + vars_pybamm, + CasadiObserver(vars_casadi), + solution, + time_integral=time_integral, + ) + + @staticmethod + def _model_leaf( + solution, + model, + var_pybamm, + time_integral, + inputs, + ys_shape, + cache_key, + ): + """One model's CasADi leaf, memoised on the model unless time-integrated.""" + _var_casadi = model._variables_casadi.get(cache_key) + if _var_casadi is not None: + return _var_casadi, var_pybamm, time_integral + + var_casadi, var_pybamm, time_integral = solution._convert_to_casadi( + var_pybamm, inputs, ys_shape + ) + + # Only cache if it's not a time integral + if time_integral is None: + model._variables_casadi[cache_key] = var_casadi + return var_casadi, var_pybamm, time_integral + + +class OutputAssembly: + """The row layout of an ``output_variables`` payload, and its attachment. + + A solve run with ``output_variables`` propagates no state trajectory: the + solver evaluates the requested variables itself and returns one row per + non-zero of each variable's CasADi function, in variable order, so a vector + variable spans several consecutive rows. + + Parameters + ---------- + names : list of str + The output variables, in row order. + casadi_fns : dict + Map of name to the CasADi function ``f(t, y, p)`` the rows were + evaluated by. Its non-zeros set the rows each variable owns. + time_integrals : dict, optional + Map of name to :class:`pybamm.ProcessedVariableTimeIntegral` for the + outputs whose rows carry an integrand rather than the variable itself. + """ + + def __init__( + self, + names: Sequence[str], + casadi_fns: Mapping[str, casadi.Function], + *, + time_integrals: Mapping[str, pybamm.ProcessedVariableTimeIntegral] + | None = None, + ): + self._names = tuple(names) + self._casadi_fns = {name: casadi_fns[name] for name in self._names} + self._time_integrals = dict(time_integrals or {}) + lens = [] + # name -> sparsity of the variables whose rows are sparse + self._sparse = {} + for name in self._names: + sparsity = self._casadi_fns[name].sparsity_out(0) + lens.append(sparsity.nnz()) + if sparsity.nnz() != sparsity.numel(): + self._sparse[name] = sparsity + offsets = list(accumulate(lens, initial=0)) + self._n_rows = offsets[-1] + # The layout itself: one slice of payload rows per output variable. + self._rows = tuple(slice(start, end) for start, end in pairwise(offsets)) + + @property + def n_rows(self) -> int: + """Rows in one time point of the payload.""" + return self._n_rows + + def attach( + self, + solution: pybamm.Solution, + data: npt.ArrayLike, + *, + sensitivities: npt.ArrayLike | None = None, + sensitivity_names: Sequence[str] = (), + ) -> None: + """Populate ``solution``'s variables from one outputs-only payload. + + Parameters + ---------- + solution : :class:`pybamm.Solution` + The solution to populate, built with ``variables_returned=True``. + data : array-like + Output trajectory of shape ``(n_t, n_rows)``, time-outer. + sensitivities : array-like, optional + Output sensitivities of shape ``(n_t, n_rows, n_p)``. Omit when the + solve carried none, which leaves every variable's sensitivities empty + rather than lazily recomputed: an outputs-only solve keeps no state + to recompute them from. + sensitivity_names : list of str, optional + Sensitivity-parameter names, in ``sensitivities``' column order. + + Raises + ------ + :class:`pybamm.SolverError` + If the payload does not match this layout, or if sensitivities were + requested for a variable whose CasADi rows are sparse. + """ + data = self._checked_rows(data) + if sensitivities is not None: + sensitivities = self._checked_sensitivities( + sensitivities, data.shape[0], sensitivity_names + ) + model = solution.all_models[0] + for name, rows in zip(self._names, self._rows, strict=True): + time_integral = self._time_integrals.get(name) + entries = np.ascontiguousarray(data[:, rows]) + sparsity = self._sparse.get(name) + if sparsity is not None: + # The rows hold only the structural nonzeros; scatter them into + # every entry at their flat (column-major) indices + dense = np.zeros((entries.shape[0], sparsity.numel())) + dense[:, sparsity.find()] = entries + entries = dense + values = entries + if time_integral is not None: + # These rows are the integrand's trajectory, not the variable's. + values = time_integral.postfix( + entries.reshape(-1), solution.t, solution.all_inputs[0] + ) + variable = pybamm.ProcessedVariableComputed( + [model.get_processed_variable_or_event(name)], + [self._casadi_fns[name]], + [values], + solution, + time_integral=time_integral, + ) + variable._sensitivities = ( + {} + if sensitivities is None + else self._variable_sensitivities( + name, + entries, + sensitivities[:, rows, :], + solution, + sensitivity_names, + ) + ) + solution._variables[name] = variable + + def _checked_rows(self, data): + """``data`` as a ``(n_t, n_rows)`` array, or a complaint about its width.""" + array = np.asarray(data) + if array.ndim != 2 or array.shape[1] != self.n_rows: + raise pybamm.SolverError( + f"Output row count mismatch: expected {self.n_rows} rows (the total " + f"flattened length of {len(self._names)} output variables) but the " + f"solver returned an array of shape {array.shape}." + ) + return array + + def _checked_sensitivities(self, sensitivities, n_timesteps, sensitivity_names): + """``sensitivities`` as a ``(n_t, n_rows, n_p)`` array, or a complaint.""" + array = np.asarray(sensitivities) + expected = (n_timesteps, self.n_rows, len(sensitivity_names)) + if array.shape != expected: + raise pybamm.SolverError( + f"Output sensitivity shape mismatch: expected {expected} (times, " + f"flattened outputs, parameters) but the solver returned " + f"{array.shape}." + ) + return array + + def _variable_sensitivities( + self, name, entries, block, solution, sensitivity_names + ): + """One variable's ``"all"`` block plus a flat vector per parameter.""" + sparsity = self._sparse.get(name) + if sparsity is not None: + raise pybamm.SolverError( + f"Sensitivity of sparse variables not supported. {name} is a sparse " + f"variable with number of non-zeros {sparsity.nnz()} and shape " + f"{sparsity.shape}" + ) + n_timesteps, var_len, n_params = block.shape + all_sens = block.reshape(n_timesteps * var_len, n_params) + time_integral = self._time_integrals.get(name) + if time_integral is not None: + all_sens = time_integral.postfix_sensitivities( + name, + entries.reshape(-1), + solution.t, + solution.all_inputs[0], + sensitivity_names, + all_sens, + ) + return pack_sensitivity_dict(all_sens, sensitivity_names) + + +def join_observations(backends: Sequence[ObservationBackend]) -> ObservationBackend: + """One backend covering consecutive Solutions' segments, in order. + + Parameters + ---------- + backends : list of :class:`ObservationBackend` + The backend of each Solution being joined, in order. + + Returns + ------- + :class:`ObservationBackend` + The backend every Solution shares. + + Raises + ------ + :class:`pybamm.SolverError` + If the Solutions do not all share one backend. + """ + first, *rest = backends + for backend in rest: + if backend != first: + raise pybamm.SolverError( + "Cannot join solutions read through different observation " + f"backends: {type(first).__name__} and {type(backend).__name__}." + ) + return first + + +#: The backend a Solution carries by default. +CASADI_OBSERVATION = CasadiObservation() diff --git a/packages/pybamm/src/pybamm/solvers/processed_variable.py b/packages/pybamm/src/pybamm/solvers/processed_variable.py index 55ce35adde..2c00a7d6f1 100644 --- a/packages/pybamm/src/pybamm/solvers/processed_variable.py +++ b/packages/pybamm/src/pybamm/solvers/processed_variable.py @@ -1,15 +1,39 @@ -import bisect +import functools +import warnings -import casadi import numpy as np import xarray as xr -from pybammsolvers import idaklu import pybamm +from pybamm.solvers.variable_observer import as_observer from .base_processed_variable import BaseProcessedVariable +def _deprecated_leaves_keyword(init): + """Accept the deprecated ``base_variables_casadi`` keyword as ``observer``.""" + + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + if "base_variables_casadi" in kwargs: + if "observer" in kwargs or len(args) > 2: + raise TypeError( + "Pass the CasADi functions once, as `observer` (not also as " + "`base_variables_casadi`)." + ) + warnings.warn( + "The `base_variables_casadi` argument is deprecated and will be " + "removed in a future release; pass the CasADi functions as " + "`observer` instead.", + DeprecationWarning, + stacklevel=2, + ) + kwargs["observer"] = kwargs.pop("base_variables_casadi") + return init(self, *args, **kwargs) + + return wrapper + + class ProcessedVariable(BaseProcessedVariable): """ An object that can be evaluated at arbitrary (scalars or vectors) t and x, and @@ -27,26 +51,29 @@ class ProcessedVariable(BaseProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves. A bare list of + :class:`casadi.Function` is accepted and wrapped in a + :class:`CasadiObserver`, one function per sub-solution. Passing it as + ``base_variables_casadi`` is deprecated. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables time_integral : :class:`pybamm.ProcessedVariableTimeIntegral`, optional Not none if the variable is to be time-integrated (default is None) """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): self._name = name self.base_variables = base_variables - self.base_variables_casadi = base_variables_casadi + self._observer = as_observer(observer) self.all_ts = solution.all_ts self.all_ys = solution.all_ys @@ -85,6 +112,25 @@ def __init__( self._entries_for_interp_raw = None self._coords_raw = None + def __setstate__(self, state): + state = dict(state) + # Pickled before observers existed: the leaves are a bare attribute. + leaves = state.pop("base_variables_casadi", None) + if "_observer" not in state: + state["_observer"] = as_observer(leaves) + self.__dict__.update(state) + + @property + def base_variables_casadi(self): + """The per-sub-solution CasADi functions (deprecated).""" + warnings.warn( + "`ProcessedVariable.base_variables_casadi` is deprecated and will be " + "removed in a future release; evaluate the variable by calling it.", + DeprecationWarning, + stacklevel=2, + ) + return self._observer.leaves + def initialise(self): if self.entries_raw_initialized: return @@ -122,67 +168,11 @@ def observe_raw(self): t = self.t_pts return self._observe_postfix(self._observe_raw(), t) - def _setup_inputs(self, t, full_range): - pybamm.logger.debug("Setting up C++ interpolation inputs") - - ts = self.all_ts - ys = self.all_ys - yps = self.all_yps - inputs = self.all_inputs_stacked - - # Remove all empty ts - idxs = np.where([ti.size > 0 for ti in ts])[0] - - # Find the indices of the time points to observe - if not full_range: - ts_nonempty = [ts[idx] for idx in idxs] - idxs_subset = _find_ts_indices(ts_nonempty, t) - idxs = idxs[idxs_subset] - - # Extract the time points and inputs - ts = [ts[idx] for idx in idxs] - ys = [ys[idx] for idx in idxs] - if self.hermite_interpolation: - yps = [yps[idx] for idx in idxs] - inputs = [inputs[idx] for idx in idxs] - - is_f_contiguous = _is_f_contiguous(ys) - - ts = idaklu.VectorRealtypeNdArray(ts) - ys = idaklu.VectorRealtypeNdArray(ys) - if self.hermite_interpolation: - yps = idaklu.VectorRealtypeNdArray(yps) - else: - yps = None - inputs = idaklu.VectorRealtypeNdArray(inputs) - - # Generate the serialized C++ functions only once - funcs_unique = {} - funcs = [None] * len(idxs) - for i in range(len(idxs)): - vars = self.base_variables_casadi[idxs[i]] - if vars not in funcs_unique: - funcs_unique[vars] = vars.serialize() - funcs[i] = funcs_unique[vars] - - return ts, ys, yps, funcs, inputs, is_f_contiguous - - def _observe_hermite(self, t): - pybamm.logger.debug("Observing and Hermite interpolating the variable") - - ts, ys, yps, funcs, inputs, _ = self._setup_inputs(t, full_range=False) - shapes = self._shape(t) - return idaklu.observe_hermite_interp(t, ts, ys, yps, inputs, funcs, shapes) - def _observe_raw(self): - pybamm.logger.debug("Observing the variable raw data") - t = self.t_pts - ts, ys, _, funcs, inputs, is_f_contiguous = self._setup_inputs( - t, full_range=True - ) - shapes = self._shape(self.t_pts) + return self._observer.observe_raw(self) - return idaklu.observe(ts, ys, inputs, funcs, is_f_contiguous, shapes) + def _observe_hermite(self, t): + return self._observer.observe_hermite(self, t) def _observe_postfix(self, entries, t): return entries @@ -419,88 +409,7 @@ def sensitivities(self): def initialise_sensitivity_explicit_forward(self): "Set up the sensitivity dictionary" - - all_S_var = [] - for ts, ys, inputs, base_variable, dy_dp in zip( - self.all_ts, - self.all_ys, - self.all_inputs, - self.base_variables, - self.all_solution_sensitivities["all"], - strict=True, - ): - sensitivity_inputs = { - name: inputs[name] for name in self.sensitivity_names if name in inputs - } - sensitivity_inputs_stacked = casadi.vertcat( - *[sensitivity_inputs[name] for name in self.sensitivity_names] - ) - - # Set up symbolic variables - t_casadi = casadi.MX.sym("t") - y_casadi = casadi.MX.sym("y", ys.shape[0]) - p_casadi = { - name: casadi.MX.sym(name, value.shape[0]) - for name, value in sensitivity_inputs.items() - } - - p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) - - # Symbolic for sensitivity targets, concrete for the rest. Non-target - # inputs may still appear in the expression tree (e.g. from - # experiment steps) so they must be present for casadi conversion. - inputs_for_casadi = {**inputs, **p_casadi} - - var_casadi = base_variable.to_casadi( - t_casadi, y_casadi, inputs=inputs_for_casadi - ) - dvar_dy = casadi.jacobian(var_casadi, y_casadi) - dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) - - # Convert to functions and evaluate index-by-index - dvar_dy_func = casadi.Function( - "dvar_dy", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dy] - ) - dvar_dp_func = casadi.Function( - "dvar_dp", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dp] - ) - dvar_dy_eval = casadi.diagcat( - *[ - dvar_dy_func(t, ys[:, idx], sensitivity_inputs_stacked) - for idx, t in enumerate(ts) - ] - ) - dvar_dp_eval = casadi.vertcat( - *[ - dvar_dp_func(t, ys[:, idx], sensitivity_inputs_stacked) - for idx, t in enumerate(ts) - ] - ) - - # Compute sensitivity - S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval - - all_S_var.append(S_var) - - S_var = np.vstack(all_S_var) - if self.time_integral is not None: - # The value sums over every segment at once, so its sensitivity must too - S_var = self.time_integral.postfix_sensitivities( - self._name, - self._observe_raw(), - self.t_pts, - self.all_inputs[0], - self.sensitivity_names, - S_var, - ) - sensitivities = {"all": S_var} - - # Add the individual sensitivity - for i, name in enumerate(self.sensitivity_names): - sensitivities[name] = S_var[:, i : i + 1].reshape(-1) - - # Save attribute - self._sensitivities = sensitivities + self._sensitivities = self._observer.sensitivities(self) def _is_discrete_time_method(self): """Check if using discrete time integral method""" @@ -564,7 +473,7 @@ def __init__( cpv = pybamm.ProcessedVariableComputed( self.base_variables, - self.base_variables_casadi, + self._observer.leaves, base_data, _stub_solution(self), time_integral=self.time_integral, @@ -578,11 +487,12 @@ def __init__( class ProcessedVariable0D(ProcessedVariable): + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -590,7 +500,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -630,18 +540,19 @@ class ProcessedVariable1D(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -649,7 +560,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -715,18 +626,19 @@ class ProcessedVariable2D(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -734,7 +646,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -861,18 +773,19 @@ class ProcessedVariable2DSciKitFEM(ProcessedVariable2D): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -880,7 +793,7 @@ def __init__( super(ProcessedVariable2D, self).__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -914,11 +827,12 @@ def _observe_postfix(self, entries, t): class ProcessedVariable2DFVM(ProcessedVariable): + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -926,7 +840,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -987,11 +901,12 @@ class ProcessedVariableUnstructuredFVM(ProcessedVariable): on cell centroids; query it at arbitrary points with ``pv(t, x=, y=, z=)``. """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -1013,7 +928,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -1171,11 +1086,12 @@ class ProcessedVariableVectorFieldUnstructuredFVM: data. """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral=None, ): @@ -1189,10 +1105,11 @@ def __init__( self.internal_boundaries = [] self._component_vars = [] + leaves = as_observer(observer).leaves for k in range(vf.n_components): comp_base_k = [bv.components[k] for bv in base_variables] comp_casadi_k = [] - for bvc in base_variables_casadi: + for bvc in leaves: if isinstance(bvc, list): comp_casadi_k.append(bvc[k]) elif isinstance(bvc, pybamm.VectorField): @@ -1265,18 +1182,19 @@ class ProcessedVariable3D(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -1284,7 +1202,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -1460,18 +1378,19 @@ class ProcessedVariable3DSciKitFEM(ProcessedVariable3D): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -1479,7 +1398,7 @@ def __init__( super(ProcessedVariable3D, self).__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -1551,9 +1470,9 @@ class ProcessedVariableUnstructured(ProcessedVariable): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variables_casadi : list of :class:`casadi.Function` - A list of casadi functions. When evaluated, returns the same thing as - `base_Variable.evaluate` (but more efficiently). + observer : :class:`pybamm.solvers.variable_observer.VariableObserver` + Evaluates the variable's per-sub-solution leaves; a list of + :class:`casadi.Function` is also accepted. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables time_integral : pybamm.ProcessedVariableTimeIntegral, optional @@ -1561,11 +1480,12 @@ class ProcessedVariableUnstructured(ProcessedVariable): If provided, the processed variable will handle time integration using this object. """ + @_deprecated_leaves_keyword def __init__( self, name: str, base_variables, - base_variables_casadi, + observer, solution, time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, ): @@ -1573,7 +1493,7 @@ def __init__( super().__init__( name, base_variables, - base_variables_casadi, + observer, solution, time_integral=time_integral, ) @@ -1719,20 +1639,6 @@ def process_variable(name: str, base_variables, *args, **kwargs): raise NotImplementedError(f"Shape not recognized for {base_variables[0]}") -def _is_f_contiguous(all_ys): - """ - Check if all the ys are f-contiguous in memory - - Args: - all_ys (list of np.ndarray): list of all ys - - Returns: - bool: True if all ys are f-contiguous - """ - - return all(isinstance(y, np.ndarray) and y.data.f_contiguous for y in all_ys) - - def _is_sorted(t): """ Check if an array is sorted @@ -1744,39 +1650,3 @@ def _is_sorted(t): bool: True if array is sorted """ return np.all(t[:-1] <= t[1:]) - - -def _find_ts_indices(ts, t): - """ - Parameters: - - ts: A list of numpy arrays (each sorted) whose values are successively increasing. - - t: A sorted list or array of values to find within ts. - - Returns: - - indices: A list of indices from `ts` such that at least one value of `t` falls within ts[idx]. - """ - - indices = [] - - # Get the minimum and maximum values of the target values `t` - t_min, t_max = t[0], t[-1] - - # Step 1: Use binary search to find the range of `ts` arrays where t_min and t_max could lie - low_idx = bisect.bisect_left([ts_arr[-1] for ts_arr in ts], t_min) - high_idx = bisect.bisect_right([ts_arr[0] for ts_arr in ts], t_max) - - # Step 2: Iterate over the identified range - for idx in range(low_idx, high_idx): - ts_min, ts_max = ts[idx][0], ts[idx][-1] - - # Binary search within `t` to check if any value falls within [ts_min, ts_max] - i = bisect.bisect_left(t, ts_min) - if i < len(t) and t[i] <= ts_max: - # At least one value of t is within ts[idx] - indices.append(idx) - - # extrapolating - if (t[-1] > ts[-1][-1]) and (len(indices) == 0 or indices[-1] != len(ts) - 1): - indices.append(len(ts) - 1) - - return indices diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index 56ac46235c..a71ab570eb 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -17,6 +17,7 @@ import pybamm from pybamm.codegen.compilation import aot_compile from pybamm.solvers.base_solver import flatten_inputs +from pybamm.solvers.observation import CASADI_OBSERVATION, join_observations class NumpyEncoder(json.JSONEncoder): @@ -344,6 +345,7 @@ class Solution(SolutionBase): "_initial_start_time", "_last_state", "_observable", + "_observation", "_options", "_sensitivities", "_sub_solutions", @@ -462,6 +464,7 @@ def __init__( self._all_inputs_casadi = None self._variables = {} + self._observation = CASADI_OBSERVATION # Sub-solutions concatenated into this one. Empty list means "just # self"; storing self here would create a refcount cycle that the @@ -503,6 +506,7 @@ def __setstate__(self, state: dict | tuple) -> None: self._y0 = None self._y0_sensitivities = None self._y_event_sensitivities = None + self._observation = CASADI_OBSERVATION if isinstance(state, dict): state = { f"_{name}" if name in ("first_state", "last_state") else name: value @@ -787,6 +791,7 @@ def _build_first_state(self): ) # stacked/casadi stay lazy; built from all_inputs[:1] on first access new_sol._sub_solutions = self.sub_solutions[:1] + new_sol._observation = self._observation[:1] new_sol.solve_time = 0 new_sol.integration_time = 0 @@ -841,6 +846,7 @@ def _build_last_state(self): ) # stacked/casadi stay lazy; built from all_inputs[-1:] on first access new_sol._sub_solutions = self.sub_solutions[-1:] + new_sol._observation = self._observation[-1:] new_sol.solve_time = 0 new_sol.integration_time = 0 new_sol.solver_statistics = _NO_SOLVER_STATISTICS @@ -885,79 +891,8 @@ def update(self, variables: str | list[str]): for variable in variables: self._update_variable(variable) - def _update_model_variable( - self, - model: pybamm.BaseModel, - var_pybamm: pybamm.Symbol, - time_integral: pybamm.ProcessedVariableTimeIntegral | None, - inputs: dict, - ys_shape: tuple, - cache_key, - ): - _var_casadi = model._variables_casadi.get(cache_key) - if _var_casadi is not None: - return _var_casadi, var_pybamm, time_integral - - var_casadi, var_pybamm, time_integral = self._convert_to_casadi( - var_pybamm, inputs, ys_shape - ) - - # Only cache if it's not a time integral - if time_integral is None: - model._variables_casadi[cache_key] = var_casadi - return var_casadi, var_pybamm, time_integral - def _update_variable(self, name: str): - time_integral = None - pybamm.logger.debug(f"Post-processing {name}") - - # Iterate through all models, some may be in the list several times and - # therefore only get set up once - vars_pybamm = [ - model.get_processed_variable_or_event(name) for model in self.all_models - ] - vars_casadi = [None] * len(self.all_models) - for i, (model, ys, inputs) in enumerate( - zip(self.all_models, self.all_ys, self.all_inputs, strict=True) - ): - _var_pybamm = vars_pybamm[i] - if self.variables_returned and _var_pybamm.has_symbol_of_classes( - pybamm.expression_tree.state_vector.StateVector - ): - raise KeyError( - f"Cannot process variable '{name}' as it was not part of the " - "solve. Please re-run the solve with `output_variables` set to " - "include this variable." - ) - if isinstance(_var_pybamm, pybamm.VectorField): - comp_casadi = [] - for k, comp in enumerate(_var_pybamm.components): - cc, _, _ = self._update_model_variable( - model, - comp, - inputs=inputs, - ys_shape=ys.shape, - time_integral=None, - cache_key=f"{name}[{k}]", - ) - comp_casadi.append(cc) - vars_casadi[i] = comp_casadi - else: - var_casadi, var_pybamm, time_integral = self._update_model_variable( - model, - _var_pybamm, - inputs=inputs, - ys_shape=ys.shape, - time_integral=time_integral, - cache_key=name, - ) - vars_pybamm[i] = var_pybamm - vars_casadi[i] = var_casadi - var = pybamm.process_variable( - name, vars_pybamm, vars_casadi, self, time_integral=time_integral - ) - - self._variables[name] = var + self._variables[name] = self._observation.build_variable(self, name) def observe(self, symbol: pybamm.Symbol) -> pybamm.ProcessedVariable: """ @@ -1385,6 +1320,9 @@ def __add__(self, other): # Set sub_solutions new_sol._sub_solutions = self.sub_solutions + other.sub_solutions + new_sol._observation = join_observations( + [self._observation, other._observation] + ) # update variables which were derived at the solver stage if any([self.variables_returned, other.variables_returned]): @@ -1488,6 +1426,7 @@ def from_sub_solutions(cls, sub_solutions): new_sol._y0_sensitivities = segments[0]._y0_sensitivities # leave stacked/casadi unset; built lazily from all_inputs (casadi is costly) new_sol._sub_solutions = sub_sols + new_sol._observation = join_observations([s._observation for s in segments]) for attr in ["solve_time", "integration_time", "set_up_time"]: vals = [getattr(s, attr, None) for s in segments] @@ -1538,6 +1477,7 @@ def copy(self): new_sol.integration_time = self.integration_time new_sol.solver_statistics = self.solver_statistics new_sol.set_up_time = self.set_up_time + new_sol._observation = self._observation # copy over variables which were derived at the solver stage if self._variables and all( diff --git a/packages/pybamm/src/pybamm/solvers/variable_observer.py b/packages/pybamm/src/pybamm/solvers/variable_observer.py new file mode 100644 index 0000000000..dcf8530e70 --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/variable_observer.py @@ -0,0 +1,400 @@ +"""Per-variable observation strategies for :class:`pybamm.ProcessedVariable`. + +A :class:`VariableObserver` owns one variable's per-sub-solution leaves and +evaluates them on the solution grid, off-grid via cubic Hermite, and through +the forward chain rule for sensitivities. +""" + +from __future__ import annotations + +import bisect +from abc import ABC, abstractmethod +from collections.abc import Sequence + +import casadi +import numpy as np +import numpy.typing as npt +from pybammsolvers import idaklu + +import pybamm + + +class SegmentSelector: + """Which sub-solutions cover a set of query times. + + The non-empty segments and their end times are found once, on + construction, rather than on every query. + + Parameters + ---------- + all_ts : list of numpy.ndarray + Per-sub-solution time arrays, successively increasing. + """ + + def __init__(self, all_ts: Sequence[npt.NDArray[np.float64]]): + self.indices = np.where([ti.size > 0 for ti in all_ts])[0] + self._starts = [all_ts[idx][0] for idx in self.indices] + self._ends = [all_ts[idx][-1] for idx in self.indices] + + def select( + self, t: npt.NDArray[np.float64], full_range: bool + ) -> npt.NDArray[np.intp]: + """Indices into ``all_ts`` of the segments covering ``t``. + + Parameters + ---------- + t : numpy.ndarray + Sorted query times. + full_range : bool + Whether to keep every non-empty segment, rather than only those + whose span contains at least one of ``t`` (plus the last segment + when ``t`` extends past it). + + Returns + ------- + numpy.ndarray + Segment indices, ascending. Empty segments are always dropped. + """ + if full_range: + return self.indices + return self.indices[_find_ts_indices(self._starts, self._ends, t)] + + +class VariableObserver(ABC): + """How one variable's leaves are evaluated over a solution's segments. + + Implementations read only these attributes of the ``variable`` handed to + them: ``all_ts``, ``all_ys``, ``all_yps``, ``all_inputs``, + ``all_inputs_stacked``, ``t_pts``, ``hermite_interpolation``, + ``time_integral``, ``base_variables``, ``sensitivity_names``, + ``all_solution_sensitivities``, ``_name`` and ``_shape``. + """ + + #: Derived caches: built on demand, never pickled (see __getstate__). + _selector = None + _selector_ts = None + _serialised = None + + def __getstate__(self): + """The observer's state without its derived caches, rebuilt on demand.""" + state = self.__dict__.copy() + for key in ("_selector", "_selector_ts", "_serialised"): + state.pop(key, None) + return state + + @property + @abstractmethod + def leaves(self) -> list: + """The variable's per-sub-solution evaluable leaves, in solve order.""" + + def segments( + self, + variable: pybamm.ProcessedVariable, + t: npt.NDArray[np.float64], + full_range: bool, + ) -> npt.NDArray[np.intp]: + """Indices of ``variable``'s sub-solutions covering ``t``. + + See :meth:`SegmentSelector.select`. + """ + # Keyed to the time arrays it was built from: an observer can be + # shared by variables of different solutions. + if self._selector_ts is not variable.all_ts: + self._selector = SegmentSelector(variable.all_ts) + self._selector_ts = variable.all_ts + return self._selector.select(t, full_range) + + @abstractmethod + def observe_raw( + self, variable: pybamm.ProcessedVariable + ) -> npt.NDArray[np.float64]: + """Evaluate on the solution's own time points, shaped by ``_shape``.""" + + @abstractmethod + def observe_hermite( + self, variable: pybamm.ProcessedVariable, t: npt.NDArray[np.float64] + ) -> npt.NDArray[np.float64]: + """Evaluate at arbitrary sorted times ``t``, cubic-Hermite in state.""" + + @abstractmethod + def sensitivities( + self, variable: pybamm.ProcessedVariable + ) -> dict[str, npt.NDArray[np.float64]]: + """``{"all": (N, n_p), param: (N,)}`` forward sensitivities of the variable.""" + + +class CasadiObserver(VariableObserver): + """Observation through serialised CasADi functions and the IDAKLU kernels. + + Parameters + ---------- + leaves : list of :class:`casadi.Function` + One function per sub-solution, evaluating ``(t, y, p_stacked)``. + """ + + def __init__(self, leaves: list): + self._leaves = leaves + + @property + def leaves(self) -> list: + return self._leaves + + def _serialise(self, idxs): + """Serialised leaves for `idxs`, memoised by leaf identity. + + CasADi functions are immutable and serialising one is about half the + cost of an observe call, so the bytes are built once per leaf. + """ + if self._serialised is None: + self._serialised = {} + serialised = self._serialised + funcs = [None] * len(idxs) + for i, idx in enumerate(idxs): + leaf = self._leaves[idx] + key = id(leaf) + if key not in serialised: + serialised[key] = leaf.serialize() + funcs[i] = serialised[key] + return funcs + + def _setup(self, variable, t, full_range): + """Per-segment IDAKLU inputs: ``(ts, ys, yps, funcs, inputs, is_f_contiguous)``.""" + pybamm.logger.debug("Setting up C++ interpolation inputs") + idxs = self.segments(variable, t, full_range) + hermite = variable.hermite_interpolation + all_ts, all_ys = variable.all_ts, variable.all_ys + + ts = [all_ts[idx] for idx in idxs] + ys = [all_ys[idx] for idx in idxs] + yps = [variable.all_yps[idx] for idx in idxs] if hermite else None + inputs = [variable.all_inputs_stacked[idx] for idx in idxs] + + is_f_contiguous = _is_f_contiguous(ys) + + ts = idaklu.VectorRealtypeNdArray(ts) + ys = idaklu.VectorRealtypeNdArray(ys) + yps = idaklu.VectorRealtypeNdArray(yps) if hermite else None + inputs = idaklu.VectorRealtypeNdArray(inputs) + + return ts, ys, yps, self._serialise(idxs), inputs, is_f_contiguous + + def observe_raw(self, variable): + pybamm.logger.debug("Observing the variable raw data") + t = variable.t_pts + ts, ys, _, funcs, inputs, is_f_contiguous = self._setup( + variable, t, full_range=True + ) + return idaklu.observe( + ts, ys, inputs, funcs, is_f_contiguous, variable._shape(t) + ) + + def observe_hermite(self, variable, t): + pybamm.logger.debug("Observing and Hermite interpolating the variable") + ts, ys, yps, funcs, inputs, _ = self._setup(variable, t, full_range=False) + return idaklu.observe_hermite_interp( + t, ts, ys, yps, inputs, funcs, variable._shape(t) + ) + + def sensitivities(self, variable): + sensitivity_names = variable.sensitivity_names + all_S_var = [] + for ts, ys, inputs, base_variable, dy_dp in zip( + variable.all_ts, + variable.all_ys, + variable.all_inputs, + variable.base_variables, + variable.all_solution_sensitivities["all"], + strict=True, + ): + sensitivity_inputs = { + name: inputs[name] for name in sensitivity_names if name in inputs + } + sensitivity_inputs_stacked = casadi.vertcat( + *[sensitivity_inputs[name] for name in sensitivity_names] + ) + + # Set up symbolic variables + t_casadi = casadi.MX.sym("t") + y_casadi = casadi.MX.sym("y", ys.shape[0]) + p_casadi = { + name: casadi.MX.sym(name, value.shape[0]) + for name, value in sensitivity_inputs.items() + } + + p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) + + # Non-target inputs can still appear in the tree (e.g. from experiment + # steps), so they stay concrete while the targets go symbolic. + inputs_for_casadi = {**inputs, **p_casadi} + + var_casadi = base_variable.to_casadi( + t_casadi, y_casadi, inputs=inputs_for_casadi + ) + dvar_dy = casadi.jacobian(var_casadi, y_casadi) + dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) + + # Convert to functions and evaluate index-by-index + dvar_dy_func = casadi.Function( + "dvar_dy", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dy] + ) + dvar_dp_func = casadi.Function( + "dvar_dp", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dp] + ) + dvar_dy_eval = casadi.diagcat( + *[ + dvar_dy_func(t, ys[:, idx], sensitivity_inputs_stacked) + for idx, t in enumerate(ts) + ] + ) + dvar_dp_eval = casadi.vertcat( + *[ + dvar_dp_func(t, ys[:, idx], sensitivity_inputs_stacked) + for idx, t in enumerate(ts) + ] + ) + + # Compute sensitivity + S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval + + all_S_var.append(S_var) + + S_var = np.vstack(all_S_var) + if variable.time_integral is not None: + # The value sums over every segment at once, so its sensitivity must too + S_var = variable.time_integral.postfix_sensitivities( + variable._name, + self.observe_raw(variable), + variable.t_pts, + variable.all_inputs[0], + sensitivity_names, + S_var, + ) + return pack_sensitivity_dict(S_var, sensitivity_names) + + +def as_observer(leaves: VariableObserver | list) -> VariableObserver: + """Coerce leaves to an observer. + + Parameters + ---------- + leaves : :class:`VariableObserver` or list of :class:`casadi.Function` + An observer, returned unchanged, or one CasADi function per + sub-solution, wrapped in a :class:`CasadiObserver`. + + Returns + ------- + :class:`VariableObserver` + """ + if isinstance(leaves, VariableObserver): + return leaves + return CasadiObserver(leaves) + + +def check_variable_in_solve( + solution: pybamm.Solution, name: str, var_pybamm: pybamm.Symbol +) -> None: + """Reject a state-dependent variable an outputs-only solve did not store. + + Parameters + ---------- + solution : :class:`pybamm.Solution` + The solution the variable is read from. + name : str + The variable's name. + var_pybamm : :class:`pybamm.Symbol` + The variable's discretised expression. + + Raises + ------ + KeyError + If the solve returned variables only and ``var_pybamm`` reads states. + """ + if not solution.variables_returned: + return + if var_pybamm.has_symbol_of_classes( + pybamm.expression_tree.state_vector.StateVector + ): + raise KeyError( + f"Cannot process variable '{name}' as it was not part of the " + "solve. Please re-run the solve with `output_variables` set to " + "include this variable." + ) + + +def pack_sensitivity_dict( + sensitivity_matrix: npt.NDArray[np.float64], sensitivity_names: Sequence[str] +) -> dict[str, npt.NDArray[np.float64]]: + """Split a sensitivity matrix into its per-parameter vectors. + + Parameters + ---------- + sensitivity_matrix : numpy.ndarray + Sensitivities of shape ``(N, n_p)``, one column per parameter. + sensitivity_names : list of str + Parameter names, in column order. + + Returns + ------- + dict + The whole matrix under ``"all"``, plus one flat ``(N,)`` vector per + parameter, keyed by its name. + """ + sensitivities = {"all": sensitivity_matrix} + for i, name in enumerate(sensitivity_names): + sensitivities[name] = sensitivity_matrix[:, i : i + 1].reshape(-1) + return sensitivities + + +def _is_f_contiguous(all_ys): + """Whether every array in ``all_ys`` is Fortran-contiguous. + + Parameters + ---------- + all_ys : list of numpy.ndarray + Per-sub-solution state arrays. + + Returns + ------- + bool + """ + + return all(isinstance(y, np.ndarray) and y.data.f_contiguous for y in all_ys) + + +def _find_ts_indices(starts, ends, t): + """Segments containing at least one of ``t``. + + Parameters + ---------- + starts, ends : list of float + First and last time of each segment, successively increasing. + t : array-like + Sorted query times. + + Returns + ------- + list of int + Positions in ``starts``/``ends``, plus the last segment when ``t`` + extends past it. + """ + + indices = [] + + # Get the minimum and maximum values of the target values `t` + t_min, t_max = t[0], t[-1] + + # Binary search for the range of segments where t_min and t_max could lie + low_idx = bisect.bisect_left(ends, t_min) + high_idx = bisect.bisect_right(starts, t_max) + + for idx in range(low_idx, high_idx): + # Binary search within `t` to check if any value falls within the segment + i = bisect.bisect_left(t, starts[idx]) + if i < len(t) and t[i] <= ends[idx]: + indices.append(idx) + + # Past the last segment, extrapolate from it + if (t_max > ends[-1]) and (len(indices) == 0 or indices[-1] != len(ends) - 1): + indices.append(len(ends) - 1) + + return indices diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_legacy_mutation.py b/packages/pybamm/tests/unit/test_expression_tree/test_legacy_mutation.py index 665f4e852f..5ae7ff0ed3 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_legacy_mutation.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_legacy_mutation.py @@ -43,7 +43,6 @@ def run_compatibility_script(source): # writes to attributes named like symbol slots whose receiver is not a symbol _NON_SYMBOL_WRITES = { - ("pybamm/solvers/idaklu_solver.py", "sol.y"), ("pybamm/spatial_methods/scikit_finite_element_3d.py", "M._shape"), ( ( diff --git a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py index 6464bafa36..587d89e955 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -1421,6 +1421,43 @@ def test_with_output_variables_and_sensitivities(self): err_msg=f"Failed for '{varname}', sensitivity '{key}'", ) + def test_output_variables_sensitivities_are_keyed_by_sensitivity_parameters( + self, + ): + # An input that no sensitivity was requested for takes no key or column. + parameter_values = pybamm.ParameterValues("Chen2020") + inputs = { + "Current function [A]": 0.68, + "Electrode height [m]": parameter_values["Electrode height [m]"], + } + parameter_values.update({key: "[input]" for key in inputs}) + name = "Voltage [V]" + t_interp = np.linspace(0, 600, 11) + + def voltage_sensitivities(output_variables): + sim = pybamm.Simulation( + pybamm.lithium_ion.SPM(), + parameter_values=parameter_values, + solver=pybamm.IDAKLUSolver(output_variables=output_variables), + ) + solution = sim.solve( + [0, 600], + t_interp=t_interp, + inputs=inputs, + calculate_sensitivities=["Current function [A]"], + ) + return solution[name].sensitivities + + full = voltage_sensitivities(None) + outputs_only = voltage_sensitivities([name]) + assert set(outputs_only) == {"all", "Current function [A]"} + np.testing.assert_allclose( + outputs_only["Current function [A]"], + full["Current function [A]"], + rtol=1e-6, + atol=1e-10, + ) + def test_with_output_variables_and_event_termination(self): model = pybamm.lithium_ion.DFN() parameter_values = pybamm.ParameterValues("Chen2020") diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py b/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py new file mode 100644 index 0000000000..b6b892943d --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py @@ -0,0 +1,228 @@ +"""Tests for observation backends and variable observers.""" + +import pickle +from itertools import pairwise +from types import SimpleNamespace + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import ( + CASADI_OBSERVATION, + CasadiObservation, + join_observations, +) +from pybamm.solvers.variable_observer import ( + CasadiObserver, + SegmentSelector, + as_observer, + pack_sensitivity_dict, +) + + +class _RecordingObservation(CasadiObservation): + """A CasADi backend that records the segment slices taken of it.""" + + def __init__(self): + self.keys = [] + + def __getitem__(self, key): + self.keys.append(key) + return self + + +@pytest.fixture(scope="module") +def spm_solution(): + model = pybamm.lithium_ion.SPM() + return pybamm.Simulation(model).solve([0, 600]) + + +def _split(solution, *boundaries): + """``solution``'s trajectory as consecutive Solutions, cut at ``boundaries``.""" + model = solution.all_models[0] + edges = [0, *boundaries, len(solution.t)] + return [ + pybamm.Solution( + solution.t[start:end], + solution.y[:, start:end], + model, + {}, + all_yps=solution.all_yps[0][:, start:end], + ) + for start, end in pairwise(edges) + ] + + +class TestSegmentSlicing: + def test_the_casadi_backend_is_a_shared_stateless_value(self): + assert isinstance(CASADI_OBSERVATION, CasadiObservation) + assert CASADI_OBSERVATION[:1] is CASADI_OBSERVATION + assert CASADI_OBSERVATION[-1:] is CASADI_OBSERVATION + + def test_casadi_backends_are_interchangeable(self): + # A Solution unpickled from a derived one carries its own instance. + restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) + assert restored is not CASADI_OBSERVATION + assert restored == CASADI_OBSERVATION + assert hash(restored) == hash(CASADI_OBSERVATION) + + +class TestJoin: + def test_runs_sharing_a_backend_join_to_it(self): + joined = join_observations([CASADI_OBSERVATION, CASADI_OBSERVATION]) + assert joined is CASADI_OBSERVATION + + def test_equal_backends_join_to_the_first(self): + restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) + assert join_observations([restored, CASADI_OBSERVATION]) is restored + + def test_different_backends_do_not_join(self): + with pytest.raises(pybamm.SolverError, match=r"different observation"): + join_observations([CASADI_OBSERVATION, _RecordingObservation()]) + + +class TestObservationCarriedThroughDerivedSolutions: + @pytest.fixture + def recorded(self, spm_solution): + first, second = _split(spm_solution, 5) + backend = _RecordingObservation() + first._observation = backend + second._observation = backend + return first, second, backend + + def test_first_and_last_state_slice_it(self, recorded): + first, _, backend = recorded + assert first.first_state._observation is backend + assert first.last_state._observation is backend + assert backend.keys == [slice(None, 1), slice(-1, None)] + + def test_addition_joins_it(self, recorded): + first, second, backend = recorded + assert (first + second)._observation is backend + + def test_from_sub_solutions_joins_it(self, recorded): + first, second, backend = recorded + joined = pybamm.Solution.from_sub_solutions([first, second]) + assert joined._observation is backend + + def test_copy_keeps_it(self, recorded): + first, _, backend = recorded + assert first.copy()._observation is backend + + def test_solutions_read_differently_cannot_be_added(self, recorded): + first, second, _ = recorded + second._observation = CASADI_OBSERVATION + with pytest.raises(pybamm.SolverError, match=r"different observation"): + first + second + + +class TestSegmentSelector: + def test_full_range_keeps_every_nonempty_segment(self): + selector = SegmentSelector( + [np.array([]), np.array([0.0, 1.0]), np.array([2.0])] + ) + np.testing.assert_array_equal( + selector.select(np.array([0.0]), full_range=True), [1, 2] + ) + + def test_restricted_range_keeps_only_covering_segments(self): + selector = SegmentSelector([np.array([0.0, 1.0]), np.array([2.0, 3.0])]) + np.testing.assert_array_equal( + selector.select(np.array([2.5]), full_range=False), [1] + ) + + def test_extrapolating_past_the_end_keeps_the_last_segment(self): + selector = SegmentSelector([np.array([0.0, 1.0]), np.array([2.0, 3.0])]) + np.testing.assert_array_equal( + selector.select(np.array([4.0]), full_range=False), [1] + ) + + +class TestObserverSegments: + def test_selection_follows_the_variable_asked_about(self): + observer = CasadiObserver([None, None]) + early = SimpleNamespace(all_ts=[np.array([0.0, 1.0]), np.array([1.0, 2.0])]) + late = SimpleNamespace(all_ts=[np.array([0.0, 5.0]), np.array([5.0, 9.0])]) + t = np.array([3.0]) + + np.testing.assert_array_equal(observer.segments(early, t, False), [1]) + np.testing.assert_array_equal(observer.segments(late, t, False), [0]) + np.testing.assert_array_equal(observer.segments(early, t, False), [1]) + + def test_an_observer_shared_across_solutions_reads_each_correctly( + self, spm_solution + ): + name = "Voltage [V]" + early = pybamm.Solution.from_sub_solutions(_split(spm_solution, 3)) + late = pybamm.Solution.from_sub_solutions(_split(spm_solution, 12)) + t = np.array([spm_solution.t[6]]) + base = [m.get_processed_variable_or_event(name) for m in early.all_models] + observer = early[name]._observer + early[name](t) + + shared = pybamm.process_variable(name, base, observer, late) + + np.testing.assert_allclose(shared(t), late[name](t), rtol=1e-12) + + +class TestObserverCoercion: + def test_a_bare_casadi_list_becomes_a_casadi_observer(self): + observer = as_observer([None, None]) + assert isinstance(observer, CasadiObserver) + assert observer.leaves == [None, None] + + def test_an_observer_is_passed_through(self): + observer = CasadiObserver([None]) + assert as_observer(observer) is observer + + def test_process_variable_accepts_either_form(self, spm_solution): + name = "Terminal voltage [V]" + base = [ + m.get_processed_variable_or_event(name) for m in spm_solution.all_models + ] + leaves = spm_solution[name]._observer.leaves + + direct = pybamm.process_variable(name, base, leaves, spm_solution) + wrapped = pybamm.process_variable(name, base, as_observer(leaves), spm_solution) + np.testing.assert_allclose(direct.entries, wrapped.entries) + + +class TestCasadiObserverCaches: + def test_casadi_leaves_serialise_once_across_calls(self, spm_solution): + variable = spm_solution.copy()["Terminal voltage [V]"] + + variable.entries + serialised = dict(variable._observer._serialised) + assert serialised + variable(np.linspace(0, 600, 11)) + variable(np.linspace(0, 600, 7)) + for key, value in variable._observer._serialised.items(): + assert value is serialised[key] + + def test_pickling_drops_the_caches(self, spm_solution): + variable = spm_solution.copy()["Terminal voltage [V]"] + t = np.linspace(0, 600, 11) + expected = variable(t) + observer = variable._observer + assert observer._serialised is not None + assert observer._selector is not None + + restored = pickle.loads(pickle.dumps(observer)) + + for cache in ("_selector", "_selector_ts", "_serialised"): + assert cache not in restored.__dict__ + assert cache in observer.__dict__ + variable._observer = restored + np.testing.assert_allclose(variable(t), expected) + + +class TestPackSensitivityDict: + def test_all_block_plus_one_flat_vector_per_parameter(self): + sensitivity_matrix = np.arange(6.0).reshape(3, 2) + packed = pack_sensitivity_dict(sensitivity_matrix, ["a", "b"]) + + assert set(packed) == {"all", "a", "b"} + np.testing.assert_array_equal(packed["all"], sensitivity_matrix) + np.testing.assert_array_equal(packed["a"], [0.0, 2.0, 4.0]) + np.testing.assert_array_equal(packed["b"], [1.0, 3.0, 5.0]) diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py new file mode 100644 index 0000000000..f5cec510ee --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py @@ -0,0 +1,159 @@ +"""Tests for the row layout of ``output_variables`` payloads.""" + +import casadi +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import OutputAssembly + + +@pytest.fixture(scope="module") +def spm_solution(): + return pybamm.Simulation(pybamm.lithium_ion.SPM()).solve([0, 3600]) + + +@pytest.fixture +def solution(spm_solution): + # A copy, so variables attached by one test do not leak into the next. + return spm_solution.copy() + + +def _dense_function(size): + """A CasADi ``f(t, y, p)`` with ``size`` dense rows.""" + t, y, p = (casadi.MX.sym(name) for name in "typ") + return casadi.Function("dense", [t, y, p], [casadi.repmat(t + y + p, size, 1)]) + + +class TestLayoutContract: + def test_0d_time_major(self, solution): + # One (n_times, output_size) array per sub-solution; for 0D + # output_size == 1 and data[t, 0] is the value at t. + base = [pybamm.StateVector(slice(0, 1))] + n_t = len(solution.t) + values = np.linspace(3.0, 4.2, n_t) + data = [values.reshape(n_t, 1)] + pvc = pybamm.ProcessedVariableComputed( + base, [_dense_function(1)], data, solution + ) + np.testing.assert_allclose(pvc.entries.reshape(-1), values) + + def test_1d_time_major_unrolls_to_space_by_time(self, solution): + # A 1D variable's (n_times, len_space) array unrolls to + # (len_space, n_times) via reshape((n_times, len_space)).transpose(). + base_pv = solution["X-averaged negative particle concentration [mol.m-3]"] + var = base_pv.base_variables[0] + len_space = var.shape[0] + n_t = len(solution.t) + rows = np.arange(n_t * len_space, dtype=float).reshape(n_t, len_space) + pvc = pybamm.ProcessedVariableComputed( + [var], [_dense_function(len_space)], [rows], solution + ) + assert pvc.entries.shape[0] == len_space + np.testing.assert_allclose(pvc.entries[:, 0], rows[0, :]) + np.testing.assert_allclose(pvc.entries[:, -1], rows[-1, :]) + + +class TestOutputAssembly: + """A scalar, a 20-component vector and a second scalar. + + An ordinal-indexed reader would return one component for the vector and + shift the scalar after it. + """ + + _NAMES = [ + "Voltage [V]", + "X-averaged negative particle concentration [mol.m-3]", + "Current [A]", + ] + + @staticmethod + def _assembly(solution, names=None): + """An assembly over ``names`` and a payload whose entries are row indices.""" + names = names or TestOutputAssembly._NAMES + model = solution.all_models[0] + casadi_fns = { + name: _dense_function( + int(np.prod(model.get_processed_variable_or_event(name).shape)) + ) + for name in names + } + assembly = OutputAssembly(names, casadi_fns) + data = np.tile(np.arange(assembly.n_rows, dtype=float), (len(solution.t), 1)) + return assembly, data + + def test_rows_are_sliced_by_component_count_not_ordinal(self, solution): + assembly, data = self._assembly(solution) + assembly.attach(solution, data) + + np.testing.assert_allclose(solution["Voltage [V]"].entries, 0.0) + np.testing.assert_allclose( + solution[self._NAMES[1]].entries[:, 0], np.arange(1.0, 21.0) + ) + # 21, not 2: the vector consumed rows 1--20 rather than row 1 alone. + np.testing.assert_allclose(solution["Current [A]"].entries, 21.0) + + def test_a_payload_of_the_wrong_width_is_rejected(self, solution): + assembly, data = self._assembly(solution) + with pytest.raises(pybamm.SolverError, match=r"Output row count mismatch"): + assembly.attach(solution, data[:, :-1]) + + def test_sensitivities_are_named_and_flattened_per_parameter(self, solution): + assembly, data = self._assembly(solution, ["Voltage [V]"]) + n_t = len(solution.t) + sensitivities = np.arange(n_t * 2, dtype=float).reshape(n_t, 1, 2) + + assembly.attach( + solution, + data, + sensitivities=sensitivities, + sensitivity_names=["a", "b"], + ) + + # Read the field, not the property: the property short-circuits to {} on + # this input-free SPM solve. + attached = solution["Voltage [V]"]._sensitivities + assert set(attached) == {"all", "a", "b"} + assert attached["all"].shape == (n_t, 2) + np.testing.assert_allclose(attached["a"], sensitivities[:, 0, 0]) + np.testing.assert_allclose(attached["b"], sensitivities[:, 0, 1]) + + def test_sensitivities_of_the_wrong_shape_are_rejected(self, solution): + assembly, data = self._assembly(solution, ["Voltage [V]"]) + n_t = len(solution.t) + with pytest.raises( + pybamm.SolverError, match=r"Output sensitivity shape mismatch" + ): + assembly.attach( + solution, + data, + sensitivities=np.zeros((n_t, 1, 1)), + sensitivity_names=["a", "b"], + ) + + def test_a_solve_without_sensitivities_leaves_an_empty_mapping(self, solution): + # Not None: an outputs-only solve retains no state to compute them from, + # so the answer is "there are none", not "ask again later". + assembly, data = self._assembly(solution, ["Voltage [V]"]) + assembly.attach(solution, data) + assert solution["Voltage [V]"]._sensitivities == {} + + def test_a_sparse_variable_owns_only_its_nonzero_rows(self, solution): + t, y, p = (casadi.MX.sym(name) for name in "typ") + sparse = casadi.Function( + "sparse", [t, y, p], [casadi.vertcat(t, casadi.MX(1, 1))] + ) + assembly = OutputAssembly( + ["Voltage [V]", "Current [A]"], + {"Voltage [V]": sparse, "Current [A]": _dense_function(1)}, + ) + assert assembly.n_rows == 2 + + n_t = len(solution.t) + with pytest.raises(pybamm.SolverError, match=r"sparse variable"): + assembly.attach( + solution, + np.zeros((n_t, 2)), + sensitivities=np.zeros((n_t, 2, 1)), + sensitivity_names=["a"], + ) diff --git a/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py b/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py index 259f78f91a..4208b1a93b 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py +++ b/packages/pybamm/tests/unit/test_solvers/test_processed_variable.py @@ -5,6 +5,7 @@ import pybamm import tests +from pybamm.solvers.processed_variable import ProcessedVariable0D _hermite_args = [True, False] @@ -2165,6 +2166,46 @@ def test_process_variable_unstructured_detection(self): assert isinstance(processed_var, pybamm.ProcessedVariableUnstructured) + @staticmethod + def _scalar_variable(): + """A scalar variable, its CasADi function and a solution to read it from.""" + var = pybamm.t * pybamm.StateVector(slice(0, 1)) + t_sol = np.linspace(0, 1) + y_sol = np.array([np.linspace(0, 5)]) + solution = pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}) + return var, to_casadi(var, y_sol), solution + + @pytest.mark.parametrize( + "constructor", + [pybamm.ProcessedVariable, ProcessedVariable0D, pybamm.process_variable], + ) + def test_base_variables_casadi_keyword_is_deprecated(self, constructor): + var, var_casadi, solution = self._scalar_variable() + with pytest.warns(DeprecationWarning, match=r"base_variables_casadi"): + processed_var = constructor( + "test", [var], base_variables_casadi=[var_casadi], solution=solution + ) + assert processed_var._observer.leaves == [var_casadi] + + def test_base_variables_casadi_keyword_cannot_repeat_the_observer(self): + var, var_casadi, solution = self._scalar_variable() + with pytest.raises(TypeError, match=r"once"): + pybamm.ProcessedVariable( + "test", + [var], + [var_casadi], + base_variables_casadi=[var_casadi], + solution=solution, + ) + + def test_base_variables_casadi_attribute_is_deprecated(self): + var, var_casadi, solution = self._scalar_variable() + processed_var = pybamm.process_variable("test", [var], [var_casadi], solution) + with pytest.warns(DeprecationWarning, match=r"base_variables_casadi"): + assert processed_var.base_variables_casadi == [var_casadi] + with pytest.raises(AttributeError): + processed_var.base_variables_casadi = [var_casadi] + class TestProcessedVariableUnstructuredFVM: @staticmethod diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index 0d36c51bf1..a9997eecb2 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -1,6 +1,7 @@ # # Tests for the Solution class # +import copyreg import io import json import logging @@ -16,10 +17,42 @@ from scipy.io import loadmat import pybamm +from pybamm.solvers.observation import CASADI_OBSERVATION from pybamm.solvers.solution import _DEFAULT_SOLUTION_OPTIONS, make_cycle_solution from tests import get_discretisation_for_testing +class _PicklerBeforeObservationBackends(pickle.Pickler): + """Pickles Solutions and ProcessedVariables in their earlier layout. + + That layout has no ``Solution._observation``, and keeps a processed + variable's CasADi functions as ``base_variables_casadi``. + """ + + def reducer_override(self, obj): + if isinstance(obj, pybamm.Solution): + slots = { + name: getattr(obj, name) + for cls in type(obj).__mro__ + for name in getattr(cls, "__slots__", ()) + if not name.startswith("__") and hasattr(obj, name) + } + slots.pop("_observation", None) + state = (None, slots) + elif isinstance(obj, pybamm.ProcessedVariable): + state = dict(obj.__dict__) + state["base_variables_casadi"] = state.pop("_observer").leaves + else: + return NotImplemented + return copyreg.__newobj__, (type(obj),), state + + +def _round_trip_before_observation_backends(obj): + buffer = io.BytesIO() + _PicklerBeforeObservationBackends(buffer, pickle.HIGHEST_PROTOCOL).dump(obj) + return pickle.loads(buffer.getvalue()) + + class TestSolution: def test_init(self): t = np.linspace(0, 1) @@ -1165,6 +1198,83 @@ def test_pickle_first_states_across_processes(self, tmp_path): assert "DATA" in save_result.stdout assert save_result.stdout == load_result.stdout + def test_load_solution_pickled_before_observation_backends(self): + sim = pybamm.Simulation( + pybamm.lithium_ion.SPM(), + experiment=pybamm.Experiment( + ["Discharge at 1C for 1 minute", "Rest for 1 minute"] + ), + ) + solution = sim.solve(calc_esoh=False) + t = np.linspace(0, 120, 7) + solution["Voltage [V]"].entries + + loaded = _round_trip_before_observation_backends(solution) + + # An unpickled backend is a new instance, so this one is the default + assert loaded._observation is CASADI_OBSERVATION + np.testing.assert_allclose( + loaded["Voltage [V]"](t), solution["Voltage [V]"](t), rtol=1e-12 + ) + np.testing.assert_allclose( + loaded["Current [A]"](t), solution["Current [A]"](t), rtol=1e-12 + ) + np.testing.assert_allclose( + loaded.first_state["Voltage [V]"].entries, + solution.first_state["Voltage [V]"].entries, + ) + joined = loaded.sub_solutions[0] + loaded.sub_solutions[1] + np.testing.assert_allclose( + joined["Voltage [V]"](t), solution["Voltage [V]"](t), rtol=1e-12 + ) + summary = solution.summary_variables.get_summary_variables() + loaded_summary = loaded.summary_variables.get_summary_variables() + assert loaded_summary.keys() == summary.keys() + for name, value in summary.items(): + np.testing.assert_allclose(loaded_summary[name], value, rtol=1e-12) + + def test_load_processed_variable_pickled_before_observers(self): + parameter_values = pybamm.ParameterValues("Chen2020") + parameter_values.update({"Current function [A]": "[input]"}) + sim = pybamm.Simulation( + pybamm.lithium_ion.SPM(), parameter_values=parameter_values + ) + solution = sim.solve( + [0, 600], + inputs={"Current function [A]": 1.0}, + calculate_sensitivities=True, + ) + voltage = solution["Voltage [V]"] + t = np.linspace(0, 600, 7) + + loaded = _round_trip_before_observation_backends(voltage) + + assert "base_variables_casadi" not in loaded.__dict__ + np.testing.assert_allclose(loaded(t), voltage(t), rtol=1e-12) + np.testing.assert_allclose( + loaded.sensitivities["Current function [A]"], + voltage.sensitivities["Current function [A]"], + rtol=1e-12, + ) + np.testing.assert_allclose( + loaded.as_computed().entries, voltage.entries, rtol=1e-12 + ) + + def test_update_rebuilds_every_listed_variable(self): + model = pybamm.BaseModel() + c = pybamm.Variable("c") + model.rhs = {c: -c} + model.initial_conditions = {c: 1} + model.variables = {"c": c, "2c": 2 * c} + solution = pybamm.IDAKLUSolver().solve(model, [0, 1]) + built = solution["c"] + + solution.update(["c", "2c"]) + + assert solution["c"] is not built + np.testing.assert_allclose(solution["c"].entries, built.entries) + assert "2c" in solution._variables + def test_solution_evals_with_inputs(self): model = pybamm.lithium_ion.SPM() geometry = model.default_geometry From facb8e5688782be119c26286bb8690e1932c6eb5 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Thu, 24 Sep 2026 14:38:07 +0100 Subject: [PATCH 2/7] test: mark the new pickle round trips nosec, as main's tests do --- .../tests/unit/test_solvers/test_observation_backend.py | 8 ++++---- packages/pybamm/tests/unit/test_solvers/test_solution.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py b/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py index b6b892943d..196a0cc2fc 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py @@ -1,6 +1,6 @@ """Tests for observation backends and variable observers.""" -import pickle +import pickle # nosec B403 - used in tests with trusted input from itertools import pairwise from types import SimpleNamespace @@ -62,7 +62,7 @@ def test_the_casadi_backend_is_a_shared_stateless_value(self): def test_casadi_backends_are_interchangeable(self): # A Solution unpickled from a derived one carries its own instance. - restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) + restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) # nosec B301 assert restored is not CASADI_OBSERVATION assert restored == CASADI_OBSERVATION assert hash(restored) == hash(CASADI_OBSERVATION) @@ -74,7 +74,7 @@ def test_runs_sharing_a_backend_join_to_it(self): assert joined is CASADI_OBSERVATION def test_equal_backends_join_to_the_first(self): - restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) + restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) # nosec B301 assert join_observations([restored, CASADI_OBSERVATION]) is restored def test_different_backends_do_not_join(self): @@ -208,7 +208,7 @@ def test_pickling_drops_the_caches(self, spm_solution): assert observer._serialised is not None assert observer._selector is not None - restored = pickle.loads(pickle.dumps(observer)) + restored = pickle.loads(pickle.dumps(observer)) # nosec B301 for cache in ("_selector", "_selector_ts", "_serialised"): assert cache not in restored.__dict__ diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index a9997eecb2..150ff53273 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -50,7 +50,7 @@ def reducer_override(self, obj): def _round_trip_before_observation_backends(obj): buffer = io.BytesIO() _PicklerBeforeObservationBackends(buffer, pickle.HIGHEST_PROTOCOL).dump(obj) - return pickle.loads(buffer.getvalue()) + return pickle.loads(buffer.getvalue()) # nosec B301 class TestSolution: From 125681d69af128f85701bd83410a56667be6cff8 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Fri, 2 Oct 2026 16:41:26 +0100 Subject: [PATCH 3/7] refactor: read each segment through its own model's observer cache A Solution carried an observation backend, the module-level CASADI_OBSERVATION, which first_state, last_state, +, from_sub_solutions and copy sliced and joined. The backend held no state, so every slice and join returned it unchanged, while the CasADi leaves it compiled were cached on each model. A Solution now holds no observation state. build_variable reads each segment through its own model's ObserverCache, which replaces model._variables_casadi and keeps its rules: keyed by variable name, copied by new_copy, and skipping time integrals. A model pickled before the cache existed builds one on first read. ObservationBackend, CasadiObservation, CASADI_OBSERVATION and join_observations are removed with their tests. --- .../pybamm/src/pybamm/models/base_model.py | 8 +- .../pybamm/src/pybamm/solvers/observation.py | 243 +++++++----------- .../pybamm/src/pybamm/solvers/solution.py | 14 +- ...rvation_backend.py => test_observation.py} | 132 +++++----- .../tests/unit/test_solvers/test_solution.py | 38 +-- 5 files changed, 180 insertions(+), 255 deletions(-) rename packages/pybamm/tests/unit/test_solvers/{test_observation_backend.py => test_observation.py} (65%) diff --git a/packages/pybamm/src/pybamm/models/base_model.py b/packages/pybamm/src/pybamm/models/base_model.py index b225a095fd..ec13e25d69 100644 --- a/packages/pybamm/src/pybamm/models/base_model.py +++ b/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() diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py index e4b81da87d..cbc8c4a113 100644 --- a/packages/pybamm/src/pybamm/solvers/observation.py +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -1,9 +1,9 @@ -"""Observation backends: how a :class:`pybamm.Solution` reads its variables. +"""How a :class:`pybamm.Solution` reads its variables. -An :class:`ObservationBackend` turns a variable name into a processed variable -over a run of a Solution's sub-solutions, its *segments*. A *leaf* is one -segment's compiled form of one variable, so a variable spanning ``n`` -sub-solutions has ``n`` leaves. +:func:`build_variable` turns a variable name into a processed variable over a +Solution's sub-solutions, its *segments*, reading each segment through its own +model. A *leaf* is one segment's compiled form of one variable; each model keeps +its leaves in an :class:`ObserverCache`, so every solution of it reuses them. :class:`OutputAssembly` is the ``output_variables`` counterpart: the solver has already computed those variables, so it owns the payload's row layout and @@ -12,7 +12,6 @@ from __future__ import annotations -from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from itertools import accumulate, pairwise @@ -29,129 +28,117 @@ ) -class ObservationBackend(ABC): - """How a Solution turns a variable name into something it can evaluate. +class ObserverCache: + """The compiled leaves of one model's variables, keyed by variable name. - A backend covers an ordered run of a Solution's sub-solutions, one - *segment* each. ``backend[key]`` indexes *segments*, matching the slice the - Solution applies to its sub-solutions; :func:`join_observations` combines - the backends of consecutive Solutions. + A name is assumed to keep one expression for the model's lifetime. """ - @abstractmethod - def __getitem__(self, key: slice) -> ObservationBackend: - """This backend restricted to a slice of the Solution's segments.""" + def __init__(self): + self._casadi_leaves: dict[str, casadi.Function] = {} - @abstractmethod - def build_variable( - self, solution: pybamm.Solution, name: str - ) -> BaseProcessedVariable: - """The processed variable for ``name``, ready to evaluate. + @classmethod + def of(cls, model: pybamm.BaseModel) -> ObserverCache: + """The cache of ``model``, created on first use.""" + if model._observer_cache is None: + model._observer_cache = cls() + return model._observer_cache + + def copy(self) -> ObserverCache: + """A cache that starts with these leaves and grows independently.""" + new = type(self)() + new._casadi_leaves = self._casadi_leaves.copy() + return new + + def casadi_leaf( + self, + solution: pybamm.Solution, + key: str, + var_pybamm: pybamm.Symbol, + inputs: dict, + ys_shape: tuple[int, ...], + ) -> tuple[ + casadi.Function, pybamm.Symbol, pybamm.ProcessedVariableTimeIntegral | None + ]: + """One segment's CasADi leaf for ``var_pybamm``, converted on first use. Parameters ---------- solution : :class:`pybamm.Solution` - The solution being observed; supplies the trajectories, models and - inputs. Its segments are 1:1 with this backend's. - name : str - Variable name, as registered on the models. + The solution being read, whose options the conversion uses. + key : str + The cache key: the variable name, or a vector field component's. + var_pybamm : :class:`pybamm.Symbol` + The variable's discretised expression. + inputs : dict + The segment's inputs, which fix the leaf's parameter layout. + ys_shape : tuple of int + The shape of the segment's states. Returns ------- - :class:`pybamm.solvers.base_processed_variable.BaseProcessedVariable` - Evaluable over this backend's whole run of segments. + tuple + The CasADi function, the expression it evaluates (a time integral's + integrand) and the time integral, or None. """ + leaf = self._casadi_leaves.get(key) + if leaf is not None: + return leaf, var_pybamm, None + leaf, var_pybamm, time_integral = solution._convert_to_casadi( + var_pybamm, inputs, ys_shape + ) + # A hit returns no time integral, so caching one would read its integrand + if time_integral is None: + self._casadi_leaves[key] = leaf + return leaf, var_pybamm, time_integral -class CasadiObservation(ObservationBackend): - """Variables converted to CasADi and evaluated by the IDAKLU kernels. - - Stateless: every per-segment artifact it needs is reached through the - Solution it is handed, so any two instances are interchangeable and one, - :data:`CASADI_OBSERVATION`, covers any run of segments. - """ - - def __getitem__(self, key): - return self - - def __eq__(self, other): - return type(other) is type(self) - - def __hash__(self): - return hash(type(self)) - - def build_variable(self, solution, name): - time_integral = None - pybamm.logger.debug(f"Post-processing {name}") +def build_variable(solution: pybamm.Solution, name: str) -> BaseProcessedVariable: + """The processed variable ``name`` of ``solution``, ready to evaluate. - # Iterate through all models, some may be in the list several times and - # therefore only get set up once - vars_pybamm = [ - model.get_processed_variable_or_event(name) for model in solution.all_models - ] - vars_casadi = [None] * len(solution.all_models) - for i, (model, ys, inputs) in enumerate( - zip(solution.all_models, solution.all_ys, solution.all_inputs, strict=True) - ): - _var_pybamm = vars_pybamm[i] - check_variable_in_solve(solution, name, _var_pybamm) - if isinstance(_var_pybamm, pybamm.VectorField): - comp_casadi = [] - for k, comp in enumerate(_var_pybamm.components): - cc, _, _ = self._model_leaf( - solution, - model, - comp, - inputs=inputs, - ys_shape=ys.shape, - time_integral=None, - cache_key=f"{name}[{k}]", - ) - comp_casadi.append(cc) - vars_casadi[i] = comp_casadi - else: - var_casadi, var_pybamm, time_integral = self._model_leaf( - solution, - model, - _var_pybamm, - inputs=inputs, - ys_shape=ys.shape, - time_integral=time_integral, - cache_key=name, - ) - vars_pybamm[i] = var_pybamm - vars_casadi[i] = var_casadi - return pybamm.process_variable( - name, - vars_pybamm, - CasadiObserver(vars_casadi), - solution, - time_integral=time_integral, - ) + Parameters + ---------- + solution : :class:`pybamm.Solution` + The solution to read; each segment is read through its own model. + name : str + Variable name, as registered on the models. - @staticmethod - def _model_leaf( - solution, - model, - var_pybamm, - time_integral, - inputs, - ys_shape, - cache_key, + Returns + ------- + :class:`pybamm.solvers.base_processed_variable.BaseProcessedVariable` + Evaluable over every segment of ``solution``. + """ + pybamm.logger.debug(f"Post-processing {name}") + time_integral = None + vars_pybamm = [ + model.get_processed_variable_or_event(name) for model in solution.all_models + ] + vars_casadi = [None] * len(solution.all_models) + for i, (model, ys, inputs) in enumerate( + zip(solution.all_models, solution.all_ys, solution.all_inputs, strict=True) ): - """One model's CasADi leaf, memoised on the model unless time-integrated.""" - _var_casadi = model._variables_casadi.get(cache_key) - if _var_casadi is not None: - return _var_casadi, var_pybamm, time_integral - - var_casadi, var_pybamm, time_integral = solution._convert_to_casadi( - var_pybamm, inputs, ys_shape - ) - - # Only cache if it's not a time integral - if time_integral is None: - model._variables_casadi[cache_key] = var_casadi - return var_casadi, var_pybamm, time_integral + var_pybamm = vars_pybamm[i] + check_variable_in_solve(solution, name, var_pybamm) + cache = ObserverCache.of(model) + if isinstance(var_pybamm, pybamm.VectorField): + vars_casadi[i] = [ + cache.casadi_leaf( + solution, f"{name}[{k}]", component, inputs, ys.shape + )[0] + for k, component in enumerate(var_pybamm.components) + ] + else: + vars_casadi[i], vars_pybamm[i], segment_time_integral = cache.casadi_leaf( + solution, name, var_pybamm, inputs, ys.shape + ) + time_integral = segment_time_integral or time_integral + return pybamm.process_variable( + name, + vars_pybamm, + CasadiObserver(vars_casadi), + solution, + time_integral=time_integral, + ) class OutputAssembly: @@ -322,35 +309,3 @@ def _variable_sensitivities( all_sens, ) return pack_sensitivity_dict(all_sens, sensitivity_names) - - -def join_observations(backends: Sequence[ObservationBackend]) -> ObservationBackend: - """One backend covering consecutive Solutions' segments, in order. - - Parameters - ---------- - backends : list of :class:`ObservationBackend` - The backend of each Solution being joined, in order. - - Returns - ------- - :class:`ObservationBackend` - The backend every Solution shares. - - Raises - ------ - :class:`pybamm.SolverError` - If the Solutions do not all share one backend. - """ - first, *rest = backends - for backend in rest: - if backend != first: - raise pybamm.SolverError( - "Cannot join solutions read through different observation " - f"backends: {type(first).__name__} and {type(backend).__name__}." - ) - return first - - -#: The backend a Solution carries by default. -CASADI_OBSERVATION = CasadiObservation() diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index a71ab570eb..f4afb37791 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -17,7 +17,7 @@ import pybamm from pybamm.codegen.compilation import aot_compile from pybamm.solvers.base_solver import flatten_inputs -from pybamm.solvers.observation import CASADI_OBSERVATION, join_observations +from pybamm.solvers.observation import build_variable class NumpyEncoder(json.JSONEncoder): @@ -345,7 +345,6 @@ class Solution(SolutionBase): "_initial_start_time", "_last_state", "_observable", - "_observation", "_options", "_sensitivities", "_sub_solutions", @@ -464,7 +463,6 @@ def __init__( self._all_inputs_casadi = None self._variables = {} - self._observation = CASADI_OBSERVATION # Sub-solutions concatenated into this one. Empty list means "just # self"; storing self here would create a refcount cycle that the @@ -506,7 +504,6 @@ def __setstate__(self, state: dict | tuple) -> None: self._y0 = None self._y0_sensitivities = None self._y_event_sensitivities = None - self._observation = CASADI_OBSERVATION if isinstance(state, dict): state = { f"_{name}" if name in ("first_state", "last_state") else name: value @@ -791,7 +788,6 @@ def _build_first_state(self): ) # stacked/casadi stay lazy; built from all_inputs[:1] on first access new_sol._sub_solutions = self.sub_solutions[:1] - new_sol._observation = self._observation[:1] new_sol.solve_time = 0 new_sol.integration_time = 0 @@ -846,7 +842,6 @@ def _build_last_state(self): ) # stacked/casadi stay lazy; built from all_inputs[-1:] on first access new_sol._sub_solutions = self.sub_solutions[-1:] - new_sol._observation = self._observation[-1:] new_sol.solve_time = 0 new_sol.integration_time = 0 new_sol.solver_statistics = _NO_SOLVER_STATISTICS @@ -892,7 +887,7 @@ def update(self, variables: str | list[str]): self._update_variable(variable) def _update_variable(self, name: str): - self._variables[name] = self._observation.build_variable(self, name) + self._variables[name] = build_variable(self, name) def observe(self, symbol: pybamm.Symbol) -> pybamm.ProcessedVariable: """ @@ -1320,9 +1315,6 @@ def __add__(self, other): # Set sub_solutions new_sol._sub_solutions = self.sub_solutions + other.sub_solutions - new_sol._observation = join_observations( - [self._observation, other._observation] - ) # update variables which were derived at the solver stage if any([self.variables_returned, other.variables_returned]): @@ -1426,7 +1418,6 @@ def from_sub_solutions(cls, sub_solutions): new_sol._y0_sensitivities = segments[0]._y0_sensitivities # leave stacked/casadi unset; built lazily from all_inputs (casadi is costly) new_sol._sub_solutions = sub_sols - new_sol._observation = join_observations([s._observation for s in segments]) for attr in ["solve_time", "integration_time", "set_up_time"]: vals = [getattr(s, attr, None) for s in segments] @@ -1477,7 +1468,6 @@ def copy(self): new_sol.integration_time = self.integration_time new_sol.solver_statistics = self.solver_statistics new_sol.set_up_time = self.set_up_time - new_sol._observation = self._observation # copy over variables which were derived at the solver stage if self._variables and all( diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py b/packages/pybamm/tests/unit/test_solvers/test_observation.py similarity index 65% rename from packages/pybamm/tests/unit/test_solvers/test_observation_backend.py rename to packages/pybamm/tests/unit/test_solvers/test_observation.py index 196a0cc2fc..d1871ee7d7 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_observation_backend.py +++ b/packages/pybamm/tests/unit/test_solvers/test_observation.py @@ -1,4 +1,4 @@ -"""Tests for observation backends and variable observers.""" +"""Tests for solution observation and variable observers.""" import pickle # nosec B403 - used in tests with trusted input from itertools import pairwise @@ -8,11 +8,7 @@ import pytest import pybamm -from pybamm.solvers.observation import ( - CASADI_OBSERVATION, - CasadiObservation, - join_observations, -) +from pybamm.solvers.observation import ObserverCache from pybamm.solvers.variable_observer import ( CasadiObserver, SegmentSelector, @@ -21,17 +17,6 @@ ) -class _RecordingObservation(CasadiObservation): - """A CasADi backend that records the segment slices taken of it.""" - - def __init__(self): - self.keys = [] - - def __getitem__(self, key): - self.keys.append(key) - return self - - @pytest.fixture(scope="module") def spm_solution(): model = pybamm.lithium_ion.SPM() @@ -54,67 +39,74 @@ def _split(solution, *boundaries): ] -class TestSegmentSlicing: - def test_the_casadi_backend_is_a_shared_stateless_value(self): - assert isinstance(CASADI_OBSERVATION, CasadiObservation) - assert CASADI_OBSERVATION[:1] is CASADI_OBSERVATION - assert CASADI_OBSERVATION[-1:] is CASADI_OBSERVATION +def _decay_model(factor): + """A discretised model whose "w" is ``factor`` times its decaying state.""" + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: -u} + model.initial_conditions = {u: 1} + model.variables = {"u": u, "w": factor * u} + pybamm.Discretisation().process_model(model) + return model - def test_casadi_backends_are_interchangeable(self): - # A Solution unpickled from a derived one carries its own instance. - restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) # nosec B301 - assert restored is not CASADI_OBSERVATION - assert restored == CASADI_OBSERVATION - assert hash(restored) == hash(CASADI_OBSERVATION) +class TestObserverCache: + def test_each_segment_reads_through_its_own_model(self): + doubled, tripled = _decay_model(2), _decay_model(3) + first = pybamm.IDAKLUSolver().solve(doubled, [0, 1]) + second = pybamm.IDAKLUSolver().solve(tripled, [1, 2]) -class TestJoin: - def test_runs_sharing_a_backend_join_to_it(self): - joined = join_observations([CASADI_OBSERVATION, CASADI_OBSERVATION]) - assert joined is CASADI_OBSERVATION + joined = first + second - def test_equal_backends_join_to_the_first(self): - restored = pickle.loads(pickle.dumps(CASADI_OBSERVATION)) # nosec B301 - assert join_observations([restored, CASADI_OBSERVATION]) is restored + np.testing.assert_allclose(joined["w"](0.5), 2 * first["u"](0.5), rtol=1e-6) + np.testing.assert_allclose(joined["w"](1.5), 3 * second["u"](1.5), rtol=1e-6) + leaves = joined["w"]._observer.leaves + assert leaves[0] is first["w"]._observer.leaves[0] + assert leaves[1] is second["w"]._observer.leaves[0] + assert leaves[0] is not leaves[1] - def test_different_backends_do_not_join(self): - with pytest.raises(pybamm.SolverError, match=r"different observation"): - join_observations([CASADI_OBSERVATION, _RecordingObservation()]) + def test_solving_a_model_again_reuses_its_leaves(self): + model = _decay_model(2) + solver = pybamm.IDAKLUSolver() + leaf = solver.solve(model, [0, 1])["w"]._observer.leaves[0] + assert solver.solve(model, [0, 2])["w"]._observer.leaves[0] is leaf -class TestObservationCarriedThroughDerivedSolutions: - @pytest.fixture - def recorded(self, spm_solution): + def test_derived_solutions_reuse_the_models_leaves(self, spm_solution): + name = "Voltage [V]" + leaf = spm_solution[name]._observer.leaves[0] first, second = _split(spm_solution, 5) - backend = _RecordingObservation() - first._observation = backend - second._observation = backend - return first, second, backend - - def test_first_and_last_state_slice_it(self, recorded): - first, _, backend = recorded - assert first.first_state._observation is backend - assert first.last_state._observation is backend - assert backend.keys == [slice(None, 1), slice(-1, None)] - - def test_addition_joins_it(self, recorded): - first, second, backend = recorded - assert (first + second)._observation is backend - - def test_from_sub_solutions_joins_it(self, recorded): - first, second, backend = recorded - joined = pybamm.Solution.from_sub_solutions([first, second]) - assert joined._observation is backend - - def test_copy_keeps_it(self, recorded): - first, _, backend = recorded - assert first.copy()._observation is backend - - def test_solutions_read_differently_cannot_be_added(self, recorded): - first, second, _ = recorded - second._observation = CASADI_OBSERVATION - with pytest.raises(pybamm.SolverError, match=r"different observation"): - first + second + + for derived in ( + spm_solution.first_state, + spm_solution.last_state, + spm_solution.copy(), + first + second, + pybamm.Solution.from_sub_solutions([first, second]), + ): + assert all(each is leaf for each in derived[name]._observer.leaves) + + def test_a_model_copy_starts_with_its_leaves_and_grows_apart(self): + model = _decay_model(2) + leaf = pybamm.IDAKLUSolver().solve(model, [0, 1])["w"]._observer.leaves[0] + model_copy = model.new_copy() + + solution = pybamm.IDAKLUSolver().solve(model_copy, [0, 1]) + assert solution["w"]._observer.leaves[0] is leaf + solution["u"] + assert "u" in ObserverCache.of(model_copy)._casadi_leaves + assert "u" not in ObserverCache.of(model)._casadi_leaves + + def test_a_model_pickled_before_the_cache_existed_builds_one(self): + model = _decay_model(2) + del model._observer_cache + + restored = pickle.loads(pickle.dumps(model)) # nosec B301 + + assert restored._observer_cache is None + solution = pybamm.IDAKLUSolver().solve(restored, [0, 1]) + np.testing.assert_allclose(solution["w"].entries, 2 * solution["u"].entries) + assert "w" in ObserverCache.of(restored)._casadi_leaves class TestSegmentSelector: diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index 150ff53273..b44b13abc6 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -17,39 +17,25 @@ from scipy.io import loadmat import pybamm -from pybamm.solvers.observation import CASADI_OBSERVATION from pybamm.solvers.solution import _DEFAULT_SOLUTION_OPTIONS, make_cycle_solution from tests import get_discretisation_for_testing -class _PicklerBeforeObservationBackends(pickle.Pickler): - """Pickles Solutions and ProcessedVariables in their earlier layout. - - That layout has no ``Solution._observation``, and keeps a processed - variable's CasADi functions as ``base_variables_casadi``. - """ +class _PicklerBeforeVariableObservers(pickle.Pickler): + """Pickles ProcessedVariables in their earlier layout, which keeps the CasADi + functions as ``base_variables_casadi``.""" def reducer_override(self, obj): - if isinstance(obj, pybamm.Solution): - slots = { - name: getattr(obj, name) - for cls in type(obj).__mro__ - for name in getattr(cls, "__slots__", ()) - if not name.startswith("__") and hasattr(obj, name) - } - slots.pop("_observation", None) - state = (None, slots) - elif isinstance(obj, pybamm.ProcessedVariable): - state = dict(obj.__dict__) - state["base_variables_casadi"] = state.pop("_observer").leaves - else: + if not isinstance(obj, pybamm.ProcessedVariable): return NotImplemented + state = dict(obj.__dict__) + state["base_variables_casadi"] = state.pop("_observer").leaves return copyreg.__newobj__, (type(obj),), state -def _round_trip_before_observation_backends(obj): +def _round_trip_before_variable_observers(obj): buffer = io.BytesIO() - _PicklerBeforeObservationBackends(buffer, pickle.HIGHEST_PROTOCOL).dump(obj) + _PicklerBeforeVariableObservers(buffer, pickle.HIGHEST_PROTOCOL).dump(obj) return pickle.loads(buffer.getvalue()) # nosec B301 @@ -1198,7 +1184,7 @@ def test_pickle_first_states_across_processes(self, tmp_path): assert "DATA" in save_result.stdout assert save_result.stdout == load_result.stdout - def test_load_solution_pickled_before_observation_backends(self): + def test_load_solution_pickled_before_variable_observers(self): sim = pybamm.Simulation( pybamm.lithium_ion.SPM(), experiment=pybamm.Experiment( @@ -1209,10 +1195,8 @@ def test_load_solution_pickled_before_observation_backends(self): t = np.linspace(0, 120, 7) solution["Voltage [V]"].entries - loaded = _round_trip_before_observation_backends(solution) + loaded = _round_trip_before_variable_observers(solution) - # An unpickled backend is a new instance, so this one is the default - assert loaded._observation is CASADI_OBSERVATION np.testing.assert_allclose( loaded["Voltage [V]"](t), solution["Voltage [V]"](t), rtol=1e-12 ) @@ -1247,7 +1231,7 @@ def test_load_processed_variable_pickled_before_observers(self): voltage = solution["Voltage [V]"] t = np.linspace(0, 600, 7) - loaded = _round_trip_before_observation_backends(voltage) + loaded = _round_trip_before_variable_observers(voltage) assert "base_variables_casadi" not in loaded.__dict__ np.testing.assert_allclose(loaded(t), voltage(t), rtol=1e-12) From b2bc7b7d44b6ec5e616308376b26686749dd7bf1 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Fri, 2 Oct 2026 16:45:17 +0100 Subject: [PATCH 4/7] perf: convert each time integral to CasADi once per model ObserverCache skipped time integrals: it cached only the CasADi leaf, so a hit lost the integral and would have read its integrand. Every read of a time integral therefore converted it again for each segment, 20 times for a 20-step experiment, and again for every solve of the model. The cache now keeps a leaf with the expression it evaluates and its time integral. The integral's analysis depends only on its expression, and nothing changes it after conversion, so solutions of one model can share it. Reading an ExplicitTimeIntegral from a 20-step experiment drops from 1.9 ms to 0.6 ms on first read and from 1.5 ms to 0.01 ms on later solves of the model, with identical values. --- CHANGELOG.md | 1 + .../pybamm/src/pybamm/solvers/observation.py | 21 ++++++-------- .../unit/test_solvers/test_observation.py | 28 +++++++++++++++++++ 3 files changed, 37 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 20836598b1..4f6c148419 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -48,6 +48,7 @@ - `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)) +- A time integral (`ExplicitTimeIntegral` or `DiscreteTimeSum`) is converted to CasADi once per model instead of once per solution segment: its first read from a 20-step experiment drops from 1.9 ms to 0.6 ms. ([#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 diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py index cbc8c4a113..3c52f6838a 100644 --- a/packages/pybamm/src/pybamm/solvers/observation.py +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -35,7 +35,8 @@ class ObserverCache: """ def __init__(self): - self._casadi_leaves: dict[str, casadi.Function] = {} + # Key -> (leaf, the expression it evaluates, time integral or None) + self._casadi_leaves: dict[str, tuple] = {} @classmethod def of(cls, model: pybamm.BaseModel) -> ObserverCache: @@ -81,16 +82,11 @@ def casadi_leaf( The CasADi function, the expression it evaluates (a time integral's integrand) and the time integral, or None. """ - leaf = self._casadi_leaves.get(key) - if leaf is not None: - return leaf, var_pybamm, None - leaf, var_pybamm, time_integral = solution._convert_to_casadi( - var_pybamm, inputs, ys_shape - ) - # A hit returns no time integral, so caching one would read its integrand - if time_integral is None: - self._casadi_leaves[key] = leaf - return leaf, var_pybamm, time_integral + entry = self._casadi_leaves.get(key) + if entry is None: + entry = solution._convert_to_casadi(var_pybamm, inputs, ys_shape) + self._casadi_leaves[key] = entry + return entry def build_variable(solution: pybamm.Solution, name: str) -> BaseProcessedVariable: @@ -128,10 +124,9 @@ def build_variable(solution: pybamm.Solution, name: str) -> BaseProcessedVariabl for k, component in enumerate(var_pybamm.components) ] else: - vars_casadi[i], vars_pybamm[i], segment_time_integral = cache.casadi_leaf( + vars_casadi[i], vars_pybamm[i], time_integral = cache.casadi_leaf( solution, name, var_pybamm, inputs, ys.shape ) - time_integral = segment_time_integral or time_integral return pybamm.process_variable( name, vars_pybamm, diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation.py b/packages/pybamm/tests/unit/test_solvers/test_observation.py index d1871ee7d7..5829a70040 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_observation.py +++ b/packages/pybamm/tests/unit/test_solvers/test_observation.py @@ -3,9 +3,11 @@ import pickle # nosec B403 - used in tests with trusted input from itertools import pairwise from types import SimpleNamespace +from unittest import mock import numpy as np import pytest +import scipy import pybamm from pybamm.solvers.observation import ObserverCache @@ -108,6 +110,32 @@ def test_a_model_pickled_before_the_cache_existed_builds_one(self): np.testing.assert_allclose(solution["w"].entries, 2 * solution["u"].entries) assert "w" in ObserverCache.of(restored)._casadi_leaves + def test_a_time_integral_is_converted_once_per_model(self): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: -u} + model.initial_conditions = {u: 1} + model.variables = { + "u": u, + "Q": pybamm.ExplicitTimeIntegral(u, pybamm.Scalar(1)), + "Q squared": pybamm.ExplicitTimeIntegral(u, pybamm.Scalar(0)) ** 2, + } + pybamm.Discretisation().process_model(model) + solver = pybamm.IDAKLUSolver() + convert = pybamm.Solution._convert_to_casadi + + with mock.patch.object( + pybamm.Solution, "_convert_to_casadi", autospec=True, side_effect=convert + ) as spy: + for t_end in (1, 2): + solution = solver.solve(model, [0, t_end]) + integral = scipy.integrate.trapezoid(solution["u"].entries, solution.t) + np.testing.assert_allclose(solution["Q"].entries, [1 + integral]) + np.testing.assert_allclose(solution["Q squared"].entries, [integral**2]) + + # u, Q and Q squared, each once across both solves + assert spy.call_count == 3 + class TestSegmentSelector: def test_full_range_keeps_every_nonempty_segment(self): From 05cda8327d465b14d52af5b7b899c2698e85559d Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Fri, 2 Oct 2026 16:48:39 +0100 Subject: [PATCH 5/7] docs: trim the changelog entries and comments Each changelog entry is now one line, leaving the detail to the PR, and the performance entry moves from Features to Optimizations. Eight two-line comments become one line, and four that restated the code they moved with are removed. VariableObserver holds its leaves rather than owning them: the model's ObserverCache shares them. --- CHANGELOG.md | 6 +++--- packages/pybamm/src/pybamm/solvers/idaklu_solver.py | 3 +-- packages/pybamm/src/pybamm/solvers/observation.py | 5 ++--- .../pybamm/src/pybamm/solvers/variable_observer.py | 12 +++--------- .../unit/test_solvers/test_observation_layout.py | 12 ++++-------- 5 files changed, 13 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4f6c148419..cd33c418e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +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)) -- The `base_variables_casadi` argument and attribute of `ProcessedVariable` are deprecated. Pass the CasADi functions as the third positional argument or as `observer=`, and call the variable to evaluate it rather than reading the functions back. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) +- `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 @@ -17,7 +17,6 @@ - A multi-input `IDAKLUSolver` solve now hands each thread its next input set as soon as it is free, instead of splitting the sets into equal static blocks plus a serial remainder solved on the calling thread, so a heterogeneous sweep (for example a current sweep whose high currents stop early on a voltage cut-off) no longer waits on its slowest block. A failing set no longer hides the others: the `SolverError` names every set that failed, whether it failed from the start or part-way through, each prefixed `input set N: `, and a single-set solve's message now starts with `input set 0: `. Each `on_failure="warn"` warning names its set the same way. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782)) - Added `Solution.sensitivity_names`, the sensitivity inputs in the column order of `sensitivities["all"]`. ([#5785](https://github.com/pybamm-team/PyBaMM/pull/5785)) - An outputs-only `pybammsolvers` solution also returns `y_init`, `yS_init` and `yS_term`: its initial states and its states' sensitivities at both ends. ([#5794](https://github.com/pybamm-team/PyBaMM/pull/5794)) -- Evaluating a `ProcessedVariable` again, for example at new times, no longer re-serialises its CasADi functions or rescans every sub-solution's time span: a repeat `solution["Voltage [V]"](t)` on an SPMe discharge is about 40% faster, and finding the sub-solutions that cover the requested times, across 60 of them, takes 1 µs rather than 10 µs. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) ## Bug fixes @@ -48,7 +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)) -- A time integral (`ExplicitTimeIntegral` or `DiscreteTimeSum`) is converted to CasADi once per model instead of once per solution segment: its first read from a 20-step experiment drops from 1.9 ms to 0.6 ms. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) +- 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 diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index f32a73c166..03d5711e4e 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -973,8 +973,7 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict, t_ev sol.yS_term, sensitivity_names ) - # On this path `sol.y` carries the concatenated outputs rather than the - # states, and `sol.yS` their sensitivities as (n_t, n_rows, n_p). + # 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), diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py index 3c52f6838a..ccb6b1be5d 100644 --- a/packages/pybamm/src/pybamm/solvers/observation.py +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -177,7 +177,7 @@ def __init__( self._sparse[name] = sparsity offsets = list(accumulate(lens, initial=0)) self._n_rows = offsets[-1] - # The layout itself: one slice of payload rows per output variable. + # One slice of payload rows per output variable self._rows = tuple(slice(start, end) for start, end in pairwise(offsets)) @property @@ -226,8 +226,7 @@ def attach( entries = np.ascontiguousarray(data[:, rows]) sparsity = self._sparse.get(name) if sparsity is not None: - # The rows hold only the structural nonzeros; scatter them into - # every entry at their flat (column-major) indices + # Scatter the returned structural nonzeros to their flat indices dense = np.zeros((entries.shape[0], sparsity.numel())) dense[:, sparsity.find()] = entries entries = dense diff --git a/packages/pybamm/src/pybamm/solvers/variable_observer.py b/packages/pybamm/src/pybamm/solvers/variable_observer.py index dcf8530e70..a753ecde58 100644 --- a/packages/pybamm/src/pybamm/solvers/variable_observer.py +++ b/packages/pybamm/src/pybamm/solvers/variable_observer.py @@ -1,6 +1,6 @@ """Per-variable observation strategies for :class:`pybamm.ProcessedVariable`. -A :class:`VariableObserver` owns one variable's per-sub-solution leaves and +A :class:`VariableObserver` holds one variable's per-sub-solution leaves and evaluates them on the solution grid, off-grid via cubic Hermite, and through the forward chain rule for sensitivities. """ @@ -97,8 +97,7 @@ def segments( See :meth:`SegmentSelector.select`. """ - # Keyed to the time arrays it was built from: an observer can be - # shared by variables of different solutions. + # An observer can be shared across solutions, so key on their time arrays if self._selector_ts is not variable.all_ts: self._selector = SegmentSelector(variable.all_ts) self._selector_ts = variable.all_ts @@ -213,7 +212,6 @@ def sensitivities(self, variable): *[sensitivity_inputs[name] for name in sensitivity_names] ) - # Set up symbolic variables t_casadi = casadi.MX.sym("t") y_casadi = casadi.MX.sym("y", ys.shape[0]) p_casadi = { @@ -223,8 +221,7 @@ def sensitivities(self, variable): p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) - # Non-target inputs can still appear in the tree (e.g. from experiment - # steps), so they stay concrete while the targets go symbolic. + # Symbolic for the sensitivity inputs only; the others keep their values inputs_for_casadi = {**inputs, **p_casadi} var_casadi = base_variable.to_casadi( @@ -233,7 +230,6 @@ def sensitivities(self, variable): dvar_dy = casadi.jacobian(var_casadi, y_casadi) dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) - # Convert to functions and evaluate index-by-index dvar_dy_func = casadi.Function( "dvar_dy", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dy] ) @@ -253,7 +249,6 @@ def sensitivities(self, variable): ] ) - # Compute sensitivity S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval all_S_var.append(S_var) @@ -380,7 +375,6 @@ def _find_ts_indices(starts, ends, t): indices = [] - # Get the minimum and maximum values of the target values `t` t_min, t_max = t[0], t[-1] # Binary search for the range of segments where t_min and t_max could lie diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py index f5cec510ee..e7e88acd92 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py @@ -27,8 +27,7 @@ def _dense_function(size): class TestLayoutContract: def test_0d_time_major(self, solution): - # One (n_times, output_size) array per sub-solution; for 0D - # output_size == 1 and data[t, 0] is the value at t. + # One (n_times, output_size) array per sub-solution base = [pybamm.StateVector(slice(0, 1))] n_t = len(solution.t) values = np.linspace(3.0, 4.2, n_t) @@ -39,8 +38,7 @@ def test_0d_time_major(self, solution): np.testing.assert_allclose(pvc.entries.reshape(-1), values) def test_1d_time_major_unrolls_to_space_by_time(self, solution): - # A 1D variable's (n_times, len_space) array unrolls to - # (len_space, n_times) via reshape((n_times, len_space)).transpose(). + # A 1D variable's (n_times, len_space) rows unroll to (len_space, n_times) base_pv = solution["X-averaged negative particle concentration [mol.m-3]"] var = base_pv.base_variables[0] len_space = var.shape[0] @@ -110,8 +108,7 @@ def test_sensitivities_are_named_and_flattened_per_parameter(self, solution): sensitivity_names=["a", "b"], ) - # Read the field, not the property: the property short-circuits to {} on - # this input-free SPM solve. + # The field: the property short-circuits to {} on this input-free solve attached = solution["Voltage [V]"]._sensitivities assert set(attached) == {"all", "a", "b"} assert attached["all"].shape == (n_t, 2) @@ -132,8 +129,7 @@ def test_sensitivities_of_the_wrong_shape_are_rejected(self, solution): ) def test_a_solve_without_sensitivities_leaves_an_empty_mapping(self, solution): - # Not None: an outputs-only solve retains no state to compute them from, - # so the answer is "there are none", not "ask again later". + # Empty, not None: an outputs-only solve keeps no state to compute them from assembly, data = self._assembly(solution, ["Voltage [V]"]) assembly.attach(solution, data) assert solution["Voltage [V]"]._sensitivities == {} From e19e540afd8e5e0f1bfdf0e015843e68d57d66c1 Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Thu, 8 Oct 2026 13:30:31 +0100 Subject: [PATCH 6/7] test: read an output variable with no structural nonzeros The solver returns no rows for such a variable, so its payload is (n_t, 0). Covers it as the only output and alongside a regular one. --- .../unit/test_solvers/test_idaklu_solver.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py index 587d89e955..df5cd88664 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -1249,6 +1249,24 @@ def test_with_sparse_output_variables_and_sensitivities(self): ): sim.solve([0, 100], inputs=input_parameters, calculate_sensitivities=True) + @pytest.mark.parametrize("output_variables", [["zero"], ["zero", "u"]]) + def test_output_variable_with_no_structural_nonzeros(self, output_variables): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: -u} + model.initial_conditions = {u: 1} + # The solver returns no rows for a variable with no structural nonzeros + model.variables = {"u": u, "zero": pybamm.Matrix(csc_matrix((1, 1)))} + t_interp = np.linspace(0, 1, 5) + + sol = pybamm.IDAKLUSolver(output_variables=output_variables).solve( + model, [0, 1], t_interp=t_interp + ) + + np.testing.assert_array_equal(sol["zero"](t_interp), np.zeros(5)) + if "u" in output_variables: + np.testing.assert_allclose(sol["u"](t_interp), np.exp(-t_interp), rtol=1e-3) + def test_output_variables_sensitivities_to_some_inputs(self): model = pybamm.BaseModel() c = pybamm.Variable("c") From 2d1f76d9a2ad38eebf6568f03fc575cfba8d3e4e Mon Sep 17 00:00:00 2001 From: bradyplanden Date: Thu, 8 Oct 2026 13:30:32 +0100 Subject: [PATCH 7/7] docs: say a cached observer leaf keeps its first solve's inputs and options --- packages/pybamm/src/pybamm/solvers/observation.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py index ccb6b1be5d..baa7d40e41 100644 --- a/packages/pybamm/src/pybamm/solvers/observation.py +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -31,7 +31,9 @@ class ObserverCache: """The compiled leaves of one model's variables, keyed by variable name. - A name is assumed to keep one expression for the model's lifetime. + A name is assumed to keep one expression for the model's lifetime. Each leaf + is built for the input layout and solution options (such as ``compile`` and + ``cse``) of the first solve that reads it, and later solves reuse it as is. """ def __init__(self):