diff --git a/CHANGELOG.md b/CHANGELOG.md index 1949337530..cd33c418e7 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)) +- `ProcessedVariable`'s `base_variables_casadi` argument and attribute are deprecated; pass the CasADi functions positionally or as `observer=`, and call the variable to evaluate it. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) ## Features @@ -46,6 +47,8 @@ - `Solution` stores its attributes in `__slots__`, so an experiment, which keeps four solutions per step, retains about 40% less memory per step on Python 3.13. Solutions pickled by earlier versions still load. ([#5827](https://github.com/pybamm-team/PyBaMM/pull/5827)) - Symbol mutation is no longer guarded by intercepting every attribute write, and `children` and `domain` return the stored immutable sequences instead of building a view on each access. Writes to a symbol's slots after construction are instead rejected by a static check of PyBaMM's code, tests, examples and documentation. Under the test suite's mutation guard, constructing an expression is twice as fast and parameterising it is 25% faster, which recovers most of the `test_parameterise` benchmark regression from [#5779](https://github.com/pybamm-team/PyBaMM/pull/5779). ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825)) - Faster reads of `output_variables` results: an off-grid `solution["Voltage [V]"](t)` on an SPM drops from about 420 µs to 1 µs. ([#5785](https://github.com/pybamm-team/PyBaMM/pull/5785)) +- Repeat evaluations of a `ProcessedVariable` are faster: a repeat `solution["Voltage [V]"](t)` on an SPMe discharge takes about 40% less time. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) +- Time integrals are converted to CasADi once per model instead of once per solution segment. ([#5786](https://github.com/pybamm-team/PyBaMM/pull/5786)) # [v26.9.0.0](https://github.com/pybamm-team/PyBaMM/tree/pybamm-v26.9.0.0) - 2026-09-28 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/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index 50994a3af8..03d5711e4e 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,17 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict, t_ev sol.yS_term, sensitivity_names ) - # Populate variables and sensitivities dictionaries directly - number_of_samples = sol.y.shape[0] // number_of_timesteps - sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) - sensitivity_params = ( - model.calculate_sensitivities if model.calculate_sensitivities else [] + # sol.y holds the outputs, and sol.yS their sensitivities (n_t, n_rows, n_p) + self._setup["output_assembly"].attach( + newsol, + np.asarray(sol.y).reshape(number_of_timesteps, -1), + sensitivities=( + np.asarray(sol.yS) if number_of_sensitivity_parameters else None + ), + sensitivity_names=sensitivity_names, ) - - start_idx = 0 - for var in self.output_variables: - var_nnz, var_shape, base_variables = self._get_variable_info(model, var) - end_idx = start_idx + var_nnz - data = sol.y[:, start_idx:end_idx] - if var_nnz != math.prod(var_shape): - # Scatter the returned structural nonzeros to their flat indices - indices = base_variables[0].sparsity_out(0).find() - dense = np.zeros((number_of_timesteps, math.prod(var_shape))) - dense[:, indices] = data - data = dense - time_integral = self._time_integral_vars.get(var) - values = data - - # handle any time integral variables - if time_integral is not None: - # time integral variables should all be 1D - values = time_integral.postfix(data.reshape(-1), sol.t, inputs_dict) - - newsol._variables[var] = pybamm.ProcessedVariableComputed( - [model.get_processed_variable_or_event(var)], - base_variables, - [values], - newsol, - time_integral=time_integral, - ) - - # Add sensitivities - newsol[var]._sensitivities = {} - if sensitivity_params: - if var_nnz != math.prod(var_shape): - raise pybamm.SolverError( - f"Sensitivity of sparse variables not supported. {var} is a sparse variable with number of non-zeros {var_nnz} and shape {var_shape}" - ) - sens_data = sol.yS[:, start_idx:end_idx, :] - sens_data = sens_data.reshape( - number_of_timesteps * (end_idx - start_idx), - number_of_sensitivity_parameters, - ) - if time_integral is not None: - sens_data = time_integral.postfix_sensitivities( - var, - data.reshape(-1), - sol.t, - inputs_dict, - sensitivity_names, - sens_data, - ) - newsol[var]._sensitivities["all"] = sens_data - - # Add the individual sensitivity - for i, name in enumerate(sensitivity_names): - sens = newsol[var]._sensitivities["all"][:, i : i + 1].reshape(-1) - newsol[var]._sensitivities[name] = sens - - start_idx += var_nnz return newsol - def _get_variable_info(self, model, var) -> tuple: - """Get variable length and base variables based on model format.""" - if model.convert_to_format == "casadi": - base_var = self._setup["var_fcns"][var] - sparsity = base_var.sparsity_out(0) - var_nnz = sparsity.nnz() - var_shape = sparsity.shape - return var_nnz, var_shape, [base_var] - else: # pragma: no cover - raise pybamm.SolverError( - f"Unsupported evaluation engine for convert_to_format=" - f"{model.convert_to_format}" - ) - def _set_consistent_initialization(self, model, time, inputs_list): """ Initialize y0 and ydot0 for the solver. In addition to calculating diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py new file mode 100644 index 0000000000..baa7d40e41 --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -0,0 +1,307 @@ +"""How a :class:`pybamm.Solution` reads its variables. + +: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 +populates the Solution eagerly. +""" + +from __future__ import annotations + +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 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. 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): + # Key -> (leaf, the expression it evaluates, time integral or None) + self._casadi_leaves: dict[str, tuple] = {} + + @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 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 + ------- + tuple + The CasADi function, the expression it evaluates (a time integral's + integrand) and the time integral, or None. + """ + 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: + """The processed variable ``name`` of ``solution``, ready to evaluate. + + 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. + + 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) + ): + 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], time_integral = cache.casadi_leaf( + solution, name, var_pybamm, inputs, ys.shape + ) + return pybamm.process_variable( + name, + vars_pybamm, + CasadiObserver(vars_casadi), + solution, + time_integral=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] + # 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: + # Scatter the returned structural nonzeros to their flat 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) 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..f4afb37791 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 build_variable class NumpyEncoder(json.JSONEncoder): @@ -885,79 +886,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] = build_variable(self, name) def observe(self, symbol: pybamm.Symbol) -> pybamm.ProcessedVariable: """ 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..a753ecde58 --- /dev/null +++ b/packages/pybamm/src/pybamm/solvers/variable_observer.py @@ -0,0 +1,394 @@ +"""Per-variable observation strategies for :class:`pybamm.ProcessedVariable`. + +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. +""" + +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`. + """ + # 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 + 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] + ) + + 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 the sensitivity inputs only; the others keep their values + 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) + + 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) + ] + ) + + 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 = [] + + 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..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") @@ -1421,6 +1439,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.py b/packages/pybamm/tests/unit/test_solvers/test_observation.py new file mode 100644 index 0000000000..5829a70040 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_observation.py @@ -0,0 +1,248 @@ +"""Tests for solution observation and variable observers.""" + +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 +from pybamm.solvers.variable_observer import ( + CasadiObserver, + SegmentSelector, + as_observer, + pack_sensitivity_dict, +) + + +@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) + ] + + +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 + + +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]) + + joined = first + second + + 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_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 + + 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) + + 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 + + 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): + 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)) # nosec B301 + + 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..e7e88acd92 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_observation_layout.py @@ -0,0 +1,155 @@ +"""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 + 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) 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] + 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"], + ) + + # 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) + 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): + # 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 == {} + + 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..b44b13abc6 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 @@ -20,6 +21,24 @@ from tests import get_discretisation_for_testing +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 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_variable_observers(obj): + buffer = io.BytesIO() + _PicklerBeforeVariableObservers(buffer, pickle.HIGHEST_PROTOCOL).dump(obj) + return pickle.loads(buffer.getvalue()) # nosec B301 + + class TestSolution: def test_init(self): t = np.linspace(0, 1) @@ -1165,6 +1184,81 @@ 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_variable_observers(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_variable_observers(solution) + + 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_variable_observers(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