diff --git a/CHANGELOG.md b/CHANGELOG.md index 29d047d748..71a2d3f398 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ## Features - SEI reactions can now be asymmetric ([#2425](https://github.com/pybamm-team/PyBaMM/pull/2425)) +- New Idaklu solver options for jacobian type and linear solver, support Sundials v6 ([#2444](https://github.com/pybamm-team/PyBaMM/pull/2444)) ## Optimizations diff --git a/CMakeBuild.py b/CMakeBuild.py index 8a2ed4716e..d1b3d725b2 100644 --- a/CMakeBuild.py +++ b/CMakeBuild.py @@ -82,7 +82,10 @@ def run(self): use_python_casadi = False else: use_python_casadi = True + + build_type = os.getenv("PYBAMM_CPP_BUILD_TYPE", "Release") cmake_args = [ + "-DCMAKE_BUILD_TYPE={}".format(build_type), "-DPYTHON_EXECUTABLE={}".format(sys.executable), "-DUSE_PYTHON_CASADI={}".format("TRUE" if use_python_casadi else "FALSE"), ] diff --git a/CMakeLists.txt b/CMakeLists.txt index fc3bec444d..cc47ba99c0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -35,14 +35,20 @@ endif() add_subdirectory(${PYBIND11_DIR}) pybind11_add_module(idaklu - pybamm/solvers/c_solvers/idaklu_python.cpp - pybamm/solvers/c_solvers/idaklu_python.hpp + pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp + pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp + pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp + pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp + pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.hpp + pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp + pybamm/solvers/c_solvers/idaklu/common.hpp + pybamm/solvers/c_solvers/idaklu/python.hpp + pybamm/solvers/c_solvers/idaklu/python.cpp + pybamm/solvers/c_solvers/idaklu/solution.cpp + pybamm/solvers/c_solvers/idaklu/solution.hpp + pybamm/solvers/c_solvers/idaklu/options.hpp + pybamm/solvers/c_solvers/idaklu/options.cpp pybamm/solvers/c_solvers/idaklu.cpp - pybamm/solvers/c_solvers/idaklu.hpp - pybamm/solvers/c_solvers/idaklu_casadi.cpp - pybamm/solvers/c_solvers/idaklu_casadi.hpp - pybamm/solvers/c_solvers/solution.cpp - pybamm/solvers/c_solvers/solution.hpp ) if (NOT DEFINED USE_PYTHON_CASADI) diff --git a/FindSUNDIALS.cmake b/FindSUNDIALS.cmake index 643b75bbd4..92a9b785e9 100644 --- a/FindSUNDIALS.cmake +++ b/FindSUNDIALS.cmake @@ -8,6 +8,8 @@ # # * sundials_ida # * sundials_sunlinsolklu +# * sundials_sunlinsoldense +# * sundials_sunlinsollapackdense # * sundials_sunmatrix_sparse # * sundials_nvecserial # @@ -31,6 +33,9 @@ find_path(SUNDIALS_INCLUDE_DIR sundials/sundials_math.h sundials/sundials_types.h sunlinsol/sunlinsol_klu.h + sunlinsol/sunlinsol_dense.h + sunlinsol/sunlinsol_spbcgs.h + sunlinsol/sunlinsol_lapackdense.h sunmatrix/sunmatrix_sparse.h PATH_SUFFIXES include @@ -41,6 +46,9 @@ find_path(SUNDIALS_INCLUDE_DIR set(SUNDIALS_WANT_COMPONENTS sundials_idas sundials_sunlinsolklu + sundials_sunlinsoldense + sundials_sunlinsolspbcgs + sundials_sunlinsollapackdense sundials_sunmatrixsparse sundials_nvecserial ) diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index 9241487419..ac90172c97 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -1,455 +1,62 @@ -#include "idaklu.hpp" -#include +#include "idaklu/casadi_solver.hpp" +#include "idaklu/common.hpp" +#include "idaklu/python.hpp" -class PybammFunctions -{ -public: - int number_of_states; - int number_of_parameters; - int number_of_events; - - PybammFunctions(const residual_type &res, const jacobian_type &jac, - const sensitivities_type &sens, - const jac_get_type &get_jac_data_in, - const jac_get_type &get_jac_row_vals_in, - const jac_get_type &get_jac_col_ptrs_in, - const event_type &event, - const int n_s, int n_e, const int n_p, - const np_array &inputs) - : number_of_states(n_s), number_of_events(n_e), - number_of_parameters(n_p), - py_res(res), py_jac(jac), - py_sens(sens), - py_event(event), py_get_jac_data(get_jac_data_in), - py_get_jac_row_vals(get_jac_row_vals_in), - py_get_jac_col_ptrs(get_jac_col_ptrs_in), - inputs(inputs) - { - } - - np_array operator()(double t, np_array y, np_array yp) - { - return py_res(t, y, inputs, yp); - } - - np_array res(double t, np_array y, np_array yp) - { - return py_res(t, y, inputs, yp); - } - - void jac(double t, np_array y, double cj) - { - // this function evaluates the jacobian and sets it to be the attribute - // of a python class which can then be called by get_jac_data, - // get_jac_col_ptr, etc - py_jac(t, y, inputs, cj); - } - - void sensitivities( - std::vector& resvalS, - const double t, const np_array& y, const np_array& yp, - const std::vector& yS, const std::vector& ypS) - { - // this function evaluates the sensitivity equations required by IDAS, - // returning them in resvalS, which is preallocated as a numpy array - // of size (np, n), where n is the number of states and np is the number - // of parameters - // - // yS and ypS are also shape (np, n), y and yp are shape (n) - // - // dF/dy * s_i + dF/dyd * sd + dFdp_i for i in range(np) - py_sens(resvalS, t, y, inputs, yp, yS, ypS); - } - - np_array get_jac_data() { return py_get_jac_data(); } - - np_array get_jac_row_vals() { return py_get_jac_row_vals(); } - - np_array get_jac_col_ptrs() { return py_get_jac_col_ptrs(); } - - np_array events(double t, np_array y) { return py_event(t, y, inputs); } +#include +#include +#include +#include -private: - residual_type py_res; - sensitivities_type py_sens; - jacobian_type py_jac; - event_type py_event; - jac_get_type py_get_jac_data; - jac_get_type py_get_jac_row_vals; - jac_get_type py_get_jac_col_ptrs; - const np_array &inputs; -}; +#include -int residual(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, - void *user_data) +Function generate_function(const std::string &data) { - PybammFunctions *python_functions_ptr = - static_cast(user_data); - PybammFunctions python_functions = *python_functions_ptr; - - realtype *yval, *ypval, *rval; - yval = N_VGetArrayPointer(yy); - ypval = N_VGetArrayPointer(yp); - rval = N_VGetArrayPointer(rr); - - int n = python_functions.number_of_states; - py::array_t y_np = py::array_t(n, yval); - py::array_t yp_np = py::array_t(n, ypval); - - py::array_t r_np; - - r_np = python_functions.res(tres, y_np, yp_np); - - auto r_np_ptr = r_np.unchecked<1>(); - - // just copying data - int i; - for (i = 0; i < n; i++) - { - rval[i] = r_np_ptr[i]; - } - return 0; + return Function::deserialize(data); } -int jacobian(realtype tt, realtype cj, N_Vector yy, N_Vector yp, - N_Vector resvec, SUNMatrix JJ, void *user_data, N_Vector tempv1, - N_Vector tempv2, N_Vector tempv3) -{ - realtype *yval; - yval = N_VGetArrayPointer(yy); - - PybammFunctions *python_functions_ptr = - static_cast(user_data); - PybammFunctions python_functions = *python_functions_ptr; - - int n = python_functions.number_of_states; - py::array_t y_np = py::array_t(n, yval); - - // create pointer to jac data, column pointers, and row values - sunindextype *jac_colptrs = SUNSparseMatrix_IndexPointers(JJ); - sunindextype *jac_rowvals = SUNSparseMatrix_IndexValues(JJ); - realtype *jac_data = SUNSparseMatrix_Data(JJ); +namespace py = pybind11; - py::array_t jac_np_array; +PYBIND11_MAKE_OPAQUE(std::vector); - python_functions.jac(tt, y_np, cj); - - np_array jac_np_data = python_functions.get_jac_data(); - int n_data = jac_np_data.request().size; - auto jac_np_data_ptr = jac_np_data.unchecked<1>(); - - // just copy across data - int i; - for (i = 0; i < n_data; i++) - { - jac_data[i] = jac_np_data_ptr[i]; - } - - np_array jac_np_row_vals = python_functions.get_jac_row_vals(); - int n_row_vals = jac_np_row_vals.request().size; - - auto jac_np_row_vals_ptr = jac_np_row_vals.unchecked<1>(); - // just copy across row vals (this might be unneeded) - for (i = 0; i < n_row_vals; i++) - { - jac_rowvals[i] = jac_np_row_vals_ptr[i]; - } - - np_array jac_np_col_ptrs = python_functions.get_jac_col_ptrs(); - int n_col_ptrs = jac_np_col_ptrs.request().size; - auto jac_np_col_ptrs_ptr = jac_np_col_ptrs.unchecked<1>(); - - // just copy across col ptrs (this might be unneeded) - for (i = 0; i < n_col_ptrs; i++) - { - jac_colptrs[i] = jac_np_col_ptrs_ptr[i]; - } - - return (0); -} - -int events(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, - void *user_data) +PYBIND11_MODULE(idaklu, m) { - realtype *yval; - yval = N_VGetArrayPointer(yy); - - PybammFunctions *python_functions_ptr = - static_cast(user_data); - PybammFunctions python_functions = *python_functions_ptr; - - int number_of_events = python_functions.number_of_events; - int number_of_states = python_functions.number_of_states; - py::array_t y_np = py::array_t(number_of_states, yval); - - py::array_t events_np_array; - - events_np_array = python_functions.events(t, y_np); - - auto events_np_data_ptr = events_np_array.unchecked<1>(); - - // just copying data (figure out how to pass pointers later) - int i; - for (i = 0; i < number_of_events; i++) - { - events_ptr[i] = events_np_data_ptr[i]; - } - - return (0); + m.doc() = "sundials solvers"; // optional module docstring + + py::bind_vector>(m, "VectorNdArray"); + + m.def("solve_python", &solve_python, + "The solve function for python evaluators", py::arg("t"), py::arg("y0"), + py::arg("yp0"), py::arg("res"), py::arg("jac"), py::arg("sens"), + py::arg("get_jac_data"), py::arg("get_jac_row_vals"), + py::arg("get_jac_col_ptr"), py::arg("nnz"), py::arg("events"), + py::arg("number_of_events"), py::arg("use_jacobian"), + py::arg("rhs_alg_id"), py::arg("atol"), py::arg("rtol"), + py::arg("inputs"), py::arg("number_of_sensitivity_parameters"), + py::return_value_policy::take_ownership); + + py::class_(m, "CasadiSolver") + .def("solve", &CasadiSolver::solve, "perform a solve", py::arg("t"), + py::arg("y0"), py::arg("yp0"), py::arg("inputs"), + py::return_value_policy::take_ownership); + + m.def("create_casadi_solver", &create_casadi_solver, + "Create a casadi idaklu solver object", py::arg("number_of_states"), + py::arg("number_of_parameters"), py::arg("rhs_alg"), + py::arg("jac_times_cjmass"), py::arg("jac_times_cjmass_colptrs"), + py::arg("jac_times_cjmass_rowvals"), py::arg("jac_times_cjmass_nnz"), + py::arg("jac_action"), py::arg("mass_action"), py::arg("sens"), + py::arg("events"), py::arg("number_of_events"), py::arg("rhs_alg_id"), + py::arg("atol"), py::arg("rtol"), py::arg("inputs"), py::arg("options"), + py::return_value_policy::take_ownership); + + m.def("generate_function", &generate_function, "Generate a casadi function", + py::arg("string"), py::return_value_policy::take_ownership); + + py::class_(m, "Function"); + + py::class_(m, "solution") + .def_readwrite("t", &Solution::t) + .def_readwrite("y", &Solution::y) + .def_readwrite("yS", &Solution::yS) + .def_readwrite("flag", &Solution::flag); } - -int sensitivities(int Ns, realtype t, N_Vector yy, N_Vector yp, - N_Vector resval, N_Vector *yS, N_Vector *ypS, N_Vector *resvalS, - void *user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3) { -// This function computes the sensitivity residual for all sensitivity -// equations. It must compute the vectors -// (∂F/∂y)s i (t)+(∂F/∂ ẏ) ṡ i (t)+(∂F/∂p i ) and store them in resvalS[i]. -// Ns is the number of sensitivities. -// t is the current value of the independent variable. -// yy is the current value of the state vector, y(t). -// yp is the current value of ẏ(t). -// resval contains the current value F of the original DAE residual. -// yS contains the current values of the sensitivities s i . -// ypS contains the current values of the sensitivity derivatives ṡ i . -// resvalS contains the output sensitivity residual vectors. -// Memory allocation for resvalS is handled within idas. -// user data is a pointer to user data. -// tmp1, tmp2, tmp3 are N Vectors of length N which can be used as -// temporary storage. -// -// Return value An IDASensResFn should return 0 if successful, -// a positive value if a recoverable error -// occurred (in which case idas will attempt to correct), -// or a negative value if it failed unrecoverably (in which case the integration is halted and IDA SRES FAIL is returned) -// - PybammFunctions *python_functions_ptr = - static_cast(user_data); - PybammFunctions python_functions = *python_functions_ptr; - - int n = python_functions.number_of_states; - int np = python_functions.number_of_parameters; - - // memory managed by sundials, so pass a destructor that does nothing - auto state_vector_shape = std::vector{n, 1}; - np_array y_np = np_array(state_vector_shape, N_VGetArrayPointer(yy), - py::capsule(&yy, [](void* p) {})); - np_array yp_np = np_array(state_vector_shape, N_VGetArrayPointer(yp), - py::capsule(&yp, [](void* p) {})); - - std::vector yS_np(np); - for (int i = 0; i < np; i++) { - auto capsule = py::capsule(yS + i, [](void* p) {}); - yS_np[i] = np_array(state_vector_shape, N_VGetArrayPointer(yS[i]), capsule); - } - - std::vector ypS_np(np); - for (int i = 0; i < np; i++) { - auto capsule = py::capsule(ypS + i, [](void* p) {}); - ypS_np[i] = np_array(state_vector_shape, N_VGetArrayPointer(ypS[i]), capsule); - } - - std::vector resvalS_np(np); - for (int i = 0; i < np; i++) { - auto capsule = py::capsule(resvalS + i, [](void* p) {}); - resvalS_np[i] = np_array(state_vector_shape, - N_VGetArrayPointer(resvalS[i]), capsule); - } - - realtype *ptr1 = static_cast(resvalS_np[0].request().ptr); - const realtype* resvalSval = N_VGetArrayPointer(resvalS[0]); - - python_functions.sensitivities(resvalS_np, t, y_np, yp_np, yS_np, ypS_np); - - return 0; -} - -/* main program */ -Solution solve_python(np_array t_np, np_array y0_np, np_array yp0_np, - residual_type res, jacobian_type jac, - sensitivities_type sens, - jac_get_type gjd, jac_get_type gjrv, jac_get_type gjcp, - int nnz, event_type event, - int number_of_events, int use_jacobian, np_array rhs_alg_id, - np_array atol_np, double rel_tol, np_array inputs, - int number_of_parameters) -{ - auto t = t_np.unchecked<1>(); - auto y0 = y0_np.unchecked<1>(); - auto yp0 = yp0_np.unchecked<1>(); - auto atol = atol_np.unchecked<1>(); - - int number_of_states = y0_np.request().size; - int number_of_timesteps = t_np.request().size; - void *ida_mem; // pointer to memory - N_Vector yy, yp, avtol; // y, y', and absolute tolerance - N_Vector *yyS, *ypS; // y, y' for sensitivities - realtype rtol, *yval, *ypval, *atval; - std::vector ySval(number_of_parameters); - int retval; - SUNMatrix J; - SUNLinearSolver LS; - - // allocate vectors - yy = N_VNew_Serial(number_of_states); - yp = N_VNew_Serial(number_of_states); - avtol = N_VNew_Serial(number_of_states); - - if (number_of_parameters > 0) { - yyS = N_VCloneVectorArray(number_of_parameters, yy); - ypS = N_VCloneVectorArray(number_of_parameters, yp); - } - - // set initial value - yval = N_VGetArrayPointer(yy); - ypval = N_VGetArrayPointer(yp); - atval = N_VGetArrayPointer(avtol); - int i; - for (i = 0; i < number_of_states; i++) - { - yval[i] = y0[i]; - ypval[i] = yp0[i]; - atval[i] = atol[i]; - } - - for (int is = 0 ; is < number_of_parameters; is++) { - ySval[is] = N_VGetArrayPointer(yyS[is]); - N_VConst(RCONST(0.0), yyS[is]); - N_VConst(RCONST(0.0), ypS[is]); - } - - // allocate memory for solver - ida_mem = IDACreate(); - - // initialise solver - realtype t0 = RCONST(t(0)); - IDAInit(ida_mem, residual, t0, yy, yp); - - // set tolerances - rtol = RCONST(rel_tol); - - IDASVtolerances(ida_mem, rtol, avtol); - - // set events - IDARootInit(ida_mem, number_of_events, events); - - // set pybamm functions by passing pointer to it - PybammFunctions pybamm_functions(res, jac, sens, gjd, gjrv, gjcp, event, - number_of_states, number_of_events, - number_of_parameters, inputs); - void *user_data = &pybamm_functions; - IDASetUserData(ida_mem, user_data); - - // set linear solver - J = SUNSparseMatrix(number_of_states, number_of_states, nnz, CSR_MAT); - - LS = SUNLinSol_KLU(yy, J); - IDASetLinearSolver(ida_mem, LS, J); - - if (use_jacobian == 1) - { - IDASetJacFn(ida_mem, jacobian); - } - - if (number_of_parameters > 0) - { - IDASensInit(ida_mem, number_of_parameters, - IDA_SIMULTANEOUS, sensitivities, yyS, ypS); - IDASensEEtolerances(ida_mem); - } - - int t_i = 1; - realtype tret; - realtype t_next; - realtype t_final = t(number_of_timesteps - 1); - - // set return vectors - std::vector t_return(number_of_timesteps); - std::vector y_return(number_of_timesteps * number_of_states); - std::vector yS_return(number_of_parameters * number_of_timesteps * number_of_states); - - t_return[0] = t(0); - for (int j = 0; j < number_of_states; j++) - { - y_return[j] = yval[j]; - } - for (int j = 0; j < number_of_parameters; j++) { - const int base_index = j * number_of_timesteps * number_of_states; - for (int k = 0; k < number_of_states; k++) { - yS_return[base_index + k] = ySval[j][k]; - } - } - - // calculate consistent initial conditions - N_Vector id; - auto id_np_val = rhs_alg_id.unchecked<1>(); - id = N_VNew_Serial(number_of_states); - realtype *id_val; - id_val = N_VGetArrayPointer(id); - - int ii; - for (ii = 0; ii < number_of_states; ii++) - { - id_val[ii] = id_np_val[ii]; - } - - IDASetId(ida_mem, id); - IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); - - while (true) - { - t_next = t(t_i); - IDASetStopTime(ida_mem, t_next); - retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); - - if (retval == IDA_TSTOP_RETURN || retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) - { - if (number_of_parameters > 0) { - IDAGetSens(ida_mem, &tret, yyS); - } - - t_return[t_i] = tret; - 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++) { - yS_return[base_index + k] = ySval[j][k]; - } - } - t_i += 1; - if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { - break; - } - - } - } - - /* Free memory */ - if (number_of_parameters > 0) { - IDASensFree(ida_mem); - } - IDAFree(&ida_mem); - SUNLinSolFree(LS); - SUNMatDestroy(J); - N_VDestroy(avtol); - N_VDestroy(yp); - if (number_of_parameters > 0) { - N_VDestroyVectorArray(yyS, number_of_parameters); - N_VDestroyVectorArray(ypS, number_of_parameters); - } - - np_array t_ret = np_array(t_i, &t_return[0]); - np_array y_ret = np_array(t_i * number_of_states, &y_return[0]); - np_array yS_ret = np_array( - std::vector{number_of_parameters, number_of_timesteps, number_of_states}, - &yS_return[0] - ); - - Solution sol(retval, t_ret, y_ret, yS_ret); - - return sol; -} - diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp new file mode 100644 index 0000000000..a2de2e7089 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -0,0 +1,69 @@ +#include "casadi_functions.hpp" + +CasadiFunction::CasadiFunction(const Function &f) : m_func(f) +{ + size_t sz_arg; + size_t sz_res; + size_t sz_iw; + size_t sz_w; + m_func.sz_work(sz_arg, sz_res, sz_iw, sz_w); + // std::cout << "name = "<< m_func.name() << " arg = " << sz_arg << " res = " + // << sz_res << " iw = " << sz_iw << " w = " << sz_w << std::endl; for (int i + // = 0; i < sz_arg; i++) { + // std::cout << "Sparsity for input " << i << std::endl; + // const Sparsity& sparsity = m_func.sparsity_in(i); + // } + // for (int i = 0; i < sz_res; i++) { + // std::cout << "Sparsity for output " << i << std::endl; + // const Sparsity& sparsity = m_func.sparsity_out(i); + // } + m_arg.resize(sz_arg); + m_res.resize(sz_res); + m_iw.resize(sz_iw); + m_w.resize(sz_w); +} + +// only call this once m_arg and m_res have been set appropriatelly +void CasadiFunction::operator()() +{ + int mem = m_func.checkout(); + m_func(m_arg.data(), m_res.data(), m_iw.data(), m_w.data(), mem); + m_func.release(mem); +} + +CasadiFunctions::CasadiFunctions( + const Function &rhs_alg, const Function &jac_times_cjmass, + const int jac_times_cjmass_nnz, + const np_array_int &jac_times_cjmass_rowvals_arg, + const np_array_int &jac_times_cjmass_colptrs_arg, + const int inputs_length, const Function &jac_action, + const Function &mass_action, const Function &sens, const Function &events, + const int n_s, int n_e, const int n_p, const Options& options) + : number_of_states(n_s), number_of_events(n_e), number_of_parameters(n_p), + number_of_nnz(jac_times_cjmass_nnz), rhs_alg(rhs_alg), + jac_times_cjmass(jac_times_cjmass), jac_action(jac_action), + mass_action(mass_action), sens(sens), events(events), + tmp(number_of_states), + options(options) +{ + + // copy across numpy array values + const int n_row_vals = jac_times_cjmass_rowvals_arg.request().size; + auto p_jac_times_cjmass_rowvals = jac_times_cjmass_rowvals_arg.unchecked<1>(); + jac_times_cjmass_rowvals.resize(n_row_vals); + for (int i = 0; i < n_row_vals; i++) { + jac_times_cjmass_rowvals[i] = p_jac_times_cjmass_rowvals[i]; + } + + const int n_col_ptrs = jac_times_cjmass_colptrs_arg.request().size; + auto p_jac_times_cjmass_colptrs = jac_times_cjmass_colptrs_arg.unchecked<1>(); + jac_times_cjmass_colptrs.resize(n_col_ptrs); + for (int i = 0; i < n_col_ptrs; i++) { + jac_times_cjmass_colptrs[i] = p_jac_times_cjmass_colptrs[i]; + } + + inputs.resize(inputs_length); + +} + +realtype *CasadiFunctions::get_tmp() { return tmp.data(); } diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp new file mode 100644 index 0000000000..2e3b6beb8d --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -0,0 +1,60 @@ +#ifndef PYBAMM_IDAKLU_CASADI_FUNCTIONS_HPP +#define PYBAMM_IDAKLU_CASADI_FUNCTIONS_HPP + +#include "common.hpp" +#include "options.hpp" +#include "solution.hpp" +#include + +using Function = casadi::Function; + +class CasadiFunction +{ +public: + explicit CasadiFunction(const Function &f); + +public: + std::vector m_arg; + std::vector m_res; + void operator()(); + +private: + const Function &m_func; + std::vector m_iw; + std::vector m_w; +}; + +class CasadiFunctions +{ +public: + int number_of_states; + int number_of_parameters; + int number_of_events; + int number_of_nnz; + CasadiFunction rhs_alg; + CasadiFunction sens; + CasadiFunction jac_times_cjmass; + std::vector jac_times_cjmass_rowvals; + std::vector jac_times_cjmass_colptrs; + std::vector inputs; + CasadiFunction jac_action; + CasadiFunction mass_action; + CasadiFunction events; + Options options; + + CasadiFunctions(const Function &rhs_alg, const Function &jac_times_cjmass, + const int jac_times_cjmass_nnz, + const np_array_int &jac_times_cjmass_rowvals, + const np_array_int &jac_times_cjmass_colptrs, + const int inputs_length, const Function &jac_action, + const Function &mass_action, const Function &sens, + const Function &events, const int n_s, int n_e, + const int n_p, const Options& options); + + realtype *get_tmp(); + +private: + std::vector tmp; +}; + +#endif // PYBAMM_IDAKLU_CASADI_FUNCTIONS_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp new file mode 100644 index 0000000000..fa701b41a2 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -0,0 +1,450 @@ +#include "casadi_solver.hpp" +#include "casadi_sundials_functions.hpp" +#include "common.hpp" +#include + +CasadiSolver * +create_casadi_solver(int number_of_states, int number_of_parameters, + const Function &rhs_alg, const Function &jac_times_cjmass, + const np_array_int &jac_times_cjmass_colptrs, + const np_array_int &jac_times_cjmass_rowvals, + const int jac_times_cjmass_nnz, const Function &jac_action, + const Function &mass_action, const Function &sens, + const Function &events, const int number_of_events, + np_array rhs_alg_id, np_array atol_np, double rel_tol, + int inputs_length, py::dict options) +{ + auto options_cpp = Options(options); + auto functions = std::make_unique( + rhs_alg, jac_times_cjmass, jac_times_cjmass_nnz, jac_times_cjmass_rowvals, + jac_times_cjmass_colptrs, inputs_length, jac_action, mass_action, sens, + events, number_of_states, number_of_events, number_of_parameters, + options_cpp); + + return new CasadiSolver(atol_np, rel_tol, rhs_alg_id, number_of_parameters, + number_of_events, jac_times_cjmass_nnz, + std::move(functions), options_cpp); +} + +CasadiSolver::CasadiSolver(np_array atol_np, double rel_tol, + np_array rhs_alg_id, int number_of_parameters, + int number_of_events, int jac_times_cjmass_nnz, + std::unique_ptr functions_arg, + const Options &options) + : 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), + functions(std::move(functions_arg)), options(options) +{ + DEBUG("CasadiSolver::CasadiSolver"); + auto atol = atol_np.unchecked<1>(); + + // allocate memory for solver +#if SUNDIALS_VERSION_MAJOR >= 6 + SUNContext_Create(NULL, &sunctx); + ida_mem = IDACreate(sunctx); +#else + ida_mem = IDACreate(); +#endif + + // allocate vectors +#if SUNDIALS_VERSION_MAJOR >= 6 + yy = N_VNew_Serial(number_of_states, sunctx); + yp = N_VNew_Serial(number_of_states, sunctx); + avtol = N_VNew_Serial(number_of_states, sunctx); + id = N_VNew_Serial(number_of_states, sunctx); +#else + yy = N_VNew_Serial(number_of_states); + yp = N_VNew_Serial(number_of_states); + avtol = N_VNew_Serial(number_of_states); + id = N_VNew_Serial(number_of_states); +#endif + + if (number_of_parameters > 0) + { + yyS = N_VCloneVectorArray(number_of_parameters, yy); + ypS = N_VCloneVectorArray(number_of_parameters, yp); + } + + // set initial value + realtype *atval = N_VGetArrayPointer(avtol); + for (int i = 0; i < number_of_states; i++) + { + atval[i] = atol[i]; + } + + for (int is = 0; is < number_of_parameters; is++) + { + N_VConst(RCONST(0.0), yyS[is]); + N_VConst(RCONST(0.0), ypS[is]); + } + + // initialise solver + + IDAInit(ida_mem, residual_casadi, 0, yy, yp); + + // set tolerances + rtol = RCONST(rel_tol); + + IDASVtolerances(ida_mem, rtol, avtol); + + // set events + IDARootInit(ida_mem, number_of_events, events_casadi); + + void *user_data = functions.get(); + IDASetUserData(ida_mem, user_data); + + // set matrix + if (options.jacobian == "sparse") + { + DEBUG("\tsetting sparse matrix"); +#if SUNDIALS_VERSION_MAJOR >= 6 + J = SUNSparseMatrix(number_of_states, number_of_states, + jac_times_cjmass_nnz, CSC_MAT, sunctx); +#else + J = SUNSparseMatrix(number_of_states, number_of_states, + jac_times_cjmass_nnz, CSC_MAT); +#endif + } + else if (options.jacobian == "dense" || options.jacobian == "none") + { + DEBUG("\tsetting dense matrix"); +#if SUNDIALS_VERSION_MAJOR >= 6 + J = SUNDenseMatrix(number_of_states, number_of_states, sunctx); +#else + J = SUNDenseMatrix(number_of_states, number_of_states); +#endif + } + else if (options.jacobian == "matrix-free") + { + DEBUG("\tsetting matrix-free"); + J = NULL; + } + + #if SUNDIALS_VERSION_MAJOR >= 6 + int precon_type = SUN_PREC_NONE; + if (options.preconditioner != "none") { + precon_type = SUN_PREC_LEFT; + } + #else + int precon_type = PREC_NONE; + if (options.preconditioner != "none") { + precon_type = PREC_LEFT; + } + #endif + + // set linear solver + if (options.linear_solver == "SUNLinSol_Dense") + { + DEBUG("\tsetting SUNLinSol_Dense linear solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_Dense(yy, J, sunctx); +#else + LS = SUNLinSol_Dense(yy, J); +#endif + } + else if (options.linear_solver == "SUNLinSol_LapackDense") + { + DEBUG("\tsetting SUNLinSol_LapackDense linear solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_LapackDense(yy, J, sunctx); +#else + LS = SUNLinSol_LapackDense(yy, J); +#endif + } + else if (options.linear_solver == "SUNLinSol_KLU") + { + DEBUG("\tsetting SUNLinSol_KLU linear solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_KLU(yy, J, sunctx); +#else + LS = SUNLinSol_KLU(yy, J); +#endif + } + else if (options.linear_solver == "SUNLinSol_SPBCGS") + { + DEBUG("\tsetting SUNLinSol_SPBCGS_linear solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_SPBCGS(yy, precon_type, options.linsol_max_iterations, + sunctx); +#else + LS = SUNLinSol_SPBCGS(yy, precon_type, options.linsol_max_iterations); +#endif + } + else if (options.linear_solver == "SUNLinSol_SPFGMR") + { + DEBUG("\tsetting SUNLinSol_SPFGMR_linear solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_SPFGMR(yy, precon_type, options.linsol_max_iterations, + sunctx); +#else + LS = SUNLinSol_SPFGMR(yy, precon_type, options.linsol_max_iterations); +#endif + } + else if (options.linear_solver == "SUNLinSol_SPGMR") + { + DEBUG("\tsetting SUNLinSol_SPGMR solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_SPGMR(yy, precon_type, options.linsol_max_iterations, + sunctx); +#else + LS = SUNLinSol_SPGMR(yy, precon_type, options.linsol_max_iterations); +#endif + } + else if (options.linear_solver == "SUNLinSol_SPTFQMR") + { + DEBUG("\tsetting SUNLinSol_SPGMR solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_SPTFQMR(yy, precon_type, options.linsol_max_iterations, + sunctx); +#else + LS = SUNLinSol_SPTFQMR(yy, precon_type, options.linsol_max_iterations); +#endif + } + + + + IDASetLinearSolver(ida_mem, LS, J); + + if (options.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_casadi_approx, NULL); + } + + if (options.jacobian == "matrix-free") + { + IDASetJacTimes(ida_mem, NULL, jtimes_casadi); + } + else if (options.jacobian != "none") + { + IDASetJacFn(ida_mem, jacobian_casadi); + } + + if (number_of_parameters > 0) + { + IDASensInit(ida_mem, number_of_parameters, IDA_SIMULTANEOUS, + sensitivities_casadi, yyS, ypS); + IDASensEEtolerances(ida_mem); + } + + 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++) + { + id_val[ii] = id_np_val[ii]; + } + + IDASetId(ida_mem, id); +} + +CasadiSolver::~CasadiSolver() +{ + + /* Free memory */ + if (number_of_parameters > 0) + { + IDASensFree(ida_mem); + } + SUNLinSolFree(LS); + SUNMatDestroy(J); + N_VDestroy(avtol); + N_VDestroy(yy); + N_VDestroy(yp); + N_VDestroy(id); + if (number_of_parameters > 0) + { + N_VDestroyVectorArray(yyS, number_of_parameters); + N_VDestroyVectorArray(ypS, number_of_parameters); + } + + IDAFree(&ida_mem); +#if SUNDIALS_VERSION_MAJOR >= 6 + SUNContext_Free(&sunctx); +#endif +} + +Solution CasadiSolver::solve(np_array t_np, np_array y0_np, np_array yp0_np, + np_array_dense inputs) +{ + DEBUG("CasadiSolver::solve"); + int number_of_timesteps = t_np.request().size; + + // set inputs + auto p_inputs = inputs.unchecked<2>(); + for (int i = 0; i < functions->inputs.size(); i++) + { + functions->inputs[i] = p_inputs(i, 0); + } + + realtype *yval = N_VGetArrayPointer(yy); + realtype *ypval = N_VGetArrayPointer(yp); + std::vector ySval(number_of_parameters); + for (int is = 0 ; is < number_of_parameters; is++) { + ySval[is] = N_VGetArrayPointer(yyS[is]); + N_VConst(RCONST(0.0), yyS[is]); + N_VConst(RCONST(0.0), ypS[is]); + } + + auto t = t_np.unchecked<1>(); + auto y0 = y0_np.unchecked<1>(); + auto yp0 = yp0_np.unchecked<1>(); + for (int i = 0; i < number_of_states; i++) + { + yval[i] = y0[i]; + ypval[i] = yp0[i]; + } + + realtype t0 = RCONST(t(0)); + IDAReInit(ida_mem, t0, yy, yp); + + int t_i = 1; + realtype tret; + realtype t_next; + realtype t_final = t(number_of_timesteps - 1); + + // set return vectors + realtype *t_return = new realtype[number_of_timesteps]; + realtype *y_return = new realtype[number_of_timesteps * number_of_states]; + realtype *yS_return = new realtype[number_of_parameters * + number_of_timesteps * number_of_states]; + + py::capsule free_t_when_done(t_return, + [](void *f) + { + realtype *vect = + reinterpret_cast(f); + delete[] vect; + }); + py::capsule free_y_when_done(y_return, + [](void *f) + { + realtype *vect = + reinterpret_cast(f); + delete[] vect; + }); + py::capsule free_yS_when_done(yS_return, + [](void *f) + { + realtype *vect = + reinterpret_cast(f); + delete[] vect; + }); + + t_return[0] = t(0); + for (int j = 0; j < number_of_states; j++) + { + y_return[j] = yval[j]; + } + for (int j = 0; j < number_of_parameters; j++) + { + const int base_index = j * number_of_timesteps * number_of_states; + for (int k = 0; k < number_of_states; k++) + { + yS_return[base_index + k] = ySval[j][k]; + } + } + + // calculate consistent initial conditions + DEBUG("IDACalcIC"); + IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); + + int retval; + while (true) + { + 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 || retval == IDA_SUCCESS || + retval == IDA_ROOT_RETURN) + { + if (number_of_parameters > 0) + { + IDAGetSens(ida_mem, &tret, yyS); + } + + t_return[t_i] = tret; + 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++) + { + yS_return[base_index + k] = ySval[j][k]; + } + } + t_i += 1; + if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) + { + break; + } + } + else + { + // failed + break; + } + } + + np_array t_ret = np_array(t_i, &t_return[0], free_t_when_done); + np_array y_ret = + np_array(t_i * number_of_states, &y_return[0], free_y_when_done); + np_array yS_ret = np_array( + std::vector{number_of_parameters, number_of_timesteps, number_of_states}, + &yS_return[0], free_yS_when_done); + + Solution sol(retval, t_ret, y_ret, yS_ret); + + if (options.print_stats) + { + long nsteps, nrevals, nlinsetups, netfails; + int klast, kcur; + realtype hinused, hlast, hcur, tcur; + + IDAGetIntegratorStats(ida_mem, &nsteps, &nrevals, &nlinsetups, &netfails, + &klast, &kcur, &hinused, &hlast, &hcur, &tcur); + + long nniters, nncfails; + IDAGetNonlinSolvStats(ida_mem, &nniters, &nncfails); + + long int ngevalsBBDP = 0; + if (options.using_iterative_solver) + { + IDABBDPrecGetNumGfnEvals(ida_mem, &ngevalsBBDP); + } + + py::print("Solver Stats:"); + py::print("\tNumber of steps =", nsteps); + py::print("\tNumber of calls to residual function =", nrevals); + py::print("\tNumber of calls to residual function in preconditioner =", + ngevalsBBDP); + py::print("\tNumber of linear solver setup calls =", nlinsetups); + py::print("\tNumber of error test failures =", netfails); + py::print("\tMethod order used on last step =", klast); + py::print("\tMethod order used on next step =", kcur); + py::print("\tInitial step size =", hinused); + py::print("\tStep size on last step =", hlast); + py::print("\tStep size on next step =", hcur); + py::print("\tCurrent internal time reached =", tcur); + py::print("\tNumber of nonlinear iterations performed =", nniters); + py::print("\tNumber of nonlinear convergence failures =", nncfails); + } + + return sol; +} diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp new file mode 100644 index 0000000000..3eed122e04 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -0,0 +1,57 @@ +#ifndef PYBAMM_IDAKLU_CASADI_SOLVER_HPP +#define PYBAMM_IDAKLU_CASADI_SOLVER_HPP + +#include +using Function = casadi::Function; + +#include "casadi_functions.hpp" +#include "common.hpp" +#include "options.hpp" +#include "solution.hpp" + +class CasadiSolver +{ +public: + CasadiSolver(np_array atol_np, double rel_tol, np_array rhs_alg_id, + int number_of_parameters, int number_of_events, + int jac_times_cjmass_nnz, + std::unique_ptr functions, const Options& options); + ~CasadiSolver(); + + void *ida_mem; // pointer to memory + +#if SUNDIALS_VERSION_MAJOR >= 6 + SUNContext sunctx; +#endif + + int number_of_states; + int number_of_parameters; + int number_of_events; + N_Vector yy, yp, avtol; // y, y', and absolute tolerance + N_Vector *yyS, *ypS; // y, y' for sensitivities + N_Vector id; // rhs_alg_id + realtype rtol; + const int jac_times_cjmass_nnz; + + SUNMatrix J; + SUNLinearSolver LS; + + std::unique_ptr functions; + Options options; + + Solution solve(np_array t_np, np_array y0_np, np_array yp0_np, + np_array_dense inputs); +}; + +CasadiSolver * +create_casadi_solver(int number_of_states, int number_of_parameters, + const Function &rhs_alg, const Function &jac_times_cjmass, + const np_array_int &jac_times_cjmass_colptrs, + const np_array_int &jac_times_cjmass_rowvals, + const int jac_times_cjmass_nnz, const Function &jac_action, + const Function &mass_action, const Function &sens, + const Function &event, const int number_of_events, + np_array rhs_alg_id, np_array atol_np, + double rel_tol, int inputs_length, py::dict options); + +#endif // PYBAMM_IDAKLU_CASADI_SOLVER_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp new file mode 100644 index 0000000000..ce2a892725 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -0,0 +1,287 @@ +#include "casadi_sundials_functions.hpp" +#include "casadi_functions.hpp" +#include "common.hpp" + +int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, + void *user_data) +{ + DEBUG("residual_casadi"); + CasadiFunctions *p_python_functions = + static_cast(user_data); + + p_python_functions->rhs_alg.m_arg[0] = &tres; + p_python_functions->rhs_alg.m_arg[1] = NV_DATA_S(yy); + p_python_functions->rhs_alg.m_arg[2] = p_python_functions->inputs.data(); + p_python_functions->rhs_alg.m_res[0] = NV_DATA_S(rr); + p_python_functions->rhs_alg(); + + realtype *tmp = p_python_functions->get_tmp(); + p_python_functions->mass_action.m_arg[0] = NV_DATA_S(yp); + p_python_functions->mass_action.m_res[0] = tmp; + p_python_functions->mass_action(); + + // AXPY: y <- a*x + y + const int ns = p_python_functions->number_of_states; + casadi::casadi_axpy(ns, -1., tmp, NV_DATA_S(rr)); + + DEBUG_VECTOR(yy); + DEBUG_VECTOR(yp); + DEBUG_VECTOR(rr); + + // now rr has rhs_alg(t, y) - mass_matrix * yp + return 0; +} + +// This Gres function computes G(t, y, yp). It loads the vector gval as a +// function of tt, yy, and yp. +// +// Arguments: +// Nlocal – is the local vector length. +// +// tt – is the value of the independent variable. +// +// yy – is the dependent variable. +// +// yp – is the derivative of the dependent variable. +// +// gval – is the output vector. +// +// user_data – is a pointer to user data, the same as the user_data parameter +// passed to IDASetUserData(). +// +// Return value: +// +// An IDABBDLocalFn function type should return 0 to indicate success, 1 for a +// recoverable error, or -1 for a non-recoverable error. +// +// Notes: +// +// This function must assume that all inter-processor communication of data +// needed to calculate gval has already been done, and this data is accessible +// within user_data. +// +// The case where G is mathematically identical to F is allowed. +int residual_casadi_approx(sunindextype Nlocal, realtype tt, N_Vector yy, + N_Vector yp, N_Vector gval, void *user_data) +{ + DEBUG("residual_casadi_approx"); + + // Just use true residual for now + int result = residual_casadi(tt, yy, yp, gval, user_data); + return result; +} + +// Purpose This function computes the product Jv of the DAE system Jacobian J +// (or an approximation to it) and a given vector v, where J is defined by Eq. +// (2.6). +// J = ∂F/∂y + cj ∂F/∂y˙ +// Arguments tt is the current value of the independent variable. +// yy is the current value of the dependent variable vector, y(t). +// yp is the current value of ˙y(t). +// rr is the current value of the residual vector F(t, y, y˙). +// v is the vector by which the Jacobian must be multiplied to the right. +// Jv is the computed output vector. +// cj is the scalar in the system Jacobian, proportional to the inverse of +// the step +// size (α in Eq. (2.6) ). +// user data is a pointer to user data, the same as the user data parameter +// passed to +// IDASetUserData. +// tmp1 +// tmp2 are pointers to memory allocated for variables of type N Vector +// which can +// be used by IDALsJacTimesVecFn as temporary storage or work space. +int jtimes_casadi(realtype tt, N_Vector yy, N_Vector yp, N_Vector rr, + N_Vector v, N_Vector Jv, realtype cj, void *user_data, + N_Vector tmp1, N_Vector tmp2) +{ + DEBUG("jtimes_casadi"); + CasadiFunctions *p_python_functions = + static_cast(user_data); + + // Jv has ∂F/∂y v + p_python_functions->jac_action.m_arg[0] = &tt; + p_python_functions->jac_action.m_arg[1] = NV_DATA_S(yy); + p_python_functions->jac_action.m_arg[2] = p_python_functions->inputs.data(); + p_python_functions->jac_action.m_arg[3] = NV_DATA_S(v); + p_python_functions->jac_action.m_res[0] = NV_DATA_S(Jv); + p_python_functions->jac_action(); + + // tmp has -∂F/∂y˙ v + realtype *tmp = p_python_functions->get_tmp(); + p_python_functions->mass_action.m_arg[0] = NV_DATA_S(v); + p_python_functions->mass_action.m_res[0] = tmp; + p_python_functions->mass_action(); + + // AXPY: y <- a*x + y + // Jv has ∂F/∂y v + cj ∂F/∂y˙ v + const int ns = p_python_functions->number_of_states; + casadi::casadi_axpy(ns, -cj, tmp, NV_DATA_S(Jv)); + + return 0; +} + +// Arguments tt is the current value of the independent variable t. +// cj is the scalar in the system Jacobian, proportional to the inverse of the +// step +// size (α in Eq. (2.6) ). +// yy is the current value of the dependent variable vector, y(t). +// yp is the current value of ˙y(t). +// rr is the current value of the residual vector F(t, y, y˙). +// Jac is the output (approximate) Jacobian matrix (of type SUNMatrix), J = +// ∂F/∂y + cj ∂F/∂y˙. +// user data is a pointer to user data, the same as the user data parameter +// passed to +// IDASetUserData. +// tmp1 +// tmp2 +// tmp3 are pointers to memory allocated for variables of type N Vector which +// can +// be used by IDALsJacFn function as temporary storage or work space. +int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, + N_Vector resvec, SUNMatrix JJ, void *user_data, + N_Vector tempv1, N_Vector tempv2, N_Vector tempv3) +{ + DEBUG("jacobian_casadi"); + + CasadiFunctions *p_python_functions = + static_cast(user_data); + + // create pointer to jac data, column pointers, and row values + sunindextype *jac_colptrs; + sunindextype *jac_rowvals; + realtype *jac_data; + if (p_python_functions->options.using_sparse_matrix) + { + jac_colptrs = SUNSparseMatrix_IndexPointers(JJ); + jac_rowvals = SUNSparseMatrix_IndexValues(JJ); + jac_data = SUNSparseMatrix_Data(JJ); + } + else + { + jac_data = SUNDenseMatrix_Data(JJ); + } + + // args are t, y, cj, put result in jacobian data matrix + p_python_functions->jac_times_cjmass.m_arg[0] = &tt; + p_python_functions->jac_times_cjmass.m_arg[1] = NV_DATA_S(yy); + p_python_functions->jac_times_cjmass.m_arg[2] = + p_python_functions->inputs.data(); + p_python_functions->jac_times_cjmass.m_arg[3] = &cj; + p_python_functions->jac_times_cjmass.m_res[0] = jac_data; + p_python_functions->jac_times_cjmass(); + + if (p_python_functions->options.using_sparse_matrix) + { + // row vals and col ptrs + const int n_row_vals = p_python_functions->jac_times_cjmass_rowvals.size(); + auto p_jac_times_cjmass_rowvals = + p_python_functions->jac_times_cjmass_rowvals.data(); + + // just copy across row vals (do I need to do this every time?) + // (or just in the setup?) + for (int i = 0; i < n_row_vals; i++) + { + jac_rowvals[i] = p_jac_times_cjmass_rowvals[i]; + } + + const int n_col_ptrs = p_python_functions->jac_times_cjmass_colptrs.size(); + auto p_jac_times_cjmass_colptrs = + p_python_functions->jac_times_cjmass_colptrs.data(); + + // just copy across col ptrs (do I need to do this every time?) + for (int i = 0; i < n_col_ptrs; i++) + { + jac_colptrs[i] = p_jac_times_cjmass_colptrs[i]; + } + } + + return (0); +} + +int events_casadi(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, + void *user_data) +{ + CasadiFunctions *p_python_functions = + static_cast(user_data); + + // args are t, y, put result in events_ptr + p_python_functions->events.m_arg[0] = &t; + p_python_functions->events.m_arg[1] = NV_DATA_S(yy); + p_python_functions->events.m_arg[2] = p_python_functions->inputs.data(); + p_python_functions->events.m_res[0] = events_ptr; + p_python_functions->events(); + + return (0); +} + +// This function computes the sensitivity residual for all sensitivity +// equations. It must compute the vectors +// (∂F/∂y)s i (t)+(∂F/∂ ẏ) ṡ i (t)+(∂F/∂p i ) and store them in resvalS[i]. +// Ns is the number of sensitivities. +// t is the current value of the independent variable. +// yy is the current value of the state vector, y(t). +// yp is the current value of ẏ(t). +// resval contains the current value F of the original DAE residual. +// yS contains the current values of the sensitivities s i . +// ypS contains the current values of the sensitivity derivatives ṡ i . +// resvalS contains the output sensitivity residual vectors. +// Memory allocation for resvalS is handled within idas. +// user data is a pointer to user data. +// tmp1, tmp2, tmp3 are N Vectors of length N which can be used as +// temporary storage. +// +// Return value An IDASensResFn should return 0 if successful, +// a positive value if a recoverable error +// occurred (in which case idas will attempt to correct), +// or a negative value if it failed unrecoverably (in which case the integration +// is halted and IDA SRES FAIL is returned) +// +int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp, + N_Vector resval, N_Vector *yS, N_Vector *ypS, + N_Vector *resvalS, void *user_data, N_Vector tmp1, + N_Vector tmp2, N_Vector tmp3) +{ + + CasadiFunctions *p_python_functions = + static_cast(user_data); + + const int np = p_python_functions->number_of_parameters; + + // args are t, y put result in rr + p_python_functions->sens.m_arg[0] = &t; + p_python_functions->sens.m_arg[1] = NV_DATA_S(yy); + p_python_functions->sens.m_arg[2] = p_python_functions->inputs.data(); + for (int i = 0; i < np; i++) + { + p_python_functions->sens.m_res[i] = NV_DATA_S(resvalS[i]); + } + // resvalsS now has (∂F/∂p i ) + p_python_functions->sens(); + + for (int i = 0; i < np; i++) + { + // put (∂F/∂y)s i (t) in tmp + realtype *tmp = p_python_functions->get_tmp(); + p_python_functions->jac_action.m_arg[0] = &t; + p_python_functions->jac_action.m_arg[1] = NV_DATA_S(yy); + p_python_functions->jac_action.m_arg[2] = p_python_functions->inputs.data(); + p_python_functions->jac_action.m_arg[3] = NV_DATA_S(yS[i]); + p_python_functions->jac_action.m_res[0] = tmp; + p_python_functions->jac_action(); + + const int ns = p_python_functions->number_of_states; + casadi::casadi_axpy(ns, 1., tmp, NV_DATA_S(resvalS[i])); + + // put -(∂F/∂ ẏ) ṡ i (t) in tmp2 + p_python_functions->mass_action.m_arg[0] = NV_DATA_S(ypS[i]); + p_python_functions->mass_action.m_res[0] = tmp; + p_python_functions->mass_action(); + + // (∂F/∂y)s i (t)+(∂F/∂ ẏ) ṡ i (t)+(∂F/∂p i ) + // AXPY: y <- a*x + y + casadi::casadi_axpy(ns, -1., tmp, NV_DATA_S(resvalS[i])); + } + + return 0; +} diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.hpp new file mode 100644 index 0000000000..a2192030b4 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.hpp @@ -0,0 +1,27 @@ +#ifndef PYBAMM_IDAKLU_CASADI_SUNDIALS_FUNCTIONS_HPP +#define PYBAMM_IDAKLU_CASADI_SUNDIALS_FUNCTIONS_HPP + +#include "common.hpp" + +int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, + void *user_data); + +int jtimes_casadi(realtype tt, N_Vector yy, N_Vector yp, N_Vector rr, + N_Vector v, N_Vector Jv, realtype cj, void *user_data, + N_Vector tmp1, N_Vector tmp2); + +int events_casadi(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, + void *user_data); + +int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp, + N_Vector resval, N_Vector *yS, N_Vector *ypS, + N_Vector *resvalS, void *user_data, N_Vector tmp1, + N_Vector tmp2, N_Vector tmp3); + +int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, + N_Vector resvec, SUNMatrix JJ, void *user_data, + N_Vector tempv1, N_Vector tempv2, N_Vector tempv3); + +int residual_casadi_approx(sunindextype Nlocal, realtype tt, N_Vector yy, + N_Vector yp, N_Vector gval, void *user_data); +#endif // PYBAMM_IDAKLU_CASADI_SUNDIALS_FUNCTIONS_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/common.hpp b/pybamm/solvers/c_solvers/idaklu/common.hpp new file mode 100644 index 0000000000..0c69971918 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/common.hpp @@ -0,0 +1,60 @@ +#ifndef PYBAMM_IDAKLU_COMMON_HPP +#define PYBAMM_IDAKLU_COMMON_HPP + +#include /* prototypes for IDAS fcts., consts. */ +#include /* access to IDABBDPRE preconditioner */ + +#include /* access to serial N_Vector */ +#include /* defs. of SUNRabs, SUNRexp, etc. */ +#include /* defs. of SUNRabs, SUNRexp, etc. */ +#include /* defs. of realtype, sunindextype */ + + +#if SUNDIALS_VERSION_MAJOR >= 6 + #include +#endif + +#include /* access to KLU linear solver */ +#include /* access to dense linear solver */ +#include /* access to lapack linear solver */ +#include /* access to spbcgs iterative linear solver */ +#include +#include +#include + +#include /* access to sparse SUNMatrix */ +#include /* access to dense SUNMatrix */ + + + +#include +#include + +namespace py = pybind11; +using np_array = py::array_t; +using np_array_dense = py::array_t; +using np_array_int = py::array_t; + +#ifdef NDEBUG +#define DEBUG(x) +#else +#define DEBUG(x) do { std::cerr << __FILE__ << ':' << __LINE__ << ' ' << x << std::endl; } while (0) +#endif + +#ifdef NDEBUG +#define DEBUG_VECTOR(vector) +#else +#define DEBUG_VECTOR(vector) {\ + std::cout << #vector << " = ["; \ + auto array_ptr = N_VGetArrayPointer(vector); \ + auto N = N_VGetLength(vector); \ + for (int i = 0; i < N; i++) { \ + std::cout << array_ptr[i]; \ + if (i < N-1) { \ + std::cout << ", "; \ + } \ + } \ + std::cout << "]" << std::endl; } +#endif + +#endif // PYBAMM_IDAKLU_COMMON_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/options.cpp b/pybamm/solvers/c_solvers/idaklu/options.cpp new file mode 100644 index 0000000000..283fd48501 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/options.cpp @@ -0,0 +1,103 @@ +#include "options.hpp" +#include + + +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()) +{ + + using_sparse_matrix = true; + if (jacobian == "sparse") + { + } + else if (jacobian == "dense" || jacobian == "none") + { + using_sparse_matrix = false; + } + else if (jacobian == "matrix-free") + { + } + else + { + throw std::domain_error( + "Unknown jacobian type \""s + jacobian + + "\". Should be one of \"sparse\", \"dense\", \"matrix-free\" or \"none\"."s + ); + } + + using_iterative_solver = false; + if (linear_solver == "SUNLinSol_Dense" && (jacobian == "dense" || jacobian == "none")) + { + } + else if (linear_solver == "SUNLinSol_LapackDense" && (jacobian == "dense" || jacobian == "none")) + { + } + else if (linear_solver == "SUNLinSol_KLU" && jacobian == "sparse") + { + } + else if ((linear_solver == "SUNLinSol_SPBCGS" || + linear_solver == "SUNLinSol_SPFGMR" || + linear_solver == "SUNLinSol_SPGMR" || + linear_solver == "SUNLinSol_SPTFQMR") && + (jacobian == "sparse" || jacobian == "matrix-free")) + { + using_iterative_solver = true; + } + else if (jacobian == "sparse") + { + throw std::domain_error( + "Unknown linear solver or incompatible options: " + "jacobian = \"" + jacobian + "\" linear solver = \"" + linear_solver + + "\". For a sparse jacobian " + "please use the SUNLinSol_KLU linear solver" + ); + } + else if (jacobian == "matrix-free") + { + throw std::domain_error( + "Unknown linear solver or incompatible options. " + "jacobian = \"" + jacobian + "\" linear solver = \"" + linear_solver + + "\". For a matrix-free jacobian " + "please use one of the iterative linear solvers: \"SUNLinSol_SPBCGS\", " + "\"SUNLinSol_SPFGMR\", \"SUNLinSol_SPGMR\", or \"SUNLinSol_SPTFQMR\"." + ); + } + else if (jacobian == "none") + { + throw std::domain_error( + "Unknown linear solver or incompatible options: " + "jacobian = \"" + jacobian + "\" linear solver = \"" + linear_solver + + "\". For no jacobian please use the SUNLinSol_Dense solver" + ); + } + else + { + throw std::domain_error( + "Unknown linear solver or incompatible options. " + "jacobian = \"" + jacobian + "\" linear solver = \"" + linear_solver + "\"" + ); + } + + if (using_iterative_solver) + { + if (preconditioner != "none" && preconditioner != "BBDP") + { + throw std::domain_error( + "Unknown preconditioner \""s + preconditioner + + "\", use one of \"BBDP\" or \"none\""s + ); + } + } + else + { + preconditioner = "none"; + } +} diff --git a/pybamm/solvers/c_solvers/idaklu/options.hpp b/pybamm/solvers/c_solvers/idaklu/options.hpp new file mode 100644 index 0000000000..2fc807e48f --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/options.hpp @@ -0,0 +1,20 @@ +#ifndef PYBAMM_OPTIONS_HPP +#define PYBAMM_OPTIONS_HPP + +#include "common.hpp" + +struct Options { + bool print_stats; + bool using_sparse_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; + explicit Options(py::dict options); + +}; + +#endif // PYBAMM_OPTIONS_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/python.cpp b/pybamm/solvers/c_solvers/idaklu/python.cpp new file mode 100644 index 0000000000..a1803988d4 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/python.cpp @@ -0,0 +1,478 @@ +#include "common.hpp" +#include "python.hpp" +#include + +class PybammFunctions +{ +public: + int number_of_states; + int number_of_parameters; + int number_of_events; + + PybammFunctions(const residual_type &res, const jacobian_type &jac, + const sensitivities_type &sens, + const jac_get_type &get_jac_data_in, + const jac_get_type &get_jac_row_vals_in, + const jac_get_type &get_jac_col_ptrs_in, + const event_type &event, + const int n_s, int n_e, const int n_p, + const np_array &inputs) + : number_of_states(n_s), number_of_events(n_e), + number_of_parameters(n_p), + py_res(res), py_jac(jac), + py_sens(sens), + py_event(event), py_get_jac_data(get_jac_data_in), + py_get_jac_row_vals(get_jac_row_vals_in), + py_get_jac_col_ptrs(get_jac_col_ptrs_in), + inputs(inputs) + { + } + + np_array operator()(double t, np_array y, np_array yp) + { + return py_res(t, y, inputs, yp); + } + + np_array res(double t, np_array y, np_array yp) + { + return py_res(t, y, inputs, yp); + } + + void jac(double t, np_array y, double cj) + { + // this function evaluates the jacobian and sets it to be the attribute + // of a python class which can then be called by get_jac_data, + // get_jac_col_ptr, etc + py_jac(t, y, inputs, cj); + } + + void sensitivities( + std::vector& resvalS, + const double t, const np_array& y, const np_array& yp, + const std::vector& yS, const std::vector& ypS) + { + // this function evaluates the sensitivity equations required by IDAS, + // returning them in resvalS, which is preallocated as a numpy array + // of size (np, n), where n is the number of states and np is the number + // of parameters + // + // yS and ypS are also shape (np, n), y and yp are shape (n) + // + // dF/dy * s_i + dF/dyd * sd + dFdp_i for i in range(np) + py_sens(resvalS, t, y, inputs, yp, yS, ypS); + } + + np_array get_jac_data() { return py_get_jac_data(); } + + np_array get_jac_row_vals() { return py_get_jac_row_vals(); } + + np_array get_jac_col_ptrs() { return py_get_jac_col_ptrs(); } + + np_array events(double t, np_array y) { return py_event(t, y, inputs); } + +private: + residual_type py_res; + sensitivities_type py_sens; + jacobian_type py_jac; + event_type py_event; + jac_get_type py_get_jac_data; + jac_get_type py_get_jac_row_vals; + jac_get_type py_get_jac_col_ptrs; + const np_array &inputs; +}; + +int residual(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, + void *user_data) +{ + PybammFunctions *python_functions_ptr = + static_cast(user_data); + PybammFunctions python_functions = *python_functions_ptr; + + realtype *yval, *ypval, *rval; + yval = N_VGetArrayPointer(yy); + ypval = N_VGetArrayPointer(yp); + rval = N_VGetArrayPointer(rr); + + int n = python_functions.number_of_states; + py::array_t y_np = py::array_t(n, yval); + py::array_t yp_np = py::array_t(n, ypval); + + py::array_t r_np; + + r_np = python_functions.res(tres, y_np, yp_np); + + auto r_np_ptr = r_np.unchecked<1>(); + + // just copying data + int i; + for (i = 0; i < n; i++) + { + rval[i] = r_np_ptr[i]; + } + return 0; +} + +int jacobian(realtype tt, realtype cj, N_Vector yy, N_Vector yp, + N_Vector resvec, SUNMatrix JJ, void *user_data, N_Vector tempv1, + N_Vector tempv2, N_Vector tempv3) +{ + realtype *yval; + yval = N_VGetArrayPointer(yy); + + PybammFunctions *python_functions_ptr = + static_cast(user_data); + PybammFunctions python_functions = *python_functions_ptr; + + int n = python_functions.number_of_states; + py::array_t y_np = py::array_t(n, yval); + + // create pointer to jac data, column pointers, and row values + sunindextype *jac_colptrs = SUNSparseMatrix_IndexPointers(JJ); + sunindextype *jac_rowvals = SUNSparseMatrix_IndexValues(JJ); + realtype *jac_data = SUNSparseMatrix_Data(JJ); + + py::array_t jac_np_array; + + python_functions.jac(tt, y_np, cj); + + np_array jac_np_data = python_functions.get_jac_data(); + int n_data = jac_np_data.request().size; + auto jac_np_data_ptr = jac_np_data.unchecked<1>(); + + // just copy across data + int i; + for (i = 0; i < n_data; i++) + { + jac_data[i] = jac_np_data_ptr[i]; + } + + np_array jac_np_row_vals = python_functions.get_jac_row_vals(); + int n_row_vals = jac_np_row_vals.request().size; + + auto jac_np_row_vals_ptr = jac_np_row_vals.unchecked<1>(); + // just copy across row vals (this might be unneeded) + for (i = 0; i < n_row_vals; i++) + { + jac_rowvals[i] = jac_np_row_vals_ptr[i]; + } + + np_array jac_np_col_ptrs = python_functions.get_jac_col_ptrs(); + int n_col_ptrs = jac_np_col_ptrs.request().size; + auto jac_np_col_ptrs_ptr = jac_np_col_ptrs.unchecked<1>(); + + // just copy across col ptrs (this might be unneeded) + for (i = 0; i < n_col_ptrs; i++) + { + jac_colptrs[i] = jac_np_col_ptrs_ptr[i]; + } + + return (0); +} + +int events(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, + void *user_data) +{ + realtype *yval; + yval = N_VGetArrayPointer(yy); + + PybammFunctions *python_functions_ptr = + static_cast(user_data); + PybammFunctions python_functions = *python_functions_ptr; + + int number_of_events = python_functions.number_of_events; + int number_of_states = python_functions.number_of_states; + py::array_t y_np = py::array_t(number_of_states, yval); + + py::array_t events_np_array; + + events_np_array = python_functions.events(t, y_np); + + auto events_np_data_ptr = events_np_array.unchecked<1>(); + + // just copying data (figure out how to pass pointers later) + int i; + for (i = 0; i < number_of_events; i++) + { + events_ptr[i] = events_np_data_ptr[i]; + } + + return (0); +} + +int sensitivities(int Ns, realtype t, N_Vector yy, N_Vector yp, + N_Vector resval, N_Vector *yS, N_Vector *ypS, N_Vector *resvalS, + void *user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3) { +// This function computes the sensitivity residual for all sensitivity +// equations. It must compute the vectors +// (∂F/∂y)s i (t)+(∂F/∂ ẏ) ṡ i (t)+(∂F/∂p i ) and store them in resvalS[i]. +// Ns is the number of sensitivities. +// t is the current value of the independent variable. +// yy is the current value of the state vector, y(t). +// yp is the current value of ẏ(t). +// resval contains the current value F of the original DAE residual. +// yS contains the current values of the sensitivities s i . +// ypS contains the current values of the sensitivity derivatives ṡ i . +// resvalS contains the output sensitivity residual vectors. +// Memory allocation for resvalS is handled within idas. +// user data is a pointer to user data. +// tmp1, tmp2, tmp3 are N Vectors of length N which can be used as +// temporary storage. +// +// Return value An IDASensResFn should return 0 if successful, +// a positive value if a recoverable error +// occurred (in which case idas will attempt to correct), +// or a negative value if it failed unrecoverably (in which case the integration is halted and IDA SRES FAIL is returned) +// + PybammFunctions *python_functions_ptr = + static_cast(user_data); + PybammFunctions python_functions = *python_functions_ptr; + + int n = python_functions.number_of_states; + int np = python_functions.number_of_parameters; + + // memory managed by sundials, so pass a destructor that does nothing + auto state_vector_shape = std::vector{n, 1}; + np_array y_np = np_array(state_vector_shape, N_VGetArrayPointer(yy), + py::capsule(&yy, [](void* p) {})); + np_array yp_np = np_array(state_vector_shape, N_VGetArrayPointer(yp), + py::capsule(&yp, [](void* p) {})); + + std::vector yS_np(np); + for (int i = 0; i < np; i++) { + auto capsule = py::capsule(yS + i, [](void* p) {}); + yS_np[i] = np_array(state_vector_shape, N_VGetArrayPointer(yS[i]), capsule); + } + + std::vector ypS_np(np); + for (int i = 0; i < np; i++) { + auto capsule = py::capsule(ypS + i, [](void* p) {}); + ypS_np[i] = np_array(state_vector_shape, N_VGetArrayPointer(ypS[i]), capsule); + } + + std::vector resvalS_np(np); + for (int i = 0; i < np; i++) { + auto capsule = py::capsule(resvalS + i, [](void* p) {}); + resvalS_np[i] = np_array(state_vector_shape, + N_VGetArrayPointer(resvalS[i]), capsule); + } + + realtype *ptr1 = static_cast(resvalS_np[0].request().ptr); + const realtype* resvalSval = N_VGetArrayPointer(resvalS[0]); + + python_functions.sensitivities(resvalS_np, t, y_np, yp_np, yS_np, ypS_np); + + return 0; +} + +/* main program */ +Solution solve_python(np_array t_np, np_array y0_np, np_array yp0_np, + residual_type res, jacobian_type jac, + sensitivities_type sens, + jac_get_type gjd, jac_get_type gjrv, jac_get_type gjcp, + int nnz, event_type event, + int number_of_events, int use_jacobian, np_array rhs_alg_id, + np_array atol_np, double rel_tol, np_array inputs, + int number_of_parameters) +{ + auto t = t_np.unchecked<1>(); + auto y0 = y0_np.unchecked<1>(); + auto yp0 = yp0_np.unchecked<1>(); + auto atol = atol_np.unchecked<1>(); + + int number_of_states = y0_np.request().size; + int number_of_timesteps = t_np.request().size; + void *ida_mem; // pointer to memory + N_Vector yy, yp, avtol; // y, y', and absolute tolerance + N_Vector *yyS, *ypS; // y, y' for sensitivities + N_Vector id; + realtype rtol, *yval, *ypval, *atval; + std::vector ySval(number_of_parameters); + int retval; + SUNMatrix J; + SUNLinearSolver LS; + +#if SUNDIALS_VERSION_MAJOR >= 6 + SUNContext sunctx; + SUNContext_Create(NULL, &sunctx); + + // allocate memory for solver + ida_mem = IDACreate(sunctx); + + // allocate vectors + yy = N_VNew_Serial(number_of_states, sunctx); + yp = N_VNew_Serial(number_of_states, sunctx); + avtol = N_VNew_Serial(number_of_states, sunctx); + id = N_VNew_Serial(number_of_states, sunctx); +#else + // allocate memory for solver + ida_mem = IDACreate(); + + // allocate vectors + yy = N_VNew_Serial(number_of_states); + yp = N_VNew_Serial(number_of_states); + avtol = N_VNew_Serial(number_of_states); + id = N_VNew_Serial(number_of_states); +#endif + + if (number_of_parameters > 0) { + yyS = N_VCloneVectorArray(number_of_parameters, yy); + ypS = N_VCloneVectorArray(number_of_parameters, yp); + } + + // set initial value + yval = N_VGetArrayPointer(yy); + ypval = N_VGetArrayPointer(yp); + atval = N_VGetArrayPointer(avtol); + int i; + for (i = 0; i < number_of_states; i++) + { + yval[i] = y0[i]; + ypval[i] = yp0[i]; + atval[i] = atol[i]; + } + + for (int is = 0 ; is < number_of_parameters; is++) { + ySval[is] = N_VGetArrayPointer(yyS[is]); + N_VConst(RCONST(0.0), yyS[is]); + N_VConst(RCONST(0.0), ypS[is]); + } + + // initialise solver + realtype t0 = RCONST(t(0)); + IDAInit(ida_mem, residual, t0, yy, yp); + + // set tolerances + rtol = RCONST(rel_tol); + + IDASVtolerances(ida_mem, rtol, avtol); + + // set events + IDARootInit(ida_mem, number_of_events, events); + + // set pybamm functions by passing pointer to it + PybammFunctions pybamm_functions(res, jac, sens, gjd, gjrv, gjcp, event, + number_of_states, number_of_events, + number_of_parameters, inputs); + void *user_data = &pybamm_functions; + IDASetUserData(ida_mem, user_data); + + // set linear solver +#if SUNDIALS_VERSION_MAJOR >= 6 + J = SUNSparseMatrix(number_of_states, number_of_states, nnz, CSR_MAT, sunctx); + LS = SUNLinSol_KLU(yy, J, sunctx); +#else + J = SUNSparseMatrix(number_of_states, number_of_states, nnz, CSR_MAT); + LS = SUNLinSol_KLU(yy, J); +#endif + + IDASetLinearSolver(ida_mem, LS, J); + + if (use_jacobian == 1) + { + IDASetJacFn(ida_mem, jacobian); + } + + if (number_of_parameters > 0) + { + IDASensInit(ida_mem, number_of_parameters, + IDA_SIMULTANEOUS, sensitivities, yyS, ypS); + IDASensEEtolerances(ida_mem); + } + + int t_i = 1; + realtype tret; + realtype t_next; + realtype t_final = t(number_of_timesteps - 1); + + // set return vectors + std::vector t_return(number_of_timesteps); + std::vector y_return(number_of_timesteps * number_of_states); + std::vector yS_return(number_of_parameters * number_of_timesteps * number_of_states); + + t_return[0] = t(0); + for (int j = 0; j < number_of_states; j++) + { + y_return[j] = yval[j]; + } + for (int j = 0; j < number_of_parameters; j++) { + const int base_index = j * number_of_timesteps * number_of_states; + for (int k = 0; k < number_of_states; k++) { + yS_return[base_index + k] = ySval[j][k]; + } + } + + // calculate consistent initial conditions + 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++) + { + id_val[ii] = id_np_val[ii]; + } + + IDASetId(ida_mem, id); + IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); + + while (true) + { + t_next = t(t_i); + IDASetStopTime(ida_mem, t_next); + retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); + + if (retval == IDA_TSTOP_RETURN || retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) + { + if (number_of_parameters > 0) { + IDAGetSens(ida_mem, &tret, yyS); + } + + t_return[t_i] = tret; + 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++) { + yS_return[base_index + k] = ySval[j][k]; + } + } + t_i += 1; + if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { + break; + } + + } + } + + /* Free memory */ + if (number_of_parameters > 0) { + IDASensFree(ida_mem); + } + IDAFree(&ida_mem); + SUNLinSolFree(LS); + SUNMatDestroy(J); + N_VDestroy(avtol); + N_VDestroy(yp); + if (number_of_parameters > 0) { + N_VDestroyVectorArray(yyS, number_of_parameters); + N_VDestroyVectorArray(ypS, number_of_parameters); + } +#if SUNDIALS_VERSION_MAJOR >= 6 + SUNContext_Free(&sunctx); +#endif + + np_array t_ret = np_array(t_i, &t_return[0]); + np_array y_ret = np_array(t_i * number_of_states, &y_return[0]); + np_array yS_ret = np_array( + std::vector{number_of_parameters, number_of_timesteps, number_of_states}, + &yS_return[0] + ); + + Solution sol(retval, t_ret, y_ret, yS_ret); + + return sol; +} + diff --git a/pybamm/solvers/c_solvers/idaklu.hpp b/pybamm/solvers/c_solvers/idaklu/python.hpp similarity index 97% rename from pybamm/solvers/c_solvers/idaklu.hpp rename to pybamm/solvers/c_solvers/idaklu/python.hpp index a39f149252..8c29bbc496 100644 --- a/pybamm/solvers/c_solvers/idaklu.hpp +++ b/pybamm/solvers/c_solvers/idaklu/python.hpp @@ -1,7 +1,7 @@ #ifndef PYBAMM_IDAKLU_HPP #define PYBAMM_IDAKLU_HPP -#include "idaklu_python.hpp" +#include "common.hpp" #include "solution.hpp" #include diff --git a/pybamm/solvers/c_solvers/solution.cpp b/pybamm/solvers/c_solvers/idaklu/solution.cpp similarity index 100% rename from pybamm/solvers/c_solvers/solution.cpp rename to pybamm/solvers/c_solvers/idaklu/solution.cpp diff --git a/pybamm/solvers/c_solvers/solution.hpp b/pybamm/solvers/c_solvers/idaklu/solution.hpp similarity index 62% rename from pybamm/solvers/c_solvers/solution.hpp rename to pybamm/solvers/c_solvers/idaklu/solution.hpp index 08192f3c49..047ae6ef8e 100644 --- a/pybamm/solvers/c_solvers/solution.hpp +++ b/pybamm/solvers/c_solvers/idaklu/solution.hpp @@ -1,7 +1,7 @@ -#ifndef PYBAMM_SOLUTION_HPP -#define PYBAMM_SOLUTION_HPP +#ifndef PYBAMM_IDAKLU_SOLUTION_HPP +#define PYBAMM_IDAKLU_SOLUTION_HPP -#include "idaklu_python.hpp" +#include "common.hpp" class Solution { @@ -17,4 +17,4 @@ class Solution np_array yS; }; -#endif // PYBAMM_SOLUTION_HPP +#endif // PYBAMM_IDAKLU_COMMON_HPP diff --git a/pybamm/solvers/c_solvers/idaklu_casadi.cpp b/pybamm/solvers/c_solvers/idaklu_casadi.cpp deleted file mode 100644 index f611c6ebc3..0000000000 --- a/pybamm/solvers/c_solvers/idaklu_casadi.cpp +++ /dev/null @@ -1,582 +0,0 @@ -#include "idaklu_casadi.hpp" -#include "idaklu_python.hpp" - -#include - -#include -using casadi::casadi_axpy; -using Sparsity = casadi::Sparsity; - -class CasadiFunction { -public: - explicit CasadiFunction(const Function &f):m_func(f) { - size_t sz_arg; - size_t sz_res; - size_t sz_iw; - size_t sz_w; - m_func.sz_work(sz_arg, sz_res, sz_iw, sz_w); - m_arg.resize(sz_arg); - m_res.resize(sz_res); - m_iw.resize(sz_iw); - m_w.resize(sz_w); - } - - // only call this once m_arg and m_res have been set appropriatelly - void operator()() { - int mem = m_func.checkout(); - m_func(m_arg.data(), m_res.data(), m_iw.data(), m_w.data(), mem); - m_func.release(mem); - } - -public: - std::vector m_arg; - std::vector m_res; - -private: - const Function &m_func; - std::vector m_iw; - std::vector m_w; -}; - - -class CasadiFunctions { -public: - int number_of_states; - int number_of_parameters; - int number_of_events; - int number_of_nnz; - CasadiFunction rhs_alg; - CasadiFunction sens; - CasadiFunction jac_times_cjmass; - const np_array_int &jac_times_cjmass_rowvals; - const np_array_int &jac_times_cjmass_colptrs; - const np_array_dense &inputs; - CasadiFunction jac_action; - CasadiFunction mass_action; - CasadiFunction events; - - CasadiFunctions(const Function &rhs_alg, - const Function &jac_times_cjmass, - const int jac_times_cjmass_nnz, - const np_array_int &jac_times_cjmass_rowvals, - const np_array_int &jac_times_cjmass_colptrs, - const np_array_dense &inputs, - const Function &jac_action, - const Function &mass_action, - const Function &sens, - const Function &events, - const int n_s, int n_e, const int n_p) - : number_of_states(n_s), number_of_events(n_e), - number_of_parameters(n_p), - number_of_nnz(jac_times_cjmass_nnz), - rhs_alg(rhs_alg), - jac_times_cjmass(jac_times_cjmass), - jac_times_cjmass_rowvals(jac_times_cjmass_rowvals), - jac_times_cjmass_colptrs(jac_times_cjmass_colptrs), - inputs(inputs), - jac_action(jac_action), - mass_action(mass_action), - sens(sens), - events(events), - tmp(number_of_states) - {} - - realtype *get_tmp() { - return tmp.data(); - } - -private: - std::vector tmp; -}; - -int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, - void *user_data) -{ - CasadiFunctions *p_python_functions = - static_cast(user_data); - - py::buffer_info input_buf = p_python_functions->inputs.request(); - p_python_functions->rhs_alg.m_arg[0] = &tres; - p_python_functions->rhs_alg.m_arg[1] = NV_DATA_S(yy); - p_python_functions->rhs_alg.m_arg[2] = static_cast(input_buf.ptr); - p_python_functions->rhs_alg.m_res[0] = NV_DATA_S(rr); - p_python_functions->rhs_alg(); - - realtype *tmp = p_python_functions->get_tmp(); - - // args is yp, put result in tmp - p_python_functions->mass_action.m_arg[0] = NV_DATA_S(yp); - p_python_functions->mass_action.m_res[0] = tmp; - p_python_functions->mass_action(); - - // AXPY: y <- a*x + y - const int ns = p_python_functions->number_of_states; - casadi_axpy(ns, -1., tmp, NV_DATA_S(rr)); - - // now rr has rhs_alg(t, y) - mass_matrix * yp - return 0; -} - -// Purpose This function computes the product Jv of the DAE system Jacobian J -// (or an approximation to it) and a given vector v, where J is defined by Eq. (2.6). -// J = ∂F/∂y + cj ∂F/∂y˙ -// Arguments tt is the current value of the independent variable. -// yy is the current value of the dependent variable vector, y(t). -// yp is the current value of ˙y(t). -// rr is the current value of the residual vector F(t, y, y˙). -// v is the vector by which the Jacobian must be multiplied to the right. -// Jv is the computed output vector. -// cj is the scalar in the system Jacobian, proportional to the inverse of the step -// size (α in Eq. (2.6) ). -// user data is a pointer to user data, the same as the user data parameter passed to -// IDASetUserData. -// tmp1 -// tmp2 are pointers to memory allocated for variables of type N Vector which can -// be used by IDALsJacTimesVecFn as temporary storage or work space. -int jtimes_casadi(realtype tt, N_Vector yy, N_Vector yp, N_Vector rr, - N_Vector v, N_Vector Jv, realtype cj, void *user_data, - N_Vector tmp1, N_Vector tmp2) { - CasadiFunctions *p_python_functions = - static_cast(user_data); - - // rr has ∂F/∂y v - py::buffer_info input_buf = p_python_functions->inputs.request(); - p_python_functions->jac_action.m_arg[0] = &tt; - p_python_functions->jac_action.m_arg[1] = NV_DATA_S(yy); - p_python_functions->jac_action.m_arg[2] = static_cast(input_buf.ptr); - p_python_functions->jac_action.m_arg[3] = NV_DATA_S(v); - p_python_functions->jac_action.m_res[0] = NV_DATA_S(rr); - p_python_functions->jac_action(); - - // tmp has -∂F/∂y˙ v - realtype *tmp = p_python_functions->get_tmp(); - p_python_functions->mass_action.m_arg[0] = NV_DATA_S(v); - p_python_functions->mass_action.m_res[0] = tmp; - p_python_functions->mass_action(); - - // AXPY: y <- a*x + y - // rr has ∂F/∂y v + cj ∂F/∂y˙ v - const int ns = p_python_functions->number_of_states; - casadi_axpy(ns, -cj, tmp, NV_DATA_S(rr)); - - return 0; -} - - -// Arguments tt is the current value of the independent variable t. -// cj is the scalar in the system Jacobian, proportional to the inverse of the step -// size (α in Eq. (2.6) ). -// yy is the current value of the dependent variable vector, y(t). -// yp is the current value of ˙y(t). -// rr is the current value of the residual vector F(t, y, y˙). -// Jac is the output (approximate) Jacobian matrix (of type SUNMatrix), J = -// ∂F/∂y + cj ∂F/∂y˙. -// user data is a pointer to user data, the same as the user data parameter passed to -// IDASetUserData. -// tmp1 -// tmp2 -// tmp3 are pointers to memory allocated for variables of type N Vector which can -// be used by IDALsJacFn function as temporary storage or work space. -int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, - N_Vector resvec, SUNMatrix JJ, void *user_data, N_Vector tempv1, - N_Vector tempv2, N_Vector tempv3) { - - CasadiFunctions *p_python_functions = - static_cast(user_data); - - // create pointer to jac data, column pointers, and row values - sunindextype *jac_colptrs = SUNSparseMatrix_IndexPointers(JJ); - sunindextype *jac_rowvals = SUNSparseMatrix_IndexValues(JJ); - realtype *jac_data = SUNSparseMatrix_Data(JJ); - - // args are t, y, cj, put result in jacobian data matrix - py::buffer_info input_buf = p_python_functions->inputs.request(); - p_python_functions->jac_times_cjmass.m_arg[0] = &tt; - p_python_functions->jac_times_cjmass.m_arg[1] = NV_DATA_S(yy); - p_python_functions->jac_times_cjmass.m_arg[2] = static_cast(input_buf.ptr); - p_python_functions->jac_times_cjmass.m_arg[3] = &cj; - p_python_functions->jac_times_cjmass.m_res[0] = jac_data; - p_python_functions->jac_times_cjmass(); - - // row vals and col ptrs - const np_array &jac_times_cjmass_rowvals = p_python_functions->jac_times_cjmass_rowvals; - const int n_row_vals = jac_times_cjmass_rowvals.request().size; - auto p_jac_times_cjmass_rowvals = jac_times_cjmass_rowvals.unchecked<1>(); - - // just copy across row vals (do I need to do this every time?) - // (or just in the setup?) - for (int i = 0; i < n_row_vals; i++) { - jac_rowvals[i] = p_jac_times_cjmass_rowvals[i]; - } - - const np_array &jac_times_cjmass_colptrs = p_python_functions->jac_times_cjmass_colptrs; - const int n_col_ptrs = jac_times_cjmass_colptrs.request().size; - auto p_jac_times_cjmass_colptrs = jac_times_cjmass_colptrs.unchecked<1>(); - - // just copy across col ptrs (do I need to do this every time?) - for (int i = 0; i < n_col_ptrs; i++) { - jac_colptrs[i] = p_jac_times_cjmass_colptrs[i]; - } - - return (0); -} - -int events_casadi(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, - void *user_data) -{ - CasadiFunctions *p_python_functions = - static_cast(user_data); - - // args are t, y, put result in events_ptr - py::buffer_info input_buf = p_python_functions->inputs.request(); - p_python_functions->events.m_arg[0] = &t; - p_python_functions->events.m_arg[1] = NV_DATA_S(yy); - p_python_functions->events.m_arg[2] = static_cast(input_buf.ptr); - p_python_functions->events.m_res[0] = events_ptr; - p_python_functions->events(); - - return (0); -} - -// This function computes the sensitivity residual for all sensitivity -// equations. It must compute the vectors -// (∂F/∂y)s i (t)+(∂F/∂ ẏ) ṡ i (t)+(∂F/∂p i ) and store them in resvalS[i]. -// Ns is the number of sensitivities. -// t is the current value of the independent variable. -// yy is the current value of the state vector, y(t). -// yp is the current value of ẏ(t). -// resval contains the current value F of the original DAE residual. -// yS contains the current values of the sensitivities s i . -// ypS contains the current values of the sensitivity derivatives ṡ i . -// resvalS contains the output sensitivity residual vectors. -// Memory allocation for resvalS is handled within idas. -// user data is a pointer to user data. -// tmp1, tmp2, tmp3 are N Vectors of length N which can be used as -// temporary storage. -// -// Return value An IDASensResFn should return 0 if successful, -// a positive value if a recoverable error -// occurred (in which case idas will attempt to correct), -// or a negative value if it failed unrecoverably (in which case the integration is halted and IDA SRES FAIL is returned) -// -int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp, - N_Vector resval, N_Vector *yS, N_Vector *ypS, N_Vector *resvalS, - void *user_data, N_Vector tmp1, N_Vector tmp2, N_Vector tmp3) { - - CasadiFunctions *p_python_functions = - static_cast(user_data); - - const int np = p_python_functions->number_of_parameters; - - // args are t, y put result in rr - py::buffer_info input_buf = p_python_functions->inputs.request(); - p_python_functions->sens.m_arg[0] = &t; - p_python_functions->sens.m_arg[1] = NV_DATA_S(yy); - p_python_functions->sens.m_arg[2] = static_cast(input_buf.ptr); - for (int i = 0; i < np; i++) { - p_python_functions->sens.m_res[i] = NV_DATA_S(resvalS[i]); - } - // resvalsS now has (∂F/∂p i ) - p_python_functions->sens(); - - for (int i = 0; i < np; i++) { - // put (∂F/∂y)s i (t) in tmp - realtype *tmp = p_python_functions->get_tmp(); - p_python_functions->jac_action.m_arg[0] = &t; - p_python_functions->jac_action.m_arg[1] = NV_DATA_S(yy); - p_python_functions->jac_action.m_arg[2] = static_cast(input_buf.ptr); - p_python_functions->jac_action.m_arg[3] = NV_DATA_S(yS[i]); - p_python_functions->jac_action.m_res[0] = tmp; - p_python_functions->jac_action(); - - const int ns = p_python_functions->number_of_states; - casadi_axpy(ns, 1., tmp, NV_DATA_S(resvalS[i])); - - // put -(∂F/∂ ẏ) ṡ i (t) in tmp2 - p_python_functions->mass_action.m_arg[0] = NV_DATA_S(ypS[i]); - p_python_functions->mass_action.m_res[0] = tmp; - p_python_functions->mass_action(); - - // (∂F/∂y)s i (t)+(∂F/∂ ẏ) ṡ i (t)+(∂F/∂p i ) - // AXPY: y <- a*x + y - casadi_axpy(ns, -1., tmp, NV_DATA_S(resvalS[i])); - } - - return 0; -} - - - -/* main program */ - -Solution solve_casadi(np_array t_np, np_array y0_np, np_array yp0_np, - const Function &rhs_alg, - const Function &jac_times_cjmass, - const np_array_int &jac_times_cjmass_colptrs, - const np_array_int &jac_times_cjmass_rowvals, - const int jac_times_cjmass_nnz, - const Function &jac_action, - const Function &mass_action, - const Function &sens, - const Function &events, - const int number_of_events, - int use_jacobian, - np_array rhs_alg_id, - np_array atol_np, double rel_tol, - np_array_dense inputs, - int number_of_parameters) -{ - - auto t = t_np.unchecked<1>(); - auto y0 = y0_np.unchecked<1>(); - auto yp0 = yp0_np.unchecked<1>(); - auto atol = atol_np.unchecked<1>(); - - int number_of_states = y0_np.request().size; - int number_of_timesteps = t_np.request().size; - void *ida_mem; // pointer to memory - N_Vector yy, yp, avtol; // y, y', and absolute tolerance - N_Vector *yyS, *ypS; // y, y' for sensitivities - realtype rtol, *yval, *ypval, *atval; - std::vector ySval(number_of_parameters); - int retval; - SUNMatrix J; - SUNLinearSolver LS; - - // allocate vectors - yy = N_VNew_Serial(number_of_states); - yp = N_VNew_Serial(number_of_states); - avtol = N_VNew_Serial(number_of_states); - - if (number_of_parameters > 0) { - yyS = N_VCloneVectorArray(number_of_parameters, yy); - ypS = N_VCloneVectorArray(number_of_parameters, yp); - } - - // set initial value - yval = N_VGetArrayPointer(yy); - ypval = N_VGetArrayPointer(yp); - atval = N_VGetArrayPointer(avtol); - int i; - for (i = 0; i < number_of_states; i++) - { - yval[i] = y0[i]; - ypval[i] = yp0[i]; - atval[i] = atol[i]; - } - - for (int is = 0 ; is < number_of_parameters; is++) { - ySval[is] = N_VGetArrayPointer(yyS[is]); - N_VConst(RCONST(0.0), yyS[is]); - N_VConst(RCONST(0.0), ypS[is]); - } - - // allocate memory for solver - ida_mem = IDACreate(); - - // initialise solver - realtype t0 = RCONST(t(0)); - IDAInit(ida_mem, residual_casadi, t0, yy, yp); - - // set tolerances - rtol = RCONST(rel_tol); - - IDASVtolerances(ida_mem, rtol, avtol); - - // set events - IDARootInit(ida_mem, number_of_events, events_casadi); - - // set pybamm functions by passing pointer to it - CasadiFunctions pybamm_functions = CasadiFunctions( - rhs_alg, - jac_times_cjmass, - jac_times_cjmass_nnz, - jac_times_cjmass_rowvals, - jac_times_cjmass_colptrs, - inputs, - jac_action, mass_action, - sens, events, - number_of_states, number_of_events, - number_of_parameters); - - void *user_data = &pybamm_functions; - IDASetUserData(ida_mem, user_data); - - // set linear solver - if (use_jacobian == 1) { - J = SUNSparseMatrix(number_of_states, number_of_states, jac_times_cjmass_nnz, CSC_MAT); - LS = SUNLinSol_KLU(yy, J); - } else { - J = SUNDenseMatrix(number_of_states, number_of_states); - LS = SUNLinSol_Dense(yy, J); - } - - IDASetLinearSolver(ida_mem, LS, J); - - if (use_jacobian == 1) { - IDASetJacFn(ida_mem, jacobian_casadi); - } - - - if (number_of_parameters > 0) { - IDASensInit(ida_mem, number_of_parameters, - IDA_SIMULTANEOUS, sensitivities_casadi, yyS, ypS); - IDASensEEtolerances(ida_mem); - } - - SUNLinSolInitialize(LS); - - int t_i = 1; - realtype tret; - realtype t_next; - realtype t_final = t(number_of_timesteps - 1); - - // set return vectors - realtype* t_return = new realtype[number_of_timesteps]; - realtype* y_return = new realtype[number_of_timesteps * number_of_states]; - realtype* yS_return = new realtype[number_of_parameters * number_of_timesteps * number_of_states]; - - py::capsule free_t_when_done(t_return, [](void *f) { - realtype *vect = reinterpret_cast(f); - delete[] vect; - }); - py::capsule free_y_when_done(y_return, [](void *f) { - realtype *vect = reinterpret_cast(f); - delete[] vect; - }); - py::capsule free_yS_when_done(yS_return, [](void *f) { - realtype *vect = reinterpret_cast(f); - delete[] vect; - }); - - t_return[0] = t(0); - for (int j = 0; j < number_of_states; j++) - { - y_return[j] = yval[j]; - } - for (int j = 0; j < number_of_parameters; j++) { - const int base_index = j * number_of_timesteps * number_of_states; - for (int k = 0; k < number_of_states; k++) { - yS_return[base_index + k] = ySval[j][k]; - } - } - - // calculate consistent initial conditions - N_Vector id; - auto id_np_val = rhs_alg_id.unchecked<1>(); - id = N_VNew_Serial(number_of_states); - realtype *id_val; - id_val = N_VGetArrayPointer(id); - - int ii; - for (ii = 0; ii < number_of_states; ii++) - { - id_val[ii] = id_np_val[ii]; - } - - IDASetId(ida_mem, id); - IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); - - while (true) - { - t_next = t(t_i); - IDASetStopTime(ida_mem, t_next); - retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); - - if (retval == IDA_TSTOP_RETURN || retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { - if (number_of_parameters > 0) { - IDAGetSens(ida_mem, &tret, yyS); - } - - t_return[t_i] = tret; - 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++) { - yS_return[base_index + k] = ySval[j][k]; - } - } - t_i += 1; - if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { - break; - } - } else { - // failed - break; - } - } - - np_array t_ret = np_array(t_i, &t_return[0], free_t_when_done); - np_array y_ret = np_array(t_i * number_of_states, &y_return[0], free_y_when_done); - np_array yS_ret = np_array( - std::vector{number_of_parameters, number_of_timesteps, number_of_states}, - &yS_return[0], free_yS_when_done - ); - - Solution sol(retval, t_ret, y_ret, yS_ret); - - // TODO config input to choose stuff like this - const bool print_stats = false; - if (print_stats) { - long nsteps, nrevals, nlinsetups, netfails; - int klast, kcur; - realtype hinused, hlast, hcur, tcur; - - IDAGetIntegratorStats(ida_mem, - &nsteps, - &nrevals, - &nlinsetups, - &netfails, - &klast, - &kcur, - &hinused, - &hlast, - &hcur, - &tcur - ); - - long nniters, nncfails; - IDAGetNonlinSolvStats(ida_mem, &nniters, &nncfails); - - std::cout << "Solver Stats: \n" - << " Number of steps = " << nsteps << "\n" - << " Number of calls to residual function = " << nrevals << "\n" - << " Number of linear solver setup calls = " << nlinsetups << "\n" - << " Number of error test failures = " << netfails << "\n" - << " Method order used on last step = " << klast << "\n" - << " Method order used on next step = " << kcur << "\n" - << " Initial step size = " << hinused << "\n" - << " Step size on last step = " << hlast << "\n" - << " Step size on next step = " << hcur << "\n" - << " Current internal time reached = " << tcur << "\n" - << " Number of nonlinear iterations performed = " << nniters << "\n" - << " Number of nonlinear convergence failures = " << nncfails << "\n" - << std::endl; - } - - - - /* Free memory */ - if (number_of_parameters > 0) { - IDASensFree(ida_mem); - } - SUNLinSolFree(LS); - SUNMatDestroy(J); - N_VDestroy(avtol); - N_VDestroy(yy); - N_VDestroy(yp); - N_VDestroy(id); - if (number_of_parameters > 0) { - N_VDestroyVectorArray(yyS, number_of_parameters); - N_VDestroyVectorArray(ypS, number_of_parameters); - } - - IDAFree(&ida_mem); - - return sol; -} - diff --git a/pybamm/solvers/c_solvers/idaklu_casadi.hpp b/pybamm/solvers/c_solvers/idaklu_casadi.hpp deleted file mode 100644 index 6aa655392b..0000000000 --- a/pybamm/solvers/c_solvers/idaklu_casadi.hpp +++ /dev/null @@ -1,32 +0,0 @@ - -#ifndef PYBAMM_IDAKLU_CASADI_HPP -#define PYBAMM_IDAKLU_CASADI_HPP - -#include - -using Function = casadi::Function; - -#include "solution.hpp" - -Solution solve_casadi(np_array t_np, np_array y0_np, np_array yp0_np, - const Function &rhs_alg, - const Function &jac_times_cjmass, - const np_array_int &jac_times_cjmass_colptrs, - const np_array_int &jac_times_cjmass_rowvals, - const int jac_times_cjmass_nnz, - const Function &jac_action, - const Function &mass_action, - const Function &sens, - const Function &event, - const int number_of_events, - int use_jacobian, - np_array rhs_alg_id, - np_array atol_np, - double rel_tol, - np_array_dense inputs, - int number_of_parameters); - - - - -#endif // PYBAMM_IDAKLU_CASADI_HPP diff --git a/pybamm/solvers/c_solvers/idaklu_python.cpp b/pybamm/solvers/c_solvers/idaklu_python.cpp deleted file mode 100644 index 6e8446ace7..0000000000 --- a/pybamm/solvers/c_solvers/idaklu_python.cpp +++ /dev/null @@ -1,68 +0,0 @@ - -#include "idaklu.hpp" - -#include "idaklu_casadi.hpp" -#include "idaklu_python.hpp" - -#include -#include -#include -#include - -#include - -Function generate_function(const std::string& data) { - return Function::deserialize(data); -} - -namespace py = pybind11; - -PYBIND11_MAKE_OPAQUE(std::vector); - -PYBIND11_MODULE(idaklu, m) -{ - m.doc() = "sundials solvers"; // optional module docstring - - py::bind_vector>(m, "VectorNdArray"); - - m.def("solve_python", &solve_python, "The solve function for python evaluators", - py::arg("t"), py::arg("y0"), - py::arg("yp0"), py::arg("res"), py::arg("jac"), py::arg("sens"), - py::arg("get_jac_data"), - py::arg("get_jac_row_vals"), py::arg("get_jac_col_ptr"), py::arg("nnz"), - py::arg("events"), py::arg("number_of_events"), py::arg("use_jacobian"), - py::arg("rhs_alg_id"), py::arg("atol"), py::arg("rtol"), py::arg("inputs"), - py::arg("number_of_sensitivity_parameters"), - py::return_value_policy::take_ownership); - - m.def("solve_casadi", &solve_casadi, "The solve function for casadi evaluators", - py::arg("t"), py::arg("y0"), py::arg("yp0"), - py::arg("rhs_alg"), - py::arg("jac_times_cjmass"), - py::arg("jac_times_cjmass_colptrs"), - py::arg("jac_times_cjmass_rowvals"), - py::arg("jac_times_cjmass_nnz"), - py::arg("jac_action"), - py::arg("mass_action"), - py::arg("sens"), - py::arg("events"), py::arg("number_of_events"), - py::arg("use_jacobian"), - py::arg("rhs_alg_id"), - py::arg("atol"), py::arg("rtol"), py::arg("inputs"), - py::arg("number_of_sensitivity_parameters"), - py::return_value_policy::take_ownership); - - m.def("generate_function", &generate_function, "Generate a casadi function", - py::arg("string"), - py::return_value_policy::take_ownership); - - py::class_(m, "Function"); - - py::class_(m, "solution") - .def_readwrite("t", &Solution::t) - .def_readwrite("y", &Solution::y) - .def_readwrite("yS", &Solution::yS) - .def_readwrite("flag", &Solution::flag); -} - - diff --git a/pybamm/solvers/c_solvers/idaklu_python.hpp b/pybamm/solvers/c_solvers/idaklu_python.hpp deleted file mode 100644 index e7583da371..0000000000 --- a/pybamm/solvers/c_solvers/idaklu_python.hpp +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef PYBAMM_IDAKLU_PYTHON_HPP -#define PYBAMM_IDAKLU_PYTHON_HPP - -#include /* prototypes for IDAS fcts., consts. */ -#include /* access to serial N_Vector */ -#include /* defs. of SUNRabs, SUNRexp, etc. */ -#include /* defs. of realtype, sunindextype */ -#include /* access to KLU linear solver */ -#include /* access to dense linear solver */ -#include /* access to sparse SUNMatrix */ -#include /* access to dense SUNMatrix */ - -#include - -namespace py = pybind11; -using np_array = py::array_t; -using np_array_dense = py::array_t; -using np_array_int = py::array_t; - - -#endif // PYBAMM_IDAKLU_PYTHON_HPP diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 94316e049a..8998a3f6a6 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -43,6 +43,41 @@ class IDAKLUSolver(pybamm.BaseSolver): The tolerance for the initial-condition solver (default is 1e-6). extrap_tol : float, optional The tolerance to assert whether extrapolation occurs or not (default is 0). + options: dict, optional + Addititional options to pass to the solver, by default: + + .. code-block:: python + + options = { + # print statistics of the solver after every solve + "print_stats": False, + + # jacobian form, can be "none", "dense", "sparse", "matrix-free" + "jacobian": "sparse", + + # name of sundials linear solver to use options are: "SUNLinSol_KLU", + # "SUNLinSol_Dense", "SUNLinSol_LapackDense" "SUNLinSol_SPBCGS", + # "SUNLinSol_SPFGMR", "SUNLinSol_SPGMR", "SUNLinSol_SPTFQMR", + "linear_solver": "SUNLinSol_KLU", + + # 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 + # approximate jacobian + "precon_half_bandwidth": 5, + + # for iterative linear solver preconditioner, bandwidth of + # approximate jacobian that is kept + "precon_half_bandwidth_keep": 5 + } + + Note: These options only have an effect if model.convert_to_format == 'casadi' + + """ def __init__( @@ -52,8 +87,28 @@ def __init__( root_method="casadi", root_tol=1e-6, extrap_tol=0, + options=None, ): + # set default options, + # (only if user does not supply) + 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, + } + if options is None: + options = default_options + else: + for key, value in default_options.items(): + if key not in options: + options[key] = value + self._options = options + if idaklu_spec is None: # pragma: no cover raise ImportError("KLU is not installed") @@ -177,7 +232,10 @@ def inputs_to_dict(inputs): if model.convert_to_format == "jax": mass_matrix = model.mass_matrix.entries.toarray() elif model.convert_to_format == "casadi": - mass_matrix = casadi.DM(model.mass_matrix.entries) + if self._options["jacobian"] == "dense": + mass_matrix = casadi.DM(model.mass_matrix.entries.toarray()) + else: + mass_matrix = casadi.DM(model.mass_matrix.entries) else: mass_matrix = model.mass_matrix.entries @@ -195,7 +253,6 @@ def resfn(t, y, inputs, ydot): if not model.use_jacobian: raise pybamm.SolverError("KLU requires the Jacobian") - use_jac = 1 # need to provide jacobian_rhs_alg - cj * mass_matrix if model.convert_to_format == "casadi": @@ -372,6 +429,14 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): for i, dFdp_i in enumerate(dFdp.values()): resvalS[i][:] = dFdy @ yS[i] - dFdyd @ ypS[i] + dFdp_i + try: + atol = model.atol + except AttributeError: + atol = self.atol + + rtol = self.rtol + atol = self._check_atol_type(atol, y0.size) + if model.convert_to_format == "casadi": rhs_algebraic = idaklu.generate_function(rhs_algebraic.serialize()) jac_times_cjmass = idaklu.generate_function(jac_times_cjmass.serialize()) @@ -381,6 +446,7 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): rootfn = idaklu.generate_function(rootfn.serialize()) mass_action = idaklu.generate_function(mass_action.serialize()) sensfn = idaklu.generate_function(sensfn.serialize()) + self._setup = { "rhs_algebraic": rhs_algebraic, "jac_times_cjmass": jac_times_cjmass, @@ -392,11 +458,32 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): "sensfn": sensfn, "rootfn": rootfn, "num_of_events": num_of_events, - "use_jac": use_jac, "ids": ids, "sensitivity_names": sensitivity_names, "number_of_sensitivity_parameters": number_of_sensitivity_parameters, } + + solver = idaklu.create_casadi_solver( + len(y0), + self._setup["number_of_sensitivity_parameters"], + self._setup["rhs_algebraic"], + self._setup["jac_times_cjmass"], + self._setup["jac_times_cjmass_colptrs"], + self._setup["jac_times_cjmass_rowvals"], + self._setup["jac_times_cjmass_nnz"], + self._setup["jac_rhs_algebraic_action"], + self._setup["mass_action"], + self._setup["sensfn"], + self._setup["rootfn"], + self._setup["num_of_events"], + self._setup["ids"], + atol, + rtol, + len(inputs), + self._options, + ) + + self._setup["solver"] = solver else: self._setup = { "resfn": resfn, @@ -404,7 +491,7 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): "sensfn": sensfn, "rootfn": rootfn, "num_of_events": num_of_events, - "use_jac": use_jac, + "use_jac": 1, "ids": ids, "sensitivity_names": sensitivity_names, "number_of_sensitivity_parameters": number_of_sensitivity_parameters, @@ -452,26 +539,11 @@ def _integrate(self, model, t_eval, inputs_dict=None): timer = pybamm.Timer() if model.convert_to_format == "casadi": - sol = idaklu.solve_casadi( + sol = self._setup["solver"].solve( t_eval, y0, ydot0, - self._setup["rhs_algebraic"], - self._setup["jac_times_cjmass"], - self._setup["jac_times_cjmass_colptrs"], - self._setup["jac_times_cjmass_rowvals"], - self._setup["jac_times_cjmass_nnz"], - self._setup["jac_rhs_algebraic_action"], - self._setup["mass_action"], - self._setup["sensfn"], - self._setup["rootfn"], - self._setup["num_of_events"], - self._setup["use_jac"], - self._setup["ids"], - atol, - rtol, inputs, - self._setup["number_of_sensitivity_parameters"], ) else: sol = idaklu.solve_python( diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index 932dc7876e..b080f14ca5 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -1,9 +1,13 @@ # # Tests for the KLU Solver class # -import pybamm -import numpy as np +from contextlib import redirect_stdout +import io import unittest + +import numpy as np + +import pybamm from tests import get_discretisation_for_testing @@ -250,7 +254,7 @@ def test_sensitivities_with_events(self): # this test implements a python version of the ida Roberts # example provided in sundials # see sundials ida examples pdf - for form in ["python", "casadi", "jax"]: + for form in ["casadi", "python", "jax"]: if form == "jax" and not pybamm.have_jax(): continue if form == "casadi": @@ -443,6 +447,90 @@ def test_dae_solver_algebraic_model(self): solution = solver.solve(model, t_eval) np.testing.assert_array_equal(solution.y, -1) + def test_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) + + # test print_stats + solver = pybamm.IDAKLUSolver(options={"print_stats": True}) + f = io.StringIO() + with redirect_stdout(f): + solver.solve(model, t_eval) + s = f.getvalue() + self.assertIn("Solver Stats", s) + + solver = pybamm.IDAKLUSolver(options={"print_stats": False}) + f = io.StringIO() + with redirect_stdout(f): + solver.solve(model, t_eval) + s = f.getvalue() + self.assertEqual(len(s), 0) + + # test everything else + for jacobian in ["none", "dense", "sparse", "matrix-free", "garbage"]: + for linear_solver in [ + "SUNLinSol_SPBCGS", + "SUNLinSol_Dense", + "SUNLinSol_LapackDense", + "SUNLinSol_KLU", + "SUNLinSol_SPFGMR", + "SUNLinSol_SPGMR", + "SUNLinSol_SPTFQMR", + "garbage", + ]: + for precon in ["none", "BBDP"]: + options = { + "jacobian": jacobian, + "linear_solver": linear_solver, + "preconditioner": precon, + } + solver = pybamm.IDAKLUSolver(options=options) + if ( + jacobian == "none" + and ( + linear_solver == "SUNLinSol_Dense" + or linear_solver == "SUNLinSol_LapackDense" + ) + or jacobian == "dense" + and ( + linear_solver == "SUNLinSol_Dense" + or linear_solver == "SUNLinSol_LapackDense" + ) + or jacobian == "sparse" + and ( + linear_solver != "SUNLinSol_Dense" + and linear_solver != "SUNLinSol_LapackDense" + and linear_solver != "garbage" + ) + or jacobian == "matrix-free" + and ( + linear_solver != "SUNLinSol_KLU" + and linear_solver != "SUNLinSol_Dense" + and linear_solver != "SUNLinSol_LapackDense" + and linear_solver != "garbage" + ) + ): + works = True + else: + works = False + + if works: + soln = solver.solve(model, t_eval) + np.testing.assert_array_almost_equal(soln.y, soln_base.y, 5) + else: + with self.assertRaises(ValueError): + soln = solver.solve(model, t_eval) + if __name__ == "__main__": print("Add -v for more debug output")