diff --git a/CHANGELOG.md b/CHANGELOG.md index c3dd6846a8..aa5d5e381c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Features +- `pybamm.base_solver.solve` function can take a list of input parameters to calculate the sensitivities of the solution with respect to. Alternatively, it can be set to `True` to calculate the sensitivities for all input parameters ([#1552](https://github.com/pybamm-team/PyBaMM/pull/1552)) - Addes UDDS and WLTC drive cycles ([#1601](https://github.com/pybamm-team/PyBaMM/pull/1601)) - Added capability for `quaternary` domains (in addition to `primary`, `secondary` and `tertiary`), increasing the maximum number of domains that a `Symbol` can have to 4. ([#1580](https://github.com/pybamm-team/PyBaMM/pull/1580)) - Tabs can now be placed at the bottom of the cell in 1+1D thermal models ([#1581](https://github.com/pybamm-team/PyBaMM/pull/1581)) @@ -42,6 +43,7 @@ ## Breaking changes +- Changed sensitivity API. Removed `ProcessedSymbolicVariable`, all sensitivity now handled within the solvers and `ProcessedVariable` () ([#1552](https://github.com/pybamm-team/PyBaMM/pull/1552)) - The `Yang2017` parameter set has been removed as the complete parameter set is not publicly available in the literature ([#1577](https://github.com/pybamm-team/PyBaMM/pull/1577)) - Changed how options are specified for the "loss of active material" and "particle cracking" submodels. "loss of active material" can now be one of "none", "stress-driven", or "reaction-driven", or a 2-tuple for different options in negative and positive electrode. Similarly "particle cracking" (now called "particle mechanics") can now be "none", "swelling only", "swelling and cracking", or a 2-tuple ([#1490](https://github.com/pybamm-team/PyBaMM/pull/1490)) - Changed the variable in the full diffusion model from "Electrolyte concentration" to "Porosity times concentration" ([#1476](https://github.com/pybamm-team/PyBaMM/pull/1476)) @@ -189,7 +191,6 @@ This release adds new operators for more complex models, some basic sensitivity (e.g. `standard_parameters_lithium_ion` is now `LithiumIonParameters`) ([#1120](https://github.com/pybamm-team/PyBaMM/pull/1120)) - Renamed `quick_plot_vars` to `output_variables` in `Simulation` to be consistent with `QuickPlot`. Passing `quick_plot_vars` to `Simulation.plot()` has been deprecated and `output_variables` should be passed instead ([#1099](https://github.com/pybamm-team/PyBaMM/pull/1099)) - # [v0.2.3](https://github.com/pybamm-team/PyBaMM/tree/v0.2.3) - 2020-07-01 This release enables the use of [Google Colab](https://colab.research.google.com/github/pybamm-team/PyBaMM/blob/main/) for running example notebooks, and adds some small new features and bug fixes. diff --git a/FindSUNDIALS.cmake b/FindSUNDIALS.cmake index 7b866056b8..7b3051d903 100644 --- a/FindSUNDIALS.cmake +++ b/FindSUNDIALS.cmake @@ -27,7 +27,7 @@ # find the SUNDIALS include directories find_path(SUNDIALS_INCLUDE_DIR NAMES - ida/ida.h + idas/idas.h sundials/sundials_math.h sundials/sundials_types.h sunlinsol/sunlinsol_klu.h @@ -39,7 +39,7 @@ find_path(SUNDIALS_INCLUDE_DIR ) set(SUNDIALS_WANT_COMPONENTS - sundials_ida + sundials_idas sundials_sunlinsolklu sundials_sunmatrixsparse sundials_nvecserial diff --git a/docs/source/solvers/processed_variable.rst b/docs/source/solvers/processed_variable.rst index 4b4296061b..9f1d62f2a7 100644 --- a/docs/source/solvers/processed_variable.rst +++ b/docs/source/solvers/processed_variable.rst @@ -3,6 +3,3 @@ Post-Process Variables .. autoclass:: pybamm.ProcessedVariable :members: - -.. autoclass:: pybamm.ProcessedSymbolicVariable - :members: diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py index 34e1c654b3..50617a0560 100644 --- a/pybamm/discretisations/discretisation.py +++ b/pybamm/discretisations/discretisation.py @@ -209,6 +209,11 @@ def process_model(self, model, inplace=True, check_model=True): model_disc.rhs, model_disc.concatenated_rhs = rhs, concat_rhs model_disc.algebraic, model_disc.concatenated_algebraic = alg, concat_alg + # Save length of rhs and algebraic + model_disc.len_rhs = model_disc.concatenated_rhs.size + model_disc.len_alg = model_disc.concatenated_algebraic.size + model_disc.len_rhs_and_alg = model_disc.len_rhs + model_disc.len_alg + # Process events processed_events = [] pybamm.logger.verbose("Discretise events for {}".format(model.name)) diff --git a/pybamm/expression_tree/concatenations.py b/pybamm/expression_tree/concatenations.py index 670dbcef94..e0ba4083e4 100644 --- a/pybamm/expression_tree/concatenations.py +++ b/pybamm/expression_tree/concatenations.py @@ -45,6 +45,18 @@ def __str__(self): out = out[:-2] + ")" return out + def _diff(self, variable): + """ See :meth:`pybamm.Symbol._diff()`. """ + children_diffs = [ + child.diff(variable) for child in self.cached_children + ] + if len(children_diffs) == 1: + diff = children_diffs[0] + else: + diff = self.__class__(*children_diffs) + + return diff + def get_children_domains(self, children): # combine domains from children domain = [] diff --git a/pybamm/expression_tree/operations/evaluate_python.py b/pybamm/expression_tree/operations/evaluate_python.py index c8005c5625..ff8bf92853 100644 --- a/pybamm/expression_tree/operations/evaluate_python.py +++ b/pybamm/expression_tree/operations/evaluate_python.py @@ -559,7 +559,7 @@ def __init__(self, symbol): constants[symbol_id] = jax.device_put(constants[symbol_id]) # get a list of constant arguments to input to the function - arg_list = [ + self._arg_list = [ id_to_python_variable(symbol_id, True) for symbol_id in constants.keys() ] @@ -580,9 +580,11 @@ def __init__(self, symbol): # add function def to first line args = "t=None, y=None, y_dot=None, inputs=None, known_evals=None" - if arg_list: - args = ",".join(arg_list) + ", " + args - python_str = "def evaluate_jax({}):\n".format(args) + python_str + if self._arg_list: + args = ",".join(self._arg_list) + ", " + args + python_str = ( + "def evaluate_jax({}):\n".format(args) + python_str + ) # calculate the final variable that will output the result of calling `evaluate` # on `symbol` @@ -606,17 +608,32 @@ def __init__(self, symbol): compiled_function = compile(python_str, result_var, "exec") exec(compiled_function) - n = len(arg_list) - static_argnums = tuple(static_argnums) - self._jit_evaluate = jax.jit(self._evaluate_jax, static_argnums=static_argnums) + self._static_argnums = tuple(static_argnums) + self._jit_evaluate = jax.jit(self._evaluate_jax, + static_argnums=self._static_argnums) + + def get_jacobian(self): + n = len(self._arg_list) - # store a jit version of evaluate_jax's jacobian + # forward mode autodiff wrt y, which is argument 1 after arg_list jacobian_evaluate = jax.jacfwd(self._evaluate_jax, argnums=1 + n) - self._jac_evaluate = jax.jit(jacobian_evaluate, static_argnums=static_argnums) - def get_jacobian(self): + self._jac_evaluate = jax.jit(jacobian_evaluate, + static_argnums=self._static_argnums) + return EvaluatorJaxJacobian(self._jac_evaluate, self._constants) + def get_sensitivities(self): + n = len(self._arg_list) + + # forward mode autodiff wrt inputs, which is argument 3 after arg_list + jacobian_evaluate = jax.jacfwd(self._evaluate_jax, argnums=3 + n) + + self._sens_evaluate = jax.jit(jacobian_evaluate, + static_argnums=self._static_argnums) + + return EvaluatorJaxSensitivities(self._sens_evaluate, self._constants) + def debug(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None): # generated code assumes y is a column vector if y is not None and y.ndim == 1: @@ -668,7 +685,28 @@ def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None): result = self._jac_evaluate(*self._constants, t, y, y_dot, inputs, known_evals) result = result.reshape(result.shape[0], -1) - # don't need known_evals, but need to reproduce Symbol.evaluate signature + if known_evals is not None: + return result, known_evals + else: + return result + + +class EvaluatorJaxSensitivities: + def __init__(self, jac_evaluate, constants): + self._jac_evaluate = jac_evaluate + self._constants = constants + + def evaluate(self, t=None, y=None, y_dot=None, inputs=None, known_evals=None): + """ + Acts as a drop-in replacement for :func:`pybamm.Symbol.evaluate` + """ + # generated code assumes y is a column vector + if y is not None and y.ndim == 1: + y = y.reshape(-1, 1) + + # execute code + result = self._jac_evaluate(*self._constants, t, y, y_dot, inputs, known_evals) + if known_evals is not None: return result, known_evals else: diff --git a/pybamm/models/base_model.py b/pybamm/models/base_model.py index e94868457b..26bf77c613 100644 --- a/pybamm/models/base_model.py +++ b/pybamm/models/base_model.py @@ -273,6 +273,19 @@ def timescale(self, value): """Set the timescale""" self._timescale = value + @property + def length_scales(self): + "Length scales of model" + return self._length_scale + + @length_scales.setter + def length_scales(self, values): + "Set the length scale, converting any numbers to pybamm.Scalar" + for domain, scale in values.items(): + if isinstance(scale, numbers.Number): + values[domain] = pybamm.Scalar(scale) + self._length_scale = values + @property def parameters(self): """Returns all the parameters in the model""" diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py index 5971a45de1..04f32ab325 100644 --- a/pybamm/solvers/base_solver.py +++ b/pybamm/solvers/base_solver.py @@ -3,6 +3,7 @@ # import copy import itertools +from scipy.sparse import block_diag import multiprocessing as mp import numbers import sys @@ -138,7 +139,6 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): Any input parameters to pass to the model when solving t_eval : numeric type, optional The times (in seconds) at which to compute the solution - """ pybamm.logger.info("Start solver set-up") @@ -201,16 +201,50 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): ) model.convert_to_format = "casadi" + # find all the input parameters in the model + input_parameters = {} + for equation in [model.concatenated_rhs, + model.concatenated_algebraic, + model.concatenated_initial_conditions]: + input_parameters.update({ + symbol._id: symbol for symbol in equation.pre_order() + if isinstance(symbol, pybamm.InputParameter) + }) + + # set default calculate sensitivities on model + if not hasattr(model, 'calculate_sensitivities'): + model.calculate_sensitivities = [] + + # see if we need to form the explicit sensitivity equations + calculate_sensitivities_explicit = False + if model.calculate_sensitivities and not isinstance(self, pybamm.IDAKLUSolver): + calculate_sensitivities_explicit = True + + # if we are calculating sensitivities explicitly then the number of + # equations will change + if calculate_sensitivities_explicit: + num_parameters = 0 + for name in model.calculate_sensitivities: + # if not a number, assume its a vector + if isinstance(inputs[name], numbers.Number): + num_parameters += 1 + else: + num_parameters += len(inputs[name]) + model.len_rhs_sens = model.len_rhs * num_parameters + model.len_alg_sens = model.len_alg * num_parameters + if model.convert_to_format != "casadi": # Create Jacobian from concatenated rhs and algebraic - y = pybamm.StateVector(slice(0, model.concatenated_initial_conditions.size)) + y = pybamm.StateVector(slice(0, model.len_rhs_and_alg)) # set up Jacobian object, for re-use of dict jacobian = pybamm.Jacobian() + else: # Convert model attributes to casadi t_casadi = casadi.MX.sym("t") - y_diff = casadi.MX.sym("y_diff", model.concatenated_rhs.size) - y_alg = casadi.MX.sym("y_alg", model.concatenated_algebraic.size) + # Create the symbolic state vectors + y_diff = casadi.MX.sym("y_diff", model.len_rhs) + y_alg = casadi.MX.sym("y_alg", model.len_alg) y_casadi = casadi.vertcat(y_diff, y_alg) p_casadi = {} for name, value in inputs.items(): @@ -219,6 +253,16 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): else: p_casadi[name] = casadi.MX.sym(name, value.shape[0]) p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) + # sensitivity vectors + if calculate_sensitivities_explicit: + pS_casadi_stacked = casadi.vertcat( + *[p_casadi[name] for name in model.calculate_sensitivities] + ) + S_x = casadi.MX.sym("S_x", model.len_rhs_sens) + S_z = casadi.MX.sym("S_z", model.len_alg_sens) + y_and_S = casadi.vertcat(y_diff, S_x, y_alg, S_z) + else: + y_and_S = y_casadi def process(func, name, use_jacobian=None): def report(string): @@ -228,12 +272,52 @@ def report(string): if use_jacobian is None: use_jacobian = model.use_jacobian - if model.convert_to_format != "casadi": - # Process with pybamm functions - if model.convert_to_format == "jax": - report(f"Converting {name} to jax") - jax_func = pybamm.EvaluatorJax(func) + if model.convert_to_format == "jax": + report(f"Converting {name} to jax") + func = pybamm.EvaluatorJax(func) + jacp = None + if model.calculate_sensitivities: + report(( + f"Calculating sensitivities for {name} with respect " + f"to parameters {model.calculate_sensitivities} using jax" + )) + jacp = func.get_sensitivities() + jacp = jacp.evaluate + if use_jacobian: + report(f"Calculating jacobian for {name} using jax") + jac = func.get_jacobian() + jac = jac.evaluate + else: + jac = None + + func = func.evaluate + + elif model.convert_to_format != "casadi": + # Process with pybamm functions, optionally converting + # to python evaluator + if model.calculate_sensitivities: + report(( + f"Calculating sensitivities for {name} with respect " + f"to parameters {model.calculate_sensitivities}" + )) + jacp_dict = { + p: func.diff(pybamm.InputParameter(p)) + for p in model.calculate_sensitivities + } + if model.convert_to_format == "python": + report(f"Converting sensitivities for {name} to python") + jacp_dict = { + p: pybamm.EvaluatorPython(jacp) + for p, jacp in jacp_dict.items() + } + + # jacp should be a function that returns a dict of sensitivities + def jacp(*args, **kwargs): + return {k: v.evaluate(*args, **kwargs) + for k, v in jacp_dict.items()} + else: + jacp = None if use_jacobian: report(f"Calculating jacobian for {name}") @@ -241,9 +325,6 @@ def report(string): if model.convert_to_format == "python": report(f"Converting jacobian for {name} to python") jac = pybamm.EvaluatorPython(jac) - elif model.convert_to_format == "jax": - report(f"Converting jacobian for {name} to jax") - jac = jax_func.get_jacobian() jac = jac.evaluate else: jac = None @@ -251,9 +332,6 @@ def report(string): if model.convert_to_format == "python": report(f"Converting {name} to python") func = pybamm.EvaluatorPython(func) - if model.convert_to_format == "jax": - report(f"Converting {name} to jax") - func = jax_func func = func.evaluate @@ -261,16 +339,101 @@ def report(string): # Process with CasADi report(f"Converting {name} to CasADi") func = func.to_casadi(t_casadi, y_casadi, inputs=p_casadi) + # Add sensitivity vectors to the rhs and algebraic equations + jacp = None + if calculate_sensitivities_explicit: + # The formulation is as per Park, S., Kato, D., Gima, Z., Klein, R., + # & Moura, S. (2018). Optimal experimental design for + # parameterization of an electrochemical lithium-ion battery model. + # Journal of The Electrochemical Society, 165(7), A1309.". See #1100 + # for details + if name == "RHS" and model.len_rhs > 0: + report( + "Creating explicit forward sensitivity equations " + "for rhs using CasADi" + ) + df_dx = casadi.jacobian(func, y_diff) + df_dp = casadi.jacobian(func, pS_casadi_stacked) + S_x_mat = S_x.reshape( + (model.len_rhs, pS_casadi_stacked.shape[0]) + ) + if model.len_alg == 0: + S_rhs = (df_dx @ S_x_mat + df_dp).reshape((-1, 1)) + else: + df_dz = casadi.jacobian(func, y_alg) + S_z_mat = S_z.reshape( + (model.len_alg, pS_casadi_stacked.shape[0]) + ) + S_rhs = (df_dx @ S_x_mat + df_dz @ S_z_mat + df_dp).reshape( + (-1, 1) + ) + func = casadi.vertcat(func, S_rhs) + if name == "algebraic" and model.len_alg > 0: + report( + "Creating explicit forward sensitivity equations " + "for algebraic using CasADi" + ) + dg_dz = casadi.jacobian(func, y_alg) + dg_dp = casadi.jacobian(func, pS_casadi_stacked) + S_z_mat = S_z.reshape( + (model.len_alg, pS_casadi_stacked.shape[0]) + ) + if model.len_rhs == 0: + S_alg = (dg_dz @ S_z_mat + dg_dp).reshape((-1, 1)) + else: + dg_dx = casadi.jacobian(func, y_diff) + S_x_mat = S_x.reshape( + (model.len_rhs, pS_casadi_stacked.shape[0]) + ) + S_alg = (dg_dx @ S_x_mat + dg_dz @ S_z_mat + dg_dp).reshape( + (-1, 1) + ) + func = casadi.vertcat(func, S_alg) + if name == "initial_conditions": + if model.len_rhs == 0 or model.len_alg == 0: + S_0 = casadi.jacobian(func, pS_casadi_stacked).reshape( + (-1, 1) + ) + func = casadi.vertcat(func, S_0) + else: + x0 = func[: model.len_rhs] + z0 = func[model.len_rhs :] + Sx_0 = casadi.jacobian(x0, pS_casadi_stacked).reshape( + (-1, 1) + ) + Sz_0 = casadi.jacobian(z0, pS_casadi_stacked).reshape( + (-1, 1) + ) + func = casadi.vertcat(x0, Sx_0, z0, Sz_0) + elif model.calculate_sensitivities: + report(( + f"Calculating sensitivities for {name} with respect " + f"to parameters {model.calculate_sensitivities} using CasADi" + )) + jacp_dict = {} + for pname in model.calculate_sensitivities: + p_diff = casadi.jacobian(func, p_casadi[pname]) + jacp_dict[pname] = casadi.Function( + name, [t_casadi, y_casadi, p_casadi_stacked], + [p_diff] + ) + # jacp should be a function that returns a dict of sensitivities + + def jacp(*args, **kwargs): + return {k: v(*args, **kwargs) + for k, v in jacp_dict.items()} + if use_jacobian: report(f"Calculating jacobian for {name} using CasADi") - jac_casadi = casadi.jacobian(func, y_casadi) + jac_casadi = casadi.jacobian(func, y_and_S) jac = casadi.Function( - name, [t_casadi, y_casadi, p_casadi_stacked], [jac_casadi] + name, [t_casadi, y_and_S, p_casadi_stacked], [jac_casadi] ) else: jac = None + func = casadi.Function( - name, [t_casadi, y_casadi, p_casadi_stacked], [func] + name, [t_casadi, y_and_S, p_casadi_stacked], [func] ) if name == "residuals": func_call = Residuals(func, name, model) @@ -280,7 +443,14 @@ def report(string): jac_call = SolverCallable(jac, name + "_jac", model) else: jac_call = None - return func, func_call, jac_call + if jacp is not None: + jacp_call = SensitivityCallable( + jacp, name + "_sensitivity_wrt_inputs", model, + model.convert_to_format + ) + else: + jacp_call = None + return func, func_call, jac_call, jacp_call # Process initial conditions initial_conditions = process( @@ -289,8 +459,19 @@ def report(string): use_jacobian=False, )[0] init_eval = InitialConditions(initial_conditions, model) - model.init_eval = init_eval + + if calculate_sensitivities_explicit: + y0_total_size = ( + model.len_rhs + model.len_rhs_sens + + model.len_alg + model.len_alg_sens + ) + init_eval.y_dummy = np.zeros((y0_total_size, 1)) + else: + init_eval.y_dummy = np.zeros((model.len_rhs_and_alg, 1)) + + # Calculate initial conditions model.y0 = init_eval(inputs) + model.init_eval = init_eval if not ics_only: # Check for heaviside and modulo functions in rhs and algebraic and add @@ -364,8 +545,8 @@ def report(string): ) # Process rhs, algebraic and event expressions - rhs, rhs_eval, jac_rhs = process(model.concatenated_rhs, "RHS") - algebraic, algebraic_eval, jac_algebraic = process( + rhs, rhs_eval, jac_rhs, jacp_rhs = process(model.concatenated_rhs, "RHS") + algebraic, algebraic_eval, jac_algebraic, jacp_algebraic = process( model.concatenated_algebraic, "algebraic" ) casadi_terminate_events = [] @@ -425,6 +606,52 @@ def report(string): interpolant_extrapolation_events_eval ) + # if we have changed the equations to include the explicit sensitivity + # equations, then we also need to update the mass matrix and bounds + if calculate_sensitivities_explicit: + if model.len_rhs != 0: + n_inputs = model.len_rhs_sens // model.len_rhs + elif model.len_alg != 0: + n_inputs = model.len_alg_sens // model.len_alg + if model.bounds[0].shape[0] == model.len_rhs_and_alg: + model.bounds = ( + np.repeat(model.bounds[0], n_inputs + 1), + np.repeat(model.bounds[1], n_inputs + 1), + ) + if (model.mass_matrix is not None + and model.mass_matrix.shape[0] == model.len_rhs_and_alg): + + if model.mass_matrix_inv is not None: + model.mass_matrix_inv = pybamm.Matrix( + block_diag( + [model.mass_matrix_inv.entries] * (n_inputs + 1), + format="csr" + ) + ) + model.mass_matrix = pybamm.Matrix( + block_diag( + [model.mass_matrix.entries] * (n_inputs + 1), format="csr" + ) + ) + else: + # take care if calculate_sensitivites used then not used + if model.bounds[0].shape[0] > model.len_rhs_and_alg: + model.bounds = ( + model.bounds[0][:model.len_rhs_and_alg], + model.bounds[1][:model.len_rhs_and_alg], + ) + if (model.mass_matrix is not None and + model.mass_matrix.shape[0] > model.len_rhs_and_alg): + if model.mass_matrix_inv is not None: + model.mass_matrix_inv = pybamm.Matrix( + model.mass_matrix_inv.entries[:model.len_rhs, + :model.len_rhs] + ) + model.mass_matrix = pybamm.Matrix( + model.mass_matrix.entries[:model.len_rhs_and_alg, + :model.len_rhs_and_alg] + ) + # Save CasADi functions for the CasADi solver # Note: when we pass to casadi the ode part of the problem must be in # explicit form so we pre-multiply by the inverse of the mass matrix @@ -435,29 +662,36 @@ def report(string): if len(model.rhs) > 0: mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries) explicit_rhs = mass_matrix_inv @ rhs( - t_casadi, y_casadi, p_casadi_stacked + t_casadi, y_and_S, p_casadi_stacked ) model.casadi_rhs = casadi.Function( - "rhs", [t_casadi, y_casadi, p_casadi_stacked], [explicit_rhs] + "rhs", [t_casadi, y_and_S, p_casadi_stacked], [explicit_rhs] ) model.casadi_algebraic = algebraic + model.casadi_sensitivities_rhs = jacp_rhs + model.casadi_sensitivities_algebraic = jacp_algebraic if len(model.rhs) == 0: # No rhs equations: residuals is algebraic only model.residuals_eval = Residuals(algebraic, "residuals", model) model.jacobian_eval = jac_algebraic + model.sensitivities_eval = jacp_algebraic elif len(model.algebraic) == 0: # No algebraic equations: residuals is rhs only model.residuals_eval = Residuals(rhs, "residuals", model) model.jacobian_eval = jac_rhs + model.sensitivities_eval = jacp_rhs # Calculate consistent initial conditions for the algebraic equations else: all_states = pybamm.NumpyConcatenation( model.concatenated_rhs, model.concatenated_algebraic ) # Process again, uses caching so should be quick - residuals_eval, jacobian_eval = process(all_states, "residuals")[1:] + residuals_eval, jacobian_eval, jacobian_wrtp_eval = process( + all_states, "residuals" + )[1:] model.residuals_eval = residuals_eval model.jacobian_eval = jacobian_eval + model.sensitivities_eval = jacobian_wrtp_eval pybamm.logger.info("Finish solver set-up") @@ -478,13 +712,14 @@ def _set_initial_conditions(self, model, inputs, update_rhs): Whether to update the rhs. True for 'solve', False for 'step'. """ + if self.algebraic_solver is True: # Don't update model.y0 return None elif len(model.algebraic) == 0: if update_rhs is True: # Recalculate initial conditions for the rhs equations - model.y0 = model.init_eval(inputs) + y0 = model.init_eval(inputs) else: # Don't update model.y0 return None @@ -494,7 +729,7 @@ def _set_initial_conditions(self, model, inputs, update_rhs): y0_from_inputs = model.init_eval(inputs) # Reuse old solution for algebraic equations y0_from_model = model.y0 - len_rhs = model.concatenated_rhs.size + len_rhs = model.len_rhs # update model.y0, which is used for initialising the algebraic solver if len_rhs == 0: model.y0 = y0_from_model @@ -502,7 +737,9 @@ def _set_initial_conditions(self, model, inputs, update_rhs): model.y0 = casadi.vertcat( y0_from_inputs[:len_rhs], y0_from_model[len_rhs:] ) - model.y0 = self.calculate_consistent_state(model, 0, inputs) + y0 = self.calculate_consistent_state(model, 0, inputs) + # Make y0 a function of inputs if doing symbolic with casadi + model.y0 = y0 def calculate_consistent_state(self, model, time=0, inputs=None): """ @@ -529,12 +766,13 @@ def calculate_consistent_state(self, model, time=0, inputs=None): if self.root_method is None: return model.y0 try: - root_sol = self.root_method._integrate(model, [time], inputs) + root_sol = self.root_method._integrate(model, np.array([time]), inputs) except pybamm.SolverError as e: raise pybamm.SolverError( "Could not find consistent states: {}".format(e.args[0]) ) pybamm.logger.debug("Found consistent states") + y0 = root_sol.all_ys[0] if isinstance(y0, np.ndarray): y0 = y0.flatten() @@ -548,6 +786,7 @@ def solve( inputs=None, initial_conditions=None, nproc=None, + calculate_sensitivities=False ): """ Execute the solver setup and calculate the solution of the model at @@ -573,6 +812,10 @@ def solve( nproc : int, optional Number of processes to use when solving for more than one set of input parameters. Defaults to value returned by "os.cpu_count()". + calculate_sensitivites : list of str or bool + If true, solver calculates sensitivities of all input parameters. + If only a subset of sensitivities are required, can also pass a + list of input parameter names Returns ------- @@ -589,6 +832,15 @@ def solve( """ pybamm.logger.info("Start solving {} with {}".format(model.name, self.name)) + # get a list-only version of calculate_sensitivities + if isinstance(calculate_sensitivities, bool): + if calculate_sensitivities: + calculate_sensitivities_list = [p for p in inputs.keys()] + else: + calculate_sensitivities_list = [] + else: + calculate_sensitivities_list = calculate_sensitivities + # Make sure model isn't empty if len(model.rhs) == 0 and len(model.algebraic) == 0: if not isinstance(self, pybamm.DummySolver): @@ -630,7 +882,8 @@ def solve( # If "inputs" is a single dict, "inputs_list" is a list of only one dict. inputs_list = inputs if isinstance(inputs, list) else [inputs] ext_and_inputs_list = [ - self._set_up_ext_and_inputs(model, external_variables, inputs) + self._set_up_ext_and_inputs(model, external_variables, inputs, + calculate_sensitivities_list) for inputs in inputs_list ] @@ -641,6 +894,21 @@ def solve( 'when model in format "jax".' ) + # Check that calculate_sensitivites have not been updated + calculate_sensitivities_list.sort() + if not hasattr(model, 'calculate_sensitivities'): + model.calculate_sensitivities = [] + model.calculate_sensitivities.sort() + if (calculate_sensitivities_list != model.calculate_sensitivities): + self.models_set_up.pop(model, None) + # CasadiSolver caches its integrators using model, so delete this too + if isinstance(self, pybamm.CasadiSolver): + self.integrators.pop(model, None) + + # save sensitivity parameters so we can identify them later on + # (FYI: this is used in the Solution class) + model.calculate_sensitivities = calculate_sensitivities_list + # Set up (if not done already) timer = pybamm.Timer() if model not in self.models_set_up: @@ -662,6 +930,7 @@ def solve( self.models_set_up[model][ "initial conditions" ] = model.concatenated_initial_conditions + set_up_time = timer.time() timer.reset() @@ -908,6 +1177,9 @@ def step( save : bool Turn on to store the solution of all previous timesteps + + + Raises ------ :class:`pybamm.ModelError` @@ -971,6 +1243,7 @@ def step( pybamm.logger.verbose( "Start stepping {} with {}".format(model.name, self.name) ) + self.set_up(model, ext_and_inputs) self.models_set_up.update( {model: {"initial conditions": model.concatenated_initial_conditions}} @@ -978,6 +1251,7 @@ def step( t = 0.0 elif model not in self.models_set_up: # Run set up if the model has changed + self.set_up(model, ext_and_inputs) self.models_set_up.update( {model: {"initial conditions": model.concatenated_initial_conditions}} @@ -1159,7 +1433,8 @@ def check_extrapolation(self, solution, events): return [k for k, v in extrap_events.items() if v] - def _set_up_ext_and_inputs(self, model, external_variables, inputs): + def _set_up_ext_and_inputs(self, model, external_variables, inputs, + calculate_sensitivities): """Set up external variables and input parameters""" inputs = inputs or {} @@ -1198,12 +1473,13 @@ def __init__(self, function, name, model): self.timescale = self.model.timescale_eval def __call__(self, t, y, inputs): - if self.name in ["RHS", "algebraic", "residuals"]: - pybamm.logger.debug( - "Evaluating {} for {} at t={}".format( - self.name, self.model.name, t * self.timescale - ) + pybamm.logger.debug( + "Evaluating {} for {} at t={}".format( + self.name, self.model.name, t * self.timescale ) + ) + if self.name in ["RHS", "algebraic", "residuals", "event"]: + return self.function(t, y, inputs).flatten() else: return self.function(t, y, inputs) @@ -1220,6 +1496,35 @@ def function(self, t, y, inputs): return self._function(t, y, inputs=inputs, known_evals={})[0] +class SensitivityCallable: + """A class that will be called by the solver when integrating""" + + def __init__(self, function, name, model, form): + self._function = function + self.form = form + self.name = name + self.model = model + self.timescale = self.model.timescale_eval + + def __call__(self, t, y, inputs): + pybamm.logger.debug( + "Evaluating sensitivities of {} for {} at t={}".format( + self.name, self.model.name, t * self.timescale + ) + ) + return self.function(t, y, inputs) + + def function(self, t, y, inputs): + if self.form == "casadi": + return self._function(t, y, inputs) + elif self.form == "jax": + return self._function(t, y, inputs=inputs) + else: + ret_with_known_evals = \ + self._function(t, y, inputs=inputs, known_evals={}) + return {k: v[0] for k, v in ret_with_known_evals.items()} + + class Residuals(SolverCallable): """Returns information about residuals at time t and state y""" @@ -1238,7 +1543,6 @@ class InitialConditions(SolverCallable): def __init__(self, function, model): super().__init__(function, "initial conditions", model) - self.y_dummy = np.zeros(model.concatenated_initial_conditions.shape) def __call__(self, inputs): if self.form == "casadi": diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index 55704e2161..d08796ba95 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -1,7 +1,7 @@ #include #include -#include /* prototypes for IDA fcts., consts. */ +#include /* prototypes for IDAS fcts., consts. */ #include /* access to serial N_Vector */ #include /* defs. of SUNRabs, SUNRexp, etc. */ #include /* defs. of realtype, sunindextype */ @@ -11,15 +11,24 @@ #include #include #include +#include + +//#include namespace py = pybind11; -using residual_type = std::function( - double, py::array_t, py::array_t)>; -using jacobian_type = - std::function(double, py::array_t, double)>; + +using np_array = py::array_t; +PYBIND11_MAKE_OPAQUE(std::vector); +using residual_type = std::function; +using sensitivities_type = std::function&, realtype, const np_array&, + const np_array&, const std::vector&, + const std::vector& + )>; +using jacobian_type = std::function; + using event_type = - std::function(double, py::array_t)>; -using np_array = py::array_t; + std::function; using jac_get_type = std::function; @@ -27,14 +36,20 @@ class PybammFunctions { public: int number_of_states; + int number_of_parameters; int number_of_events; PybammFunctions(const residual_type &res, const jacobian_type &jac, + const sensitivities_type &sens, const jac_get_type &get_jac_data_in, const jac_get_type &get_jac_row_vals_in, const jac_get_type &get_jac_col_ptrs_in, - const event_type &event, const int n_s, int n_e) - : number_of_states(n_s), number_of_events(n_e), py_res(res), py_jac(jac), + const event_type &event, + const int n_s, int n_e, const int n_p) + : number_of_states(n_s), number_of_events(n_e), + number_of_parameters(n_p), + py_res(res), py_jac(jac), + py_sens(sens), py_event(event), py_get_jac_data(get_jac_data_in), py_get_jac_row_vals(get_jac_row_vals_in), py_get_jac_col_ptrs(get_jac_col_ptrs_in) @@ -61,6 +76,22 @@ class PybammFunctions py_jac(t, y, cj); } + void sensitivities( + std::vector& resvalS, + const double t, const np_array& y, const np_array& yp, + const std::vector& yS, const std::vector& ypS) + { + // this function evaluates the sensitivity equations required by IDAS, + // returning them in resvalS, which is preallocated as a numpy array + // of size (np, n), where n is the number of states and np is the number + // of parameters + // + // yS and ypS are also shape (np, n), y and yp are shape (n) + // + // dF/dy * s_i + dF/dyd * sd + dFdp_i for i in range(np) + py_sens(resvalS, t, y, yp, yS, ypS); + } + np_array get_jac_data() { return py_get_jac_data(); } np_array get_jac_row_vals() { return py_get_jac_row_vals(); } @@ -71,6 +102,7 @@ class PybammFunctions private: residual_type py_res; + sensitivities_type py_sens; jacobian_type py_jac; event_type py_event; jac_get_type py_get_jac_data; @@ -196,39 +228,105 @@ int events(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, return (0); } +int sensitivities(int Ns, realtype t, N_Vector yy, N_Vector yp, + N_Vector resval, N_Vector *yS, N_Vector *ypS, N_Vector *resvalS, + void *user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3) { +// This function computes the sensitivity residual for all sensitivity +// equations. It must compute the vectors +// (∂F/∂y)s i (t)+(∂F/∂ ẏ) ṡ i (t)+(∂F/∂p i ) and store them in resvalS[i]. +// Ns is the number of sensitivities. +// t is the current value of the independent variable. +// yy is the current value of the state vector, y(t). +// yp is the current value of ẏ(t). +// resval contains the current value F of the original DAE residual. +// yS contains the current values of the sensitivities s i . +// ypS contains the current values of the sensitivity derivatives ṡ i . +// resvalS contains the output sensitivity residual vectors. +// Memory allocation for resvalS is handled within idas. +// user data is a pointer to user data. +// tmp1, tmp2, tmp3 are N Vectors of length N which can be used as +// temporary storage. +// +// Return value An IDASensResFn should return 0 if successful, +// a positive value if a recoverable error +// occurred (in which case idas will attempt to correct), +// or a negative value if it failed unrecoverably (in which case the integration is halted and IDA SRES FAIL is returned) +// + PybammFunctions *python_functions_ptr = + static_cast(user_data); + PybammFunctions python_functions = *python_functions_ptr; + + int n = python_functions.number_of_states; + int np = python_functions.number_of_parameters; + + // memory managed by sundials, so pass a destructor that does nothing + auto state_vector_shape = std::vector{n, 1}; + np_array y_np = np_array(state_vector_shape, N_VGetArrayPointer(yy), + py::capsule(&yy, [](void* p) {})); + np_array yp_np = np_array(state_vector_shape, N_VGetArrayPointer(yp), + py::capsule(&yp, [](void* p) {})); + + std::vector yS_np(np); + for (int i = 0; i < np; i++) { + auto capsule = py::capsule(yS + i, [](void* p) {}); + yS_np[i] = np_array(state_vector_shape, N_VGetArrayPointer(yS[i]), capsule); + } + + std::vector ypS_np(np); + for (int i = 0; i < np; i++) { + auto capsule = py::capsule(ypS + i, [](void* p) {}); + ypS_np[i] = np_array(state_vector_shape, N_VGetArrayPointer(ypS[i]), capsule); + } + + std::vector resvalS_np(np); + for (int i = 0; i < np; i++) { + auto capsule = py::capsule(resvalS + i, [](void* p) {}); + resvalS_np[i] = np_array(state_vector_shape, + N_VGetArrayPointer(resvalS[i]), capsule); + } + + realtype *ptr1 = static_cast(resvalS_np[0].request().ptr); + const realtype* resvalSval = N_VGetArrayPointer(resvalS[0]); + + python_functions.sensitivities(resvalS_np, t, y_np, yp_np, yS_np, ypS_np); + + return 0; +} + class Solution { public: - Solution(int retval, np_array t_np, np_array y_np) - : flag(retval), t(t_np), y(y_np) + Solution(int retval, np_array t_np, np_array y_np, np_array yS_np) + : flag(retval), t(t_np), y(y_np), yS(yS_np) { } int flag; np_array t; np_array y; + np_array yS; }; /* main program */ Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, - residual_type res, jacobian_type jac, jac_get_type gjd, - jac_get_type gjrv, jac_get_type gjcp, int nnz, event_type event, + residual_type res, jacobian_type jac, + sensitivities_type sens, + jac_get_type gjd, jac_get_type gjrv, jac_get_type gjcp, + int nnz, event_type event, int number_of_events, int use_jacobian, np_array rhs_alg_id, - np_array atol_np, double rel_tol) + np_array atol_np, double rel_tol, int number_of_parameters) { auto t = t_np.unchecked<1>(); auto y0 = y0_np.unchecked<1>(); auto yp0 = yp0_np.unchecked<1>(); auto atol = atol_np.unchecked<1>(); - int number_of_states; - number_of_states = y0_np.request().size; - int number_of_timesteps; - number_of_timesteps = t_np.request().size; - + int number_of_states = y0_np.request().size; + int number_of_timesteps = t_np.request().size; void *ida_mem; // pointer to memory N_Vector yy, yp, avtol; // y, y', and absolute tolerance - realtype rtol, *yval, *ypval, *atval; + N_Vector *yyS, *ypS; // y, y' for sensitivities + realtype rtol, *yval, *ypval, *atval, *ySval; int retval; SUNMatrix J; SUNLinearSolver LS; @@ -238,8 +336,16 @@ Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, yp = N_VNew_Serial(number_of_states); avtol = N_VNew_Serial(number_of_states); + if (number_of_parameters > 0) { + yyS = N_VCloneVectorArray(number_of_parameters, yy); + ypS = N_VCloneVectorArray(number_of_parameters, yp); + } + // set initial value yval = N_VGetArrayPointer(yy); + if (number_of_parameters > 0) { + ySval = N_VGetArrayPointer(yyS[0]); + } ypval = N_VGetArrayPointer(yp); atval = N_VGetArrayPointer(avtol); int i; @@ -250,6 +356,11 @@ Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, atval[i] = atol[i]; } + for (int is = 0 ; is < number_of_parameters; is++) { + N_VConst(RCONST(0.0), yyS[is]); + N_VConst(RCONST(0.0), ypS[is]); + } + // allocate memory for solver ida_mem = IDACreate(); @@ -266,8 +377,9 @@ Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, IDARootInit(ida_mem, number_of_events, events); // set pybamm functions by passing pointer to it - PybammFunctions pybamm_functions(res, jac, gjd, gjrv, gjcp, event, - number_of_states, number_of_events); + PybammFunctions pybamm_functions(res, jac, sens, gjd, gjrv, gjcp, event, + number_of_states, number_of_events, + number_of_parameters); void *user_data = &pybamm_functions; IDASetUserData(ida_mem, user_data); @@ -282,6 +394,13 @@ Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, IDASetJacFn(ida_mem, jacobian); } + if (number_of_parameters > 0) + { + IDASensInit(ida_mem, number_of_parameters, + IDA_SIMULTANEOUS, sensitivities, yyS, ypS); + IDASensEEtolerances(ida_mem); + } + int t_i = 1; realtype tret; realtype t_next; @@ -290,13 +409,19 @@ Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, // set return vectors std::vector t_return(number_of_timesteps); std::vector y_return(number_of_timesteps * number_of_states); + std::vector yS_return(number_of_parameters * number_of_timesteps * number_of_states); t_return[0] = t(0); - int j; - for (j = 0; j < number_of_states; j++) + for (int j = 0; j < number_of_states; j++) { y_return[j] = yval[j]; } + for (int j = 0; j < number_of_parameters; j++) { + const int base_index = j * number_of_timesteps * number_of_states; + for (int k = 0; k < number_of_states; k++) { + yS_return[base_index + k] = ySval[j * number_of_states + k]; + } + } // calculate consistent initial conditions N_Vector id; @@ -320,39 +445,54 @@ Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, IDASetStopTime(ida_mem, t_next); retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); - if (retval == IDA_TSTOP_RETURN) + if (retval == IDA_TSTOP_RETURN || retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { - t_return[t_i] = tret; - for (j = 0; j < number_of_states; j++) - { - y_return[t_i * number_of_states + j] = yval[j]; + if (number_of_parameters > 0) { + IDAGetSens(ida_mem, &tret, yyS); } - t_i += 1; - } - if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) - { t_return[t_i] = tret; - for (j = 0; j < number_of_states; j++) + for (int j = 0; j < number_of_states; j++) { y_return[t_i * number_of_states + j] = yval[j]; } - break; + for (int j = 0; j < number_of_parameters; j++) { + const int base_index = j * number_of_timesteps * number_of_states + + t_i * number_of_states; + for (int k = 0; k < number_of_states; k++) { + yS_return[base_index + k] = ySval[j * number_of_states + k]; + } + } + t_i += 1; + if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { + break; + } + } } /* Free memory */ + if (number_of_parameters > 0) { + IDASensFree(ida_mem); + } IDAFree(&ida_mem); SUNLinSolFree(LS); SUNMatDestroy(J); N_VDestroy(avtol); N_VDestroy(yp); + if (number_of_parameters > 0) { + N_VDestroyVectorArray(yyS, number_of_parameters); + N_VDestroyVectorArray(ypS, number_of_parameters); + } - py::array_t t_ret = py::array_t((t_i + 1), &t_return[0]); - py::array_t y_ret = - py::array_t((t_i + 1) * number_of_states, &y_return[0]); + np_array t_ret = np_array(t_i, &t_return[0]); + np_array y_ret = np_array(t_i * number_of_states, &y_return[0]); + np_array yS_ret = np_array( + std::vector{number_of_parameters, t_i, number_of_states}, + &yS_return[0] + ); - Solution sol(retval, t_ret, y_ret); + Solution sol(retval, t_ret, y_ret, yS_ret); return sol; } @@ -361,15 +501,20 @@ PYBIND11_MODULE(idaklu, m) { m.doc() = "sundials solvers"; // optional module docstring + py::bind_vector>(m, "VectorNdArray"); + m.def("solve", &solve, "The solve function", py::arg("t"), py::arg("y0"), - py::arg("yp0"), py::arg("res"), py::arg("jac"), py::arg("get_jac_data"), + py::arg("yp0"), py::arg("res"), py::arg("jac"), py::arg("sens"), + py::arg("get_jac_data"), py::arg("get_jac_row_vals"), py::arg("get_jac_col_ptr"), py::arg("nnz"), py::arg("events"), py::arg("number_of_events"), py::arg("use_jacobian"), py::arg("rhs_alg_id"), py::arg("atol"), py::arg("rtol"), + py::arg("number_of_sensitivity_parameters"), py::return_value_policy::take_ownership); py::class_(m, "solution") .def_readwrite("t", &Solution::t) .def_readwrite("y", &Solution::y) + .def_readwrite("yS", &Solution::yS) .def_readwrite("flag", &Solution::flag); } diff --git a/pybamm/solvers/casadi_algebraic_solver.py b/pybamm/solvers/casadi_algebraic_solver.py index 3bcac0bcf7..5125c7c4e0 100644 --- a/pybamm/solvers/casadi_algebraic_solver.py +++ b/pybamm/solvers/casadi_algebraic_solver.py @@ -21,6 +21,7 @@ class CasadiAlgebraicSolver(pybamm.BaseSolver): Any options to pass to the CasADi rootfinder. Please consult `CasADi documentation `_ for details. + """ def __init__(self, tol=1e-6, extra_options=None): @@ -78,7 +79,11 @@ def _integrate(self, model, t_eval, inputs_dict=None): y0_diff = casadi.DM() y0_alg = y0 else: - len_rhs = model.concatenated_rhs.size + # Check y0 to see if it includes sensitivities + if model.len_rhs_and_alg == y0.shape[0]: + len_rhs = model.len_rhs + else: + len_rhs = model.len_rhs * (inputs.shape[0] + 1) y0_diff = y0[:len_rhs] y0_alg = y0[len_rhs:] @@ -141,6 +146,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): "constraints": list(constraints[len_rhs:]), }, ) + timer = pybamm.Timer() integration_time = 0 for idx, t in enumerate(t_eval): @@ -196,9 +202,17 @@ def _integrate(self, model, t_eval, inputs_dict=None): # Concatenate differential part y_diff = casadi.horzcat(*[y0_diff] * len(t_eval)) y_sol = casadi.vertcat(y_diff, y_alg) + # Return solution object (no events, so pass None to t_event, y_event) + + try: + explicit_sensitivities = bool(model.calculate_sensitivities) + except AttributeError: + explicit_sensitivities = False + sol = pybamm.Solution( - [t_eval], y_sol, model, inputs_dict, termination="success" + [t_eval], y_sol, model, inputs_dict, termination="success", + sensitivities=explicit_sensitivities ) sol.integration_time = integration_time return sol diff --git a/pybamm/solvers/casadi_solver.py b/pybamm/solvers/casadi_solver.py index 5117e2930b..8890e05d7a 100644 --- a/pybamm/solvers/casadi_solver.py +++ b/pybamm/solvers/casadi_solver.py @@ -62,7 +62,6 @@ class CasadiSolver(pybamm.BaseSolver): Any options to pass to the CasADi integrator when calling the integrator. Please consult `CasADi documentation `_ for details. - """ def __init__( @@ -79,7 +78,12 @@ def __init__( extra_options_call=None, ): super().__init__( - "problem dependent", rtol, atol, root_method, root_tol, extrap_tol + "problem dependent", + rtol, + atol, + root_method, + root_tol, + extrap_tol, ) if mode in ["safe", "fast", "fast with events", "safe without grid"]: self.mode = mode @@ -101,6 +105,7 @@ def __init__( # Initialize self.integrators = {} self.integrator_specs = {} + self.y_sols = {} pybamm.citations.register("Andersson2019") @@ -117,6 +122,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): inputs_dict : dict, optional Any external variables or input parameters to pass to the model when solving """ + # Record whether there are any symbolic inputs inputs_dict = inputs_dict or {} has_symbolic_inputs = any( @@ -147,11 +153,10 @@ def _integrate(self, model, t_eval, inputs_dict=None): # Create integrator without grid to avoid having to create several times self.create_integrator(model, inputs) solution = self._run_integrator( - model, model.y0, inputs_dict, inputs, t_eval, use_grid=False + model, model.y0, inputs_dict, inputs, t_eval, use_grid=False, ) - solution.termination = "final time" - return solution - elif self.mode in ["fast", "fast with events"] or not model.events: + + if self.mode in ["fast", "fast with events"] or not model.events: if not model.events: pybamm.logger.info("No events found, running fast mode") if self.mode == "fast with events": @@ -187,7 +192,10 @@ def _integrate(self, model, t_eval, inputs_dict=None): # to avoid having to create several times self.create_integrator(model, inputs) # Initialize solution - solution = pybamm.Solution(np.array([t]), y0, model, inputs_dict) + solution = pybamm.Solution( + np.array([t]), y0, model, inputs_dict, + sensitivities=False, + ) solution.solve_time = 0 solution.integration_time = 0 use_grid = False @@ -229,7 +237,8 @@ def _integrate(self, model, t_eval, inputs_dict=None): # halve the step size and try again. try: current_step_sol = self._run_integrator( - model, y0, inputs_dict, inputs, t_window, use_grid=use_grid + model, y0, inputs_dict, inputs, t_window, use_grid=use_grid, + extract_sensitivities_in_solution=False, ) solved = True except pybamm.SolverError: @@ -262,6 +271,11 @@ def _integrate(self, model, t_eval, inputs_dict=None): t = t_window[-1] # update y0 y0 = solution.all_ys[-1][:, -1] + + # now we extract sensitivities from the solution + if (bool(model.calculate_sensitivities)): + solution.sensitivities = True + return solution def _solve_for_event(self, coarse_solution, init_event_signs): @@ -274,6 +288,7 @@ def _solve_for_event(self, coarse_solution, init_event_signs): so that only the times up to the event are returned """ pybamm.logger.debug("Solving for events") + model = coarse_solution.all_models[-1] inputs_dict = coarse_solution.all_inputs[-1] inputs = casadi.vertcat(*[x for x in inputs_dict.values()]) @@ -430,6 +445,7 @@ def integer_bisect(): np.array([t_event]), y_event[:, np.newaxis], "event", + sensitivities=bool(model.calculate_sensitivities) ) solution.integration_time = ( coarse_solution.integration_time + dense_step_sol.integration_time @@ -479,6 +495,7 @@ def create_integrator(self, model, inputs, t_eval=None, use_event_switch=False): Otherwise, the integrator has grid [0,1]. """ pybamm.logger.debug("Creating CasADi integrator") + # Use grid if t_eval is given use_grid = not (t_eval is None) if use_grid is True: @@ -502,7 +519,6 @@ def create_integrator(self, model, inputs, t_eval=None, use_event_switch=False): self.integrators[model][t_eval_shifted_rounded] = integrator return integrator else: - y0 = model.y0 rhs = model.casadi_rhs algebraic = model.casadi_algebraic @@ -525,6 +541,8 @@ def create_integrator(self, model, inputs, t_eval=None, use_event_switch=False): # set up and solve t = casadi.MX.sym("t") p = casadi.MX.sym("p", inputs.shape[0]) + y0 = model.y0 + y_diff = casadi.MX.sym("y_diff", rhs(0, y0, p).shape[0]) y_alg = casadi.MX.sym("y_alg", algebraic(0, y0, p).shape[0]) y_full = casadi.vertcat(y_diff, y_alg) @@ -580,17 +598,37 @@ def create_integrator(self, model, inputs, t_eval=None, use_event_switch=False): return integrator - def _run_integrator(self, model, y0, inputs_dict, inputs, t_eval, use_grid=True): + def _run_integrator(self, model, y0, inputs_dict, + inputs, t_eval, use_grid=True, + extract_sensitivities_in_solution=None, + ): pybamm.logger.debug("Running CasADi integrator") + + # are we solving explicit forward equations? + explicit_sensitivities = bool(model.calculate_sensitivities) + + # by default we extract sensitivities in the solution if we + # are calculating the sensitivities + if extract_sensitivities_in_solution is None: + extract_sensitivities_in_solution = explicit_sensitivities + if use_grid is True: t_eval_shifted = t_eval - t_eval[0] t_eval_shifted_rounded = np.round(t_eval_shifted, decimals=12).tobytes() integrator = self.integrators[model][t_eval_shifted_rounded] else: integrator = self.integrators[model]["no grid"] + len_rhs = model.concatenated_rhs.size + + # Check y0 to see if it includes sensitivities + if explicit_sensitivities: + num_parameters = model.len_rhs_sens // model.len_rhs + len_rhs = len_rhs * (num_parameters + 1) + y0_diff = y0[:len_rhs] y0_alg = y0[len_rhs:] + # Solve try: # Try solving if use_grid is True: @@ -603,7 +641,10 @@ def _run_integrator(self, model, y0, inputs_dict, inputs, t_eval, use_grid=True) ) integration_time = timer.time() y_sol = casadi.vertcat(casadi_sol["xf"], casadi_sol["zf"]) - sol = pybamm.Solution(t_eval, y_sol, model, inputs_dict) + sol = pybamm.Solution( + t_eval, y_sol, model, inputs_dict, + sensitivities=extract_sensitivities_in_solution + ) sol.integration_time = integration_time return sol else: @@ -627,13 +668,16 @@ def _run_integrator(self, model, y0, inputs_dict, inputs, t_eval, use_grid=True) if not z.is_empty(): y_alg = casadi.horzcat(y_alg, z) if z.is_empty(): - sol = pybamm.Solution(t_eval, y_diff, model, inputs_dict) + y_sol = y_diff else: y_sol = casadi.vertcat(y_diff, y_alg) - sol = pybamm.Solution(t_eval, y_sol, model, inputs_dict) - sol.integration_time = integration_time - return sol + sol = pybamm.Solution( + t_eval, y_sol, model, inputs_dict, + sensitivities=extract_sensitivities_in_solution + ) + sol.integration_time = integration_time + return sol except RuntimeError as e: # If it doesn't work raise error raise pybamm.SolverError(e.args[0]) diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index f6b8362702..d18213f5b7 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -53,11 +53,17 @@ def __init__( max_steps="deprecated", ): - if idaklu_spec is None: + if idaklu_spec is None: # pragma: no cover raise ImportError("KLU is not installed") super().__init__( - "ida", rtol, atol, root_method, root_tol, extrap_tol, max_steps + "ida", + rtol, + atol, + root_method, + root_tol, + extrap_tol, + max_steps, ) self.name = "IDA KLU solver" @@ -134,7 +140,7 @@ def _check_atol_type(self, atol, size): if atol.size != size: raise pybamm.SolverError( """Absolute tolerances must be either a scalar or a numpy arrray - of the same shape at y0""" + of the same shape as y0 ({})""".format(size) ) return atol @@ -152,6 +158,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): inputs_dict : dict, optional Any external variables or input parameters to pass to the model when solving """ + inputs_dict = inputs_dict or {} if model.rhs_eval.form == "casadi": # stack inputs inputs = casadi.vertcat(*[x for x in inputs_dict.values()]) @@ -180,7 +187,14 @@ def _integrate(self, model, t_eval, inputs_dict=None): rtol = self._rtol atol = self._check_atol_type(atol, y0.size) - mass_matrix = model.mass_matrix.entries + if model.convert_to_format == "jax": + mass_matrix = model.mass_matrix.entries.toarray() + else: + mass_matrix = model.mass_matrix.entries + + # construct residuals function by binding inputs + def resfn(t, y, ydot): + return model.residuals_eval(t, y, ydot, inputs) if model.jacobian_eval: jac_y0_t0 = model.jacobian_eval(t_eval[0], y0, inputs) @@ -236,18 +250,58 @@ def rootfn(t, y): return return_root # get ids of rhs and algebraic variables - rhs_ids = np.ones(model.rhs_eval(0, y0, inputs).shape) + rhs_ids = np.ones(model.rhs_eval(0, y0, inputs).shape[0]) alg_ids = np.zeros(len(y0) - len(rhs_ids)) ids = np.concatenate((rhs_ids, alg_ids)) + number_of_sensitivity_parameters = 0 + if model.sensitivities_eval is not None: + sens0 = model.sensitivities_eval(t=0, y=y0, inputs=inputs) + number_of_sensitivity_parameters = len(sens0.keys()) + + def sensfn(resvalS, t, y, yp, yS, ypS): + """ + this function evaluates the sensitivity equations required by IDAS, + returning them in resvalS, which is preallocated as a numpy array of size + (np, n), where n is the number of states and np is the number of parameters + + The equations returned are: + + dF/dy * s_i + dF/dyd * sd_i + dFdp_i for i in range(np) + + Parameters + ---------- + resvalS: ndarray of shape (np, n) + returns the sensitivity equations in this preallocated array + t: number + time value + y: ndarray of shape (n) + current state vector + yp: list (np) of ndarray of shape (n) + current time derivative of state vector + yS: list (np) of ndarray of shape (n) + current state vector of sensitivity equations + ypS: list (np) of ndarray of shape (n) + current time derivative of state vector of sensitivity equations + + """ + + dFdy = model.jacobian_eval(t, y, inputs) + dFdyd = mass_matrix + dFdp = model.sensitivities_eval(t, y, inputs) + + for i, dFdp_i in enumerate(dFdp.values()): + resvalS[i][:] = dFdy @ yS[i] - dFdyd @ ypS[i] + dFdp_i + # solve timer = pybamm.Timer() sol = idaklu.solve( t_eval, y0, ydot0, - lambda t, y, ydot: model.residuals_eval(t, y, ydot, inputs), + resfn, jac_class.jac_res, + sensfn, jac_class.get_jac_data, jac_class.get_jac_row_vals, jac_class.get_jac_col_ptrs, @@ -258,6 +312,7 @@ def rootfn(t, y): ids, atol, rtol, + number_of_sensitivity_parameters, ) integration_time = timer.time() @@ -266,7 +321,14 @@ def rootfn(t, y): number_of_states = y0.size y_out = sol.y.reshape((number_of_timesteps, number_of_states)) - # return solution, we need to tranpose y to match scipy's interface + # return sensitivity solution, we need to flatten yS to + # (#timesteps * #states,) to match format used by Solution + if number_of_sensitivity_parameters != 0: + yS_out = { + name: sol.yS[i].reshape(-1, 1) for i, name in enumerate(sens0.keys()) + } + else: + yS_out = False if sol.flag in [0, 2]: # 0 = solved for all t_eval if sol.flag == 0: @@ -283,6 +345,7 @@ def rootfn(t, y): t[-1], np.transpose(y_out[-1])[:, np.newaxis], termination, + sensitivities=yS_out, ) sol.integration_time = integration_time return sol diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py index 3b2b45cf10..c6f62dd710 100644 --- a/pybamm/solvers/processed_variable.py +++ b/pybamm/solvers/processed_variable.py @@ -1,6 +1,7 @@ # # Processed Variable class # +import casadi import numbers import numpy as np import pybamm @@ -38,6 +39,7 @@ def __init__(self, base_variables, base_variables_casadi, solution, warn=True): self.all_ts = solution.all_ts self.all_ys = solution.all_ys + self.all_inputs = solution.all_inputs self.all_inputs_casadi = solution.all_inputs_casadi self.mesh = base_variables[0].mesh @@ -45,6 +47,12 @@ def __init__(self, base_variables, base_variables_casadi, solution, warn=True): self.auxiliary_domains = base_variables[0].auxiliary_domains self.warn = warn + self.symbolic_inputs = solution.has_symbolic_inputs + + # Sensitivity starts off uninitialized, only set when called + self._sensitivities = None + self.solution_sensitivities = solution.sensitivities + # Set timescale self.timescale = solution.timescale_eval self.t_pts = solution.t * self.timescale @@ -520,6 +528,84 @@ def data(self): """Same as entries, but different name""" return self.entries + @property + def sensitivities(self): + """ + Returns a dictionary of sensitivities for each input parameter. + The keys are the input parameters, and the value is a matrix of size + (n_x * n_t, n_p), where n_x is the number of states, n_t is the number of time + points, and n_p is the size of the input parameter + """ + # No sensitivities if there are no inputs + if len(self.all_inputs[0]) == 0: + return {} + # Otherwise initialise and return sensitivities + if self._sensitivities is None: + if self.solution_sensitivities != {}: + self.initialise_sensitivity_explicit_forward() + else: + raise ValueError( + "Cannot compute sensitivities. The 'sensitivities' argument of the " + "solver.solve should be changed from 'None' to allow sensitivities " + "calculations. Check solver documentation for details." + ) + return self._sensitivities + + def initialise_sensitivity_explicit_forward(self): + "Set up the sensitivity dictionary" + inputs_stacked = self.all_inputs_casadi[0] + + # Set up symbolic variables + t_casadi = casadi.MX.sym("t") + y_casadi = casadi.MX.sym("y", self.all_ys[0].shape[0]) + p_casadi = { + name: casadi.MX.sym(name, value.shape[0]) + for name, value in self.all_inputs[0].items() + } + p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) + + # Convert variable to casadi format for differentiating + var_casadi = self.base_variables[0].to_casadi( + t_casadi, y_casadi, inputs=p_casadi + ) + dvar_dy = casadi.jacobian(var_casadi, y_casadi) + dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) + + # Convert to functions and evaluate index-by-index + dvar_dy_func = casadi.Function( + "dvar_dy", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dy] + ) + dvar_dp_func = casadi.Function( + "dvar_dp", [t_casadi, y_casadi, p_casadi_stacked], [dvar_dp] + ) + for index, (ts, ys) in enumerate(zip(self.all_ts, self.all_ys)): + for idx, t in enumerate(ts): + u = ys[:, idx] + next_dvar_dy_eval = dvar_dy_func(t, u, inputs_stacked) + next_dvar_dp_eval = dvar_dp_func(t, u, inputs_stacked) + if index == 0 and idx == 0: + dvar_dy_eval = next_dvar_dy_eval + dvar_dp_eval = next_dvar_dp_eval + else: + dvar_dy_eval = casadi.diagcat(dvar_dy_eval, next_dvar_dy_eval) + dvar_dp_eval = casadi.vertcat(dvar_dp_eval, next_dvar_dp_eval) + + # Compute sensitivity + dy_dp = self.solution_sensitivities["all"] + S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval + + sensitivities = {"all": S_var} + + # Add the individual sensitivity + start = 0 + for name, inp in self.all_inputs[0].items(): + end = start + inp.shape[0] + sensitivities[name] = S_var[:, start:end] + start = end + + # Save attribute + self._sensitivities = sensitivities + class Interpolant0D: def __init__(self, entries): diff --git a/pybamm/solvers/scikits_ode_solver.py b/pybamm/solvers/scikits_ode_solver.py index ddf393744d..0c8913b9c8 100644 --- a/pybamm/solvers/scikits_ode_solver.py +++ b/pybamm/solvers/scikits_ode_solver.py @@ -83,6 +83,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): Any input parameters to pass to the model when solving """ + inputs_dict = inputs_dict or {} if model.rhs_eval.form == "casadi": inputs = casadi.vertcat(*[x for x in inputs_dict.values()]) else: diff --git a/pybamm/solvers/scipy_solver.py b/pybamm/solvers/scipy_solver.py index 7a67e59726..38b30103ac 100644 --- a/pybamm/solvers/scipy_solver.py +++ b/pybamm/solvers/scipy_solver.py @@ -28,9 +28,19 @@ class ScipySolver(pybamm.BaseSolver): """ def __init__( - self, method="BDF", rtol=1e-6, atol=1e-6, extrap_tol=0, extra_options=None + self, + method="BDF", + rtol=1e-6, + atol=1e-6, + extrap_tol=0, + extra_options=None, ): - super().__init__(method, rtol, atol, extrap_tol=extrap_tol) + super().__init__( + method=method, + rtol=rtol, + atol=atol, + extrap_tol=extrap_tol, + ) self.ode_solver = True self.extra_options = extra_options or {} self.name = "Scipy solver ({})".format(method) @@ -56,6 +66,8 @@ def _integrate(self, model, t_eval, inputs_dict=None): various diagnostic messages. """ + # Save inputs dictionary, and if necessary convert inputs to a casadi vector + inputs_dict = inputs_dict or {} if model.convert_to_format == "casadi": inputs = casadi.vertcat(*[x for x in inputs_dict.values()]) else: @@ -116,7 +128,8 @@ def event_fn(t, y): t_event = None y_event = np.array(None) sol = pybamm.Solution( - sol.t, sol.y, model, inputs_dict, t_event, y_event, termination + sol.t, sol.y, model, inputs_dict, t_event, y_event, termination, + sensitivities=bool(model.calculate_sensitivities) ) sol.integration_time = integration_time return sol diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py index 9514d1f937..5c0ef66893 100644 --- a/pybamm/solvers/solution.py +++ b/pybamm/solvers/solution.py @@ -42,6 +42,11 @@ class Solution(object): termination : str String to indicate why the solution terminated + sensitivities: bool or dict + True if sensitivities included as the solution of the explicit forwards + equations. False if no sensitivities included/wanted. Dict if sensitivities are + provided as a dict of {parameter: sensitivities} pairs. + """ def __init__( @@ -53,6 +58,7 @@ def __init__( t_event=None, y_event=None, termination="final time", + sensitivities=False ): if not isinstance(all_ts, list): all_ts = [all_ts] @@ -62,21 +68,27 @@ def __init__( all_models = [all_models] self._all_ts = all_ts self._all_ys = all_ys + self._all_ys_and_sens = all_ys self._all_models = all_models + # Set up inputs + if not isinstance(all_inputs, list): + all_inputs_copy = dict(all_inputs) + for key, value in all_inputs_copy.items(): + if isinstance(value, numbers.Number): + all_inputs_copy[key] = np.array([value]) + self.all_inputs = [all_inputs_copy] + else: + self.all_inputs = all_inputs + + self.sensitivities = sensitivities + self._t_event = t_event self._y_event = y_event self._termination = termination - # Set up inputs - if not isinstance(all_inputs, list): - for key, value in all_inputs.items(): - if isinstance(value, numbers.Number): - all_inputs[key] = np.array([value]) - all_inputs = [all_inputs] - self.all_inputs = all_inputs self.has_symbolic_inputs = any( - isinstance(v, casadi.MX) for v in all_inputs[0].values() + isinstance(v, casadi.MX) for v in self.all_inputs[0].values() ) # Copy the timescale_eval and lengthscale_evals if they exist @@ -93,6 +105,12 @@ def __init__( for domain, scale in all_models[0].length_scales.items() } + # Events + self._t_event = t_event + self._y_event = y_event + self._termination = termination + + # Initialize times self.set_up_time = None self.solve_time = None self.integration_time = None @@ -113,6 +131,119 @@ def __init__( # Solution now uses CasADi pybamm.citations.register("Andersson2019") + def extract_explicit_sensitivities(self): + # if we got here, we havn't set y yet + self.set_y() + + # extract sensitivities from full y solution + self._y, self._sensitivities = \ + self._extract_explicit_sensitivities( + self.all_models[0], self.y, self.t, self.all_inputs[0] + ) + + # make sure we remove all sensitivities from all_ys + for index, (model, ys, ts, inputs) in enumerate( + zip(self.all_models, self.all_ys, self.all_ts, + self.all_inputs) + ): + self._all_ys[index], _ = \ + self._extract_explicit_sensitivities( + model, ys, ts, inputs + ) + + def _extract_explicit_sensitivities(self, model, y, t_eval, inputs): + """ + given a model and a solution y, extracts the sensitivities + + Parameters + -------- + model : :class:`pybamm.BaseModel` + A model that has been already setup by this base solver + y: ndarray + The solution of the full explicit sensitivity equations + t_eval: ndarray + The evaluation times + inputs: dict + parameter inputs + + Returns + ------- + y: ndarray + The solution of the ode/dae in model + sensitivities: dict of (string: ndarray) + A dictionary of parameter names, and the corresponding solution of + the sensitivity equations + """ + + n_states = model.len_rhs_and_alg + n_rhs = model.len_rhs + n_alg = model.len_alg + # Get the point where the algebraic equations start + if model.len_rhs != 0: + n_p = model.len_rhs_sens // model.len_rhs + else: + n_p = model.len_alg_sens // model.len_alg + len_rhs_and_sens = model.len_rhs + model.len_rhs_sens + + n_t = len(t_eval) + # y gets the part of the solution vector that correspond to the + # actual ODE/DAE solution + + # save sensitivities as a dictionary + # first save the whole sensitivity matrix + # reshape using Fortran order to get the right array: + # t0_x0_p0, t0_x0_p1, ..., t0_x0_pn + # t0_x1_p0, t0_x1_p1, ..., t0_x1_pn + # ... + # t0_xn_p0, t0_xn_p1, ..., t0_xn_pn + # t1_x0_p0, t1_x0_p1, ..., t1_x0_pn + # t1_x1_p0, t1_x1_p1, ..., t1_x1_pn + # ... + # t1_xn_p0, t1_xn_p1, ..., t1_xn_pn + # ... + # tn_x0_p0, tn_x0_p1, ..., tn_x0_pn + # tn_x1_p0, tn_x1_p1, ..., tn_x1_pn + # ... + # tn_xn_p0, tn_xn_p1, ..., tn_xn_pn + # 1, Extract rhs and alg sensitivities and reshape into 3D matrices + # with shape (n_p, n_states, n_t) + if isinstance(y, casadi.DM): + y_full = y.full() + else: + y_full = y + ode_sens = y_full[n_rhs:len_rhs_and_sens, :].reshape(n_p, n_rhs, n_t) + alg_sens = y_full[len_rhs_and_sens + n_alg :, :].reshape( + n_p, n_alg, n_t + ) + # 2. Concatenate into a single 3D matrix with shape (n_p, n_states, n_t) + # i.e. along first axis + full_sens_matrix = np.concatenate([ode_sens, alg_sens], axis=1) + # Transpose and reshape into a (n_states * n_t, n_p) matrix + full_sens_matrix = full_sens_matrix.transpose(2, 1, 0).reshape( + n_t * n_states, n_p + ) + + # Save the full sensitivity matrix + sensitivity = {"all": full_sens_matrix} + + # also save the sensitivity wrt each parameter (read the columns of the + # sensitivity matrix) + start = 0 + for name in model.calculate_sensitivities: + inp = inputs[name] + input_size = inp.shape[0] + end = start + input_size + sensitivity[name] = full_sens_matrix[:, start:end] + start = end + + y_dae = np.vstack( + [ + y[: model.len_rhs, :], + y[len_rhs_and_sens : len_rhs_and_sens + model.len_alg, :], + ] + ) + return y_dae, sensitivity + @property def t(self): """Times at which the solution is evaluated""" @@ -134,8 +265,31 @@ def y(self): return self._y except AttributeError: self.set_y() + + # if y is evaluated before sensitivities then need to extract them + if isinstance(self._sensitivities, bool) and self._sensitivities: + self.extract_explicit_sensitivities() + return self._y + @property + def sensitivities(self): + """Values of the sensitivities. Returns a dict of param_name: np_array""" + if isinstance(self._sensitivities, bool): + if self._sensitivities: + self.extract_explicit_sensitivities() + else: + self._sensitivities = {} + return self._sensitivities + + @sensitivities.setter + def sensitivities(self, value): + """Updates the sensitivity""" + # sensitivities must be a dict or bool + if not isinstance(value, (bool, dict)): + raise TypeError('sensitivities arg needs to be a bool or dict') + self._sensitivities = value + def set_y(self): try: if isinstance(self.all_ys[0], (casadi.DM, casadi.MX)): @@ -279,6 +433,10 @@ def set_summary_variables(self, all_summary_variables): def update(self, variables): """Add ProcessedVariables to the dictionary of variables in the solution""" + # make sure that sensitivities are extracted if required + if isinstance(self._sensitivities, bool) and self._sensitivities: + self.extract_explicit_sensitivities() + # Convert single entry to list if isinstance(variables, str): variables = [variables] diff --git a/tests/integration/test_models/standard_model_tests.py b/tests/integration/test_models/standard_model_tests.py index a6dbf520d5..5fa5fa2fa1 100644 --- a/tests/integration/test_models/standard_model_tests.py +++ b/tests/integration/test_models/standard_model_tests.py @@ -62,7 +62,8 @@ def test_processing_disc(self, disc=None): # Model should still be well-posed after processing self.model.check_well_posedness(post_discretisation=True) - def test_solving(self, solver=None, t_eval=None): + def test_solving(self, solver=None, t_eval=None, inputs=None, + calculate_sensitivities=False): # Overwrite solver if given if solver is not None: self.solver = solver @@ -80,7 +81,9 @@ def test_solving(self, solver=None, t_eval=None): if t_eval is None: t_eval = np.linspace(0, 3600 / Crate, 100) - self.solution = self.solver.solve(self.model, t_eval) + self.solution = self.solver.solve( + self.model, t_eval, inputs=inputs, + ) def test_outputs(self): # run the standard output tests @@ -89,6 +92,52 @@ def test_outputs(self): ) std_out_test.test_all() + def test_sensitivities(self, param_name, param_value, + output_name='Terminal voltage [V]'): + + self.parameter_values.update({param_name: param_value}) + Crate = abs( + self.parameter_values["Current function [A]"] + / self.parameter_values["Nominal cell capacity [A.h]"] + ) + t_eval = np.linspace(0, 3600 / Crate, 100) + + # make param_name an input + self.parameter_values.update({param_name: "[input]"}) + inputs = {param_name: param_value} + + self.test_processing_parameters() + self.test_processing_disc() + + # Use tighter default tolerances for testing + self.solver.rtol = 1e-8 + self.solver.atol = 1e-8 + + self.solution = self.solver.solve( + self.model, t_eval, inputs=inputs, + calculate_sensitivities=True + ) + output_sens = self.solution[output_name].sensitivities[param_name] + + # check via finite differencing + h = 1e-6 * param_value + inputs_plus = {param_name: (param_value + 0.5 * h)} + inputs_neg = {param_name: (param_value - 0.5 * h)} + sol_plus = self.solver.solve( + self.model, t_eval, inputs=inputs_plus, + ) + output_plus = sol_plus[output_name](t=t_eval) + sol_neg = self.solver.solve( + self.model, t_eval, inputs=inputs_neg + ) + output_neg = sol_neg[output_name](t=t_eval) + fd = ((np.array(output_plus) - np.array(output_neg)) / h) + fd = fd.transpose().reshape(-1, 1) + np.testing.assert_allclose( + output_sens, fd, + rtol=1e-2, atol=1e-6, + ) + def test_all( self, param=None, disc=None, solver=None, t_eval=None, skip_output_tests=False ): diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py index b0b65673d6..9f52b6a7d0 100644 --- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py +++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py @@ -21,6 +21,20 @@ def test_basic_processing(self): ) modeltest.test_all() + def test_sensitivities(self): + options = {"thermal": "isothermal"} + model = pybamm.lithium_ion.DFN(options) + # use Ecker parameters for nonlinear diffusion + param = pybamm.ParameterValues(chemistry=pybamm.parameter_sets.Ecker2015) + var = pybamm.standard_spatial_vars + var_pts = {var.x_n: 10, var.x_s: 10, var.x_p: 10, var.r_n: 5, var.r_p: 5} + modeltest = tests.StandardModelTest( + model, parameter_values=param, var_pts=var_pts + ) + modeltest.test_sensitivities( + 'Current function [A]', 0.15652, + ) + def test_basic_processing_1plus1D(self): options = {"current collector": "potential pair", "dimensionality": 1} model = pybamm.lithium_ion.DFN(options) diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spm.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spm.py index ed2ab9b447..7fc44d4e94 100644 --- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spm.py +++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spm.py @@ -17,6 +17,15 @@ def test_basic_processing(self): modeltest = tests.StandardModelTest(model, parameter_values=param) modeltest.test_all() + def test_sensitivities(self): + options = {"thermal": "isothermal"} + model = pybamm.lithium_ion.SPM(options) + param = pybamm.ParameterValues(chemistry=pybamm.parameter_sets.Ecker2015) + modeltest = tests.StandardModelTest(model, parameter_values=param) + modeltest.test_sensitivities( + 'Current function [A]', 0.15652, + ) + def test_basic_processing_1plus1D(self): options = {"current collector": "potential pair", "dimensionality": 1} diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spme.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spme.py index 472abd8541..62cec28bcd 100644 --- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spme.py +++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_spme.py @@ -18,6 +18,16 @@ def test_basic_processing(self): modeltest = tests.StandardModelTest(model, parameter_values=param) modeltest.test_all() + def test_sensitivities(self): + options = {"thermal": "isothermal"} + model = pybamm.lithium_ion.SPMe(options) + # use Ecker parameters for nonlinear diffusion + param = pybamm.ParameterValues(chemistry=pybamm.parameter_sets.Ecker2015) + modeltest = tests.StandardModelTest(model, parameter_values=param) + modeltest.test_sensitivities( + 'Current function [A]', 0.15652, + ) + def test_basic_processing_python(self): options = {"thermal": "isothermal"} model = pybamm.lithium_ion.SPMe(options) diff --git a/tests/integration/test_solvers/test_idaklu.py b/tests/integration/test_solvers/test_idaklu.py index d61249a08c..805c6ffbe1 100644 --- a/tests/integration/test_solvers/test_idaklu.py +++ b/tests/integration/test_solvers/test_idaklu.py @@ -19,6 +19,49 @@ def test_on_spme(self): solution = pybamm.IDAKLUSolver().solve(model, t_eval) np.testing.assert_array_less(1, solution.t.size) + def test_on_spme_sensitivities(self): + param_name = 'Typical current [A]' + param_value = 0.15652 + model = pybamm.lithium_ion.SPMe() + geometry = model.default_geometry + param = model.default_parameter_values + param.update({param_name: "[input]"}) + inputs = {param_name: param_value} + param.process_model(model) + param.process_geometry(geometry) + mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + t_eval = np.linspace(0, 3600, 100) + solver = pybamm.IDAKLUSolver(rtol=1e-10, atol=1e-10) + solution = solver.solve( + model, t_eval, + inputs=inputs, + calculate_sensitivities=True, + ) + np.testing.assert_array_less(1, solution.t.size) + + # evaluate the sensitivities using idas + dyda_ida = solution.sensitivities[param_name] + + # evaluate the sensitivities using finite difference + h = 1e-6 + sol_plus = solver.solve( + model, t_eval, + inputs={param_name: param_value + 0.5 * h} + ) + sol_neg = solver.solve( + model, t_eval, + inputs={param_name: param_value - 0.5 * h} + ) + dyda_fd = (sol_plus.y - sol_neg.y) / h + dyda_fd = dyda_fd.transpose().reshape(-1, 1) + + np.testing.assert_allclose( + dyda_ida, dyda_fd, + rtol=1e-2, atol=1e-3, + ) + def test_set_tol_by_variable(self): model = pybamm.lithium_ion.SPMe() geometry = model.default_geometry diff --git a/tests/unit/test_models/test_base_model.py b/tests/unit/test_models/test_base_model.py index fb23b0a34c..22f2c3f457 100644 --- a/tests/unit/test_models/test_base_model.py +++ b/tests/unit/test_models/test_base_model.py @@ -101,6 +101,12 @@ def test_boundary_conditions_set_get(self): with self.assertRaisesRegex(pybamm.ModelError, "boundary condition"): model.boundary_conditions = bad_bcs + def test_length_scales(self): + model = pybamm.BaseModel() + model.length_scales = {"a": 1.3} + self.assertIsInstance(model.length_scales["a"], pybamm.Scalar) + self.assertEqual(model.length_scales["a"].value, 1.3) + def test_variables_set_get(self): model = pybamm.BaseModel() variables = {"c": "alpha", "d": "beta"} diff --git a/tests/unit/test_solvers/test_algebraic_solver.py b/tests/unit/test_solvers/test_algebraic_solver.py index 9882ae2f4f..15e8a7bbd4 100644 --- a/tests/unit/test_solvers/test_algebraic_solver.py +++ b/tests/unit/test_solvers/test_algebraic_solver.py @@ -38,13 +38,14 @@ def test_wrong_solver(self): def test_simple_root_find(self): # Simple system: a single algebraic equation - class Model: + class Model(pybamm.BaseModel): y0 = np.array([2]) rhs = {} timescale_eval = 1 length_scales = {} jac_algebraic_eval = None convert_to_format = "python" + len_rhs_and_alg = 1 def algebraic_eval(self, t, y, inputs): return y + 2 @@ -61,13 +62,14 @@ def algebraic_eval(self, t, y, inputs): self.assertNotEqual(solution.y, -2) def test_root_find_fail(self): - class Model: + class Model(pybamm.BaseModel): y0 = np.array([2]) rhs = {} timescale_eval = 1 length_scales = {} jac_algebraic_eval = None convert_to_format = "casadi" + len_rhs_and_alg = 1 def algebraic_eval(self, t, y, inputs): # algebraic equation has no real root @@ -92,12 +94,13 @@ def test_with_jacobian(self): A = np.array([[4, 3], [1, -1]]) b = np.array([0, 7]) - class Model: + class Model(pybamm.BaseModel): y0 = np.zeros(2) rhs = {} timescale_eval = 1 length_scales = {} convert_to_format = "python" + len_rhs_and_alg = 2 def algebraic_eval(self, t, y, inputs): return A @ y - b diff --git a/tests/unit/test_solvers/test_base_solver.py b/tests/unit/test_solvers/test_base_solver.py index 12a4530626..7967c96949 100644 --- a/tests/unit/test_solvers/test_base_solver.py +++ b/tests/unit/test_solvers/test_base_solver.py @@ -125,6 +125,7 @@ def __init__(self): ) self.convert_to_format = "casadi" self.bounds = (np.array([-np.inf]), np.array([np.inf])) + self.len_rhs_and_alg = 1 self.interpolant_extrapolation_events_eval = [] def rhs_eval(self, t, y, inputs): @@ -162,6 +163,8 @@ def __init__(self): ) self.convert_to_format = "casadi" self.bounds = (-np.inf * np.ones(4), np.inf * np.ones(4)) + self.len_rhs = 1 + self.len_rhs_and_alg = 4 self.interpolant_extrapolation_events_eval = [] def rhs_eval(self, t, y, inputs): @@ -322,6 +325,58 @@ def test_extrapolation_warnings(self): with self.assertWarns(pybamm.SolverWarning): solver.solve(model, t_eval=[0, 1]) + @unittest.skipIf(not pybamm.have_idaklu(), "idaklu solver is not installed") + def test_sensitivities(self): + + def exact_diff_a(y, a, b): + return np.array([ + [y[0]**2 + 2 * a], + [y[0]] + ]) + + def exact_diff_b(y, a, b): + return np.array([[y[0]], [0]]) + + for convert_to_format in ['', 'python', 'casadi', 'jax']: + model = pybamm.BaseModel() + v = pybamm.Variable("v") + u = pybamm.Variable("u") + a = pybamm.InputParameter("a") + b = pybamm.InputParameter("b") + model.rhs = {v: a * v**2 + b * v + a**2} + model.algebraic = {u: a * v - u} + model.initial_conditions = {v: 1, u: a * 1} + model.convert_to_format = convert_to_format + solver = pybamm.IDAKLUSolver(root_method='lm') + model.calculate_sensitivities = ['a', 'b'] + solver.set_up(model, inputs={'a': 0, 'b': 0}) + all_inputs = [] + for v_value in [0.1, -0.2, 1.5, 8.4]: + for u_value in [0.13, -0.23, 1.3, 13.4]: + for a_value in [0.12, 1.5]: + for b_value in [0.82, 1.9]: + y = np.array([v_value, u_value]) + t = 0 + inputs = {'a': a_value, 'b': b_value} + all_inputs.append((t, y, inputs)) + for t, y, inputs in all_inputs: + if model.convert_to_format == 'casadi': + use_inputs = casadi.vertcat(*[x for x in inputs.values()]) + else: + use_inputs = inputs + + sens = model.sensitivities_eval( + t, y, use_inputs + ) + np.testing.assert_allclose( + sens['a'], + exact_diff_a(y, inputs['a'], inputs['b']) + ) + np.testing.assert_allclose( + sens['b'], + exact_diff_b(y, inputs['a'], inputs['b']) + ) + if __name__ == "__main__": print("Add -v for more debug output") diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py index 6a3f286b4b..c98d4f2e63 100644 --- a/tests/unit/test_solvers/test_casadi_solver.py +++ b/tests/unit/test_solvers/test_casadi_solver.py @@ -6,7 +6,6 @@ import numpy as np from tests import get_mesh_for_testing, get_discretisation_for_testing from scipy.sparse import eye -from scipy.optimize import least_squares class TestCasadiSolver(unittest.TestCase): @@ -538,71 +537,115 @@ def test_casadi_safe_no_termination(self): solver.solve(model, t_eval=[0, 1]) -class TestCasadiSolverSensitivity(unittest.TestCase): - def test_solve_with_symbolic_input(self): - # Simple system: a single differential equation - var = pybamm.Variable("var") +class TestCasadiSolverODEsWithForwardSensitivityEquations(unittest.TestCase): + def test_solve_sensitivity_scalar_var_scalar_input(self): + # Create model model = pybamm.BaseModel() - model.rhs = {var: pybamm.InputParameter("param")} - model.initial_conditions = {var: 2} - model.variables = {"var": var} - - # create discretisation - disc = pybamm.Discretisation() - disc.process_model(model) + var = pybamm.Variable("var") + p = pybamm.InputParameter("p") + model.rhs = {var: p * var} + model.initial_conditions = {var: 1} + model.variables = {"var squared": var ** 2} # Solve - solver = pybamm.CasadiSolver() - t_eval = np.linspace(0, 1) - solution = solver.solve(model, t_eval) - np.testing.assert_array_almost_equal( - solution["var"].value({"param": 7}).full().flatten(), 2 + 7 * t_eval + # Make sure that passing in extra options works + solver = pybamm.CasadiSolver( + mode="fast", rtol=1e-10, atol=1e-10 ) - np.testing.assert_array_almost_equal( - solution["var"].value({"param": -3}).full().flatten(), 2 - 3 * t_eval + t_eval = np.linspace(0, 1, 80) + solution = solver.solve(model, t_eval, inputs={"p": 0.1}, + calculate_sensitivities=True) + np.testing.assert_array_equal(solution.t, t_eval) + np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t)) + np.testing.assert_allclose( + solution.sensitivities["p"], + (solution.t * np.exp(0.1 * solution.t))[:, np.newaxis], ) - np.testing.assert_array_almost_equal( - solution["var"].sensitivity({"param": 3}).full().flatten(), t_eval + np.testing.assert_allclose( + solution["var squared"].data, np.exp(0.1 * solution.t) ** 2 + ) + np.testing.assert_allclose( + solution["var squared"].sensitivities["p"], + (2 * np.exp(0.1 * solution.t) * solution.t * np.exp(0.1 * solution.t))[ + :, np.newaxis + ], ) - def test_least_squares_fit(self): - # Simple system: a single algebraic equation - var1 = pybamm.Variable("var1", domain="negative electrode") - var2 = pybamm.Variable("var2", domain="negative electrode") + # More complicated model + # Create model model = pybamm.BaseModel() + var = pybamm.Variable("var") p = pybamm.InputParameter("p") q = pybamm.InputParameter("q") - model.rhs = {var1: -var1} - model.algebraic = {var2: (var2 - p)} - model.initial_conditions = {var1: 1, var2: 3} - model.variables = {"objective": (var2 - q) ** 2 + (p - 3) ** 2} - - # create discretisation - disc = get_discretisation_for_testing() - disc.process_model(model) + r = pybamm.InputParameter("r") + s = pybamm.InputParameter("s") + model.rhs = {var: p * q} + model.initial_conditions = {var: r} + model.variables = {"var times s": var * s} # Solve - solver = pybamm.CasadiSolver() - solution = solver.solve(model, np.linspace(0, 1)) - sol_var = solution["objective"] - - def objective(x): - return sol_var.value({"p": x[0], "q": x[1]}).full().flatten() - - # without jacobian - lsq_sol = least_squares(objective, [2, 2], method="lm") - np.testing.assert_array_almost_equal(lsq_sol.x, [3, 3], decimal=3) - - def jac(x): - return sol_var.sensitivity({"p": x[0], "q": x[1]}) - - # with jacobian - lsq_sol = least_squares(objective, [2, 2], jac=jac, method="lm") - np.testing.assert_array_almost_equal(lsq_sol.x, [3, 3], decimal=3) + # Make sure that passing in extra options works + solver = pybamm.CasadiSolver( + rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve( + model, t_eval, inputs={"p": 0.1, "q": 2, "r": -1, "s": 0.5}, + calculate_sensitivities=True, + ) + np.testing.assert_allclose(solution.y[0], -1 + 0.2 * solution.t) + np.testing.assert_allclose( + solution.sensitivities["p"], (2 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution.sensitivities["q"], (0.1 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose(solution.sensitivities["r"], 1) + np.testing.assert_allclose(solution.sensitivities["s"], 0) + np.testing.assert_allclose( + solution.sensitivities["all"], + np.hstack( + [ + solution.sensitivities["p"], + solution.sensitivities["q"], + solution.sensitivities["r"], + solution.sensitivities["s"], + ] + ), + ) + np.testing.assert_allclose( + solution["var times s"].data, 0.5 * (-1 + 0.2 * solution.t) + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["p"], + 0.5 * (2 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["q"], + 0.5 * (0.1 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose(solution["var times s"].sensitivities["r"], 0.5) + np.testing.assert_allclose( + solution["var times s"].sensitivities["s"], + (-1 + 0.2 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["all"], + np.hstack( + [ + solution["var times s"].sensitivities["p"], + solution["var times s"].sensitivities["q"], + solution["var times s"].sensitivities["r"], + solution["var times s"].sensitivities["s"], + ] + ), + ) - def test_solve_with_symbolic_input_1D_scalar_input(self): + def test_solve_sensitivity_vector_var_scalar_input(self): var = pybamm.Variable("var", "negative electrode") model = pybamm.BaseModel() + # Set length scales to avoid warning + model.length_scales = {"negative electrode": 1} param = pybamm.InputParameter("param") model.rhs = {var: -param * var} model.initial_conditions = {var: 2} @@ -611,127 +654,270 @@ def test_solve_with_symbolic_input_1D_scalar_input(self): # create discretisation disc = get_discretisation_for_testing() disc.process_model(model) + n = disc.mesh["negative electrode"].npts # Solve - scalar input solver = pybamm.CasadiSolver() t_eval = np.linspace(0, 1) - solution = solver.solve(model, t_eval) + solution = solver.solve(model, t_eval, inputs={"param": 7}, + calculate_sensitivities=["param"]) np.testing.assert_array_almost_equal( - solution["var"].value({"param": 7}), - np.repeat(2 * np.exp(-7 * t_eval), 40)[:, np.newaxis], - decimal=4, + solution["var"].data, np.tile(2 * np.exp(-7 * t_eval), (n, 1)), decimal=4, ) np.testing.assert_array_almost_equal( - solution["var"].value({"param": 3}), - np.repeat(2 * np.exp(-3 * t_eval), 40)[:, np.newaxis], + solution["var"].sensitivities["param"], + np.repeat(-2 * t_eval * np.exp(-7 * t_eval), n)[:, np.newaxis], decimal=4, ) - np.testing.assert_array_almost_equal( - solution["var"].sensitivity({"param": 3}), - np.repeat( - -2 * t_eval * np.exp(-3 * t_eval), disc.mesh["negative electrode"].npts - )[:, np.newaxis], - decimal=4, + + # More complicated model + # Create model + model = pybamm.BaseModel() + # Set length scales to avoid warning + model.length_scales = {"negative electrode": 1} + var = pybamm.Variable("var", "negative electrode") + p = pybamm.InputParameter("p") + q = pybamm.InputParameter("q") + r = pybamm.InputParameter("r") + s = pybamm.InputParameter("s") + model.rhs = {var: p * q} + model.initial_conditions = {var: r} + model.variables = {"var times s": var * s} + + # Discretise + disc.process_model(model) + + # Solve + # Make sure that passing in extra options works + solver = pybamm.CasadiSolver( + rtol=1e-10, atol=1e-10, + ) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve( + model, t_eval, inputs={"p": 0.1, "q": 2, "r": -1, "s": 0.5}, + calculate_sensitivities=True, + ) + np.testing.assert_allclose(solution.y, np.tile(-1 + 0.2 * solution.t, (n, 1))) + np.testing.assert_allclose( + solution.sensitivities["p"], np.repeat(2 * solution.t, n)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution.sensitivities["q"], np.repeat(0.1 * solution.t, n)[:, np.newaxis], + ) + np.testing.assert_allclose(solution.sensitivities["r"], 1) + np.testing.assert_allclose(solution.sensitivities["s"], 0) + np.testing.assert_allclose( + solution.sensitivities["all"], + np.hstack( + [ + solution.sensitivities["p"], + solution.sensitivities["q"], + solution.sensitivities["r"], + solution.sensitivities["s"], + ] + ), + ) + np.testing.assert_allclose( + solution["var times s"].data, np.tile(0.5 * (-1 + 0.2 * solution.t), (n, 1)) + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["p"], + np.repeat(0.5 * (2 * solution.t), n)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["q"], + np.repeat(0.5 * (0.1 * solution.t), n)[:, np.newaxis], + ) + np.testing.assert_allclose(solution["var times s"].sensitivities["r"], 0.5) + np.testing.assert_allclose( + solution["var times s"].sensitivities["s"], + np.repeat(-1 + 0.2 * solution.t, n)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["all"], + np.hstack( + [ + solution["var times s"].sensitivities["p"], + solution["var times s"].sensitivities["q"], + solution["var times s"].sensitivities["r"], + solution["var times s"].sensitivities["s"], + ] + ), ) - def test_solve_with_symbolic_input_1D_vector_input(self): + def test_solve_sensitivity_scalar_var_vector_input(self): var = pybamm.Variable("var", "negative electrode") model = pybamm.BaseModel() + # Set length scales to avoid warning + model.length_scales = {"negative electrode": 1} + param = pybamm.InputParameter("param", "negative electrode") model.rhs = {var: -param * var} model.initial_conditions = {var: 2} - model.variables = {"var": var} + model.variables = { + "var": var, + "integral of var": pybamm.Integral(var, pybamm.standard_spatial_vars.x_n), + } # create discretisation - disc = get_discretisation_for_testing() + mesh = get_mesh_for_testing(xpts=5) + spatial_methods = {"macroscale": pybamm.FiniteVolume()} + disc = pybamm.Discretisation(mesh, spatial_methods) disc.process_model(model) - - # Solve - scalar input - solver = pybamm.CasadiSolver() - solution = solver.solve(model, np.linspace(0, 1)) n = disc.mesh["negative electrode"].npts - solver = pybamm.CasadiSolver() + # Solve - constant input + solver = pybamm.CasadiSolver( + mode="fast", rtol=1e-10, atol=1e-10 + ) t_eval = np.linspace(0, 1) - solution = solver.solve(model, t_eval) - p = np.linspace(0, 1, n)[:, np.newaxis] + solution = solver.solve(model, t_eval, inputs={"param": 7 * np.ones(n)}, + calculate_sensitivities=True) + l_n = mesh["negative electrode"].edges[-1] np.testing.assert_array_almost_equal( - solution["var"].value({"param": 3 * np.ones(n)}), - np.repeat(2 * np.exp(-3 * t_eval), 40)[:, np.newaxis], - decimal=4, + solution["var"].data, np.tile(2 * np.exp(-7 * t_eval), (n, 1)), decimal=4, ) + np.testing.assert_array_almost_equal( - solution["var"].value({"param": 2 * p}), - 2 * np.exp(-2 * p * t_eval).T.reshape(-1, 1), - decimal=4, + solution["var"].sensitivities["param"], + np.vstack([np.eye(n) * -2 * t * np.exp(-7 * t) for t in t_eval]), ) np.testing.assert_array_almost_equal( - solution["var"].sensitivity({"param": 3 * np.ones(n)}), - np.kron(-2 * t_eval * np.exp(-3 * t_eval), np.eye(40)).T, - decimal=4, + solution["integral of var"].data, 2 * np.exp(-7 * t_eval) * l_n, decimal=4, + ) + np.testing.assert_array_almost_equal( + solution["integral of var"].sensitivities["param"], + np.tile(-2 * t_eval * np.exp(-7 * t_eval) * l_n / n, (n, 1)).T, ) - sens = solution["var"].sensitivity({"param": p}).full() - for idx, t in enumerate(t_eval): - np.testing.assert_array_almost_equal( - sens[40 * idx : 40 * (idx + 1), :], - -2 * t * np.exp(-p * t) * np.eye(40), - decimal=4, - ) - - def test_solve_with_symbolic_input_in_initial_conditions(self): - # Simple system: a single algebraic equation - var = pybamm.Variable("var") - model = pybamm.BaseModel() - model.rhs = {var: -var} - model.initial_conditions = {var: pybamm.InputParameter("param")} - model.variables = {"var": var} - - # create discretisation - disc = pybamm.Discretisation() - disc.process_model(model) - - # Solve - solver = pybamm.CasadiSolver(atol=1e-10, rtol=1e-10) - t_eval = np.linspace(0, 1) - solution = solver.solve(model, t_eval) + # Solve - linspace input + p_eval = np.linspace(1, 2, n) + solution = solver.solve(model, t_eval, inputs={"param": p_eval}, + calculate_sensitivities=True) + l_n = mesh["negative electrode"].edges[-1] + np.testing.assert_array_almost_equal( + solution["var"].data, 2 * np.exp(-p_eval[:, np.newaxis] * t_eval), decimal=4 + ) np.testing.assert_array_almost_equal( - solution["var"].value({"param": 7}), 7 * np.exp(-t_eval)[np.newaxis, :] + solution["var"].sensitivities["param"], + np.vstack([np.diag(-2 * t * np.exp(-p_eval * t)) for t in t_eval]), ) + np.testing.assert_array_almost_equal( - solution["var"].value({"param": 3}), 3 * np.exp(-t_eval)[np.newaxis, :] + solution["integral of var"].data, + np.sum( + 2 + * np.exp(-p_eval[:, np.newaxis] * t_eval) + * mesh["negative electrode"].d_edges[:, np.newaxis], + axis=0, + ), ) np.testing.assert_array_almost_equal( - solution["var"].sensitivity({"param": 3}), np.exp(-t_eval)[:, np.newaxis] + solution["integral of var"].sensitivities["param"], + np.vstack([-2 * t * np.exp(-p_eval * t) * l_n / n for t in t_eval]), ) - def test_least_squares_fit_input_in_initial_conditions(self): - # Simple system: a single algebraic equation - var1 = pybamm.Variable("var1", domain="negative electrode") - var2 = pybamm.Variable("var2", domain="negative electrode") + def test_solve_sensitivity_then_no_sensitivity(self): + # Create model model = pybamm.BaseModel() + var = pybamm.Variable("var") p = pybamm.InputParameter("p") - q = pybamm.InputParameter("q") - model.rhs = {var1: -var1} - model.algebraic = {var2: (var2 - p)} - model.initial_conditions = {var1: 1, var2: p} - model.variables = {"objective": (var2 - q) ** 2 + (p - 3) ** 2} + model.rhs = {var: p * var} + model.initial_conditions = {var: 1} + model.variables = {"var squared": var ** 2} - # create discretisation - disc = get_discretisation_for_testing() - disc.process_model(model) + # Solve + # Make sure that passing in extra options works + solver = pybamm.CasadiSolver( + mode="fast", rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve(model, t_eval, inputs={"p": 0.1}, + calculate_sensitivities=True) + + # check sensitivities + np.testing.assert_allclose( + solution.sensitivities["p"], + (solution.t * np.exp(0.1 * solution.t))[:, np.newaxis], + ) + + solution = solver.solve(model, t_eval, inputs={"p": 0.1}) + + np.testing.assert_array_equal(solution.t, t_eval) + np.testing.assert_allclose(solution.y, np.exp(0.1 * solution.t).reshape(1, -1)) + np.testing.assert_allclose( + solution["var squared"].data, np.exp(0.1 * solution.t) ** 2 + ) + + +class TestCasadiSolverDAEsWithForwardSensitivityEquations(unittest.TestCase): + def test_solve_sensitivity_scalar_var_scalar_input(self): + # Create model + model = pybamm.BaseModel() + var1 = pybamm.Variable("var1") + p = pybamm.InputParameter("p") + var1 = pybamm.Variable("var1") + var2 = pybamm.Variable("var2") + model.rhs = {var1: p * var1} + model.algebraic = {var2: 2 * var1 - var2} + model.initial_conditions = {var1: 1, var2: 2} + model.variables = {"var2 squared": var2 ** 2} # Solve - solver = pybamm.CasadiSolver() - solution = solver.solve(model, np.linspace(0, 1)) - sol_var = solution["objective"] + # Make sure that passing in extra options works + solver = pybamm.CasadiSolver( + mode="fast", rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve(model, t_eval, inputs={"p": 0.1}, + calculate_sensitivities=True) + np.testing.assert_array_equal(solution.t, t_eval) + np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t)) + np.testing.assert_allclose( + solution.sensitivities["p"], + np.stack(( + solution.t * np.exp(0.1 * solution.t), + 2 * solution.t * np.exp(0.1 * solution.t), + )).transpose().reshape(-1, 1), + atol=1e-7 + ) + np.testing.assert_allclose( + solution["var2 squared"].data, 4 * np.exp(2 * 0.1 * solution.t) + ) + np.testing.assert_allclose( + solution["var2 squared"].sensitivities["p"], + (8 * solution.t * np.exp(2 * 0.1 * solution.t))[:, np.newaxis], + atol=1e-7 + ) - def objective(x): - return sol_var.value({"p": x[0], "q": x[1]}).full().flatten() + def test_solve_sensitivity_algebraic(self): + # Create model + model = pybamm.BaseModel() + var = pybamm.Variable("var") + p = pybamm.InputParameter("p") + model.algebraic = {var: var - p * pybamm.t} + model.initial_conditions = {var: 0} + model.variables = {"var squared": var ** 2} - # without jacobian - lsq_sol = least_squares(objective, [2, 2], method="lm") - np.testing.assert_array_almost_equal(lsq_sol.x, [3, 3], decimal=3) + # Solve + # Make sure that passing in extra options works + solver = pybamm.CasadiAlgebraicSolver(tol=1e-10) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve(model, t_eval, inputs={"p": 0.1}, + calculate_sensitivities=True) + np.testing.assert_array_equal(solution.t, t_eval) + np.testing.assert_allclose(solution.y[0], 0.1 * solution.t) + np.testing.assert_allclose( + solution.sensitivities["p"], solution.t.reshape(-1, 1), atol=1e-7 + ) + np.testing.assert_allclose( + solution["var squared"].data, (0.1 * solution.t) ** 2 + ) + np.testing.assert_allclose( + solution["var squared"].sensitivities["p"], + (2 * 0.1 * solution.t ** 2).reshape(-1, 1), + atol=1e-7 + ) if __name__ == "__main__": diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index 773f9b9e30..ca66e0c9c5 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -46,6 +46,76 @@ def test_ida_roberts_klu(self): true_solution = 0.1 * solution.t np.testing.assert_array_almost_equal(solution.y[0, :], true_solution) + def test_ida_roberts_klu_sensitivities(self): + # this test implements a python version of the ida Roberts + # example provided in sundials + # see sundials ida examples pdf + for form in ["python", "casadi", "jax"]: + print(form) + model = pybamm.BaseModel() + model.convert_to_format = form + u = pybamm.Variable("u") + v = pybamm.Variable("v") + a = pybamm.InputParameter("a") + model.rhs = {u: a * v} + model.algebraic = {v: 1 - v} + model.initial_conditions = {u: 0, v: 1} + + disc = pybamm.Discretisation() + disc.process_model(model) + + solver = pybamm.IDAKLUSolver(root_method="lm") + + t_eval = np.linspace(0, 3, 100) + a_value = 0.1 + + # solve first without sensitivities + sol = solver.solve( + model, t_eval, inputs={"a": a_value}, + ) + + # test that y[1] remains constant + np.testing.assert_array_almost_equal( + sol.y[1, :], np.ones(sol.t.shape) + ) + + # test that y[0] = to true solution + true_solution = a_value * sol.t + np.testing.assert_array_almost_equal(sol.y[0, :], true_solution) + + # should be no sensitivities calculated + with self.assertRaises(KeyError): + print(sol.sensitivities["a"]) + + # now solve with sensitivities (this should cause set_up to be run again) + sol = solver.solve( + model, t_eval, inputs={"a": a_value}, + calculate_sensitivities=True + ) + + # test that y[1] remains constant + np.testing.assert_array_almost_equal( + sol.y[1, :], np.ones(sol.t.shape) + ) + + # test that y[0] = to true solution + true_solution = a_value * sol.t + np.testing.assert_array_almost_equal(sol.y[0, :], true_solution) + + # evaluate the sensitivities using idas + dyda_ida = sol.sensitivities["a"] + + # evaluate the sensitivities using finite difference + h = 1e-6 + sol_plus = solver.solve(model, t_eval, inputs={"a": a_value + 0.5 * h}) + sol_neg = solver.solve(model, t_eval, inputs={"a": a_value - 0.5 * h}) + dyda_fd = (sol_plus.y - sol_neg.y) / h + dyda_fd = dyda_fd.transpose().reshape(-1, 1) + + np.testing.assert_array_almost_equal( + dyda_ida, dyda_fd + ) + def test_set_atol(self): model = pybamm.lithium_ion.DFN() geometry = model.default_geometry @@ -60,6 +130,31 @@ def test_set_atol(self): variable_tols = {"Porosity times concentration": 1e-3} solver.set_atol_by_variable(variable_tols, model) + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: -0.1 * u} + model.initial_conditions = {u: 1} + t_eval = np.linspace(0, 3, 100) + + disc = pybamm.Discretisation() + disc.process_model(model) + + # numpy array atol + atol = np.zeros(1) + solver = pybamm.IDAKLUSolver(root_method="lm", atol=atol) + solver.solve(model, t_eval) + + # list atol + atol = [1] + solver = pybamm.IDAKLUSolver(root_method="lm", atol=atol) + solver.solve(model, t_eval) + + # wrong size (should fail) + atol = [1, 2] + solver = pybamm.IDAKLUSolver(root_method="lm", atol=atol) + with self.assertRaisesRegex(pybamm.SolverError, 'Absolute tolerances'): + solver.solve(model, t_eval) + def test_failures(self): # this test implements a python version of the ida Roberts # example provided in sundials @@ -79,6 +174,23 @@ def test_failures(self): with self.assertRaisesRegex(pybamm.SolverError, "KLU requires the Jacobian"): solver.solve(model, t_eval) + model = pybamm.BaseModel() + u = pybamm.Variable("u") + model.rhs = {u: -0.1 * u} + model.initial_conditions = {u: 1} + + disc = pybamm.Discretisation() + disc.process_model(model) + + solver = pybamm.IDAKLUSolver(root_method="lm") + + # will give solver error + t_eval = np.linspace(0, -3, 100) + with self.assertRaisesRegex( + pybamm.SolverError, 't_eval must increase monotonically' + ): + solver.solve(model, t_eval) + def test_dae_solver_algebraic_model(self): model = pybamm.BaseModel() var = pybamm.Variable("var") diff --git a/tests/unit/test_solvers/test_processed_variable.py b/tests/unit/test_solvers/test_processed_variable.py index 35f111e7be..87cb715dfc 100644 --- a/tests/unit/test_solvers/test_processed_variable.py +++ b/tests/unit/test_solvers/test_processed_variable.py @@ -89,6 +89,48 @@ def test_processed_variable_0D(self): ) np.testing.assert_array_equal(processed_var.entries, y_sol[0]) + # check empty sensitivity works + + def test_processed_variable_0D_no_sensitivity(self): + # without space + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + var = t * y + var.mesh = None + t_sol = np.linspace(0, 1) + y_sol = np.array([np.linspace(0, 5)]) + var_casadi = to_casadi(var, y_sol) + processed_var = pybamm.ProcessedVariable( + [var], + [var_casadi], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + warn=False, + ) + + # test no inputs (i.e. no sensitivity) + self.assertDictEqual(processed_var.sensitivities, {}) + + # with parameter + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + a = pybamm.InputParameter('a') + var = t * y * a + var.mesh = None + t_sol = np.linspace(0, 1) + y_sol = np.array([np.linspace(0, 5)]) + inputs = {'a': np.array([1.0])} + var_casadi = to_casadi(var, y_sol, inputs=inputs) + processed_var = pybamm.ProcessedVariable( + [var], + [var_casadi], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), inputs), + warn=False, + ) + + # test no sensitivity raises error + with self.assertRaisesRegex(ValueError, 'Cannot compute sensitivities'): + print(processed_var.sensitivities) + def test_processed_variable_1D(self): t = pybamm.t var = pybamm.Variable("var", domain=["negative electrode", "separator"]) diff --git a/tests/unit/test_solvers/test_scikits_solvers.py b/tests/unit/test_solvers/test_scikits_solvers.py index 599a673023..de4f496ccb 100644 --- a/tests/unit/test_solvers/test_scikits_solvers.py +++ b/tests/unit/test_solvers/test_scikits_solvers.py @@ -99,6 +99,7 @@ class Model: timescale_eval = 1 length_scales = {} convert_to_format = "python" + len_rhs_and_alg = 2 def residuals_eval(self, t, y, ydot, inputs): return np.array( diff --git a/tests/unit/test_solvers/test_scipy_solver.py b/tests/unit/test_solvers/test_scipy_solver.py index 9c586372ef..b525d95eac 100644 --- a/tests/unit/test_solvers/test_scipy_solver.py +++ b/tests/unit/test_solvers/test_scipy_solver.py @@ -1,10 +1,9 @@ -# # Tests for the Scipy Solver class # import pybamm import unittest import numpy as np -from tests import get_mesh_for_testing +from tests import get_mesh_for_testing, get_discretisation_for_testing import warnings import sys from platform import system, version @@ -505,6 +504,301 @@ def test_model_solver_manually_update_initial_conditions(self): ) +class TestScipySolverWithSensitivity(unittest.TestCase): + def test_solve_sensitivity_scalar_var_scalar_input(self): + # Create model + model = pybamm.BaseModel() + var = pybamm.Variable("var") + p = pybamm.InputParameter("p") + model.rhs = {var: p * var} + model.initial_conditions = {var: 1} + model.variables = {"var squared": var ** 2} + + # Solve + # Make sure that passing in extra options works + solver = pybamm.ScipySolver( + rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve(model, t_eval, inputs={"p": 0.1}, + calculate_sensitivities=True) + np.testing.assert_array_equal(solution.t, t_eval) + np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t)) + np.testing.assert_allclose( + solution.sensitivities["p"], + (solution.t * np.exp(0.1 * solution.t))[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var squared"].data, np.exp(0.1 * solution.t) ** 2 + ) + np.testing.assert_allclose( + solution["var squared"].sensitivities["p"], + (2 * np.exp(0.1 * solution.t) * solution.t * np.exp(0.1 * solution.t))[ + :, np.newaxis + ], + ) + + # More complicated model + # Create model + model = pybamm.BaseModel() + var = pybamm.Variable("var") + p = pybamm.InputParameter("p") + q = pybamm.InputParameter("q") + r = pybamm.InputParameter("r") + s = pybamm.InputParameter("s") + model.rhs = {var: p * q} + model.initial_conditions = {var: r} + model.variables = {"var times s": var * s} + + # Solve + # Make sure that passing in extra options works + solver = pybamm.ScipySolver( + rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve( + model, t_eval, inputs={"p": 0.1, "q": 2, "r": -1, "s": 0.5}, + calculate_sensitivities=True + ) + np.testing.assert_allclose(solution.y[0], -1 + 0.2 * solution.t) + np.testing.assert_allclose( + solution.sensitivities["p"], + (2 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution.sensitivities["q"], + (0.1 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose(solution.sensitivities["r"], 1) + np.testing.assert_allclose(solution.sensitivities["s"], 0) + np.testing.assert_allclose( + solution.sensitivities["all"], + np.hstack( + [ + solution.sensitivities["p"], + solution.sensitivities["q"], + solution.sensitivities["r"], + solution.sensitivities["s"], + ] + ), + ) + np.testing.assert_allclose( + solution["var times s"].data, 0.5 * (-1 + 0.2 * solution.t) + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["p"], + 0.5 * (2 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["q"], + 0.5 * (0.1 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose(solution["var times s"].sensitivities["r"], 0.5) + np.testing.assert_allclose( + solution["var times s"].sensitivities["s"], + (-1 + 0.2 * solution.t)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["all"], + np.hstack( + [ + solution["var times s"].sensitivities["p"], + solution["var times s"].sensitivities["q"], + solution["var times s"].sensitivities["r"], + solution["var times s"].sensitivities["s"], + ] + ), + ) + + def test_solve_sensitivity_vector_var_scalar_input(self): + var = pybamm.Variable("var", "negative electrode") + model = pybamm.BaseModel() + # Set length scales to avoid warning + model.length_scales = {"negative electrode": 1} + param = pybamm.InputParameter("param") + model.rhs = {var: -param * var} + model.initial_conditions = {var: 2} + model.variables = {"var": var} + + # create discretisation + disc = get_discretisation_for_testing() + disc.process_model(model) + n = disc.mesh["negative electrode"].npts + + # Solve - scalar input + solver = pybamm.ScipySolver() + t_eval = np.linspace(0, 1) + solution = solver.solve(model, t_eval, inputs={"param": 7}, + calculate_sensitivities=True) + np.testing.assert_array_almost_equal( + solution["var"].data, + np.tile(2 * np.exp(-7 * t_eval), (n, 1)), + decimal=4, + ) + np.testing.assert_array_almost_equal( + solution["var"].sensitivities["param"], + np.repeat(-2 * t_eval * np.exp(-7 * t_eval), n)[:, np.newaxis], + decimal=4, + ) + + # More complicated model + # Create model + model = pybamm.BaseModel() + # Set length scales to avoid warning + model.length_scales = {"negative electrode": 1} + var = pybamm.Variable("var", "negative electrode") + p = pybamm.InputParameter("p") + q = pybamm.InputParameter("q") + r = pybamm.InputParameter("r") + s = pybamm.InputParameter("s") + model.rhs = {var: p * q} + model.initial_conditions = {var: r} + model.variables = {"var times s": var * s} + + # Discretise + disc.process_model(model) + + # Solve + # Make sure that passing in extra options works + solver = pybamm.ScipySolver( + rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1, 80) + solution = solver.solve( + model, t_eval, inputs={"p": 0.1, "q": 2, "r": -1, "s": 0.5}, + calculate_sensitivities=True + ) + np.testing.assert_allclose(solution.y, np.tile(-1 + 0.2 * solution.t, (n, 1))) + np.testing.assert_allclose( + solution.sensitivities["p"], + np.repeat(2 * solution.t, n)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution.sensitivities["q"], + np.repeat(0.1 * solution.t, n)[:, np.newaxis], + ) + np.testing.assert_allclose(solution.sensitivities["r"], 1) + np.testing.assert_allclose(solution.sensitivities["s"], 0) + np.testing.assert_allclose( + solution.sensitivities["all"], + np.hstack( + [ + solution.sensitivities["p"], + solution.sensitivities["q"], + solution.sensitivities["r"], + solution.sensitivities["s"], + ] + ), + ) + np.testing.assert_allclose( + solution["var times s"].data, np.tile(0.5 * (-1 + 0.2 * solution.t), (n, 1)) + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["p"], + np.repeat(0.5 * (2 * solution.t), n)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["q"], + np.repeat(0.5 * (0.1 * solution.t), n)[:, np.newaxis], + ) + np.testing.assert_allclose(solution["var times s"].sensitivities["r"], 0.5) + np.testing.assert_allclose( + solution["var times s"].sensitivities["s"], + np.repeat(-1 + 0.2 * solution.t, n)[:, np.newaxis], + ) + np.testing.assert_allclose( + solution["var times s"].sensitivities["all"], + np.hstack( + [ + solution["var times s"].sensitivities["p"], + solution["var times s"].sensitivities["q"], + solution["var times s"].sensitivities["r"], + solution["var times s"].sensitivities["s"], + ] + ), + ) + + def test_solve_sensitivity_vector_var_vector_input(self): + var = pybamm.Variable("var", "negative electrode") + model = pybamm.BaseModel() + # Set length scales to avoid warning + model.length_scales = {"negative electrode": 1} + + param = pybamm.InputParameter("param", "negative electrode") + model.rhs = {var: -param * var} + model.initial_conditions = {var: 2} + model.variables = { + "var": var, + "integral of var": pybamm.Integral(var, pybamm.standard_spatial_vars.x_n), + } + + # create discretisation + mesh = get_mesh_for_testing() + spatial_methods = {"macroscale": pybamm.FiniteVolume()} + disc = pybamm.Discretisation(mesh, spatial_methods) + disc.process_model(model) + n = disc.mesh["negative electrode"].npts + + # Solve - constant input + solver = pybamm.ScipySolver( + rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1) + solution = solver.solve(model, t_eval, inputs={"param": 7 * np.ones(n)}, + calculate_sensitivities=True) + l_n = mesh["negative electrode"].edges[-1] + np.testing.assert_array_almost_equal( + solution["var"].data, + np.tile(2 * np.exp(-7 * t_eval), (n, 1)), + decimal=4, + ) + + np.testing.assert_array_almost_equal( + solution["var"].sensitivities["param"], + np.vstack([np.eye(n) * -2 * t * np.exp(-7 * t) for t in t_eval]), + ) + np.testing.assert_array_almost_equal( + solution["integral of var"].data, + 2 * np.exp(-7 * t_eval) * l_n, + decimal=4, + ) + np.testing.assert_array_almost_equal( + solution["integral of var"].sensitivities["param"], + np.tile(-2 * t_eval * np.exp(-7 * t_eval) * l_n / n, (n, 1)).T, + ) + + # Solve - linspace input + solver = pybamm.ScipySolver( + rtol=1e-10, atol=1e-10 + ) + t_eval = np.linspace(0, 1) + p_eval = np.linspace(1, 2, n) + solution = solver.solve(model, t_eval, inputs={"param": p_eval}, + calculate_sensitivities=True) + l_n = mesh["negative electrode"].edges[-1] + np.testing.assert_array_almost_equal( + solution["var"].data, 2 * np.exp(-p_eval[:, np.newaxis] * t_eval), decimal=4 + ) + np.testing.assert_array_almost_equal( + solution["var"].sensitivities["param"], + np.vstack([np.diag(-2 * t * np.exp(-p_eval * t)) for t in t_eval]), + ) + + np.testing.assert_array_almost_equal( + solution["integral of var"].data, + np.sum( + 2 + * np.exp(-p_eval[:, np.newaxis] * t_eval) + * mesh["negative electrode"].d_edges[:, np.newaxis], + axis=0, + ), + ) + np.testing.assert_array_almost_equal( + solution["integral of var"].sensitivities["param"], + np.vstack([-2 * t * np.exp(-p_eval * t) * l_n / n for t in t_eval]), + ) + + if __name__ == "__main__": print("Add -v for more debug output") diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py index 40ba22a24d..3edf888f2f 100644 --- a/tests/unit/test_solvers/test_solution.py +++ b/tests/unit/test_solvers/test_solution.py @@ -22,6 +22,12 @@ def test_init(self): self.assertEqual(sol.all_inputs, [{}]) self.assertIsInstance(sol.all_models[0], pybamm.BaseModel) + def test_sensitivities(self): + t = np.linspace(0, 1) + y = np.tile(t, (20, 1)) + with self.assertRaises(TypeError): + pybamm.Solution(t, y, pybamm.BaseModel(), {}, sensitivities=1.0) + def test_errors(self): bad_ts = [np.array([1, 2, 3]), np.array([3, 4, 5])] sol = pybamm.Solution( @@ -46,10 +52,10 @@ def test_add_solutions(self): sol2 = pybamm.Solution(t2, y2, pybamm.BaseModel(), {"a": 2}) sol2.solve_time = 1 sol2.integration_time = 0.5 + sol_sum = sol1 + sol2 # Test - self.assertEqual(sol_sum.solve_time, 2.5) self.assertEqual(sol_sum.integration_time, 0.8) np.testing.assert_array_equal(sol_sum.t, np.concatenate([t1, t2[1:]])) np.testing.assert_array_equal(