diff --git a/CHANGELOG.md b/CHANGELOG.md index f4d68a666e..2c4f304934 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ - `model.convert_to_format = "rust"` evaluates a model's rhs, algebraic, Jacobian and event functions through compiled Rust evaluators, for `ScipySolver`, `AlgebraicSolver` and EIS. ([#5960](https://github.com/pybamm-team/PyBaMM/pull/5960)) - pybammsolvers gains an IDAKLU backend that evaluates a Rust `CompiledModel` through its C ABI, resolved at runtime so IDAKLU still loads without the Rust extension, with a Rust-backed Newton solver and output sensitivities projected natively in-solver. ([#5961](https://github.com/pybamm-team/PyBaMM/pull/5961)) - `IDAKLUSolver` solves `convert_to_format = "rust"` models through a Rust `CompiledModel`, one evaluator per thread, so `num_threads` parallelises input sets. `output_variables` support sensitivities and time integrals, and `NonlinearSolver` compiles its own model. ([#5962](https://github.com/pybamm-team/PyBaMM/pull/5962)) +- `IDAKLUSolver` observes rust-model solutions natively: variables and their sensitivities are evaluated on tapes compiled into the solve's Rust graph, with no CasADi function, and `Solution.update([...])` compiles a set of variables into one tape per segment. ([#5963](https://github.com/pybamm-team/PyBaMM/pull/5963)) ## Bug fixes diff --git a/packages/pybamm-rust/pybamm-python/src/model.rs b/packages/pybamm-rust/pybamm-python/src/model.rs index 41920d49c7..d82ac0a307 100644 --- a/packages/pybamm-rust/pybamm-python/src/model.rs +++ b/packages/pybamm-rust/pybamm-python/src/model.rs @@ -464,6 +464,12 @@ impl CompiledModel { .map(|f| f.clone_ref(py)) } + /// The retained derivation graph; roots lowered into it share the solve's indices. + #[getter] + fn graph(&self, py: Python<'_>) -> Py { + self.graph.clone_ref(py) + } + /// Pure df/dy (cj = 0, no mass): the bundle's composed `JacobianData`. #[getter] fn jacobian(&self, py: Python<'_>) -> PyResult> { diff --git a/packages/pybamm/src/pybamm/plotting/quick_plot.py b/packages/pybamm/src/pybamm/plotting/quick_plot.py index b2161b23d2..cc6f5a895f 100644 --- a/packages/pybamm/src/pybamm/plotting/quick_plot.py +++ b/packages/pybamm/src/pybamm/plotting/quick_plot.py @@ -300,6 +300,11 @@ def set_output_variables(self, output_variables, solutions): # Calculate subplot positions based on number of variables supplied self.subplot_positions = {} + # Every name below is read from every solution, so build them as one set + names = [var for variable_tuple in output_variables for var in variable_tuple] + for solution in solutions: + solution._build_missing(names) + for k, variable_tuple in enumerate(output_variables): # Prepare list of variables variables = [None] * len(solutions) diff --git a/packages/pybamm/src/pybamm/rust/_core.pyi b/packages/pybamm/src/pybamm/rust/_core.pyi index 542c9025f1..254fab3430 100644 --- a/packages/pybamm/src/pybamm/rust/_core.pyi +++ b/packages/pybamm/src/pybamm/rust/_core.pyi @@ -402,6 +402,10 @@ class CompiledModel: def rhs(self) -> CompiledFunction: """Primal ``f(t, y, p)`` as a shareable view.""" + @property + def graph(self) -> ExprGraph: + """The retained derivation arena (observation lowers new roots into it).""" + @property def jacobian(self) -> CompiledJacobian: """Pure ``df/dy`` (``cj = 0``, no mass), composed over the bundle's tapes.""" diff --git a/packages/pybamm/src/pybamm/solvers/base_solver.py b/packages/pybamm/src/pybamm/solvers/base_solver.py index 3029a7e7f7..58ad664a2e 100644 --- a/packages/pybamm/src/pybamm/solvers/base_solver.py +++ b/packages/pybamm/src/pybamm/solvers/base_solver.py @@ -67,6 +67,9 @@ class BaseSolver: # Set by solvers integrating a whole-model CompiledModel, to skip per-group lowering _integrates_via_compiled_model = False + # Whether set-up binds a rust model's CompiledModel to its ObserverCache + _observes_via_compiled_model = False + def __init__( self, method=None, diff --git a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py index ce84e6c695..3b1cec717b 100644 --- a/packages/pybamm/src/pybamm/solvers/idaklu_solver.py +++ b/packages/pybamm/src/pybamm/solvers/idaklu_solver.py @@ -18,7 +18,7 @@ stack_inputs, validate_rust_sensitivity_widths, ) -from pybamm.solvers.observation import OutputAssembly +from pybamm.solvers.observation import ObserverCache, OutputAssembly from pybamm.solvers.rust_lowering import RustModelLowering _UNSET = object() @@ -323,6 +323,7 @@ class IDAKLUSolver(pybamm.BaseSolver): """ _integrates_via_compiled_model = True + _observes_via_compiled_model = True class StateID(IntEnum): ALGEBRAIC = 0 @@ -570,6 +571,8 @@ def _set_up_rust(self, model, inputs_dict, y0): n_inputs = lowering.n_inputs rust_model = lowering.compile() lowering.bind_generic_evaluators(rust_model, root_method=self.root_method) + if self._observes_via_compiled_model: + ObserverCache.of(model).bind_rust_model(rust_model) colptrs, rowinds = rust_model.csc_sparsity_pattern() nnz = rust_model.nnz diff --git a/packages/pybamm/src/pybamm/solvers/observation.py b/packages/pybamm/src/pybamm/solvers/observation.py index 14ff17792c..44b1805c02 100644 --- a/packages/pybamm/src/pybamm/solvers/observation.py +++ b/packages/pybamm/src/pybamm/solvers/observation.py @@ -14,6 +14,7 @@ from collections.abc import Mapping, Sequence from itertools import accumulate, pairwise +from typing import cast import casadi import numpy as np @@ -23,6 +24,9 @@ from pybamm.solvers.base_processed_variable import BaseProcessedVariable from pybamm.solvers.variable_observer import ( CasadiObserver, + NativeObserver, + SharedOutputLeaf, + SharedOutputs, check_variable_in_solve, pack_sensitivity_dict, ) @@ -31,14 +35,25 @@ 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 + A name is assumed to keep one expression for the model's lifetime. Each CasADi + 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. + Native leaves compile into the graph bound by :meth:`bind_rust_model`. """ def __init__(self): # Key -> (leaf, the expression it evaluates, time integral or None) self._casadi_leaves: dict[str, tuple] = {} + self._rust_model = None + self._clear_native() + + def _clear_native(self) -> None: + self._native_leaves: dict[str, tuple] = {} + # Names, in row order -> one tape evaluating the whole set + self._native_sets: dict[tuple[str, ...], object] = {} + self._post_sum_leaves: dict[str, object] = {} + # Cached, as from_pybamm_var walks the whole tree + self._time_integrals: dict[str, pybamm.ProcessedVariableTimeIntegral] = {} @classmethod def of(cls, model: pybamm.BaseModel) -> ObserverCache: @@ -51,8 +66,96 @@ def copy(self) -> ObserverCache: """A cache that starts with these leaves and grows independently.""" new = type(self)() new._casadi_leaves = self._casadi_leaves.copy() + new._rust_model = self._rust_model + new._native_leaves = self._native_leaves.copy() + new._native_sets = self._native_sets.copy() + new._post_sum_leaves = self._post_sum_leaves.copy() + new._time_integrals = self._time_integrals.copy() return new + @property + def rust_model(self) -> object | None: + """The :class:`pybamm.rust.CompiledModel` native leaves compile into, or None.""" + return self._rust_model + + def bind_rust_model(self, rust_model) -> None: + """Compile native leaves into ``rust_model``'s graph from now on. + + Rebinding drops the previous model's native leaves. + """ + if rust_model is self._rust_model: + return + self._rust_model = rust_model + self._clear_native() + + def time_integral( + self, name: str, var_pybamm: pybamm.Symbol + ) -> pybamm.ProcessedVariableTimeIntegral | None: + """The time integral ``var_pybamm`` reads, or None, found on first use.""" + if name not in self._time_integrals: + self._time_integrals[name] = ( + pybamm.ProcessedVariableTimeIntegral.from_pybamm_var( + var_pybamm, self._rust_model.n_states + ) + ) + return self._time_integrals[name] + + def native_leaf( + self, name: str, var_pybamm: pybamm.Symbol + ) -> tuple[object, pybamm.Symbol, pybamm.ProcessedVariableTimeIntegral | None]: + """One segment's native leaf for ``var_pybamm``, compiled on first use.""" + entry = self._native_leaves.get(name) + if entry is None: + time_integral = self.time_integral(name, var_pybamm) + integrand = ( + var_pybamm if time_integral is None else time_integral.sum_node.child + ) + graph = self._rust_model.graph + leaf = graph.compile( + integrand.to_rust(graph), name=name, n_states=self._rust_model.n_states + ) + entry = (leaf, integrand, time_integral) + self._native_leaves[name] = entry + return entry + + def native_set(self, vars_pybamm: Mapping[str, pybamm.Symbol]) -> object: + """One tape over a set of variables, compiled on first use. + + The set shares one symbol memo, so common subexpressions are lowered once. + """ + key = tuple(vars_pybamm) + function = self._native_sets.get(key) + if function is None: + graph = self._rust_model.graph + rust_symbols = {} + function = graph.compile_outputs( + { + name: var.to_rust(graph, rust_symbols) + for name, var in vars_pybamm.items() + }, + n_states=self._rust_model.n_states, + ) + self._native_sets[key] = function + return function + + def post_sum_leaf( + self, name: str, time_integral: pybamm.ProcessedVariableTimeIntegral + ) -> object: + """A time integral's post-sum tape, whose states are the integral's value.""" + function = self._post_sum_leaves.get(name) + if function is None: + n_integral = time_integral.sum_node.evaluate_for_shape().shape[0] + # _post_sum_leaf only asks for an integral with a post-sum + post_sum_node = cast(pybamm.Symbol, time_integral.post_sum_node) + graph = self._rust_model.graph + function = graph.compile( + post_sum_node.to_rust(graph), + name="post_sum", + n_states=n_integral, + ) + self._post_sum_leaves[name] = function + return function + def casadi_leaf( self, solution: pybamm.Solution, @@ -111,6 +214,11 @@ def build_variable(solution: pybamm.Solution, name: str) -> BaseProcessedVariabl vars_pybamm = [ model.get_processed_variable_or_event(name) for model in solution.all_models ] + caches = _native_caches(solution) + if caches is not None and not any( + isinstance(var, pybamm.VectorField) for var in vars_pybamm + ): + return _build_native_variable(solution, name, vars_pybamm, caches) 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) @@ -138,6 +246,112 @@ def build_variable(solution: pybamm.Solution, name: str) -> BaseProcessedVariabl ) +def build_variables( + solution: pybamm.Solution, names: Sequence[str] +) -> dict[str, BaseProcessedVariable]: + """The processed variables ``names`` of ``solution``, built as one set. + + On a natively read solution, non-time-integral variables share one tape per + segment, so their common subexpressions are evaluated once. + """ + caches = _native_caches(solution) + if caches is None: + return {name: build_variable(solution, name) for name in names} + vars_pybamm = { + name: [ + model.get_processed_variable_or_event(name) for model in solution.all_models + ] + for name in names + } + shared = [ + name + for name in names + if not any( + isinstance(var, pybamm.VectorField) + or cache.time_integral(name, var) is not None + for var, cache in zip(vars_pybamm[name], caches, strict=True) + ) + ] + if len(shared) < 2: + return {name: build_variable(solution, name) for name in names} + variables = { + name: build_variable(solution, name) for name in names if name not in shared + } + variables.update(_build_native_set(solution, shared, vars_pybamm, caches)) + return {name: variables[name] for name in names} + + +def _native_caches(solution: pybamm.Solution) -> list[ObserverCache] | None: + """Each segment's cache, when every segment reads through a bound Rust model.""" + caches = [] + for model in solution.all_models: + cache = ObserverCache.of(model) + if model.convert_to_format != "rust" or cache.rust_model is None: + return None + caches.append(cache) + return caches + + +def _placeholder_states(solution: pybamm.Solution) -> list[int] | None: + """Each segment's state count when the solve stored no states, else None.""" + if not solution.variables_returned: + return None + return [model.len_rhs_and_alg for model in solution.all_models] + + +def _post_sum_leaf(cache, name, time_integral) -> object | None: + """The post-sum tape a time integral's sensitivities need, or None.""" + if time_integral is None or time_integral.post_sum_node is None: + return None + return cache.post_sum_leaf(name, time_integral) + + +def _build_native_variable( + solution, name, vars_pybamm, caches +) -> BaseProcessedVariable: + """A lazily evaluated variable whose leaves are native tapes.""" + time_integral = None + leaves = [None] * len(caches) + for i, cache in enumerate(caches): + check_variable_in_solve(solution, name, vars_pybamm[i]) + leaves[i], vars_pybamm[i], time_integral = cache.native_leaf( + name, vars_pybamm[i] + ) + observer = NativeObserver( + leaves, + _placeholder_states(solution), + _post_sum_leaf(caches[0], name, time_integral), + ) + return pybamm.process_variable( + name, vars_pybamm, observer, solution, time_integral=time_integral + ) + + +def _build_native_set( + solution, names, vars_pybamm, caches +) -> dict[str, BaseProcessedVariable]: + """Lazily evaluated variables reading one shared native tape per segment.""" + segments = [] + for i, cache in enumerate(caches): + for name in names: + check_variable_in_solve(solution, name, vars_pybamm[name][i]) + function = cache.native_set({name: vars_pybamm[name][i] for name in names}) + segments.append(SharedOutputs(function)) + placeholder_states = _placeholder_states(solution) + return { + name: pybamm.process_variable( + name, + vars_pybamm[name], + NativeObserver( + [SharedOutputLeaf(shared, name) for shared in segments], + placeholder_states, + ), + solution, + ) + for name in names + } + + class OutputAssembly: """The row layout of an ``output_variables`` payload, and its attachment. diff --git a/packages/pybamm/src/pybamm/solvers/processed_variable.py b/packages/pybamm/src/pybamm/solvers/processed_variable.py index fb6767f89e..0f77694b5d 100644 --- a/packages/pybamm/src/pybamm/solvers/processed_variable.py +++ b/packages/pybamm/src/pybamm/solvers/processed_variable.py @@ -508,9 +508,12 @@ def __init__( def _observe_postfix(self, entries, t): if self.time_integral is None: return entries - return self.time_integral.postfix( - entries, self.t_pts, self.all_inputs_stacked[0] - ) + # A CasADi post-sum takes stacked inputs; the symbolic fallback takes the dict + if self.time_integral.post_sum is None: + inputs = self.all_inputs[0] + else: + inputs = self.all_inputs_stacked[0] + return self.time_integral.postfix(entries, self.t_pts, inputs) def _interp_setup(self, entries, t): # save attributes for interpolation diff --git a/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py b/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py index d6ade6718c..6b27f117c3 100644 --- a/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py +++ b/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py @@ -14,13 +14,13 @@ @dataclass class ProcessedVariableTimeIntegral: method: Literal["discrete", "continuous"] - sum_node: pybamm.Symbol + sum_node: pybamm.ExplicitTimeIntegral | pybamm.DiscreteTimeSum initial_condition: npt.NDArray[np.float64] | float discrete_times: npt.NDArray[np.float64] | None post_sum_node: pybamm.Symbol | None = None post_sum: casadi.Function | None = None - def _sum_over_time(self, entries, t_pts) -> np.ndarray: + def sum_over_time(self, entries, t_pts) -> np.ndarray: """Sum or integrate the entries over the time points.""" if self.method == "discrete": return np.sum(entries, axis=0, keepdims=True) @@ -28,7 +28,7 @@ def _sum_over_time(self, entries, t_pts) -> np.ndarray: return np.array([trapezoid(entries, t_pts, axis=0)]) def postfix_sum(self, entries, t_pts) -> np.ndarray: - return self._sum_over_time(entries, t_pts) + float(self.initial_condition) + return self.sum_over_time(entries, t_pts) + float(self.initial_condition) def postfix(self, entries, t_pts, inputs) -> np.ndarray: """ @@ -45,6 +45,23 @@ def postfix(self, entries, t_pts, inputs) -> np.ndarray: ret = self.post_sum(0.0, the_integral, inputs).full().reshape(-1) return ret + def check_discrete_times(self, var_name, t_pts) -> None: + """Raise if a discrete-time sum's ``t_pts`` are not its discrete times.""" + if self.method != "discrete": + return + discrete_times = self.discrete_times + if ( + discrete_times is not None + and len(t_pts) == len(discrete_times) + and np.allclose(t_pts, discrete_times, atol=1e-10) + ): + return + raise pybamm.SolverError( + f'Processing discrete-time-sum variable "{var_name}": solution times ' + "and discrete times of the time integral are not equal. Set 't_interp=discrete_sum_times' to " + f"ensure the correct times are used.\nSolution times: {t_pts}\nDiscrete Sum times: {self.discrete_times}" + ) + def postfix_sensitivities( self, var_name, @@ -77,23 +94,10 @@ def postfix_sensitivities( np.ndarray The sensitivities of the postfix value. """ - # post fix for discrete time integral won't give correct result - # if ts are not equal to the discrete times. Raise error - # in this case - discrete_times = self.discrete_times - if self.method == "discrete" and not ( - discrete_times is not None - and len(t_pts) == len(discrete_times) - and np.allclose(t_pts, discrete_times, atol=1e-10) - ): - raise pybamm.SolverError( - f'Processing discrete-time-sum variable "{var_name}": solution times ' - "and discrete times of the time integral are not equal. Set 't_interp=discrete_sum_times' to " - f"ensure the correct times are used.\nSolution times: {t_pts}\nDiscrete Sum times: {self.discrete_times}" - ) + self.check_discrete_times(var_name, t_pts) # The initial condition is a constant, so it adds nothing to a sensitivity - integral_sensitivities = self._sum_over_time(sensitivities, t_pts) + integral_sensitivities = self.sum_over_time(sensitivities, t_pts) if self.post_sum_node is None: return integral_sensitivities diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index e8adb41781..cfdddddc29 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 build_variable +from pybamm.solvers.observation import build_variable, build_variables class NumpyEncoder(json.JSONEncoder): @@ -877,14 +877,25 @@ def update_summary_variables(self, all_summary_variables): ) def update(self, variables: str | list[str]): - """Add ProcessedVariables to the dictionary of variables in the solution""" - # Single variable + """Add ProcessedVariables to the dictionary of variables in the solution. + + Every listed variable is rebuilt; a list is built as one set, and stored + only once every variable in it has built. + """ if isinstance(variables, str): - variables = [variables] + self._update_variable(variables) + return + names = list(dict.fromkeys(variables)) + self._variables.update(build_variables(self, names)) - # Process - for variable in variables: - self._update_variable(variable) + def _build_missing(self, names: list[str]) -> None: + """Build, as one set, the variables in ``names`` not yet built. + + Built variables are kept: one computed during the solve has no states. + """ + missing = [name for name in dict.fromkeys(names) if name not in self._variables] + if missing: + self._variables.update(build_variables(self, missing)) def _update_variable(self, name: str): self._variables[name] = build_variable(self, name) @@ -1059,6 +1070,7 @@ def get_data_dict(self, variables=None, short_names=None, cycles_and_steps=True) else: if isinstance(variables, str): variables = [variables] + self._build_missing(variables) # otherwise, save only the variables specified data_long_names = {} for name in variables: diff --git a/packages/pybamm/src/pybamm/solvers/variable_observer.py b/packages/pybamm/src/pybamm/solvers/variable_observer.py index 615bd4b8b3..4176f8b0d1 100644 --- a/packages/pybamm/src/pybamm/solvers/variable_observer.py +++ b/packages/pybamm/src/pybamm/solvers/variable_observer.py @@ -8,8 +8,10 @@ from __future__ import annotations import bisect +import threading from abc import ABC, abstractmethod from collections.abc import Sequence +from typing import cast import casadi import numpy as np @@ -269,6 +271,256 @@ def sensitivities(self, variable): return pack_sensitivity_dict(S_var, sensitivity_names) +class NativeObserver(VariableObserver): + """Observation through compiled Rust tapes lowered into the retained graph. + + ``placeholder_states`` holds each segment's state count for outputs-only + solves, which store no states; ``None`` when states are real. + """ + + def __init__(self, leaves, placeholder_states=None, post_sum_leaf=None) -> None: + self._leaves = leaves + self._placeholder_states = placeholder_states + self._post_sum_leaf = post_sum_leaf + + @property + def leaves(self) -> list: + return self._leaves + + def _states(self, variable, idxs) -> list[npt.NDArray[np.float64]]: + """The states of segments ``idxs``; zeros of that shape if none are stored.""" + if self._placeholder_states is None: + return [variable.all_ys[idx] for idx in idxs] + return [ + np.zeros((self._placeholder_states[idx], variable.all_ts[idx].size)) + for idx in idxs + ] + + def _setup( + self, variable, t, full_range + ) -> tuple[list, list, list | None, list, list]: + """Per-segment ``(ts, ys, yps, inputs, leaves)``; inputs stay dicts.""" + idxs = self.segments(variable, t, full_range) + ts = [variable.all_ts[idx] for idx in idxs] + ys = self._states(variable, idxs) + yps = ( + [variable.all_yps[idx] for idx in idxs] + if variable.hermite_interpolation + else None + ) + inputs = [variable.all_inputs[idx] for idx in idxs] + leaves = [self._leaves[idx] for idx in idxs] + return ts, ys, yps, inputs, leaves + + def observe_raw(self, variable) -> npt.NDArray[np.float64]: + # eval_trajectory returns (output_len, n_t), so reshape in order="F" + t = variable.t_pts + ts, ys, _, inputs, leaves = self._setup(variable, t, full_range=True) + cols = [ + np.asarray(leaf.eval_trajectory(t_i, y_i, inp_i)) + for leaf, t_i, y_i, inp_i in zip(leaves, ts, ys, inputs, strict=True) + ] + return np.concatenate(cols, axis=1).reshape(variable._shape(t), order="F") + + def observe_hermite(self, variable, t) -> npt.NDArray[np.float64]: + # Route query times to segments as observe.cpp's knot-window scan does + ts, ys, yps, inputs, leaves = self._setup(variable, t, full_range=False) + # Only a solution that stored y' is Hermite-interpolated + yps = cast(list, yps) + n_segments = len(ts) + cols = [] + i = 0 + n = len(t) + for seg_idx, (leaf, t_i, y_i, yp_i, inp_i) in enumerate( + zip(leaves, ts, ys, yps, inputs, strict=True) + ): + if i >= n: + break + is_last_segment = seg_idx == n_segments - 1 + if is_last_segment: + j = n + else: + j = i + np.searchsorted(t[i:], t_i[-1], side="right") + if j <= i: + continue + # Copy: the caller may mutate t, and a shared leaf holds its arguments + query = np.array(t[i:j]) + if t_i.size < 2: + # A single knot: hold its state constant and evaluate directly + y_query = np.repeat(y_i, len(query), axis=1) + cols.append(np.asarray(leaf.eval_trajectory(query, y_query, inp_i))) + else: + cols.append( + np.asarray( + leaf.eval_trajectory_hermite(query, t_i, y_i, yp_i, inp_i) + ) + ) + i = j + return np.concatenate(cols, axis=1).reshape(variable._shape(t), order="F") + + def sensitivities(self, variable) -> dict[str, npt.NDArray[np.float64]]: + sensitivity_names = variable.sensitivity_names + all_ys = self._states(variable, range(len(self._leaves))) + all_dy_dp = variable.all_solution_sensitivities["all"] + if self._placeholder_states is not None: + # Placeholder states are constants, so they carry no sensitivity + all_dy_dp = [np.zeros((ys.size, len(sensitivity_names))) for ys in all_ys] + segments = list( + zip(self._leaves, variable.all_ts, all_ys, variable.all_inputs, strict=True) + ) + sensitivities = native_sensitivities(segments, all_dy_dp, sensitivity_names) + if variable.time_integral is not None: + # The value sums over every segment at once, so its sensitivity must too + sensitivities = native_postfix_sensitivities( + variable.time_integral, + variable._name, + variable.t_pts, + self.observe_raw(variable), + sensitivities, + variable.all_inputs[0], + sensitivity_names, + self._post_sum_leaf, + ) + return pack_sensitivity_dict(sensitivities, sensitivity_names) + + +def _row_spans(function) -> dict[str, slice]: + """``{name: slice}`` over the rows of a multi-output compiled tape.""" + spans = {} + offset = 0 + for name, length in zip(function.output_names, function.output_lens, strict=True): + spans[name] = slice(offset, offset + length) + offset += length + return spans + + +def _same_argument(previous, current) -> bool: + """Whether two call arguments select the same evaluation; unknowns differ.""" + if previous is current: + return True + if isinstance(previous, np.ndarray) and isinstance(current, np.ndarray): + return previous.shape == current.shape and np.array_equal(previous, current) + if isinstance(previous, dict) and isinstance(current, dict): + return previous.keys() == current.keys() and all( + _same_argument(previous[key], current[key]) for key in previous + ) + if isinstance(previous, list | tuple) and isinstance(current, list | tuple): + return len(previous) == len(current) and all( + _same_argument(was, now) for was, now in zip(previous, current, strict=True) + ) + if np.isscalar(previous) and np.isscalar(current): + return previous == current + return False + + +class SharedOutputs: + """One segment's tape over several variables, evaluated once for all of them. + + A result is held until every leaf has read its rows. Held arguments must be + arrays PyBaMM owns and never mutates; copy a caller's array first. Concurrent + reads are safe: each caller gets the result of its own arguments. + """ + + def __init__(self, function) -> None: + self.function = function + self._spans = _row_spans(function) + self._lock = threading.Lock() + self._call = None + self._result = None + self._unread = 0 + + def __getstate__(self) -> dict: + """Pickle the tape, never a held result.""" + return {"function": self.function} + + def __setstate__(self, state) -> None: + self.__init__(state["function"]) + + def span(self, name) -> slice: + """The rows ``name`` owns in this tape's output.""" + return self._spans[name] + + def _hold(self, call, compute) -> npt.NDArray[np.float64]: + """``compute()``'s result, reused while an identical call is still held.""" + with self._lock: + if self._call is not None and _same_argument(self._call, call): + # _result is set and cleared together with _call + result = cast(npt.NDArray[np.float64], self._result) + self._keep_for(self._unread - 1) + return result + # Unlocked, so calls on other arguments run in parallel; each keeps its own + result = compute() + with self._lock: + self._call = call + self._result = result + self._keep_for(len(self._spans) - 1) + return result + + def _keep_for(self, unread: int) -> None: + """Hold the current call for ``unread`` more reads, releasing it at zero.""" + self._unread = unread + if unread <= 0: + self._call = None + self._result = None + + def evaluate(self, method, args) -> npt.NDArray[np.float64]: + """``method(*args)`` on the shared tape, reusing an identical held call.""" + return self._hold( + (method, args), + lambda: np.asarray(getattr(self.function, method)(*args)), + ) + + def chain_rule(self, ts, ys, inputs, dy_dp, sens_names) -> npt.NDArray[np.float64]: + """The whole set's ``dvar/dp`` for one sub-solution, swept once. + + Shape ``(n_t * function.output_len, n_p)``, time-outer / output-inner. + """ + args = (ts, ys, inputs, dy_dp, sens_names) + return self._hold( + ("chain_rule", args), lambda: _sweep_sensitivities(self.function, *args) + ) + + +class SharedOutputLeaf: + """One variable's rows of a :class:`SharedOutputs` tape, used as its own tape.""" + + def __init__(self, shared, name) -> None: + self._shared = shared + self._name = name + self._rows = shared.span(name) + + @property + def output_len(self) -> int: + return self._rows.stop - self._rows.start + + @property + def input_names(self) -> tuple[str, ...]: + return self._shared.function.input_names + + @property + def n_inputs(self) -> int: + return self._shared.function.n_inputs + + def eval_trajectory(self, ts, ys, inputs) -> npt.NDArray[np.float64]: + return self._shared.evaluate("eval_trajectory", (ts, ys, inputs))[self._rows] + + def eval_trajectory_hermite( + self, t_query, ts, ys, yps, inputs + ) -> npt.NDArray[np.float64]: + return self._shared.evaluate( + "eval_trajectory_hermite", (t_query, ts, ys, yps, inputs) + )[self._rows] + + def chain_rule(self, ts, ys, inputs, dy_dp, sens_names) -> npt.NDArray[np.float64]: + """This variable's rows of the set's chain-rule sweep.""" + full = self._shared.chain_rule(ts, ys, inputs, dy_dp, sens_names) + n_p = len(sens_names) + # The tape's width, not -1, which an empty segment leaves undetermined + n_rows = self._shared.function.output_len + rows = full.reshape(len(ts), n_rows, n_p)[:, self._rows, :] + return rows.reshape(len(ts) * self.output_len, n_p) + + def as_observer(leaves: VariableObserver | list) -> VariableObserver: """Coerce leaves to an observer. @@ -318,6 +570,91 @@ def check_variable_in_solve( ) +def native_sensitivities( + segments, dy_dp_segments, sensitivity_names +) -> npt.NDArray[np.float64]: + """Forward sensitivities ``(N, n_p)`` of a natively observed variable. + + ``segments`` holds one ``(leaf, ts, ys, inputs)`` per sub-solution, stacked in + solve order. + """ + return np.vstack( + [ + chain_rule_sensitivities(leaf, ts, ys, inputs, dy_dp, sensitivity_names) + for (leaf, ts, ys, inputs), dy_dp in zip( + segments, dy_dp_segments, strict=True + ) + ] + ) + + +def native_postfix_sensitivities( + time_integral, + var_name, + t_pts, + entries, + sensitivities, + inputs, + sensitivity_names, + post_sum_leaf, +) -> npt.NDArray[np.float64]: + """Sensitivities ``(n_integral, n_p)`` of a time integral's value, natively.""" + time_integral.check_discrete_times(var_name, t_pts) + + # The initial condition is a constant, so it adds nothing to a sensitivity + integral_sensitivities = time_integral.sum_over_time(sensitivities, t_pts) + if time_integral.post_sum_node is None: + return integral_sensitivities + + integral = np.ascontiguousarray( + time_integral.postfix_sum(entries, t_pts), dtype=np.float64 + ).ravel() + p_stacked = post_sum_leaf.pack(inputs) + dpost_dy = post_sum_leaf.jacobian("y")(0.0, integral, p_stacked).toarray() + dpost_dp_all = post_sum_leaf.jacobian("p")(0.0, integral, p_stacked).toarray() + input_names = list(post_sum_leaf.input_names) + dpost_dp = np.zeros((dpost_dy.shape[0], len(sensitivity_names))) + for k, name in enumerate(sensitivity_names): + if name in input_names: + dpost_dp[:, k] = dpost_dp_all[:, input_names.index(name)] + return dpost_dy @ integral_sensitivities + dpost_dp + + +def chain_rule_sensitivities( + cf, ts, ys, inputs, dy_dp, sens_names +) -> npt.NDArray[np.float64]: + """Variable sensitivities ``dvar/dp`` for one sub-solution via jvp_trajectory. + + ``dy_dp`` is ``(n_t * n_states, n_p)`` and the result ``(n_t * output_len, n_p)``, + both time-outer, with columns ordered as ``sens_names``. + """ + if isinstance(cf, SharedOutputLeaf): + return cf.chain_rule(ts, ys, inputs, dy_dp, sens_names) + return _sweep_sensitivities(cf, ts, ys, inputs, dy_dp, sens_names) + + +def _sweep_sensitivities( + cf, ts, ys, inputs, dy_dp, sens_names +) -> npt.NDArray[np.float64]: + """One forward sweep per sensitivity parameter over ``cf``'s whole output.""" + n_states, n_t = ys.shape + out_len = cf.output_len + input_names = cf.input_names + sensitivities = np.empty((n_t * out_len, len(sens_names))) + for k, name in enumerate(sens_names): + vy_k = np.ascontiguousarray(dy_dp[:, k].reshape(n_t, n_states).T) + if name in input_names: + vp = np.zeros(cf.n_inputs) + vp[input_names.index(name)] = 1.0 + col = cf.jvp_trajectory(ts, ys, inputs, vy_k, vp=vp) + else: + # variable has no direct dependence on this parameter; dvar_dp == 0 + col = cf.jvp_trajectory(ts, ys, inputs, vy_k) + # (out_len, n_t) -> time-outer / output-inner flat column + sensitivities[:, k] = np.asarray(col).T.reshape(-1) + return sensitivities + + def pack_sensitivity_dict( sensitivity_matrix: npt.NDArray[np.float64], sensitivity_names: Sequence[str] ) -> dict[str, npt.NDArray[np.float64]]: diff --git a/packages/pybamm/tests/integration/test_rust_observation.py b/packages/pybamm/tests/integration/test_rust_observation.py new file mode 100644 index 0000000000..f317e91940 --- /dev/null +++ b/packages/pybamm/tests/integration/test_rust_observation.py @@ -0,0 +1,356 @@ +"""Parity: native (Rust) observation vs CasADi observation of a rust model.""" + +import numpy as np +import pytest + +import pybamm +from pybamm.solvers.observation import _native_caches + + +def _solve(solver_cls, native, model_cls=pybamm.lithium_ion.SPM, **opts): + model = model_cls() + model.convert_to_format = "rust" + solver = solver_cls(**opts) + # The instance attribute survives the shallow solver.copy() Simulation performs. + solver._observes_via_compiled_model = native + sim = pybamm.Simulation(model, solver=solver) + return sim.solve([0, 3600]) + + +def _native_idaklu(native=True): + """An IDAKLUSolver observing natively (``native=True``) or through CasADi.""" + solver = pybamm.IDAKLUSolver() + solver._observes_via_compiled_model = native + return solver + + +SOLVERS = [ + ("idaklu", pybamm.IDAKLUSolver, {}), +] +NATIVE_IDS = [s[0] for s in SOLVERS] + +MODELS = [ + ("spm", pybamm.lithium_ion.SPM), + ("spme", pybamm.lithium_ion.SPMe), +] +MODEL_IDS = [m[0] for m in MODELS] + +# 0D scalars plus 1D fields on two distinct axes (r and x). +DIRECT_VARS = [ + "Terminal voltage [V]", + "Current [A]", + "Negative particle concentration [mol.m-3]", + "Electrolyte concentration [mol.m-3]", +] + + +class TestDirectVariableParity: + @pytest.mark.parametrize( + "solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS + ) + @pytest.mark.parametrize("model_name,model_cls", MODELS, ids=MODEL_IDS) + @pytest.mark.parametrize("var", DIRECT_VARS) + def test_direct_variable_parity( + self, solver_name, solver_cls, native_opts, model_name, model_cls, var + ): + native_sol = _solve(solver_cls, native=True, model_cls=model_cls, **native_opts) + casadi_sol = _solve(solver_cls, native=False, model_cls=model_cls) + np.testing.assert_allclose( + native_sol[var].entries, casadi_sol[var].entries, rtol=1e-9, atol=1e-11 + ) + + +class TestZeroCopyStateTrajectory: + """Non-F-contiguous `all_ys` would make every native observe copy the states.""" + + @pytest.mark.parametrize( + "solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS + ) + def test_full_state_trajectory_is_f_contiguous( + self, solver_name, solver_cls, native_opts + ): + sol = _solve(solver_cls, native=True, **native_opts) + assert _native_caches(sol) is not None + for ys in sol.all_ys: + arr = np.asarray(ys) + assert arr.ndim == 2 + assert arr.flags["F_CONTIGUOUS"], ( + f"{solver_name}: all_ys segment is not F-contiguous " + f"(shape={arr.shape}, strides={arr.strides}) — eval_trajectory " + "would gather-copy the whole state matrix on every observe." + ) + + def test_composed_trajectory_segments_are_f_contiguous(self): + # Each stepped segment is consumed by its own eval_trajectory. + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=_native_idaklu()) + for _ in range(3): + sim.step(360) + for ys in sim.solution.all_ys: + arr = np.asarray(ys) + assert arr.flags["F_CONTIGUOUS"], ( + f"composed all_ys segment is not F-contiguous " + f"(shape={arr.shape}, strides={arr.strides})" + ) + + +class TestTimeIntegralParity: + @pytest.mark.parametrize( + "solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS + ) + def test_time_integral_parity(self, solver_name, solver_cls, native_opts): + native_sol = _solve(solver_cls, native=True, **native_opts) + casadi_sol = _solve(solver_cls, native=False) + np.testing.assert_allclose( + native_sol["Discharge capacity [A.h]"].entries, + casadi_sol["Discharge capacity [A.h]"].entries, + rtol=1e-8, + atol=1e-10, + ) + + @pytest.mark.parametrize( + "solver_name,solver_cls,native_opts", SOLVERS, ids=NATIVE_IDS + ) + def test_explicit_time_integral_parity(self, solver_name, solver_cls, native_opts): + # Unlike "Discharge capacity [A.h]", this drives the native time-integral read. + def _build(native): + model = pybamm.lithium_ion.SPM() + model.variables["Integrated current [A.s]"] = pybamm.ExplicitTimeIntegral( + model.variables["Current [A]"], pybamm.Scalar(0.0) + ) + solver = solver_cls(**native_opts) + model.convert_to_format = "rust" + solver._observes_via_compiled_model = native + return pybamm.Simulation(model, solver=solver).solve([0, 3600]) + + native_sol = _build(native=True) + casadi_sol = _build(native=False) + assert _native_caches(native_sol) is not None + native_entries = native_sol["Integrated current [A.s]"].entries + assert native_entries.shape == (1,), ( + f"Expected scalar shape (1,) for time-integral output, got {native_entries.shape}" + ) + np.testing.assert_allclose( + native_entries, + casadi_sol["Integrated current [A.s]"].entries, + rtol=1e-8, + atol=1e-10, + ) + + +class TestSteppedSolutionParity: + def test_stepping_addition_parity(self): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=_native_idaklu()) + for _ in range(5): + sim.step(360) + native_v = sim.solution["Terminal voltage [V]"].entries + + cmodel = pybamm.lithium_ion.SPM() + cmodel.convert_to_format = "rust" + csim = pybamm.Simulation(cmodel, solver=_native_idaklu(native=False)) + for _ in range(5): + csim.step(360) + casadi_v = csim.solution["Terminal voltage [V]"].entries + np.testing.assert_allclose(native_v, casadi_v, rtol=1e-8, atol=1e-10) + + def test_first_last_state_parity(self): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation(model, solver=_native_idaklu()) + sol = sim.solve([0, 3600]) + # Observation must still work through first/last_state, which carry the context. + assert sol.first_state["Terminal voltage [V]"].entries.size >= 1 + assert sol.last_state["Terminal voltage [V]"].entries.size >= 1 + + def test_experiment_parity(self): + experiment = pybamm.Experiment( + ["Discharge at 1C until 3.0 V", "Charge at 1C until 4.0 V"] + ) + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + sim = pybamm.Simulation( + model, + experiment=experiment, + solver=_native_idaklu(), + ) + sol = sim.solve() + assert _native_caches(sol) is not None, ( + "Experiment solution does not read natively" + ) + native_v = sol["Terminal voltage [V]"].entries + + cmodel = pybamm.lithium_ion.SPM() + cmodel.convert_to_format = "rust" + csim = pybamm.Simulation( + cmodel, experiment=experiment, solver=_native_idaklu(native=False) + ) + csol = csim.solve() + np.testing.assert_allclose( + native_v, csol["Terminal voltage [V]"].entries, rtol=1e-7, atol=1e-9 + ) + + +def test_multistep_time_integral_parity(): + # A time integral over a multi-segment solution integrates once, not per segment. + def _build(native): + model = pybamm.lithium_ion.SPM() + model.variables["Integrated current [A.s]"] = pybamm.ExplicitTimeIntegral( + model.variables["Current [A]"], pybamm.Scalar(0.0) + ) + solver = pybamm.IDAKLUSolver() + model.convert_to_format = "rust" + solver._observes_via_compiled_model = native + sim = pybamm.Simulation(model, solver=solver) + for _ in range(3): + sim.step(360) + return sim.solution + + native_sol = _build(native=True) + casadi_sol = _build(native=False) + assert _native_caches(native_sol) is not None + nv = native_sol["Integrated current [A.s]"].entries + cv = casadi_sol["Integrated current [A.s]"].entries + assert nv.shape == (1,), f"expected single scalar, got shape {nv.shape}" + np.testing.assert_allclose(nv, cv, rtol=1e-8, atol=1e-10) + + +class TestDiscreteTimeSumParity: + def test_discrete_time_sum_parity(self): + # Both paths sum over solution.t, so pin it to the discrete times. + discrete_times = np.linspace(0.0, 3600.0, 11) + data = pybamm.DiscreteTimeData(discrete_times, np.zeros(11), "td") + + def _build(native): + model = pybamm.lithium_ion.SPM() + model.variables["DT sum"] = pybamm.DiscreteTimeSum( + model.variables["Voltage [V]"] - data + ) + solver = pybamm.IDAKLUSolver() + model.convert_to_format = "rust" + solver._observes_via_compiled_model = native + return pybamm.Simulation(model, solver=solver).solve( + t_eval=[0, 3600], t_interp=discrete_times + ) + + native_sol = _build(native=True) + casadi_sol = _build(native=False) + assert _native_caches(native_sol) is not None + nv = native_sol["DT sum"].entries + cv = casadi_sol["DT sum"].entries + assert nv.shape == (1,), f"expected a single scalar, got shape {nv.shape}" + np.testing.assert_allclose(nv, cv, rtol=1e-8, atol=1e-10) + + +class TestNativeInputParity: + def test_native_observed_variable_sensitivities_require_calculate_sensitivities( + self, + ): + def _build(calculate_sensitivities): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + param = model.default_parameter_values + param["Current function [A]"] = pybamm.InputParameter("I") + sim = pybamm.Simulation( + model, parameter_values=param, solver=_native_idaklu() + ) + return sim.solve( + [0, 3600], + inputs={"I": 0.5}, + calculate_sensitivities=calculate_sensitivities, + ) + + sol = _build(calculate_sensitivities=False) + assert _native_caches(sol) is not None + v = sol["Terminal voltage [V]"] + assert len(sol.all_inputs[0]) > 0 + with pytest.raises(ValueError, match=r"calculate_sensitivities"): + _ = v.sensitivities + + sens_sol = _build(calculate_sensitivities=True) + assert sens_sol["Terminal voltage [V]"].sensitivities != {} + + @pytest.mark.parametrize("input_val", [1, np.int64(1), np.float32(1.0)]) + def test_native_observation_nonfloat_input_parity(self, input_val): + def build(native): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + param = model.default_parameter_values + param["Current function [A]"] = pybamm.InputParameter("I") + solver = pybamm.IDAKLUSolver() + solver._observes_via_compiled_model = native + sim = pybamm.Simulation(model, parameter_values=param, solver=solver) + return sim.solve([0, 3600], inputs={"I": input_val}) + + native_sol = build(native=True) + casadi_sol = build(native=False) + assert _native_caches(native_sol) is not None + np.testing.assert_allclose( + native_sol["Terminal voltage [V]"].entries, + casadi_sol["Terminal voltage [V]"].entries, + rtol=1e-8, + atol=1e-10, + ) + + +@pytest.mark.parametrize("solver_cls", [pybamm.IDAKLUSolver]) +def test_dense_trajectory_parity(solver_cls): + t_interp = np.linspace(0, 3600, 500) + + def build(native): + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "rust" + solver = solver_cls() + solver._observes_via_compiled_model = native + return pybamm.Simulation(model, solver=solver).solve( + t_eval=[0, 3600], t_interp=t_interp + ) + + native_sol = build(native=True) + casadi_sol = build(native=False) + assert _native_caches(native_sol) is not None + assert native_sol["Terminal voltage [V]"].entries.shape[-1] == len(t_interp) + for var in ("Terminal voltage [V]", "Electrolyte concentration [mol.m-3]"): + np.testing.assert_allclose( + native_sol[var].entries, casadi_sol[var].entries, rtol=1e-8, atol=1e-10 + ) + + +class TestFieldLayoutParity: + def test_2d_variable_native_parity(self): + # Locks the native column-major -> .T -> reshape contract for a 2D (r, x) field. + var = "Negative particle concentration [mol.m-3]" + native_sol = _solve( + pybamm.IDAKLUSolver, native=True, model_cls=pybamm.lithium_ion.DFN + ) + casadi_sol = _solve( + pybamm.IDAKLUSolver, native=False, model_cls=pybamm.lithium_ion.DFN + ) + native_entries = native_sol[var].entries + casadi_entries = casadi_sol[var].entries + assert native_entries.ndim == 3, ( + f"expected a 2D field (r, x, t), got ndim={native_entries.ndim}" + ) + assert native_entries.shape == casadi_entries.shape + # This test guards the 2D layout contract, not numerical equality. + np.testing.assert_allclose(native_entries, casadi_entries, rtol=1e-4, atol=1e-3) + + def test_3d_variable_native_parity(self): + # Pins the unroll_3D order for a 3D (r, R, x) field. + var = "Negative particle concentration distribution [mol.m-3]" + native_sol = _solve( + pybamm.IDAKLUSolver, native=True, model_cls=pybamm.lithium_ion.MPM + ) + casadi_sol = _solve( + pybamm.IDAKLUSolver, native=False, model_cls=pybamm.lithium_ion.MPM + ) + native_entries = native_sol[var].entries + casadi_entries = casadi_sol[var].entries + assert native_entries.ndim == 4, ( + f"expected a 3D field (r, R, x, t), got ndim={native_entries.ndim}" + ) + assert native_entries.shape == casadi_entries.shape + # Layout contract: a wrong 3D reshape would mismatch by orders of magnitude. + np.testing.assert_allclose(native_entries, casadi_entries, rtol=1e-4, atol=1e-3) 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 38b1214daa..307b14eb41 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_idaklu_solver.py @@ -11,6 +11,7 @@ import warnings from contextlib import redirect_stdout +import casadi import numpy as np import pandas as pd import pytest @@ -20,6 +21,8 @@ import pybamm from pybamm.solvers.idaklu_solver import _csc_bandwidths +from pybamm.solvers.observation import ObserverCache +from pybamm.solvers.variable_observer import NativeObserver from tests import ( get_broken_input_model, get_discretisation_for_testing, @@ -2748,6 +2751,26 @@ def test_spy(self, convert_to_format): assert "nnz" in ax.get_title() plt.close("all") + def test_reduce_solution_keeps_native_observation(self): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver(rtol=1e-6, atol=1e-8) + sol = pybamm.Simulation(model, solver=solver).solve([0, 3600]) + leaf = sol["Terminal voltage [V]"]._observer.leaves[0] + + reduced = solver.reduce_solution(sol, hermite_reduction_factor=2.0) + + assert reduced.all_models == sol.all_models + assert isinstance(reduced["Terminal voltage [V]"]._observer, NativeObserver) + assert reduced["Terminal voltage [V]"]._observer.leaves[0] is leaf + t = np.linspace(0, 3600, 51) + np.testing.assert_allclose( + reduced["Terminal voltage [V]"](t), + sol["Terminal voltage [V]"](t), + rtol=0, + atol=1e-4, + ) + def test_reduce_solution_vs_online(self): """Compare post-hoc reduce_solution with online knot reduction on a drive cycle. @@ -2950,6 +2973,62 @@ def test_offgrid_values_match_casadi(self): rust[name](t_interp), cas[name](t_interp), rtol=1e-6, atol=1e-6 ) + def test_native_processed_variable_matches_casadi_direct(self): + t_interp = np.linspace(0, 3600, 41) + cas, _ = self._solve("casadi", calculate_sensitivities=True, t_interp=t_interp) + sol, _ = self._solve("rust", calculate_sensitivities=True, t_interp=t_interp) + + name = "Terminal voltage [V]" + vars_pybamm = [m.get_processed_variable_or_event(name) for m in sol.all_models] + rust_fns = [ + ObserverCache.of(m).native_leaf(name, vp)[0] + for m, vp in zip(sol.all_models, vars_pybamm, strict=True) + ] + pv = pybamm.process_variable( + name, vars_pybamm, NativeObserver(rust_fns), sol, time_integral=None + ) + + np.testing.assert_allclose(pv.entries, cas[name].entries, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose( + pv(t_interp), cas[name](t_interp), rtol=1e-6, atol=1e-6 + ) + np.testing.assert_allclose( + pv.sensitivities["Current function [A]"], + cas[name].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_flip_compiles_no_casadi_observe(self): + rust, _ = self._solve("rust") + v = rust["Terminal voltage [V]"] + _ = v.entries # force observation + assert isinstance(v._observer, NativeObserver) + assert not any(isinstance(f, casadi.Function) for f in v._observer.leaves) + + def test_offgrid_values_native_backed(self): + rust, _ = self._solve("rust") + for name in [ + "Terminal voltage [V]", + "X-averaged negative particle concentration [mol.m-3]", + ]: + v = rust[name] + _ = v.entries + assert isinstance(v._observer, NativeObserver) + assert not any(isinstance(f, casadi.Function) for f in v._observer.leaves) + + def test_sensitivities_match_casadi(self): + cas, _ = self._solve("casadi", calculate_sensitivities=True) + rust, _ = self._solve("rust", calculate_sensitivities=True) + for name in ["Terminal voltage [V]"]: + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name].sensitivities["Current function [A]"], + cas[name].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + def _solve_output_vars( self, convert_to_format, output_variables, model=None, inputs=None, var_pts=None ): @@ -3077,6 +3156,38 @@ def test_output_variables_sensitivities_no_casadi_compiled(self): assert "dvar_dp_idaklu_fcns" not in solver._setup assert rust.variables_returned is True + @staticmethod + def _discretise_dfn(convert_to_format): + model = pybamm.lithium_ion.DFN() + geometry = model.default_geometry + param = model.default_parameter_values + param.process_model(model) + param.process_geometry(geometry) + var_pts = {"x_n": 20, "x_s": 20, "x_p": 20, "r_n": 10, "r_p": 10} + mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + model.convert_to_format = convert_to_format + return model + + def test_spatial_1d_2d_variables_match_casadi(self): + t_interp = np.linspace(0, 3600, 51) + cas = pybamm.IDAKLUSolver().solve( + self._discretise_dfn("casadi"), [0, 3600], t_interp=t_interp + ) + rust = pybamm.IDAKLUSolver().solve( + self._discretise_dfn("rust"), [0, 3600], t_interp=t_interp + ) + off = t_interp[:-1] + np.diff(t_interp) / 3 + for name in [ + "Electrolyte concentration [mol.m-3]", # 1D in x + "Negative particle concentration [mol.m-3]", # 2D (r, x) + ]: + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name](off), cas[name](off), rtol=1e-6, atol=1e-6 + ) + @staticmethod def _discretise_spm(convert_to_format, as_input=False): model = pybamm.lithium_ion.SPM() @@ -3092,6 +3203,92 @@ def _discretise_spm(convert_to_format, as_input=False): model.convert_to_format = convert_to_format return model + def test_multi_segment_experiment_matches_casadi(self): + def stepped(convert_to_format): + model = self._discretise_spm(convert_to_format) + solver = pybamm.IDAKLUSolver() + sol = None + for _ in range(2): + sol = solver.step(sol, model, dt=1800) + return sol + + cas = stepped("casadi") + rust = stepped("rust") + assert len(rust.all_ys) == 2 + off = np.linspace(0, rust.t[-1], 73)[1:-1] + 5.0 + off = off[off < rust.t[-1]] + for name in ["Terminal voltage [V]"]: + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name](off), cas[name](off), rtol=1e-6, atol=1e-6 + ) + + def test_event_termination_matches_casadi(self): + cas = pybamm.IDAKLUSolver().solve(self._discretise_spm("casadi"), [0, 100000]) + rust = pybamm.IDAKLUSolver().solve(self._discretise_spm("rust"), [0, 100000]) + assert rust.termination.startswith("event:") + name = "Terminal voltage [V]" + assert isinstance(rust[name]._observer, NativeObserver) + t_end = min(rust.t[-1], cas.t[-1]) + off = np.linspace(0, t_end, 40)[1:-1] + 1.0 + off = off[off < t_end] + np.testing.assert_allclose( + rust[name](off), cas[name](off), rtol=1e-6, atol=1e-6 + ) + + def test_hermite_off_still_native_no_casadi(self): + solver_kwargs = {"options": {"hermite_interpolation": False}} + cas = pybamm.IDAKLUSolver(**solver_kwargs).solve( + self._discretise_spm("casadi"), [0, 3600] + ) + rust = pybamm.IDAKLUSolver(**solver_kwargs).solve( + self._discretise_spm("rust"), [0, 3600] + ) + assert not rust.hermite_interpolation + v = rust["Terminal voltage [V]"] + assert isinstance(v._observer, NativeObserver) + assert not any(isinstance(f, casadi.Function) for f in v._observer.leaves) + np.testing.assert_allclose( + v.entries, cas["Terminal voltage [V]"].entries, rtol=1e-6, atol=1e-6 + ) + + def test_time_integral_sensitivities_match_casadi(self): + inputs = {"Current function [A]": 0.68} + cas = pybamm.IDAKLUSolver().solve( + self._discretise_spm("casadi", as_input=True), + [0, 3600], + inputs=inputs, + calculate_sensitivities=True, + ) + rust = pybamm.IDAKLUSolver().solve( + self._discretise_spm("rust", as_input=True), + [0, 3600], + inputs=inputs, + calculate_sensitivities=True, + ) + name = "Discharge capacity [A.h]" + assert isinstance(rust[name]._observer, NativeObserver) + np.testing.assert_allclose( + rust[name].sensitivities["Current function [A]"], + cas[name].sensitivities["Current function [A]"], + rtol=1e-5, + atol=1e-6, + ) + + def test_output_variables_first_last_state_observe_natively(self): + def solve(output_variables): + model = self._discretise_spm("rust") + solver = pybamm.IDAKLUSolver(output_variables=output_variables) + return solver.solve(model, [0, 3600]) + + full = solve(None) + outputs_only = solve(["Voltage [V]"]) + name = "Total lithium in electrolyte [mol]" + assert isinstance(outputs_only.last_state[name]._observer, NativeObserver) + np.testing.assert_array_equal( + outputs_only.last_state[name].data, full.last_state[name].data + ) + def test_output_variables_stay_aligned_across_the_batch_window(self): # 200 points cross the Rust FFI's 128-point batched-evaluation window. t_interp = np.linspace(0, 3600, 200) diff --git a/packages/pybamm/tests/unit/test_solvers/test_native_observation.py b/packages/pybamm/tests/unit/test_solvers/test_native_observation.py new file mode 100644 index 0000000000..be9cce10e7 --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_native_observation.py @@ -0,0 +1,629 @@ +"""Unit tests for native (Rust) observation: accessor, producer, compile-once.""" + +import numpy as np + +from pybamm.rust import ExprGraph +from pybamm.solvers.observation import ObserverCache +from pybamm.solvers.variable_observer import chain_rule_sensitivities + +# Default tolerances split the backends' step sequences on D_n; 1e-11 converges both. +_TWO_PARAM_SOLVER_TOL = 1e-11 + + +def _make_compiled_model(): + """Minimal 2-state model: rhs = [a * y1, -y0], identity mass, one input 'a'.""" + from pybamm.rust import CompiledModel + + g = ExprGraph() + a = g.input_parameter("a") + y0 = g.state_vector(0, 1) + y1 = g.state_vector(1, 2) + rhs = g.concat([g.mul(a, y1), g.neg(y0)]) + mass = np.array([1.0, 1.0]) + indptr = np.array([0, 1, 2], dtype=np.int64) + indices = np.array([0, 1], dtype=np.int64) + model = CompiledModel.from_expr(g, rhs, mass, indptr, indices, n_inputs=1) + return g, model + + +class TestGraphAccessor: + def test_graph_returns_exprgraph(self): + _, model = _make_compiled_model() + assert isinstance(model.graph, ExprGraph) + + def test_graph_compiles_a_new_observation_root(self): + # The new root has no inputs, so p is the empty stacked array. + _, model = _make_compiled_model() + graph = model.graph + new_root = graph.mul(graph.scalar(2.0), graph.state_vector(0, 1)) + fn = graph.compile(new_root, name="obs", n_states=model.n_states) + out = fn(0.0, np.array([3.0, 5.0]), np.array([1.0])) + np.testing.assert_allclose(out, [6.0]) + + +def _native_idaklu_sim(model, **solver_kwargs): + """Simulation on idaklu of ``model`` as a rust model, so it reads natively.""" + import pybamm + + model.convert_to_format = "rust" + return pybamm.Simulation(model, solver=pybamm.IDAKLUSolver(**solver_kwargs)) + + +class TestSingleVariableNative: + def test_native_matches_casadi_0d(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + sol = _native_idaklu_sim(model).solve([0, 3600]) + v_native = sol["Terminal voltage [V]"].entries + + casadi_model = pybamm.lithium_ion.SPM() + casadi_model.convert_to_format = "casadi" + sim_c = pybamm.Simulation(casadi_model, solver=pybamm.IDAKLUSolver()) + sol_c = sim_c.solve([0, 3600]) + v_casadi = sol_c["Terminal voltage [V]"].entries + + np.testing.assert_allclose(v_native, v_casadi, rtol=1e-10, atol=1e-12) + + def test_hermite_native_matches_casadi_multi_tile(self): + import pybamm + + # Off-knot grid spanning several Hermite lane tiles plus a ragged tail. + model = pybamm.lithium_ion.SPM() + sol = _native_idaklu_sim(model).solve([0, 3600]) + + casadi_model = pybamm.lithium_ion.SPM() + casadi_model.convert_to_format = "casadi" + sol_c = pybamm.Simulation(casadi_model, solver=pybamm.IDAKLUSolver()).solve( + [0, 3600] + ) + + t_query = np.linspace(0, 3600, 137) + v_native = sol["Terminal voltage [V]"](t_query) + v_casadi = sol_c["Terminal voltage [V]"](t_query) + assert v_native.shape == t_query.shape + np.testing.assert_allclose(v_native, v_casadi, rtol=1e-6, atol=1e-8) + + def test_compile_once(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + sol = _native_idaklu_sim(model).solve([0, 600]) + leaf = sol["Terminal voltage [V]"]._observer.leaves[0] + cache = ObserverCache.of(sol.all_models[0]) + assert cache._native_leaves["Terminal voltage [V]"][0] is leaf + sol.update("Terminal voltage [V]") + assert sol["Terminal voltage [V]"]._observer.leaves[0] is leaf + + +def _make_var_fn(): + """g(t, y, p) = [a*y0*y1, sin(y0)*b]; 2 states, 2 inputs ('a','b').""" + g = ExprGraph() + a = g.input_parameter("a") + b = g.input_parameter("b") + y0 = g.state_vector(0, 1) + y1 = g.state_vector(1, 2) + expr = g.concat([g.mul(g.mul(a, y0), y1), g.mul(g.sin(y0), b)]) + return g, g.compile(expr, name="var") + + +class TestNativeChainRuleKernel: + def _trajectory(self): + n_t, n_states, n_p = 8, 2, 2 + ts = np.linspace(0.0, 1.0, n_t) + ys = np.vstack([np.linspace(0.1, 1.0, n_t), np.linspace(2.0, 3.0, n_t)]) + rng = np.random.default_rng(0) + dy_dp = rng.standard_normal((n_t * n_states, n_p)) + p = np.array([3.0, 4.0]) + return ts, ys, dy_dp, p, n_t, n_states, n_p + + def test_matches_numpy_jacobian_oracle(self): + # Independent oracle: per timestep, dvar_dy @ yS_k + dvar_dp @ e_k. + _, cf = _make_var_fn() + ts, ys, dy_dp, p, n_t, n_states, n_p = self._trajectory() + out_len = cf.output_len + Jy = cf.jacobian("y") + Jp = cf.jacobian("p") + S = chain_rule_sensitivities(cf, ts, ys, p, dy_dp, ["a", "b"]) + assert S.shape == (n_t * out_len, n_p) + for k in range(n_p): + e_k = np.zeros(n_p) + e_k[k] = 1.0 + for j in range(n_t): + dvar_dy = Jy(ts[j], ys[:, j].copy(), p).toarray() + dvar_dp = Jp(ts[j], ys[:, j].copy(), p).toarray() + yS_jk = dy_dp[j * n_states : (j + 1) * n_states, k] + oracle = dvar_dy @ yS_jk + dvar_dp @ e_k + block = S[j * out_len : (j + 1) * out_len, k] + np.testing.assert_allclose(block, oracle, rtol=1e-12, atol=1e-12) + + def test_preserves_zero_derivative_width(self): + g = ExprGraph() + cv = g.array(np.array([1.0, 2.0, 3.0])) + cf = g.compile(g.add(cv, g.input_parameter("a")), name="zeroy", n_states=2) + assert cf.output_len == 3 + n_t, n_states = 5, 2 + ts = np.linspace(0.0, 1.0, n_t) + ys = np.zeros((n_states, n_t)) + dy_dp = np.ones((n_t * n_states, 1)) + S = chain_rule_sensitivities(cf, ts, ys, np.array([5.0]), dy_dp, ["a"]) + assert S.shape == (n_t * 3, 1) + np.testing.assert_array_equal(S, np.ones((n_t * 3, 1))) + + +class TestNativeSensitivities: + def _solve_both(self, output_name, calc, extra_inputs=None, solver_tol=None): + import pybamm + + params = pybamm.ParameterValues("Chen2020") + params["Current function [A]"] = pybamm.InputParameter("I") + inputs = {"I": 0.5} + if extra_inputs: + for pybamm_name, (input_name, value) in extra_inputs.items(): + params[pybamm_name] = pybamm.InputParameter(input_name) + inputs[input_name] = value + t_eval = np.linspace(0, 100, 15) + tol = {} if solver_tol is None else {"rtol": solver_tol, "atol": solver_tol} + + m_native = pybamm.lithium_ion.SPM() + m_native.events = [] + m_native.convert_to_format = "rust" + solver_n = pybamm.IDAKLUSolver(**tol) + sol_n = pybamm.Simulation( + m_native, parameter_values=params, solver=solver_n + ).solve( + t_eval, + inputs=inputs, + calculate_sensitivities=calc, + t_interp=t_eval, + ) + + m_casadi = pybamm.lithium_ion.SPM() + m_casadi.events = [] + m_casadi.convert_to_format = "casadi" + sol_c = pybamm.Simulation( + m_casadi, parameter_values=params, solver=pybamm.IDAKLUSolver(**tol) + ).solve(t_eval, inputs=inputs, calculate_sensitivities=calc, t_interp=t_eval) + return sol_n[output_name], sol_c[output_name] + + def test_sensitivities_match_casadi_0d(self): + var_n, var_c = self._solve_both("Terminal voltage [V]", ["I"]) + sens_n = var_n.sensitivities + sens_c = var_c.sensitivities + assert set(sens_n) == set(sens_c) + for key in sens_c: + np.testing.assert_allclose(sens_n[key], sens_c[key], rtol=1e-5, atol=1e-8) + + def test_sensitivities_match_casadi_spatial(self): + # Exercises the time-outer/output-inner flatten against CasADi's layout. + name = "Negative particle concentration [mol.m-3]" + var_n, var_c = self._solve_both(name, ["I"]) + assert var_c.entries.shape[0] > 1, ( + f"{name!r} resolved to 0D (shape {var_c.entries.shape}); " + "test premise violated" + ) + sens_n = var_n.sensitivities + sens_c = var_c.sensitivities + assert set(sens_n) == set(sens_c) + for key in sens_c: + assert sens_n[key].shape == sens_c[key].shape + np.testing.assert_allclose(sens_n[key], sens_c[key], rtol=1e-5, atol=1e-8) + + def test_sensitivities_all_block_two_params(self): + # "I" plus a diffusivity input; sorted column order is ["D_n", "I"]. + var_n, var_c = self._solve_both( + "Terminal voltage [V]", + ["D_n", "I"], + extra_inputs={"Negative particle diffusivity [m2.s-1]": ("D_n", 3.3e-14)}, + solver_tol=_TWO_PARAM_SOLVER_TOL, + ) + a_n = var_n.sensitivities["all"] + a_c = var_c.sensitivities["all"] + assert a_n.shape == a_c.shape + assert a_n.shape[1] == 2 + assert np.any(var_n.sensitivities["D_n"] != 0) + np.testing.assert_allclose( + a_n[:, 0], var_n.sensitivities["D_n"], rtol=1e-5, atol=1e-8 + ) + np.testing.assert_allclose(a_n, a_c, rtol=1e-5, atol=1e-8) + + def test_no_sensitivities_when_not_requested(self): + import pybamm + + model = pybamm.lithium_ion.SPM() + sol = _native_idaklu_sim(model).solve([0, 600]) + assert sol["Terminal voltage [V]"].sensitivities == {} + + def test_outputs_only_solve_reads_other_variables_sensitivities(self): + import pybamm + from pybamm.solvers.variable_observer import NativeObserver + + # "2a" and "a t" read no states, so an outputs-only solve still serves them + def solve(convert_to_format): + model = pybamm.BaseModel(name="decay") + c = pybamm.Variable("c") + a = pybamm.InputParameter("a") + model.rhs = {c: -a * c} + model.initial_conditions = {c: 1} + model.variables = {"c": c, "2a": 2 * a, "a t": a * pybamm.t} + model.convert_to_format = convert_to_format + solver = pybamm.IDAKLUSolver(output_variables=["c"]) + # The lowerings step differently, so t_interp puts both on one grid + return solver.solve( + model, + [0, 1], + t_interp=np.linspace(0, 1, 11), + inputs={"a": 1.5}, + calculate_sensitivities=True, + ) + + native, casadi = solve("rust"), solve("casadi") + for name in ("2a", "a t"): + assert isinstance(native[name]._observer, NativeObserver) + np.testing.assert_allclose( + native[name].sensitivities["a"], + casadi[name].sensitivities["a"], + rtol=1e-12, + atol=1e-12, + ) + + +class TestObservationCacheAcrossSolves: + @staticmethod + def _sim(): + # events removed so termination is "final time" + import pybamm + + model = pybamm.lithium_ion.SPM() + model.events = [] + return _native_idaklu_sim(model) + + def test_compiled_leaf_reused_across_solves(self): + sim = self._sim() + sol1 = sim.solve([0, 600]) + v1 = sol1["Terminal voltage [V]"] + sol2 = sim.solve([0, 600]) + v2 = sol2["Terminal voltage [V]"] + + cache = ObserverCache.of(sol1.all_models[0]) + assert sol2.all_models[0] is sol1.all_models[0] + assert cache.rust_model is sim._solver._setup["rust_model"] + assert v2._observer.leaves[0] is v1._observer.leaves[0] + + def test_retained_graph_stops_growing(self): + sim = self._sim() + sol1 = sim.solve([0, 600]) + _ = sol1["Terminal voltage [V]"] + cache = ObserverCache.of(sol1.all_models[0])._native_leaves + nodes_after_first = sim._solver._setup["rust_model"].graph.n_nodes + cache_len_after_first = len(cache) + + sol2 = sim.solve([0, 600]) + _ = sol2["Terminal voltage [V]"] + assert sim._solver._setup["rust_model"].graph.n_nodes == nodes_after_first + assert len(cache) == cache_len_after_first + + def test_discrete_time_sum_across_two_solves(self): + import pybamm + + # The 0D model auto-discretises, so the solver is driven directly. + data_times = np.linspace(0, 1, 10) + data = pybamm.DiscreteTimeData(data_times, np.zeros_like(data_times), "zeros") + + def _build_model(fmt): + m = pybamm.BaseModel(name="dts_model") + c = pybamm.Variable("c") + m.rhs = {c: -2 * c} + m.initial_conditions = {c: 1} + m.variables["c"] = c + # (c - 0)^2 summed over the data times -> sum(exp(-4 t)) + m.variables["dts"] = pybamm.DiscreteTimeSum((c - data) ** 2) + m.convert_to_format = fmt + return m + + model = _build_model("rust") + solver = pybamm.IDAKLUSolver() + sol1 = solver.solve(model, t_eval=[0, 1], t_interp=data_times) + val1 = sol1["dts"]() + sol2 = solver.solve(model, t_eval=[0, 1], t_interp=data_times) + val2 = sol2["dts"]() + + np.testing.assert_allclose(val2, val1, rtol=1e-12, atol=0) + model_c = _build_model("casadi") + sol_c = pybamm.IDAKLUSolver().solve(model_c, t_eval=[0, 1], t_interp=data_times) + np.testing.assert_allclose(val1, sol_c["dts"](), rtol=1e-6, atol=1e-8) + + time_integral = ObserverCache.of(model)._time_integrals["dts"] + assert sol2["dts"].time_integral is time_integral + assert sol1["dts"].time_integral is time_integral + + def test_pickle_round_trip_after_observe(self): + import pickle + + sim = self._sim() + sol = sim.solve([0, 600]) + expected = sol["Terminal voltage [V]"].entries # populate the cache + + restored = pickle.loads(pickle.dumps(sol)) + np.testing.assert_allclose( + restored["Terminal voltage [V]"].entries, expected, rtol=1e-10, atol=1e-12 + ) + + +class TestBatchedNativeObservation: + """`update([...])` shares one tape per segment across the whole set.""" + + NAMES = [ + "Terminal voltage [V]", + "Current [A]", + "X-averaged negative particle surface concentration [mol.m-3]", + "Electrolyte concentration [mol.m-3]", + ] + + def _experiment_solution(self): + import pybamm + + model = pybamm.lithium_ion.SPMe() + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver() + experiment = pybamm.Experiment( + [("Discharge at 1C until 3.4 V", "Charge at 0.5C until 3.9 V")] + ) + sim = pybamm.Simulation(model, solver=solver, experiment=experiment) + return sim.solve() + + def test_batched_matches_per_name_bitwise(self): + batched = self._experiment_solution() + per_name = self._experiment_solution() + assert len(batched.all_ys) > 1 + + batched.update(self.NAMES) + for name in self.NAMES: + np.testing.assert_array_equal(batched[name].entries, per_name[name].entries) + + def test_batched_hermite_matches_per_name_bitwise(self): + batched = self._experiment_solution() + per_name = self._experiment_solution() + batched.update(self.NAMES) + + # off the solver's own knots, so every read Hermite-reconstructs + t = np.linspace(batched.t[0] + 1.0, batched.t[-1] - 1.0, 61) + for name in self.NAMES: + np.testing.assert_array_equal(batched[name](t), per_name[name](t)) + + def test_one_tape_per_segment_serves_every_variable(self): + from pybamm.solvers.variable_observer import SharedOutputLeaf + + solution = self._experiment_solution() + solution.update(self.NAMES) + observers = [solution[name]._observer for name in self.NAMES] + + leaves = [observer.leaves[0] for observer in observers] + assert all(isinstance(leaf, SharedOutputLeaf) for leaf in leaves) + shared = {id(leaf._shared) for leaf in leaves} + assert len(shared) == 1 + assert leaves[0]._shared.function.output_names == tuple(self.NAMES) + + def test_the_set_is_compiled_once_and_reused(self): + solution = self._experiment_solution() + cache = ObserverCache.of(solution.all_models[0]) + graph = cache.rust_model.graph + + solution.update(self.NAMES) + sets = dict(cache._native_sets) + nodes = graph.n_nodes + assert tuple(self.NAMES) in sets + + solution.update(self.NAMES) + assert cache._native_sets == sets + assert graph.n_nodes == nodes + + def test_a_time_integral_keeps_the_per_name_path(self): + import pybamm + from pybamm.solvers.variable_observer import SharedOutputLeaf + + # A time integral is not a row of the set's tape, so it falls back alone. + model = pybamm.BaseModel(name="integral_model") + c = pybamm.Variable("c") + model.rhs = {c: -2 * c} + model.initial_conditions = {c: 1} + model.variables = { + "c": c, + "2c": 2 * c, + "integral": pybamm.ExplicitTimeIntegral(c, 0), + } + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver() + solution = solver.solve(model, t_eval=[0, 1]) + + per_name = solver.solve(model, t_eval=[0, 1]) + expected = {name: per_name[name].entries for name in ("c", "2c", "integral")} + + solution.update(["c", "2c", "integral"]) + assert not isinstance( + solution["integral"]._observer.leaves[0], SharedOutputLeaf + ) + assert isinstance(solution["c"]._observer.leaves[0], SharedOutputLeaf) + for name, want in expected.items(): + np.testing.assert_array_equal(solution[name].entries, want) + + def test_a_single_name_is_not_batched(self): + from pybamm.solvers.variable_observer import SharedOutputLeaf + + solution = self._experiment_solution() + solution.update([self.NAMES[0]]) + leaf = solution[self.NAMES[0]]._observer.leaves[0] + assert not isinstance(leaf, SharedOutputLeaf) + + def test_a_data_export_prepares_the_requested_set(self): + from pybamm.solvers.variable_observer import SharedOutputLeaf + + solution = self._experiment_solution() + data = solution.get_data_dict(self.NAMES, cycles_and_steps=False) + + assert set(data) == set(self.NAMES) + leaves = [solution[name]._observer.leaves[0] for name in self.NAMES] + assert all(isinstance(leaf, SharedOutputLeaf) for leaf in leaves) + + def test_an_already_built_variable_is_left_alone(self): + solution = self._experiment_solution() + built = solution[self.NAMES[0]] + solution._build_missing(self.NAMES) + assert solution[self.NAMES[0]] is built + + def test_batched_sensitivities_match_per_name_bitwise(self): + import pybamm + + def solve(): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver() + parameters = pybamm.ParameterValues("Chen2020") + parameters.update({"Negative particle diffusivity [m2.s-1]": "[input]"}) + sim = pybamm.Simulation(model, parameter_values=parameters, solver=solver) + return sim.solve( + [0, 3600], + inputs={"Negative particle diffusivity [m2.s-1]": 3.3e-14}, + calculate_sensitivities=True, + ) + + names = [ + "Terminal voltage [V]", + "X-averaged negative particle surface concentration [mol.m-3]", + ] + batched, per_name = solve(), solve() + batched.update(names) + for name in names: + np.testing.assert_array_equal( + batched[name].sensitivities["all"], + per_name[name].sensitivities["all"], + ) + + def test_batched_sensitivities_match_per_name_over_several_parameters(self): + import pybamm + + def solve(): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver() + parameters = pybamm.ParameterValues("Chen2020") + parameters.update( + { + "Negative particle diffusivity [m2.s-1]": "[input]", + "Current function [A]": "[input]", + } + ) + sim = pybamm.Simulation(model, parameter_values=parameters, solver=solver) + return sim.solve( + [0, 3600], + inputs={ + "Negative particle diffusivity [m2.s-1]": 3.3e-14, + "Current function [A]": 0.68, + }, + calculate_sensitivities=True, + ) + + names = [ + "Terminal voltage [V]", + "Negative particle concentration [mol.m-3]", + ] + batched, per_name = solve(), solve() + batched.update(names) + + spatial = per_name[names[1]].sensitivities["all"] + assert spatial.shape[1] == 2 + assert spatial.shape[0] > len(per_name.t) + for name in names: + np.testing.assert_array_equal( + batched[name].sensitivities["all"], + per_name[name].sensitivities["all"], + ) + + def test_a_mutated_query_array_is_not_read_from_the_held_tape(self): + import pybamm + from pybamm.solvers.variable_observer import SharedOutputLeaf + + # Mutating the last call's array must not make the next call match it. + model = pybamm.BaseModel(name="decay") + c = pybamm.Variable("c") + model.rhs = {c: -c} + model.initial_conditions = {c: 1} + model.variables = {"c": c, "2c": 2 * c} + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver() + solution = solver.solve(model, t_eval=[0, 1]) + solution.update(["c", "2c"]) + assert isinstance(solution["c"]._observer.leaves[0], SharedOutputLeaf) + assert solution["c"].hermite_interpolation + + t = np.array([0.2, 0.3]) + solution["c"](t) + t[:] = [0.7, 0.8] + np.testing.assert_allclose(solution["2c"](t), 2 * np.exp(-t), rtol=1e-3) + + def test_batched_sensitivities_over_an_empty_segment(self): + import pybamm + + def joined(): + model = pybamm.BaseModel(name="decay") + c = pybamm.Variable("c") + a = pybamm.InputParameter("a") + model.rhs = {c: -a * c} + model.initial_conditions = {c: 1} + model.variables = {"c": c, "2c": 2 * c} + model.convert_to_format = "rust" + solver = pybamm.IDAKLUSolver() + kwargs = {"inputs": {"a": 1.0}, "calculate_sensitivities": True} + first = solver.solve(model, [0, 1], **kwargs) + second = solver.solve(model, [1, 2], **kwargs) + # The last state's single point is dropped as a repeat, leaving it empty + return first + (first.last_state + second) + + batched, per_name = joined(), joined() + assert any(t.size == 0 for t in batched.all_ts) + batched.update(["c", "2c"]) + for name in ("c", "2c"): + np.testing.assert_array_equal( + batched[name].sensitivities["all"], + per_name[name].sensitivities["all"], + ) + + def test_concurrent_reads_each_get_their_own_rows(self): + import sys + import threading + from concurrent.futures import ThreadPoolExecutor + + import pybamm + + model = pybamm.BaseModel(name="decay") + c = pybamm.Variable("c") + model.rhs = {c: -c} + model.initial_conditions = {c: 1} + model.variables = {"c": c, "2c": 2 * c, "3c": 3 * c} + model.convert_to_format = "rust" + solution = pybamm.IDAKLUSolver().solve(model, t_eval=[0, 1]) + names = ["c", "2c", "3c"] + solution.update(names) + + # Every thread reads the whole set at its own times, so held calls differ + n_threads = 8 + queries = [np.linspace(0.01 + 1e-3 * k, 0.99, 40) for k in range(n_threads)] + expected = [{name: solution[name](t) for name in names} for t in queries] + barrier = threading.Barrier(n_threads) + + def read(k): + barrier.wait() + for _ in range(50): + for name in names: + np.testing.assert_array_equal( + solution[name](queries[k]), expected[k][name] + ) + + # Frequent thread switches expose any unguarded window in the shared hold + switch_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + try: + with ThreadPoolExecutor(n_threads) as pool: + list(pool.map(read, range(n_threads))) + finally: + sys.setswitchinterval(switch_interval) diff --git a/packages/pybamm/tests/unit/test_solvers/test_observation.py b/packages/pybamm/tests/unit/test_solvers/test_observation.py index 71ef100907..ebcc7ca26d 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_observation.py +++ b/packages/pybamm/tests/unit/test_solvers/test_observation.py @@ -13,15 +13,22 @@ from pybamm.solvers.observation import ObserverCache from pybamm.solvers.variable_observer import ( CasadiObserver, + NativeObserver, SegmentSelector, + SharedOutputLeaf, + SharedOutputs, + _sweep_sensitivities, as_observer, + chain_rule_sensitivities, pack_sensitivity_dict, ) @pytest.fixture(scope="module") def spm_solution(): + # CasADi whatever the default, since these tests read CasADi leaves model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" return pybamm.Simulation(model).solve([0, 600]) @@ -137,6 +144,317 @@ def test_a_time_integral_is_converted_once_per_model(self): assert spy.call_count == 3 +def _rust_decay_model(factor): + model = _decay_model(factor) + model.convert_to_format = "rust" + return model + + +class TestNativeReading: + def test_a_rust_solve_binds_its_compiled_model(self): + model = _rust_decay_model(2) + solver = pybamm.IDAKLUSolver() + solution = solver.solve(model, [0, 1]) + + assert ObserverCache.of(model).rust_model is solver._setup["rust_model"] + assert isinstance(solution["w"]._observer, NativeObserver) + np.testing.assert_allclose(solution["w"].entries, 2 * solution["u"].entries) + + def test_derived_solutions_read_through_the_same_tapes(self): + solution = pybamm.IDAKLUSolver().solve(_rust_decay_model(2), [0, 1]) + leaf = solution["w"]._observer.leaves[0] + + for derived in (solution.copy(), solution.first_state, solution.last_state): + assert derived["w"]._observer.leaves == [leaf] + + def test_solutions_of_two_rust_models_join_natively(self): + first = pybamm.IDAKLUSolver().solve(_rust_decay_model(2), [0, 1]) + second = pybamm.IDAKLUSolver().solve(_rust_decay_model(3), [1, 2]) + + joined = first + second + + assert isinstance(joined["w"]._observer, NativeObserver) + 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) + + def test_a_join_with_a_casadi_segment_reads_through_casadi(self): + casadi_model = _decay_model(3) + casadi_model.convert_to_format = "casadi" + first = pybamm.IDAKLUSolver().solve(_rust_decay_model(2), [0, 1]) + second = pybamm.IDAKLUSolver().solve(casadi_model, [1, 2]) + + joined = first + second + + assert isinstance(joined["w"]._observer, CasadiObserver) + 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) + + def test_a_model_switched_back_to_casadi_reads_through_casadi(self): + model = _rust_decay_model(2) + solution = pybamm.IDAKLUSolver().solve(model, [0, 1]) + + model.convert_to_format = "casadi" + + assert isinstance(solution["w"]._observer, CasadiObserver) + + def test_setting_up_again_drops_the_previous_graph_s_tapes(self): + model = _rust_decay_model(2) + pybamm.IDAKLUSolver().solve(model, [0, 1])["w"] + cache = ObserverCache.of(model) + cache.bind_rust_model(cache.rust_model) + assert "w" in cache._native_leaves + + solver = pybamm.IDAKLUSolver() + solver.set_up(model) + + assert cache.rust_model is solver._setup["rust_model"] + assert "w" not in cache._native_leaves + + def test_rebinding_leaves_how_a_solution_without_yps_is_read(self): + model = _rust_decay_model(2) + solver = pybamm.IDAKLUSolver(options={"hermite_interpolation": False}) + solution = solver.solve(model, [0, 1]) + assert not solution.hermite_interpolation + before = solution["w"].entries + + pybamm.IDAKLUSolver().set_up(model) + solution.update("w") + + assert isinstance(solution["w"]._observer, NativeObserver) + np.testing.assert_array_equal(solution["w"].entries, before) + + def test_a_model_copy_keeps_the_binding_and_its_tapes(self): + model = _rust_decay_model(2) + leaf = pybamm.IDAKLUSolver().solve(model, [0, 1])["w"]._observer.leaves[0] + + cache = ObserverCache.of(model.new_copy()) + + assert cache.rust_model is ObserverCache.of(model).rust_model + assert cache._native_leaves["w"][0] is leaf + + +class _Variable: + """The slice of ProcessedVariable an observer is allowed to read.""" + + def __init__(self, all_ts, all_ys, all_yps=None): + self.all_ts = all_ts + self.all_ys = all_ys + self.all_yps = all_yps + self.all_inputs = [{} for _ in all_ts] + self.t_pts = np.concatenate(all_ts) + self.time_integral = None + + @property + def hermite_interpolation(self): + return self.all_yps is not None + + def _shape(self, t): + return [len(t)] + + +class _Leaf: + """Stand-in for a compiled tape: records its calls, returns 2*y0.""" + + def __init__(self): + self.trajectory_calls = [] + self.hermite_calls = [] + + def eval_trajectory(self, ts, ys, inputs): + self.trajectory_calls.append((ts, ys)) + return np.asfortranarray(2.0 * np.asarray(ys)[:1, :]) + + def eval_trajectory_hermite(self, t_query, ts, ys, yps, inputs): + self.hermite_calls.append(t_query) + return np.asfortranarray(np.zeros((1, len(t_query)))) + + +class TestNativeObserverStandalone: + def test_observe_raw_evaluates_every_segment(self): + leaves = [_Leaf(), _Leaf()] + observer = NativeObserver(leaves) + variable = _Variable( + [np.array([0.0, 1.0]), np.array([2.0])], + [np.array([[1.0, 2.0]]), np.array([[3.0]])], + ) + + np.testing.assert_allclose(observer.observe_raw(variable), [2.0, 4.0, 6.0]) + assert observer.leaves is leaves + + def test_outputs_only_solves_evaluate_on_shaped_zeros(self): + leaf = _Leaf() + observer = NativeObserver([leaf], placeholder_states=[3]) + variable = _Variable([np.array([0.0, 1.0])], [np.zeros((0, 0))]) + + np.testing.assert_allclose(observer.observe_raw(variable), [0.0, 0.0]) + _, ys = leaf.trajectory_calls[0] + assert ys.shape == (3, 2) + + def test_empty_segments_are_skipped(self): + leaves = [_Leaf(), _Leaf()] + observer = NativeObserver(leaves) + variable = _Variable( + [np.array([]), np.array([0.0, 1.0])], + [np.zeros((1, 0)), np.array([[1.0, 2.0]])], + ) + + np.testing.assert_allclose(observer.observe_raw(variable), [2.0, 4.0]) + assert leaves[0].trajectory_calls == [] + + def test_the_query_times_a_leaf_is_handed_are_the_observer_s_own(self): + # A shared tape keeps its last call, so a leaf must not see the caller's array. + leaf = _Leaf() + observer = NativeObserver([leaf]) + variable = _Variable( + [np.array([0.0, 1.0])], + [np.array([[1.0, 2.0]])], + all_yps=[np.array([[0.0, 0.0]])], + ) + + t = np.array([0.2, 0.7]) + observer.observe_hermite(variable, t) + t[:] = [0.9, 0.95] + np.testing.assert_array_equal(leaf.hermite_calls[0], [0.2, 0.7]) + + def test_segment_selection_is_computed_once_per_observer(self): + observer = NativeObserver([_Leaf()]) + variable = _Variable([np.array([0.0, 1.0])], [np.array([[1.0, 2.0]])]) + + observer.observe_raw(variable) + selector = observer._selector + observer.observe_raw(variable) + assert observer._selector is selector + + +class _SharedFunction: + """Stand-in for a multi-output CompiledFunction: records its calls.""" + + output_names = ("a", "b") + output_lens = (2, 1) + input_names = ("p",) + n_inputs = 1 + + output_len = 3 + + def __init__(self): + self.calls = [] + self.jvp_calls = [] + + def eval_trajectory(self, ts, ys, inputs): + self.calls.append(("eval_trajectory", ts)) + return np.vstack( + [ + np.asarray(ys)[:1] * 2.0, + np.asarray(ys)[:1] * 3.0, + np.asarray(ys)[:1] * 5.0, + ] + ) + + def jvp_trajectory(self, ts, ys, inputs, vy_traj, vp=None): + self.jvp_calls.append((vy_traj, vp)) + seed = 0.0 if vp is None else float(vp[0]) + row = np.asarray(vy_traj)[:1] + return np.vstack([row * 2.0, row * 3.0, row * 5.0]) + seed + + +class TestSharedOutputs: + def _shared(self): + function = _SharedFunction() + shared = SharedOutputs(function) + return ( + function, + shared, + SharedOutputLeaf(shared, "a"), + SharedOutputLeaf(shared, "b"), + ) + + def test_each_leaf_sees_only_its_own_rows(self): + _, _, leaf_a, leaf_b = self._shared() + ys = np.array([[1.0, 2.0]]) + np.testing.assert_allclose( + leaf_a.eval_trajectory(np.array([0.0, 1.0]), ys, {}), + [[2.0, 4.0], [3.0, 6.0]], + ) + np.testing.assert_allclose( + leaf_b.eval_trajectory(np.array([0.0, 1.0]), ys, {}), [[5.0, 10.0]] + ) + assert leaf_a.output_len == 2 + assert leaf_b.output_len == 1 + + def test_one_evaluation_serves_every_leaf_of_a_call(self): + function, _, leaf_a, leaf_b = self._shared() + ts, ys = np.array([0.0, 1.0]), np.array([[1.0, 2.0]]) + leaf_a.eval_trajectory(ts, ys, {}) + leaf_b.eval_trajectory(ts, ys, {}) + assert len(function.calls) == 1 + + def test_a_different_call_evaluates_afresh(self): + function, _, leaf_a, leaf_b = self._shared() + ys = np.array([[1.0, 2.0]]) + leaf_a.eval_trajectory(np.array([0.0, 1.0]), ys, {}) + leaf_b.eval_trajectory(np.array([2.0, 3.0]), ys, {}) + assert len(function.calls) == 2 + + def test_equal_arrays_from_different_objects_still_share(self): + function, _, leaf_a, leaf_b = self._shared() + leaf_a.eval_trajectory(np.array([0.0, 1.0]), np.array([[1.0, 2.0]]), {}) + leaf_b.eval_trajectory(np.array([0.0, 1.0]), np.array([[1.0, 2.0]]), {}) + assert len(function.calls) == 1 + + def test_the_held_result_is_released_once_every_leaf_has_read_it(self): + _, shared, leaf_a, leaf_b = self._shared() + ts, ys = np.array([0.0, 1.0]), np.array([[1.0, 2.0]]) + leaf_a.eval_trajectory(ts, ys, {}) + assert shared._result is not None + leaf_b.eval_trajectory(ts, ys, {}) + assert shared._result is None + + @staticmethod + def _sweep_arguments(): + """``(ts, ys, dy_dp)`` for a two-time, one-state, two-parameter sweep.""" + return ( + np.array([0.0, 1.0]), + np.array([[1.0, 2.0]]), + np.array([[1.0, 0.5], [0.25, 2.0]]), + ) + + def test_one_sweep_per_parameter_serves_the_whole_set(self): + function, _, leaf_a, leaf_b = self._shared() + ts, ys, dy_dp = self._sweep_arguments() + + # sensitivity_names is rebuilt per variable: equal, but not the same list. + chain_rule_sensitivities(leaf_a, ts, ys, {}, dy_dp, ["p", "q"]) + chain_rule_sensitivities(leaf_b, ts, ys, {}, dy_dp, ["p", "q"]) + assert len(function.jvp_calls) == 2 + + def test_each_leaf_takes_its_own_rows_of_the_set_s_sweep(self): + function, _, leaf_a, leaf_b = self._shared() + ts, ys, dy_dp = self._sweep_arguments() + names = ["p", "q"] + + blocks = _sweep_sensitivities(function, ts, ys, {}, dy_dp, names).reshape( + len(ts), function.output_len, len(names) + ) + np.testing.assert_array_equal( + chain_rule_sensitivities(leaf_a, ts, ys, {}, dy_dp, names), + blocks[:, 0:2, :].reshape(2 * len(ts), len(names)), + ) + np.testing.assert_array_equal( + chain_rule_sensitivities(leaf_b, ts, ys, {}, dy_dp, names), + blocks[:, 2:3, :].reshape(len(ts), len(names)), + ) + + def test_pickling_drops_the_held_result(self): + import pickle + + function = _SharedFunction() + shared = SharedOutputs(function) + leaf = SharedOutputLeaf(shared, "a") + leaf.eval_trajectory(np.array([0.0]), np.array([[1.0]]), {}) + restored = pickle.loads(pickle.dumps(shared)) + assert restored._result is None + assert restored.span("b") == slice(2, 3) + + class TestSegmentSelector: def test_full_range_keeps_every_nonempty_segment(self): selector = SegmentSelector( @@ -196,6 +514,10 @@ def test_an_observer_is_passed_through(self): observer = CasadiObserver([None]) assert as_observer(observer) is observer + def test_a_native_observer_is_passed_through(self): + observer = NativeObserver([_Leaf()]) + assert as_observer(observer) is observer + def test_process_variable_accepts_either_form(self, spm_solution): name = "Terminal voltage [V]" base = [ diff --git a/packages/pybamm/tests/unit/test_solvers/test_solution.py b/packages/pybamm/tests/unit/test_solvers/test_solution.py index 26d5967fff..1234b11b64 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_solution.py +++ b/packages/pybamm/tests/unit/test_solvers/test_solution.py @@ -1258,6 +1258,25 @@ def test_update_rebuilds_every_listed_variable(self): np.testing.assert_allclose(solution["c"].entries, built.entries) assert "2c" in solution._variables + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) + def test_update_stores_nothing_when_a_listed_variable_fails( + self, convert_to_format + ): + model = pybamm.BaseModel() + c = pybamm.Variable("c") + model.rhs = {c: -c} + model.initial_conditions = {c: 1} + model.variables = {"c": c, "2c": 2 * c} + model.convert_to_format = convert_to_format + solution = pybamm.IDAKLUSolver().solve(model, [0, 1]) + built = solution["c"] + + with pytest.raises(KeyError, match=r"missing"): + solution.update(["c", "2c", "missing"]) + + assert solution._variables["c"] is built + assert "2c" not in solution._variables + def test_solution_save_load_rust_backend(self, tmp_path): model = pybamm.lithium_ion.SPM() model.convert_to_format = "rust" @@ -1531,10 +1550,14 @@ def test_explicit_time_integral(self, solver_class, use_post_sum, use_output_var ) assert isinstance(sol["integral"].sensitivities["a"], np.ndarray) + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) @pytest.mark.parametrize("use_output_var", [False, True]) - def test_explicit_time_integral_initial_condition_sensitivity(self, use_output_var): + def test_explicit_time_integral_initial_condition_sensitivity( + self, use_output_var, convert_to_format + ): times = np.linspace(0, 1, 10) model = pybamm.BaseModel() + model.convert_to_format = convert_to_format c = pybamm.Variable("c") a = pybamm.InputParameter("a") model.rhs = {c: -a * c} @@ -1559,6 +1582,7 @@ def test_explicit_time_integral_initial_condition_sensitivity(self, use_output_v sol["integral"].sensitivities["a"], [expected], rtol=1e-3 ) + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) @pytest.mark.parametrize("use_post_sum", [False, True]) @pytest.mark.parametrize( ("t_later", "integral", "integral_sensitivity"), @@ -1566,9 +1590,10 @@ def test_explicit_time_integral_initial_condition_sensitivity(self, use_output_v ids=["shared-boundary", "gap"], ) def test_explicit_time_integral_sensitivity_over_joined_solutions( - self, t_later, integral, integral_sensitivity, use_post_sum + self, t_later, integral, integral_sensitivity, use_post_sum, convert_to_format ): model = pybamm.BaseModel() + model.convert_to_format = convert_to_format y = pybamm.Variable("y") a = pybamm.InputParameter("a") model.rhs = {y: 0 * y} @@ -1592,12 +1617,14 @@ def test_explicit_time_integral_sensitivity_over_joined_solutions( np.testing.assert_allclose(joined.entries, [value], rtol=1e-6) np.testing.assert_allclose(joined.sensitivities["a"], [sensitivity], rtol=1e-6) + @pytest.mark.parametrize("convert_to_format", ["casadi", "rust"]) @pytest.mark.parametrize("use_output_var", [False, True]) @pytest.mark.parametrize("sensitivity_name", ["a", "b"]) def test_explicit_time_integral_post_sum_sensitivity_to_one_input( - self, sensitivity_name, use_output_var + self, sensitivity_name, use_output_var, convert_to_format ): model = pybamm.BaseModel() + model.convert_to_format = convert_to_format c = pybamm.Variable("c") a = pybamm.InputParameter("a") b = pybamm.InputParameter("b") @@ -1626,6 +1653,7 @@ def test_explicit_time_integral_post_sum_sensitivity_to_one_input( "a": 2 * integral * 2.0 * dintegral_da, "b": integral**2, }[sensitivity_name] + np.testing.assert_allclose(sol["value"].entries, [integral**2 * 2.0], rtol=1e-6) sensitivities = sol["value"].sensitivities assert sorted(sensitivities) == sorted(["all", sensitivity_name]) np.testing.assert_allclose(sensitivities["all"], [[expected]], rtol=1e-6)