diff --git a/CHANGELOG.md b/CHANGELOG.md index decbacf529..97276a275b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # [Unreleased](https://github.com/pybamm-team/PyBaMM/) +## Features + +- Added additional user-configurable options to the (`IDAKLUSolver`) and adjusted the default values to improve performance. ([#4282](https://github.com/pybamm-team/PyBaMM/pull/4282)) + # [v24.5rc2](https://github.com/pybamm-team/PyBaMM/tree/v24.5rc2) - 2024-07-12 ## Features diff --git a/pybamm/solvers/c_solvers/idaklu/Expressions/Base/ExpressionSet.hpp b/pybamm/solvers/c_solvers/idaklu/Expressions/Base/ExpressionSet.hpp index a32f906a38..13c746a37d 100644 --- a/pybamm/solvers/c_solvers/idaklu/Expressions/Base/ExpressionSet.hpp +++ b/pybamm/solvers/c_solvers/idaklu/Expressions/Base/ExpressionSet.hpp @@ -31,7 +31,7 @@ class ExpressionSet const int n_s, const int n_e, const int n_p, - const Options& options) + const SetupOptions& options) : number_of_states(n_s), number_of_events(n_e), number_of_parameters(n_p), @@ -46,7 +46,7 @@ class ExpressionSet events(events), tmp_state_vector(number_of_states), tmp_sparse_jacobian_data(jac_times_cjmass_nnz), - options(options) + setup_opts(options) {}; int number_of_states; @@ -73,7 +73,7 @@ class ExpressionSet std::vector jac_times_cjmass_colptrs; // cppcheck-suppress unusedStructMember std::vector inputs; // cppcheck-suppress unusedStructMember - Options options; + SetupOptions setup_opts; virtual realtype *get_tmp_state_vector() = 0; virtual realtype *get_tmp_sparse_jacobian_data() = 0; diff --git a/pybamm/solvers/c_solvers/idaklu/Expressions/Casadi/CasadiFunctions.hpp b/pybamm/solvers/c_solvers/idaklu/Expressions/Casadi/CasadiFunctions.hpp index cc4b7cbb63..64db2e6106 100644 --- a/pybamm/solvers/c_solvers/idaklu/Expressions/Casadi/CasadiFunctions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/Expressions/Casadi/CasadiFunctions.hpp @@ -76,7 +76,7 @@ class CasadiFunctions : public ExpressionSet const std::vector& var_fcns, const std::vector& dvar_dy_fcns, const std::vector& dvar_dp_fcns, - const Options& options + const SetupOptions& setup_opts ) : rhs_alg_casadi(rhs_alg), jac_times_cjmass_casadi(jac_times_cjmass), @@ -98,7 +98,7 @@ class CasadiFunctions : public ExpressionSet static_cast(&sens_casadi), static_cast(&events_casadi), n_s, n_e, n_p, - options) + setup_opts) { // convert BaseFunctionType list to CasadiFunction list // NOTE: You must allocate ALL std::vector elements before taking references diff --git a/pybamm/solvers/c_solvers/idaklu/Expressions/IREE/IREEFunctions.hpp b/pybamm/solvers/c_solvers/idaklu/Expressions/IREE/IREEFunctions.hpp index cc864f4046..9a3169a46e 100644 --- a/pybamm/solvers/c_solvers/idaklu/Expressions/IREE/IREEFunctions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/Expressions/IREE/IREEFunctions.hpp @@ -59,7 +59,7 @@ class IREEFunctions : public ExpressionSet const std::vector& var_fcns, const std::vector& dvar_dy_fcns, const std::vector& dvar_dp_fcns, - const Options& options + const SetupOptions& setup_opts ) : iree_init_status(iree_init()), rhs_alg_iree(rhs_alg), @@ -82,7 +82,7 @@ class IREEFunctions : public ExpressionSet static_cast(&sens_iree), static_cast(&events_iree), n_s, n_e, n_p, - options) + setup_opts) { // convert BaseFunctionType list to IREEFunction list // NOTE: You must allocate ALL std::vector elements before taking references diff --git a/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp index 8c49069b30..98148a3c9f 100644 --- a/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp +++ b/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.hpp @@ -64,7 +64,8 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver std::vector res; std::vector res_dvar_dy; std::vector res_dvar_dp; - Options options; + SetupOptions setup_opts; + SolverOptions solver_opts; #if SUNDIALS_VERSION_MAJOR >= 6 SUNContext sunctx; @@ -84,7 +85,9 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver int jac_bandwidth_lower, int jac_bandwidth_upper, std::unique_ptr functions, - const Options& options); + const SetupOptions &setup_opts, + const SolverOptions &solver_opts + ); /** * @brief Destructor @@ -139,6 +142,16 @@ class IDAKLUSolverOpenMP : public IDAKLUSolver * @brief Allocate memory for matrices (noting appropriate matrix format/types) */ void SetMatrix(); + + /** + * @brief Apply user-configurable IDA options + */ + void SetSolverOptions(); + + /** + * @brief Check the return flag for errors + */ + void CheckErrors(int const & flag); }; #include "IDAKLUSolverOpenMP.inl" diff --git a/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl b/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl index 383037e2ca..340baa2d30 100644 --- a/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl +++ b/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl @@ -3,27 +3,29 @@ template IDAKLUSolverOpenMP::IDAKLUSolverOpenMP( - np_array atol_np, + np_array atol_np_input, double rel_tol, - np_array rhs_alg_id, - int number_of_parameters, - int number_of_events, - int jac_times_cjmass_nnz, - int jac_bandwidth_lower, - int jac_bandwidth_upper, + np_array rhs_alg_id_input, + int number_of_parameters_input, + int number_of_events_input, + int jac_times_cjmass_nnz_input, + int jac_bandwidth_lower_input, + int jac_bandwidth_upper_input, std::unique_ptr functions_arg, - const Options &options + const SetupOptions &setup_input, + const SolverOptions &solver_input ) : - atol_np(atol_np), - rhs_alg_id(rhs_alg_id), - number_of_states(atol_np.request().size), - number_of_parameters(number_of_parameters), - number_of_events(number_of_events), - jac_times_cjmass_nnz(jac_times_cjmass_nnz), - jac_bandwidth_lower(jac_bandwidth_lower), - jac_bandwidth_upper(jac_bandwidth_upper), + atol_np(atol_np_input), + rhs_alg_id(rhs_alg_id_input), + number_of_states(atol_np_input.request().size), + number_of_parameters(number_of_parameters_input), + number_of_events(number_of_events_input), + jac_times_cjmass_nnz(jac_times_cjmass_nnz_input), + jac_bandwidth_lower(jac_bandwidth_lower_input), + jac_bandwidth_upper(jac_bandwidth_upper_input), functions(std::move(functions_arg)), - options(options) + setup_opts(setup_input), + solver_opts(solver_input) { // Construction code moved to Initialize() which is called from the // (child) IDAKLUSolver_* class constructors. @@ -38,17 +40,17 @@ IDAKLUSolverOpenMP::IDAKLUSolverOpenMP( // create the vector of initial values AllocateVectors(); - if (number_of_parameters > 0) - { + if (number_of_parameters > 0) { yyS = N_VCloneVectorArray(number_of_parameters, yy); ypS = N_VCloneVectorArray(number_of_parameters, yp); } // set initial values realtype *atval = N_VGetArrayPointer(avtol); - for (int i = 0; i < number_of_states; i++) + for (int i = 0; i < number_of_states; i++) { atval[i] = atol[i]; - for (int is = 0; is < number_of_parameters; is++) - { + } + + for (int is = 0; is < number_of_parameters; is++) { N_VConst(RCONST(0.0), yyS[is]); N_VConst(RCONST(0.0), ypS[is]); } @@ -63,14 +65,16 @@ IDAKLUSolverOpenMP::IDAKLUSolverOpenMP( rtol = RCONST(rel_tol); IDASVtolerances(ida_mem, rtol, avtol); - // set events + // Set events IDARootInit(ida_mem, number_of_events, events_eval); + + // Set user data void *user_data = functions.get(); IDASetUserData(ida_mem, user_data); - // specify preconditioner type + // Specify preconditioner type precon_type = SUN_PREC_NONE; - if (options.preconditioner != "none") { + if (this->setup_opts.preconditioner != "none") { precon_type = SUN_PREC_LEFT; } } @@ -78,17 +82,83 @@ IDAKLUSolverOpenMP::IDAKLUSolverOpenMP( template void IDAKLUSolverOpenMP::AllocateVectors() { // Create vectors - yy = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); - yp = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); - avtol = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); - id = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); + yy = N_VNew_OpenMP(number_of_states, setup_opts.num_threads, sunctx); + yp = 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); } +template +void IDAKLUSolverOpenMP::SetSolverOptions() { + // Maximum order of the linear multistep method + CheckErrors(IDASetMaxOrd(ida_mem, solver_opts.max_order_bdf)); + + // Maximum number of steps to be taken by the solver in its attempt to reach + // the next output time + CheckErrors(IDASetMaxNumSteps(ida_mem, solver_opts.max_num_steps)); + + // Initial step size + CheckErrors(IDASetInitStep(ida_mem, solver_opts.dt_init)); + + // Maximum absolute step size + CheckErrors(IDASetMaxStep(ida_mem, solver_opts.dt_max)); + + // Maximum number of error test failures in attempting one step + CheckErrors(IDASetMaxErrTestFails(ida_mem, solver_opts.max_error_test_failures)); + + // Maximum number of nonlinear solver iterations at one step + CheckErrors(IDASetMaxNonlinIters(ida_mem, solver_opts.max_nonlinear_iterations)); + + // Maximum number of nonlinear solver convergence failures at one step + CheckErrors(IDASetMaxConvFails(ida_mem, solver_opts.max_convergence_failures)); + + // Safety factor in the nonlinear convergence test + CheckErrors(IDASetNonlinConvCoef(ida_mem, solver_opts.nonlinear_convergence_coefficient)); + + // Suppress algebraic variables from error test + CheckErrors(IDASetSuppressAlg(ida_mem, solver_opts.suppress_algebraic_error)); + + // Positive constant in the Newton iteration convergence test within the initial + // condition calculation + CheckErrors(IDASetNonlinConvCoefIC(ida_mem, solver_opts.nonlinear_convergence_coefficient_ic)); + + // Maximum number of steps allowed when icopt=IDA_YA_YDP_INIT in IDACalcIC + CheckErrors(IDASetMaxNumStepsIC(ida_mem, solver_opts.max_num_steps_ic)); + + // Maximum number of the approximate Jacobian or preconditioner evaluations + // allowed when the Newton iteration appears to be slowly converging + CheckErrors(IDASetMaxNumJacsIC(ida_mem, solver_opts.max_num_jacobians_ic)); + + // Maximum number of Newton iterations allowed in any one attempt to solve + // the initial conditions calculation problem + CheckErrors(IDASetMaxNumItersIC(ida_mem, solver_opts.max_num_iterations_ic)); + + // Maximum number of linesearch backtracks allowed in any Newton iteration, + // when solving the initial conditions calculation problem + CheckErrors(IDASetMaxBacksIC(ida_mem, solver_opts.max_linesearch_backtracks_ic)); + + // Turn off linesearch + CheckErrors(IDASetLineSearchOffIC(ida_mem, solver_opts.linesearch_off_ic)); + + // Ratio between linear and nonlinear tolerances + CheckErrors(IDASetEpsLin(ida_mem, solver_opts.epsilon_linear_tolerance)); + + // Increment factor used in DQ Jv approximation + CheckErrors(IDASetIncrementFactor(ida_mem, solver_opts.increment_factor)); + + int LS_type = SUNLinSolGetType(LS); + if (LS_type == SUNLINEARSOLVER_DIRECT || LS_type == SUNLINEARSOLVER_MATRIX_ITERATIVE) { + // Enable or disable linear solution scaling + CheckErrors(IDASetLinearSolutionScaling(ida_mem, solver_opts.linear_solution_scaling)); + } +} + + + template void IDAKLUSolverOpenMP::SetMatrix() { // Create Matrix object - if (options.jacobian == "sparse") - { + if (setup_opts.jacobian == "sparse") { DEBUG("\tsetting sparse matrix"); J = SUNSparseMatrix( number_of_states, @@ -97,8 +167,7 @@ void IDAKLUSolverOpenMP::SetMatrix() { CSC_MAT, sunctx ); - } - else if (options.jacobian == "banded") { + } else if (setup_opts.jacobian == "banded") { DEBUG("\tsetting banded matrix"); J = SUNBandMatrix( number_of_states, @@ -106,22 +175,19 @@ void IDAKLUSolverOpenMP::SetMatrix() { jac_bandwidth_lower, sunctx ); - } else if (options.jacobian == "dense" || options.jacobian == "none") - { + } else if (setup_opts.jacobian == "dense" || setup_opts.jacobian == "none") { DEBUG("\tsetting dense matrix"); J = SUNDenseMatrix( number_of_states, number_of_states, sunctx ); - } - else if (options.jacobian == "matrix-free") - { + } else if (setup_opts.jacobian == "matrix-free") { DEBUG("\tsetting matrix-free"); J = NULL; - } - else + } else { throw std::invalid_argument("Unsupported matrix requested"); + } } template @@ -129,61 +195,64 @@ void IDAKLUSolverOpenMP::Initialize() { // Call after setting the solver // attach the linear solver - if (LS == nullptr) + if (LS == nullptr) { throw std::invalid_argument("Linear solver not set"); - IDASetLinearSolver(ida_mem, LS, J); + } + CheckErrors(IDASetLinearSolver(ida_mem, LS, J)); - if (options.preconditioner != "none") - { + if (setup_opts.preconditioner != "none") { DEBUG("\tsetting IDADDB preconditioner"); // setup preconditioner - IDABBDPrecInit( - ida_mem, number_of_states, options.precon_half_bandwidth, - options.precon_half_bandwidth, options.precon_half_bandwidth_keep, - options.precon_half_bandwidth_keep, 0.0, residual_eval_approx, NULL); + CheckErrors(IDABBDPrecInit( + ida_mem, number_of_states, setup_opts.precon_half_bandwidth, + setup_opts.precon_half_bandwidth, setup_opts.precon_half_bandwidth_keep, + setup_opts.precon_half_bandwidth_keep, 0.0, residual_eval_approx, NULL)); } - if (options.jacobian == "matrix-free") - IDASetJacTimes(ida_mem, NULL, jtimes_eval); - else if (options.jacobian != "none") - IDASetJacFn(ida_mem, jacobian_eval); + if (setup_opts.jacobian == "matrix-free") { + CheckErrors(IDASetJacTimes(ida_mem, NULL, jtimes_eval)); + } else if (setup_opts.jacobian != "none") { + CheckErrors(IDASetJacFn(ida_mem, jacobian_eval)); + } - if (number_of_parameters > 0) - { - IDASensInit(ida_mem, number_of_parameters, IDA_SIMULTANEOUS, - sensitivities_eval, yyS, ypS); - IDASensEEtolerances(ida_mem); + if (number_of_parameters > 0) { + CheckErrors(IDASensInit(ida_mem, number_of_parameters, IDA_SIMULTANEOUS, + sensitivities_eval, yyS, ypS)); + CheckErrors(IDASensEEtolerances(ida_mem)); } - SUNLinSolInitialize(LS); + CheckErrors(SUNLinSolInitialize(LS)); auto id_np_val = rhs_alg_id.unchecked<1>(); realtype *id_val; id_val = N_VGetArrayPointer(id); int ii; - for (ii = 0; ii < number_of_states; ii++) + for (ii = 0; ii < number_of_states; ii++) { id_val[ii] = id_np_val[ii]; + } - IDASetId(ida_mem, id); + // Variable types: differential (1) and algebraic (0) + CheckErrors(IDASetId(ida_mem, id)); } template -IDAKLUSolverOpenMP::~IDAKLUSolverOpenMP() -{ +IDAKLUSolverOpenMP::~IDAKLUSolverOpenMP() { + bool sensitivity = number_of_parameters > 0; // Free memory - if (number_of_parameters > 0) - IDASensFree(ida_mem); + if (sensitivity) { + IDASensFree(ida_mem); + } + + CheckErrors(SUNLinSolFree(LS)); - SUNLinSolFree(LS); SUNMatDestroy(J); N_VDestroy(avtol); N_VDestroy(yy); N_VDestroy(yp); N_VDestroy(id); - if (number_of_parameters > 0) - { + if (sensitivity) { N_VDestroyVectorArray(yyS, number_of_parameters); N_VDestroyVectorArray(ypS, number_of_parameters); } @@ -209,8 +278,9 @@ void IDAKLUSolverOpenMP::CalcVars( for (auto& var_fcn : functions->var_fcns) { (*var_fcn)({tret, yval, functions->inputs.data()}, {&res[0]}); // store in return vector - for (size_t jj=0; jjnnz_out(); jj++) + for (size_t jj=0; jjnnz_out(); jj++) { y_return[t_i*length_of_return_vector + j++] = res[jj]; + } } // calculate sensitivities CalcVarsSensitivities(tret, yval, ySval, yS_return, ySk); @@ -235,15 +305,18 @@ void IDAKLUSolverOpenMP::CalcVarsSensitivities( (*dvar_dy)({tret, yval, functions->inputs.data()}, {&res_dvar_dy[0]}); // Calculate dvar/dp and convert to dense array for indexing (*dvar_dp)({tret, yval, functions->inputs.data()}, {&res_dvar_dp[0]}); - for(int k=0; knnz_out(); k++) + } + for (int k=0; knnz_out(); k++) { dens_dvar_dp[dvar_dp->get_row()[k]] = res_dvar_dp[k]; + } // Calculate sensitivities - for(int paramk=0; paramknnz_out(); spk++) + for (int spk=0; spknnz_out(); spk++) { yS_return[*ySk] += res_dvar_dy[spk] * ySval[paramk][dvar_dy->get_col()[spk]]; + } (*ySk)++; } } @@ -265,21 +338,25 @@ Solution IDAKLUSolverOpenMP::solve( auto y0 = y0_np.unchecked<1>(); auto yp0 = yp0_np.unchecked<1>(); auto n_coeffs = number_of_states + number_of_parameters * number_of_states; + bool const sensitivity = number_of_parameters > 0; - if (y0.size() != n_coeffs) + if (y0.size() != n_coeffs) { throw std::domain_error( "y0 has wrong size. Expected " + std::to_string(n_coeffs) + " but got " + std::to_string(y0.size())); + } - if (yp0.size() != n_coeffs) + if (yp0.size() != n_coeffs) { throw std::domain_error( "yp0 has wrong size. Expected " + std::to_string(n_coeffs) + " but got " + std::to_string(yp0.size())); + } // set inputs auto p_inputs = inputs.unchecked<2>(); - for (int i = 0; i < functions->inputs.size(); i++) + for (int i = 0; i < functions->inputs.size(); i++) { functions->inputs[i] = p_inputs(i, 0); + } // set initial conditions realtype *yval = N_VGetArrayPointer(yy); @@ -295,21 +372,29 @@ Solution IDAKLUSolverOpenMP::solve( } } - for (int i = 0; i < number_of_states; i++) - { + for (int i = 0; i < number_of_states; i++) { yval[i] = y0[i]; ypval[i] = yp0[i]; } - IDAReInit(ida_mem, t0, yy, yp); - if (number_of_parameters > 0) - IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS); + SetSolverOptions(); + + CheckErrors(IDAReInit(ida_mem, t0, yy, yp)); + if (sensitivity) { + CheckErrors(IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS)); + } // correct initial values - DEBUG("IDACalcIC"); - IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); - if (number_of_parameters > 0) - IDAGetSens(ida_mem, &t0, yyS); + int const init_type = solver_opts.init_all_y_ic ? IDA_Y_INIT : IDA_YA_YDP_INIT; + if (solver_opts.calc_ic) { + DEBUG("IDACalcIC"); + // IDACalcIC will throw a warning if it fails to find initial conditions + IDACalcIC(ida_mem, init_type, t(1)); + } + + if (sensitivity) { + CheckErrors(IDAGetSens(ida_mem, &t0, yyS)); + } realtype tret; realtype t_final = t(number_of_timesteps - 1); @@ -323,10 +408,12 @@ Solution IDAKLUSolverOpenMP::solve( for (auto& var_fcn : functions->var_fcns) { max_res_size = std::max(max_res_size, size_t(var_fcn->out_shape(0))); length_of_return_vector += var_fcn->nnz_out(); - for (auto& dvar_fcn : functions->dvar_dy_fcns) + for (auto& dvar_fcn : functions->dvar_dy_fcns) { max_res_dvar_dy = std::max(max_res_dvar_dy, size_t(dvar_fcn->out_shape(0))); - for (auto& dvar_fcn : functions->dvar_dp_fcns) + } + for (auto& dvar_fcn : functions->dvar_dp_fcns) { max_res_dvar_dp = std::max(max_res_dvar_dp, size_t(dvar_fcn->out_shape(0))); + } } } else { // Return full y state-vector @@ -375,63 +462,62 @@ Solution IDAKLUSolverOpenMP::solve( &tret, yval, ySval, yS_return, &ySk); } else { // Retain complete copy of the state vector - for (int j = 0; j < number_of_states; j++) + for (int j = 0; j < number_of_states; j++) { y_return[j] = yval[j]; - for (int j = 0; j < number_of_parameters; j++) - { + } + for (int j = 0; j < number_of_parameters; j++) { const int base_index = j * number_of_timesteps * number_of_states; - for (int k = 0; k < number_of_states; k++) + for (int k = 0; k < number_of_states; k++) { yS_return[base_index + k] = ySval[j][k]; + } } } // Subsequent states (t_i>0) int retval; t_i = 1; - while (true) - { + while (true) { realtype t_next = t(t_i); IDASetStopTime(ida_mem, t_next); DEBUG("IDASolve"); retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); - if (retval == IDA_TSTOP_RETURN || + if (!(retval == IDA_TSTOP_RETURN || retval == IDA_SUCCESS || - retval == IDA_ROOT_RETURN) - { - if (number_of_parameters > 0) - IDAGetSens(ida_mem, &tret, yyS); - - // Evaluate and store results for the time step - t_return[t_i] = tret; - if (functions->var_fcns.size() > 0) { - // Evaluate functions for each requested variable and store - // NOTE: Indexing of yS_return is (time:var:param) - CalcVars(y_return, length_of_return_vector, t_i, - &tret, yval, ySval, yS_return, &ySk); - } else { - // Retain complete copy of the state vector - for (int j = 0; j < number_of_states; j++) - y_return[t_i * number_of_states + j] = yval[j]; - for (int j = 0; j < number_of_parameters; j++) - { - const int base_index = - j * number_of_timesteps * number_of_states + - t_i * number_of_states; - for (int k = 0; k < number_of_states; k++) - // NOTE: Indexing of yS_return is (time:param:yvec) - yS_return[base_index + k] = ySval[j][k]; + retval == IDA_ROOT_RETURN)) { + // failed + break; + } + + if (sensitivity) { + CheckErrors(IDAGetSens(ida_mem, &tret, yyS)); + } + + // Evaluate and store results for the time step + t_return[t_i] = tret; + if (functions->var_fcns.size() > 0) { + // Evaluate functions for each requested variable and store + // NOTE: Indexing of yS_return is (time:var:param) + CalcVars(y_return, length_of_return_vector, t_i, + &tret, yval, ySval, yS_return, &ySk); + } else { + // Retain complete copy of the state vector + for (int j = 0; j < number_of_states; j++) { + y_return[t_i * number_of_states + j] = yval[j]; + } + for (int j = 0; j < number_of_parameters; j++) { + const int base_index = + j * number_of_timesteps * number_of_states + + t_i * number_of_states; + for (int k = 0; k < number_of_states; k++) { + // NOTE: Indexing of yS_return is (time:param:yvec) + yS_return[base_index + k] = ySval[j][k]; } } - t_i += 1; - - if (retval == IDA_SUCCESS || - retval == IDA_ROOT_RETURN) - break; } - else - { - // failed + t_i += 1; + + if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { break; } } @@ -473,13 +559,12 @@ Solution IDAKLUSolverOpenMP::solve( Solution sol(retval, t_ret, y_ret, yS_ret); - if (options.print_stats) - { + if (solver_opts.print_stats) { long nsteps, nrevals, nlinsetups, netfails; int klast, kcur; realtype hinused, hlast, hcur, tcur; - IDAGetIntegratorStats( + CheckErrors(IDAGetIntegratorStats( ida_mem, &nsteps, &nrevals, @@ -491,14 +576,15 @@ Solution IDAKLUSolverOpenMP::solve( &hlast, &hcur, &tcur - ); + )); long nniters, nncfails; - IDAGetNonlinSolvStats(ida_mem, &nniters, &nncfails); + CheckErrors(IDAGetNonlinSolvStats(ida_mem, &nniters, &nncfails)); long int ngevalsBBDP = 0; - if (options.using_iterative_solver) - IDABBDPrecGetNumGfnEvals(ida_mem, &ngevalsBBDP); + if (setup_opts.using_iterative_solver) { + CheckErrors(IDABBDPrecGetNumGfnEvals(ida_mem, &ngevalsBBDP)); + } py::print("Solver Stats:"); py::print("\tNumber of steps =", nsteps); @@ -519,3 +605,12 @@ Solution IDAKLUSolverOpenMP::solve( return sol; } + +template +void IDAKLUSolverOpenMP::CheckErrors(int const & flag) { + if (flag < 0) { + auto message = (std::string("IDA failed with flag ") + std::to_string(flag)).c_str(); + py::set_error(PyExc_ValueError, message); + throw py::error_already_set(); + } +} diff --git a/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP_solvers.hpp b/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP_solvers.hpp index ebeb543232..5f6f29b47b 100644 --- a/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP_solvers.hpp +++ b/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP_solvers.hpp @@ -61,7 +61,7 @@ class IDAKLUSolverOpenMP_SPBCGS : public IDAKLUSolverOpenMP { Base::LS = SUNLinSol_SPBCGS( Base::yy, Base::precon_type, - Base::options.linsol_max_iterations, + Base::setup_opts.linsol_max_iterations, Base::sunctx ); Base::Initialize(); @@ -81,7 +81,7 @@ class IDAKLUSolverOpenMP_SPFGMR : public IDAKLUSolverOpenMP { Base::LS = SUNLinSol_SPFGMR( Base::yy, Base::precon_type, - Base::options.linsol_max_iterations, + Base::setup_opts.linsol_max_iterations, Base::sunctx ); Base::Initialize(); @@ -101,7 +101,7 @@ class IDAKLUSolverOpenMP_SPGMR : public IDAKLUSolverOpenMP { Base::LS = SUNLinSol_SPGMR( Base::yy, Base::precon_type, - Base::options.linsol_max_iterations, + Base::setup_opts.linsol_max_iterations, Base::sunctx ); Base::Initialize(); @@ -121,7 +121,7 @@ class IDAKLUSolverOpenMP_SPTFQMR : public IDAKLUSolverOpenMP { Base::LS = SUNLinSol_SPTFQMR( Base::yy, Base::precon_type, - Base::options.linsol_max_iterations, + Base::setup_opts.linsol_max_iterations, Base::sunctx ); Base::Initialize(); diff --git a/pybamm/solvers/c_solvers/idaklu/Options.cpp b/pybamm/solvers/c_solvers/idaklu/Options.cpp index 684ab47f33..b6a33e016e 100644 --- a/pybamm/solvers/c_solvers/idaklu/Options.cpp +++ b/pybamm/solvers/c_solvers/idaklu/Options.cpp @@ -5,15 +5,14 @@ using namespace std::string_literals; -Options::Options(py::dict options) - : print_stats(options["print_stats"].cast()), - jacobian(options["jacobian"].cast()), - preconditioner(options["preconditioner"].cast()), - linsol_max_iterations(options["linsol_max_iterations"].cast()), - linear_solver(options["linear_solver"].cast()), - precon_half_bandwidth(options["precon_half_bandwidth"].cast()), - precon_half_bandwidth_keep(options["precon_half_bandwidth_keep"].cast()), - num_threads(options["num_threads"].cast()) +SetupOptions::SetupOptions(py::dict &py_opts) + : jacobian(py_opts["jacobian"].cast()), + preconditioner(py_opts["preconditioner"].cast()), + precon_half_bandwidth(py_opts["precon_half_bandwidth"].cast()), + precon_half_bandwidth_keep(py_opts["precon_half_bandwidth_keep"].cast()), + num_threads(py_opts["num_threads"].cast()), + linear_solver(py_opts["linear_solver"].cast()), + linsol_max_iterations(py_opts["linsol_max_iterations"].cast()) { using_sparse_matrix = true; @@ -119,3 +118,30 @@ Options::Options(py::dict options) preconditioner = "none"; } } + +SolverOptions::SolverOptions(py::dict &py_opts) + : print_stats(py_opts["print_stats"].cast()), + // IDA main solver + max_order_bdf(py_opts["max_order_bdf"].cast()), + max_num_steps(py_opts["max_num_steps"].cast()), + dt_init(RCONST(py_opts["dt_init"].cast())), + dt_max(RCONST(py_opts["dt_max"].cast())), + max_error_test_failures(py_opts["max_error_test_failures"].cast()), + max_nonlinear_iterations(py_opts["max_nonlinear_iterations"].cast()), + max_convergence_failures(py_opts["max_convergence_failures"].cast()), + 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()), + // IDA initial conditions calculation + calc_ic(py_opts["calc_ic"].cast()), + init_all_y_ic(py_opts["init_all_y_ic"].cast()), + max_num_steps_ic(py_opts["max_num_steps_ic"].cast()), + max_num_jacobians_ic(py_opts["max_num_jacobians_ic"].cast()), + max_num_iterations_ic(py_opts["max_num_iterations_ic"].cast()), + max_linesearch_backtracks_ic(py_opts["max_linesearch_backtracks_ic"].cast()), + linesearch_off_ic(py_opts["linesearch_off_ic"].cast()), + // IDALS linear solver interface + linear_solution_scaling(py_opts["linear_solution_scaling"].cast()), + epsilon_linear_tolerance(RCONST(py_opts["epsilon_linear_tolerance"].cast())), + increment_factor(RCONST(py_opts["increment_factor"].cast())) +{} diff --git a/pybamm/solvers/c_solvers/idaklu/Options.hpp b/pybamm/solvers/c_solvers/idaklu/Options.hpp index b70d0f4a30..66a175cfff 100644 --- a/pybamm/solvers/c_solvers/idaklu/Options.hpp +++ b/pybamm/solvers/c_solvers/idaklu/Options.hpp @@ -4,22 +4,52 @@ #include "common.hpp" /** - * @brief Options passed to the idaklu solver by pybamm + * @brief SetupOptions passed to the idaklu setup by pybamm */ -struct Options { - bool print_stats; +struct SetupOptions { bool using_sparse_matrix; bool using_banded_matrix; bool using_iterative_solver; std::string jacobian; - std::string linear_solver; // klu, lapack, spbcg std::string preconditioner; // spbcg - int linsol_max_iterations; int precon_half_bandwidth; int precon_half_bandwidth_keep; int num_threads; - explicit Options(py::dict options); + // IDALS linear solver interface + std::string linear_solver; // klu, lapack, spbcg + int linsol_max_iterations; + explicit SetupOptions(py::dict &py_opts); +}; +/** + * @brief SolverOptions passed to the idaklu solver by pybamm + */ +struct SolverOptions { + bool print_stats; + // IDA main solver + int max_order_bdf; + int max_num_steps; + double dt_init; + double dt_max; + int max_error_test_failures; + int max_nonlinear_iterations; + int max_convergence_failures; + double nonlinear_convergence_coefficient; + double nonlinear_convergence_coefficient_ic; + sunbooleantype suppress_algebraic_error; + // IDA initial conditions calculation + bool calc_ic; + bool init_all_y_ic; + int max_num_steps_ic; + int max_num_jacobians_ic; + int max_num_iterations_ic; + int max_linesearch_backtracks_ic; + sunbooleantype linesearch_off_ic; + // IDALS linear solver interface + sunbooleantype linear_solution_scaling; + double epsilon_linear_tolerance; + double increment_factor; + explicit SolverOptions(py::dict &py_opts); }; #endif // PYBAMM_OPTIONS_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/idaklu_solver.hpp b/pybamm/solvers/c_solvers/idaklu/idaklu_solver.hpp index a53b167ac4..ce1765aa82 100644 --- a/pybamm/solvers/c_solvers/idaklu/idaklu_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/idaklu_solver.hpp @@ -33,9 +33,10 @@ IDAKLUSolver *create_idaklu_solver( const std::vector& var_fcns, const std::vector& dvar_dy_fcns, const std::vector& dvar_dp_fcns, - py::dict options + py::dict py_opts ) { - auto options_cpp = Options(options); + auto setup_opts = SetupOptions(py_opts); + auto solver_opts = SolverOptions(py_opts); auto functions = std::make_unique( rhs_alg, jac_times_cjmass, @@ -55,13 +56,13 @@ IDAKLUSolver *create_idaklu_solver( var_fcns, dvar_dy_fcns, dvar_dp_fcns, - options_cpp + setup_opts ); IDAKLUSolver *idakluSolver = nullptr; // Instantiate solver class - if (options_cpp.linear_solver == "SUNLinSol_Dense") + if (setup_opts.linear_solver == "SUNLinSol_Dense") { DEBUG("\tsetting SUNLinSol_Dense linear solver"); idakluSolver = new IDAKLUSolverOpenMP_Dense( @@ -74,10 +75,11 @@ IDAKLUSolver *create_idaklu_solver( jac_bandwidth_lower, jac_bandwidth_upper, std::move(functions), - options_cpp + setup_opts, + solver_opts ); } - else if (options_cpp.linear_solver == "SUNLinSol_KLU") + else if (setup_opts.linear_solver == "SUNLinSol_KLU") { DEBUG("\tsetting SUNLinSol_KLU linear solver"); idakluSolver = new IDAKLUSolverOpenMP_KLU( @@ -90,10 +92,11 @@ IDAKLUSolver *create_idaklu_solver( jac_bandwidth_lower, jac_bandwidth_upper, std::move(functions), - options_cpp + setup_opts, + solver_opts ); } - else if (options_cpp.linear_solver == "SUNLinSol_Band") + else if (setup_opts.linear_solver == "SUNLinSol_Band") { DEBUG("\tsetting SUNLinSol_Band linear solver"); idakluSolver = new IDAKLUSolverOpenMP_Band( @@ -106,10 +109,11 @@ IDAKLUSolver *create_idaklu_solver( jac_bandwidth_lower, jac_bandwidth_upper, std::move(functions), - options_cpp + setup_opts, + solver_opts ); } - else if (options_cpp.linear_solver == "SUNLinSol_SPBCGS") + else if (setup_opts.linear_solver == "SUNLinSol_SPBCGS") { DEBUG("\tsetting SUNLinSol_SPBCGS_linear solver"); idakluSolver = new IDAKLUSolverOpenMP_SPBCGS( @@ -122,10 +126,11 @@ IDAKLUSolver *create_idaklu_solver( jac_bandwidth_lower, jac_bandwidth_upper, std::move(functions), - options_cpp + setup_opts, + solver_opts ); } - else if (options_cpp.linear_solver == "SUNLinSol_SPFGMR") + else if (setup_opts.linear_solver == "SUNLinSol_SPFGMR") { DEBUG("\tsetting SUNLinSol_SPFGMR_linear solver"); idakluSolver = new IDAKLUSolverOpenMP_SPFGMR( @@ -138,10 +143,11 @@ IDAKLUSolver *create_idaklu_solver( jac_bandwidth_lower, jac_bandwidth_upper, std::move(functions), - options_cpp + setup_opts, + solver_opts ); } - else if (options_cpp.linear_solver == "SUNLinSol_SPGMR") + else if (setup_opts.linear_solver == "SUNLinSol_SPGMR") { DEBUG("\tsetting SUNLinSol_SPGMR solver"); idakluSolver = new IDAKLUSolverOpenMP_SPGMR( @@ -154,10 +160,11 @@ IDAKLUSolver *create_idaklu_solver( jac_bandwidth_lower, jac_bandwidth_upper, std::move(functions), - options_cpp + setup_opts, + solver_opts ); } - else if (options_cpp.linear_solver == "SUNLinSol_SPTFQMR") + else if (setup_opts.linear_solver == "SUNLinSol_SPTFQMR") { DEBUG("\tsetting SUNLinSol_SPGMR solver"); idakluSolver = new IDAKLUSolverOpenMP_SPTFQMR( @@ -170,7 +177,8 @@ IDAKLUSolver *create_idaklu_solver( jac_bandwidth_lower, jac_bandwidth_upper, std::move(functions), - options_cpp + setup_opts, + solver_opts ); } diff --git a/pybamm/solvers/c_solvers/idaklu/sundials_functions.inl b/pybamm/solvers/c_solvers/idaklu/sundials_functions.inl index 532995dfb4..98257275a8 100644 --- a/pybamm/solvers/c_solvers/idaklu/sundials_functions.inl +++ b/pybamm/solvers/c_solvers/idaklu/sundials_functions.inl @@ -161,11 +161,11 @@ int jacobian_eval(realtype tt, realtype cj, N_Vector yy, N_Vector yp, // create pointer to jac data, column pointers, and row values realtype *jac_data; - if (p_python_functions->options.using_sparse_matrix) + if (p_python_functions->setup_opts.using_sparse_matrix) { jac_data = SUNSparseMatrix_Data(JJ); } - else if (p_python_functions->options.using_banded_matrix) { + else if (p_python_functions->setup_opts.using_banded_matrix) { jac_data = p_python_functions->get_tmp_sparse_jacobian_data(); } else @@ -191,7 +191,7 @@ int jacobian_eval(realtype tt, realtype cj, N_Vector yy, N_Vector yp, DEBUG("cj = " << cj); DEBUG_v(jac_data, 100); - if (p_python_functions->options.using_banded_matrix) + if (p_python_functions->setup_opts.using_banded_matrix) { // copy data from temporary matrix to the banded matrix auto jac_colptrs = p_python_functions->jac_times_cjmass_colptrs.data(); @@ -207,7 +207,7 @@ int jacobian_eval(realtype tt, realtype cj, N_Vector yy, N_Vector yp, } } } - else if (p_python_functions->options.using_sparse_matrix) + else if (p_python_functions->setup_opts.using_sparse_matrix) { if (SUNSparseMatrix_SparseType(JJ) == CSC_MAT) { diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index f1f32b1e63..8741a595f8 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -75,29 +75,83 @@ class IDAKLUSolver(pybamm.BaseSolver): .. code-block:: python options = { - # print statistics of the solver after every solve + # Print statistics of the solver after every solve "print_stats": False, - # jacobian form, can be "none", "dense", - # "banded", "sparse", "matrix-free" - "jacobian": "sparse", + # Number of threads available for OpenMP + "num_threads": 1, + # Evaluation engine to use for jax, can be 'jax'(native) or 'iree' + "jax_evaluator": "jax", + ## Linear solver interface # name of sundials linear solver to use options are: "SUNLinSol_KLU", # "SUNLinSol_Dense", "SUNLinSol_Band", "SUNLinSol_SPBCGS", # "SUNLinSol_SPFGMR", "SUNLinSol_SPGMR", "SUNLinSol_SPTFQMR", "linear_solver": "SUNLinSol_KLU", - # preconditioner for iterative solvers, can be "none", "BBDP" + # Jacobian form, can be "none", "dense", + # "banded", "sparse", "matrix-free" + "jacobian": "sparse", + # Preconditioner for iterative solvers, can be "none", "BBDP" "preconditioner": "BBDP", - # for iterative linear solvers, max number of iterations - "linsol_max_iterations": 5, - # for iterative linear solver preconditioner, bandwidth of + # For iterative linear solver preconditioner, bandwidth of # approximate jacobian "precon_half_bandwidth": 5, - # for iterative linear solver preconditioner, bandwidth of + # For iterative linear solver preconditioner, bandwidth of # approximate jacobian that is kept "precon_half_bandwidth_keep": 5, - # Number of threads available for OpenMP - "num_threads": 1, - # Evaluation engine to use for jax, can be 'jax'(native) or 'iree' - "jax_evaluator": "jax", + # For iterative linear solvers, max number of iterations + "linsol_max_iterations": 5, + # Ratio between linear and nonlinear tolerances + "epsilon_linear_tolerance": 0.05, + # Increment factor used in DQ Jacobian-vector product approximation + "increment_factor": 1.0, + # Enable or disable linear solution scaling + "linear_solution_scaling": True, + ## Main solver + # Maximum order of the linear multistep method + "max_order_bdf": 5, + # Maximum number of steps to be taken by the solver in its attempt to + # reach the next output time. + # Note: this value differs from the IDA default of 500 + "max_num_steps": 100000, + # Initial step size. The solver default is used if this is left at 0.0 + "dt_init": 0.0, + # Maximum absolute step size. The solver default is used if this is + # left at 0.0 + "dt_max": 0.0, + # Maximum number of error test failures in attempting one step + "max_error_test_failures": 10, + # Maximum number of nonlinear solver iterations at one step + "max_nonlinear_iterations": 4, + # Maximum number of nonlinear solver convergence failures at one step + "max_convergence_failures": 10, + # Safety factor in the nonlinear convergence test + "nonlinear_convergence_coefficient": 0.33, + # Suppress algebraic variables from error test + "suppress_algebraic_error": False, + ## Initial conditions calculation + # Positive constant in the Newton iteration convergence test within the + # initial condition calculation + "nonlinear_convergence_coefficient_ic": 0.0033, + # Maximum number of steps allowed when `init_all_y_ic = False` + "max_num_steps_ic": 5, + # Maximum number of the approximate Jacobian or preconditioner evaluations + # allowed when the Newton iteration appears to be slowly converging + # Note: this value differs from the IDA default of 4 + "max_num_jacobians_ic": 40, + # Maximum number of Newton iterations allowed in any one attempt to solve + # the initial conditions calculation problem + # Note: this value differs from the IDA default of 10 + "max_num_iterations_ic": 100, + # Maximum number of linesearch backtracks allowed in any Newton iteration, + # when solving the initial conditions calculation problem + "max_linesearch_backtracks_ic": 100, + # Turn off linesearch + "linesearch_off_ic": False, + # How to calculate the initial conditions. + # "True": calculate all y0 given ydot0 + # "False": calculate y_alg0 and ydot_diff0 given y_diff0 + "init_all_y_ic": False, + # Calculate consistent initial conditions + "calc_ic": True, } Note: These options only have an effect if model.convert_to_format == 'casadi' @@ -107,7 +161,7 @@ class IDAKLUSolver(pybamm.BaseSolver): def __init__( self, - rtol=1e-6, + rtol=1e-4, atol=1e-6, root_method="casadi", root_tol=1e-6, @@ -120,13 +174,33 @@ def __init__( default_options = { "print_stats": False, "jacobian": "sparse", - "linear_solver": "SUNLinSol_KLU", "preconditioner": "BBDP", - "linsol_max_iterations": 5, "precon_half_bandwidth": 5, "precon_half_bandwidth_keep": 5, "num_threads": 1, "jax_evaluator": "jax", + "linear_solver": "SUNLinSol_KLU", + "linsol_max_iterations": 5, + "epsilon_linear_tolerance": 0.05, + "increment_factor": 1.0, + "linear_solution_scaling": True, + "max_order_bdf": 5, + "max_num_steps": 100000, + "dt_init": 0.0, + "dt_max": 0.0, + "max_error_test_failures": 10, + "max_nonlinear_iterations": 40, + "max_convergence_failures": 100, + "nonlinear_convergence_coefficient": 0.33, + "suppress_algebraic_error": False, + "nonlinear_convergence_coefficient_ic": 0.0033, + "max_num_steps_ic": 5, + "max_num_jacobians_ic": 4, + "max_num_iterations_ic": 10, + "max_linesearch_backtracks_ic": 100, + "linesearch_off_ic": False, + "init_all_y_ic": False, + "calc_ic": True, } if options is None: options = default_options @@ -912,73 +986,71 @@ def _integrate(self, model, t_eval, inputs_dict=None): else: yS_out = False - if sol.flag in [0, 2]: - # 0 = solved for all t_eval - if sol.flag == 0: - termination = "final time" - # 2 = found root(s) - elif sol.flag == 2: - termination = "event" - - newsol = pybamm.Solution( - sol.t, - np.transpose(y_out), - model, - inputs_dict, - np.array([t[-1]]), - np.transpose(y_out[-1])[:, np.newaxis], - termination, - sensitivities=yS_out, - ) - newsol.integration_time = integration_time - if self.output_variables: - # Populate variables and sensititivies dictionaries directly - number_of_samples = sol.y.shape[0] // number_of_timesteps - sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) - startk = 0 - for _, var in enumerate(self.output_variables): - # ExplicitTimeIntegral's are not computed as part of the solver and - # do not need to be converted - if isinstance( - model.variables_and_events[var], pybamm.ExplicitTimeIntegral - ): - continue - if model.convert_to_format == "casadi": - len_of_var = ( - self._setup["var_fcns"][var](0.0, 0.0, 0.0).sparsity().nnz() - ) - base_variables = [self._setup["var_fcns"][var]] - elif ( - model.convert_to_format == "jax" - and self._options["jax_evaluator"] == "iree" - ): - idx = self.output_variables.index(var) - len_of_var = self._setup["var_idaklu_fcns"][idx].nnz - base_variables = [self._setup["var_idaklu_fcns"][idx]] - else: # pragma: no cover - raise pybamm.SolverError( - "Unsupported evaluation engine for convert_to_format=" - + f"{model.convert_to_format} " - + f"(jax_evaluator={self._options['jax_evaluator']})" - ) - newsol._variables[var] = pybamm.ProcessedVariableComputed( - [model.variables_and_events[var]], - base_variables, - [sol.y[:, startk : (startk + len_of_var)]], - newsol, - ) - # Add sensitivities - newsol[var]._sensitivities = {} - if model.calculate_sensitivities: - for paramk, param in enumerate(inputs_dict.keys()): - newsol[var].add_sensitivity( - param, - [sol.yS[:, startk : (startk + len_of_var), paramk]], - ) - startk += len_of_var - return newsol + # 0 = solved for all t_eval + if sol.flag == 0: + termination = "final time" + # 2 = found root(s) + elif sol.flag == 2: + termination = "event" else: raise pybamm.SolverError("idaklu solver failed") + newsol = pybamm.Solution( + sol.t, + np.transpose(y_out), + model, + inputs_dict, + np.array([t[-1]]), + np.transpose(y_out[-1])[:, np.newaxis], + termination, + sensitivities=yS_out, + ) + newsol.integration_time = integration_time + if self.output_variables: + # Populate variables and sensititivies dictionaries directly + number_of_samples = sol.y.shape[0] // number_of_timesteps + sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) + startk = 0 + for var in self.output_variables: + # ExplicitTimeIntegral's are not computed as part of the solver and + # do not need to be converted + if isinstance( + model.variables_and_events[var], pybamm.ExplicitTimeIntegral + ): + continue + if model.convert_to_format == "casadi": + len_of_var = ( + self._setup["var_fcns"][var](0.0, 0.0, 0.0).sparsity().nnz() + ) + base_variables = [self._setup["var_fcns"][var]] + elif ( + model.convert_to_format == "jax" + and self._options["jax_evaluator"] == "iree" + ): + idx = self.output_variables.index(var) + len_of_var = self._setup["var_idaklu_fcns"][idx].nnz + base_variables = [self._setup["var_idaklu_fcns"][idx]] + else: # pragma: no cover + raise pybamm.SolverError( + "Unsupported evaluation engine for convert_to_format=" + + f"{model.convert_to_format} " + + f"(jax_evaluator={self._options['jax_evaluator']})" + ) + newsol._variables[var] = pybamm.ProcessedVariableComputed( + [model.variables_and_events[var]], + base_variables, + [sol.y[:, startk : (startk + len_of_var)]], + newsol, + ) + # Add sensitivities + newsol[var]._sensitivities = {} + if model.calculate_sensitivities: + for paramk, param in enumerate(inputs_dict.keys()): + newsol[var].add_sensitivity( + param, + [sol.yS[:, startk : (startk + len_of_var), paramk]], + ) + startk += len_of_var + return newsol def jaxify( self, diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/base_lithium_ion_tests.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/base_lithium_ion_tests.py index 4db5ddea61..1fff91b476 100644 --- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/base_lithium_ion_tests.py +++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/base_lithium_ion_tests.py @@ -20,10 +20,12 @@ def test_basic_processing(self): def test_sensitivities(self): model = self.model() param = pybamm.ParameterValues("Ecker2015") + rtol = 1e-6 + atol = 1e-6 if pybamm.have_idaklu(): - solver = pybamm.IDAKLUSolver() + solver = pybamm.IDAKLUSolver(rtol=rtol, atol=atol) else: - solver = pybamm.CasadiSolver() + solver = pybamm.CasadiSolver(rtol=rtol, atol=atol) modeltest = tests.StandardModelTest( model, parameter_values=param, solver=solver ) diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index a80ab74b9e..562dde27a5 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -241,6 +241,8 @@ def test_sensitivities_initial_condition(self): disc = pybamm.Discretisation() disc.process_model(model) solver = pybamm.IDAKLUSolver( + rtol=1e-6, + atol=1e-6, root_method=root_method, output_variables=output_variables, options={"jax_evaluator": "iree"} if form == "iree" else {}, @@ -539,7 +541,7 @@ def test_banded(self): np.testing.assert_array_almost_equal(soln.y, soln_banded.y, 5) - def test_options(self): + def test_setup_options(self): model = pybamm.BaseModel() u = pybamm.Variable("u") v = pybamm.Variable("v") @@ -584,8 +586,13 @@ def test_options(self): "jacobian": jacobian, "linear_solver": linear_solver, "preconditioner": precon, + "max_num_steps": 10000, } - solver = pybamm.IDAKLUSolver(options=options) + solver = pybamm.IDAKLUSolver( + atol=1e-8, + rtol=1e-8, + options=options, + ) if ( jacobian == "none" and (linear_solver == "SUNLinSol_Dense") @@ -614,6 +621,70 @@ def test_options(self): with self.assertRaises(ValueError): soln = solver.solve(model, t_eval) + def test_solver_options(self): + model = pybamm.BaseModel() + u = pybamm.Variable("u") + v = pybamm.Variable("v") + model.rhs = {u: -0.1 * u} + model.algebraic = {v: v - u} + model.initial_conditions = {u: 1, v: 1} + disc = pybamm.Discretisation() + disc.process_model(model) + + t_eval = np.linspace(0, 1) + solver = pybamm.IDAKLUSolver() + soln_base = solver.solve(model, t_eval) + + options_success = { + "max_order_bdf": 4, + "max_num_steps": 490, + "dt_init": 0.01, + "dt_max": 1000.9, + "max_error_test_failures": 11, + "max_nonlinear_iterations": 5, + "max_convergence_failures": 11, + "nonlinear_convergence_coefficient": 1.0, + "suppress_algebraic_error": True, + "nonlinear_convergence_coefficient_ic": 0.01, + "max_num_steps_ic": 6, + "max_num_jacobians_ic": 5, + "max_num_iterations_ic": 11, + "max_linesearch_backtracks_ic": 101, + "linesearch_off_ic": True, + "init_all_y_ic": False, + "linear_solver": "SUNLinSol_KLU", + "linsol_max_iterations": 6, + "epsilon_linear_tolerance": 0.06, + "increment_factor": 0.99, + "linear_solution_scaling": False, + } + + # test everything works + for option in options_success: + options = {option: options_success[option]} + solver = pybamm.IDAKLUSolver(rtol=1e-6, atol=1e-6, options=options) + soln = solver.solve(model, t_eval) + + np.testing.assert_array_almost_equal(soln.y, soln_base.y, 5) + + options_fail = { + "max_order_bdf": -1, + "max_num_steps_ic": -1, + "max_num_jacobians_ic": -1, + "max_num_iterations_ic": -1, + "max_linesearch_backtracks_ic": -1, + "epsilon_linear_tolerance": -1.0, + "increment_factor": -1.0, + } + + # test that the solver throws a warning + for option in options_fail: + options = {option: options_fail[option]} + solver = pybamm.IDAKLUSolver(options=options) + + with self.assertRaises(ValueError): + solver.solve(model, t_eval) + def test_with_output_variables(self): # Construct a model and solve for all variables, then test # the 'output_variables' option for each variable in turn, confirming