From 2da8a10aadeb231285a9a1bc0f5dc5be841a32bd Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Sun, 28 Jun 2026 22:07:40 +0000 Subject: [PATCH 01/10] feat: draft out a loss function solver --- .../expression_tree/operations/diffsl.py | 61 ++++-- .../src/pybamm/simulation/loss_function.py | 41 ++++ .../src/pybamm/simulation/loss_solver.py | 203 ++++++++++++++++++ .../pybamm/src/pybamm/solvers/base_solver.py | 1 - .../processed_variable_time_integral.py | 28 ++- .../pybamm/src/pybamm/solvers/solution.py | 4 +- .../integration/test_solvers/test_diffsl.py | 2 +- 7 files changed, 307 insertions(+), 33 deletions(-) create mode 100644 packages/pybamm/src/pybamm/simulation/loss_function.py create mode 100644 packages/pybamm/src/pybamm/simulation/loss_solver.py diff --git a/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py b/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py index 9c4ce80935..ae2dd8d1cb 100644 --- a/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py +++ b/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py @@ -50,6 +50,8 @@ def __init__( raise ValueError("float_precision must be a positive integer") self.float_precision = float_precision self._schedule_to_model_branch_order = None + self._input_names = None + self._nstates = None self._preprocess_model() def _preprocess_model(self) -> None: @@ -710,7 +712,8 @@ def to_diffeq(self, outputs: list[str]) -> str: new_line = "\n" # inputs, find all pybamm.InputParameters in the model - inputs = [vn for _, vn in self._collect_input_names(all_vars, outputs)] + self._input_names = self._collect_input_names(all_vars, outputs) + inputs = [vn for _, vn in self._input_names] if len(inputs) > 0: lines = ["in_i {"] @@ -795,6 +798,7 @@ def to_diffeq(self, outputs: list[str]) -> str: + new_line + "}" ) + self._nstates = start_index # diff of state vector u if not is_ode: @@ -1070,7 +1074,41 @@ def to_diffeq(self, outputs: list[str]) -> str: return "\n".join(all_lines) + "\n" - def map_inputs(self, inputs: dict, outputs: list[str] | None = None) -> np.ndarray: + def default_inputs(self) -> dict: + """ + Return a dict of default input values for the model. + + Returns + ------- + dict + PyBaMM-style parameter dict mapping parameter names to scalar values. + The keys are the original PyBaMM parameter names (e.g. ``"Lower voltage cut-off [V]"``). + """ + return {original_name: 1.0 for original_name, _ in self._input_names} + + def input_names(self) -> list[tuple[str, str]]: + """ + Return a list of input parameter names for the model. + + Returns + ------- + list[tuple[str, str]] + List of tuples containing the original PyBaMM parameter name and its diffsl-transformed form (e.g. ``("Lower voltage cut-off [V]", "lowervoltagecutoffv")``). + """ + return self._input_names + + def nstates(self) -> int | None: + """ + Return the number of states in the model. + + Returns + ------- + int + The number of states in the model. + """ + return self._nstates + + def map_inputs(self, inputs: dict) -> np.ndarray: """ Map a PyBaMM inputs dict to the ordered array expected by the DiffSL model. @@ -1099,24 +1137,7 @@ def map_inputs(self, inputs: dict, outputs: list[str] | None = None) -> np.ndarr KeyError If a required input parameter is not present in *inputs*. """ - if outputs is None: - outputs = [] - - model = self.model - - # Build all_vars for the output expressions (same logic as to_diffeq) - all_vars = self._all_vars.copy() - for out in outputs: - if out not in all_vars: - raise ValueError(f"output {out} not in model") - if model.symbol_processor: - try: - all_vars[out] = model.get_processed_variable(out) - except KeyError: # pragma: no cover - pass - - # Reconstruct the ordered input list the same way to_diffeq does - ordered_names = self._collect_input_names(all_vars, outputs) + ordered_names = self._input_names if not ordered_names: return np.array([], dtype=float) diff --git a/packages/pybamm/src/pybamm/simulation/loss_function.py b/packages/pybamm/src/pybamm/simulation/loss_function.py new file mode 100644 index 0000000000..cb74f115fb --- /dev/null +++ b/packages/pybamm/src/pybamm/simulation/loss_function.py @@ -0,0 +1,41 @@ +import numpy as np +import pandas as pd + +import pybamm + +bdp_to_pybamm_mapping = { + "Voltage / V": pybamm.Variable("Voltage [V]"), + "Current / A": pybamm.Variable("Current [A]"), +} + + +def _data_comparison(data: pd.DataFame): + data_times = data["Test Time / s"] + data_values_list = [] + variables_list = [] + for column in data.columns: + variable = bdp_to_pybamm_mapping.get(column, None) + if variable is not None: + variables_list.append(variable) + data_values_list.append(data[column]) + if not variables_list: + raise ValueError( + "No variables found in the data. Please ensure that the data " + "contains columns with any of the following names: " + f"{list(bdp_to_pybamm_mapping.keys())}" + ) + + data_values = np.hstack(data_values_list) + variables = pybamm.NumpyConcatenation(*variables_list) + + data_values = data["value"] + data = pybamm.DiscreteTimeData(data_times, data_values, "sum of squares data") + return data, variables + + +def sum_of_squares(data: pd.DataFrame): + """ + A method to create a loss function with a sum-of-squared-error loss function, given some data in BDF format (https://battery-data-alliance.github.io/battery-data-format/) to fit against. + """ + data, variables = _data_comparison(data) + return pybamm.DiscreteTimeSum((data - variables) ** 2) diff --git a/packages/pybamm/src/pybamm/simulation/loss_solver.py b/packages/pybamm/src/pybamm/simulation/loss_solver.py new file mode 100644 index 0000000000..bf36ef9e80 --- /dev/null +++ b/packages/pybamm/src/pybamm/simulation/loss_solver.py @@ -0,0 +1,203 @@ +from enum import Enum + +import casadi +import numpy as np +import pydiffsol as ds + +import pybamm + + +class LossSolver: + """ + A solver defined by a PyBaMM time-series model (i.e. dy/dt = f(y, t, p)) + and a scalar loss function L(p) = g(y(t), p). Has the functionality to: + (a) calculate the loss function for a given set of parameters p, by solving the ODE and evaluating L at the solution; + (c) calculate the solution of the ODE y(t) for a given set of parameters p, by solving the ODE; + (b) calculate the gradient of the loss function with respect to the parameters p, by: + (i) finitely-differencing the loss function with respect to p, which requires multiple ODE solves; or + (ii) solving the forward sensitivity equations, which requires a single ODE solve of an augmented system of equations. + (iii) solving the adjoint equations, which requires a single ODE solve of an augmented system of equations backwards in time. + (d) batching the above calculations across multiple parameter sets p + (c) pickle and unpickle the solver so it can be saved and loaded in multiple contexts (e.g. training and inference workflows) + By definition, the loss function can only vary with the parameters p, and it must depend on the solution y(t) for the loss to be well-defined, + so there are some restrictions on the form of the loss function g(y(t), p): + (1) g must include either a pybamm.DiscreteSumSum or a pybamm.ExplicitTimeIntegral node in the expression tree that integrates/sums over time and the solution y(t). + (2) any instances of the state variables or time in g must be contained within the scope of the aforementioned time-sum/integral node + (3) the output shape of g must be a scalar + For convenience, a sum-of-squared-error loss function factory method is provided to create a LossSolver from data in BDF format + """ + + INNER_LOSS_FUNCTION_NAME = "inner loss function" + + def __init__( + self, sim: pybamm.Simulation, loss_function: pybamm.Symbol, final_time: float + ): + self._sim = sim + self._processed_loss = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var( + loss_function, final_time + ) + if self._processed_loss is None: + raise ValueError( + "Loss function must contain either a DiscreteSum or an ExplicitTimeIntegral node" + ) + self._final_time = final_time + + # add inner function as the output of the model + if self.INNER_LOSS_FUNCTION_NAME in self._sim.model.variables: + raise ValueError( + f"Model already contains a variable named {self.INNER_LOSS_FUNCTION_NAME}. " + "Please rename the inner loss function to avoid conflicts." + ) + self._sim.model.variables[self.INNER_LOSS_FUNCTION_NAME] = self._inner + + self._exporter = pybamm.DiffSLExport(self._sim) + code = self._exporter.to_diffeq([self.INNER_LOSS_FUNCTION_NAME]) + self._ode = ds.Ode( + code, + matrix_type=ds.faer_sparse, + scalar_type=ds.f64, + linear_solver=ds.lu, + ode_solver=ds.bdf, + ) + self._ode.integrate_out = self._processed_loss.method == "continuous" + self._ode.rtol = self._sim.solver.rtol + self._ode.atol = self._sim.solver.atol + + # generate casadi functions for post-sum node and its sensitivities + inputs = self._exporter.default_inputs() + post_sum_node = self._processed_loss.post_sum_node + (self._post_sum, self._post_sum_sens) = self._post_sum( + post_sum_node, inputs, self._exporter.input_names() + ) + + def _post_sum(self, var_pybamm, inputs, input_names): + t_casadi = casadi.MX.sym("t") + sum_casadi = casadi.MX.sym("sum", 1) + p_casadi = {name: casadi.MX.sym(name, value) for name, value in inputs.items()} + post_sum_casadi = var_pybamm.to_casadi(t_casadi, sum_casadi, inputs=p_casadi) + + p_casadi_stacked = casadi.vertcat(*[p_casadi[name] for name, _ in input_names]) + sens_casadi = casadi.MX.sym("sens", len(input_names)) + dpost_dy = casadi.jacobian(post_sum_casadi, sum_casadi) + dpost_dp = casadi.jacobian(post_sum_casadi, p_casadi_stacked) + sens = dpost_dy * sens_casadi + dpost_dp + post_sum_sens_casadi = casadi.Function( + "sens_fun", + [t_casadi, sum_casadi, p_casadi, sens_casadi], + [sens], + ) + return post_sum_casadi, post_sum_sens_casadi + + def inputs_to_parameters(self, inputs: list[dict]) -> np.ndarray: + """Converts a standard set of pybamm input dictionaries to a 2D parameter array (n_batch, n_params) for use in the functions below.""" + # TODO: add batching + return self._exporter.map_inputs( + inputs, outputs=[self.INNER_LOSS_FUNCTION_NAME] + ) + + def parameters_to_inputs(self, p: np.ndarray) -> list[dict]: + """Converts a 2D parameter array (n_batch, n_params) to a standard set of pybamm input dictionaries.""" + # TODO: implement inverse_map_inputs in DiffSLExporter to allow this to work + # TODO: add batching + return self._exporter.inverse_map_inputs( + p, outputs=[self.INNER_LOSS_FUNCTION_NAME] + ) + + def predict(self, p: np.ndarray) -> list[pybamm.Solution]: + """Calculate the solution of the ODE for each set of parameters in inputs.""" + # TODO: add batching + inputs = self.parameters_to_inputs(p) + return self._sim.solve(inputs=inputs) + + def _discrete_sum_to_loss(self, sol: ds.Solution, inputs: dict) -> np.ndarray: + """Calculate the loss function for a discrete sum loss function.""" + the_integral = np.sum(sol.ys, axis=1) + if self.post_sum_node is None: + ret = the_integral + else: + ret = self._post_sum(0.0, the_integral, inputs).full() + return ret + + def _explicit_time_integral_to_loss( + self, sol: ds.Solution, inputs: dict + ) -> np.ndarray: + """Calculate the loss function for an explicit time integral loss function.""" + the_integral = sol.ys[:, -1] + if self.post_sum_node is None: + ret = the_integral + else: + ret = self._post_sum(0.0, the_integral, inputs).full() + return ret + + def loss(self, p: np.ndarray) -> np.ndarray: + """ + Calculate the loss function for each set of parameters in inputs. + Returns a 1D array of loss values of length n_batch. + """ + # TODO: add batching + if self._processed_loss.method == "discrete": + sol = self._ode.solve_dense(p, self._processed_loss.discrete_times) + return self._discrete_sum_to_loss(sol, self.parameters_to_inputs(p)) + elif self._processed_loss.method == "continuous": + sol = self._ode.solve(p, self._final_time) + return self._explicit_time_integral_to_loss( + sol, self.parameters_to_inputs(p) + ) + + def finite_difference_gradient(self, p: np.ndarray, h: float = 1e-5) -> np.ndarray: + """ + Calculate the gradient of the loss function with respect to the parameters for each set of parameters in inputs using finite differencing. + Returns a 2D array of gradients, with shape (n_batch, n_params), where each row corresponds to the gradient for a given input parameter set. + """ + raise NotImplementedError("LossSolver is not yet implemented") + + def _discrete_sum_to_gradient( + self, sol: ds.Solution, inputs: dict + ) -> tuple[np.ndarray, np.ndarray]: + ys_sum = np.sum(sol.ys, axis=1) + sens_sum = np.array([np.sum(s, axis=1) for s in sol.sens]) + if self.post_sum_node is None: + return ys_sum, sens_sum + else: + loss = self.post_sum_node.evaluate(0.0, ys_sum, None, inputs) + gradient = self._post_sum_sens(0.0, ys_sum, inputs, sens_sum) + return loss, gradient + + def loss_and_gradient( + self, p: np.ndarray, mode: "LossSolverGradientMode" + ) -> tuple[np.ndarray, np.ndarray]: + """ + Calculate the loss and gradient of the loss function with respect to the parameters for each set of parameters in inputs. + Returns a tuple of arrays, where the first contains the loss values as a 1D array of length n_batch and the second contains the gradients as a 2D array of shape (n_batch, n_params) + """ + # TODO: add batching + if self._processed_loss.method == "discrete": + times = self._processed_loss.discrete_times + if mode == self.LossSolverGradientMode.FORWARD_SENSITIVITY: + sol = self._ode.solve_fwd_sens(p, times) + return self._discrete_sum_to_gradient(sol, self.parameters_to_inputs(p)) + else: + raise NotImplementedError( + "Adjoint sensitivity for discrete sum is not yet implemented" + ) + + elif self._processed_loss.method == "continuous": + if mode == self.LossSolverGradientMode.FORWARD_SENSITIVITY: + raise NotImplementedError( + "Forward sensitivity for explicit time integral is not yet implemented" + ) + else: + integral, integral_sens = self._ode.solve_continuous_adjoint( + p, self._final_time + ) + if self._post_sum is None: + return integral, integral_sens + else: + inputs = self.parameters_to_inputs(p) + loss = self._post_sum(0.0, integral, inputs) + gradient = self._post_sum_sens(0.0, integral, inputs, integral_sens) + return loss, gradient + + class LossSolverGradientMode(Enum): + FORWARD_SENSITIVITY = "forward_sensitivity" + ADJOINT_SENSITIVITY = "adjoint_sensitivity" diff --git a/packages/pybamm/src/pybamm/solvers/base_solver.py b/packages/pybamm/src/pybamm/solvers/base_solver.py index 3d53af9890..761b5d2ea8 100644 --- a/packages/pybamm/src/pybamm/solvers/base_solver.py +++ b/packages/pybamm/src/pybamm/solvers/base_solver.py @@ -369,7 +369,6 @@ def set_up( processed_time_integral = ( pybamm.ProcessedVariableTimeIntegral.from_pybamm_var( model.get_processed_variable_or_event(key), - model.len_rhs_and_alg, ) ) # We will evaluate the sum node in the solver and sum it afterwards 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 48db840e01..16982914ec 100644 --- a/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py +++ b/packages/pybamm/src/pybamm/solvers/processed_variable_time_integral.py @@ -19,6 +19,7 @@ class ProcessedVariableTimeIntegral: discrete_times: npt.NDArray[np.float64] | None post_sum_node: pybamm.Symbol | None = None post_sum: casadi.Function | None = None + sens_fun: casadi.Function | None = None def postfix_sum(self, entries, t_pts) -> np.ndarray: if self.method == "discrete": @@ -70,9 +71,14 @@ def postfix_sensitivities( the_integral = self.postfix_sum(sensitivities, t_pts) if self.post_sum_node is None: return the_integral + sens_fun = self.generate_sens_fun(inputs, the_integral.shape) + inputs_stacked = casadi.vertcat(*[v for v in inputs.values()]) + sens_values = sens_fun(0.0, entries, inputs_stacked, the_integral) + return sens_values.full() - y_casadi = casadi.MX.sym("y", entries.shape[0]) - sens_casadi = casadi.MX.sym("s_var", the_integral.shape) + def generate_sens_fun(self, inputs, sens_shape): + y_casadi = casadi.MX.sym("y", sens_shape[0]) + sens_casadi = casadi.MX.sym("s_var", sens_shape) t_casadi = casadi.MX.sym("t") p_casadi = { name: casadi.MX.sym( @@ -81,21 +87,22 @@ def postfix_sensitivities( for name, value in inputs.items() } p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) - inputs_stacked = casadi.vertcat(*[v for v in inputs.values()]) post_sum_casadi = self.post_sum_node.to_casadi( t_casadi, y_casadi, inputs=p_casadi ) + # post = (y) + # dpost_dy = (y, y) + # dpost_dp = (p, y) dpost_dy = casadi.jacobian(post_sum_casadi, y_casadi) dpost_dp = casadi.jacobian(post_sum_casadi, p_casadi_stacked) + # sens = (y, y) @ (s_var, y) + (p, y) = (y) sens = dpost_dy @ sens_casadi + dpost_dp - sens_fun = casadi.Function( + return casadi.Function( "sens_fun", [t_casadi, y_casadi, p_casadi_stacked, sens_casadi], [sens], ) - sens_values = sens_fun(0.0, entries, inputs_stacked, the_integral) - return sens_values.full() @staticmethod def to_post_sum_expr( @@ -123,7 +130,7 @@ def to_post_sum_expr( @staticmethod def from_pybamm_var( var: pybamm.Symbol, - nstates: int, + final_time: float | None = None, ) -> ProcessedVariableTimeIntegral | None: sum_node = None for symbol in var.pre_order(): @@ -145,12 +152,17 @@ def from_pybamm_var( var, sum_node, sum_y_len ) if isinstance(sum_node, pybamm.DiscreteTimeSum): + discrete_times = ( + sum_node.sum_times[sum_node.sum_times <= final_time] + if final_time is not None + else sum_node.sum_times + ) return ProcessedVariableTimeIntegral( method="discrete", post_sum_node=post_sum_node, sum_node=sum_node, initial_condition=0.0, - discrete_times=sum_node.sum_times, + discrete_times=discrete_times, ) elif isinstance(sum_node, pybamm.ExplicitTimeIntegral): if isinstance(sum_node.initial_condition, pybamm.Symbol): diff --git a/packages/pybamm/src/pybamm/solvers/solution.py b/packages/pybamm/src/pybamm/solvers/solution.py index 3a5119a6df..ac238faa52 100644 --- a/packages/pybamm/src/pybamm/solvers/solution.py +++ b/packages/pybamm/src/pybamm/solvers/solution.py @@ -791,9 +791,7 @@ def observe(self, symbol: pybamm.Symbol) -> pybamm.ProcessedVariable: return self[name] def _convert_to_casadi(self, var_pybamm, inputs, ys_shape): - time_integral = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var( - var_pybamm, ys_shape[0] - ) + time_integral = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var(var_pybamm) if time_integral is not None: var_pybamm = time_integral.sum_node.child if time_integral.post_sum_node is not None: diff --git a/packages/pybamm/tests/integration/test_solvers/test_diffsl.py b/packages/pybamm/tests/integration/test_solvers/test_diffsl.py index 97dc806c12..fa9c54a6b6 100644 --- a/packages/pybamm/tests/integration/test_solvers/test_diffsl.py +++ b/packages/pybamm/tests/integration/test_solvers/test_diffsl.py @@ -212,7 +212,7 @@ def test_simulations(self, model, inputs, experiment, request, timing_results): map_inputs_dict = dict(pv_inputs) if experiment is not None: map_inputs_dict["Ambient temperature [K]"] = pv["Ambient temperature [K]"] - ds_inputs = exporter.map_inputs(map_inputs_dict, outputs=[output_variable]) + ds_inputs = exporter.map_inputs(map_inputs_dict) logger = logging.getLogger() logger.info(f"DiffSL export time: {time.perf_counter() - t0:.5f} seconds") From 1ef43caadfdf741ab2ac267a069b518c4e5ed5a0 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Mon, 29 Jun 2026 22:28:40 +0000 Subject: [PATCH 02/10] feat: fill out loss solver impl --- .../expression_tree/operations/diffsl.py | 42 ++- .../src/pybamm/simulation/loss_solver.py | 333 ++++++++++++------ .../integration/test_solvers/test_diffsl.py | 2 +- .../test_operations/test_diffsl_export.py | 17 +- .../unit/test_solvers/test_loss_solver.py | 280 +++++++++++++++ 5 files changed, 563 insertions(+), 111 deletions(-) create mode 100644 packages/pybamm/tests/unit/test_solvers/test_loss_solver.py diff --git a/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py b/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py index ae2dd8d1cb..deaef55b2c 100644 --- a/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py +++ b/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py @@ -1112,8 +1112,9 @@ def map_inputs(self, inputs: dict) -> np.ndarray: """ Map a PyBaMM inputs dict to the ordered array expected by the DiffSL model. - The ordering matches the ``in_i {}`` block produced by :meth:`to_diffeq`. - Each key in ``inputs`` may be either the original PyBaMM parameter name + The ordering matches the ``in_i {}`` block produced by the most recent + call to :meth:`to_diffeq`, which must be called first. Each key in + ``inputs`` may be either the original PyBaMM parameter name (e.g. ``"Lower voltage cut-off [V]"``) or its DiffSL-transformed form (e.g. ``"lowervoltagecutoffv"``). @@ -1121,9 +1122,6 @@ def map_inputs(self, inputs: dict) -> np.ndarray: ---------- inputs : dict PyBaMM-style parameter dict mapping parameter names to scalar values. - outputs : list[str] or None, optional - Output variable names, used to scan for ``InputParameter`` nodes that - appear only inside output expressions. Defaults to an empty list. Returns ------- @@ -1156,6 +1154,40 @@ def map_inputs(self, inputs: dict) -> np.ndarray: ) return np.array(values, dtype=float) + def inverse_map_inputs(self, values: np.ndarray) -> dict: + """ + Map an ordered DiffSL parameter array back to a PyBaMM inputs dict. + + This is the inverse of :meth:`map_inputs`: the ``i``-th value is assigned + to the ``i``-th input parameter in the order of the ``in_i {}`` block + produced by :meth:`to_diffeq`. + + Parameters + ---------- + values : array_like + 1-D array of parameter values ordered to match the ``in_i {}`` block. + + Returns + ------- + dict + PyBaMM-style parameter dict keyed by the original parameter names. + + Raises + ------ + ValueError + If the number of values does not match the number of input parameters. + """ + ordered_names = self._input_names + values = np.asarray(values, dtype=float).reshape(-1) + if len(values) != len(ordered_names): + raise ValueError( + f"Expected {len(ordered_names)} parameter values, got {len(values)}." + ) + return { + original_name: float(value) + for (original_name, _), value in zip(ordered_names, values, strict=True) + } + def _equation_to_diffeq( equation: pybamm.Symbol, diff --git a/packages/pybamm/src/pybamm/simulation/loss_solver.py b/packages/pybamm/src/pybamm/simulation/loss_solver.py index bf36ef9e80..f23bb6e4d8 100644 --- a/packages/pybamm/src/pybamm/simulation/loss_solver.py +++ b/packages/pybamm/src/pybamm/simulation/loss_solver.py @@ -1,8 +1,10 @@ +import concurrent.futures +import multiprocessing +import pickle from enum import Enum import casadi import numpy as np -import pydiffsol as ds import pybamm @@ -21,18 +23,41 @@ class LossSolver: (c) pickle and unpickle the solver so it can be saved and loaded in multiple contexts (e.g. training and inference workflows) By definition, the loss function can only vary with the parameters p, and it must depend on the solution y(t) for the loss to be well-defined, so there are some restrictions on the form of the loss function g(y(t), p): - (1) g must include either a pybamm.DiscreteSumSum or a pybamm.ExplicitTimeIntegral node in the expression tree that integrates/sums over time and the solution y(t). + (1) g must include either a pybamm.DiscreteTimeSum or a pybamm.ExplicitTimeIntegral node in the expression tree that integrates/sums over time and the solution y(t). (2) any instances of the state variables or time in g must be contained within the scope of the aforementioned time-sum/integral node (3) the output shape of g must be a scalar For convenience, a sum-of-squared-error loss function factory method is provided to create a LossSolver from data in BDF format + + Parameters + ---------- + sim : pybamm.Simulation + The simulation wrapping the time-series model. + loss_function : pybamm.Symbol + The scalar loss function expression. + final_time : float + The final time to integrate the model to. + max_workers : int or None, optional + If greater than 1, the batched methods (``loss``, ``loss_and_gradient``, + ``finite_difference_gradient``) distribute the per-parameter-set solves + across a process pool of this many workers. Defaults to ``None`` + (sequential). ``predict`` is unaffected; its parallelism comes from the + wrapped solver's ``num_threads`` option. """ INNER_LOSS_FUNCTION_NAME = "inner loss function" def __init__( - self, sim: pybamm.Simulation, loss_function: pybamm.Symbol, final_time: float + self, + sim: pybamm.Simulation, + loss_function: pybamm.Symbol, + final_time: float, + max_workers: int | None = None, ): + ds = pybamm.import_optional_dependency("pydiffsol") + self._sim = sim + self._max_workers = max_workers + self._pool = None self._processed_loss = pybamm.ProcessedVariableTimeIntegral.from_pybamm_var( loss_function, final_time ) @@ -42,6 +67,9 @@ def __init__( ) self._final_time = final_time + # the inner function is the integrand that is summed/integrated over time + self._inner = self._processed_loss.sum_node.orphans[0] + # add inner function as the output of the model if self.INNER_LOSS_FUNCTION_NAME in self._sim.model.variables: raise ValueError( @@ -63,105 +91,163 @@ def __init__( self._ode.rtol = self._sim.solver.rtol self._ode.atol = self._sim.solver.atol - # generate casadi functions for post-sum node and its sensitivities - inputs = self._exporter.default_inputs() - post_sum_node = self._processed_loss.post_sum_node - (self._post_sum, self._post_sum_sens) = self._post_sum( - post_sum_node, inputs, self._exporter.input_names() + # generate casadi functions for the post-sum node and its sensitivities + self._post_sum, self._post_sum_sens = self._make_post_sum_functions( + self._processed_loss.post_sum_node ) - def _post_sum(self, var_pybamm, inputs, input_names): + # set up the process pool for the batched methods, if requested + if self._max_workers is not None and self._max_workers > 1: + # spawn (not fork) avoids deadlocks with the native threadpools used + # by the matrix backend; pickling self ships the compiled Ode cheaply. + self._pool = concurrent.futures.ProcessPoolExecutor( + max_workers=self._max_workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=_init_worker, + initargs=(pickle.dumps(self),), + ) + + def _make_post_sum_functions( + self, post_sum_node: pybamm.Symbol | None + ) -> tuple[casadi.Function | None, casadi.Function | None]: + """Build casadi functions for the post-sum node and its sensitivities.""" + if post_sum_node is None: + return None, None + + input_names = self._exporter.input_names() + sum_size = self._processed_loss.sum_node.evaluate_for_shape().shape[0] + n_params = len(input_names) + t_casadi = casadi.MX.sym("t") - sum_casadi = casadi.MX.sym("sum", 1) - p_casadi = {name: casadi.MX.sym(name, value) for name, value in inputs.items()} - post_sum_casadi = var_pybamm.to_casadi(t_casadi, sum_casadi, inputs=p_casadi) + sum_casadi = casadi.MX.sym("sum", sum_size) + p_casadi = { + original: casadi.MX.sym(variable, 1) for original, variable in input_names + } + post_sum_casadi = post_sum_node.to_casadi(t_casadi, sum_casadi, inputs=p_casadi) + p_casadi_stacked = casadi.vertcat( + *[p_casadi[original] for original, _ in input_names] + ) + + post_sum = casadi.Function( + "post_sum", + [t_casadi, sum_casadi, p_casadi_stacked], + [post_sum_casadi], + ) - p_casadi_stacked = casadi.vertcat(*[p_casadi[name] for name, _ in input_names]) - sens_casadi = casadi.MX.sym("sens", len(input_names)) + sens_casadi = casadi.MX.sym("sens", sum_size, n_params) dpost_dy = casadi.jacobian(post_sum_casadi, sum_casadi) dpost_dp = casadi.jacobian(post_sum_casadi, p_casadi_stacked) - sens = dpost_dy * sens_casadi + dpost_dp - post_sum_sens_casadi = casadi.Function( - "sens_fun", - [t_casadi, sum_casadi, p_casadi, sens_casadi], + sens = dpost_dy @ sens_casadi + dpost_dp + post_sum_sens = casadi.Function( + "post_sum_sens", + [t_casadi, sum_casadi, p_casadi_stacked, sens_casadi], [sens], ) - return post_sum_casadi, post_sum_sens_casadi + return post_sum, post_sum_sens def inputs_to_parameters(self, inputs: list[dict]) -> np.ndarray: """Converts a standard set of pybamm input dictionaries to a 2D parameter array (n_batch, n_params) for use in the functions below.""" - # TODO: add batching - return self._exporter.map_inputs( - inputs, outputs=[self.INNER_LOSS_FUNCTION_NAME] - ) + return np.array([self._exporter.map_inputs(single) for single in inputs]) def parameters_to_inputs(self, p: np.ndarray) -> list[dict]: """Converts a 2D parameter array (n_batch, n_params) to a standard set of pybamm input dictionaries.""" - # TODO: implement inverse_map_inputs in DiffSLExporter to allow this to work - # TODO: add batching - return self._exporter.inverse_map_inputs( - p, outputs=[self.INNER_LOSS_FUNCTION_NAME] - ) + return [self._exporter.inverse_map_inputs(row) for row in np.atleast_2d(p)] def predict(self, p: np.ndarray) -> list[pybamm.Solution]: - """Calculate the solution of the ODE for each set of parameters in inputs.""" - # TODO: add batching - inputs = self.parameters_to_inputs(p) - return self._sim.solve(inputs=inputs) - - def _discrete_sum_to_loss(self, sol: ds.Solution, inputs: dict) -> np.ndarray: - """Calculate the loss function for a discrete sum loss function.""" - the_integral = np.sum(sol.ys, axis=1) - if self.post_sum_node is None: - ret = the_integral - else: - ret = self._post_sum(0.0, the_integral, inputs).full() - return ret + """Calculate the solution of the ODE for each set of parameters in inputs. - def _explicit_time_integral_to_loss( - self, sol: ds.Solution, inputs: dict + The parameter sets are solved together in a single batched solve, so any + ``num_threads`` parallelism configured on the wrapped solver is used. + """ + solutions = self._sim.solve( + t_eval=[0, self._final_time], inputs=self.parameters_to_inputs(p) + ) + # a single-element batch is returned as a bare Solution + if not isinstance(solutions, list): + solutions = [solutions] + return solutions + + def _apply_post_sum( + self, the_integral: np.ndarray, params: np.ndarray ) -> np.ndarray: - """Calculate the loss function for an explicit time integral loss function.""" - the_integral = sol.ys[:, -1] - if self.post_sum_node is None: - ret = the_integral - else: - ret = self._post_sum(0.0, the_integral, inputs).full() - return ret + """Apply the post-sum node (if any) to the summed/integrated inner output. - def loss(self, p: np.ndarray) -> np.ndarray: + ``params`` is the ordered parameter row, which is exactly the stacked + input vector expected by the post-sum casadi function. """ - Calculate the loss function for each set of parameters in inputs. - Returns a 1D array of loss values of length n_batch. - """ - # TODO: add batching + if self._post_sum is None: + return the_integral + return self._post_sum(0.0, the_integral, params).full().reshape(-1) + + def _single_loss(self, params: np.ndarray) -> float: + """Calculate the scalar loss for a single parameter set.""" + params = np.asarray(params, dtype=float) if self._processed_loss.method == "discrete": - sol = self._ode.solve_dense(p, self._processed_loss.discrete_times) - return self._discrete_sum_to_loss(sol, self.parameters_to_inputs(p)) - elif self._processed_loss.method == "continuous": - sol = self._ode.solve(p, self._final_time) - return self._explicit_time_integral_to_loss( - sol, self.parameters_to_inputs(p) + sol = self._ode.solve_dense(params, self._processed_loss.discrete_times) + the_integral = np.sum(sol.ys, axis=1) + else: + sol = self._ode.solve(params, self._final_time) + the_integral = sol.ys[:, -1] + value = self._apply_post_sum(the_integral, params) + return float(np.asarray(value).reshape(-1)[0]) + + def _single_loss_and_gradient( + self, params: np.ndarray, mode: "LossSolverGradientMode" + ) -> tuple[float, np.ndarray]: + """Calculate the scalar loss and (n_params,) gradient for a single parameter set.""" + params = np.asarray(params, dtype=float) + if self._processed_loss.method == "discrete": + if mode == self.LossSolverGradientMode.FORWARD_SENSITIVITY: + sol = self._ode.solve_fwd_sens( + params, self._processed_loss.discrete_times + ) + return self._discrete_sum_to_gradient(sol, params) + raise NotImplementedError( + "Adjoint sensitivity for discrete sum is not yet implemented" ) - def finite_difference_gradient(self, p: np.ndarray, h: float = 1e-5) -> np.ndarray: - """ - Calculate the gradient of the loss function with respect to the parameters for each set of parameters in inputs using finite differencing. - Returns a 2D array of gradients, with shape (n_batch, n_params), where each row corresponds to the gradient for a given input parameter set. - """ - raise NotImplementedError("LossSolver is not yet implemented") + if mode == self.LossSolverGradientMode.FORWARD_SENSITIVITY: + raise NotImplementedError( + "Forward sensitivity for explicit time integral is not yet implemented" + ) + integral, integral_sens = self._ode.solve_continuous_adjoint( + params, self._final_time + ) + if self._post_sum is None: + loss = float(np.asarray(integral).reshape(-1)[0]) + gradient = np.asarray(integral_sens).reshape(-1) + return loss, gradient + loss = float(self._post_sum(0.0, integral, params).full().reshape(-1)[0]) + gradient = ( + self._post_sum_sens(0.0, integral, params, integral_sens).full().reshape(-1) + ) + return loss, gradient def _discrete_sum_to_gradient( - self, sol: ds.Solution, inputs: dict - ) -> tuple[np.ndarray, np.ndarray]: + self, sol, params: np.ndarray + ) -> tuple[float, np.ndarray]: ys_sum = np.sum(sol.ys, axis=1) sens_sum = np.array([np.sum(s, axis=1) for s in sol.sens]) - if self.post_sum_node is None: - return ys_sum, sens_sum + if self._post_sum is None: + loss = float(np.asarray(ys_sum).reshape(-1)[0]) + return loss, sens_sum.reshape(-1) + loss = float(self._post_sum(0.0, ys_sum, params).full().reshape(-1)[0]) + gradient = ( + self._post_sum_sens(0.0, ys_sum, params, sens_sum.T).full().reshape(-1) + ) + return loss, gradient + + def loss(self, p: np.ndarray) -> np.ndarray: + """ + Calculate the loss function for each set of parameters in inputs. + Returns a 1D array of loss values of length n_batch. + """ + rows = list(np.atleast_2d(p)) + if self._pool is not None: + values = list(self._pool.map(_worker_loss, rows)) else: - loss = self.post_sum_node.evaluate(0.0, ys_sum, None, inputs) - gradient = self._post_sum_sens(0.0, ys_sum, inputs, sens_sum) - return loss, gradient + values = [self._single_loss(row) for row in rows] + return np.array(values) def loss_and_gradient( self, p: np.ndarray, mode: "LossSolverGradientMode" @@ -170,34 +256,81 @@ def loss_and_gradient( Calculate the loss and gradient of the loss function with respect to the parameters for each set of parameters in inputs. Returns a tuple of arrays, where the first contains the loss values as a 1D array of length n_batch and the second contains the gradients as a 2D array of shape (n_batch, n_params) """ - # TODO: add batching - if self._processed_loss.method == "discrete": - times = self._processed_loss.discrete_times - if mode == self.LossSolverGradientMode.FORWARD_SENSITIVITY: - sol = self._ode.solve_fwd_sens(p, times) - return self._discrete_sum_to_gradient(sol, self.parameters_to_inputs(p)) - else: - raise NotImplementedError( - "Adjoint sensitivity for discrete sum is not yet implemented" - ) + rows = list(np.atleast_2d(p)) + if self._pool is not None: + items = [(row, mode.value) for row in rows] + results = list(self._pool.map(_worker_loss_and_gradient, items)) + else: + results = [self._single_loss_and_gradient(row, mode) for row in rows] + losses = np.array([loss for loss, _ in results]) + gradients = np.array([gradient for _, gradient in results]) + return losses, gradients - elif self._processed_loss.method == "continuous": - if mode == self.LossSolverGradientMode.FORWARD_SENSITIVITY: - raise NotImplementedError( - "Forward sensitivity for explicit time integral is not yet implemented" - ) - else: - integral, integral_sens = self._ode.solve_continuous_adjoint( - p, self._final_time - ) - if self._post_sum is None: - return integral, integral_sens - else: - inputs = self.parameters_to_inputs(p) - loss = self._post_sum(0.0, integral, inputs) - gradient = self._post_sum_sens(0.0, integral, inputs, integral_sens) - return loss, gradient + def finite_difference_gradient(self, p: np.ndarray, h: float = 1e-5) -> np.ndarray: + """ + Calculate the gradient of the loss function with respect to the parameters for each set of parameters in inputs using finite differencing. + Returns a 2D array of gradients, with shape (n_batch, n_params), where each row corresponds to the gradient for a given input parameter set. + """ + p = np.atleast_2d(p).astype(float) + n_batch, n_params = p.shape + gradient = np.zeros((n_batch, n_params)) + for j in range(n_params): + p_plus = p.copy() + p_plus[:, j] += h + p_minus = p.copy() + p_minus[:, j] -= h + gradient[:, j] = (self.loss(p_plus) - self.loss(p_minus)) / (2 * h) + return gradient + + def close(self) -> None: + """Shut down the process pool, if one was created.""" + if self._pool is not None: + self._pool.shutdown(wait=True) + self._pool = None + + def __enter__(self) -> "LossSolver": + return self + + def __exit__(self, *exc) -> None: + self.close() + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + # the live process pool and casadi functions are not picklable; the + # unpickled copy runs sequentially (this also keeps the copies sent to + # worker processes pool-less), and the casadi functions are rebuilt + state["_pool"] = None + state["_post_sum_node_present"] = self._post_sum is not None + state["_post_sum"] = None + state["_post_sum_sens"] = None + return state + + def __setstate__(self, state: dict) -> None: + rebuild = state.pop("_post_sum_node_present", False) + self.__dict__.update(state) + if rebuild: + self._post_sum, self._post_sum_sens = self._make_post_sum_functions( + self._processed_loss.post_sum_node + ) class LossSolverGradientMode(Enum): FORWARD_SENSITIVITY = "forward_sensitivity" ADJOINT_SENSITIVITY = "adjoint_sensitivity" + + +_WORKER_SOLVER = None + + +def _init_worker(solver_bytes: bytes) -> None: + global _WORKER_SOLVER + _WORKER_SOLVER = pickle.loads(solver_bytes) + + +def _worker_loss(params: np.ndarray) -> float: + return _WORKER_SOLVER._single_loss(params) + + +def _worker_loss_and_gradient(item: tuple) -> tuple[float, np.ndarray]: + params, mode_value = item + mode = type(_WORKER_SOLVER).LossSolverGradientMode(mode_value) + return _WORKER_SOLVER._single_loss_and_gradient(params, mode) diff --git a/packages/pybamm/tests/integration/test_solvers/test_diffsl.py b/packages/pybamm/tests/integration/test_solvers/test_diffsl.py index fa9c54a6b6..57a8e3e35f 100644 --- a/packages/pybamm/tests/integration/test_solvers/test_diffsl.py +++ b/packages/pybamm/tests/integration/test_solvers/test_diffsl.py @@ -102,7 +102,7 @@ def test_models(self, model, inputs, request, timing_results): t0 = time.perf_counter() exporter = pybamm.DiffSLExport(model_disc) diffsl_code = exporter.to_diffeq(outputs=[output_variable]) - ds_inputs = exporter.map_inputs(pv_inputs, outputs=[output_variable]) + ds_inputs = exporter.map_inputs(pv_inputs) logger.info(f"DiffSL export time: {time.perf_counter() - t0:.5f} seconds") t0 = time.perf_counter() diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py index 867bec8942..04ea03075d 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py @@ -349,6 +349,7 @@ def test_output_specific_input_parameter(self): def test_map_inputs_basic(self, model): exporter = pybamm.DiffSLExport(model) + exporter.to_diffeq(outputs=["x"]) result = exporter.map_inputs({"p": 3.14}) assert isinstance(result, np.ndarray) assert result.shape == (1,) @@ -356,7 +357,8 @@ def test_map_inputs_basic(self, model): def test_map_inputs_no_outputs(self, model): exporter = pybamm.DiffSLExport(model) - result = exporter.map_inputs({"p": 1.0}, outputs=None) + exporter.to_diffeq(outputs=["x"]) + result = exporter.map_inputs({"p": 1.0}) assert result[0] == 1.0 def test_map_inputs_empty_model(self): @@ -372,6 +374,7 @@ def test_map_inputs_empty_model(self): def test_map_inputs_missing_key_raises(self, model): exporter = pybamm.DiffSLExport(model) + exporter.to_diffeq(outputs=["x"]) with pytest.raises(KeyError, match="not found in inputs dict"): exporter.map_inputs({}) @@ -385,7 +388,8 @@ def test_map_inputs_output_specific(self): disc = pybamm.Discretisation() disc.process_model(model) exporter = pybamm.DiffSLExport(model) - result = exporter.map_inputs({"extra_param": 2.0}, outputs=["extra_out"]) + exporter.to_diffeq(outputs=["extra_out"]) + result = exporter.map_inputs({"extra_param": 2.0}) assert result[0] == 2.0 def test_reg_power_with_non_scalar_exponent(self): @@ -404,11 +408,12 @@ def test_reg_power_with_non_scalar_exponent(self): def test_map_inputs_invalid_output_raises(self, model): exporter = pybamm.DiffSLExport(model) with pytest.raises(ValueError, match="output nonexistent not in model"): - exporter.map_inputs({}, outputs=["nonexistent"]) + exporter.to_diffeq(outputs=["nonexistent"]) def test_map_inputs_processed_variable_path(self, model): exporter = pybamm.DiffSLExport(model) - result = exporter.map_inputs({"p": 1.0}, outputs=["x"]) + exporter.to_diffeq(outputs=["x"]) + result = exporter.map_inputs({"p": 1.0}) assert result[0] == 1.0 def test_map_inputs_with_symbol_processor(self): @@ -418,7 +423,8 @@ def test_map_inputs_with_symbol_processor(self): ) sim.build() exporter = pybamm.DiffSLExport(sim) - result = exporter.map_inputs({}, outputs=["Terminal voltage [V]"]) + exporter.to_diffeq(outputs=["Terminal voltage [V]"]) + result = exporter.map_inputs({}) assert len(result) >= 0 def test_map_inputs_diffsl_transformed_name(self): @@ -431,6 +437,7 @@ def test_map_inputs_diffsl_transformed_name(self): disc = pybamm.Discretisation() disc.process_model(model) exporter = pybamm.DiffSLExport(model) + exporter.to_diffeq(outputs=["x"]) result = exporter.map_inputs({"testparam": 2.0}) assert result[0] == 2.0 diff --git a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py new file mode 100644 index 0000000000..f1b765896c --- /dev/null +++ b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py @@ -0,0 +1,280 @@ +# +# Tests for the LossSolver class +# +import importlib.util +import pickle + +import numpy as np +import pytest + +import pybamm +from pybamm.simulation.loss_solver import LossSolver + +has_pydiffsol = importlib.util.find_spec("pydiffsol") is not None + +K_TRUE = 0.5 +K_OTHER = 0.8 +FINAL_TIME = 2.0 +N_DATA = 11 +Y0 = 1.0 + +discrete_not_supported = pytest.mark.xfail( + reason="DiffSLExport cannot export the DiscreteTimeData interpolant; " + "the discrete sum LossSolver path is not yet supported", +) + + +def _decay_model(): + """dy/dt = -k * y, y(0) = 1, with analytic solution y(t) = exp(-k * t).""" + model = pybamm.BaseModel("exponential decay") + y = pybamm.Variable("y") + k = pybamm.InputParameter("k") + model.rhs = {y: -k * y} + model.initial_conditions = {y: pybamm.Scalar(Y0)} + model.variables = {"y": y} + return model + + +def _data_times(): + return np.linspace(0, FINAL_TIME, N_DATA) + + +def _analytic_solution(k, t): + return Y0 * np.exp(-k * np.asarray(t, dtype=float)) + + +def _discrete_loss(k): + t = _data_times() + residual = _analytic_solution(k, t) - _analytic_solution(K_TRUE, t) + return np.sum(residual**2) + + +def _discrete_loss_gradient(k): + t = _data_times() + residual = _analytic_solution(k, t) - _analytic_solution(K_TRUE, t) + dydk = -t * _analytic_solution(k, t) + return np.sum(2 * residual * dydk) + + +def _continuous_loss(k): + return (1 - np.exp(-2 * k * FINAL_TIME)) / (2 * k) + + +def _continuous_loss_gradient(k): + e = np.exp(-2 * k * FINAL_TIME) + return (2 * k * FINAL_TIME * e - (1 - e)) / (2 * k**2) + + +def _discrete_loss_function(): + t = _data_times() + data = pybamm.DiscreteTimeData(t, _analytic_solution(K_TRUE, t), "decay data") + return pybamm.DiscreteTimeSum((data - pybamm.Variable("y")) ** 2) + + +def _continuous_loss_function(): + return pybamm.ExplicitTimeIntegral(pybamm.Variable("y") ** 2, pybamm.Scalar(0)) + + +def _make_loss_solver(loss_function, max_workers=None): + sim = pybamm.Simulation( + _decay_model(), solver=pybamm.IDAKLUSolver(rtol=1e-9, atol=1e-9) + ) + return LossSolver(sim, loss_function, FINAL_TIME, max_workers=max_workers) + + +@pytest.mark.skipif(not has_pydiffsol, reason="pydiffsol is not installed") +class TestLossSolver: + @pytest.fixture + def continuous_solver(self): + return _make_loss_solver(_continuous_loss_function()) + + def test_init_raises_without_time_integral(self): + sim = pybamm.Simulation(_decay_model()) + with pytest.raises(ValueError, match=r"DiscreteSum or an ExplicitTimeIntegral"): + LossSolver(sim, pybamm.Variable("y"), FINAL_TIME) + + def test_init_raises_on_duplicate_inner_name(self): + sim = pybamm.Simulation(_decay_model()) + sim.model.variables[LossSolver.INNER_LOSS_FUNCTION_NAME] = pybamm.Scalar(0) + with pytest.raises(ValueError, match=r"already contains a variable named"): + LossSolver(sim, _continuous_loss_function(), FINAL_TIME) + + def test_inputs_to_parameters(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) + assert p.shape == (1, 1) + np.testing.assert_allclose(p, [[K_TRUE]]) + + def test_parameters_to_inputs_round_trip(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) + assert continuous_solver.parameters_to_inputs(p) == [{"k": K_TRUE}] + + def test_predict_matches_analytic(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) + solution = continuous_solver.predict(p)[0] + t = _data_times() + np.testing.assert_allclose( + solution["y"](t), _analytic_solution(K_TRUE, t), rtol=1e-3, atol=1e-4 + ) + + def test_loss_continuous_matches_analytic(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) + np.testing.assert_allclose( + continuous_solver.loss(p), [_continuous_loss(K_TRUE)], rtol=1e-3 + ) + + def test_finite_difference_gradient_matches_analytic(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) + gradient = continuous_solver.finite_difference_gradient(p) + np.testing.assert_allclose( + gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=1e-2 + ) + + def test_loss_and_gradient_continuous_adjoint(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) + loss, gradient = continuous_solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + ) + np.testing.assert_allclose(loss, [_continuous_loss(K_TRUE)], rtol=1e-3) + np.testing.assert_allclose( + gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=1e-2 + ) + + def test_loss_and_gradient_continuous_forward_not_implemented( + self, continuous_solver + ): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) + with pytest.raises(NotImplementedError): + continuous_solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.FORWARD_SENSITIVITY + ) + + def test_inputs_to_parameters_batch(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + assert p.shape == (2, 1) + np.testing.assert_allclose(p, [[K_TRUE], [K_OTHER]]) + + def test_parameters_to_inputs_batch_round_trip(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + assert continuous_solver.parameters_to_inputs(p) == [ + {"k": K_TRUE}, + {"k": K_OTHER}, + ] + + def test_predict_batch_matches_analytic(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + solutions = continuous_solver.predict(p) + assert len(solutions) == 2 + t = _data_times() + for solution, k in zip(solutions, (K_TRUE, K_OTHER), strict=True): + np.testing.assert_allclose( + solution["y"](t), _analytic_solution(k, t), rtol=1e-3, atol=1e-4 + ) + + def test_loss_batch_matches_analytic(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + loss = continuous_solver.loss(p) + assert loss.shape == (2,) + np.testing.assert_allclose( + loss, [_continuous_loss(K_TRUE), _continuous_loss(K_OTHER)], rtol=1e-3 + ) + + def test_finite_difference_gradient_batch(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + gradient = continuous_solver.finite_difference_gradient(p) + assert gradient.shape == (2, 1) + np.testing.assert_allclose( + gradient, + [[_continuous_loss_gradient(K_TRUE)], [_continuous_loss_gradient(K_OTHER)]], + rtol=1e-2, + ) + + def test_loss_and_gradient_batch_adjoint(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + loss, gradient = continuous_solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + ) + assert loss.shape == (2,) + assert gradient.shape == (2, 1) + np.testing.assert_allclose( + loss, [_continuous_loss(K_TRUE), _continuous_loss(K_OTHER)], rtol=1e-3 + ) + np.testing.assert_allclose( + gradient, + [[_continuous_loss_gradient(K_TRUE)], [_continuous_loss_gradient(K_OTHER)]], + rtol=1e-2, + ) + + def test_pickle_round_trip(self, continuous_solver): + p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + expected = continuous_solver.loss(p) + restored = pickle.loads(pickle.dumps(continuous_solver)) + np.testing.assert_allclose(restored.loss(p), expected) + + def test_parallel_matches_sequential(self): + inputs = [{"k": k} for k in (0.4, 0.6, 0.8, 1.0)] + sequential = _make_loss_solver(_continuous_loss_function()) + parallel = _make_loss_solver(_continuous_loss_function(), max_workers=2) + mode = LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + try: + p = sequential.inputs_to_parameters(inputs) + np.testing.assert_allclose(parallel.loss(p), sequential.loss(p)) + seq_loss, seq_grad = sequential.loss_and_gradient(p, mode) + par_loss, par_grad = parallel.loss_and_gradient(p, mode) + np.testing.assert_allclose(par_loss, seq_loss) + np.testing.assert_allclose(par_grad, seq_grad) + np.testing.assert_allclose( + parallel.finite_difference_gradient(p), + sequential.finite_difference_gradient(p), + ) + finally: + parallel.close() + + @discrete_not_supported + def test_loss_discrete_zero_at_true(self): + solver = _make_loss_solver(_discrete_loss_function()) + p = solver.inputs_to_parameters([{"k": K_TRUE}]) + np.testing.assert_allclose(solver.loss(p), [0.0], atol=1e-5) + + @discrete_not_supported + def test_loss_discrete_matches_analytic_off_true(self): + solver = _make_loss_solver(_discrete_loss_function()) + p_true = solver.inputs_to_parameters([{"k": K_TRUE}]) + p_other = solver.inputs_to_parameters([{"k": K_OTHER}]) + loss_other = solver.loss(p_other) + np.testing.assert_allclose( + loss_other, [_discrete_loss(K_OTHER)], rtol=1e-3, atol=1e-5 + ) + assert loss_other[0] > solver.loss(p_true)[0] + + @discrete_not_supported + def test_loss_and_gradient_discrete_forward(self): + solver = _make_loss_solver(_discrete_loss_function()) + p = solver.inputs_to_parameters([{"k": K_OTHER}]) + loss, gradient = solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.FORWARD_SENSITIVITY + ) + np.testing.assert_allclose(loss, solver.loss(p), rtol=1e-3, atol=1e-5) + np.testing.assert_allclose( + gradient, [[_discrete_loss_gradient(K_OTHER)]], rtol=1e-2 + ) + np.testing.assert_allclose( + gradient, solver.finite_difference_gradient(p), rtol=1e-2 + ) + + @discrete_not_supported + def test_loss_and_gradient_discrete_forward_zero_at_true(self): + solver = _make_loss_solver(_discrete_loss_function()) + p = solver.inputs_to_parameters([{"k": K_TRUE}]) + _, gradient = solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.FORWARD_SENSITIVITY + ) + np.testing.assert_allclose(gradient, [[0.0]], atol=1e-4) + + @discrete_not_supported + def test_loss_and_gradient_discrete_adjoint_not_implemented(self): + solver = _make_loss_solver(_discrete_loss_function()) + p = solver.inputs_to_parameters([{"k": K_TRUE}]) + with pytest.raises(NotImplementedError): + solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + ) From d8efb584105a67d7ac7c50ad98e886d0210f5245 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Tue, 30 Jun 2026 12:12:03 +0000 Subject: [PATCH 03/10] feat: broadcasting for loss solver --- .../src/pybamm/simulation/loss_solver.py | 217 +++++++++--------- .../unit/test_solvers/test_loss_solver.py | 16 ++ 2 files changed, 129 insertions(+), 104 deletions(-) diff --git a/packages/pybamm/src/pybamm/simulation/loss_solver.py b/packages/pybamm/src/pybamm/simulation/loss_solver.py index f23bb6e4d8..05fbc221c9 100644 --- a/packages/pybamm/src/pybamm/simulation/loss_solver.py +++ b/packages/pybamm/src/pybamm/simulation/loss_solver.py @@ -98,14 +98,109 @@ def __init__( # set up the process pool for the batched methods, if requested if self._max_workers is not None and self._max_workers > 1: - # spawn (not fork) avoids deadlocks with the native threadpools used - # by the matrix backend; pickling self ships the compiled Ode cheaply. - self._pool = concurrent.futures.ProcessPoolExecutor( - max_workers=self._max_workers, - mp_context=multiprocessing.get_context("spawn"), - initializer=_init_worker, - initargs=(pickle.dumps(self),), + self._pool = self._start_pool() + + def inputs_to_parameters(self, inputs: list[dict]) -> np.ndarray: + """Converts a standard set of pybamm input dictionaries to a 2D parameter array (n_batch, n_params) for use in the functions below.""" + return np.array([self._exporter.map_inputs(single) for single in inputs]) + + def parameters_to_inputs(self, p: np.ndarray) -> list[dict]: + """Converts a 2D parameter array (n_batch, n_params) to a standard set of pybamm input dictionaries.""" + return [self._exporter.inverse_map_inputs(row) for row in np.atleast_2d(p)] + + def predict(self, p: np.ndarray) -> list[pybamm.Solution]: + """Calculate the solution of the ODE for each set of parameters in inputs. + + The parameter sets are solved together in a single batched solve, so any + ``num_threads`` parallelism configured on the wrapped solver is used. + """ + solutions = self._sim.solve( + t_eval=[0, self._final_time], inputs=self.parameters_to_inputs(p) + ) + # a single-element batch is returned as a bare Solution + if not isinstance(solutions, list): + solutions = [solutions] + return solutions + + def loss(self, p: np.ndarray) -> np.ndarray: + """ + Calculate the loss function for each set of parameters in inputs. + Returns a 1D array of loss values of length n_batch. + """ + rows = list(np.atleast_2d(p)) + if self._pool is not None: + values = list(self._pool.map(_worker_loss, rows)) + else: + values = [self._single_loss(row) for row in rows] + return np.array(values) + + def loss_and_gradient( + self, p: np.ndarray, mode: "LossSolverGradientMode" + ) -> tuple[np.ndarray, np.ndarray]: + """ + Calculate the loss and gradient of the loss function with respect to the parameters for each set of parameters in inputs. + Returns a tuple of arrays, where the first contains the loss values as a 1D array of length n_batch and the second contains the gradients as a 2D array of shape (n_batch, n_params) + """ + rows = list(np.atleast_2d(p)) + if self._pool is not None: + items = [(row, mode.value) for row in rows] + results = list(self._pool.map(_worker_loss_and_gradient, items)) + else: + results = [self._single_loss_and_gradient(row, mode) for row in rows] + losses = np.array([loss for loss, _ in results]) + gradients = np.array([gradient for _, gradient in results]) + return losses, gradients + + def finite_difference_gradient(self, p: np.ndarray, h: float = 1e-5) -> np.ndarray: + """ + Calculate the gradient of the loss function with respect to the parameters for each set of parameters in inputs using finite differencing. + Returns a 2D array of gradients, with shape (n_batch, n_params), where each row corresponds to the gradient for a given input parameter set. + """ + p = np.atleast_2d(p).astype(float) + n_batch, n_params = p.shape + gradient = np.zeros((n_batch, n_params)) + for j in range(n_params): + p_plus = p.copy() + p_plus[:, j] += h + p_minus = p.copy() + p_minus[:, j] -= h + gradient[:, j] = (self.loss(p_plus) - self.loss(p_minus)) / (2 * h) + return gradient + + def close(self) -> None: + """Shut down the process pool, if one was created.""" + if self._pool is not None: + self._pool.shutdown(wait=True) + self._pool = None + + def __enter__(self) -> "LossSolver": + return self + + def __exit__(self, *exc) -> None: + self.close() + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + # the live process pool and casadi functions are not picklable; the pool + # is recreated on unpickle (except inside worker processes) and the + # casadi functions are rebuilt + state["_pool"] = None + state["_post_sum_node_present"] = self._post_sum is not None + state["_post_sum"] = None + state["_post_sum_sens"] = None + return state + + def __setstate__(self, state: dict) -> None: + rebuild = state.pop("_post_sum_node_present", False) + self.__dict__.update(state) + if rebuild: + self._post_sum, self._post_sum_sens = self._make_post_sum_functions( + self._processed_loss.post_sum_node ) + # recreate the pool on unpickle, but never inside a worker process, + # which would spawn nested pools + if not _IN_WORKER and self._max_workers and self._max_workers > 1: + self._pool = self._start_pool() def _make_post_sum_functions( self, post_sum_node: pybamm.Symbol | None @@ -145,27 +240,15 @@ def _make_post_sum_functions( ) return post_sum, post_sum_sens - def inputs_to_parameters(self, inputs: list[dict]) -> np.ndarray: - """Converts a standard set of pybamm input dictionaries to a 2D parameter array (n_batch, n_params) for use in the functions below.""" - return np.array([self._exporter.map_inputs(single) for single in inputs]) - - def parameters_to_inputs(self, p: np.ndarray) -> list[dict]: - """Converts a 2D parameter array (n_batch, n_params) to a standard set of pybamm input dictionaries.""" - return [self._exporter.inverse_map_inputs(row) for row in np.atleast_2d(p)] - - def predict(self, p: np.ndarray) -> list[pybamm.Solution]: - """Calculate the solution of the ODE for each set of parameters in inputs. - - The parameter sets are solved together in a single batched solve, so any - ``num_threads`` parallelism configured on the wrapped solver is used. - """ - solutions = self._sim.solve( - t_eval=[0, self._final_time], inputs=self.parameters_to_inputs(p) + def _start_pool(self) -> concurrent.futures.ProcessPoolExecutor: + # spawn (not fork) avoids deadlocks with the native threadpools used by + # the matrix backend; pickling self ships the compiled Ode cheaply. + return concurrent.futures.ProcessPoolExecutor( + max_workers=self._max_workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=_init_worker, + initargs=(pickle.dumps(self),), ) - # a single-element batch is returned as a bare Solution - if not isinstance(solutions, list): - solutions = [solutions] - return solutions def _apply_post_sum( self, the_integral: np.ndarray, params: np.ndarray @@ -237,92 +320,18 @@ def _discrete_sum_to_gradient( ) return loss, gradient - def loss(self, p: np.ndarray) -> np.ndarray: - """ - Calculate the loss function for each set of parameters in inputs. - Returns a 1D array of loss values of length n_batch. - """ - rows = list(np.atleast_2d(p)) - if self._pool is not None: - values = list(self._pool.map(_worker_loss, rows)) - else: - values = [self._single_loss(row) for row in rows] - return np.array(values) - - def loss_and_gradient( - self, p: np.ndarray, mode: "LossSolverGradientMode" - ) -> tuple[np.ndarray, np.ndarray]: - """ - Calculate the loss and gradient of the loss function with respect to the parameters for each set of parameters in inputs. - Returns a tuple of arrays, where the first contains the loss values as a 1D array of length n_batch and the second contains the gradients as a 2D array of shape (n_batch, n_params) - """ - rows = list(np.atleast_2d(p)) - if self._pool is not None: - items = [(row, mode.value) for row in rows] - results = list(self._pool.map(_worker_loss_and_gradient, items)) - else: - results = [self._single_loss_and_gradient(row, mode) for row in rows] - losses = np.array([loss for loss, _ in results]) - gradients = np.array([gradient for _, gradient in results]) - return losses, gradients - - def finite_difference_gradient(self, p: np.ndarray, h: float = 1e-5) -> np.ndarray: - """ - Calculate the gradient of the loss function with respect to the parameters for each set of parameters in inputs using finite differencing. - Returns a 2D array of gradients, with shape (n_batch, n_params), where each row corresponds to the gradient for a given input parameter set. - """ - p = np.atleast_2d(p).astype(float) - n_batch, n_params = p.shape - gradient = np.zeros((n_batch, n_params)) - for j in range(n_params): - p_plus = p.copy() - p_plus[:, j] += h - p_minus = p.copy() - p_minus[:, j] -= h - gradient[:, j] = (self.loss(p_plus) - self.loss(p_minus)) / (2 * h) - return gradient - - def close(self) -> None: - """Shut down the process pool, if one was created.""" - if self._pool is not None: - self._pool.shutdown(wait=True) - self._pool = None - - def __enter__(self) -> "LossSolver": - return self - - def __exit__(self, *exc) -> None: - self.close() - - def __getstate__(self) -> dict: - state = self.__dict__.copy() - # the live process pool and casadi functions are not picklable; the - # unpickled copy runs sequentially (this also keeps the copies sent to - # worker processes pool-less), and the casadi functions are rebuilt - state["_pool"] = None - state["_post_sum_node_present"] = self._post_sum is not None - state["_post_sum"] = None - state["_post_sum_sens"] = None - return state - - def __setstate__(self, state: dict) -> None: - rebuild = state.pop("_post_sum_node_present", False) - self.__dict__.update(state) - if rebuild: - self._post_sum, self._post_sum_sens = self._make_post_sum_functions( - self._processed_loss.post_sum_node - ) - class LossSolverGradientMode(Enum): FORWARD_SENSITIVITY = "forward_sensitivity" ADJOINT_SENSITIVITY = "adjoint_sensitivity" _WORKER_SOLVER = None +_IN_WORKER = False def _init_worker(solver_bytes: bytes) -> None: - global _WORKER_SOLVER + global _WORKER_SOLVER, _IN_WORKER + _IN_WORKER = True _WORKER_SOLVER = pickle.loads(solver_bytes) diff --git a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py index f1b765896c..50ec56f797 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py @@ -210,6 +210,22 @@ def test_pickle_round_trip(self, continuous_solver): restored = pickle.loads(pickle.dumps(continuous_solver)) np.testing.assert_allclose(restored.loss(p), expected) + def test_pickle_round_trip_restores_pool(self): + sequential = _make_loss_solver(_continuous_loss_function()) + p = sequential.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) + expected = sequential.loss(p) + + parallel = _make_loss_solver(_continuous_loss_function(), max_workers=2) + blob = pickle.dumps(parallel) + parallel.close() + + restored = pickle.loads(blob) + try: + assert restored._pool is not None + np.testing.assert_allclose(restored.loss(p), expected) + finally: + restored.close() + def test_parallel_matches_sequential(self): inputs = [{"k": k} for k in (0.4, 0.6, 0.8, 1.0)] sequential = _make_loss_solver(_continuous_loss_function()) From 36d43edf1fc51b2bc89fb59faa4f45cb70147a3e Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Wed, 1 Jul 2026 19:59:18 +0000 Subject: [PATCH 04/10] update to diffsl 0.6.0, implement discrete time sum --- packages/pybamm/pyproject.toml | 194 +++++++++--------- .../expression_tree/operations/diffsl.py | 41 ++++ .../src/pybamm/simulation/loss_solver.py | 46 ++++- .../test_operations/test_diffsl_export.py | 59 ++++++ .../unit/test_solvers/test_loss_solver.py | 71 ++++--- uv.lock | 15 +- 6 files changed, 283 insertions(+), 143 deletions(-) diff --git a/packages/pybamm/pyproject.toml b/packages/pybamm/pyproject.toml index 9ea43afdd5..a929aa4e7a 100644 --- a/packages/pybamm/pyproject.toml +++ b/packages/pybamm/pyproject.toml @@ -12,34 +12,34 @@ maintainers = [{ name = "The PyBaMM Team", email = "pybamm@pybamm.org" }] requires-python = ">=3.10, <3.15" readme = { file = "README.md", content-type = "text/markdown" } classifiers = [ - "Development Status :: 5 - Production/Stable", - "Intended Audience :: Developers", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: BSD License", - "Programming Language :: Python", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - "Programming Language :: Python :: 3.12", - "Programming Language :: Python :: 3.13", - "Programming Language :: Python :: 3.14", - "Topic :: Scientific/Engineering", + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: BSD License", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering", ] dependencies = [ - "pybammsolvers>=0.8.0,<0.9.0", - "black", - "numpy", - "scipy>=1.11.4", - "xarray>=2022.6.0", - "anytree>=2.8.0", - "sympy>=1.12", - "typing-extensions>=4.10.0", - "pandas>=1.5.0", - "pooch>=1.8.1", - "posthog", - "pyyaml", - "platformdirs", + "pybammsolvers>=0.8.0,<0.9.0", + "black", + "numpy", + "scipy>=1.11.4", + "xarray>=2022.6.0", + "anytree>=2.8.0", + "sympy>=1.12", + "typing-extensions>=4.10.0", + "pandas>=1.5.0", + "pooch>=1.8.1", + "posthog", + "pyyaml", + "platformdirs", ] [project.urls] @@ -53,66 +53,68 @@ Changelog = "https://github.com/pybamm-team/PyBaMM/blob/main/CHANGELOG.md" # For example notebooks examples = ["jupyter"] plot = [ - # Note: matplotlib is loaded for debug plots, but to ensure PyBaMM runs - # on systems without an attached display, it should never be imported - # outside of plot() methods. - "matplotlib>=3.6.0", + # Note: matplotlib is loaded for debug plots, but to ensure PyBaMM runs + # on systems without an attached display, it should never be imported + # outside of plot() methods. + "matplotlib>=3.6.0", ] cite = ["pybtex>=0.25.0"] # For diffsl export integration tests -pydiffsol = ["pydiffsol>=0.5.2"] +pydiffsol = ["pydiffsol>=0.6.0"] # Battery Parameter eXchange format bpx = ["bpx>=1.1.0,<1.2.0"] # Low-overhead progress bars tqdm = ["tqdm"] -jax = ["jax>=0.7.0, <0.9.0; python_version >= '3.11' and (sys_platform != 'darwin' or platform_machine != 'x86_64')"] +jax = [ + "jax>=0.7.0, <0.9.0; python_version >= '3.11' and (sys_platform != 'darwin' or platform_machine != 'x86_64')", +] # Contains all optional dependencies, except for jax, and dev dependencies all = [ - "scikit-fem>=8.1.0", - "meshio>=5.3.0", - "pybamm[examples,plot,cite,bpx,tqdm]", + "scikit-fem>=8.1.0", + "meshio>=5.3.0", + "pybamm[examples,plot,cite,bpx,tqdm]", ] [dependency-groups] docs = [ - "sphinx>=6", - "sphinx_rtd_theme>=0.5", - "pydata-sphinx-theme", - "sphinx_design", - "sphinx-copybutton", - "myst-parser", - "sphinx-inline-tabs", - "sphinxcontrib-bibtex", - "sphinx-autobuild", - "sphinx-last-updated-by-git", - "nbsphinx", - "ipykernel", - "ipywidgets", - "sphinx-gallery", - "sphinx-docsearch", + "sphinx>=6", + "sphinx_rtd_theme>=0.5", + "pydata-sphinx-theme", + "sphinx_design", + "sphinx-copybutton", + "myst-parser", + "sphinx-inline-tabs", + "sphinxcontrib-bibtex", + "sphinx-autobuild", + "sphinx-last-updated-by-git", + "nbsphinx", + "ipykernel", + "ipywidgets", + "sphinx-gallery", + "sphinx-docsearch", ] dev = [ - # For working with pre-commit hooks - "pre-commit", - # For code style checks: linting and auto-formatting - "ruff", - # For running testing sessions - "nox", - # For coverage - "pytest-cov", - # For doctest - "pytest-doctestplus", - # pytest and its plugins - "pytest>=9.0", - "pytest-xdist", - "pytest-mock", - "pytest-snapshot", - # For testing Jupyter notebooks - "nbmake", - # To access the metadata for python packages - "importlib-metadata; python_version < '3.10'", - # For property based testing - "hypothesis", + # For working with pre-commit hooks + "pre-commit", + # For code style checks: linting and auto-formatting + "ruff", + # For running testing sessions + "nox", + # For coverage + "pytest-cov", + # For doctest + "pytest-doctestplus", + # pytest and its plugins + "pytest>=9.0", + "pytest-xdist", + "pytest-mock", + "pytest-snapshot", + # For testing Jupyter notebooks + "nbmake", + # To access the metadata for python packages + "importlib-metadata; python_version < '3.10'", + # For property based testing + "hypothesis", ] [project.entry-points."pybamm_parameter_sets"] @@ -159,21 +161,29 @@ build.hooks.vcs.version-file = "src/pybamm/_version.py" # pyproject. Without this, setuptools_scm looks for .git in packages/pybamm/, # fails, and falls back to 0.0.0. root = "../.." -git_describe_command = ["git", "describe", "--dirty", "--tags", "--long", "--match", "pybamm-v*"] +git_describe_command = [ + "git", + "describe", + "--dirty", + "--tags", + "--long", + "--match", + "pybamm-v*", +] tag_regex = '^pybamm-v(?P[0-9].*)$' [tool.repo-review] ignore = [ - "MY", # ignore all MyPy setting checks (MY101, MY103, MY104, MY105, MY106) - "PP004", # ignore the check that prevents an upper cap on Python requires - "PP308", # ignore the explicit pytest summary flag (-ra) check - "PC140", # ignore missing type checker in pre-commit - "PC160", # ignore missing spell checker in pre-commit - "PC180", # ignore missing markdown formatter in pre-commit - "RF102", # ignore the missing isort selection in Ruff configuration - "PP006", # ignore dev dependency group should be defined - "NOX201", # ignore set a script block with dependencies in your noxfile - "NOX202", # ignore noxfile has a shebang line + "MY", # ignore all MyPy setting checks (MY101, MY103, MY104, MY105, MY106) + "PP004", # ignore the check that prevents an upper cap on Python requires + "PP308", # ignore the explicit pytest summary flag (-ra) check + "PC140", # ignore missing type checker in pre-commit + "PC160", # ignore missing spell checker in pre-commit + "PC180", # ignore missing markdown formatter in pre-commit + "RF102", # ignore the missing isort selection in Ruff configuration + "PP006", # ignore dev dependency group should be defined + "NOX201", # ignore set a script block with dependencies in your noxfile + "NOX202", # ignore noxfile has a shebang line ] # NOTE: The shared [tool.ruff] config was hoisted to the repo-root pyproject.toml @@ -188,14 +198,14 @@ testpaths = ["tests"] console_output_style = "progress" xfail_strict = true filterwarnings = [ - "error", - # ignore internal nbmake warnings - 'ignore:unclosed \ dict: state["_post_sum_node_present"] = self._post_sum is not None state["_post_sum"] = None state["_post_sum_sens"] = None + state["_post_sum_dy"] = None return state def __setstate__(self, state: dict) -> None: rebuild = state.pop("_post_sum_node_present", False) self.__dict__.update(state) if rebuild: - self._post_sum, self._post_sum_sens = self._make_post_sum_functions( - self._processed_loss.post_sum_node + self._post_sum, self._post_sum_sens, self._post_sum_dy = ( + self._make_post_sum_functions(self._processed_loss.post_sum_node) ) # recreate the pool on unpickle, but never inside a worker process, # which would spawn nested pools @@ -204,10 +211,10 @@ def __setstate__(self, state: dict) -> None: def _make_post_sum_functions( self, post_sum_node: pybamm.Symbol | None - ) -> tuple[casadi.Function | None, casadi.Function | None]: + ) -> tuple[casadi.Function | None, casadi.Function | None, casadi.Function | None]: """Build casadi functions for the post-sum node and its sensitivities.""" if post_sum_node is None: - return None, None + return None, None, None input_names = self._exporter.input_names() sum_size = self._processed_loss.sum_node.evaluate_for_shape().shape[0] @@ -238,7 +245,12 @@ def _make_post_sum_functions( [t_casadi, sum_casadi, p_casadi_stacked, sens_casadi], [sens], ) - return post_sum, post_sum_sens + post_sum_dy = casadi.Function( + "post_sum_dy", + [t_casadi, sum_casadi, p_casadi_stacked], + [dpost_dy], + ) + return post_sum, post_sum_sens, post_sum_dy def _start_pool(self) -> concurrent.futures.ProcessPoolExecutor: # spawn (not fork) avoids deadlocks with the native threadpools used by @@ -285,9 +297,7 @@ def _single_loss_and_gradient( params, self._processed_loss.discrete_times ) return self._discrete_sum_to_gradient(sol, params) - raise NotImplementedError( - "Adjoint sensitivity for discrete sum is not yet implemented" - ) + return self._discrete_adjoint_gradient(params) if mode == self.LossSolverGradientMode.FORWARD_SENSITIVITY: raise NotImplementedError( @@ -320,6 +330,22 @@ def _discrete_sum_to_gradient( ) return loss, gradient + def _discrete_adjoint_gradient( + self, params: np.ndarray + ) -> tuple[float, np.ndarray]: + times = self._processed_loss.discrete_times + sol, checkpoint = self._ode.solve_adjoint_fwd(params, times) + the_sum = np.sum(sol.ys, axis=1) + loss = float(self._apply_post_sum(the_sum, params).reshape(-1)[0]) + + if self._post_sum is None: + dgdu = np.ones((sol.ys.shape[0], len(sol.ts))) + else: + df_dS = self._post_sum_dy(0.0, the_sum, params).full() # (1, sum_size) + dgdu = np.tile(df_dS.T, (1, len(sol.ts))) # (sum_size, n_discrete) + gradient = self._ode.solve_adjoint_bkwd(sol, checkpoint, dgdu) + return loss, np.asarray(gradient).reshape(-1) + class LossSolverGradientMode(Enum): FORWARD_SENSITIVITY = "forward_sensitivity" ADJOINT_SENSITIVITY = "adjoint_sensitivity" diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py index 04ea03075d..f8eb3523e0 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py @@ -641,3 +641,62 @@ def test_unified_experiment_step_value_tensor_values(self): pos0 = export.index(s0) pos1 = export.index(s1) assert pos0 < pos1 + + def test_interpolant_linear_exports_interp1d(self): + model = pybamm.BaseModel() + x = pybamm.Variable("x") + model.rhs = {x: -x} + model.initial_conditions = {x: pybamm.Scalar(1)} + model.variables = { + "data": pybamm.DiscreteTimeData( + np.array([0.0, 0.5, 1.0]), + np.array([1.0, 2.0, 3.0]), + "test", + ), + } + disc = pybamm.Discretisation() + disc.process_model(model) + export = pybamm.DiffSLExport(model).to_diffeq(outputs=["data"]) + + assert "interp1d(constant" in export + assert "(0:1): 0," in export + assert "(1:2): 0.5," in export + assert "(2:3): 1," in export + + def test_interpolant_pchip_raises(self): + model = pybamm.BaseModel() + x = pybamm.Variable("x") + model.rhs = {x: -x} + model.initial_conditions = {x: pybamm.Scalar(1)} + model.variables = { + "out": pybamm.Interpolant( + np.array([0, 1, 2]), + np.array([1, 2, 3]), + pybamm.t, + interpolator="pchip", + ), + } + disc = pybamm.Discretisation() + disc.process_model(model) + msg = r"DiffSL export only supports 'linear' interpolants" + with pytest.raises(ValueError, match=msg): + pybamm.DiffSLExport(model).to_diffeq(outputs=["out"]) + + def test_interpolant_cubic_raises(self): + model = pybamm.BaseModel() + x = pybamm.Variable("x") + model.rhs = {x: -x} + model.initial_conditions = {x: pybamm.Scalar(1)} + model.variables = { + "out": pybamm.Interpolant( + np.array([0, 1, 2]), + np.array([1, 2, 3]), + pybamm.t, + interpolator="cubic", + ), + } + disc = pybamm.Discretisation() + disc.process_model(model) + msg = r"DiffSL export only supports 'linear' interpolants" + with pytest.raises(ValueError, match=msg): + pybamm.DiffSLExport(model).to_diffeq(outputs=["out"]) diff --git a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py index 50ec56f797..4d6bb91d6c 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py @@ -18,11 +18,6 @@ N_DATA = 11 Y0 = 1.0 -discrete_not_supported = pytest.mark.xfail( - reason="DiffSLExport cannot export the DiscreteTimeData interpolant; " - "the discrete sum LossSolver path is not yet supported", -) - def _decay_model(): """dy/dt = -k * y, y(0) = 1, with analytic solution y(t) = exp(-k * t).""" @@ -113,20 +108,20 @@ def test_predict_matches_analytic(self, continuous_solver): solution = continuous_solver.predict(p)[0] t = _data_times() np.testing.assert_allclose( - solution["y"](t), _analytic_solution(K_TRUE, t), rtol=1e-3, atol=1e-4 + solution["y"](t), _analytic_solution(K_TRUE, t), atol=1e-4 ) def test_loss_continuous_matches_analytic(self, continuous_solver): p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) np.testing.assert_allclose( - continuous_solver.loss(p), [_continuous_loss(K_TRUE)], rtol=1e-3 + continuous_solver.loss(p), [_continuous_loss(K_TRUE)], atol=1e-6 ) def test_finite_difference_gradient_matches_analytic(self, continuous_solver): p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) gradient = continuous_solver.finite_difference_gradient(p) np.testing.assert_allclose( - gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=1e-2 + gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=3e-5, atol=1e-6 ) def test_loss_and_gradient_continuous_adjoint(self, continuous_solver): @@ -134,9 +129,9 @@ def test_loss_and_gradient_continuous_adjoint(self, continuous_solver): loss, gradient = continuous_solver.loss_and_gradient( p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY ) - np.testing.assert_allclose(loss, [_continuous_loss(K_TRUE)], rtol=1e-3) + np.testing.assert_allclose(loss, [_continuous_loss(K_TRUE)], atol=1e-6) np.testing.assert_allclose( - gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=1e-2 + gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=1e-3, atol=1e-5 ) def test_loss_and_gradient_continuous_forward_not_implemented( @@ -167,7 +162,7 @@ def test_predict_batch_matches_analytic(self, continuous_solver): t = _data_times() for solution, k in zip(solutions, (K_TRUE, K_OTHER), strict=True): np.testing.assert_allclose( - solution["y"](t), _analytic_solution(k, t), rtol=1e-3, atol=1e-4 + solution["y"](t), _analytic_solution(k, t), atol=1e-4 ) def test_loss_batch_matches_analytic(self, continuous_solver): @@ -175,7 +170,7 @@ def test_loss_batch_matches_analytic(self, continuous_solver): loss = continuous_solver.loss(p) assert loss.shape == (2,) np.testing.assert_allclose( - loss, [_continuous_loss(K_TRUE), _continuous_loss(K_OTHER)], rtol=1e-3 + loss, [_continuous_loss(K_TRUE), _continuous_loss(K_OTHER)], atol=1e-6 ) def test_finite_difference_gradient_batch(self, continuous_solver): @@ -185,7 +180,8 @@ def test_finite_difference_gradient_batch(self, continuous_solver): np.testing.assert_allclose( gradient, [[_continuous_loss_gradient(K_TRUE)], [_continuous_loss_gradient(K_OTHER)]], - rtol=1e-2, + rtol=3e-5, + atol=1e-6, ) def test_loss_and_gradient_batch_adjoint(self, continuous_solver): @@ -196,12 +192,13 @@ def test_loss_and_gradient_batch_adjoint(self, continuous_solver): assert loss.shape == (2,) assert gradient.shape == (2, 1) np.testing.assert_allclose( - loss, [_continuous_loss(K_TRUE), _continuous_loss(K_OTHER)], rtol=1e-3 + loss, [_continuous_loss(K_TRUE), _continuous_loss(K_OTHER)], atol=1e-6 ) np.testing.assert_allclose( gradient, [[_continuous_loss_gradient(K_TRUE)], [_continuous_loss_gradient(K_OTHER)]], - rtol=1e-2, + rtol=1e-3, + atol=1e-5, ) def test_pickle_round_trip(self, continuous_solver): @@ -245,52 +242,62 @@ def test_parallel_matches_sequential(self): finally: parallel.close() - @discrete_not_supported def test_loss_discrete_zero_at_true(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_TRUE}]) np.testing.assert_allclose(solver.loss(p), [0.0], atol=1e-5) - @discrete_not_supported def test_loss_discrete_matches_analytic_off_true(self): solver = _make_loss_solver(_discrete_loss_function()) p_true = solver.inputs_to_parameters([{"k": K_TRUE}]) p_other = solver.inputs_to_parameters([{"k": K_OTHER}]) loss_other = solver.loss(p_other) - np.testing.assert_allclose( - loss_other, [_discrete_loss(K_OTHER)], rtol=1e-3, atol=1e-5 - ) + np.testing.assert_allclose(loss_other, [_discrete_loss(K_OTHER)], atol=1e-6) assert loss_other[0] > solver.loss(p_true)[0] - @discrete_not_supported def test_loss_and_gradient_discrete_forward(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_OTHER}]) loss, gradient = solver.loss_and_gradient( p, LossSolver.LossSolverGradientMode.FORWARD_SENSITIVITY ) - np.testing.assert_allclose(loss, solver.loss(p), rtol=1e-3, atol=1e-5) + np.testing.assert_allclose(loss, solver.loss(p), atol=1e-6) np.testing.assert_allclose( - gradient, [[_discrete_loss_gradient(K_OTHER)]], rtol=1e-2 + gradient, [[_discrete_loss_gradient(K_OTHER)]], rtol=3e-5, atol=1e-6 ) np.testing.assert_allclose( - gradient, solver.finite_difference_gradient(p), rtol=1e-2 + gradient, solver.finite_difference_gradient(p), rtol=3e-5, atol=1e-6 ) - @discrete_not_supported def test_loss_and_gradient_discrete_forward_zero_at_true(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_TRUE}]) _, gradient = solver.loss_and_gradient( p, LossSolver.LossSolverGradientMode.FORWARD_SENSITIVITY ) - np.testing.assert_allclose(gradient, [[0.0]], atol=1e-4) + np.testing.assert_allclose(gradient, [[0.0]], atol=5e-4) - @discrete_not_supported - def test_loss_and_gradient_discrete_adjoint_not_implemented(self): + def test_loss_and_gradient_discrete_adjoint(self): + solver = _make_loss_solver(_discrete_loss_function()) + p = solver.inputs_to_parameters([{"k": K_OTHER}]) + loss, gradient = solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + ) + np.testing.assert_allclose(loss, solver.loss(p), atol=1e-6) + np.testing.assert_allclose( + gradient, [[_discrete_loss_gradient(K_OTHER)]], rtol=1e-3, atol=1e-5 + ) + np.testing.assert_allclose( + gradient, + solver.finite_difference_gradient(p), + rtol=1e-3, + atol=1e-5, + ) + + def test_loss_and_gradient_discrete_adjoint_zero_at_true(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_TRUE}]) - with pytest.raises(NotImplementedError): - solver.loss_and_gradient( - p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY - ) + _, gradient = solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + ) + np.testing.assert_allclose(gradient, [[0.0]], atol=5e-4) diff --git a/uv.lock b/uv.lock index d595e23a23..0288c3c583 100644 --- a/uv.lock +++ b/uv.lock @@ -3075,7 +3075,7 @@ requires-dist = [ { name = "pybammsolvers", editable = "packages/pybammsolvers" }, { name = "pybtex", marker = "extra == 'all'", specifier = ">=0.25.0" }, { name = "pybtex", marker = "extra == 'cite'", specifier = ">=0.25.0" }, - { name = "pydiffsol", marker = "extra == 'pydiffsol'", specifier = ">=0.5.2" }, + { name = "pydiffsol", marker = "extra == 'pydiffsol'", specifier = ">=0.6.0" }, { name = "pyyaml" }, { name = "scikit-fem", marker = "extra == 'all'", specifier = ">=8.1.0" }, { name = "scipy", specifier = ">=1.11.4" }, @@ -3359,20 +3359,17 @@ wheels = [ [[package]] name = "pydiffsol" -version = "0.5.2" +version = "0.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/90/b2/bd3f30be470bcb5a514e37a02a3cc594fe2ed0d9e69560d17e01156c2b9c/pydiffsol-0.5.2.tar.gz", hash = "sha256:b5affb6b2b6189a74dba8c38d2200d740d61144d17ec8fd08c8e19e89ee377cd", size = 236219, upload-time = "2026-05-20T16:02:41.434Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/2e/f4f2d601b9adf57abd77f31410204fc966eb293969afda4560eb8279d8be/pydiffsol-0.6.0.tar.gz", hash = "sha256:131e75a7f84684a05208b8e0b0995f2bb8961908596d0a7dd3bab91469f5d59e", size = 241146, upload-time = "2026-07-01T16:20:00.691Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/70/76/f6d658363ce5a915349b15fa5caa268e06a94d22f9393087a24b3fdbd39f/pydiffsol-0.5.2-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:e6278ff74bba4ae51e930720120a5ff7e00f947d3d1e010474b92ab08de6f459", size = 36875919, upload-time = "2026-05-20T16:02:35.787Z" }, - { url = "https://files.pythonhosted.org/packages/10/89/1c5f1538d135c0c99808428ecf9ab83a5cbb1285189d289f9ed8f48b4d95/pydiffsol-0.5.2-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:2645c2a2e7717e67d4ea793959d13d369d70b97b670736195bc1b5e81eff9cc1", size = 47205400, upload-time = "2026-05-20T16:02:43.936Z" }, - { url = "https://files.pythonhosted.org/packages/f9/99/5088e1a81f2c1b980698acb28540cb5ebfbf46df07857930a3b9d63d984e/pydiffsol-0.5.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c2c95c821359950bcda22e3caa53deedd4a1641ba9fa1bc6eaa8e1a6faf9e69e", size = 4132015, upload-time = "2026-05-20T16:02:50.312Z" }, - { url = "https://files.pythonhosted.org/packages/d6/eb/3e63722d89d3cc1040ce3b0e24018bc15d506e251283a2ff634b9b310bc4/pydiffsol-0.5.2-cp39-abi3-macosx_14_0_arm64.whl", hash = "sha256:db6d926df360dc2e107f3a524b5bf139d2f186951406fe238c7dfabe251138c7", size = 36883932, upload-time = "2026-05-20T16:02:38.962Z" }, - { url = "https://files.pythonhosted.org/packages/e6/ca/901bc57399b5bd47899479ab709368dfd338da59a41c7f1b20a510793258/pydiffsol-0.5.2-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:66370286565951cdb4a104929a0a4ce05265634a1d95099894542cf436cfc325", size = 47206738, upload-time = "2026-05-20T16:02:47.092Z" }, - { url = "https://files.pythonhosted.org/packages/26/fd/b5a6c2dc755d495c31d62834750e844ab00d462ae20dd3640cac247c4e2b/pydiffsol-0.5.2-cp39-abi3-win_amd64.whl", hash = "sha256:024ac97b2096f4bab76811cbb97cf6912acc07031685f716d3002faeaf2d352a", size = 4138061, upload-time = "2026-05-20T16:02:51.845Z" }, + { url = "https://files.pythonhosted.org/packages/c9/6a/948b1c8c28f9470c1219428479a9d7afd9cdf1fe2de70e7d6c139c48ff1e/pydiffsol-0.6.0-cp310-abi3-macosx_14_0_arm64.whl", hash = "sha256:41fee4f0017cfc77912d0663e0e73fb26b230105dacaa1abb2212fbfda0d314f", size = 36853699, upload-time = "2026-07-01T16:19:57.783Z" }, + { url = "https://files.pythonhosted.org/packages/a4/ed/54632e673365c196b0bedfe5da9535b4a3ed9cbe46900c43842444476f5d/pydiffsol-0.6.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:57d0cceca92e066b196fa6edd54d4159b274020688e05611977153ced7f60bfd", size = 46887687, upload-time = "2026-07-01T16:20:02.033Z" }, + { url = "https://files.pythonhosted.org/packages/e2/13/1d5be91d143ea6b22ee4ef0e99dc29bec39331892c549124a9c53c616f8a/pydiffsol-0.6.0-cp310-abi3-win_amd64.whl", hash = "sha256:772a7b5fe4583cd7dd0b6f5242294f3e599b51e24dca6a97176f5cc2611f26b1", size = 4508733, upload-time = "2026-07-01T16:20:04.547Z" }, ] [[package]] From a4b28505986625b384cab631d3f5c7339a148884 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Wed, 15 Jul 2026 12:40:23 +0000 Subject: [PATCH 05/10] fix: diffsl reserved words and update pydiffsol --- packages/pybamm/pyproject.toml | 2 +- .../expression_tree/operations/diffsl.py | 7 ++++++ .../src/pybamm/simulation/loss_solver.py | 23 ++++++++++++------ .../test_operations/test_diffsl_export.py | 24 +++++++++++++++++++ uv.lock | 13 +++++----- 5 files changed, 55 insertions(+), 14 deletions(-) diff --git a/packages/pybamm/pyproject.toml b/packages/pybamm/pyproject.toml index a929aa4e7a..6c9d78060c 100644 --- a/packages/pybamm/pyproject.toml +++ b/packages/pybamm/pyproject.toml @@ -60,7 +60,7 @@ plot = [ ] cite = ["pybtex>=0.25.0"] # For diffsl export integration tests -pydiffsol = ["pydiffsol>=0.6.0"] +pydiffsol = ["pydiffsol>=0.6.1"] # Battery Parameter eXchange format bpx = ["bpx>=1.1.0,<1.2.0"] # Low-overhead progress bars diff --git a/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py b/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py index c230779150..f73f51db85 100644 --- a/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py +++ b/packages/pybamm/src/pybamm/expression_tree/operations/diffsl.py @@ -1397,6 +1397,11 @@ def equation_to_diffeq( ) +_DIFFSL_RESERVED = frozenset( + {"in", "u", "dudt", "M", "F", "out", "stop", "reset", "constant", "varying"} +) + + def to_variable_name(name: str) -> str: """Convert a name to a valid diffeq variable name""" if name == pybamm.Simulation._STEP_VALUE_INPUT: @@ -1405,6 +1410,8 @@ def to_variable_name(name: str) -> str: name = name.lower() for char in convert_to_underscore: name = name.replace(char, "") + if name in _DIFFSL_RESERVED: + name = f"x{name}" return name diff --git a/packages/pybamm/src/pybamm/simulation/loss_solver.py b/packages/pybamm/src/pybamm/simulation/loss_solver.py index 0e30d425e5..c08fe3d478 100644 --- a/packages/pybamm/src/pybamm/simulation/loss_solver.py +++ b/packages/pybamm/src/pybamm/simulation/loss_solver.py @@ -42,6 +42,11 @@ class LossSolver: across a process pool of this many workers. Defaults to ``None`` (sequential). ``predict`` is unaffected; its parallelism comes from the wrapped solver's ``num_threads`` option. + ode_solver : str, optional + The diffsol ODE solver to use. One of ``"bdf"``, ``"tr_bdf2"``, + ``"esdirk34"``, ``"tsit45"``. Defaults to ``"bdf"``. ``"tr_bdf2"`` + or ``"esdirk34"`` can be more stable than BDF for large DAE systems + with forward sensitivity. """ INNER_LOSS_FUNCTION_NAME = "inner loss function" @@ -52,6 +57,7 @@ def __init__( loss_function: pybamm.Symbol, final_time: float, max_workers: int | None = None, + ode_solver: str = "bdf", ): ds = pybamm.import_optional_dependency("pydiffsol") @@ -85,17 +91,20 @@ def __init__( matrix_type=ds.faer_sparse, scalar_type=ds.f64, linear_solver=ds.lu, - ode_solver=ds.bdf, + ode_solver=getattr(ds, ode_solver, ds.bdf), ) self._ode.integrate_out = self._processed_loss.method == "continuous" self._ode.rtol = self._sim.solver.rtol self._ode.atol = self._sim.solver.atol - self._ode.sens_rtol = self._sim.solver.rtol - self._ode.sens_atol = self._sim.solver.atol - self._ode.out_rtol = self._sim.solver.rtol - self._ode.out_atol = self._sim.solver.atol - self._ode.param_rtol = self._sim.solver.rtol - self._ode.param_atol = self._sim.solver.atol + self._ode.sens_rtol = None + self._ode.sens_atol = None + + # match PyBaMM's Newton solver resilience for large DAE systems + self._ode.options.max_nonlinear_solver_iterations = 100 + self._ode.options.max_error_test_failures = 200 + self._ode.ic_options.max_newton_iterations = 100 + self._ode.ic_options.max_linear_solver_setups = 8 + self._ode.ic_options.max_linesearch_iterations = 20 # generate casadi functions for the post-sum node and its sensitivities self._post_sum, self._post_sum_sens, self._post_sum_dy = ( diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py index f8eb3523e0..963b1b3399 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_operations/test_diffsl_export.py @@ -700,3 +700,27 @@ def test_interpolant_cubic_raises(self): msg = r"DiffSL export only supports 'linear' interpolants" with pytest.raises(ValueError, match=msg): pybamm.DiffSLExport(model).to_diffeq(outputs=["out"]) + + @pytest.mark.parametrize( + "reserved_name", + ["u", "dudt", "in", "out"], + ) + def test_export_handles_reserved_state_names(self, reserved_name): + """DiffSL export renames states with reserved names so the + generated code compiles without name collisions.""" + model = pybamm.BaseModel() + x = pybamm.Variable(reserved_name) + k = pybamm.InputParameter("k") + model.rhs = {x: -k * x} + model.initial_conditions = {x: pybamm.Scalar(1.0)} + model.variables = {reserved_name: x} + + disc = pybamm.Discretisation() + disc.process_model(model) + + export = pybamm.DiffSLExport(model).to_diffeq(outputs=[reserved_name]) + + renamed = f"x{reserved_name}" + assert renamed in export + assert f"{renamed}_i" in export + assert f" {reserved_name} =" not in export diff --git a/uv.lock b/uv.lock index 0288c3c583..adcdf59439 100644 --- a/uv.lock +++ b/uv.lock @@ -3075,7 +3075,7 @@ requires-dist = [ { name = "pybammsolvers", editable = "packages/pybammsolvers" }, { name = "pybtex", marker = "extra == 'all'", specifier = ">=0.25.0" }, { name = "pybtex", marker = "extra == 'cite'", specifier = ">=0.25.0" }, - { name = "pydiffsol", marker = "extra == 'pydiffsol'", specifier = ">=0.6.0" }, + { name = "pydiffsol", marker = "extra == 'pydiffsol'", specifier = ">=0.6.1" }, { name = "pyyaml" }, { name = "scikit-fem", marker = "extra == 'all'", specifier = ">=8.1.0" }, { name = "scipy", specifier = ">=1.11.4" }, @@ -3359,17 +3359,18 @@ wheels = [ [[package]] name = "pydiffsol" -version = "0.6.0" +version = "0.6.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ba/2e/f4f2d601b9adf57abd77f31410204fc966eb293969afda4560eb8279d8be/pydiffsol-0.6.0.tar.gz", hash = "sha256:131e75a7f84684a05208b8e0b0995f2bb8961908596d0a7dd3bab91469f5d59e", size = 241146, upload-time = "2026-07-01T16:20:00.691Z" } +sdist = { url = "https://files.pythonhosted.org/packages/38/28/810dacc6082e7a0c5819ba6763ae42a5e8137019aca89a3eb9fb364daaee/pydiffsol-0.6.1.tar.gz", hash = "sha256:f665d3be7864b76b2b51b14d38ff24d53270a4a5448b34f1aebfc7f52e962c34", size = 240958, upload-time = "2026-07-10T10:12:59.491Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c9/6a/948b1c8c28f9470c1219428479a9d7afd9cdf1fe2de70e7d6c139c48ff1e/pydiffsol-0.6.0-cp310-abi3-macosx_14_0_arm64.whl", hash = "sha256:41fee4f0017cfc77912d0663e0e73fb26b230105dacaa1abb2212fbfda0d314f", size = 36853699, upload-time = "2026-07-01T16:19:57.783Z" }, - { url = "https://files.pythonhosted.org/packages/a4/ed/54632e673365c196b0bedfe5da9535b4a3ed9cbe46900c43842444476f5d/pydiffsol-0.6.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:57d0cceca92e066b196fa6edd54d4159b274020688e05611977153ced7f60bfd", size = 46887687, upload-time = "2026-07-01T16:20:02.033Z" }, - { url = "https://files.pythonhosted.org/packages/e2/13/1d5be91d143ea6b22ee4ef0e99dc29bec39331892c549124a9c53c616f8a/pydiffsol-0.6.0-cp310-abi3-win_amd64.whl", hash = "sha256:772a7b5fe4583cd7dd0b6f5242294f3e599b51e24dca6a97176f5cc2611f26b1", size = 4508733, upload-time = "2026-07-01T16:20:04.547Z" }, + { url = "https://files.pythonhosted.org/packages/d2/cf/b5c7154c4f8cb314abb02181ca1a2648de6a0862a3ecee3fcb29cf9f4d5a/pydiffsol-0.6.1-cp310-abi3-macosx_14_0_arm64.whl", hash = "sha256:fa49c985d5d7f41f433b835ad0b4c3b40904882c633ddba3dd3dbc526ed3f786", size = 44040271, upload-time = "2026-07-10T10:12:57.207Z" }, + { url = "https://files.pythonhosted.org/packages/f2/f8/056bf2f4b8823e7d145a3ff298ab982a0a395bd9fe7f259bbb22e10302fe/pydiffsol-0.6.1-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:52cad8a9db84f5a226bad5312e42c15804191c793d2f448c7211eb4be5e20d48", size = 50969068, upload-time = "2026-07-10T10:13:01.162Z" }, + { url = "https://files.pythonhosted.org/packages/7e/38/ac3480f5802830a3bb8c4b49aa0d3df5a8f3176e8544e08bd394f2a984ee/pydiffsol-0.6.1-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:9b7568cdda902c904dada2d3f0a7e5e08979472e425f8b55641a39d3cc37e31a", size = 54362150, upload-time = "2026-07-10T10:13:04.603Z" }, + { url = "https://files.pythonhosted.org/packages/ef/c2/53c44583e97ab14b1ec417a22977d7f983558825fd8d7cdeab8129ee2e87/pydiffsol-0.6.1-cp310-abi3-win_amd64.whl", hash = "sha256:20eb8ba8ee45dee27a8d2277da0e307d1dd07c9374543577194863b6bc4cf35f", size = 4506091, upload-time = "2026-07-10T10:13:07.376Z" }, ] [[package]] From b95c87bf9171aed51621af0d0dee64e7a9b2263d Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Wed, 15 Jul 2026 12:56:14 +0000 Subject: [PATCH 06/10] add CHANGELOG.md --- CHANGELOG.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c21a0bb397..edb5311a36 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ - PyBaMM and `pybammsolvers` now develop in a single repository — a UV workspace under `packages/` — while continuing to release independently to PyPI. Release tags are namespaced (`pybamm-v*` and `pybammsolvers-v*`), and PyBaMM's CI now tests against the in-repo solver on every platform. The published `pybamm` package and its dependency on `pybammsolvers` are unchanged for users. See `RELEASE.md` for the release model. ([#5512](https://github.com/pybamm-team/PyBaMM/issues/5512)) - Legacy BPX v0.x files/objects now load again: `bpx` itself detects and converts them to the v1.x schema on a best-effort basis (with a `UserWarning`), so `ParameterValues.create_from_bpx`/`create_from_bpx_obj` no longer raise a `ValidationError`. PyBaMM officially supports `bpx>=1`. ([#5574](https://github.com/pybamm-team/PyBaMM/pull/5574)) - `create_from_bpx`/`create_from_bpx_obj` now also accept BPX files that omit `State` fields (or the whole `State` section): the ambient/initial temperatures default to the reference temperature and the initial electrolyte concentration to 1000 mol.m-3 (logged), while opt-in fields (initial hysteresis state, heat transfer coefficient) are left for the model to default. ([#5574](https://github.com/pybamm-team/PyBaMM/pull/5574)) +- add `pybamm.LossSolver` that constructs a loss function from a `pybamm.Simulation` [(#5625)](https://github.com/pybamm-team/PyBaMM/pull/5652) ## Bug fixes @@ -240,6 +241,7 @@ as initial conditions. ([#5311](https://github.com/pybamm-team/PyBaMM/pull/5311) - Fixed a bug in 2D concatenatations for quantities that vary in the `tb` direction ([#5310](https://github.com/pybamm-team/PyBaMM/pull/5310)) # Breaking changes + - Removes default constants added to `ParameterValues` on construction. **Only breaking if you rely on this functionality in custom models, parameters, etc.** ([#5336](https://github.com/pybamm-team/PyBaMM/pull/5336)) # [v25.10.2](https://github.com/pybamm-team/PyBaMM/tree/v25.10.2) - 2025-11-27 @@ -291,12 +293,12 @@ as initial conditions. ([#5311](https://github.com/pybamm-team/PyBaMM/pull/5311) - Fix Bruggeman coefficient computation from BPX porosity and transport efficiency instead of hard-coding, remove redundant values, and add a unit test for verification. ([#5196](https://github.com/pybamm-team/PyBaMM/pull/5196)) ## Breaking changes + - Updates the hysteresis decay rate parameters to a "true" hysteresis decay rate which changes the interpretation of the units of the hysteresis decay rate parameters. ([#5217](https://github.com/pybamm-team/PyBaMM/pull/5217)) - Changed fundamental variable for all SEI models from thickness to concentration ([#4869](https://github.com/pybamm-team/PyBaMM/pull/4869)) # [v25.8.0](https://github.com/pybamm-team/PyBaMM/tree/v25.8.0) - 2025-08-04 - ## Features - Added `plot_3d_cross_section` & `plot_3d_heatmap` functions to support plotting for 3D thermal simulations. ([#5130](https://github.com/pybamm-team/PyBaMM/pull/5130)) @@ -321,7 +323,7 @@ as initial conditions. ([#5311](https://github.com/pybamm-team/PyBaMM/pull/5311) - Fixed non-deterministic plotting CI issues ([#5150](https://github.com/pybamm-team/PyBaMM/pull/5150)) - Fix non-deterministic ShapeError in 3D FEM gradient method ([#5143](https://github.com/pybamm-team/PyBaMM/pull/5143)) - Fixes negative electrode boundary values for half-cell voltage contributions. ([#5139](https://github.com/pybamm-team/PyBaMM/pull/5139)) -- Makes `A_cc` L_z * L_y * number of layers ([#5138](https://github.com/pybamm-team/PyBaMM/pull/5138)) +- Makes `A_cc` L_z *L_y* number of layers ([#5138](https://github.com/pybamm-team/PyBaMM/pull/5138)) - Fixes `TimeIntegral` expression node summation when dependent on an input parameter. ([#5119](https://github.com/pybamm-team/PyBaMM/pull/5119)) - Fixed a bug that ignored the default duration of drive cycles for `CRate` steps and a bug that overwrote custom `period` arguments for drive cycles. ([#5090](https://github.com/pybamm-team/PyBaMM/pull/5090)) - Converts sensitivities to numpy objects, fixing bug in `DiscreteTimeSum` sensitivity calculation ([#5037](https://github.com/pybamm-team/PyBaMM/pull/5037)) @@ -330,7 +332,6 @@ as initial conditions. ([#5311](https://github.com/pybamm-team/PyBaMM/pull/5311) - Fixed a bug where simplifications cause heavisides to evaluate as booleans ([#4893](https://github.com/pybamm-team/PyBaMM/pull/4893)) - Fixed a bug in the `WyciskOpenCircuitPotential` model where the differential capacity was not being evaluated correctly. ([#4893](https://github.com/pybamm-team/PyBaMM/pull/4893)) - ## Breaking changes - Changed behavior of drive cycle steps in `pybamm.Experiment`s to treat each time point as a discontinuity, consistent with how input interpolants work. This ensures more accurate simulation of drive cycles with rapid changes. ([#5141](https://github.com/pybamm-team/PyBaMM/pull/5141)) @@ -339,7 +340,6 @@ as initial conditions. ([#5311](https://github.com/pybamm-team/PyBaMM/pull/5311) - Removed support for Python 3.9 ([#5052](https://github.com/pybamm-team/PyBaMM/pull/5052)) - In OCP hysteresis models, users need to explicitly give the equilibrium, delithiation, and lithiation OCPs when using a hysteresis model. E.g., you must provide all three of "Negative electrode OCP [V]", "Negative electrode delithiation OCP [V]", and "Negative electrode lithiation OCP [V]". ([#4893](https://github.com/pybamm-team/PyBaMM/pull/4893)) - # [v25.6.0](https://github.com/pybamm-team/PyBaMM/tree/v25.6.0) - 2025-05-27 ## Features @@ -497,6 +497,7 @@ package to install PyBaMM with only the required dependencies. ([conda-forge/pyb - Removed the `start_step_offset` setting and disabled minimum `dt` warnings for drive cycles with the (`IDAKLUSolver`). ([#4416](https://github.com/pybamm-team/PyBaMM/pull/4416)) ## Bug Fixes + - Added error for binary operators on two concatenations with different numbers of children. Previously, the extra children were dropped. Also fixed bug where Q_rxn was dropped from the total heating term in half-cell models. ([#4562](https://github.com/pybamm-team/PyBaMM/pull/4562)) - Fixed bug where Q_rxn was set to 0 for the negative electrode in half-cell models. ([#4557](https://github.com/pybamm-team/PyBaMM/pull/4557)) - Fixed bug in post-processing solutions with infeasible experiments using the (`IDAKLUSolver`). ([#4541](https://github.com/pybamm-team/PyBaMM/pull/4541)) @@ -1264,7 +1265,7 @@ This release introduces: - Added `NewmanTobias` li-ion battery model ([#1423](https://github.com/pybamm-team/PyBaMM/pull/1423)) - Added `plot_voltage_components` to easily plot the component overpotentials that make up the voltage ([#1419](https://github.com/pybamm-team/PyBaMM/pull/1419)) - Made `QuickPlot` more customizable and added an example ([#1419](https://github.com/pybamm-team/PyBaMM/pull/1419)) -- `Solution` objects can now be created by stepping _different_ models ([#1408](https://github.com/pybamm-team/PyBaMM/pull/1408)) +- `Solution` objects can now be created by stepping *different* models ([#1408](https://github.com/pybamm-team/PyBaMM/pull/1408)) - Added Yang et al 2017 model that couples irreversible lithium plating, SEI growth and change in porosity which produces a transition from linear to nonlinear degradation pattern of lithium-ion battery over extended cycles([#1398](https://github.com/pybamm-team/PyBaMM/pull/1398)) - Added support for Python 3.9 and dropped support for Python 3.6. Python 3.6 may still work but is now untested ([#1370](https://github.com/pybamm-team/PyBaMM/pull/1370)) - Added the electrolyte overpotential and Ohmic losses for full conductivity, including surface form ([#1350](https://github.com/pybamm-team/PyBaMM/pull/1350)) From c5000d5185e6e09036414838fc2a1c8f30da1bca Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Wed, 15 Jul 2026 13:06:54 +0000 Subject: [PATCH 07/10] remove loss_function.py --- .../src/pybamm/simulation/loss_function.py | 41 ------------------- 1 file changed, 41 deletions(-) delete mode 100644 packages/pybamm/src/pybamm/simulation/loss_function.py diff --git a/packages/pybamm/src/pybamm/simulation/loss_function.py b/packages/pybamm/src/pybamm/simulation/loss_function.py deleted file mode 100644 index cb74f115fb..0000000000 --- a/packages/pybamm/src/pybamm/simulation/loss_function.py +++ /dev/null @@ -1,41 +0,0 @@ -import numpy as np -import pandas as pd - -import pybamm - -bdp_to_pybamm_mapping = { - "Voltage / V": pybamm.Variable("Voltage [V]"), - "Current / A": pybamm.Variable("Current [A]"), -} - - -def _data_comparison(data: pd.DataFame): - data_times = data["Test Time / s"] - data_values_list = [] - variables_list = [] - for column in data.columns: - variable = bdp_to_pybamm_mapping.get(column, None) - if variable is not None: - variables_list.append(variable) - data_values_list.append(data[column]) - if not variables_list: - raise ValueError( - "No variables found in the data. Please ensure that the data " - "contains columns with any of the following names: " - f"{list(bdp_to_pybamm_mapping.keys())}" - ) - - data_values = np.hstack(data_values_list) - variables = pybamm.NumpyConcatenation(*variables_list) - - data_values = data["value"] - data = pybamm.DiscreteTimeData(data_times, data_values, "sum of squares data") - return data, variables - - -def sum_of_squares(data: pd.DataFrame): - """ - A method to create a loss function with a sum-of-squared-error loss function, given some data in BDF format (https://battery-data-alliance.github.io/battery-data-format/) to fit against. - """ - data, variables = _data_comparison(data) - return pybamm.DiscreteTimeSum((data - variables) ** 2) From e2f232419eefe5ff1db6eb1a9cbbe14c643662c8 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Thu, 16 Jul 2026 12:10:30 +0000 Subject: [PATCH 08/10] turn on pydiffsol in unit tests --- noxfile.py | 2 +- packages/pybamm/pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/noxfile.py b/noxfile.py index 1665d2aef7..e6cba4f2a2 100644 --- a/noxfile.py +++ b/noxfile.py @@ -147,7 +147,7 @@ def run_doctests(session): def run_unit(session): """Run the unit tests.""" set_environment_variables(PYBAMM_ENV, session=session) - install_locked(session, extras=["all", "jax"], groups=["dev"]) + install_locked(session, extras=["all", "jax", "pydiffsol"], groups=["dev"]) session.run("python", "-m", "pytest", "-m", "unit", "packages/pybamm/tests") diff --git a/packages/pybamm/pyproject.toml b/packages/pybamm/pyproject.toml index d93042b5bf..34e6de014c 100644 --- a/packages/pybamm/pyproject.toml +++ b/packages/pybamm/pyproject.toml @@ -59,7 +59,7 @@ plot = [ "matplotlib>=3.6.0", ] cite = ["pybtex>=0.25.0"] -# For diffsl export integration tests +# For diffsl tests pydiffsol = ["pydiffsol>=0.6.1"] # Battery Parameter eXchange format bpx = ["bpx>=1.1.1,<1.2.0"] From 06ef175d5b7f48be404b7bd6f4a74e3faf939d10 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Thu, 16 Jul 2026 12:41:27 +0000 Subject: [PATCH 09/10] fix: turn of sensitivity and pickle tests for windows --- .../unit/test_solvers/test_loss_solver.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py index 4d6bb91d6c..931a6ded8b 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py @@ -3,6 +3,7 @@ # import importlib.util import pickle +import sys import numpy as np import pytest @@ -11,6 +12,7 @@ from pybamm.simulation.loss_solver import LossSolver has_pydiffsol = importlib.util.find_spec("pydiffsol") is not None +is_windows = sys.platform == "win32" K_TRUE = 0.5 K_OTHER = 0.8 @@ -124,6 +126,9 @@ def test_finite_difference_gradient_matches_analytic(self, continuous_solver): gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=3e-5, atol=1e-6 ) + @pytest.mark.skipif( + is_windows, reason="adjoint sensitivity not available on Windows" + ) def test_loss_and_gradient_continuous_adjoint(self, continuous_solver): p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) loss, gradient = continuous_solver.loss_and_gradient( @@ -134,6 +139,9 @@ def test_loss_and_gradient_continuous_adjoint(self, continuous_solver): gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=1e-3, atol=1e-5 ) + @pytest.mark.skipif( + is_windows, reason="forward sensitivity not available on Windows" + ) def test_loss_and_gradient_continuous_forward_not_implemented( self, continuous_solver ): @@ -184,6 +192,9 @@ def test_finite_difference_gradient_batch(self, continuous_solver): atol=1e-6, ) + @pytest.mark.skipif( + is_windows, reason="adjoint sensitivity not available on Windows" + ) def test_loss_and_gradient_batch_adjoint(self, continuous_solver): p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) loss, gradient = continuous_solver.loss_and_gradient( @@ -201,12 +212,14 @@ def test_loss_and_gradient_batch_adjoint(self, continuous_solver): atol=1e-5, ) + @pytest.mark.skipif(is_windows, reason="pickling not supported on Windows") def test_pickle_round_trip(self, continuous_solver): p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) expected = continuous_solver.loss(p) restored = pickle.loads(pickle.dumps(continuous_solver)) np.testing.assert_allclose(restored.loss(p), expected) + @pytest.mark.skipif(is_windows, reason="pickling not supported on Windows") def test_pickle_round_trip_restores_pool(self): sequential = _make_loss_solver(_continuous_loss_function()) p = sequential.inputs_to_parameters([{"k": K_TRUE}, {"k": K_OTHER}]) @@ -223,6 +236,9 @@ def test_pickle_round_trip_restores_pool(self): finally: restored.close() + @pytest.mark.skipif( + is_windows, reason="adjoint sensitivity not available on Windows" + ) def test_parallel_matches_sequential(self): inputs = [{"k": k} for k in (0.4, 0.6, 0.8, 1.0)] sequential = _make_loss_solver(_continuous_loss_function()) @@ -255,6 +271,9 @@ def test_loss_discrete_matches_analytic_off_true(self): np.testing.assert_allclose(loss_other, [_discrete_loss(K_OTHER)], atol=1e-6) assert loss_other[0] > solver.loss(p_true)[0] + @pytest.mark.skipif( + is_windows, reason="forward sensitivity not available on Windows" + ) def test_loss_and_gradient_discrete_forward(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_OTHER}]) @@ -269,6 +288,9 @@ def test_loss_and_gradient_discrete_forward(self): gradient, solver.finite_difference_gradient(p), rtol=3e-5, atol=1e-6 ) + @pytest.mark.skipif( + is_windows, reason="forward sensitivity not available on Windows" + ) def test_loss_and_gradient_discrete_forward_zero_at_true(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_TRUE}]) @@ -277,6 +299,9 @@ def test_loss_and_gradient_discrete_forward_zero_at_true(self): ) np.testing.assert_allclose(gradient, [[0.0]], atol=5e-4) + @pytest.mark.skipif( + is_windows, reason="adjoint sensitivity not available on Windows" + ) def test_loss_and_gradient_discrete_adjoint(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_OTHER}]) @@ -294,6 +319,9 @@ def test_loss_and_gradient_discrete_adjoint(self): atol=1e-5, ) + @pytest.mark.skipif( + is_windows, reason="adjoint sensitivity not available on Windows" + ) def test_loss_and_gradient_discrete_adjoint_zero_at_true(self): solver = _make_loss_solver(_discrete_loss_function()) p = solver.inputs_to_parameters([{"k": K_TRUE}]) From d413162ab43711e0783720b0a25d514bc2af2da8 Mon Sep 17 00:00:00 2001 From: martinjrobins Date: Mon, 20 Jul 2026 14:52:04 +0000 Subject: [PATCH 10/10] test: add coverage for post sum loss solver --- noxfile.py | 2 +- .../src/pybamm/simulation/loss_solver.py | 9 ++- .../unit/test_solvers/test_loss_solver.py | 68 +++++++++++++++++++ 3 files changed, 75 insertions(+), 4 deletions(-) diff --git a/noxfile.py b/noxfile.py index e6cba4f2a2..e95e437710 100644 --- a/noxfile.py +++ b/noxfile.py @@ -105,7 +105,7 @@ def install_locked(session, *, extras=None, groups=None): def run_coverage(session): """Run the coverage tests and generate an XML report.""" set_environment_variables(PYBAMM_ENV, session=session) - install_locked(session, extras=["all", "jax"], groups=["dev"]) + install_locked(session, extras=["all", "jax", "pydiffsol"], groups=["dev"]) # Using plugin here since coverage runs unit tests on linux with latest python version. if "CI" in os.environ: session.install("pytest-github-actions-annotate-failures") diff --git a/packages/pybamm/src/pybamm/simulation/loss_solver.py b/packages/pybamm/src/pybamm/simulation/loss_solver.py index c08fe3d478..48a663b36d 100644 --- a/packages/pybamm/src/pybamm/simulation/loss_solver.py +++ b/packages/pybamm/src/pybamm/simulation/loss_solver.py @@ -364,17 +364,20 @@ class LossSolverGradientMode(Enum): _IN_WORKER = False -def _init_worker(solver_bytes: bytes) -> None: +# Spawned-worker coverage is not collected; exercised by test_parallel_matches_sequential. +def _init_worker(solver_bytes: bytes) -> None: # pragma: no cover global _WORKER_SOLVER, _IN_WORKER _IN_WORKER = True _WORKER_SOLVER = pickle.loads(solver_bytes) -def _worker_loss(params: np.ndarray) -> float: +def _worker_loss(params: np.ndarray) -> float: # pragma: no cover return _WORKER_SOLVER._single_loss(params) -def _worker_loss_and_gradient(item: tuple) -> tuple[float, np.ndarray]: +def _worker_loss_and_gradient( + item: tuple, +) -> tuple[float, np.ndarray]: # pragma: no cover params, mode_value = item mode = type(_WORKER_SOLVER).LossSolverGradientMode(mode_value) return _WORKER_SOLVER._single_loss_and_gradient(params, mode) diff --git a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py index 931a6ded8b..2a73d3a616 100644 --- a/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py +++ b/packages/pybamm/tests/unit/test_solvers/test_loss_solver.py @@ -72,6 +72,14 @@ def _continuous_loss_function(): return pybamm.ExplicitTimeIntegral(pybamm.Variable("y") ** 2, pybamm.Scalar(0)) +def _continuous_post_sum_loss_function(): + return _continuous_loss_function() ** 0.5 + + +def _discrete_post_sum_loss_function(): + return _discrete_loss_function() ** 0.5 + + def _make_loss_solver(loss_function, max_workers=None): sim = pybamm.Simulation( _decay_model(), solver=pybamm.IDAKLUSolver(rtol=1e-9, atol=1e-9) @@ -119,6 +127,11 @@ def test_loss_continuous_matches_analytic(self, continuous_solver): continuous_solver.loss(p), [_continuous_loss(K_TRUE)], atol=1e-6 ) + def test_context_manager(self): + with _make_loss_solver(_continuous_loss_function()) as solver: + p = solver.inputs_to_parameters([{"k": K_TRUE}]) + np.testing.assert_allclose(solver.loss(p), [_continuous_loss(K_TRUE)]) + def test_finite_difference_gradient_matches_analytic(self, continuous_solver): p = continuous_solver.inputs_to_parameters([{"k": K_TRUE}]) gradient = continuous_solver.finite_difference_gradient(p) @@ -139,6 +152,22 @@ def test_loss_and_gradient_continuous_adjoint(self, continuous_solver): gradient, [[_continuous_loss_gradient(K_TRUE)]], rtol=1e-3, atol=1e-5 ) + @pytest.mark.skipif( + is_windows, reason="adjoint sensitivity not available on Windows" + ) + def test_loss_and_gradient_continuous_post_sum_adjoint(self): + solver = _make_loss_solver(_continuous_post_sum_loss_function()) + p = solver.inputs_to_parameters([{"k": K_OTHER}]) + loss, gradient = solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + ) + expected_loss = np.sqrt(_continuous_loss(K_OTHER)) + expected_gradient = _continuous_loss_gradient(K_OTHER) / (2 * expected_loss) + np.testing.assert_allclose(loss, [expected_loss], atol=1e-6) + np.testing.assert_allclose( + gradient, [[expected_gradient]], rtol=1e-3, atol=1e-5 + ) + @pytest.mark.skipif( is_windows, reason="forward sensitivity not available on Windows" ) @@ -219,6 +248,13 @@ def test_pickle_round_trip(self, continuous_solver): restored = pickle.loads(pickle.dumps(continuous_solver)) np.testing.assert_allclose(restored.loss(p), expected) + @pytest.mark.skipif(is_windows, reason="pickling not supported on Windows") + def test_pickle_round_trip_post_sum(self): + solver = _make_loss_solver(_continuous_post_sum_loss_function()) + p = solver.inputs_to_parameters([{"k": K_OTHER}]) + restored = pickle.loads(pickle.dumps(solver)) + np.testing.assert_allclose(restored.loss(p), solver.loss(p)) + @pytest.mark.skipif(is_windows, reason="pickling not supported on Windows") def test_pickle_round_trip_restores_pool(self): sequential = _make_loss_solver(_continuous_loss_function()) @@ -288,6 +324,22 @@ def test_loss_and_gradient_discrete_forward(self): gradient, solver.finite_difference_gradient(p), rtol=3e-5, atol=1e-6 ) + @pytest.mark.skipif( + is_windows, reason="forward sensitivity not available on Windows" + ) + def test_loss_and_gradient_discrete_post_sum_forward(self): + solver = _make_loss_solver(_discrete_post_sum_loss_function()) + p = solver.inputs_to_parameters([{"k": K_OTHER}]) + loss, gradient = solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.FORWARD_SENSITIVITY + ) + expected_loss = np.sqrt(_discrete_loss(K_OTHER)) + expected_gradient = _discrete_loss_gradient(K_OTHER) / (2 * expected_loss) + np.testing.assert_allclose(loss, [expected_loss], atol=1e-6) + np.testing.assert_allclose( + gradient, [[expected_gradient]], rtol=3e-5, atol=1e-6 + ) + @pytest.mark.skipif( is_windows, reason="forward sensitivity not available on Windows" ) @@ -319,6 +371,22 @@ def test_loss_and_gradient_discrete_adjoint(self): atol=1e-5, ) + @pytest.mark.skipif( + is_windows, reason="adjoint sensitivity not available on Windows" + ) + def test_loss_and_gradient_discrete_post_sum_adjoint(self): + solver = _make_loss_solver(_discrete_post_sum_loss_function()) + p = solver.inputs_to_parameters([{"k": K_OTHER}]) + loss, gradient = solver.loss_and_gradient( + p, LossSolver.LossSolverGradientMode.ADJOINT_SENSITIVITY + ) + expected_loss = np.sqrt(_discrete_loss(K_OTHER)) + expected_gradient = _discrete_loss_gradient(K_OTHER) / (2 * expected_loss) + np.testing.assert_allclose(loss, [expected_loss], atol=1e-6) + np.testing.assert_allclose( + gradient, [[expected_gradient]], rtol=1e-3, atol=1e-5 + ) + @pytest.mark.skipif( is_windows, reason="adjoint sensitivity not available on Windows" )