diff --git a/CHANGELOG.md b/CHANGELOG.md index de37f21b3b..bf032d3e09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,8 +2,9 @@ ## Features -- Added temperature control to experiment class. ([#2518])(https://github.com/pybamm-team/PyBaMM/pull/2518) +- Added an option for using a banded jacobian and sundials banded solvers for the IDAKLU solve [#2677](https://github.com/pybamm-team/PyBaMM/pull/2677) - The "particle size" option can now be a tuple to allow different behaviour in each electrode([#2672](https://github.com/pybamm-team/PyBaMM/pull/2672)). +- Added temperature control to experiment class. [#2518](https://github.com/pybamm-team/PyBaMM/pull/2518) ## Bug fixes diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index ac90172c97..132e8883f4 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -44,6 +44,7 @@ PYBIND11_MODULE(idaklu, m) 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_bandwidth_lower"), py::arg("jac_bandwidth_upper"), 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"), diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index a2de2e7089..e56b0902b2 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -34,16 +34,20 @@ void CasadiFunction::operator()() CasadiFunctions::CasadiFunctions( const Function &rhs_alg, const Function &jac_times_cjmass, const int jac_times_cjmass_nnz, + const int jac_bandwidth_lower, const int jac_bandwidth_upper, 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), + number_of_nnz(jac_times_cjmass_nnz), + jac_bandwidth_lower(jac_bandwidth_lower), jac_bandwidth_upper(jac_bandwidth_upper), + 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), + tmp_state_vector(number_of_states), + tmp_sparse_jacobian_data(jac_times_cjmass_nnz), options(options) { @@ -66,4 +70,5 @@ CasadiFunctions::CasadiFunctions( } -realtype *CasadiFunctions::get_tmp() { return tmp.data(); } +realtype *CasadiFunctions::get_tmp_state_vector() { return tmp_state_vector.data(); } +realtype *CasadiFunctions::get_tmp_sparse_jacobian_data() { return tmp_sparse_jacobian_data.data(); } diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 2e3b6beb8d..03264a8478 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -31,6 +31,8 @@ class CasadiFunctions int number_of_parameters; int number_of_events; int number_of_nnz; + int jac_bandwidth_lower; + int jac_bandwidth_upper; CasadiFunction rhs_alg; CasadiFunction sens; CasadiFunction jac_times_cjmass; @@ -44,6 +46,7 @@ class CasadiFunctions CasadiFunctions(const Function &rhs_alg, const Function &jac_times_cjmass, const int jac_times_cjmass_nnz, + const int jac_bandwidth_lower, const int jac_bandwidth_upper, const np_array_int &jac_times_cjmass_rowvals, const np_array_int &jac_times_cjmass_colptrs, const int inputs_length, const Function &jac_action, @@ -51,10 +54,12 @@ class CasadiFunctions const Function &events, const int n_s, int n_e, const int n_p, const Options& options); - realtype *get_tmp(); + realtype *get_tmp_state_vector(); + realtype *get_tmp_sparse_jacobian_data(); private: - std::vector tmp; + std::vector tmp_state_vector; + std::vector tmp_sparse_jacobian_data; }; #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 index d1bb76ea68..d10d0bdbf6 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -8,7 +8,9 @@ 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 int jac_times_cjmass_nnz, + const int jac_bandwidth_lower, const int jac_bandwidth_upper, + 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, @@ -16,19 +18,21 @@ create_casadi_solver(int number_of_states, int number_of_parameters, { auto options_cpp = Options(options); auto functions = std::make_unique( - rhs_alg, jac_times_cjmass, jac_times_cjmass_nnz, jac_times_cjmass_rowvals, + rhs_alg, jac_times_cjmass, jac_times_cjmass_nnz, jac_bandwidth_lower, jac_bandwidth_upper, 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, + number_of_events, jac_times_cjmass_nnz, + jac_bandwidth_lower, jac_bandwidth_upper, 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, + int jac_bandwidth_lower, int jac_bandwidth_upper, std::unique_ptr functions_arg, const Options &options) : number_of_states(atol_np.request().size), @@ -107,7 +111,14 @@ CasadiSolver::CasadiSolver(np_array atol_np, double rel_tol, jac_times_cjmass_nnz, CSC_MAT); #endif } - else if (options.jacobian == "dense" || options.jacobian == "none") + else if (options.jacobian == "banded") { + DEBUG("\tsetting banded matrix"); + #if SUNDIALS_VERSION_MAJOR >= 6 + J = SUNBandMatrix(number_of_states, jac_bandwidth_upper, jac_bandwidth_lower, sunctx); + #else + J = SUNBandMatrix(number_of_states, jac_bandwidth_upper, jac_bandwidth_lower); + #endif + } else if (options.jacobian == "dense" || options.jacobian == "none") { DEBUG("\tsetting dense matrix"); #if SUNDIALS_VERSION_MAJOR >= 6 @@ -151,6 +162,15 @@ CasadiSolver::CasadiSolver(np_array atol_np, double rel_tol, LS = SUNLinSol_KLU(yy, J, sunctx); #else LS = SUNLinSol_KLU(yy, J); +#endif + } + else if (options.linear_solver == "SUNLinSol_Band") + { + DEBUG("\tsetting SUNLinSol_Band linear solver"); +#if SUNDIALS_VERSION_MAJOR >= 6 + LS = SUNLinSol_Band(yy, J, sunctx); +#else + LS = SUNLinSol_Band(yy, J); #endif } else if (options.linear_solver == "SUNLinSol_SPBCGS") diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index 3eed122e04..09c4434d5b 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -14,7 +14,7 @@ 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, + int jac_times_cjmass_nnz, int jac_bandwidth_lower, int jac_bandwidth_upper, std::unique_ptr functions, const Options& options); ~CasadiSolver(); @@ -48,7 +48,9 @@ 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 int jac_times_cjmass_nnz, + const int jac_bandwidth_lower, const int jac_bandwidth_upper, + 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, diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index ce2a892725..031ef67d20 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -15,7 +15,7 @@ int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, p_python_functions->rhs_alg.m_res[0] = NV_DATA_S(rr); p_python_functions->rhs_alg(); - realtype *tmp = p_python_functions->get_tmp(); + realtype *tmp = p_python_functions->get_tmp_state_vector(); 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(); @@ -108,7 +108,7 @@ int jtimes_casadi(realtype tt, N_Vector yy, N_Vector yp, N_Vector rr, p_python_functions->jac_action(); // tmp has -∂F/∂y˙ v - realtype *tmp = p_python_functions->get_tmp(); + realtype *tmp = p_python_functions->get_tmp_state_vector(); 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(); @@ -148,15 +148,14 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, 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 if (p_python_functions->options.using_banded_matrix) { + jac_data = p_python_functions->get_tmp_sparse_jacobian_data(); + } else { jac_data = SUNDenseMatrix_Data(JJ); @@ -169,10 +168,31 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, 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) + + if (p_python_functions->options.using_banded_matrix) { + // copy data from temporary matrix to the banded matrix + auto jac_colptrs = p_python_functions->jac_times_cjmass_colptrs.data(); + auto jac_rowvals = p_python_functions->jac_times_cjmass_rowvals.data(); + int ncols = p_python_functions->number_of_states; + for (int col_ij = 0; col_ij < ncols; col_ij++) { + realtype *banded_col = SM_COLUMN_B(JJ, col_ij); + for (auto data_i = jac_colptrs[col_ij]; data_i < jac_colptrs[col_ij+1]; data_i++) { + auto row_ij = jac_rowvals[data_i]; + const realtype value_ij = jac_data[data_i]; + DEBUG("(" << row_ij << ", " << col_ij << ") = " << value_ij); + SM_COLUMN_ELEMENT_B(banded_col, row_ij, col_ij) = value_ij; + } + } + } + else if (p_python_functions->options.using_sparse_matrix) + { + + sunindextype *jac_colptrs = SUNSparseMatrix_IndexPointers(JJ); + sunindextype *jac_rowvals = SUNSparseMatrix_IndexValues(JJ); // row vals and col ptrs const int n_row_vals = p_python_functions->jac_times_cjmass_rowvals.size(); auto p_jac_times_cjmass_rowvals = @@ -262,7 +282,7 @@ int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp, for (int i = 0; i < np; i++) { // put (∂F/∂y)s i (t) in tmp - realtype *tmp = p_python_functions->get_tmp(); + realtype *tmp = p_python_functions->get_tmp_state_vector(); 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(); diff --git a/pybamm/solvers/c_solvers/idaklu/common.hpp b/pybamm/solvers/c_solvers/idaklu/common.hpp index 5bac325fc8..b1947654ea 100644 --- a/pybamm/solvers/c_solvers/idaklu/common.hpp +++ b/pybamm/solvers/c_solvers/idaklu/common.hpp @@ -16,6 +16,7 @@ #include /* access to KLU linear solver */ #include /* access to dense linear solver */ +#include /* access to dense linear solver */ #include /* access to spbcgs iterative linear solver */ #include #include diff --git a/pybamm/solvers/c_solvers/idaklu/options.cpp b/pybamm/solvers/c_solvers/idaklu/options.cpp index c3c7cb3583..f5d1d8c79e 100644 --- a/pybamm/solvers/c_solvers/idaklu/options.cpp +++ b/pybamm/solvers/c_solvers/idaklu/options.cpp @@ -1,4 +1,5 @@ #include "options.hpp" +#include #include @@ -15,9 +16,14 @@ Options::Options(py::dict options) { using_sparse_matrix = true; + using_banded_matrix = false; if (jacobian == "sparse") { } + else if (jacobian == "banded") { + using_banded_matrix = true; + using_sparse_matrix = false; + } else if (jacobian == "dense" || jacobian == "none") { using_sparse_matrix = false; @@ -29,7 +35,7 @@ Options::Options(py::dict options) { throw std::domain_error( "Unknown jacobian type \""s + jacobian + - "\". Should be one of \"sparse\", \"dense\", \"matrix-free\" or \"none\"."s + "\". Should be one of \"sparse\", \"banded\", \"dense\", \"matrix-free\" or \"none\"."s ); } @@ -40,6 +46,17 @@ Options::Options(py::dict options) else if (linear_solver == "SUNLinSol_KLU" && jacobian == "sparse") { } + else if (linear_solver == "SUNLinSol_Band" && jacobian == "banded") + { + } + else if (jacobian == "banded") { + throw std::domain_error( + "Unknown linear solver or incompatible options: " + "jacobian = \"" + jacobian + "\" linear solver = \"" + linear_solver + + "\". For a banded jacobian " + "please use the SUNLinSol_Band linear solver" + ); + } else if ((linear_solver == "SUNLinSol_SPBCGS" || linear_solver == "SUNLinSol_SPFGMR" || linear_solver == "SUNLinSol_SPGMR" || diff --git a/pybamm/solvers/c_solvers/idaklu/options.hpp b/pybamm/solvers/c_solvers/idaklu/options.hpp index 2fc807e48f..ecaffe8307 100644 --- a/pybamm/solvers/c_solvers/idaklu/options.hpp +++ b/pybamm/solvers/c_solvers/idaklu/options.hpp @@ -6,6 +6,7 @@ struct Options { bool print_stats; bool using_sparse_matrix; + bool using_banded_matrix; bool using_iterative_solver; std::string jacobian; std::string linear_solver; // klu, lapack, spbcg diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 864457acb9..f29d7faf29 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -52,11 +52,12 @@ class IDAKLUSolver(pybamm.BaseSolver): # print statistics of the solver after every solve "print_stats": False, - # jacobian form, can be "none", "dense", "sparse", "matrix-free" + # jacobian form, can be "none", "dense", + # "banded", "sparse", "matrix-free" "jacobian": "sparse", # name of sundials linear solver to use options are: "SUNLinSol_KLU", - # "SUNLinSol_Dense", "SUNLinSol_SPBCGS", + # "SUNLinSol_Dense", "SUNLinSol_Band", "SUNLinSol_SPBCGS", # "SUNLinSol_SPFGMR", "SUNLinSol_SPGMR", "SUNLinSol_SPTFQMR", "linear_solver": "SUNLinSol_KLU", @@ -274,7 +275,10 @@ def resfn(t, y, inputs, ydot): - cj_casadi * mass_matrix ], ) + jac_times_cjmass_sparsity = jac_times_cjmass.sparsity_out(0) + jac_bw_lower = jac_times_cjmass_sparsity.bw_lower() + jac_bw_upper = jac_times_cjmass_sparsity.bw_upper() jac_times_cjmass_nnz = jac_times_cjmass_sparsity.nnz() jac_times_cjmass_colptrs = np.array( jac_times_cjmass_sparsity.colind(), dtype=np.int64 @@ -447,6 +451,8 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): sensfn = idaklu.generate_function(sensfn.serialize()) self._setup = { + "jac_bandwidth_upper": jac_bw_upper, + "jac_bandwidth_lower": jac_bw_lower, "rhs_algebraic": rhs_algebraic, "jac_times_cjmass": jac_times_cjmass, "jac_times_cjmass_colptrs": jac_times_cjmass_colptrs, @@ -470,6 +476,8 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): self._setup["jac_times_cjmass_colptrs"], self._setup["jac_times_cjmass_rowvals"], self._setup["jac_times_cjmass_nnz"], + jac_bw_lower, + jac_bw_upper, self._setup["jac_rhs_algebraic_action"], self._setup["mass_action"], self._setup["sensfn"], diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index 9788d0d278..c834c5ca62 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -447,6 +447,30 @@ def test_dae_solver_algebraic_model(self): solution = solver.solve(model, t_eval) np.testing.assert_array_equal(solution.y, -1) + def test_banded(self): + model = pybamm.lithium_ion.SPM() + model.convert_to_format = "casadi" + param = model.default_parameter_values + param.process_model(model) + geometry = model.default_geometry + param.process_geometry(geometry) + mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + + t_eval = np.linspace(0, 3600, 100) + solver = pybamm.IDAKLUSolver() + soln = solver.solve(model, t_eval) + + options = { + "jacobian": "banded", + "linear_solver": "SUNLinSol_Band", + } + solver_banded = pybamm.IDAKLUSolver(options=options) + soln_banded = solver_banded.solve(model, t_eval) + + np.testing.assert_array_almost_equal(soln.y, soln_banded.y, 5) + def test_options(self): model = pybamm.BaseModel() u = pybamm.Variable("u")