diff --git a/CHANGELOG.md b/CHANGELOG.md index f70aae0272..fd590fbc12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Features +- Added Hermite interpolation to the (`IDAKLUSolver`) that improves the accuracy and performance of post-processing variables. ([#4464](https://github.com/pybamm-team/PyBaMM/pull/4464)) - Added sensitivity calculation support for `pybamm.Simulation` and `pybamm.Experiment` ([#4415](https://github.com/pybamm-team/PyBaMM/pull/4415)) - Added OpenMP parallelization to IDAKLU solver for lists of input parameters ([#4449](https://github.com/pybamm-team/PyBaMM/pull/4449)) - Added phase-dependent particle options to LAM diff --git a/CMakeLists.txt b/CMakeLists.txt index a7f68ce7a0..ec594e5ca5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -105,6 +105,8 @@ pybind11_add_module(idaklu src/pybamm/solvers/c_solvers/idaklu/Expressions/Base/Expression.hpp src/pybamm/solvers/c_solvers/idaklu/Expressions/Base/ExpressionSet.hpp src/pybamm/solvers/c_solvers/idaklu/Expressions/Base/ExpressionTypes.hpp + src/pybamm/solvers/c_solvers/idaklu/observe.hpp + src/pybamm/solvers/c_solvers/idaklu/observe.cpp # IDAKLU expressions - concrete implementations ${IDAKLU_EXPR_CASADI_SOURCE_FILES} ${IDAKLU_EXPR_IREE_SOURCE_FILES} diff --git a/setup.py b/setup.py index 74de1baca4..8a49bfd715 100644 --- a/setup.py +++ b/setup.py @@ -327,6 +327,8 @@ def compile_KLU(): "src/pybamm/solvers/c_solvers/idaklu/Solution.hpp", "src/pybamm/solvers/c_solvers/idaklu/Options.hpp", "src/pybamm/solvers/c_solvers/idaklu/Options.cpp", + "src/pybamm/solvers/c_solvers/idaklu/observe.hpp", + "src/pybamm/solvers/c_solvers/idaklu/observe.cpp", "src/pybamm/solvers/c_solvers/idaklu.cpp", ], ) diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index 36ad0b137a..51c7f49969 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -157,7 +157,7 @@ # Solver classes from .solvers.solution import Solution, EmptySolution, make_cycle_solution -from .solvers.processed_variable import ProcessedVariable +from .solvers.processed_variable import ProcessedVariable, process_variable from .solvers.processed_variable_computed import ProcessedVariableComputed from .solvers.base_solver import BaseSolver from .solvers.dummy_solver import DummySolver diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py index 39dc974f9b..cddce58d77 100644 --- a/src/pybamm/plotting/quick_plot.py +++ b/src/pybamm/plotting/quick_plot.py @@ -419,14 +419,14 @@ def reset_axis(self): spatial_vars = self.spatial_variable_dict[key] var_min = np.min( [ - ax_min(var(self.ts_seconds[i], **spatial_vars, warn=False)) + ax_min(var(self.ts_seconds[i], **spatial_vars)) for i, variable_list in enumerate(variable_lists) for var in variable_list ] ) var_max = np.max( [ - ax_max(var(self.ts_seconds[i], **spatial_vars, warn=False)) + ax_max(var(self.ts_seconds[i], **spatial_vars)) for i, variable_list in enumerate(variable_lists) for var in variable_list ] @@ -512,7 +512,7 @@ def plot(self, t, dynamic=False): full_t = self.ts_seconds[i] (self.plots[key][i][j],) = ax.plot( full_t / self.time_scaling_factor, - variable(full_t, warn=False), + variable(full_t), color=self.colors[i], linestyle=linestyle, ) @@ -548,7 +548,7 @@ def plot(self, t, dynamic=False): linestyle = self.linestyles[j] (self.plots[key][i][j],) = ax.plot( self.first_spatial_variable[key], - variable(t_in_seconds, **spatial_vars, warn=False), + variable(t_in_seconds, **spatial_vars), color=self.colors[i], linestyle=linestyle, zorder=10, @@ -570,13 +570,13 @@ def plot(self, t, dynamic=False): y_name = next(iter(spatial_vars.keys()))[0] x = self.second_spatial_variable[key] y = self.first_spatial_variable[key] - var = variable(t_in_seconds, **spatial_vars, warn=False) + var = variable(t_in_seconds, **spatial_vars) else: x_name = next(iter(spatial_vars.keys()))[0] y_name = list(spatial_vars.keys())[1][0] x = self.first_spatial_variable[key] y = self.second_spatial_variable[key] - var = variable(t_in_seconds, **spatial_vars, warn=False).T + var = variable(t_in_seconds, **spatial_vars).T ax.set_xlabel(f"{x_name} [{self.spatial_unit}]") ax.set_ylabel(f"{y_name} [{self.spatial_unit}]") vmin, vmax = self.variable_limits[key] @@ -710,7 +710,6 @@ def slider_update(self, t): var = variable( time_in_seconds, **self.spatial_variable_dict[key], - warn=False, ) plot[i][j].set_ydata(var) var_min = min(var_min, ax_min(var)) @@ -729,11 +728,11 @@ def slider_update(self, t): if self.x_first_and_y_second[key] is False: x = self.second_spatial_variable[key] y = self.first_spatial_variable[key] - var = variable(time_in_seconds, **spatial_vars, warn=False) + var = variable(time_in_seconds, **spatial_vars) else: x = self.first_spatial_variable[key] y = self.second_spatial_variable[key] - var = variable(time_in_seconds, **spatial_vars, warn=False).T + var = variable(time_in_seconds, **spatial_vars).T # store the plot and the var data (for testing) as cant access # z data from QuadMesh or QuadContourSet object if self.is_y_z[key] is True: diff --git a/src/pybamm/solvers/c_solvers/idaklu.cpp b/src/pybamm/solvers/c_solvers/idaklu.cpp index db7147feb2..82a3cbe91c 100644 --- a/src/pybamm/solvers/c_solvers/idaklu.cpp +++ b/src/pybamm/solvers/c_solvers/idaklu.cpp @@ -9,6 +9,7 @@ #include #include "idaklu/idaklu_solver.hpp" +#include "idaklu/observe.hpp" #include "idaklu/IDAKLUSolverGroup.hpp" #include "idaklu/IdakluJax.hpp" #include "idaklu/common.hpp" @@ -27,6 +28,7 @@ casadi::Function generate_casadi_function(const std::string &data) namespace py = pybind11; PYBIND11_MAKE_OPAQUE(std::vector); +PYBIND11_MAKE_OPAQUE(std::vector); PYBIND11_MAKE_OPAQUE(std::vector); PYBIND11_MODULE(idaklu, m) @@ -34,6 +36,7 @@ PYBIND11_MODULE(idaklu, m) m.doc() = "sundials solvers"; // optional module docstring py::bind_vector>(m, "VectorNdArray"); + py::bind_vector>(m, "VectorRealtypeNdArray"); py::bind_vector>(m, "VectorSolution"); py::class_(m, "IDAKLUSolverGroup") @@ -72,6 +75,27 @@ PYBIND11_MODULE(idaklu, m) py::arg("options"), py::return_value_policy::take_ownership); + m.def("observe", &observe, + "Observe variables", + py::arg("ts"), + py::arg("ys"), + py::arg("inputs"), + py::arg("funcs"), + py::arg("is_f_contiguous"), + py::arg("shape"), + py::return_value_policy::take_ownership); + + m.def("observe_hermite_interp", &observe_hermite_interp, + "Observe and Hermite interpolate variables", + py::arg("t_interp"), + py::arg("ts"), + py::arg("ys"), + py::arg("yps"), + py::arg("inputs"), + py::arg("funcs"), + py::arg("shape"), + py::return_value_policy::take_ownership); + #ifdef IREE_ENABLE m.def("create_iree_solver_group", &create_idaklu_solver_group, "Create a group of iree idaklu solver objects", @@ -167,7 +191,9 @@ PYBIND11_MODULE(idaklu, m) py::class_(m, "solution") .def_readwrite("t", &Solution::t) .def_readwrite("y", &Solution::y) + .def_readwrite("yp", &Solution::yp) .def_readwrite("yS", &Solution::yS) + .def_readwrite("ypS", &Solution::ypS) .def_readwrite("y_term", &Solution::y_term) .def_readwrite("flag", &Solution::flag); } diff --git a/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp b/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp index 92eede3643..ee2c03abff 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp +++ b/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp @@ -54,9 +54,9 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver int const number_of_events; // cppcheck-suppress unusedStructMember int number_of_timesteps; int precon_type; // cppcheck-suppress unusedStructMember - N_Vector yy, yp, y_cache, avtol; // y, y', y cache vector, and absolute tolerance + N_Vector yy, yyp, y_cache, avtol; // y, y', y cache vector, and absolute tolerance N_Vector *yyS; // cppcheck-suppress unusedStructMember - N_Vector *ypS; // cppcheck-suppress unusedStructMember + N_Vector *yypS; // cppcheck-suppress unusedStructMember N_Vector id; // rhs_alg_id realtype rtol; int const jac_times_cjmass_nnz; // cppcheck-suppress unusedStructMember @@ -70,11 +70,14 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver vector res_dvar_dp; bool const sensitivity; // cppcheck-suppress unusedStructMember bool const save_outputs_only; // cppcheck-suppress unusedStructMember + bool save_hermite; // cppcheck-suppress unusedStructMember bool is_ODE; // cppcheck-suppress unusedStructMember int length_of_return_vector; // cppcheck-suppress unusedStructMember vector t; // cppcheck-suppress unusedStructMember vector> y; // cppcheck-suppress unusedStructMember + vector> yp; // cppcheck-suppress unusedStructMember vector>> yS; // cppcheck-suppress unusedStructMember + vector>> ypS; // cppcheck-suppress unusedStructMember SetupOptions const setup_opts; SolverOptions const solver_opts; @@ -144,6 +147,11 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver */ void InitializeStorage(int const N); + /** + * @brief Initialize the storage for Hermite interpolation + */ + void InitializeHermiteStorage(int const N); + /** * @brief Apply user-configurable IDA options */ @@ -190,13 +198,20 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver */ void ExtendAdaptiveArrays(); + /** + * @brief Extend the Hermite interpolation info by 1 + */ + void ExtendHermiteArrays(); + /** * @brief Set the step values */ void SetStep( - realtype &t_val, + realtype &tval, realtype *y_val, + realtype *yp_val, vector const &yS_val, + vector const &ypS_val, int &i_save ); @@ -211,7 +226,9 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver realtype &t_prev, realtype const &t_next, realtype *y_val, + realtype *yp_val, vector const &yS_val, + vector const &ypS_val, int &i_save ); @@ -255,6 +272,26 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver int &i_save ); + /** + * @brief Save the output function results at the requested time + */ + void SetStepHermite( + realtype &t_val, + realtype *yp_val, + const vector &ypS_val, + int &i_save + ); + + /** + * @brief Save the output function sensitivities at the requested time + */ + void SetStepHermiteSensitivities( + realtype &t_val, + realtype *yp_val, + const vector &ypS_val, + int &i_save + ); + }; #include "IDAKLUSolverOpenMP.inl" diff --git a/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl b/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl index 56f546facf..d128ae1809 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl +++ b/src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl @@ -1,7 +1,6 @@ #include "Expressions/Expressions.hpp" #include "sundials_functions.hpp" #include - #include "common.hpp" #include "SolutionData.hpp" @@ -48,7 +47,7 @@ IDAKLUSolverOpenMP::IDAKLUSolverOpenMP( AllocateVectors(); if (sensitivity) { yyS = N_VCloneVectorArray(number_of_parameters, yy); - ypS = N_VCloneVectorArray(number_of_parameters, yp); + yypS = N_VCloneVectorArray(number_of_parameters, yyp); } // set initial values realtype *atval = N_VGetArrayPointer(avtol); @@ -58,14 +57,14 @@ IDAKLUSolverOpenMP::IDAKLUSolverOpenMP( for (int is = 0; is < number_of_parameters; is++) { N_VConst(RCONST(0.0), yyS[is]); - N_VConst(RCONST(0.0), ypS[is]); + N_VConst(RCONST(0.0), yypS[is]); } // create Matrix objects SetMatrix(); // initialise solver - IDAInit(ida_mem, residual_eval, 0, yy, yp); + IDAInit(ida_mem, residual_eval, 0, yy, yyp); // set tolerances rtol = RCONST(rel_tol); @@ -87,6 +86,9 @@ IDAKLUSolverOpenMP::IDAKLUSolverOpenMP( // The default is to solve a DAE for generality. This may be changed // to an ODE during the Initialize() call is_ODE = false; + + // Will be overwritten during the solve() call + save_hermite = solver_opts.hermite_interpolation; } template @@ -95,14 +97,14 @@ void IDAKLUSolverOpenMP::AllocateVectors() { // Create vectors if (setup_opts.num_threads == 1) { yy = N_VNew_Serial(number_of_states, sunctx); - yp = N_VNew_Serial(number_of_states, sunctx); + yyp = N_VNew_Serial(number_of_states, sunctx); y_cache = N_VNew_Serial(number_of_states, sunctx); avtol = N_VNew_Serial(number_of_states, sunctx); id = N_VNew_Serial(number_of_states, sunctx); } else { DEBUG("IDAKLUSolverOpenMP::AllocateVectors OpenMP"); yy = N_VNew_OpenMP(number_of_states, setup_opts.num_threads, sunctx); - yp = N_VNew_OpenMP(number_of_states, setup_opts.num_threads, sunctx); + yyp = N_VNew_OpenMP(number_of_states, setup_opts.num_threads, sunctx); y_cache = N_VNew_OpenMP(number_of_states, setup_opts.num_threads, sunctx); avtol = N_VNew_OpenMP(number_of_states, setup_opts.num_threads, sunctx); id = N_VNew_OpenMP(number_of_states, setup_opts.num_threads, sunctx); @@ -127,6 +129,26 @@ void IDAKLUSolverOpenMP::InitializeStorage(int const N) { vector(length_of_return_vector, 0.0) ) ); + + if (save_hermite) { + InitializeHermiteStorage(N); + } +} + +template +void IDAKLUSolverOpenMP::InitializeHermiteStorage(int const N) { + yp = vector>( + N, + vector(number_of_states, 0.0) + ); + + ypS = vector>>( + N, + vector>( + number_of_parameters, + vector(number_of_states, 0.0) + ) + ); } template @@ -285,7 +307,7 @@ void IDAKLUSolverOpenMP::Initialize() { if (sensitivity) { CheckErrors(IDASensInit(ida_mem, number_of_parameters, IDA_SIMULTANEOUS, - sensitivities_eval, yyS, ypS)); + sensitivities_eval, yyS, yypS)); CheckErrors(IDASensEEtolerances(ida_mem)); } @@ -321,13 +343,13 @@ IDAKLUSolverOpenMP::~IDAKLUSolverOpenMP() { SUNMatDestroy(J); N_VDestroy(avtol); N_VDestroy(yy); - N_VDestroy(yp); + N_VDestroy(yyp); N_VDestroy(y_cache); N_VDestroy(id); if (sensitivity) { N_VDestroyVectorArray(yyS, number_of_parameters); - N_VDestroyVectorArray(ypS, number_of_parameters); + N_VDestroyVectorArray(yypS, number_of_parameters); } IDAFree(&ida_mem); @@ -349,9 +371,15 @@ SolutionData IDAKLUSolverOpenMP::solve( const int number_of_evals = t_eval.size(); const int number_of_interps = t_interp.size(); - if (t.size() < number_of_evals + number_of_interps) { - InitializeStorage(number_of_evals + number_of_interps); - } + // Hermite interpolation is only available when saving + // 1. adaptive steps and 2. the full solution + save_hermite = ( + solver_opts.hermite_interpolation && + save_adaptive_steps && + !save_outputs_only + ); + + InitializeStorage(number_of_evals + number_of_interps); int i_save = 0; @@ -378,12 +406,12 @@ SolutionData IDAKLUSolverOpenMP::solve( // Setup consistent initialization realtype *y_val = N_VGetArrayPointer(yy); - realtype *yp_val = N_VGetArrayPointer(yp); + realtype *yp_val = N_VGetArrayPointer(yyp); vector yS_val(number_of_parameters); vector ypS_val(number_of_parameters); for (int p = 0 ; p < number_of_parameters; p++) { yS_val[p] = N_VGetArrayPointer(yyS[p]); - ypS_val[p] = N_VGetArrayPointer(ypS[p]); + ypS_val[p] = N_VGetArrayPointer(yypS[p]); for (int i = 0; i < number_of_states; i++) { yS_val[p][i] = y0[i + (p + 1) * number_of_states]; ypS_val[p][i] = yp0[i + (p + 1) * number_of_states]; @@ -409,23 +437,27 @@ SolutionData IDAKLUSolverOpenMP::solve( ConsistentInitialization(t0, t_eval_next, init_type); } + // Set the initial stop time + IDASetStopTime(ida_mem, t_eval_next); + + // Progress one step. This must be done before the while loop to ensure + // that we can run IDAGetDky at t0 for dky = 1 + int retval = IDASolve(ida_mem, tf, &t_val, yy, yyp, IDA_ONE_STEP); + + // Store consistent initialization + CheckErrors(IDAGetDky(ida_mem, t0, 0, yy)); if (sensitivity) { - CheckErrors(IDAGetSensDky(ida_mem, t_val, 0, yyS)); + CheckErrors(IDAGetSensDky(ida_mem, t0, 0, yyS)); } - // Store Consistent initialization - SetStep(t0, y_val, yS_val, i_save); + SetStep(t0, y_val, yp_val, yS_val, ypS_val, i_save); - // Set the initial stop time - IDASetStopTime(ida_mem, t_eval_next); + // Reset the states at t = t_val. Sensitivities are handled in the while-loop + CheckErrors(IDAGetDky(ida_mem, t_val, 0, yy)); // Solve the system - int retval; DEBUG("IDASolve"); while (true) { - // Progress one step - retval = IDASolve(ida_mem, tf, &t_val, yy, yp, IDA_ONE_STEP); - if (retval < 0) { // failed break; @@ -448,18 +480,21 @@ SolutionData IDAKLUSolverOpenMP::solve( if (hit_tinterp) { // Save the interpolated state at t_prev < t < t_val, for all t in t_interp - SetStepInterp(i_interp, + SetStepInterp( + i_interp, t_interp_next, t_interp, t_val, t_prev, t_eval_next, y_val, + yp_val, yS_val, + ypS_val, i_save); } - if (hit_adaptive || hit_teval || hit_event) { + if (hit_adaptive || hit_teval || hit_event || hit_final_time) { if (hit_tinterp) { // Reset the states and sensitivities at t = t_val CheckErrors(IDAGetDky(ida_mem, t_val, 0, yy)); @@ -469,11 +504,16 @@ SolutionData IDAKLUSolverOpenMP::solve( } // Save the current state at t_val - if (hit_adaptive) { - // Dynamically allocate memory for the adaptive step - ExtendAdaptiveArrays(); + // First, check to make sure that the t_val is not equal to the current t value + // If it is, we don't want to save the current state twice + if (!hit_tinterp || t_val != t.back()) { + if (hit_adaptive) { + // Dynamically allocate memory for the adaptive step + ExtendAdaptiveArrays(); + } + + SetStep(t_val, y_val, yp_val, yS_val, ypS_val, i_save); } - SetStep(t_val, y_val, yS_val, i_save); } if (hit_final_time || hit_event) { @@ -481,7 +521,7 @@ SolutionData IDAKLUSolverOpenMP::solve( break; } else if (hit_teval) { // Set the next stop time - i_eval += 1; + i_eval++; t_eval_next = t_eval[i_eval]; CheckErrors(IDASetStopTime(ida_mem, t_eval_next)); @@ -491,6 +531,9 @@ SolutionData IDAKLUSolverOpenMP::solve( } t_prev = t_val; + + // Progress one step + retval = IDASolve(ida_mem, tf, &t_val, yy, yyp, IDA_ONE_STEP); } int const length_of_final_sv_slice = save_outputs_only ? number_of_states : 0; @@ -547,7 +590,50 @@ SolutionData IDAKLUSolverOpenMP::solve( } } - return SolutionData(retval, number_of_timesteps, length_of_return_vector, arg_sens0, arg_sens1, arg_sens2, length_of_final_sv_slice, t_return, y_return, yS_return, yterm_return); + realtype *yp_return = new realtype[(save_hermite ? 1 : 0) * (number_of_timesteps * number_of_states)]; + realtype *ypS_return = new realtype[(save_hermite ? 1 : 0) * (arg_sens0 * arg_sens1 * arg_sens2)]; + if (save_hermite) { + count = 0; + for (size_t i = 0; i < number_of_timesteps; i++) { + for (size_t j = 0; j < number_of_states; j++) { + yp_return[count] = yp[i][j]; + count++; + } + } + + // Sensitivity states, ypS + // Note: Ordering of vector is different if computing outputs vs returning + // the complete state vector + count = 0; + for (size_t idx0 = 0; idx0 < arg_sens0; idx0++) { + for (size_t idx1 = 0; idx1 < arg_sens1; idx1++) { + for (size_t idx2 = 0; idx2 < arg_sens2; idx2++) { + auto i = (save_outputs_only ? idx0 : idx1); + auto j = (save_outputs_only ? idx1 : idx2); + auto k = (save_outputs_only ? idx2 : idx0); + + ypS_return[count] = ypS[i][k][j]; + count++; + } + } + } + } + + return SolutionData( + retval, + number_of_timesteps, + length_of_return_vector, + arg_sens0, + arg_sens1, + arg_sens2, + length_of_final_sv_slice, + save_hermite, + t_return, + y_return, + yp_return, + yS_return, + ypS_return, + yterm_return); } template @@ -563,14 +649,30 @@ void IDAKLUSolverOpenMP::ExtendAdaptiveArrays() { if (sensitivity) { yS.emplace_back(number_of_parameters, vector(length_of_return_vector, 0.0)); } + + if (save_hermite) { + ExtendHermiteArrays(); + } +} + +template +void IDAKLUSolverOpenMP::ExtendHermiteArrays() { + DEBUG("IDAKLUSolver::ExtendHermiteArrays"); + // States + yp.emplace_back(number_of_states, 0.0); + + // Sensitivity + if (sensitivity) { + ypS.emplace_back(number_of_parameters, vector(number_of_states, 0.0)); + } } template void IDAKLUSolverOpenMP::ReinitializeIntegrator(const realtype& t_val) { DEBUG("IDAKLUSolver::ReinitializeIntegrator"); - CheckErrors(IDAReInit(ida_mem, t_val, yy, yp)); + CheckErrors(IDAReInit(ida_mem, t_val, yy, yyp)); if (sensitivity) { - CheckErrors(IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS)); + CheckErrors(IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, yypS)); } } @@ -609,14 +711,16 @@ void IDAKLUSolverOpenMP::ConsistentInitializationODE( realtype *y_cache_val = N_VGetArrayPointer(y_cache); std::memset(y_cache_val, 0, number_of_states * sizeof(realtype)); // Overwrite yp - residual_eval(t_val, yy, y_cache, yp, functions.get()); + residual_eval(t_val, yy, y_cache, yyp, functions.get()); } template void IDAKLUSolverOpenMP::SetStep( realtype &tval, realtype *y_val, + realtype *yp_val, vector const &yS_val, + vector const &ypS_val, int &i_save ) { // Set adaptive step results for y and yS @@ -629,6 +733,10 @@ void IDAKLUSolverOpenMP::SetStep( SetStepOutput(tval, y_val, yS_val, i_save); } else { SetStepFull(tval, y_val, yS_val, i_save); + + if (save_hermite) { + SetStepHermite(tval, yp_val, ypS_val, i_save); + } } i_save++; @@ -644,7 +752,9 @@ void IDAKLUSolverOpenMP::SetStepInterp( realtype &t_prev, realtype const &t_eval_next, realtype *y_val, + realtype *yp_val, vector const &yS_val, + vector const &ypS_val, int &i_save ) { // Save the state at the requested time @@ -657,7 +767,7 @@ void IDAKLUSolverOpenMP::SetStepInterp( } // Memory is already allocated for the interpolated values - SetStep(t_interp_next, y_val, yS_val, i_save); + SetStep(t_interp_next, y_val, yp_val, yS_val, ypS_val, i_save); i_interp++; if (i_interp == (t_interp.size())) { @@ -769,6 +879,50 @@ void IDAKLUSolverOpenMP::SetStepOutputSensitivities( } } +template +void IDAKLUSolverOpenMP::SetStepHermite( + realtype &tval, + realtype *yp_val, + vector const &ypS_val, + int &i_save +) { + // Set adaptive step results for yp and ypS + DEBUG("IDAKLUSolver::SetStepHermite"); + + // States + CheckErrors(IDAGetDky(ida_mem, tval, 1, yyp)); + auto &yp_back = yp[i_save]; + for (size_t j = 0; j < length_of_return_vector; ++j) { + yp_back[j] = yp_val[j]; + + } + + // Sensitivity + if (sensitivity) { + SetStepHermiteSensitivities(tval, yp_val, ypS_val, i_save); + } +} + +template +void IDAKLUSolverOpenMP::SetStepHermiteSensitivities( + realtype &tval, + realtype *yp_val, + vector const &ypS_val, + int &i_save +) { + DEBUG("IDAKLUSolver::SetStepHermiteSensitivities"); + + // Calculate sensitivities for the full ypS array + CheckErrors(IDAGetSensDky(ida_mem, tval, 1, yypS)); + for (size_t j = 0; j < number_of_parameters; ++j) { + auto &ypS_back_j = ypS[i_save][j]; + auto &ypSval_j = ypS_val[j]; + for (size_t k = 0; k < number_of_states; ++k) { + ypS_back_j[k] = ypSval_j[k]; + } + } +} + template void IDAKLUSolverOpenMP::CheckErrors(int const & flag) { if (flag < 0) { diff --git a/src/pybamm/solvers/c_solvers/idaklu/Options.cpp b/src/pybamm/solvers/c_solvers/idaklu/Options.cpp index 51544040ee..8eb605fe77 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/Options.cpp +++ b/src/pybamm/solvers/c_solvers/idaklu/Options.cpp @@ -149,6 +149,7 @@ SolverOptions::SolverOptions(py::dict &py_opts) nonlinear_convergence_coefficient(RCONST(py_opts["nonlinear_convergence_coefficient"].cast())), nonlinear_convergence_coefficient_ic(RCONST(py_opts["nonlinear_convergence_coefficient_ic"].cast())), suppress_algebraic_error(py_opts["suppress_algebraic_error"].cast()), + hermite_interpolation(py_opts["hermite_interpolation"].cast()), // IDA initial conditions calculation calc_ic(py_opts["calc_ic"].cast()), init_all_y_ic(py_opts["init_all_y_ic"].cast()), diff --git a/src/pybamm/solvers/c_solvers/idaklu/Options.hpp b/src/pybamm/solvers/c_solvers/idaklu/Options.hpp index d0c0c1d766..7418c68ec3 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/Options.hpp +++ b/src/pybamm/solvers/c_solvers/idaklu/Options.hpp @@ -38,6 +38,7 @@ struct SolverOptions { double nonlinear_convergence_coefficient; double nonlinear_convergence_coefficient_ic; sunbooleantype suppress_algebraic_error; + bool hermite_interpolation; // IDA initial conditions calculation bool calc_ic; bool init_all_y_ic; diff --git a/src/pybamm/solvers/c_solvers/idaklu/Solution.hpp b/src/pybamm/solvers/c_solvers/idaklu/Solution.hpp index a43e6a7174..8227bb9da8 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/Solution.hpp +++ b/src/pybamm/solvers/c_solvers/idaklu/Solution.hpp @@ -17,8 +17,8 @@ class Solution /** * @brief Constructor */ - Solution(int &retval, np_array &t_np, np_array &y_np, np_array &yS_np, np_array &y_term_np) - : flag(retval), t(t_np), y(y_np), yS(yS_np), y_term(y_term_np) + Solution(int &retval, np_array &t_np, np_array &y_np, np_array &yp_np, np_array &yS_np, np_array &ypS_np, np_array &y_term_np) + : flag(retval), t(t_np), y(y_np), yp(yp_np), yS(yS_np), ypS(ypS_np), y_term(y_term_np) { } @@ -30,7 +30,9 @@ class Solution int flag; np_array t; np_array y; + np_array yp; np_array yS; + np_array ypS; np_array y_term; }; diff --git a/src/pybamm/solvers/c_solvers/idaklu/SolutionData.cpp b/src/pybamm/solvers/c_solvers/idaklu/SolutionData.cpp index 00c2ddbccc..bc48c646d3 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/SolutionData.cpp +++ b/src/pybamm/solvers/c_solvers/idaklu/SolutionData.cpp @@ -29,6 +29,20 @@ Solution SolutionData::generate_solution() { free_y_when_done ); + py::capsule free_yp_when_done( + yp_return, + [](void *f) { + realtype *vect = reinterpret_cast(f); + delete[] vect; + } + ); + + np_array yp_ret = np_array( + (save_hermite ? 1 : 0) * number_of_timesteps * length_of_return_vector, + &yp_return[0], + free_yp_when_done + ); + py::capsule free_yS_when_done( yS_return, [](void *f) { @@ -47,6 +61,24 @@ Solution SolutionData::generate_solution() { free_yS_when_done ); + py::capsule free_ypS_when_done( + ypS_return, + [](void *f) { + realtype *vect = reinterpret_cast(f); + delete[] vect; + } + ); + + np_array ypS_ret = np_array( + std::vector { + (save_hermite ? 1 : 0) * arg_sens0, + arg_sens1, + arg_sens2 + }, + &ypS_return[0], + free_ypS_when_done + ); + // Final state slice, yterm py::capsule free_yterm_when_done( yterm_return, @@ -63,5 +95,5 @@ Solution SolutionData::generate_solution() { ); // Store the solution - return Solution(flag, t_ret, y_ret, yS_ret, y_term); + return Solution(flag, t_ret, y_ret, yp_ret, yS_ret, ypS_ret, y_term); } diff --git a/src/pybamm/solvers/c_solvers/idaklu/SolutionData.hpp b/src/pybamm/solvers/c_solvers/idaklu/SolutionData.hpp index 815e41daca..81ca7f5221 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/SolutionData.hpp +++ b/src/pybamm/solvers/c_solvers/idaklu/SolutionData.hpp @@ -27,9 +27,12 @@ class SolutionData int arg_sens1, int arg_sens2, int length_of_final_sv_slice, + bool save_hermite, realtype *t_return, realtype *y_return, + realtype *yp_return, realtype *yS_return, + realtype *ypS_return, realtype *yterm_return): flag(flag), number_of_timesteps(number_of_timesteps), @@ -38,9 +41,12 @@ class SolutionData arg_sens1(arg_sens1), arg_sens2(arg_sens2), length_of_final_sv_slice(length_of_final_sv_slice), + save_hermite(save_hermite), t_return(t_return), y_return(y_return), + yp_return(yp_return), yS_return(yS_return), + ypS_return(ypS_return), yterm_return(yterm_return) {} @@ -64,9 +70,12 @@ class SolutionData int arg_sens1; int arg_sens2; int length_of_final_sv_slice; + bool save_hermite; realtype *t_return; realtype *y_return; + realtype *yp_return; realtype *yS_return; + realtype *ypS_return; realtype *yterm_return; }; diff --git a/src/pybamm/solvers/c_solvers/idaklu/common.hpp b/src/pybamm/solvers/c_solvers/idaklu/common.hpp index 58be90932e..90672080b6 100644 --- a/src/pybamm/solvers/c_solvers/idaklu/common.hpp +++ b/src/pybamm/solvers/c_solvers/idaklu/common.hpp @@ -33,6 +33,7 @@ namespace py = pybind11; // note: we rely on c_style ordering for numpy arrays so don't change this! using np_array = py::array_t; +using np_array_realtype = py::array_t; using np_array_int = py::array_t; /** diff --git a/src/pybamm/solvers/c_solvers/idaklu/observe.cpp b/src/pybamm/solvers/c_solvers/idaklu/observe.cpp new file mode 100644 index 0000000000..8f1d90e55d --- /dev/null +++ b/src/pybamm/solvers/c_solvers/idaklu/observe.cpp @@ -0,0 +1,343 @@ +#include "observe.hpp" + +int _setup_len_spatial(const std::vector& shape) { + // Calculate the product of all dimensions except the last (spatial dimensions) + int size_spatial = 1; + for (size_t i = 0; i < shape.size() - 1; ++i) { + size_spatial *= shape[i]; + } + + if (size_spatial == 0 || shape.back() == 0) { + throw std::invalid_argument("output array must have at least one element"); + } + + return size_spatial; +} + +// Coupled observe and Hermite interpolation of variables +class HermiteInterpolator { +public: + HermiteInterpolator(const py::detail::unchecked_reference& t, + const py::detail::unchecked_reference& y, + const py::detail::unchecked_reference& yp) + : t(t), y(y), yp(yp) {} + + void compute_knots(const size_t j, vector& c, vector& d) const { + // Called at the start of each interval + const realtype h_full = t(j + 1) - t(j); + const realtype inv_h = 1.0 / h_full; + const realtype inv_h2 = inv_h * inv_h; + const realtype inv_h3 = inv_h2 * inv_h; + + for (size_t i = 0; i < y.shape(0); ++i) { + realtype y_ij = y(i, j); + realtype yp_ij = yp(i, j); + realtype y_ijp1 = y(i, j + 1); + realtype yp_ijp1 = yp(i, j + 1); + + c[i] = 3.0 * (y_ijp1 - y_ij) * inv_h2 - (2.0 * yp_ij + yp_ijp1) * inv_h; + d[i] = 2.0 * (y_ij - y_ijp1) * inv_h3 + (yp_ij + yp_ijp1) * inv_h2; + } + } + + void interpolate(vector& entries, + realtype t_interp, + const size_t j, + vector& c, + vector& d) const { + // Must be called after compute_knots + const realtype h = t_interp - t(j); + const realtype h2 = h * h; + const realtype h3 = h2 * h; + + for (size_t i = 0; i < entries.size(); ++i) { + realtype y_ij = y(i, j); + realtype yp_ij = yp(i, j); + entries[i] = y_ij + yp_ij * h + c[i] * h2 + d[i] * h3; + } + } + +private: + const py::detail::unchecked_reference& t; + const py::detail::unchecked_reference& y; + const py::detail::unchecked_reference& yp; +}; + +class TimeSeriesInterpolator { +public: + TimeSeriesInterpolator(const np_array_realtype& _t_interp, + const vector& _ts_data, + const vector& _ys_data, + const vector& _yps_data, + const vector& _inputs, + const vector>& _funcs, + realtype* _entries, + const int _size_spatial) + : t_interp_np(_t_interp), ts_data_np(_ts_data), ys_data_np(_ys_data), + yps_data_np(_yps_data), inputs_np(_inputs), funcs(_funcs), + entries(_entries), size_spatial(_size_spatial) {} + + void process() { + auto t_interp = t_interp_np.unchecked<1>(); + ssize_t i_interp = 0; + int i_entries = 0; + const ssize_t N_interp = t_interp.size(); + + // Main processing within bounds + process_within_bounds(i_interp, i_entries, t_interp, N_interp); + + // Extrapolation for remaining points + if (i_interp < N_interp) { + extrapolate_remaining(i_interp, i_entries, t_interp, N_interp); + } + } + + void process_within_bounds( + ssize_t& i_interp, + int& i_entries, + const py::detail::unchecked_reference& t_interp, + const ssize_t N_interp + ) { + for (size_t i = 0; i < ts_data_np.size(); i++) { + const auto& t_data = ts_data_np[i].unchecked<1>(); + const realtype t_data_final = t_data(t_data.size() - 1); + realtype t_interp_next = t_interp(i_interp); + // Continue if the next interpolation point is beyond the final data point + if (t_interp_next > t_data_final) { + continue; + } + + const auto& y_data = ys_data_np[i].unchecked<2>(); + const auto& yp_data = yps_data_np[i].unchecked<2>(); + const auto input = inputs_np[i].data(); + const auto func = *funcs[i]; + + resize_arrays(y_data.shape(0), funcs[i]); + args[1] = y_buffer.data(); + args[2] = input; + + ssize_t j = 0; + ssize_t j_prev = -1; + const auto itp = HermiteInterpolator(t_data, y_data, yp_data); + while (t_interp_next <= t_data_final) { + for (; j < t_data.size() - 2; ++j) { + if (t_data(j) <= t_interp_next && t_interp_next <= t_data(j + 1)) { + break; + } + } + + if (j != j_prev) { + // Compute c and d for the new interval + itp.compute_knots(j, c, d); + } + + itp.interpolate(y_buffer, t_interp(i_interp), j, c, d); + + args[0] = &t_interp(i_interp); + results[0] = &entries[i_entries]; + func(args.data(), results.data(), iw.data(), w.data(), 0); + + ++i_interp; + if (i_interp == N_interp) { + return; + } + t_interp_next = t_interp(i_interp); + i_entries += size_spatial; + j_prev = j; + } + } + } + + void extrapolate_remaining( + ssize_t& i_interp, + int& i_entries, + const py::detail::unchecked_reference& t_interp, + const ssize_t N_interp + ) { + const auto& t_data = ts_data_np.back().unchecked<1>(); + const auto& y_data = ys_data_np.back().unchecked<2>(); + const auto& yp_data = yps_data_np.back().unchecked<2>(); + const auto input = inputs_np.back().data(); + const auto func = *funcs.back(); + const ssize_t j = t_data.size() - 2; + + resize_arrays(y_data.shape(0), funcs.back()); + args[1] = y_buffer.data(); + args[2] = input; + + const auto itp = HermiteInterpolator(t_data, y_data, yp_data); + itp.compute_knots(j, c, d); + + for (; i_interp < N_interp; ++i_interp) { + const realtype t_interp_next = t_interp(i_interp); + itp.interpolate(y_buffer, t_interp_next, j, c, d); + + args[0] = &t_interp_next; + results[0] = &entries[i_entries]; + func(args.data(), results.data(), iw.data(), w.data(), 0); + + i_entries += size_spatial; + } + } + + void resize_arrays(const int M, std::shared_ptr func) { + args.resize(func->sz_arg()); + results.resize(func->sz_res()); + iw.resize(func->sz_iw()); + w.resize(func->sz_w()); + if (y_buffer.size() < M) { + y_buffer.resize(M); + c.resize(M); + d.resize(M); + } + } + +private: + const np_array_realtype& t_interp_np; + const vector& ts_data_np; + const vector& ys_data_np; + const vector& yps_data_np; + const vector& inputs_np; + const vector>& funcs; + realtype* entries; + const int size_spatial; + vector c; + vector d; + vector y_buffer; + vector args; + vector results; + vector iw; + vector w; +}; + +// Observe the raw data +class TimeSeriesProcessor { +public: + TimeSeriesProcessor(const vector& _ts, + const vector& _ys, + const vector& _inputs, + const vector>& _funcs, + realtype* _entries, + const bool _is_f_contiguous, + const int _size_spatial) + : ts(_ts), ys(_ys), inputs(_inputs), funcs(_funcs), + entries(_entries), is_f_contiguous(_is_f_contiguous), size_spatial(_size_spatial) {} + + void process() { + int i_entries = 0; + for (size_t i = 0; i < ts.size(); i++) { + const auto& t = ts[i].unchecked<1>(); + const auto& y = ys[i].unchecked<2>(); + const auto input = inputs[i].data(); + const auto func = *funcs[i]; + + resize_arrays(y.shape(0), funcs[i]); + args[2] = input; + + for (size_t j = 0; j < t.size(); j++) { + const realtype t_val = t(j); + const realtype* y_val = is_f_contiguous ? &y(0, j) : copy_to_buffer(y_buffer, y, j); + + args[0] = &t_val; + args[1] = y_val; + results[0] = &entries[i_entries]; + + func(args.data(), results.data(), iw.data(), w.data(), 0); + + i_entries += size_spatial; + } + } + } + +private: + const realtype* copy_to_buffer( + vector& entries, + const py::detail::unchecked_reference& y, + size_t j) { + for (size_t i = 0; i < entries.size(); ++i) { + entries[i] = y(i, j); + } + + return entries.data(); + } + + void resize_arrays(const int M, std::shared_ptr func) { + args.resize(func->sz_arg()); + results.resize(func->sz_res()); + iw.resize(func->sz_iw()); + w.resize(func->sz_w()); + if (!is_f_contiguous && y_buffer.size() < M) { + y_buffer.resize(M); + } + } + + const vector& ts; + const vector& ys; + const vector& inputs; + const vector>& funcs; + realtype* entries; + const bool is_f_contiguous; + int size_spatial; + vector y_buffer; + vector args; + vector results; + vector iw; + vector w; +}; + +const np_array_realtype observe_hermite_interp( + const np_array_realtype& t_interp_np, + const vector& ts_np, + const vector& ys_np, + const vector& yps_np, + const vector& inputs_np, + const vector& strings, + const vector& shape +) { + const int size_spatial = _setup_len_spatial(shape); + const auto& funcs = setup_casadi_funcs(strings); + py::array_t out_array(shape); + auto entries = out_array.mutable_data(); + + TimeSeriesInterpolator(t_interp_np, ts_np, ys_np, yps_np, inputs_np, funcs, entries, size_spatial).process(); + + return out_array; +} + +const np_array_realtype observe( + const vector& ts_np, + const vector& ys_np, + const vector& inputs_np, + const vector& strings, + const bool is_f_contiguous, + const vector& shape +) { + const int size_spatial = _setup_len_spatial(shape); + const auto& funcs = setup_casadi_funcs(strings); + py::array_t out_array(shape); + auto entries = out_array.mutable_data(); + + TimeSeriesProcessor(ts_np, ys_np, inputs_np, funcs, entries, is_f_contiguous, size_spatial).process(); + + return out_array; +} + +const vector> setup_casadi_funcs(const vector& strings) { + std::unordered_map> function_cache; + vector> funcs(strings.size()); + + for (size_t i = 0; i < strings.size(); ++i) { + const std::string& str = strings[i]; + + // Check if function is already in the local cache + if (function_cache.find(str) == function_cache.end()) { + // If not in the cache, create a new casadi::Function::deserialize and store it + function_cache[str] = std::make_shared(casadi::Function::deserialize(str)); + } + + // Retrieve the function from the cache as a shared pointer + funcs[i] = function_cache[str]; + } + + return funcs; +} diff --git a/src/pybamm/solvers/c_solvers/idaklu/observe.hpp b/src/pybamm/solvers/c_solvers/idaklu/observe.hpp new file mode 100644 index 0000000000..52a0cfdf84 --- /dev/null +++ b/src/pybamm/solvers/c_solvers/idaklu/observe.hpp @@ -0,0 +1,47 @@ +#ifndef PYBAMM_CREATE_OBSERVE_HPP +#define PYBAMM_CREATE_OBSERVE_HPP + +#include +#include +#include +#include "common.hpp" +#include +#include +using std::vector; + +#if defined(_MSC_VER) + #include + typedef SSIZE_T ssize_t; +#endif + +/** + * @brief Observe and Hermite interpolate ND variables + */ +const np_array_realtype observe_hermite_interp( + const np_array_realtype& t_interp, + const vector& ts, + const vector& ys, + const vector& yps, + const vector& inputs, + const vector& strings, + const vector& shape +); + + +/** + * @brief Observe ND variables + */ +const np_array_realtype observe( + const vector& ts_np, + const vector& ys_np, + const vector& inputs_np, + const vector& strings, + const bool is_f_contiguous, + const vector& shape +); + +const vector> setup_casadi_funcs(const vector& strings); + +int _setup_len_spatial(const vector& shape); + +#endif // PYBAMM_CREATE_OBSERVE_HPP diff --git a/src/pybamm/solvers/idaklu_solver.py b/src/pybamm/solvers/idaklu_solver.py index 2b2d852697..80eaffebf4 100644 --- a/src/pybamm/solvers/idaklu_solver.py +++ b/src/pybamm/solvers/idaklu_solver.py @@ -30,6 +30,7 @@ if idaklu_spec.loader: idaklu_spec.loader.exec_module(idaklu) except ImportError as e: # pragma: no cover + idaklu = None print(f"Error loading idaklu: {e}") idaklu_spec = None @@ -133,6 +134,10 @@ class IDAKLUSolver(pybamm.BaseSolver): "nonlinear_convergence_coefficient": 0.33, # Suppress algebraic variables from error test "suppress_algebraic_error": False, + # Store Hermite interpolation data for the solution. + # Note: this option is always disabled if output_variables are given + # or if t_interp values are specified + "hermite_interpolation": True, ## Initial conditions calculation # Positive constant in the Newton iteration convergence test within the # initial condition calculation @@ -201,6 +206,7 @@ def __init__( "max_convergence_failures": 100, "nonlinear_convergence_coefficient": 0.33, "suppress_algebraic_error": False, + "hermite_interpolation": True, "nonlinear_convergence_coefficient_ic": 0.0033, "max_num_steps_ic": 50, "max_num_jacobians_ic": 40, @@ -756,6 +762,16 @@ def _integrate(self, model, t_eval, inputs_list=None, t_interp=None): The times (in seconds) at which to interpolate the solution. Defaults to `None`, which returns the adaptive time-stepping times. """ + if not ( + model.convert_to_format == "casadi" + or ( + model.convert_to_format == "jax" + and self._options["jax_evaluator"] == "iree" + ) + ): # pragma: no cover + # Shouldn't ever reach this point + raise pybamm.SolverError("Unsupported IDAKLU solver configuration.") + inputs_list = inputs_list or [{}] # stack inputs so that they are a 2D array of shape (number_of_inputs, number_of_parameters) @@ -779,20 +795,13 @@ def _integrate(self, model, t_eval, inputs_list=None, t_interp=None): atol = self._check_atol_type(atol, y0full.size) timer = pybamm.Timer() - if model.convert_to_format == "casadi" or ( - model.convert_to_format == "jax" - and self._options["jax_evaluator"] == "iree" - ): - solns = self._setup["solver"].solve( - t_eval, - t_interp, - y0full, - ydot0full, - inputs, - ) - else: # pragma: no cover - # Shouldn't ever reach this point - raise pybamm.SolverError("Unsupported IDAKLU solver configuration.") + solns = self._setup["solver"].solve( + t_eval, + t_interp, + y0full, + ydot0full, + inputs, + ) integration_time = timer.time() return [ @@ -807,7 +816,8 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict): sensitivity_names = self._setup["sensitivity_names"] number_of_timesteps = sol.t.size number_of_states = model.len_rhs_and_alg - if self.output_variables: + save_outputs_only = self.output_variables + if save_outputs_only: # Substitute empty vectors for state vector 'y' y_out = np.zeros((number_of_timesteps * number_of_states, 0)) y_event = sol.y_term @@ -838,6 +848,11 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict): else: raise pybamm.SolverError(f"FAILURE {self._solver_flag(sol.flag)}") + if sol.yp.size > 0: + yp = sol.yp.reshape((number_of_timesteps, number_of_states)).T + else: + yp = None + newsol = pybamm.Solution( sol.t, np.transpose(y_out), @@ -847,10 +862,11 @@ def _post_process_solution(self, sol, model, integration_time, inputs_dict): np.transpose(y_event)[:, np.newaxis], termination, all_sensitivities=yS_out, + all_yps=yp, ) + newsol.integration_time = integration_time - if not self.output_variables: - # print((newsol.y).shape) + if not save_outputs_only: return newsol # Populate variables and sensititivies dictionaries directly diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py index 2464466348..5cf928ca7f 100644 --- a/src/pybamm/solvers/processed_variable.py +++ b/src/pybamm/solvers/processed_variable.py @@ -6,6 +6,7 @@ import pybamm from scipy.integrate import cumulative_trapezoid import xarray as xr +import bisect class ProcessedVariable: @@ -23,14 +24,11 @@ class ProcessedVariable: Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variable_casadis : list of :class:`casadi.Function` + base_variables_casadi : list of :class:`casadi.Function` A list of casadi functions. When evaluated, returns the same thing as `base_Variable.evaluate` (but more efficiently). solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables - warn : bool, optional - Whether to raise warnings when trying to evaluate time and length scales. - Default is True. """ def __init__( @@ -38,7 +36,6 @@ def __init__( base_variables, base_variables_casadi, solution, - warn=True, cumtrapz_ic=None, ): self.base_variables = base_variables @@ -46,13 +43,13 @@ def __init__( self.all_ts = solution.all_ts self.all_ys = solution.all_ys + self.all_yps = solution.all_yps self.all_inputs = solution.all_inputs self.all_inputs_casadi = solution.all_inputs_casadi self.mesh = base_variables[0].mesh self.domain = base_variables[0].domain self.domains = base_variables[0].domains - self.warn = warn self.cumtrapz_ic = cumtrapz_ic # Process spatial variables @@ -75,46 +72,431 @@ def __init__( self.base_eval_shape = self.base_variables[0].shape self.base_eval_size = self.base_variables[0].size - # xr_data_array is initialized - self._xr_data_array = None + self._xr_array_raw = None + self._entries_raw = None + self._entries_for_interp_raw = None + self._coords_raw = None - # handle 2D (in space) finite element variables differently - if ( - self.mesh - and "current collector" in self.domain - and isinstance(self.mesh, pybamm.ScikitSubMesh2D) + def initialise(self): + if self.entries_raw_initialized: + return + + entries = self.observe_raw() + + t = self.t_pts + entries_for_interp, coords = self._interp_setup(entries, t) + + self._entries_raw = entries + self._entries_for_interp_raw = entries_for_interp + self._coords_raw = coords + + def observe_and_interp(self, t, fill_value): + """ + Interpolate the variable at the given time points and y values. + t must be a sorted array of time points. + """ + + entries = self._observe_hermite_cpp(t) + processed_entries = self._observe_postfix(entries, t) + + tf = self.t_pts[-1] + if t[-1] > tf and fill_value != "extrapolate": + # fill the rest + idx = np.searchsorted(t, tf, side="right") + processed_entries[..., idx:] = fill_value + + return processed_entries + + def observe_raw(self): + """ + Evaluate the base variable at the given time points and y values. + """ + t = self.t_pts + + # For small number of points, use Python + if pybamm.has_idaklu(): + entries = self._observe_raw_cpp() + else: + # Fallback method for when IDAKLU is not available. To be removed + # when the C++ code is migrated to a new repo + entries = self._observe_raw_python() # pragma: no cover + + return self._observe_postfix(entries, t) + + def _setup_cpp_inputs(self, t, full_range): + pybamm.logger.debug("Setting up C++ interpolation inputs") + + ts = self.all_ts + ys = self.all_ys + yps = self.all_yps + inputs = self.all_inputs_casadi + # Find the indices of the time points to observe + if full_range: + idxs = range(len(ts)) + else: + idxs = _find_ts_indices(ts, t) + + if isinstance(idxs, list): + # Extract the time points and inputs + ts = [ts[idx] for idx in idxs] + ys = [ys[idx] for idx in idxs] + if self.hermite_interpolation: + yps = [yps[idx] for idx in idxs] + inputs = [self.all_inputs_casadi[idx] for idx in idxs] + + is_f_contiguous = _is_f_contiguous(ys) + + ts = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(ts) + ys = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(ys) + if self.hermite_interpolation: + yps = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(yps) + else: + yps = None + inputs = pybamm.solvers.idaklu_solver.idaklu.VectorRealtypeNdArray(inputs) + + # Generate the serialized C++ functions only once + funcs_unique = {} + funcs = [None] * len(idxs) + for i in range(len(idxs)): + vars = self.base_variables_casadi[idxs[i]] + if vars not in funcs_unique: + funcs_unique[vars] = vars.serialize() + funcs[i] = funcs_unique[vars] + + return ts, ys, yps, funcs, inputs, is_f_contiguous + + def _observe_hermite_cpp(self, t): + pybamm.logger.debug("Observing and Hermite interpolating the variable in C++") + + ts, ys, yps, funcs, inputs, _ = self._setup_cpp_inputs(t, full_range=False) + shapes = self._shape(t) + return pybamm.solvers.idaklu_solver.idaklu.observe_hermite_interp( + t, ts, ys, yps, inputs, funcs, shapes + ) + + def _observe_raw_cpp(self): + pybamm.logger.debug("Observing the variable raw data in C++") + t = self.t_pts + ts, ys, _, funcs, inputs, is_f_contiguous = self._setup_cpp_inputs( + t, full_range=True + ) + shapes = self._shape(self.t_pts) + + return pybamm.solvers.idaklu_solver.idaklu.observe( + ts, ys, inputs, funcs, is_f_contiguous, shapes + ) + + def _observe_raw_python(self): + raise NotImplementedError # pragma: no cover + + def _observe_postfix(self, entries, t): + return entries + + def _interp_setup(self, entries, t): + raise NotImplementedError # pragma: no cover + + def _shape(self, t): + raise NotImplementedError # pragma: no cover + + def _process_spatial_variable_names(self, spatial_variable): + if len(spatial_variable) == 0: + return None + + # Extract names + raw_names = [] + for var in spatial_variable: + # Ignore tabs in domain names + if var == "tabs": + continue + if isinstance(var, str): + raw_names.append(var) + else: + raw_names.append(var.name) + + # Rename battery variables to match PyBaMM convention + if all([var.startswith("r") for var in raw_names]): + return "r" + elif all([var.startswith("x") for var in raw_names]): + return "x" + elif all([var.startswith("R") for var in raw_names]): + return "R" + elif len(raw_names) == 1: + return raw_names[0] + else: + raise NotImplementedError( + f"Spatial variable name not recognized for {spatial_variable}" + ) + + def __call__( + self, + t=None, + x=None, + r=None, + y=None, + z=None, + R=None, + fill_value=np.nan, + ): + # Check to see if we are interpolating exactly onto the raw solution time points + t_observe, observe_raw = self._check_observe_raw(t) + + # Check if the time points are sorted and unique + is_sorted = observe_raw or _is_sorted(t_observe) + + # Sort them if not + if not is_sorted: + idxs_sort = np.argsort(t_observe) + t_observe = t_observe[idxs_sort] + + hermite_time_interp = ( + pybamm.has_idaklu() and self.hermite_interpolation and not observe_raw + ) + + if hermite_time_interp: + entries = self.observe_and_interp(t_observe, fill_value) + + spatial_interp = any(a is not None for a in [x, r, y, z, R]) + + xr_interp = spatial_interp or not hermite_time_interp + + if xr_interp: + if hermite_time_interp: + # Already interpolated in time + t = None + entries_for_interp, coords = self._interp_setup(entries, t_observe) + else: + self.initialise() + entries_for_interp, coords = ( + self._entries_for_interp_raw, + self._coords_raw, + ) + + processed_entries = self._xr_interpolate( + entries_for_interp, + coords, + observe_raw, + t, + x, + r, + y, + z, + R, + fill_value, + ) + else: + processed_entries = entries + + if not is_sorted: + idxs_unsort = np.zeros_like(idxs_sort) + idxs_unsort[idxs_sort] = np.arange(len(t_observe)) + + processed_entries = processed_entries[..., idxs_unsort] + + # Remove a singleton time dimension if we interpolate in time with hermite + if hermite_time_interp and t_observe.size == 1: + processed_entries = np.squeeze(processed_entries, axis=-1) + + return processed_entries + + def _xr_interpolate( + self, + entries_for_interp, + coords, + observe_raw, + t=None, + x=None, + r=None, + y=None, + z=None, + R=None, + fill_value=None, + ): + """ + Evaluate the variable at arbitrary *dimensional* t (and x, r, y, z and/or R), + using interpolation + """ + if observe_raw: + if not self.xr_array_raw_initialized: + self._xr_array_raw = xr.DataArray(entries_for_interp, coords=coords) + xr_data_array = self._xr_array_raw + else: + xr_data_array = xr.DataArray(entries_for_interp, coords=coords) + + kwargs = {"t": t, "x": x, "r": r, "y": y, "z": z, "R": R} + + # Remove any None arguments + kwargs = {key: value for key, value in kwargs.items() if value is not None} + + # Use xarray interpolation, return numpy array + out = xr_data_array.interp(**kwargs, kwargs={"fill_value": fill_value}).values + + return out + + def _check_observe_raw(self, t): + """ + Checks if the raw data should be observed exactly at the solution time points + + Args: + t (np.ndarray, list, None): time points to observe + + Returns: + t_observe (np.ndarray): time points to observe + observe_raw (bool): True if observing the raw data + """ + observe_raw = (t is None) or ( + np.asarray(t).size == len(self.t_pts) and np.all(t == self.t_pts) + ) + + if observe_raw: + t_observe = self.t_pts + elif not isinstance(t, np.ndarray): + if not isinstance(t, list): + t = [t] + t_observe = np.array(t) + else: + t_observe = t + + if t_observe[0] < self.t_pts[0]: + raise ValueError( + "The interpolation points must be greater than or equal to the initial solution time" + ) + + return t_observe, observe_raw + + @property + def entries(self): + """ + Returns the raw data entries of the processed variable. If the processed + variable has not been initialized (i.e. the entries have not been + calculated), then the processed variable is initialized first. + """ + self.initialise() + return self._entries_raw + + @property + def data(self): + """Same as entries, but different name""" + return self.entries + + @property + def entries_raw_initialized(self): + return self._entries_raw is not None + + @property + def xr_array_raw_initialized(self): + return self._xr_array_raw is not None + + @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.all_solution_sensitivities: + self.initialise_sensitivity_explicit_forward() + else: + raise ValueError( + "Cannot compute sensitivities. The 'calculate_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" + + all_S_var = [] + for ts, ys, inputs_stacked, inputs, base_variable, dy_dp in zip( + self.all_ts, + self.all_ys, + self.all_inputs_casadi, + self.all_inputs, + self.base_variables, + self.all_solution_sensitivities["all"], ): - return self.initialise_2D_scikit_fem() + # Set up symbolic variables + t_casadi = casadi.MX.sym("t") + y_casadi = casadi.MX.sym("y", ys.shape[0]) + p_casadi = { + name: casadi.MX.sym(name, value.shape[0]) + for name, value in inputs.items() + } - # check variable shape - if len(self.base_eval_shape) == 0 or self.base_eval_shape[0] == 1: - return self.initialise_0D() + p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) - n = self.mesh.npts - base_shape = self.base_eval_shape[0] - # Try some shapes that could make the variable a 1D variable - if base_shape in [n, n + 1]: - return self.initialise_1D() + # Convert variable to casadi format for differentiating + var_casadi = base_variable.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) - # Try some shapes that could make the variable a 2D variable - first_dim_nodes = self.mesh.nodes - first_dim_edges = self.mesh.edges - second_dim_pts = self.base_variables[0].secondary_mesh.nodes - if self.base_eval_size // len(second_dim_pts) in [ - len(first_dim_nodes), - len(first_dim_edges), - ]: - return self.initialise_2D() - - # Raise error for 3D variable - raise NotImplementedError( - f"Shape not recognized for {base_variables[0]}" - + "(note processing of 3D variables is not yet implemented)" + # 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 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 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 + S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval + all_S_var.append(S_var) + + S_var = casadi.vertcat(*all_S_var) + 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 + + @property + def hermite_interpolation(self): + return self.all_yps is not None + + +class ProcessedVariable0D(ProcessedVariable): + def __init__( + self, + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=None, + ): + self.dimensions = 0 + super().__init__( + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=cumtrapz_ic, ) - def initialise_0D(self): + def _observe_raw_python(self): + pybamm.logger.debug("Observing the variable raw data in Python") # initialise empty array of the correct size - entries = np.empty(len(self.t_pts)) + entries = np.empty(self._shape(self.t_pts)) idx = 0 # Evaluate the base_variable index-by-index for ts, ys, inputs, base_var_casadi in zip( @@ -126,22 +508,67 @@ def initialise_0D(self): entries[idx] = float(base_var_casadi(t, y, inputs)) idx += 1 + return entries - if self.cumtrapz_ic is not None: - entries = cumulative_trapezoid( - entries, self.t_pts, initial=float(self.cumtrapz_ic) - ) + def _observe_postfix(self, entries, _): + if self.cumtrapz_ic is None: + return entries + + return cumulative_trapezoid( + entries, self.t_pts, initial=float(self.cumtrapz_ic) + ) + def _interp_setup(self, entries, t): # save attributes for interpolation - self.entries_for_interp = entries - self.coords_for_interp = {"t": self.t_pts} + entries_for_interp = entries + coords_for_interp = {"t": t} - self.entries = entries - self.dimensions = 0 + return entries_for_interp, coords_for_interp - def initialise_1D(self, fixed_t=False): - len_space = self.base_eval_shape[0] - entries = np.empty((len_space, len(self.t_pts))) + def _shape(self, t): + return [len(t)] + + +class ProcessedVariable1D(ProcessedVariable): + """ + An object that can be evaluated at arbitrary (scalars or vectors) t and x, and + returns the (interpolated) value of the base variable at that t and x. + + Parameters + ---------- + base_variables : list of :class:`pybamm.Symbol` + A list of base variables with a method `evaluate(t,y)`, each entry of which + returns the value of that variable for that particular sub-solution. + A Solution can be comprised of sub-solutions which are the solutions of + different models. + Note that this can be any kind of node in the expression tree, not + just a :class:`pybamm.Variable`. + When evaluated, returns an array of size (m,n) + base_variables_casadi : list of :class:`casadi.Function` + A list of casadi functions. When evaluated, returns the same thing as + `base_Variable.evaluate` (but more efficiently). + solution : :class:`pybamm.Solution` + The solution object to be used to create the processed variables + """ + + def __init__( + self, + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=None, + ): + self.dimensions = 1 + super().__init__( + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=cumtrapz_ic, + ) + + def _observe_raw_python(self): + pybamm.logger.debug("Observing the variable raw data in Python") + entries = np.empty(self._shape(self.t_pts)) # Evaluate the base_variable index-by-index idx = 0 @@ -153,7 +580,9 @@ def initialise_1D(self, fixed_t=False): y = ys[:, inner_idx] entries[:, idx] = base_var_casadi(t, y, inputs).full()[:, 0] idx += 1 + return entries + def _interp_setup(self, entries, t): # Get node and edge values nodes = self.mesh.nodes edges = self.mesh.edges @@ -173,8 +602,6 @@ def initialise_1D(self, fixed_t=False): ) # assign attributes for reference (either x_sol or r_sol) - self.entries = entries - self.dimensions = 1 self.spatial_variable_names = { k: self._process_spatial_variable_names(v) for k, v in self.spatial_variables.items() @@ -189,26 +616,71 @@ def initialise_1D(self, fixed_t=False): self.first_dim_pts = edges # save attributes for interpolation - self.entries_for_interp = entries_for_interp - self.coords_for_interp = {self.first_dimension: pts_for_interp, "t": self.t_pts} + coords_for_interp = {self.first_dimension: pts_for_interp, "t": t} - def initialise_2D(self): - """ - Initialise a 2D object that depends on x and r, x and z, x and R, or R and r. - """ + return entries_for_interp, coords_for_interp + + def _shape(self, t): + t_size = len(t) + space_size = self.base_eval_shape[0] + return [space_size, t_size] + + +class ProcessedVariable2D(ProcessedVariable): + """ + An object that can be evaluated at arbitrary (scalars or vectors) t and x, and + returns the (interpolated) value of the base variable at that t and x. + + Parameters + ---------- + base_variables : list of :class:`pybamm.Symbol` + A list of base variables with a method `evaluate(t,y)`, each entry of which + returns the value of that variable for that particular sub-solution. + A Solution can be comprised of sub-solutions which are the solutions of + different models. + Note that this can be any kind of node in the expression tree, not + just a :class:`pybamm.Variable`. + When evaluated, returns an array of size (m,n) + base_variables_casadi : list of :class:`casadi.Function` + A list of casadi functions. When evaluated, returns the same thing as + `base_Variable.evaluate` (but more efficiently). + solution : :class:`pybamm.Solution` + The solution object to be used to create the processed variables + """ + + def __init__( + self, + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=None, + ): + self.dimensions = 2 + super().__init__( + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=cumtrapz_ic, + ) first_dim_nodes = self.mesh.nodes first_dim_edges = self.mesh.edges second_dim_nodes = self.base_variables[0].secondary_mesh.nodes - second_dim_edges = self.base_variables[0].secondary_mesh.edges if self.base_eval_size // len(second_dim_nodes) == len(first_dim_nodes): first_dim_pts = first_dim_nodes elif self.base_eval_size // len(second_dim_nodes) == len(first_dim_edges): first_dim_pts = first_dim_edges second_dim_pts = second_dim_nodes - first_dim_size = len(first_dim_pts) - second_dim_size = len(second_dim_pts) - entries = np.empty((first_dim_size, second_dim_size, len(self.t_pts))) + self.first_dim_size = len(first_dim_pts) + self.second_dim_size = len(second_dim_pts) + + def _observe_raw_python(self): + """ + Initialise a 2D object that depends on x and r, x and z, x and R, or R and r. + """ + pybamm.logger.debug("Observing the variable raw data in Python") + first_dim_size, second_dim_size, t_size = self._shape(self.t_pts) + entries = np.empty((first_dim_size, second_dim_size, t_size)) # Evaluate the base_variable index-by-index idx = 0 @@ -224,6 +696,22 @@ def initialise_2D(self): order="F", ) idx += 1 + return entries + + def _interp_setup(self, entries, t): + """ + Initialise a 2D object that depends on x and r, x and z, x and R, or R and r. + """ + first_dim_nodes = self.mesh.nodes + first_dim_edges = self.mesh.edges + second_dim_nodes = self.base_variables[0].secondary_mesh.nodes + second_dim_edges = self.base_variables[0].secondary_mesh.edges + if self.base_eval_size // len(second_dim_nodes) == len(first_dim_nodes): + first_dim_pts = first_dim_nodes + elif self.base_eval_size // len(second_dim_nodes) == len(first_dim_edges): + first_dim_pts = first_dim_edges + + second_dim_pts = second_dim_nodes # add points outside first dimension domain for extrapolation to # boundaries @@ -281,8 +769,6 @@ def initialise_2D(self): self.second_dimension = self.spatial_variable_names["secondary"] # assign attributes for reference - self.entries = entries - self.dimensions = 2 first_dim_pts_for_interp = first_dim_pts second_dim_pts_for_interp = second_dim_pts @@ -291,38 +777,68 @@ def initialise_2D(self): self.second_dim_pts = second_dim_edges # save attributes for interpolation - self.entries_for_interp = entries_for_interp - self.coords_for_interp = { + coords_for_interp = { self.first_dimension: first_dim_pts_for_interp, self.second_dimension: second_dim_pts_for_interp, - "t": self.t_pts, + "t": t, } - def initialise_2D_scikit_fem(self): + return entries_for_interp, coords_for_interp + + def _shape(self, t): + first_dim_size = self.first_dim_size + second_dim_size = self.second_dim_size + t_size = len(t) + return [first_dim_size, second_dim_size, t_size] + + +class ProcessedVariable2DSciKitFEM(ProcessedVariable2D): + """ + An object that can be evaluated at arbitrary (scalars or vectors) t and x, and + returns the (interpolated) value of the base variable at that t and x. + + Parameters + ---------- + base_variables : list of :class:`pybamm.Symbol` + A list of base variables with a method `evaluate(t,y)`, each entry of which + returns the value of that variable for that particular sub-solution. + A Solution can be comprised of sub-solutions which are the solutions of + different models. + Note that this can be any kind of node in the expression tree, not + just a :class:`pybamm.Variable`. + When evaluated, returns an array of size (m,n) + base_variables_casadi : list of :class:`casadi.Function` + A list of casadi functions. When evaluated, returns the same thing as + `base_Variable.evaluate` (but more efficiently). + solution : :class:`pybamm.Solution` + The solution object to be used to create the processed variables + """ + + def __init__( + self, + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=None, + ): + self.dimensions = 2 + super(ProcessedVariable2D, self).__init__( + base_variables, + base_variables_casadi, + solution, + cumtrapz_ic=cumtrapz_ic, + ) y_sol = self.mesh.edges["y"] - len_y = len(y_sol) z_sol = self.mesh.edges["z"] - len_z = len(z_sol) - entries = np.empty((len_y, len_z, len(self.t_pts))) - # Evaluate the base_variable index-by-index - idx = 0 - for ts, ys, inputs, base_var_casadi in zip( - self.all_ts, self.all_ys, self.all_inputs_casadi, self.base_variables_casadi - ): - for inner_idx, t in enumerate(ts): - t = ts[inner_idx] - y = ys[:, inner_idx] - entries[:, :, idx] = np.reshape( - base_var_casadi(t, y, inputs).full(), - [len_y, len_z], - order="C", - ) - idx += 1 + self.first_dim_size = len(y_sol) + self.second_dim_size = len(z_sol) + + def _interp_setup(self, entries, t): + y_sol = self.mesh.edges["y"] + z_sol = self.mesh.edges["z"] # assign attributes for reference - self.entries = entries - self.dimensions = 2 self.y_sol = y_sol self.z_sol = z_sol self.first_dimension = "y" @@ -331,148 +847,116 @@ def initialise_2D_scikit_fem(self): self.second_dim_pts = z_sol # save attributes for interpolation - self.entries_for_interp = entries - self.coords_for_interp = {"y": y_sol, "z": z_sol, "t": self.t_pts} + coords_for_interp = {"y": y_sol, "z": z_sol, "t": t} - def _process_spatial_variable_names(self, spatial_variable): - if len(spatial_variable) == 0: - return None + return entries, coords_for_interp - # Extract names - raw_names = [] - for var in spatial_variable: - # Ignore tabs in domain names - if var == "tabs": - continue - if isinstance(var, str): - raw_names.append(var) - else: - raw_names.append(var.name) - # Rename battery variables to match PyBaMM convention - if all([var.startswith("r") for var in raw_names]): - return "r" - elif all([var.startswith("x") for var in raw_names]): - return "x" - elif all([var.startswith("R") for var in raw_names]): - return "R" - elif len(raw_names) == 1: - return raw_names[0] - else: - raise NotImplementedError( - f"Spatial variable name not recognized for {spatial_variable}" - ) +def process_variable(base_variables, *args, **kwargs): + mesh = base_variables[0].mesh + domain = base_variables[0].domain - def _initialize_xr_data_array(self): - """ - Initialize the xarray DataArray for interpolation. We don't do this by - default as it has some overhead (~75 us) and sometimes we only need the entries - of the processed variable, not the xarray object for interpolation. - """ - entries = self.entries_for_interp - coords = self.coords_for_interp - self._xr_data_array = xr.DataArray(entries, coords=coords) + # Evaluate base variable at initial time + base_eval_shape = base_variables[0].shape + base_eval_size = base_variables[0].size - def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None, warn=True): - """ - Evaluate the variable at arbitrary *dimensional* t (and x, r, y, z and/or R), - using interpolation - """ - if self._xr_data_array is None: - self._initialize_xr_data_array() - kwargs = {"t": t, "x": x, "r": r, "y": y, "z": z, "R": R} - # Remove any None arguments - kwargs = {key: value for key, value in kwargs.items() if value is not None} - # Use xarray interpolation, return numpy array - return self._xr_data_array.interp(**kwargs).values + # handle 2D (in space) finite element variables differently + if ( + mesh + and "current collector" in domain + and isinstance(mesh, pybamm.ScikitSubMesh2D) + ): + return ProcessedVariable2DSciKitFEM(base_variables, *args, **kwargs) + + # check variable shape + if len(base_eval_shape) == 0 or base_eval_shape[0] == 1: + return ProcessedVariable0D(base_variables, *args, **kwargs) + + n = mesh.npts + base_shape = base_eval_shape[0] + # Try some shapes that could make the variable a 1D variable + if base_shape in [n, n + 1]: + return ProcessedVariable1D(base_variables, *args, **kwargs) + + # Try some shapes that could make the variable a 2D variable + first_dim_nodes = mesh.nodes + first_dim_edges = mesh.edges + second_dim_pts = base_variables[0].secondary_mesh.nodes + if base_eval_size // len(second_dim_pts) in [ + len(first_dim_nodes), + len(first_dim_edges), + ]: + return ProcessedVariable2D(base_variables, *args, **kwargs) + + # Raise error for 3D variable + raise NotImplementedError( + f"Shape not recognized for {base_variables[0]}" + + "(note processing of 3D variables is not yet implemented)" + ) + + +def _is_f_contiguous(all_ys): + """ + Check if all the ys are f-contiguous in memory - @property - def data(self): - """Same as entries, but different name""" - return self.entries + Args: + all_ys (list of np.ndarray): list of all ys - @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.all_solution_sensitivities: - self.initialise_sensitivity_explicit_forward() - else: - raise ValueError( - "Cannot compute sensitivities. The 'calculate_sensitivities' " - "argument of the solver.solve should be changed from 'None' to " - "allow sensitivities calculations. Check solver documentation for " - "details." - ) - return self._sensitivities + Returns: + bool: True if all ys are f-contiguous + """ - def initialise_sensitivity_explicit_forward(self): - "Set up the sensitivity dictionary" + return all(isinstance(y, np.ndarray) and y.data.f_contiguous for y in all_ys) - all_S_var = [] - for ts, ys, inputs_stacked, inputs, base_variable, dy_dp in zip( - self.all_ts, - self.all_ys, - self.all_inputs_casadi, - self.all_inputs, - self.base_variables, - self.all_solution_sensitivities["all"], - ): - # Set up symbolic variables - t_casadi = casadi.MX.sym("t") - y_casadi = casadi.MX.sym("y", ys.shape[0]) - p_casadi = { - name: casadi.MX.sym(name, value.shape[0]) - for name, value in inputs.items() - } - p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()]) +def _is_sorted(t): + """ + Check if an array is sorted - # Convert variable to casadi format for differentiating - var_casadi = base_variable.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) + Args: + t (np.ndarray): array to check - # 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 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 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) + Returns: + bool: True if array is sorted + """ + return np.all(t[:-1] <= t[1:]) - # Compute sensitivity - S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval - all_S_var.append(S_var) - S_var = casadi.vertcat(*all_S_var) - sensitivities = {"all": S_var} +def _find_ts_indices(ts, t): + """ + Parameters: + - ts: A list of numpy arrays (each sorted) whose values are successively increasing. + - t: A sorted list or array of values to find within ts. - # 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 + Returns: + - indices: A list of indices from `ts` such that at least one value of `t` falls within ts[idx]. + """ - # Save attribute - self._sensitivities = sensitivities + indices = [] + + # Get the minimum and maximum values of the target values `t` + t_min, t_max = t[0], t[-1] + + # Step 1: Use binary search to find the range of `ts` arrays where t_min and t_max could lie + low_idx = bisect.bisect_left([ts_arr[-1] for ts_arr in ts], t_min) + high_idx = bisect.bisect_right([ts_arr[0] for ts_arr in ts], t_max) + + # Step 2: Iterate over the identified range + for idx in range(low_idx, high_idx): + ts_min, ts_max = ts[idx][0], ts[idx][-1] + + # Binary search within `t` to check if any value falls within [ts_min, ts_max] + i = bisect.bisect_left(t, ts_min) + if i < len(t) and t[i] <= ts_max: + # At least one value of t is within ts[idx] + indices.append(idx) + + # extrapolating + if (t[-1] > ts[-1][-1]) and (len(indices) == 0 or indices[-1] != len(ts) - 1): + indices.append(len(ts) - 1) + + if len(indices) == len(ts): + # All indices are included + return range(len(ts)) + + return indices diff --git a/src/pybamm/solvers/processed_variable_computed.py b/src/pybamm/solvers/processed_variable_computed.py index a717c8b0cb..4f0cccc8c3 100644 --- a/src/pybamm/solvers/processed_variable_computed.py +++ b/src/pybamm/solvers/processed_variable_computed.py @@ -27,7 +27,7 @@ class ProcessedVariableComputed: Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variable_casadis : list of :class:`casadi.Function` + base_variables_casadi : list of :class:`casadi.Function` A list of casadi functions. When evaluated, returns the same thing as `base_Variable.evaluate` (but more efficiently). base_variable_data : list of :numpy:array @@ -45,7 +45,6 @@ def __init__( base_variables_casadi, base_variables_data, solution, - warn=True, cumtrapz_ic=None, ): self.base_variables = base_variables @@ -60,7 +59,6 @@ def __init__( self.mesh = base_variables[0].mesh self.domain = base_variables[0].domain self.domains = base_variables[0].domains - self.warn = warn self.cumtrapz_ic = cumtrapz_ic # Sensitivity starts off uninitialized, only set when called @@ -82,33 +80,35 @@ def __init__( and isinstance(self.mesh, pybamm.ScikitSubMesh2D) ): self.initialise_2D_scikit_fem() + return # check variable shape - else: - if len(self.base_eval_shape) == 0 or self.base_eval_shape[0] == 1: - self.initialise_0D() - else: - n = self.mesh.npts - base_shape = self.base_eval_shape[0] - # Try some shapes that could make the variable a 1D variable - if base_shape in [n, n + 1]: - self.initialise_1D() - else: - # Try some shapes that could make the variable a 2D variable - first_dim_nodes = self.mesh.nodes - first_dim_edges = self.mesh.edges - second_dim_pts = self.base_variables[0].secondary_mesh.nodes - if self.base_eval_size // len(second_dim_pts) in [ - len(first_dim_nodes), - len(first_dim_edges), - ]: - self.initialise_2D() - else: - # Raise error for 3D variable - raise NotImplementedError( - f"Shape not recognized for {base_variables[0]} " - + "(note processing of 3D variables is not yet implemented)" - ) + if len(self.base_eval_shape) == 0 or self.base_eval_shape[0] == 1: + self.initialise_0D() + return + + n = self.mesh.npts + base_shape = self.base_eval_shape[0] + # Try some shapes that could make the variable a 1D variable + if base_shape in [n, n + 1]: + self.initialise_1D() + return + + # Try some shapes that could make the variable a 2D variable + first_dim_nodes = self.mesh.nodes + first_dim_edges = self.mesh.edges + second_dim_pts = self.base_variables[0].secondary_mesh.nodes + if self.base_eval_size // len(second_dim_pts) not in [ + len(first_dim_nodes), + len(first_dim_edges), + ]: + # Raise error for 3D variable + raise NotImplementedError( + f"Shape not recognized for {base_variables[0]} " + + "(note processing of 3D variables is not yet implemented)" + ) + + self.initialise_2D() def add_sensitivity(self, param, data): # unroll from sparse representation into n-d matrix @@ -203,7 +203,7 @@ def initialise_0D(self): self.entries = entries self.dimensions = 0 - def initialise_1D(self, fixed_t=False): + def initialise_1D(self): entries = self.unroll_1D() # Get node and edge values @@ -422,7 +422,7 @@ def initialise_2D_scikit_fem(self): coords={"y": y_sol, "z": z_sol, "t": self.t_pts}, ) - def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None, warn=True): + def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None): """ Evaluate the variable at arbitrary *dimensional* t (and x, r, y, z and/or R), using interpolation diff --git a/src/pybamm/solvers/solution.py b/src/pybamm/solvers/solution.py index 74d9ce7baf..c884e79e34 100644 --- a/src/pybamm/solvers/solution.py +++ b/src/pybamm/solvers/solution.py @@ -75,6 +75,7 @@ def __init__( y_event=None, termination="final time", all_sensitivities=False, + all_yps=None, check_solution=True, ): if not isinstance(all_ts, list): @@ -88,6 +89,10 @@ def __init__( self._all_ys_and_sens = all_ys self._all_models = all_models + if (all_yps is not None) and not isinstance(all_yps, list): + all_yps = [all_yps] + self._all_yps = all_yps + # Set up inputs if not isinstance(all_inputs, list): all_inputs_copy = dict(all_inputs) @@ -128,7 +133,7 @@ def __init__( # initialize empty variables and data self._variables = pybamm.FuzzyDict() - self.data = pybamm.FuzzyDict() + self._data = pybamm.FuzzyDict() # Add self as sub-solution for compatibility with ProcessedVariable self._sub_solutions = [self] @@ -298,6 +303,13 @@ def y(self): return self._y + @property + def data(self): + for k, v in self._variables.items(): + if k not in self._data: + self._data[k] = v.data + return self._data + @property def sensitivities(self): """Values of the sensitivities. Returns a dict of param_name: np_array""" @@ -400,6 +412,14 @@ def all_models(self): def all_inputs_casadi(self): return [casadi.vertcat(*inp.values()) for inp in self.all_inputs] + @property + def all_yps(self): + return self._all_yps + + @property + def hermite_interpolation(self): + return self.all_yps is not None + @property def t_event(self): """Time at which the event happens""" @@ -434,6 +454,12 @@ def first_state(self): n_states = self.all_models[0].len_rhs_and_alg for key in self._all_sensitivities: sensitivities[key] = self._all_sensitivities[key][0][-n_states:, :] + + if self.all_yps is None: + all_yps = None + else: + all_yps = self.all_yps[0][:, :1] + new_sol = Solution( self.all_ts[0][:1], self.all_ys[0][:, :1], @@ -443,6 +469,7 @@ def first_state(self): None, "final time", all_sensitivities=sensitivities, + all_yps=all_yps, ) new_sol._all_inputs_casadi = self.all_inputs_casadi[:1] new_sol._sub_solutions = self.sub_solutions[:1] @@ -467,6 +494,12 @@ def last_state(self): n_states = self.all_models[-1].len_rhs_and_alg for key in self._all_sensitivities: sensitivities[key] = self._all_sensitivities[key][-1][-n_states:, :] + + if self.all_yps is None: + all_yps = None + else: + all_yps = self.all_yps[-1][:, -1:] + new_sol = Solution( self.all_ts[-1][-1:], self.all_ys[-1][:, -1:], @@ -476,6 +509,7 @@ def last_state(self): self.y_event, self.termination, all_sensitivities=sensitivities, + all_yps=all_yps, ) new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:] new_sol._sub_solutions = self.sub_solutions[-1:] @@ -528,57 +562,61 @@ def update(self, variables): if isinstance(self._all_sensitivities, bool) and self._all_sensitivities: self.extract_explicit_sensitivities() - # Convert single entry to list + # Single variable if isinstance(variables, str): variables = [variables] + # Process - for key in variables: - cumtrapz_ic = None - pybamm.logger.debug(f"Post-processing {key}") - vars_pybamm = [model.variables_and_events[key] for model in self.all_models] - - # Iterate through all models, some may be in the list several times and - # therefore only get set up once - vars_casadi = [] - for i, (model, ys, inputs, var_pybamm) in enumerate( - zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm) + for variable in variables: + self._update_variable(variable) + + def _update_variable(self, variable): + cumtrapz_ic = None + pybamm.logger.debug(f"Post-processing {variable}") + vars_pybamm = [ + model.variables_and_events[variable] for model in self.all_models + ] + + # Iterate through all models, some may be in the list several times and + # therefore only get set up once + vars_casadi = [] + for i, (model, ys, inputs, var_pybamm) in enumerate( + zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm) + ): + if ys.size == 0 and var_pybamm.has_symbol_of_classes( + pybamm.expression_tree.state_vector.StateVector ): - if ys.size == 0 and var_pybamm.has_symbol_of_classes( - pybamm.expression_tree.state_vector.StateVector - ): - raise KeyError( - f"Cannot process variable '{key}' as it was not part of the " - "solve. Please re-run the solve with `output_variables` set to " - "include this variable." - ) - elif isinstance(var_pybamm, pybamm.ExplicitTimeIntegral): - cumtrapz_ic = var_pybamm.initial_condition - cumtrapz_ic = cumtrapz_ic.evaluate() - var_pybamm = var_pybamm.child - var_casadi = self.process_casadi_var( - var_pybamm, - inputs, - ys.shape, - ) - model._variables_casadi[key] = var_casadi - vars_pybamm[i] = var_pybamm - elif key in model._variables_casadi: - var_casadi = model._variables_casadi[key] - else: - var_casadi = self.process_casadi_var( - var_pybamm, - inputs, - ys.shape, - ) - model._variables_casadi[key] = var_casadi - vars_casadi.append(var_casadi) - var = pybamm.ProcessedVariable( - vars_pybamm, vars_casadi, self, cumtrapz_ic=cumtrapz_ic - ) + raise KeyError( + f"Cannot process variable '{variable}' as it was not part of the " + "solve. Please re-run the solve with `output_variables` set to " + "include this variable." + ) + elif isinstance(var_pybamm, pybamm.ExplicitTimeIntegral): + cumtrapz_ic = var_pybamm.initial_condition + cumtrapz_ic = cumtrapz_ic.evaluate() + var_pybamm = var_pybamm.child + var_casadi = self.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) + model._variables_casadi[variable] = var_casadi + vars_pybamm[i] = var_pybamm + elif variable in model._variables_casadi: + var_casadi = model._variables_casadi[variable] + else: + var_casadi = self.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) + model._variables_casadi[variable] = var_casadi + vars_casadi.append(var_casadi) + var = pybamm.process_variable( + vars_pybamm, vars_casadi, self, cumtrapz_ic=cumtrapz_ic + ) - # Save variable and data - self._variables[key] = var - self.data[key] = var.data + self._variables[variable] = var def process_casadi_var(self, var_pybamm, inputs, ys_shape): t_MX = casadi.MX.sym("t") @@ -588,8 +626,40 @@ def process_casadi_var(self, var_pybamm, inputs, ys_shape): } inputs_MX = casadi.vertcat(*[p for p in inputs_MX_dict.values()]) var_sym = var_pybamm.to_casadi(t_MX, y_MX, inputs=inputs_MX_dict) - var_casadi = casadi.Function("variable", [t_MX, y_MX, inputs_MX], [var_sym]) - return var_casadi + + opts = { + "cse": True, + "inputs_check": False, + "is_diff_in": [False, False, False], + "is_diff_out": [False], + "regularity_check": False, + "error_on_fail": False, + "enable_jacobian": False, + } + + # Casadi has a bug where it does not correctly handle arrays with + # zeros padded at the beginning or end. To avoid this, we add and + # subtract the same number to the variable to reinforce the + # variable bounds. This does not affect the answer + epsilon = 1.0 + var_sym = (var_sym - epsilon) + epsilon + + var_casadi = casadi.Function( + "variable", + [t_MX, y_MX, inputs_MX], + [var_sym], + opts, + ) + + # Some variables, like interpolants, cannot be expanded + try: + var_casadi_out = var_casadi.expand() + except RuntimeError as error: + if "'eval_sx' not defined for" not in str(error): + raise error # pragma: no cover + var_casadi_out = var_casadi + + return var_casadi_out def __getitem__(self, key): """Read a variable from the solution. Variables are created 'just in time', i.e. @@ -818,13 +888,23 @@ def __add__(self, other): return new_sol # Update list of sub-solutions + hermite_interpolation = ( + other.hermite_interpolation and self.hermite_interpolation + ) if other.all_ts[0][0] == self.all_ts[-1][-1]: # Skip first time step if it is repeated all_ts = self.all_ts + [other.all_ts[0][1:]] + other.all_ts[1:] all_ys = self.all_ys + [other.all_ys[0][:, 1:]] + other.all_ys[1:] + if hermite_interpolation: + all_yps = self.all_yps + [other.all_yps[0][:, 1:]] + other.all_yps[1:] else: all_ts = self.all_ts + other.all_ts all_ys = self.all_ys + other.all_ys + if hermite_interpolation: + all_yps = self.all_yps + other.all_yps + + if not hermite_interpolation: + all_yps = None # sensitivities can be: # - bool if not using sensitivities or using explicit sensitivities which still @@ -859,6 +939,7 @@ def __add__(self, other): other.y_event, other.termination, all_sensitivities=all_sensitivities, + all_yps=all_yps, ) new_sol.closest_event_idx = other.closest_event_idx @@ -891,6 +972,7 @@ def copy(self): self.y_event, self.termination, self._all_sensitivities, + self.all_yps, ) new_sol._all_inputs_casadi = self.all_inputs_casadi new_sol._sub_solutions = self.sub_solutions @@ -1001,6 +1083,7 @@ def make_cycle_solution( sum_sols.y_event, sum_sols.termination, sum_sols._all_sensitivities, + sum_sols.all_yps, ) cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi cycle_solution._sub_solutions = sum_sols.sub_solutions diff --git a/tests/integration/test_models/standard_output_comparison.py b/tests/integration/test_models/standard_output_comparison.py index 4d4d16e5ca..6c56894314 100644 --- a/tests/integration/test_models/standard_output_comparison.py +++ b/tests/integration/test_models/standard_output_comparison.py @@ -66,6 +66,7 @@ def compare(self, var, atol=0, rtol=0.02): # Get variable for each model model_variables = [solution[var] for solution in self.solutions] var0 = model_variables[0] + var0.initialise() spatial_pts = {} if var0.dimensions >= 1: diff --git a/tests/unit/test_experiments/test_simulation_with_experiment.py b/tests/unit/test_experiments/test_simulation_with_experiment.py index 394d64b257..b455e72393 100644 --- a/tests/unit/test_experiments/test_simulation_with_experiment.py +++ b/tests/unit/test_experiments/test_simulation_with_experiment.py @@ -83,8 +83,12 @@ def test_run_experiment(self): assert len(sol.cycles) == 1 # Test outputs - np.testing.assert_array_equal(sol.cycles[0].steps[0]["C-rate"].data, 1 / 20) - np.testing.assert_array_equal(sol.cycles[0].steps[1]["Current [A]"].data, -1) + np.testing.assert_array_almost_equal( + sol.cycles[0].steps[0]["C-rate"].data, 1 / 20 + ) + np.testing.assert_array_almost_equal( + sol.cycles[0].steps[1]["Current [A]"].data, -1 + ) np.testing.assert_array_almost_equal( sol.cycles[0].steps[2]["Voltage [V]"].data, 4.1, decimal=5 ) diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index b049729ae3..39918d73a4 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -43,9 +43,8 @@ def test_ida_roberts_klu(self): ) # Test - t_eval = np.linspace(0, 3, 100) - t_interp = t_eval - solution = solver.solve(model, t_eval, t_interp=t_interp) + t_eval = [0, 3] + solution = solver.solve(model, t_eval) # test that final time is time of event # y = 0.1 t + y0 so y=0.2 when t=2 @@ -311,8 +310,7 @@ def test_sensitivities_initial_condition(self): options={"jax_evaluator": "iree"} if form == "iree" else {}, ) - t_interp = np.linspace(0, 3, 100) - t_eval = [t_interp[0], t_interp[-1]] + t_eval = [0, 3] a_value = 0.1 @@ -321,7 +319,6 @@ def test_sensitivities_initial_condition(self): t_eval, inputs={"a": a_value}, calculate_sensitivities=True, - t_interp=t_interp, ) np.testing.assert_array_almost_equal( @@ -638,7 +635,7 @@ def test_failures(self): solver = pybamm.IDAKLUSolver() t_eval = [0, 3] - with pytest.raises(pybamm.SolverError, match="FAILURE IDA"): + with pytest.raises(ValueError, match="std::exception"): solver.solve(model, t_eval) def test_dae_solver_algebraic_model(self): @@ -944,7 +941,7 @@ def construct_model(): # Mock a 1D current collector and initialise (none in the model) sol["x_s [m]"].domain = ["current collector"] - sol["x_s [m]"].initialise_1D() + sol["x_s [m]"].entries def test_with_output_variables_and_sensitivities(self): # Construct a model and solve for all variables, then test @@ -1040,7 +1037,7 @@ def test_with_output_variables_and_sensitivities(self): # Mock a 1D current collector and initialise (none in the model) sol["x_s [m]"].domain = ["current collector"] - sol["x_s [m]"].initialise_1D() + sol["x_s [m]"].entries def test_bad_jax_evaluator(self): model = pybamm.lithium_ion.DFN() @@ -1108,19 +1105,40 @@ def test_simulation_period(self): def test_interpolate_time_step_start_offset(self): model = pybamm.lithium_ion.SPM() - experiment = pybamm.Experiment( - [ - "Discharge at C/10 for 10 seconds", - "Charge at C/10 for 10 seconds", - ], - period="1 seconds", - ) + + def experiment_setup(period=None): + return pybamm.Experiment( + [ + "Discharge at C/10 for 10 seconds", + "Charge at C/10 for 10 seconds", + ], + period=period, + ) + + experiment_1s = experiment_setup(period="1 seconds") solver = pybamm.IDAKLUSolver() - sim = pybamm.Simulation(model, experiment=experiment, solver=solver) - sol = sim.solve() + sim_1s = pybamm.Simulation(model, experiment=experiment_1s, solver=solver) + sol_1s = sim_1s.solve() np.testing.assert_equal( - np.nextafter(sol.sub_solutions[0].t[-1], np.inf), - sol.sub_solutions[1].t[0], + np.nextafter(sol_1s.sub_solutions[0].t[-1], np.inf), + sol_1s.sub_solutions[1].t[0], + ) + + assert not sol_1s.hermite_interpolation + + experiment = experiment_setup(period=None) + sim = pybamm.Simulation(model, experiment=experiment, solver=solver) + sol = sim.solve(model) + + assert sol.hermite_interpolation + + rtol = solver.rtol + atol = solver.atol + np.testing.assert_allclose( + sol_1s["Voltage [V]"].data, + sol["Voltage [V]"](sol_1s.t), + rtol=rtol, + atol=atol, ) def test_python_idaklu_deprecation_errors(self): diff --git a/tests/unit/test_solvers/test_processed_variable.py b/tests/unit/test_solvers/test_processed_variable.py index 6cd456347d..04de88963d 100644 --- a/tests/unit/test_solvers/test_processed_variable.py +++ b/tests/unit/test_solvers/test_processed_variable.py @@ -8,6 +8,13 @@ import numpy as np import pytest +from scipy.interpolate import CubicHermiteSpline + + +if pybamm.has_idaklu(): + _hermite_args = [True, False] +else: + _hermite_args = [False] def to_casadi(var_pybamm, y, inputs=None): @@ -27,55 +34,91 @@ def to_casadi(var_pybamm, y, inputs=None): return var_casadi -def process_and_check_2D_variable( - var, first_spatial_var, second_spatial_var, disc=None, geometry_options=None -): - # first_spatial_var should be on the "smaller" domain, i.e "r" for an "r-x" variable - if geometry_options is None: - geometry_options = {} - if disc is None: - disc = tests.get_discretisation_for_testing() - disc.set_variable_slices([var]) - - first_sol = disc.process_symbol(first_spatial_var).entries[:, 0] - second_sol = disc.process_symbol(second_spatial_var).entries[:, 0] - - # Keep only the first iteration of entries - first_sol = first_sol[: len(first_sol) // len(second_sol)] - var_sol = disc.process_symbol(var) - t_sol = np.linspace(0, 1) - y_sol = np.ones(len(second_sol) * len(first_sol))[:, np.newaxis] * np.linspace(0, 5) - - var_casadi = to_casadi(var_sol, y_sol) - model = tests.get_base_model_with_battery_geometry(**geometry_options) - processed_var = pybamm.ProcessedVariable( - [var_sol], - [var_casadi], - pybamm.Solution(t_sol, y_sol, model, {}), - warn=False, - ) - np.testing.assert_array_equal( - processed_var.entries, - np.reshape(y_sol, [len(first_sol), len(second_sol), len(t_sol)]), - ) - return y_sol, first_sol, second_sol, t_sol +class TestProcessedVariable: + @staticmethod + def _get_yps(y, hermite_interp, values=1): + if hermite_interp: + yp_sol = values * np.ones_like(y) + else: + yp_sol = None + return yp_sol + + @staticmethod + def _sol_default(t_sol, y_sol, yp_sol=None, model=None, inputs=None): + if inputs is None: + inputs = {} + if model is None: + model = tests.get_base_model_with_battery_geometry() + return pybamm.Solution( + t_sol, + y_sol, + model, + inputs, + all_yps=yp_sol, + ) + + def _process_and_check_2D_variable( + self, + var, + first_spatial_var, + second_spatial_var, + disc=None, + geometry_options=None, + hermite_interp=False, + ): + # first_spatial_var should be on the "smaller" domain, i.e "r" for an "r-x" variable + if geometry_options is None: + geometry_options = {} + if disc is None: + disc = tests.get_discretisation_for_testing() + disc.set_variable_slices([var]) + first_sol = disc.process_symbol(first_spatial_var).entries[:, 0] + second_sol = disc.process_symbol(second_spatial_var).entries[:, 0] -class TestProcessedVariable: - def test_processed_variable_0D(self): + # Keep only the first iteration of entries + first_sol = first_sol[: len(first_sol) // len(second_sol)] + var_sol = disc.process_symbol(var) + t_sol = np.linspace(0, 1) + y_sol = 5 * t_sol * np.ones(len(second_sol) * len(first_sol))[:, np.newaxis] + yp_sol = self._get_yps(y_sol, hermite_interp, values=5) + + var_casadi = to_casadi(var_sol, y_sol) + model = tests.get_base_model_with_battery_geometry(**geometry_options) + processed_var = pybamm.process_variable( + [var_sol], + [var_casadi], + self._sol_default(t_sol, y_sol, yp_sol, model), + ) + np.testing.assert_array_equal( + processed_var.entries, + np.reshape(y_sol, [len(first_sol), len(second_sol), len(t_sol)]), + ) + + # check that C++ and Python give the same result + if pybamm.has_idaklu(): + np.testing.assert_array_equal( + processed_var._observe_raw_cpp(), processed_var._observe_raw_python() + ) + + return y_sol, first_sol, second_sol, t_sol, yp_sol + + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_0D(self, hermite_interp): # without space t = pybamm.t y = pybamm.StateVector(slice(0, 1)) var = t * y var.mesh = None + model = pybamm.BaseModel() t_sol = np.linspace(0, 1) y_sol = np.array([np.linspace(0, 5)]) + yp_sol = self._get_yps(y_sol, hermite_interp) var_casadi = to_casadi(var, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var], [var_casadi], - pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol, model), ) np.testing.assert_array_equal(processed_var.entries, t_sol * y_sol[0]) @@ -84,18 +127,41 @@ def test_processed_variable_0D(self): var.mesh = None t_sol = np.array([0]) y_sol = np.array([1])[:, np.newaxis] + yp_sol = np.array([1])[:, np.newaxis] + sol = self._sol_default(t_sol, y_sol, yp_sol, model) var_casadi = to_casadi(var, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var], [var_casadi], - pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, + sol, ) np.testing.assert_array_equal(processed_var.entries, y_sol[0]) - # check empty sensitivity works + # check that repeated calls return the same data + data1 = processed_var.data + + assert processed_var.entries_raw_initialized + + data2 = processed_var.data + + np.testing.assert_array_equal(data1, data2) + + data_t1 = processed_var(sol.t) + + assert processed_var.xr_array_raw_initialized - def test_processed_variable_0D_no_sensitivity(self): + data_t2 = processed_var(sol.t) + + np.testing.assert_array_equal(data_t1, data_t2) + + # check that C++ and Python give the same result + if pybamm.has_idaklu(): + np.testing.assert_array_equal( + processed_var._observe_raw_cpp(), processed_var._observe_raw_python() + ) + + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_0D_no_sensitivity(self, hermite_interp): # without space t = pybamm.t y = pybamm.StateVector(slice(0, 1)) @@ -103,12 +169,12 @@ def test_processed_variable_0D_no_sensitivity(self): var.mesh = None t_sol = np.linspace(0, 1) y_sol = np.array([np.linspace(0, 5)]) + yp_sol = self._get_yps(y_sol, hermite_interp) var_casadi = to_casadi(var, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var], [var_casadi], - pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol, pybamm.BaseModel()), ) # test no inputs (i.e. no sensitivity) @@ -124,18 +190,18 @@ def test_processed_variable_0D_no_sensitivity(self): 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( + processed_var = pybamm.process_variable( [var], [var_casadi], pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), inputs), - warn=False, ) # test no sensitivity raises error with pytest.raises(ValueError, match="Cannot compute sensitivities"): print(processed_var.sensitivities) - def test_processed_variable_1D(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_1D(self, hermite_interp): t = pybamm.t var = pybamm.Variable("var", domain=["negative electrode", "separator"]) x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"]) @@ -149,26 +215,21 @@ def test_processed_variable_1D(self): eqn_sol = disc.process_symbol(eqn) t_sol = np.linspace(0, 1) y_sol = np.ones_like(x_sol)[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(y_sol, hermite_interp, values=5) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_equal(processed_var.entries, y_sol) np.testing.assert_array_almost_equal(processed_var(t_sol, x_sol), y_sol) eqn_casadi = to_casadi(eqn_sol, y_sol) - processed_eqn = pybamm.ProcessedVariable( + processed_eqn = pybamm.process_variable( [eqn_sol], [eqn_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_almost_equal( processed_eqn(t_sol, x_sol), t_sol * y_sol + x_sol[:, np.newaxis] @@ -184,13 +245,10 @@ def test_processed_variable_1D(self): x_s_edge = pybamm.Matrix(disc.mesh["separator"].edges, domain="separator") x_s_edge.mesh = disc.mesh["separator"] x_s_casadi = to_casadi(x_s_edge, y_sol) - processed_x_s_edge = pybamm.ProcessedVariable( + processed_x_s_edge = pybamm.process_variable( [x_s_edge], [x_s_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_equal( x_s_edge.entries[:, 0], processed_x_s_edge.entries[:, 0] @@ -201,20 +259,25 @@ def test_processed_variable_1D(self): eqn_sol = disc.process_symbol(eqn) t_sol = np.array([0]) y_sol = np.ones_like(x_sol)[:, np.newaxis] + yp_sol = self._get_yps(y_sol, hermite_interp, values=0) eqn_casadi = to_casadi(eqn_sol, y_sol) - processed_eqn2 = pybamm.ProcessedVariable( + processed_eqn2 = pybamm.process_variable( [eqn_sol], [eqn_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_equal( processed_eqn2.entries, y_sol + x_sol[:, np.newaxis] ) - def test_processed_variable_1D_unknown_domain(self): + # check that C++ and Python give the same result + if pybamm.has_idaklu(): + np.testing.assert_array_equal( + processed_eqn2._observe_raw_cpp(), processed_eqn2._observe_raw_python() + ) + + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_1D_unknown_domain(self, hermite_interp): x = pybamm.SpatialVariable("x", domain="SEI layer", coord_sys="cartesian") geometry = pybamm.Geometry( {"SEI layer": {x: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}}} @@ -227,6 +290,7 @@ def test_processed_variable_1D_unknown_domain(self): nt = 100 y_sol = np.zeros((var_pts[x], nt)) + yp_sol = self._get_yps(y_sol, hermite_interp) model = tests.get_base_model_with_battery_geometry() model._geometry = geometry solution = pybamm.Solution( @@ -237,14 +301,16 @@ def test_processed_variable_1D_unknown_domain(self): np.linspace(0, 1, 1), np.zeros(var_pts[x]), "test", + all_yps=yp_sol, ) c = pybamm.StateVector(slice(0, var_pts[x]), domain=["SEI layer"]) c.mesh = mesh["SEI layer"] c_casadi = to_casadi(c, y_sol) - pybamm.ProcessedVariable([c], [c_casadi], solution, warn=False) + pybamm.process_variable([c], [c_casadi], solution) - def test_processed_variable_2D_x_r(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_x_r(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative particle"], @@ -258,9 +324,12 @@ def test_processed_variable_2D_x_r(self): ) disc = tests.get_p2d_discretisation_for_testing() - process_and_check_2D_variable(var, r, x, disc=disc) + self._process_and_check_2D_variable( + var, r, x, disc=disc, hermite_interp=hermite_interp + ) - def test_processed_variable_2D_R_x(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_R_x(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative particle size"], @@ -274,15 +343,17 @@ def test_processed_variable_2D_R_x(self): x = pybamm.SpatialVariable("x", domain=["negative electrode"]) disc = tests.get_size_distribution_disc_for_testing() - process_and_check_2D_variable( + self._process_and_check_2D_variable( var, R, x, disc=disc, geometry_options={"options": {"particle size": "distribution"}}, + hermite_interp=hermite_interp, ) - def test_processed_variable_2D_R_z(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_R_z(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative particle size"], @@ -296,15 +367,17 @@ def test_processed_variable_2D_R_z(self): z = pybamm.SpatialVariable("z", domain=["current collector"]) disc = tests.get_size_distribution_disc_for_testing() - process_and_check_2D_variable( + self._process_and_check_2D_variable( var, R, z, disc=disc, geometry_options={"options": {"particle size": "distribution"}}, + hermite_interp=hermite_interp, ) - def test_processed_variable_2D_r_R(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_r_R(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative particle"], @@ -318,15 +391,17 @@ def test_processed_variable_2D_r_R(self): R = pybamm.SpatialVariable("R", domain=["negative particle size"]) disc = tests.get_size_distribution_disc_for_testing() - process_and_check_2D_variable( + self._process_and_check_2D_variable( var, r, R, disc=disc, geometry_options={"options": {"particle size": "distribution"}}, + hermite_interp=hermite_interp, ) - def test_processed_variable_2D_x_z(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_x_z(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative electrode", "separator"], @@ -340,7 +415,9 @@ def test_processed_variable_2D_x_z(self): z = pybamm.SpatialVariable("z", domain=["current collector"]) disc = tests.get_1p1d_discretisation_for_testing() - y_sol, x_sol, z_sol, t_sol = process_and_check_2D_variable(var, x, z, disc=disc) + y_sol, x_sol, z_sol, t_sol, yp_sol = self._process_and_check_2D_variable( + var, x, z, disc=disc, hermite_interp=hermite_interp + ) del x_sol # On edges @@ -352,19 +429,17 @@ def test_processed_variable_2D_x_z(self): x_s_edge.mesh = disc.mesh["separator"] x_s_edge.secondary_mesh = disc.mesh["current collector"] x_s_casadi = to_casadi(x_s_edge, y_sol) - processed_x_s_edge = pybamm.ProcessedVariable( + processed_x_s_edge = pybamm.process_variable( [x_s_edge], [x_s_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_equal( x_s_edge.entries.flatten(), processed_x_s_edge.entries[:, :, 0].T.flatten() ) - def test_processed_variable_2D_space_only(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_space_only(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative particle"], @@ -386,22 +461,21 @@ def test_processed_variable_2D_space_only(self): var_sol = disc.process_symbol(var) t_sol = np.array([0]) y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] + yp_sol = self._get_yps(y_sol, hermite_interp) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_equal( processed_var.entries, np.reshape(y_sol, [len(r_sol), len(x_sol), len(t_sol)]), ) - def test_processed_variable_2D_scikit(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_scikit(self, hermite_interp): var = pybamm.Variable("var", domain=["current collector"]) disc = tests.get_2p1d_discretisation_for_testing() @@ -412,21 +486,20 @@ def test_processed_variable_2D_scikit(self): var_sol.mesh = disc.mesh["current collector"] t_sol = np.linspace(0, 1) u_sol = np.ones(var_sol.shape[0])[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(u_sol, hermite_interp) var_casadi = to_casadi(var_sol, u_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, u_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, u_sol, yp_sol), ) np.testing.assert_array_equal( processed_var.entries, np.reshape(u_sol, [len(y), len(z), len(t_sol)]) ) - def test_processed_variable_2D_fixed_t_scikit(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_variable_2D_fixed_t_scikit(self, hermite_interp): var = pybamm.Variable("var", domain=["current collector"]) disc = tests.get_2p1d_discretisation_for_testing() @@ -437,22 +510,23 @@ def test_processed_variable_2D_fixed_t_scikit(self): var_sol.mesh = disc.mesh["current collector"] t_sol = np.array([0]) u_sol = np.ones(var_sol.shape[0])[:, np.newaxis] + yp_sol = self._get_yps(u_sol, hermite_interp) var_casadi = to_casadi(var_sol, u_sol) model = tests.get_base_model_with_battery_geometry( options={"dimensionality": 2} ) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution(t_sol, u_sol, model, {}), - warn=False, + pybamm.Solution(t_sol, u_sol, model, {}, all_yps=yp_sol), ) np.testing.assert_array_equal( processed_var.entries, np.reshape(u_sol, [len(y), len(z), len(t_sol)]) ) - def test_processed_var_0D_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_0D_interpolation(self, hermite_interp): # without spatial dependence t = pybamm.t y = pybamm.StateVector(slice(0, 1)) @@ -462,40 +536,39 @@ def test_processed_var_0D_interpolation(self): eqn.mesh = None t_sol = np.linspace(0, 1, 1000) - y_sol = np.array([np.linspace(0, 5, 1000)]) + y_sol = np.array([5 * t_sol]) + yp_sol = self._get_yps(y_sol, hermite_interp, values=5) var_casadi = to_casadi(var, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # vector np.testing.assert_array_equal(processed_var(t_sol), y_sol[0]) # scalar - np.testing.assert_array_equal(processed_var(0.5), 2.5) - np.testing.assert_array_equal(processed_var(0.7), 3.5) + np.testing.assert_array_almost_equal(processed_var(0.5), 2.5) + np.testing.assert_array_almost_equal(processed_var(0.7), 3.5) eqn_casadi = to_casadi(eqn, y_sol) - processed_eqn = pybamm.ProcessedVariable( + processed_eqn = pybamm.process_variable( [eqn], [eqn_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_equal(processed_eqn(t_sol), t_sol * y_sol[0]) - np.testing.assert_array_almost_equal(processed_eqn(0.5), 0.5 * 2.5) + assert processed_eqn(0.5).shape == () + + np.testing.assert_array_almost_equal(processed_eqn(0.5), 0.5 * 2.5) + np.testing.assert_array_equal(processed_eqn(2, fill_value=100), 100) # Suppress warning for this test pybamm.set_logging_level("ERROR") np.testing.assert_array_equal(processed_eqn(2), np.nan) pybamm.set_logging_level("WARNING") - def test_processed_var_0D_fixed_t_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_0D_fixed_t_interpolation(self, hermite_interp): y = pybamm.StateVector(slice(0, 1)) var = y eqn = 2 * y @@ -504,17 +577,18 @@ def test_processed_var_0D_fixed_t_interpolation(self): t_sol = np.array([10]) y_sol = np.array([[100]]) + yp_sol = self._get_yps(y_sol, hermite_interp) eqn_casadi = to_casadi(eqn, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [eqn], [eqn_casadi], - pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol, pybamm.BaseModel()), ) - np.testing.assert_array_equal(processed_var(), 200) + assert processed_var() == 200 - def test_processed_var_1D_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_1D_interpolation(self, hermite_interp): t = pybamm.t var = pybamm.Variable("var", domain=["negative electrode", "separator"]) x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"]) @@ -526,36 +600,33 @@ def test_processed_var_1D_interpolation(self): var_sol = disc.process_symbol(var) eqn_sol = disc.process_symbol(eqn) t_sol = np.linspace(0, 1) - y_sol = x_sol[:, np.newaxis] * np.linspace(0, 5) + y_sol = x_sol[:, np.newaxis] * (5 * t_sol) + yp_sol = self._get_yps(y_sol, hermite_interp, values=5) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) + # 2 vectors np.testing.assert_array_almost_equal(processed_var(t_sol, x_sol), y_sol) # 1 vector, 1 scalar np.testing.assert_array_almost_equal(processed_var(0.5, x_sol), 2.5 * x_sol) - np.testing.assert_array_equal( - processed_var(t_sol, x_sol[-1]), x_sol[-1] * np.linspace(0, 5) + np.testing.assert_array_almost_equal( + processed_var(t_sol, x_sol[-1]), + x_sol[-1] * np.linspace(0, 5), ) # 2 scalars np.testing.assert_array_almost_equal( processed_var(0.5, x_sol[-1]), 2.5 * x_sol[-1] ) eqn_casadi = to_casadi(eqn_sol, y_sol) - processed_eqn = pybamm.ProcessedVariable( + processed_eqn = pybamm.process_variable( [eqn_sol], [eqn_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # 2 vectors np.testing.assert_array_almost_equal( @@ -571,13 +642,10 @@ def test_processed_var_1D_interpolation(self): x_disc = disc.process_symbol(x) x_casadi = to_casadi(x_disc, y_sol) - processed_x = pybamm.ProcessedVariable( + processed_x = pybamm.process_variable( [x_disc], [x_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_almost_equal(processed_x(t=0, x=x_sol), x_sol) @@ -587,13 +655,10 @@ def test_processed_var_1D_interpolation(self): ) r_n.mesh = disc.mesh["negative particle"] r_n_casadi = to_casadi(r_n, y_sol) - processed_r_n = pybamm.ProcessedVariable( + processed_r_n = pybamm.process_variable( [r_n], [r_n_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) np.testing.assert_array_equal(r_n.entries[:, 0], processed_r_n.entries[:, 0]) r_test = np.linspace(0, 0.5) @@ -608,17 +673,17 @@ def test_processed_var_1D_interpolation(self): model = tests.get_base_model_with_battery_geometry( options={"particle size": "distribution"} ) - processed_R_n = pybamm.ProcessedVariable( + processed_R_n = pybamm.process_variable( [R_n], [R_n_casadi], pybamm.Solution(t_sol, y_sol, model, {}), - warn=False, ) np.testing.assert_array_equal(R_n.entries[:, 0], processed_R_n.entries[:, 0]) R_test = np.linspace(0, 1) np.testing.assert_array_almost_equal(processed_R_n(0, R=R_test), R_test) - def test_processed_var_1D_fixed_t_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_1D_fixed_t_interpolation(self, hermite_interp): var = pybamm.Variable("var", domain=["negative electrode", "separator"]) x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"]) eqn = var + x @@ -629,15 +694,13 @@ def test_processed_var_1D_fixed_t_interpolation(self): eqn_sol = disc.process_symbol(eqn) t_sol = np.array([1]) y_sol = x_sol[:, np.newaxis] + yp_sol = self._get_yps(y_sol, hermite_interp) eqn_casadi = to_casadi(eqn_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [eqn_sol], [eqn_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # vector @@ -647,7 +710,8 @@ def test_processed_var_1D_fixed_t_interpolation(self): # scalar np.testing.assert_array_almost_equal(processed_var(x=0.5), 1) - def test_processed_var_wrong_spatial_variable_names(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_wrong_spatial_variable_names(self, hermite_interp): var = pybamm.Variable( "var", domain=["domain A", "domain B"], @@ -677,6 +741,7 @@ def test_processed_var_wrong_spatial_variable_names(self): var_sol = disc.process_symbol(var) t_sol = np.linspace(0, 1) y_sol = np.ones(len(a_sol) * len(b_sol))[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(y_sol, hermite_interp) var_casadi = to_casadi(var_sol, y_sol) model = pybamm.BaseModel() @@ -687,14 +752,14 @@ def test_processed_var_wrong_spatial_variable_names(self): } ) with pytest.raises(NotImplementedError, match="Spatial variable name"): - pybamm.ProcessedVariable( + pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution(t_sol, y_sol, model, {}), - warn=False, - ) + pybamm.Solution(t_sol, y_sol, model, {}, all_yps=yp_sol), + ).initialise() - def test_processed_var_2D_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_2D_interpolation(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative particle"], @@ -716,15 +781,13 @@ def test_processed_var_2D_interpolation(self): var_sol = disc.process_symbol(var) t_sol = np.linspace(0, 1) y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(y_sol, hermite_interp) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # 3 vectors np.testing.assert_array_equal( @@ -768,20 +831,18 @@ def test_processed_var_2D_interpolation(self): y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] * np.linspace(0, 5) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # 3 vectors np.testing.assert_array_equal( processed_var(t_sol, x_sol, r_sol).shape, (10, 35, 50) ) - def test_processed_var_2D_fixed_t_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_2D_fixed_t_interpolation(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative particle"], @@ -803,15 +864,13 @@ def test_processed_var_2D_fixed_t_interpolation(self): var_sol = disc.process_symbol(var) t_sol = np.array([0]) y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] + yp_sol = self._get_yps(y_sol, hermite_interp) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # 2 vectors np.testing.assert_array_equal( @@ -823,7 +882,8 @@ def test_processed_var_2D_fixed_t_interpolation(self): # 2 scalars np.testing.assert_array_equal(processed_var(t=0, x=0.2, r=0.2).shape, ()) - def test_processed_var_2D_secondary_broadcast(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_2D_secondary_broadcast(self, hermite_interp): var = pybamm.Variable("var", domain=["negative particle"]) broad_var = pybamm.SecondaryBroadcast(var, "negative electrode") x = pybamm.SpatialVariable("x", domain=["negative electrode"]) @@ -836,15 +896,13 @@ def test_processed_var_2D_secondary_broadcast(self): var_sol = disc.process_symbol(broad_var) t_sol = np.linspace(0, 1) y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(y_sol, hermite_interp) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # 3 vectors np.testing.assert_array_equal( @@ -877,22 +935,21 @@ def test_processed_var_2D_secondary_broadcast(self): var_sol = disc.process_symbol(broad_var) t_sol = np.linspace(0, 1) y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(y_sol, hermite_interp, values=5) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, y_sol, yp_sol), ) # 3 vectors np.testing.assert_array_equal( processed_var(t_sol, x_sol, r_sol).shape, (10, 35, 50) ) - def test_processed_var_2_d_scikit_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_2_d_scikit_interpolation(self, hermite_interp): var = pybamm.Variable("var", domain=["current collector"]) disc = tests.get_2p1d_discretisation_for_testing() @@ -903,15 +960,13 @@ def test_processed_var_2_d_scikit_interpolation(self): var_sol.mesh = disc.mesh["current collector"] t_sol = np.linspace(0, 1) u_sol = np.ones(var_sol.shape[0])[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(u_sol, hermite_interp) var_casadi = to_casadi(var_sol, u_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, u_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, u_sol, yp_sol), ) # 3 vectors np.testing.assert_array_equal( @@ -938,7 +993,8 @@ def test_processed_var_2_d_scikit_interpolation(self): # 3 scalars np.testing.assert_array_equal(processed_var(0.2, y=0.2, z=0.2).shape, ()) - def test_processed_var_2D_fixed_t_scikit_interpolation(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_2D_fixed_t_scikit_interpolation(self, hermite_interp): var = pybamm.Variable("var", domain=["current collector"]) disc = tests.get_2p1d_discretisation_for_testing() @@ -949,15 +1005,13 @@ def test_processed_var_2D_fixed_t_scikit_interpolation(self): var_sol.mesh = disc.mesh["current collector"] t_sol = np.array([0]) u_sol = np.ones(var_sol.shape[0])[:, np.newaxis] + yp_sol = self._get_yps(u_sol, hermite_interp) var_casadi = to_casadi(var_sol, u_sol) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, u_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, u_sol, yp_sol), ) # 2 vectors np.testing.assert_array_equal( @@ -969,7 +1023,8 @@ def test_processed_var_2D_fixed_t_scikit_interpolation(self): # 2 scalars np.testing.assert_array_equal(processed_var(t=0, y=0.2, z=0.2).shape, ()) - def test_processed_var_2D_unknown_domain(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_processed_var_2D_unknown_domain(self, hermite_interp): var = pybamm.Variable( "var", domain=["domain B"], @@ -1007,6 +1062,7 @@ def test_processed_var_2D_unknown_domain(self): var_sol = disc.process_symbol(var) t_sol = np.linspace(0, 1) y_sol = np.ones(len(x_sol) * len(z_sol))[:, np.newaxis] * np.linspace(0, 5) + yp_sol = self._get_yps(y_sol, hermite_interp, values=5) var_casadi = to_casadi(var_sol, y_sol) model = pybamm.BaseModel() @@ -1016,11 +1072,10 @@ def test_processed_var_2D_unknown_domain(self): "domain B": {z: {"min": 0, "max": 1}}, } ) - processed_var = pybamm.ProcessedVariable( + processed_var = pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution(t_sol, y_sol, model, {}), - warn=False, + pybamm.Solution(t_sol, y_sol, model, {}, all_yps=yp_sol), ) # 3 vectors np.testing.assert_array_equal( @@ -1047,7 +1102,8 @@ def test_processed_var_2D_unknown_domain(self): # 3 scalars np.testing.assert_array_equal(processed_var(t=0.2, x=0.2, z=0.2).shape, ()) - def test_3D_raises_error(self): + @pytest.mark.parametrize("hermite_interp", _hermite_args) + def test_3D_raises_error(self, hermite_interp): var = pybamm.Variable( "var", domain=["negative electrode"], @@ -1059,16 +1115,14 @@ def test_3D_raises_error(self): var_sol = disc.process_symbol(var) t_sol = np.array([0, 1, 2]) u_sol = np.ones(var_sol.shape[0] * 3)[:, np.newaxis] + yp_sol = self._get_yps(u_sol, hermite_interp, values=0) var_casadi = to_casadi(var_sol, u_sol) with pytest.raises(NotImplementedError, match="Shape not recognized"): - pybamm.ProcessedVariable( + pybamm.process_variable( [var_sol], [var_casadi], - pybamm.Solution( - t_sol, u_sol, tests.get_base_model_with_battery_geometry(), {} - ), - warn=False, + self._sol_default(t_sol, u_sol, yp_sol), ) def test_process_spatial_variable_names(self): @@ -1080,11 +1134,10 @@ def test_process_spatial_variable_names(self): 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( + processed_var = pybamm.process_variable( [var], [var_casadi], pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, ) # Test empty list returns None @@ -1108,3 +1161,110 @@ def test_process_spatial_variable_names(self): # Test error raised if spatial variable name not recognised with pytest.raises(NotImplementedError, match="Spatial variable name"): processed_var._process_spatial_variable_names(["var1", "var2"]) + + def test_hermite_interpolator(self): + if not pybamm.has_idaklu(): + pytest.skip("Cannot test Hermite interpolation without IDAKLU") + + # initialise dummy solution to access method + def solution_setup(t_sol, sign): + y_sol = np.array([sign * np.sin(t_sol)]) + yp_sol = np.array([sign * np.cos(t_sol)]) + sol = self._sol_default(t_sol, y_sol, yp_sol) + return sol + + # without spatial dependence + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + var = y + eqn = t * y + var.mesh = None + eqn.mesh = None + + sign1 = +1 + t_sol1 = np.linspace(0, 1, 100) + sol1 = solution_setup(t_sol1, sign1) + + # Discontinuity in the solution + sign2 = -1 + t_sol2 = np.linspace(np.nextafter(t_sol1[-1], np.inf), t_sol1[-1] + 3, 99) + sol2 = solution_setup(t_sol2, sign2) + + sol = sol1 + sol2 + var_casadi = to_casadi(var, sol.all_ys[0]) + processed_var = pybamm.process_variable( + [var] * len(sol.all_ts), + [var_casadi] * len(sol.all_ts), + sol, + ) + + # Ground truth spline interpolants from scipy + spls = [ + CubicHermiteSpline(t, y, yp, axis=1) + for t, y, yp in zip(sol.all_ts, sol.all_ys, sol.all_yps) + ] + + def spl(t): + t = np.array(t) + out = np.zeros(len(t)) + for i, spl in enumerate(spls): + t0 = sol.all_ts[i][0] + tf = sol.all_ts[i][-1] + + mask = t >= t0 + # Extrapolation is allowed for the final solution + if i < len(spls) - 1: + mask &= t <= tf + + out[mask] = spl(t[mask]).flatten() + return out + + t0 = sol.t[0] + tf = sol.t[-1] + + # Test extrapolation before the first solution time + t_left_extrap = t0 - 1 + with pytest.raises( + ValueError, match="interpolation points must be greater than" + ): + processed_var(t_left_extrap) + + # Test extrapolation after the last solution time + t_right_extrap = [tf + 1] + np.testing.assert_almost_equal( + spl(t_right_extrap), + processed_var(t_right_extrap, fill_value="extrapolate"), + decimal=8, + ) + + t_dense = np.linspace(t0, tf + 1, 1000) + np.testing.assert_almost_equal( + spl(t_dense), + processed_var(t_dense, fill_value="extrapolate"), + decimal=8, + ) + + t_extended = np.union1d(sol.t, sol.t[-1] + 1) + np.testing.assert_almost_equal( + spl(t_extended), + processed_var(t_extended, fill_value="extrapolate"), + decimal=8, + ) + + ## Unsorted arrays + t_unsorted = np.array([0.5, 0.4, 0.6, 0, 1]) + idxs_sort = np.argsort(t_unsorted) + t_sorted = np.sort(t_unsorted) + + y_sorted = processed_var(t_sorted) + + idxs_unsort = np.zeros_like(idxs_sort) + idxs_unsort[idxs_sort] = np.arange(len(t_unsorted)) + + # Check that the unsorted and sorted arrays are the same + assert np.all(t_sorted == t_unsorted[idxs_sort]) + + y_unsorted = processed_var(t_unsorted) + + # Check that the unsorted and sorted arrays are the same + assert np.all(y_unsorted == y_sorted[idxs_unsort]) diff --git a/tests/unit/test_solvers/test_processed_variable_computed.py b/tests/unit/test_solvers/test_processed_variable_computed.py index 59a062b199..0fa46b4414 100644 --- a/tests/unit/test_solvers/test_processed_variable_computed.py +++ b/tests/unit/test_solvers/test_processed_variable_computed.py @@ -57,7 +57,6 @@ def process_and_check_2D_variable( [var_casadi], [y_sol], pybamm.Solution(t_sol, y_sol, model, {}), - warn=False, ) # NB: ProcessedVariableComputed does not interpret y in the same way as # ProcessedVariable; a better test of equivalence is to check that the @@ -82,7 +81,6 @@ def test_processed_variable_0D(self): [var_casadi], [y_sol], pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, ) # Assert that the processed variable is the same as the solution np.testing.assert_array_equal(processed_var.entries, y_sol[0]) @@ -94,7 +92,7 @@ def test_processed_variable_0D(self): # Check cumtrapz workflow produces no errors processed_var.cumtrapz_ic = 1 - processed_var.initialise_0D() + processed_var.entries # check empty sensitivity works def test_processed_variable_0D_no_sensitivity(self): @@ -111,7 +109,6 @@ def test_processed_variable_0D_no_sensitivity(self): [var_casadi], [y_sol], pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, ) # test no inputs (i.e. no sensitivity) @@ -132,7 +129,6 @@ def test_processed_variable_0D_no_sensitivity(self): [var_casadi], [y_sol], pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), inputs), - warn=False, ) # test no sensitivity raises error @@ -157,7 +153,6 @@ def test_processed_variable_1D(self): [var_casadi], [y_sol], sol, - warn=False, ) # Ordering from idaklu with output_variables set is different to @@ -175,7 +170,7 @@ def test_processed_variable_1D(self): processed_var.mesh.edges, processed_var.mesh.nodes, ) - processed_var.initialise_1D() + processed_var.entries processed_var.mesh.nodes, processed_var.mesh.edges = ( processed_var.mesh.edges, processed_var.mesh.nodes, @@ -192,7 +187,7 @@ def test_processed_variable_1D(self): ] for domain in domain_list: processed_var.domain[0] = domain - processed_var.initialise_1D() + processed_var.entries def test_processed_variable_1D_unknown_domain(self): x = pybamm.SpatialVariable("x", domain="SEI layer", coord_sys="cartesian") @@ -220,7 +215,7 @@ def test_processed_variable_1D_unknown_domain(self): c = pybamm.StateVector(slice(0, var_pts[x]), domain=["SEI layer"]) c.mesh = mesh["SEI layer"] c_casadi = to_casadi(c, y_sol) - pybamm.ProcessedVariableComputed([c], [c_casadi], [y_sol], solution, warn=False) + pybamm.ProcessedVariableComputed([c], [c_casadi], [y_sol], solution) def test_processed_variable_2D_x_r(self): var = pybamm.Variable( @@ -330,13 +325,12 @@ def test_processed_variable_2D_x_z(self): x_s_edge.mesh = disc.mesh["separator"] x_s_edge.secondary_mesh = disc.mesh["current collector"] x_s_casadi = to_casadi(x_s_edge, y_sol) - processed_x_s_edge = pybamm.ProcessedVariable( + processed_x_s_edge = pybamm.process_variable( [x_s_edge], [x_s_casadi], pybamm.Solution( t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} ), - warn=False, ) np.testing.assert_array_equal( x_s_edge.entries.flatten(), processed_x_s_edge.entries[:, :, 0].T.flatten() @@ -371,7 +365,6 @@ def test_processed_variable_2D_space_only(self): [var_casadi], [y_sol], pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), - warn=False, ) np.testing.assert_array_equal( processed_var.entries, @@ -408,7 +401,6 @@ def test_processed_variable_2D_fixed_t_scikit(self): [var_casadi], [u_sol], pybamm.Solution(t_sol, u_sol, pybamm.BaseModel(), {}), - warn=False, ) np.testing.assert_array_equal( processed_var.entries, np.reshape(u_sol, [len(y), len(z), len(t_sol)]) @@ -434,5 +426,4 @@ def test_3D_raises_error(self): [var_casadi], [u_sol], pybamm.Solution(t_sol, u_sol, pybamm.BaseModel(), {}), - warn=False, ) diff --git a/tests/unit/test_solvers/test_solution.py b/tests/unit/test_solvers/test_solution.py index 3aff012d5b..4ac9312531 100644 --- a/tests/unit/test_solvers/test_solution.py +++ b/tests/unit/test_solvers/test_solution.py @@ -77,14 +77,16 @@ def test_add_solutions(self): # Set up first solution t1 = np.linspace(0, 1) y1 = np.tile(t1, (20, 1)) - sol1 = pybamm.Solution(t1, y1, pybamm.BaseModel(), {"a": 1}) + yp1 = np.tile(t1, (30, 1)) + sol1 = pybamm.Solution(t1, y1, pybamm.BaseModel(), {"a": 1}, all_yps=yp1) sol1.solve_time = 1.5 sol1.integration_time = 0.3 # Set up second solution t2 = np.linspace(1, 2) y2 = np.tile(t2, (20, 1)) - sol2 = pybamm.Solution(t2, y2, pybamm.BaseModel(), {"a": 2}) + yp2 = np.tile(t1, (30, 1)) + sol2 = pybamm.Solution(t2, y2, pybamm.BaseModel(), {"a": 2}, all_yps=yp2) sol2.solve_time = 1 sol2.integration_time = 0.5