From 36a7275f857655d3b91529db0b659b98e7a5692f Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 14 Jul 2023 11:39:57 +0000 Subject: [PATCH 01/38] Refactor of idaklu code --- CMakeLists.txt | 8 +- .../solvers/c_solvers/idaklu/CasadiSolver.cpp | 9 + .../solvers/c_solvers/idaklu/CasadiSolver.hpp | 26 + .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 433 ++++++++++ .../c_solvers/idaklu/CasadiSolverOpenMP.hpp | 65 ++ .../idaklu/CasadiSolverOpenMP_solvers.cpp | 53 ++ .../idaklu/CasadiSolverOpenMP_solvers.hpp | 231 ++++++ .../c_solvers/idaklu/casadi_functions.cpp | 84 +- .../c_solvers/idaklu/casadi_functions.hpp | 29 + .../c_solvers/idaklu/casadi_solver.cpp | 615 ++++---------- .../c_solvers/idaklu/casadi_solver.hpp | 87 +- pybamm/solvers/c_solvers/idaklu/common.hpp | 24 + pybamm/solvers/c_solvers/idaklu/options.cpp | 197 ++--- pybamm/solvers/c_solvers/idaklu/python.cpp | 764 +++++++++--------- .../idaklu/sundials_legacy_wrapper.hpp | 94 +++ 15 files changed, 1675 insertions(+), 1044 deletions(-) create mode 100644 pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp create mode 100644 pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp create mode 100644 pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp create mode 100644 pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp create mode 100644 pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp create mode 100644 pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp create mode 100644 pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp diff --git a/CMakeLists.txt b/CMakeLists.txt index cc47ba99c0..76e253ea8a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -39,8 +39,14 @@ pybind11_add_module(idaklu 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/CasadiSolver.cpp + pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp + pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp + pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp + pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp + pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp + pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.hpp pybamm/solvers/c_solvers/idaklu/common.hpp pybamm/solvers/c_solvers/idaklu/python.hpp pybamm/solvers/c_solvers/idaklu/python.cpp diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp new file mode 100644 index 0000000000..f45e667d68 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp @@ -0,0 +1,9 @@ +#include "CasadiSolver.hpp" +//#include "casadi_sundials_functions.hpp" + +/* + * This is an abstract base class for the Idaklu solver + */ + +CasadiSolver::CasadiSolver() {} +CasadiSolver::~CasadiSolver() {} diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp new file mode 100644 index 0000000000..be2336c412 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp @@ -0,0 +1,26 @@ +#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" +#include "sundials_legacy_wrapper.hpp" + +class CasadiSolver +{ +public: + CasadiSolver(); + ~CasadiSolver(); + virtual Solution solve( + np_array t_np, + np_array y0_np, + np_array yp0_np, + np_array_dense inputs) = 0; + virtual void Initialize() = 0; +}; + +#endif // PYBAMM_IDAKLU_CASADI_SOLVER_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp new file mode 100644 index 0000000000..95c0a307b4 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -0,0 +1,433 @@ +#include "CasadiSolverOpenMP.hpp" +#include "casadi_sundials_functions.hpp" + +/* + * This is an abstract class that implements an OpenMP solution but + * requires a linear solver to create a concrete class. Hook functions + * are intended to be overriden to support alternative solver + * approaches, as needed. + */ + +/* Skeleton workflow: + https://sundials.readthedocs.io/en/latest/ida/Usage/index.html + 1. (N/A) Initialize parallel or multi-threaded environment + 2. Create the SUNDIALS context object + 3. Create the vector of initial values + 4. Create matrix object (if appropriate) + 5. Create linear solver object + 6. (N/A) Create nonlinear solver object + 7. Create IDA object + 8. Initialize IDA solver + 9. Specify integration tolerances + 10. Attach the linear solver + 11. Set linear solver optional inputs + 12. (N/A) Attach nonlinear solver module + 13. (N/A) Set nonlinear solver optional inputs + 14. Specify rootfinding problem (optional) + 15. Set optional inputs + 16. Correct initial values (optional) + 17. Advance solution in time + 18. Get optional outputs + 19. Destroy objects + 20. (N/A) Finalize MPI +*/ + +CasadiSolverOpenMP::CasadiSolverOpenMP( + 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 +) : + atol_np(atol_np), + rhs_alg_id(rhs_alg_id), + number_of_states(atol_np.request().size), + number_of_parameters(number_of_parameters), + number_of_events(number_of_events), + jac_times_cjmass_nnz(jac_times_cjmass_nnz), + jac_bandwidth_lower(jac_bandwidth_lower), + jac_bandwidth_upper(jac_bandwidth_upper), + functions(std::move(functions_arg)), + options(options) +{ + // Construction code moved to Initialize() which is called from the + // (child) CasadiSolver_XXX class constructors. +} + +void CasadiSolverOpenMP::AllocateVectors() { + // Create vectors + yy = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); + yp = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); + avtol = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); + id = N_VNew_OpenMP(number_of_states, options.num_threads, sunctx); +} + +void CasadiSolverOpenMP::SetMatrix() { + // Create Matrix object + if (options.jacobian == "sparse") + { + DEBUG("\tsetting sparse matrix"); + J = SUNSparseMatrix( + number_of_states, + number_of_states, + jac_times_cjmass_nnz, + CSC_MAT, // CSC is used by casadi; CSR requires a conversion step + sunctx + ); + } + else if (options.jacobian == "banded") { + DEBUG("\tsetting banded matrix"); + J = SUNBandMatrix( + number_of_states, + jac_bandwidth_upper, + jac_bandwidth_lower, + sunctx + ); + } else if (options.jacobian == "dense" || options.jacobian == "none") + { + DEBUG("\tsetting dense matrix"); + J = SUNDenseMatrix( + number_of_states, + number_of_states, + sunctx + ); + } + else if (options.jacobian == "matrix-free") + { + DEBUG("\tsetting matrix-free"); + J = NULL; + } + else + throw std::invalid_argument("Unsupported matrix requested"); +} + +void CasadiSolverOpenMP::Initialize() { + DEBUG("CasadiSolverOpenMP::Initialize"); + auto atol = atol_np.unchecked<1>(); + + // create SUNDIALS context object + SUNContext_Create(NULL, &sunctx); // calls null-wrapper if Sundials Ver<6 + + // allocate memory for solver + ida_mem = IDACreate(sunctx); + + // create the vector of initial values + AllocateVectors(); + if (number_of_parameters > 0) + { + yyS = N_VCloneVectorArray(number_of_parameters, yy); + ypS = N_VCloneVectorArray(number_of_parameters, yp); + } + // set initial values + realtype *atval = N_VGetArrayPointer(avtol); + for (int i = 0; i < number_of_states; i++) + 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]); + } + + // create Matrix objects + SetMatrix(); + + // 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); + + // specify preconditioner type + precon_type = SUN_PREC_NONE; + if (options.preconditioner != "none") { + precon_type = SUN_PREC_LEFT; + } + + // create linear solver object + SetLinearSolver(); + + // attach the linear solver + 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); +} + +CasadiSolverOpenMP::~CasadiSolverOpenMP() +{ + // 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); + SUNContext_Free(&sunctx); +} + +Solution CasadiSolverOpenMP::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; + auto t = t_np.unchecked<1>(); + realtype t0 = RCONST(t(0)); + auto y0 = y0_np.unchecked<1>(); + auto yp0 = yp0_np.unchecked<1>(); + auto n_coeffs = number_of_states + number_of_parameters * number_of_states; + + if (y0.size() != n_coeffs) + throw std::domain_error( + "y0 has wrong size. Expected " + std::to_string(n_coeffs) + + " but got " + std::to_string(y0.size()) + ); + + if (yp0.size() != n_coeffs) + throw std::domain_error( + "yp0 has wrong size. Expected " + std::to_string(n_coeffs) + + " but got " + std::to_string(yp0.size())); + + // set inputs + auto p_inputs = inputs.unchecked<2>(); + for (int i = 0; i < functions->inputs.size(); i++) + functions->inputs[i] = p_inputs(i, 0); + + // set initial conditions + realtype *yval = N_VGetArrayPointer(yy); + realtype *ypval = N_VGetArrayPointer(yp); + std::vector ySval(number_of_parameters); + std::vector ypSval(number_of_parameters); + for (int p = 0 ; p < number_of_parameters; p++) { + ySval[p] = N_VGetArrayPointer(yyS[p]); + ypSval[p] = N_VGetArrayPointer(ypS[p]); + for (int i = 0; i < number_of_states; i++) { + ySval[p][i] = y0[i + (p + 1) * number_of_states]; + ypSval[p][i] = yp0[i + (p + 1) * number_of_states]; + } + } + + for (int i = 0; i < number_of_states; i++) + { + yval[i] = y0[i]; + ypval[i] = yp0[i]; + } + + IDAReInit(ida_mem, t0, yy, yp); + if (number_of_parameters > 0) + IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS); + + // correct initial values + DEBUG("IDACalcIC"); + IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); + if (number_of_parameters > 0) + IDAGetSens(ida_mem, &t0, yyS); + + int 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]; + } + + 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/CasadiSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp new file mode 100644 index 0000000000..962b156d0a --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -0,0 +1,65 @@ +#ifndef PYBAMM_IDAKLU_CASADISOLVEROPENMP_HPP +#define PYBAMM_IDAKLU_CASADISOLVEROPENMP_HPP + +#include "CasadiSolver.hpp" +#include +using Function = casadi::Function; + +#include "casadi_functions.hpp" +#include "common.hpp" +#include "options.hpp" +#include "solution.hpp" +#include "sundials_legacy_wrapper.hpp" + +class CasadiSolverOpenMP : public CasadiSolver +{ +public: + void *ida_mem; // pointer to memory + np_array atol_np; + double rel_tol; + np_array rhs_alg_id; + int number_of_states; + int number_of_parameters; + int number_of_events; + int precon_type; + 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; + int jac_bandwidth_lower; + int jac_bandwidth_upper; + SUNMatrix J; + SUNLinearSolver LS; + std::unique_ptr functions; + Options options; + +#if SUNDIALS_VERSION_MAJOR >= 6 + SUNContext sunctx; +#endif + +public: + CasadiSolverOpenMP( + 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, + const Options& options); + ~CasadiSolverOpenMP(); + Solution solve( + np_array t_np, + np_array y0_np, + np_array yp0_np, + np_array_dense inputs) override; + void Initialize() override; + void AllocateVectors(); + void SetMatrix(); + virtual void SetLinearSolver() = 0; +}; + +#endif // PYBAMM_IDAKLU_CASADISOLVEROPENMP_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp new file mode 100644 index 0000000000..574272eef5 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp @@ -0,0 +1,53 @@ +#include "CasadiSolverOpenMP_solvers.hpp" + +/* + * CasadiSolver implementations compatible with the OPENMP vector class + */ + +void CasadiSolverOpenMP_Dense::SetLinearSolver() { + LS = SUNLinSol_Dense(yy, J, sunctx); +} + +void CasadiSolverOpenMP_KLU::SetLinearSolver() { + LS = SUNLinSol_KLU(yy, J, sunctx); +} + +void CasadiSolverOpenMP_Band::SetLinearSolver() { + LS = SUNLinSol_Band(yy, J, sunctx); +} + +void CasadiSolverOpenMP_SPBCGS::SetLinearSolver() { + LS = SUNLinSol_SPBCGS( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +} + +void CasadiSolverOpenMP_SPFGMR::SetLinearSolver() { + LS = SUNLinSol_SPFGMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +} + +void CasadiSolverOpenMP_SPGMR::SetLinearSolver() { + LS = SUNLinSol_SPGMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +} + +void CasadiSolverOpenMP_SPTFQMR::SetLinearSolver() { + LS = SUNLinSol_SPTFQMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +} diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp new file mode 100644 index 0000000000..41413c1a39 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp @@ -0,0 +1,231 @@ +#ifndef PYBAMM_IDAKLU_CASADI_SOLVER_OPENMP_HPP +#define PYBAMM_IDAKLU_CASADI_SOLVER_OPENMP_HPP + +#include "CasadiSolverOpenMP.hpp" +#include "casadi_solver.hpp" + +class CasadiSolverOpenMP_Dense : public CasadiSolverOpenMP { +public: + CasadiSolverOpenMP_Dense( + 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, + const Options& options + ) : + CasadiSolverOpenMP( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options + ) + { + Initialize(); + } + void SetLinearSolver() override; +}; + +class CasadiSolverOpenMP_KLU : public CasadiSolverOpenMP { +public: + CasadiSolverOpenMP_KLU( + 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, + const Options& options + ) : + CasadiSolverOpenMP( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options + ) + { + Initialize(); + } + void SetLinearSolver() override; +}; + +class CasadiSolverOpenMP_Band : public CasadiSolverOpenMP { +public: + CasadiSolverOpenMP_Band( + 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, + const Options& options + ) : + CasadiSolverOpenMP( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options + ) + { + Initialize(); + } + void SetLinearSolver() override; +}; + +class CasadiSolverOpenMP_SPBCGS : public CasadiSolverOpenMP { +public: + CasadiSolverOpenMP_SPBCGS( + 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, + const Options& options + ) : + CasadiSolverOpenMP( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options + ) + { + Initialize(); + } + void SetLinearSolver() override; +}; + +class CasadiSolverOpenMP_SPFGMR : public CasadiSolverOpenMP { +public: + CasadiSolverOpenMP_SPFGMR( + 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, + const Options& options + ) : + CasadiSolverOpenMP( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options + ) + { + Initialize(); + } + void SetLinearSolver() override; +}; + +class CasadiSolverOpenMP_SPGMR : public CasadiSolverOpenMP { +public: + CasadiSolverOpenMP_SPGMR( + 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, + const Options& options + ) : + CasadiSolverOpenMP( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options + ) + { + Initialize(); + } + void SetLinearSolver() override; +}; + +class CasadiSolverOpenMP_SPTFQMR : public CasadiSolverOpenMP { +public: + CasadiSolverOpenMP_SPTFQMR( + 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, + const Options& options + ) : + CasadiSolverOpenMP( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options + ) + { + Initialize(); + } + void SetLinearSolver() override; +}; + +#endif // PYBAMM_IDAKLU_CASADI_SOLVER_OPENMP_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index e56b0902b2..f9a6c43523 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -2,33 +2,33 @@ 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); + 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); + 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( @@ -41,7 +41,7 @@ CasadiFunctions::CasadiFunctions( 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), + 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), @@ -51,24 +51,28 @@ CasadiFunctions::CasadiFunctions( 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]; - } + // 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]; - } + 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); - inputs.resize(inputs_length); - } -realtype *CasadiFunctions::get_tmp_state_vector() { return tmp_state_vector.data(); } -realtype *CasadiFunctions::get_tmp_sparse_jacobian_data() { return tmp_sparse_jacobian_data.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 03264a8478..24921355e4 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -6,6 +6,35 @@ #include "solution.hpp" #include +// Utility function for compressed-sparse-column (CSC) to/from +// compressed-sparse-row (CSR) matrix representation. +template +void csc_csr(realtype f[], T1 c[], T1 r[], realtype nf[], T2 nc[], T2 nr[], int N, int cols) { + int nn[cols+1]; + int rr[N]; + for (int i=0; i #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 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, - int inputs_length, py::dict options) -{ +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 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, + 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_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, - 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), - 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 - int num_threads = options.num_threads; -#if SUNDIALS_VERSION_MAJOR >= 6 - yy = N_VNew_OpenMP(number_of_states, num_threads, sunctx); - yp = N_VNew_OpenMP(number_of_states, num_threads, sunctx); - avtol = N_VNew_OpenMP(number_of_states, num_threads, sunctx); - id = N_VNew_OpenMP(number_of_states, num_threads, sunctx); -#else - yy = N_VNew_OpenMP(number_of_states, num_threads); - yp = N_VNew_OpenMP(number_of_states, num_threads); - avtol = N_VNew_OpenMP(number_of_states, num_threads); - id = N_VNew_OpenMP(number_of_states, num_threads); -#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 == "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 - 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") + 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 + ); + + CasadiSolver *casadiSolver = nullptr; + + // Instantiate solver class + if (options_cpp.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_KLU") + casadiSolver = new CasadiSolverOpenMP_Dense( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options_cpp + ); + } + else if (options_cpp.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_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") + casadiSolver = new CasadiSolverOpenMP_KLU( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options_cpp + ); + } + else if (options_cpp.linear_solver == "SUNLinSol_Band") + { + DEBUG("\tsetting SUNLinSol_Band linear solver"); + casadiSolver = new CasadiSolverOpenMP_Band( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options_cpp + ); + } + else if (options_cpp.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") + casadiSolver = new CasadiSolverOpenMP_SPBCGS( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options_cpp + ); + } + else if (options_cpp.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") + casadiSolver = new CasadiSolverOpenMP_SPFGMR( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options_cpp + ); + } + else if (options_cpp.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") + casadiSolver = new CasadiSolverOpenMP_SPGMR( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options_cpp + ); + } + else if (options_cpp.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; - auto t = t_np.unchecked<1>(); - realtype t0 = RCONST(t(0)); - auto y0 = y0_np.unchecked<1>(); - auto yp0 = yp0_np.unchecked<1>(); - - - if (y0.size() != number_of_states + number_of_parameters * number_of_states) { - throw std::domain_error( - "y0 has wrong size. Expected " + - std::to_string(number_of_states + number_of_parameters * number_of_states) + - " but got " + std::to_string(y0.size())); - } - - if (yp0.size() != number_of_states + number_of_parameters * number_of_states) { - throw std::domain_error( - "yp0 has wrong size. Expected " + - std::to_string(number_of_states + number_of_parameters * number_of_states) + - " but got " + std::to_string(yp0.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); - } - - // set initial conditions - realtype *yval = N_VGetArrayPointer(yy); - realtype *ypval = N_VGetArrayPointer(yp); - std::vector ySval(number_of_parameters); - std::vector ypSval(number_of_parameters); - for (int p = 0 ; p < number_of_parameters; p++) { - ySval[p] = N_VGetArrayPointer(yyS[p]); - ypSval[p] = N_VGetArrayPointer(ypS[p]); - for (int i = 0; i < number_of_states; i++) { - ySval[p][i] = y0[i + (p + 1) * number_of_states]; - ypSval[p][i] = yp0[i + (p + 1) * number_of_states]; - } - } - - for (int i = 0; i < number_of_states; i++) - { - yval[i] = y0[i]; - ypval[i] = yp0[i]; + casadiSolver = new CasadiSolverOpenMP_SPTFQMR( + atol_np, + rel_tol, + rhs_alg_id, + number_of_parameters, + number_of_events, + jac_times_cjmass_nnz, + jac_bandwidth_lower, + jac_bandwidth_upper, + std::move(functions), + options_cpp + ); } - IDAReInit(ida_mem, t0, yy, yp); - if (number_of_parameters > 0) { - IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS); - } - - // calculate consistent initial conditions - DEBUG("IDACalcIC"); - IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); - if (number_of_parameters > 0) - { - IDAGetSens(ida_mem, &t0, yyS); - } - - 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]; - } - } - - - - 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); + if (casadiSolver == nullptr) { + throw std::invalid_argument("Unsupported solver requested"); } - return sol; + return casadiSolver; } diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index 09c4434d5b..551a6161f5 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -1,59 +1,28 @@ -#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, int jac_bandwidth_lower, int jac_bandwidth_upper, - 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 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, - double rel_tol, int inputs_length, py::dict options); - -#endif // PYBAMM_IDAKLU_CASADI_SOLVER_HPP +#ifndef PYBAMM_IDAKLU_CREATE_CASADI_SOLVER_HPP +#define PYBAMM_IDAKLU_CREATE_CASADI_SOLVER_HPP + +#include "CasadiSolver.hpp" + +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 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, + double rel_tol, + int inputs_length, + py::dict options +); + +#endif // PYBAMM_IDAKLU_CREATE_CASADI_SOLVER_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/common.hpp b/pybamm/solvers/c_solvers/idaklu/common.hpp index 1931f8cb61..c5e8bcd7cc 100644 --- a/pybamm/solvers/c_solvers/idaklu/common.hpp +++ b/pybamm/solvers/c_solvers/idaklu/common.hpp @@ -44,7 +44,20 @@ using np_array_int = py::array_t; #ifdef NDEBUG #define DEBUG_VECTOR(vector) +#define DEBUG_VECTORn(vector) #else + +#define DEBUG_VECTORn(vector, N) {\ + std::cout << #vector << "[n=" << N << "] = ["; \ + auto array_ptr = N_VGetArrayPointer(vector); \ + for (int i = 0; i < N; i++) { \ + std::cout << array_ptr[i]; \ + if (i < N-1) { \ + std::cout << ", "; \ + } \ + } \ + std::cout << "]" << std::endl; } + #define DEBUG_VECTOR(vector) {\ std::cout << #vector << " = ["; \ auto array_ptr = N_VGetArrayPointer(vector); \ @@ -56,6 +69,17 @@ using np_array_int = py::array_t; } \ } \ std::cout << "]" << std::endl; } + +#define DEBUG_v(v, N) {\ + std::cout << #v << "[n=" << N << "] = ["; \ + for (int i = 0; i < N; i++) { \ + std::cout << v[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 index 2a5ed58d9c..efad4d5de0 100644 --- a/pybamm/solvers/c_solvers/idaklu/options.cpp +++ b/pybamm/solvers/c_solvers/idaklu/options.cpp @@ -2,7 +2,7 @@ #include #include - + using namespace std::string_literals; Options::Options(py::dict options) @@ -16,103 +16,106 @@ Options::Options(py::dict options) num_threads(options["num_threads"].cast()) { - 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; - } - else if (jacobian == "matrix-free") - { - } - else - { - throw std::domain_error( - "Unknown jacobian type \""s + jacobian + - "\". Should be one of \"sparse\", \"banded\", \"dense\", \"matrix-free\" or \"none\"."s - ); - } + 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; + } + else if (jacobian == "matrix-free") + { + } + else + { + throw std::domain_error( + "Unknown jacobian type \""s + jacobian + + "\". Should be one of \"sparse\", \"banded\", \"dense\", \"matrix-free\" or \"none\"."s + ); + } - using_iterative_solver = false; - if (linear_solver == "SUNLinSol_Dense" && (jacobian == "dense" || jacobian == "none")) - { - } - 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" || - 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 + "\"" - ); - } + using_iterative_solver = false; + if (linear_solver == "SUNLinSol_Dense" && (jacobian == "dense" || jacobian == "none")) + { + } + else if (linear_solver == "SUNLinSol_KLU" && jacobian == "sparse") + { + } + else if (linear_solver == "SUNLinSol_cuSolverSp_batchQR" && 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" || + 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") + 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 { - throw std::domain_error( - "Unknown preconditioner \""s + preconditioner + - "\", use one of \"BBDP\" or \"none\""s - ); - } - } - else - { - preconditioner = "none"; - } + preconditioner = "none"; + } } diff --git a/pybamm/solvers/c_solvers/idaklu/python.cpp b/pybamm/solvers/c_solvers/idaklu/python.cpp index a1803988d4..0ca3bf0ac7 100644 --- a/pybamm/solvers/c_solvers/idaklu/python.cpp +++ b/pybamm/solvers/c_solvers/idaklu/python.cpp @@ -5,205 +5,213 @@ 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); } + 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; + 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; + 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); + 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); + 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; + py::array_t r_np; - r_np = python_functions.res(tres, y_np, yp_np); + r_np = python_functions.res(tres, y_np, yp_np); - auto r_np_ptr = r_np.unchecked<1>(); + 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; + // 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); + realtype *yval; + yval = N_VGetArrayPointer(yy); - PybammFunctions *python_functions_ptr = - static_cast(user_data); - PybammFunctions python_functions = *python_functions_ptr; + 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); + 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); + // 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; + py::array_t jac_np_array; - python_functions.jac(tt, y_np, cj); + 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>(); + 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]; - } + // 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; + 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]; - } + 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>(); + 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]; - } + // 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); + return (0); } int events(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, void *user_data) { - realtype *yval; - yval = N_VGetArrayPointer(yy); + realtype *yval; + yval = N_VGetArrayPointer(yy); - PybammFunctions *python_functions_ptr = - static_cast(user_data); - PybammFunctions python_functions = *python_functions_ptr; + 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); + 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; + py::array_t events_np_array; - events_np_array = python_functions.events(t, y_np); + events_np_array = python_functions.events(t, y_np); - auto events_np_data_ptr = events_np_array.unchecked<1>(); + 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]; - } + // 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); + 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 +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. @@ -212,267 +220,267 @@ int sensitivities(int Ns, realtype t, N_Vector yy, N_Vector yp, // 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. +// 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 +// 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, +// 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), +// 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; + 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) + 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; + 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); + SUNContext sunctx; + SUNContext_Create(NULL, &sunctx); - // allocate memory for solver - ida_mem = IDACreate(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); + // 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); + // 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 (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); + 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); + 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]; - } - } + IDASetLinearSolver(ida_mem, LS, J); - // calculate consistent initial conditions - auto id_np_val = rhs_alg_id.unchecked<1>(); - realtype *id_val; - id_val = N_VGetArrayPointer(id); + if (use_jacobian == 1) + { + IDASetJacFn(ida_mem, jacobian); + } - int ii; - for (ii = 0; ii < number_of_states; ii++) - { - id_val[ii] = id_np_val[ii]; - } + if (number_of_parameters > 0) + { + IDASensInit(ida_mem, number_of_parameters, + IDA_SIMULTANEOUS, sensitivities, yyS, ypS); + IDASensEEtolerances(ida_mem); + } - IDASetId(ida_mem, id); - IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); + int t_i = 1; + realtype tret; + realtype t_next; + realtype t_final = t(number_of_timesteps - 1); - while (true) - { - t_next = t(t_i); - IDASetStopTime(ida_mem, t_next); - retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); + // 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); - if (retval == IDA_TSTOP_RETURN || retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) + t_return[0] = t(0); + for (int j = 0; j < number_of_states; j++) { - 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; + 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]; + yS_return[base_index + k] = ySval[j][k]; } - } - t_i += 1; - if (retval == IDA_SUCCESS || retval == IDA_ROOT_RETURN) { - break; - } + } + + // 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); } - } - - /* 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); + 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] - ); + 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); + Solution sol(retval, t_ret, y_ret, yS_ret); - return sol; + return sol; } diff --git a/pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp b/pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp new file mode 100644 index 0000000000..19bc92dcdf --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp @@ -0,0 +1,94 @@ + +#if SUNDIALS_VERSION_MAJOR < 6 + + #define SUN_PREC_NONE PREC_NONE + #define SUN_PREC_LEFT PREC_LEFT + + // Compatibility layer - wrap older sundials functions in new-style calls + void SUNContext_Create(void *comm, SUNContext *ctx) + { + // Function not available + return; + } + + int SUNContext_Free(SUNContext *ctx) + { + // Function not available + return; + } + + void* IDACreate(SUNContext sunctx) + { + return IDACreate(); + } + + N_Vector N_VNew_Serial(sunindextype vec_length, SUNContext sunctx) + { + return N_VNew_Serial(vec_length); + } + + N_Vector N_VNew_OpenMP(sunindextype vec_length, SUNContext sunctx) + { + return N_VNew_OpenMP(vec_length); + } + + N_Vector N_VNew_Cuda(sunindextype vec_length, SUNContext sunctx) + { + return N_VNew_Cuda(vec_length); + } + + SUNMatrix SUNSparseMatrix(sunindextype M, sunindextype N, sunindextype NNZ, int sparsetype, SUNContext sunctx) + { + return SUNMatrix SUNSparseMatrix(M, N, NNZ, sparsetype); + } + + SUNMatrix SUNMatrix_cuSparse_NewCSR(int M, int N, int NNZ, cusparseHandle_t cusp, SUNContext sunctx) + { + return SUNMatrix_cuSparse_NewCSR(M, N, NNZ, cusp); + } + + SUNMatrix SUNBandMatrix(sunindextype N, sunindextype mu, sunindextype ml, SUNContext sunctx) + { + return SUNMatrix SUNBandMatrix(N, mu, ml); + } + + SUNMatrix SUNDenseMatrix(sunindextype M, sunindextype N, SUNContext sunctx) + { + return SUNDenseMatrix(M, N, sunctx); + } + + SUNLinearSolver SUNLinSol_Dense(N_Vector y, SUNMatrix A, SUNContext sunctx) + { + return SUNLinSol_Dense(y, A, sunctx); + } + + SUNLinearSolver SUNLinSol_KLU(N_Vector y, SUNMatrix A, SUNContext sunctx) + { + return SUNLinSol_KLU(y, A, sunctx); + } + + SUNLinearSolver SUNLinSol_Band(N_Vector y, SUNMatrix A, SUNContext sunctx) + { + return SUNLinSol_Band(y, A, sunctx); + } + + SUNLinearSolver SUNLinSol_SPBCGS(N_Vector y, int pretype, int maxl, SUNContext sunctx) + { + return SUNLinSol_SPBCGS(y, pretype, maxl); + } + + SUNLinearSolver SUNLinSol_SPFGMR(N_Vector y, int pretype, int maxl, SUNContext sunctx) + { + return SUNLinSol_SPFGMR(y, pretype, maxl); + } + + SUNLinearSolver SUNLinSol_SPGMR(N_Vector y, int pretype, int maxl, SUNContext sunctx) + { + return SUNLinSol_SPGMR(y, pretype, maxl); + } + + SUNLinearSolver SUNLinSol_SPTFQMR(N_Vector y, int pretype, int maxl, SUNContext sunctx) + { + return SUNLinSol_SPTFQMR(y, pretype, maxl); + } +#endif From e693cefb1f544099f1a677494e646dc80e03dea5 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 14 Jul 2023 11:42:18 +0000 Subject: [PATCH 02/38] Add convenience test code (quick compile & run) --- compile | 12 ++++++++++++ compile.py | 31 +++++++++++++++++++++++++++++++ test.py | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 95 insertions(+) create mode 100755 compile create mode 100644 compile.py create mode 100644 test.py diff --git a/compile b/compile new file mode 100755 index 0000000000..60d3c3e4cb --- /dev/null +++ b/compile @@ -0,0 +1,12 @@ +source .nox/dev/bin/activate +rm pybamm/solvers/idaklu.*.so +# rm -rf build/* +mkdir build +python compile.py +cd build +cmake --build . +cp idaklu* ../pybamm/solvers +cd .. + +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/jsb/.local/lib +python test.py diff --git a/compile.py b/compile.py new file mode 100644 index 0000000000..ecff808d5b --- /dev/null +++ b/compile.py @@ -0,0 +1,31 @@ +import os +import sys +import subprocess + +# Folder containing hte file +cmake_list_dir = os.path.abspath(os.path.dirname(__file__)) + +build_dir = 'build' + +cmake_args = [ + "-DCMAKE_BUILD_TYPE=DEBUG", + "-DPYTHON_EXECUTABLE={}".format(sys.executable), + "-DUSE_PYTHON_CASADI=TRUE", +] +if True: + cmake_args.append( + "-DSuiteSparse_ROOT={}".format(os.path.abspath("/home/jsb/.local")) + ) +if True: + cmake_args.append( + "-DSUNDIALS_ROOT={}".format(os.path.abspath("/home/jsb/.local")) + ) + +build_env = os.environ +# build_env["vcpkg_root_dir"] = vcpkg_root_dir +# build_env["vcpkg_default_triplet"] = vcpkg_default_triplet +# build_env["vcpkg_feature_flags"] = vcpkg_feature_flags + +subprocess.run( + ["cmake", cmake_list_dir] + cmake_args, cwd=build_dir, env=build_env +) diff --git a/test.py b/test.py new file mode 100644 index 0000000000..80d1177af9 --- /dev/null +++ b/test.py @@ -0,0 +1,52 @@ +solver_opt = 2 +jacobian = 'sparse' # sparse, dense, band, none +num_threads = 1 + +import pybamm +import numpy as np +import importlib + +# check for loading errors +idaklu_spec = importlib.util.find_spec("pybamm.solvers.idaklu") +idaklu = importlib.util.module_from_spec(idaklu_spec) +idaklu_spec.loader.exec_module(idaklu) + +# construct model +# pybamm.set_logging_level("INFO") +model = pybamm.lithium_ion.DFN() +# model.convert_to_format = 'jax' +geometry = model.default_geometry +param = model.default_parameter_values +param.process_model(model) +param.process_geometry(geometry) +n = 100 # control the complexity of the geometry (increases number of solver states) +var_pts = {"x_n": n, "x_s": n, "x_p": n, "r_n": 10, "r_p": 10} +mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) +disc = pybamm.Discretisation(mesh, model.default_spatial_methods) +disc.process_model(model) +t_eval = np.linspace(0, 3600, 100) + +if solver_opt == 1: + linear_solver = 'SUNLinSol_Dense' +if solver_opt == 2: + linear_solver = 'SUNLinSol_KLU' +if solver_opt == 3: + linear_solver = 'SUNLinSol_Band' +if solver_opt == 4: + linear_solver = 'SUNLinSol_SPBCGS' +if solver_opt == 5: + linear_solver = 'SUNLinSol_SPFGMR' +if solver_opt == 6: + linear_solver = 'SUNLinSol_SPGMR' +if solver_opt == 7: + linear_solver = 'SUNLinSol_SPTFQMR' +if solver_opt == 8: + linear_solver = 'SUNLinSol_cuSolverSp_batchQR' + jacobian = 'cuSparse_' + +options = {'linear_solver': linear_solver, 'jacobian': jacobian, 'num_threads': num_threads} +klu_sol = pybamm.IDAKLUSolver(atol=1e-8, rtol=1e-8, options=options).solve(model, t_eval) +print(f"Solve time: {klu_sol.solve_time.value*1000} msecs") + +# plot = pybamm.QuickPlot(klu_sol) +# plot.dynamic_plot() From 97e7e0dd271ea319ae748b5032a863e92f0f3837 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 14 Jul 2023 11:52:00 +0000 Subject: [PATCH 03/38] Add CSR support to idaklu solver --- .../idaklu/casadi_sundials_functions.cpp | 73 +++++++++++++------ 1 file changed, 52 insertions(+), 21 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index eaf383cda4..3e58d65d5d 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -174,6 +174,9 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, if (p_python_functions->options.using_banded_matrix) { + if (SUNSparseMatrix_SparseType(JJ) != CSC_MAT) + throw std::runtime_error("Banded matrix only tested with CSC format"); + // 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(); @@ -190,30 +193,58 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, } 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 = - 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++) + if (SUNSparseMatrix_SparseType(JJ) == CSC_MAT) { - jac_rowvals[i] = p_jac_times_cjmass_rowvals[i]; - } + 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 = + 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(); + 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]; - } + // 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]; + } + } else if (SUNSparseMatrix_SparseType(JJ) == CSR_MAT) { + realtype newjac[SUNSparseMatrix_NNZ(JJ)]; + sunindextype *jac_ptrs = SUNSparseMatrix_IndexPointers(JJ); + sunindextype *jac_vals = SUNSparseMatrix_IndexValues(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_OMP(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] = newjac; + p_python_functions->jac_times_cjmass(); + + // convert (casadi's) CSC format to CSR + csc_csr( + newjac, + p_python_functions->jac_times_cjmass_rowvals.data(), + p_python_functions->jac_times_cjmass_colptrs.data(), + jac_data, + jac_ptrs, + jac_vals, + SUNSparseMatrix_NNZ(JJ), + SUNSparseMatrix_NP(JJ) + ); + } else + throw std::runtime_error("Unknown matrix format detected (Expected CSC or CSR)"); } return (0); From 201ec7f29275e80096506efd0896f5cbf1664a1f Mon Sep 17 00:00:00 2001 From: John Brittain Date: Sat, 15 Jul 2023 22:28:58 +0000 Subject: [PATCH 04/38] Add initial support for a list of variables' functions in idaklu --- compile | 11 +- compile.py | 2 +- pybamm/solvers/c_solvers/idaklu.cpp | 91 ++-- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 398 +++++++++--------- .../c_solvers/idaklu/casadi_functions.cpp | 124 +++--- .../c_solvers/idaklu/casadi_functions.hpp | 48 ++- .../c_solvers/idaklu/casadi_solver.cpp | 44 +- .../c_solvers/idaklu/casadi_solver.hpp | 2 + .../idaklu/casadi_sundials_functions.cpp | 41 +- pybamm/solvers/idaklu_solver.py | 69 ++- pybamm/solvers/solution.py | 17 +- test.py | 23 +- 12 files changed, 509 insertions(+), 361 deletions(-) diff --git a/compile b/compile index 60d3c3e4cb..2fdcaf2586 100755 --- a/compile +++ b/compile @@ -1,7 +1,10 @@ +#!/usr/bin/env bash + +set -eox pipefail + source .nox/dev/bin/activate -rm pybamm/solvers/idaklu.*.so -# rm -rf build/* -mkdir build +#rm pybamm/solvers/idaklu.*.so +#rm -rf build/* python compile.py cd build cmake --build . @@ -9,4 +12,4 @@ cp idaklu* ../pybamm/solvers cd .. export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/jsb/.local/lib -python test.py +python -it test.py diff --git a/compile.py b/compile.py index ecff808d5b..edb86b3d52 100644 --- a/compile.py +++ b/compile.py @@ -8,7 +8,7 @@ build_dir = 'build' cmake_args = [ - "-DCMAKE_BUILD_TYPE=DEBUG", + "-DCMAKE_BUILD_TYPE=RELEASE", "-DPYTHON_EXECUTABLE={}".format(sys.executable), "-DUSE_PYTHON_CASADI=TRUE", ] diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index 132e8883f4..2206d69460 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -17,6 +18,7 @@ Function generate_function(const std::string &data) namespace py = pybind11; PYBIND11_MAKE_OPAQUE(std::vector); +//PYBIND11_MAKE_OPAQUE(std::vector); PYBIND11_MODULE(idaklu, m) { @@ -25,39 +27,74 @@ PYBIND11_MODULE(idaklu, m) 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); + "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); + .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); + //py::bind_vector>(m, "VectorFunction"); + //py::implicitly_convertible>(); + 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_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"), - 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); + "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_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("extra_fcn"), + py::arg("var_casadi_fcns"), + 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); + .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/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index 95c0a307b4..607748b5f1 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -225,209 +225,229 @@ Solution CasadiSolverOpenMP::solve( np_array_dense inputs ) { - DEBUG("CasadiSolver::solve"); - - int number_of_timesteps = t_np.request().size; - auto t = t_np.unchecked<1>(); - realtype t0 = RCONST(t(0)); - auto y0 = y0_np.unchecked<1>(); - auto yp0 = yp0_np.unchecked<1>(); - auto n_coeffs = number_of_states + number_of_parameters * number_of_states; - - if (y0.size() != n_coeffs) - throw std::domain_error( - "y0 has wrong size. Expected " + std::to_string(n_coeffs) + - " but got " + std::to_string(y0.size()) - ); - - if (yp0.size() != n_coeffs) - throw std::domain_error( - "yp0 has wrong size. Expected " + std::to_string(n_coeffs) + - " but got " + std::to_string(yp0.size())); - - // set inputs - auto p_inputs = inputs.unchecked<2>(); - for (int i = 0; i < functions->inputs.size(); i++) - functions->inputs[i] = p_inputs(i, 0); - - // set initial conditions - realtype *yval = N_VGetArrayPointer(yy); - realtype *ypval = N_VGetArrayPointer(yp); - std::vector ySval(number_of_parameters); - std::vector ypSval(number_of_parameters); - for (int p = 0 ; p < number_of_parameters; p++) { - ySval[p] = N_VGetArrayPointer(yyS[p]); - ypSval[p] = N_VGetArrayPointer(ypS[p]); - for (int i = 0; i < number_of_states; i++) { - ySval[p][i] = y0[i + (p + 1) * number_of_states]; - ypSval[p][i] = yp0[i + (p + 1) * number_of_states]; - } + DEBUG("CasadiSolver::solve"); + + int number_of_timesteps = t_np.request().size; + auto t = t_np.unchecked<1>(); + realtype t0 = RCONST(t(0)); + auto y0 = y0_np.unchecked<1>(); + auto yp0 = yp0_np.unchecked<1>(); + auto n_coeffs = number_of_states + number_of_parameters * number_of_states; + + if (y0.size() != n_coeffs) + throw std::domain_error( + "y0 has wrong size. Expected " + std::to_string(n_coeffs) + + " but got " + std::to_string(y0.size())); + + if (yp0.size() != n_coeffs) + throw std::domain_error( + "yp0 has wrong size. Expected " + std::to_string(n_coeffs) + + " but got " + std::to_string(yp0.size())); + + // set inputs + auto p_inputs = inputs.unchecked<2>(); + for (uint i = 0; i < functions->inputs.size(); i++) + functions->inputs[i] = p_inputs(i, 0); + + // set initial conditions + realtype *yval = N_VGetArrayPointer(yy); + realtype *ypval = N_VGetArrayPointer(yp); + std::vector ySval(number_of_parameters); + std::vector ypSval(number_of_parameters); + for (int p = 0 ; p < number_of_parameters; p++) { + ySval[p] = N_VGetArrayPointer(yyS[p]); + ypSval[p] = N_VGetArrayPointer(ypS[p]); + for (int i = 0; i < number_of_states; i++) { + ySval[p][i] = y0[i + (p + 1) * number_of_states]; + ypSval[p][i] = yp0[i + (p + 1) * number_of_states]; } + } - for (int i = 0; i < number_of_states; i++) - { - yval[i] = y0[i]; - ypval[i] = yp0[i]; - } + for (int i = 0; i < number_of_states; i++) + { + yval[i] = y0[i]; + ypval[i] = yp0[i]; + } - IDAReInit(ida_mem, t0, yy, yp); - if (number_of_parameters > 0) - IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS); - - // correct initial values - DEBUG("IDACalcIC"); - IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); - if (number_of_parameters > 0) - IDAGetSens(ida_mem, &t0, yyS); - - int 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; - } - ); + IDAReInit(ida_mem, t0, yy, yp); + if (number_of_parameters > 0) + IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS); - 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]; + // correct initial values + DEBUG("IDACalcIC"); + IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1)); + if (number_of_parameters > 0) + IDAGetSens(ida_mem, &t0, yyS); + + int 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]; + } - int retval; - while (true) + 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) { - t_next = t(t_i); - IDASetStopTime(ida_mem, t_next); - DEBUG("IDASolve"); - retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); + if (number_of_parameters > 0) + IDAGetSens(ida_mem, &tret, yyS); - 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 + 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++) { - // failed - break; + 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]; } - } - 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 - ); + // Try evaluating one of the variable casadi functions + realtype *res = new realtype[number_of_states]; + functions->extra_fcn.m_arg[0] = &t_return[t_i]; + functions->extra_fcn.m_arg[1] = &y_return[t_i * number_of_states]; + functions->extra_fcn.m_arg[2] = functions->inputs.data(); + functions->extra_fcn.m_res[0] = res; + functions->extra_fcn(); + std::cout << "Calculated value [extra]: " << *res << std::endl; + int k = 0; + for (auto& var_fcn : functions->var_casadi_fcns) { + var_fcn.m_arg[0] = &t_return[t_i]; + var_fcn.m_arg[1] = &y_return[t_i * number_of_states]; + var_fcn.m_arg[2] = functions->inputs.data(); + var_fcn.m_res[0] = res; + var_fcn(); + std::cout << "Calculated value [" << k << "]: " << *res << std::endl; + k++; + } - Solution sol(retval, t_ret, y_ret, yS_ret); + t_i += 1; - if (options.print_stats) + if (retval == IDA_SUCCESS || + retval == IDA_ROOT_RETURN) + break; + } + else { - 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); + // 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; + return sol; } diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index f9a6c43523..9dfc2fab59 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -2,77 +2,85 @@ 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); + 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, nullptr); + m_res.resize(sz_res, nullptr); + m_iw.resize(sz_iw, 0); + m_w.resize(sz_w, 0); } -// only call this once m_arg and m_res have been set appropriatelly +// only call this once m_arg and m_res have been set appropriately 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); + 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 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), - 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_state_vector(number_of_states), - tmp_sparse_jacobian_data(jac_times_cjmass_nnz), - options(options) + 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 Function &extra_fcn, + const std::vector var_casadi_fcns, + 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), + 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_state_vector(number_of_states), + tmp_sparse_jacobian_data(jac_times_cjmass_nnz), + extra_fcn(extra_fcn), + options(options) { + // convert casadi::Function list to CasadiFunction list + this->var_casadi_fcns.clear(); + for (auto& var : var_casadi_fcns) { + this->var_casadi_fcns.push_back(CasadiFunction(*var)); + } + + // 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]; + } - // 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); + 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_state_vector() { - return tmp_state_vector.data(); + return tmp_state_vector.data(); } realtype *CasadiFunctions::get_tmp_sparse_jacobian_data() { - return tmp_sparse_jacobian_data.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 24921355e4..e10e75b8a7 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -5,6 +5,7 @@ #include "options.hpp" #include "solution.hpp" #include +#include // Utility function for compressed-sparse-column (CSC) to/from // compressed-sparse-row (CSR) matrix representation. @@ -47,7 +48,7 @@ class CasadiFunction std::vector m_res; void operator()(); -private: +//private: const Function &m_func; std::vector m_iw; std::vector m_w; @@ -62,30 +63,47 @@ class CasadiFunctions int number_of_nnz; int jac_bandwidth_lower; int jac_bandwidth_upper; + 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; + CasadiFunction extra_fcn; + std::vector var_casadi_fcns; + + std::vector jac_times_cjmass_rowvals; + std::vector jac_times_cjmass_colptrs; + std::vector inputs; + Options options; - - 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, - 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_state_vector(); realtype *get_tmp_sparse_jacobian_data(); +public: + 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, + const Function &mass_action, + const Function &sens, + const Function &events, + const int n_s, + const int n_e, + const int n_p, + const Function &extra_fcn, + const std::vector var_casadi_fcns, + const Options& options + ); + private: std::vector tmp_state_vector; std::vector tmp_sparse_jacobian_data; diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp index bb687f39fd..0fea3c2d44 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -7,26 +7,28 @@ #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 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, - int inputs_length, - py::dict options -) { + 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 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, + int inputs_length, + const Function &extra_fcn, + const std::vector var_casadi_fcns, + py::dict options +) { auto options_cpp = Options(options); auto functions = std::make_unique( rhs_alg, @@ -44,6 +46,8 @@ CasadiSolver *create_casadi_solver( number_of_states, number_of_events, number_of_parameters, + extra_fcn, + var_casadi_fcns, options_cpp ); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index 551a6161f5..5414983424 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -22,6 +22,8 @@ CasadiSolver *create_casadi_solver( np_array atol_np, double rel_tol, int inputs_length, + const Function &extra_fcn, + const std::vector var_casadi_fcns, py::dict options ); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index 3e58d65d5d..5e34bb97a9 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -2,6 +2,9 @@ #include "casadi_functions.hpp" #include "common.hpp" +#define NV_DATA NV_DATA_OMP +//#define NV_DATA NV_DATA_S + int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, void *user_data) { @@ -10,19 +13,19 @@ int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, static_cast(user_data); p_python_functions->rhs_alg.m_arg[0] = &tres; - p_python_functions->rhs_alg.m_arg[1] = NV_DATA_OMP(yy); + p_python_functions->rhs_alg.m_arg[1] = NV_DATA(yy); p_python_functions->rhs_alg.m_arg[2] = p_python_functions->inputs.data(); - p_python_functions->rhs_alg.m_res[0] = NV_DATA_OMP(rr); + p_python_functions->rhs_alg.m_res[0] = NV_DATA(rr); p_python_functions->rhs_alg(); realtype *tmp = p_python_functions->get_tmp_state_vector(); - p_python_functions->mass_action.m_arg[0] = NV_DATA_OMP(yp); + p_python_functions->mass_action.m_arg[0] = NV_DATA(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_OMP(rr)); + casadi::casadi_axpy(ns, -1., tmp, NV_DATA(rr)); //DEBUG_VECTOR(yy); //DEBUG_VECTOR(yp); @@ -101,22 +104,22 @@ int jtimes_casadi(realtype tt, N_Vector yy, N_Vector yp, N_Vector rr, // Jv has ∂F/∂y v p_python_functions->jac_action.m_arg[0] = &tt; - p_python_functions->jac_action.m_arg[1] = NV_DATA_OMP(yy); + p_python_functions->jac_action.m_arg[1] = NV_DATA(yy); p_python_functions->jac_action.m_arg[2] = p_python_functions->inputs.data(); - p_python_functions->jac_action.m_arg[3] = NV_DATA_OMP(v); - p_python_functions->jac_action.m_res[0] = NV_DATA_OMP(Jv); + p_python_functions->jac_action.m_arg[3] = NV_DATA(v); + p_python_functions->jac_action.m_res[0] = NV_DATA(Jv); p_python_functions->jac_action(); // tmp has -∂F/∂y˙ v realtype *tmp = p_python_functions->get_tmp_state_vector(); - p_python_functions->mass_action.m_arg[0] = NV_DATA_OMP(v); + p_python_functions->mass_action.m_arg[0] = NV_DATA(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_OMP(Jv)); + casadi::casadi_axpy(ns, -cj, tmp, NV_DATA(Jv)); return 0; } @@ -163,7 +166,7 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, // 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_OMP(yy); + p_python_functions->jac_times_cjmass.m_arg[1] = NV_DATA(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; @@ -225,7 +228,7 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, // 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_OMP(yy); + p_python_functions->jac_times_cjmass.m_arg[1] = NV_DATA(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; @@ -258,7 +261,7 @@ int events_casadi(realtype t, N_Vector yy, N_Vector yp, realtype *events_ptr, // 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_OMP(yy); + p_python_functions->events.m_arg[1] = NV_DATA(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(); @@ -302,11 +305,11 @@ int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp, // 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_OMP(yy); + p_python_functions->sens.m_arg[1] = NV_DATA(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_OMP(resvalS[i]); + p_python_functions->sens.m_res[i] = NV_DATA(resvalS[i]); } // resvalsS now has (∂F/∂p i ) p_python_functions->sens(); @@ -316,23 +319,23 @@ int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp, // put (∂F/∂y)s i (t) in 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_OMP(yy); + p_python_functions->jac_action.m_arg[1] = NV_DATA(yy); p_python_functions->jac_action.m_arg[2] = p_python_functions->inputs.data(); - p_python_functions->jac_action.m_arg[3] = NV_DATA_OMP(yS[i]); + p_python_functions->jac_action.m_arg[3] = NV_DATA(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_OMP(resvalS[i])); + casadi::casadi_axpy(ns, 1., tmp, NV_DATA(resvalS[i])); // put -(∂F/∂ ẏ) ṡ i (t) in tmp2 - p_python_functions->mass_action.m_arg[0] = NV_DATA_OMP(ypS[i]); + p_python_functions->mass_action.m_arg[0] = NV_DATA(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_OMP(resvalS[i])); + casadi::casadi_axpy(ns, -1., tmp, NV_DATA(resvalS[i])); } return 0; diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index a09dc72362..d8cf872acd 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -43,6 +43,9 @@ 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). + output_variables : list[str], optional + List of variables to calculate and return. If none are specified then + the complete state vector is returned (can be very large) (default is []) options: dict, optional Addititional options to pass to the solver, by default: @@ -91,6 +94,7 @@ def __init__( root_method="casadi", root_tol=1e-6, extrap_tol=None, + output_variables=[], options=None, ): # set default options, @@ -113,6 +117,8 @@ def __init__( options[key] = value self._options = options + self.output_variables = output_variables + if idaklu_spec is None: # pragma: no cover raise ImportError("KLU is not installed") @@ -181,6 +187,11 @@ def inputs_to_dict(inputs): # only casadi solver needs sensitivity ics if model.convert_to_format != "casadi": y0S = None + if self.output_variables: + raise SolverError( + "output_variables can only be specified " + 'with convert_to_format="casadi"' + ) if y0S is not None: if isinstance(y0S, casadi.DM): y0S = (y0S,) @@ -217,6 +228,7 @@ def resfn(t, y, inputs, ydot): raise pybamm.SolverError("KLU requires the Jacobian") # need to provide jacobian_rhs_alg - cj * mass_matrix + self.var_casadi_fcns = [] if model.convert_to_format == "casadi": t_casadi = casadi.MX.sym("t") y_casadi = casadi.MX.sym("y", model.len_rhs_and_alg) @@ -258,6 +270,21 @@ def resfn(t, y, inputs, ydot): "mass_action", [v_casadi], [casadi.densify(mass_matrix @ v_casadi)] ) + # convert 'variable' expressions to casadi functions + for key in self.output_variables: + var_casadi = casadi.Function( + "variable", + [t_casadi, y_casadi, p_casadi_stacked], + [ + model.variables_and_events[key].to_casadi( + t_casadi, y_casadi, p_casadi_stacked + ) + ], + ) + self.var_casadi_fcns.append( + idaklu.generate_function(var_casadi.serialize()) + ) + else: t0 = 0 if t_eval is None else t_eval[0] jac_y0_t0 = model.jac_rhs_algebraic_eval(t0, y0, inputs_dict) @@ -428,28 +455,32 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): "ids": ids, "sensitivity_names": sensitivity_names, "number_of_sensitivity_parameters": number_of_sensitivity_parameters, + "output_variables": self.output_variables, + "var_casadi_fcns": self.var_casadi_fcns, } 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"], - jac_bw_lower, - jac_bw_upper, - 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, + number_of_states=len(y0), + number_of_parameters=self._setup["number_of_sensitivity_parameters"], + rhs_alg=self._setup["rhs_algebraic"], + jac_times_cjmass=self._setup["jac_times_cjmass"], + jac_times_cjmass_colptrs=self._setup["jac_times_cjmass_colptrs"], + jac_times_cjmass_rowvals=self._setup["jac_times_cjmass_rowvals"], + jac_times_cjmass_nnz=self._setup["jac_times_cjmass_nnz"], + jac_bandwidth_lower=jac_bw_lower, + jac_bandwidth_upper=jac_bw_upper, + jac_action=self._setup["jac_rhs_algebraic_action"], + mass_action=self._setup["mass_action"], + sens=self._setup["sensfn"], + events=self._setup["rootfn"], + number_of_events=self._setup["num_of_events"], + rhs_alg_id=self._setup["ids"], + atol=atol, + rtol=rtol, + inputs=len(inputs), + extra_fcn=self._setup["var_casadi_fcns"][0], + var_casadi_fcns=self._setup["var_casadi_fcns"], + options=self._options, ) self._setup["solver"] = solver diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py index 411341f887..3233570808 100644 --- a/pybamm/solvers/solution.py +++ b/pybamm/solvers/solution.py @@ -458,13 +458,21 @@ def update(self, variables): cumtrapz_ic = var_pybamm.initial_condition cumtrapz_ic = cumtrapz_ic.evaluate() var_pybamm = var_pybamm.child - var_casadi = self.process_casadi_var(var_pybamm, inputs, ys) + var_casadi = Solution.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) model._variables_casadi[key] = var_casadi vars_pybamm[i] = var_pybamm elif key in model._variables_casadi: var_casadi = model._variables_casadi[key] else: - var_casadi = self.process_casadi_var(var_pybamm, inputs, ys) + var_casadi = Solution.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) model._variables_casadi[key] = var_casadi vars_casadi.append(var_casadi) var = pybamm.ProcessedVariable( @@ -475,9 +483,10 @@ def update(self, variables): self._variables[key] = var self.data[key] = var.data - def process_casadi_var(self, var_pybamm, inputs, ys): + @staticmethod + def process_casadi_var(var_pybamm, inputs, ys_shape): t_MX = casadi.MX.sym("t") - y_MX = casadi.MX.sym("y", ys.shape[0]) + y_MX = casadi.MX.sym("y", ys_shape[0]) inputs_MX_dict = { key: casadi.MX.sym("input", value.shape[0]) for key, value in inputs.items() } diff --git a/test.py b/test.py index 80d1177af9..e0f72e83d2 100644 --- a/test.py +++ b/test.py @@ -1,6 +1,12 @@ solver_opt = 2 jacobian = 'sparse' # sparse, dense, band, none num_threads = 1 +output_variables = [ + "Time [min]", + "Voltage [V]", + "Current [A]", +] +#output_variables = [] import pybamm import numpy as np @@ -44,9 +50,16 @@ linear_solver = 'SUNLinSol_cuSolverSp_batchQR' jacobian = 'cuSparse_' -options = {'linear_solver': linear_solver, 'jacobian': jacobian, 'num_threads': num_threads} -klu_sol = pybamm.IDAKLUSolver(atol=1e-8, rtol=1e-8, options=options).solve(model, t_eval) -print(f"Solve time: {klu_sol.solve_time.value*1000} msecs") +options = { + 'linear_solver': linear_solver, + 'jacobian': jacobian, + 'num_threads': num_threads, +} + +klu_sol = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + output_variables=output_variables, +).solve(model, t_eval) -# plot = pybamm.QuickPlot(klu_sol) -# plot.dynamic_plot() +print(f"Solve time: {klu_sol.solve_time.value*1000} msecs") From 74def0402ecf56b1ab438c1754a93e400588ee64 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Sun, 16 Jul 2023 15:55:52 +0000 Subject: [PATCH 05/38] Collate vars or state vectors in idaklu; refactor Solution class structure ready for separate vars and state-vector implementations --- pybamm/__init__.py | 5 +- pybamm/solvers/c_solvers/idaklu.cpp | 1 - .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 112 ++- .../c_solvers/idaklu/casadi_functions.cpp | 2 - .../c_solvers/idaklu/casadi_functions.hpp | 49 +- .../c_solvers/idaklu/casadi_solver.cpp | 2 - .../c_solvers/idaklu/casadi_solver.hpp | 1 - pybamm/solvers/idaklu_solver.py | 1 - pybamm/solvers/solution.py | 925 +----------------- pybamm/solvers/solution_base.py | 922 +++++++++++++++++ pybamm/solvers/solution_full.py | 8 + pybamm/solvers/solution_vars.py | 8 + test.py | 2 +- 13 files changed, 1045 insertions(+), 993 deletions(-) create mode 100644 pybamm/solvers/solution_base.py create mode 100644 pybamm/solvers/solution_full.py create mode 100644 pybamm/solvers/solution_vars.py diff --git a/pybamm/__init__.py b/pybamm/__init__.py index 6c2636ba51..cf70c01c7a 100644 --- a/pybamm/__init__.py +++ b/pybamm/__init__.py @@ -200,7 +200,10 @@ # # Solver classes # -from .solvers.solution import Solution, EmptySolution, make_cycle_solution +from .solvers.solution_base import SolutionBase, EmptySolution, make_cycle_solution +from .solvers.solution_full import SolutionFull +from .solvers.solution_vars import SolutionVars +from .solvers.solution import Solution from .solvers.processed_variable import ProcessedVariable from .solvers.base_solver import BaseSolver from .solvers.dummy_solver import DummySolver diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index 2206d69460..f5b8b23aed 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -80,7 +80,6 @@ PYBIND11_MODULE(idaklu, m) py::arg("atol"), py::arg("rtol"), py::arg("inputs"), - py::arg("extra_fcn"), py::arg("var_casadi_fcns"), py::arg("options"), py::return_value_policy::take_ownership); diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index 607748b5f1..ce968876c8 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -279,18 +279,29 @@ Solution CasadiSolverOpenMP::solve( if (number_of_parameters > 0) IDAGetSens(ida_mem, &t0, yyS); - int t_i = 1; realtype tret; realtype t_next; realtype t_final = t(number_of_timesteps - 1); // set return vectors + int length_of_return_vector = 0; + size_t max_res_size = 0; // maximum result size (for common result buffer) + if (functions->var_casadi_fcns.size() > 0) { + // return only the requested variables list after computation + for (auto& var_fcn : functions->var_casadi_fcns) { + max_res_size = std::max(max_res_size, var_fcn.m_res.size()); + length_of_return_vector += max_res_size - 1; + } + } else { + // Return full y state-vector + length_of_return_vector = number_of_states; + } realtype *t_return = new realtype[number_of_timesteps]; realtype *y_return = new realtype[number_of_timesteps * - number_of_states]; + length_of_return_vector]; realtype *yS_return = new realtype[number_of_parameters * number_of_timesteps * - number_of_states]; + length_of_return_vector]; py::capsule free_t_when_done( t_return, @@ -314,17 +325,37 @@ Solution CasadiSolverOpenMP::solve( } ); + // Initial state (t_i=0) 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]; + realtype *res = new realtype[max_res_size]; + if (functions->var_casadi_fcns.size() > 0) { + // Evaluate casadi functions for each requested variable and store + size_t j = 0; + for (auto& var_fcn : functions->var_casadi_fcns) { + var_fcn.m_arg[0] = &tret; + var_fcn.m_arg[1] = yval; + var_fcn.m_arg[2] = functions->inputs.data(); + var_fcn.m_res[0] = res; + var_fcn(); + // store in return vector + for (size_t jj=0; jj0) int retval; + int t_i = 1; while (true) { t_next = t(t_i); @@ -339,37 +370,36 @@ Solution CasadiSolverOpenMP::solve( if (number_of_parameters > 0) IDAGetSens(ida_mem, &tret, yyS); + // Evaluate and store results for the time step 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]; + size_t j = 0; + std::cout << "Timestep " << t_return[t_i] << std::endl; + if (functions->var_casadi_fcns.size() > 0) { + // Evaluate casadi functions for each requested variable and store + for (auto& var_fcn : functions->var_casadi_fcns) { + var_fcn.m_arg[0] = &tret; + var_fcn.m_arg[1] = yval; + var_fcn.m_arg[2] = functions->inputs.data(); + var_fcn.m_res[0] = res; + var_fcn(); + // store in return vector + for (size_t jj=0; jjextra_fcn.m_arg[0] = &t_return[t_i]; - functions->extra_fcn.m_arg[1] = &y_return[t_i * number_of_states]; - functions->extra_fcn.m_arg[2] = functions->inputs.data(); - functions->extra_fcn.m_res[0] = res; - functions->extra_fcn(); - std::cout << "Calculated value [extra]: " << *res << std::endl; - int k = 0; - for (auto& var_fcn : functions->var_casadi_fcns) { - var_fcn.m_arg[0] = &t_return[t_i]; - var_fcn.m_arg[1] = &y_return[t_i * number_of_states]; - var_fcn.m_arg[2] = functions->inputs.data(); - var_fcn.m_res[0] = res; - var_fcn(); - std::cout << "Calculated value [" << k << "]: " << *res << std::endl; - k++; - } - t_i += 1; if (retval == IDA_SUCCESS || @@ -389,7 +419,7 @@ Solution CasadiSolverOpenMP::solve( free_t_when_done ); np_array y_ret = np_array( - t_i * number_of_states, + t_i * length_of_return_vector, &y_return[0], free_y_when_done ); @@ -397,7 +427,7 @@ Solution CasadiSolverOpenMP::solve( std::vector { number_of_parameters, number_of_timesteps, - number_of_states + length_of_return_vector }, &yS_return[0], free_yS_when_done diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index 9dfc2fab59..97ec8f5ba7 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -40,7 +40,6 @@ CasadiFunctions::CasadiFunctions( 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 Function &extra_fcn, const std::vector var_casadi_fcns, const Options& options) : number_of_states(n_s), number_of_events(n_e), number_of_parameters(n_p), @@ -51,7 +50,6 @@ CasadiFunctions::CasadiFunctions( mass_action(mass_action), sens(sens), events(events), tmp_state_vector(number_of_states), tmp_sparse_jacobian_data(jac_times_cjmass_nnz), - extra_fcn(extra_fcn), options(options) { // convert casadi::Function list to CasadiFunction list diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index e10e75b8a7..956c915e71 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -3,7 +3,6 @@ #include "common.hpp" #include "options.hpp" -#include "solution.hpp" #include #include @@ -42,13 +41,13 @@ class CasadiFunction { public: explicit CasadiFunction(const Function &f); + void operator()(); public: std::vector m_arg; std::vector m_res; - void operator()(); -//private: +private: const Function &m_func; std::vector m_iw; std::vector m_w; @@ -56,6 +55,27 @@ class CasadiFunction class CasadiFunctions { +public: + 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, + const Function &mass_action, + const Function &sens, + const Function &events, + const int n_s, + const int n_e, + const int n_p, + const std::vector var_casadi_fcns, + const Options& options + ); + public: int number_of_states; int number_of_parameters; @@ -70,7 +90,6 @@ class CasadiFunctions CasadiFunction jac_action; CasadiFunction mass_action; CasadiFunction events; - CasadiFunction extra_fcn; std::vector var_casadi_fcns; std::vector jac_times_cjmass_rowvals; @@ -82,28 +101,6 @@ class CasadiFunctions realtype *get_tmp_state_vector(); realtype *get_tmp_sparse_jacobian_data(); -public: - 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, - const Function &mass_action, - const Function &sens, - const Function &events, - const int n_s, - const int n_e, - const int n_p, - const Function &extra_fcn, - const std::vector var_casadi_fcns, - const Options& options - ); - private: std::vector tmp_state_vector; std::vector tmp_sparse_jacobian_data; diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp index 0fea3c2d44..82fd9561a2 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -25,7 +25,6 @@ CasadiSolver *create_casadi_solver( np_array atol_np, double rel_tol, int inputs_length, - const Function &extra_fcn, const std::vector var_casadi_fcns, py::dict options ) { @@ -46,7 +45,6 @@ CasadiSolver *create_casadi_solver( number_of_states, number_of_events, number_of_parameters, - extra_fcn, var_casadi_fcns, options_cpp ); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index 5414983424..54cf1e0ea2 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -22,7 +22,6 @@ CasadiSolver *create_casadi_solver( np_array atol_np, double rel_tol, int inputs_length, - const Function &extra_fcn, const std::vector var_casadi_fcns, py::dict options ); diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index d8cf872acd..0e521c6a42 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -478,7 +478,6 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): atol=atol, rtol=rtol, inputs=len(inputs), - extra_fcn=self._setup["var_casadi_fcns"][0], var_casadi_fcns=self._setup["var_casadi_fcns"], options=self._options, ) diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py index 3233570808..a71c8231ce 100644 --- a/pybamm/solvers/solution.py +++ b/pybamm/solvers/solution.py @@ -1,923 +1,14 @@ # -# Solution class +# Solution interface class # -import casadi -import json -import numbers -import numpy as np -import pickle import pybamm -import pandas as pd -from scipy.io import savemat -from functools import cached_property -class NumpyEncoder(json.JSONEncoder): - """ - Numpy serialiser helper class that converts numpy arrays to a list - https://stackoverflow.com/questions/26646362/numpy-array-is-not-json-serializable - """ - - def default(self, obj): - if isinstance(obj, np.ndarray): - return obj.tolist() - # won't be called since we only need to convert numpy arrays - return json.JSONEncoder.default(self, obj) # pragma: no cover - - -class Solution(object): - """ - Class containing the solution of, and various attributes associated with, a PyBaMM - model. - - Parameters - ---------- - all_ts : :class:`numpy.array`, size (n,) (or list of these) - A one-dimensional array containing the times at which the solution is evaluated. - A list of times can be provided instead to initialize a solution with - sub-solutions. - all_ys : :class:`numpy.array`, size (m, n) (or list of these) - A two-dimensional array containing the values of the solution. y[i, :] is the - vector of solutions at time t[i]. - A list of ys can be provided instead to initialize a solution with - sub-solutions. - all_models : :class:`pybamm.BaseModel` - The model that was used to calculate the solution. - A list of models can be provided instead to initialize a solution with - sub-solutions that have been calculated using those models. - all_inputs : dict (or list of these) - The inputs that were used to calculate the solution - A list of inputs can be provided instead to initialize a solution with - sub-solutions. - t_event : :class:`numpy.array`, size (1,) - A zero-dimensional array containing the time at which the event happens. - y_event : :class:`numpy.array`, size (m,) - A one-dimensional array containing the value of the solution at the time when - the event happens. - termination : str - String to indicate why the solution terminated - - sensitivities: bool or dict - True if sensitivities included as the solution of the explicit forwards - equations. False if no sensitivities included/wanted. Dict if sensitivities are - provided as a dict of {parameter: sensitivities} pairs. - - """ - - def __init__( - self, - all_ts, - all_ys, - all_models, - all_inputs, - t_event=None, - y_event=None, - termination="final time", - sensitivities=False, - check_solution=True, - ): - if not isinstance(all_ts, list): - all_ts = [all_ts] - if not isinstance(all_ys, list): - all_ys = [all_ys] - if not isinstance(all_models, list): - all_models = [all_models] - self._all_ts = all_ts - self._all_ys = all_ys - self._all_ys_and_sens = all_ys - self._all_models = all_models - - # Set up inputs - if not isinstance(all_inputs, list): - all_inputs_copy = dict(all_inputs) - for key, value in all_inputs_copy.items(): - if isinstance(value, numbers.Number): - all_inputs_copy[key] = np.array([value]) - self.all_inputs = [all_inputs_copy] - else: - self.all_inputs = all_inputs - - self.sensitivities = sensitivities - - self._t_event = t_event - self._y_event = y_event - self._termination = termination - - # Check no ys are too large - if check_solution: - self.check_ys_are_not_too_large() - - # Events - self._t_event = t_event - self._y_event = y_event - self._termination = termination - self.closest_event_idx = None - - # Initialize times - self.set_up_time = None - self.solve_time = None - self.integration_time = None - - # initialize empty variables and data - self._variables = pybamm.FuzzyDict() - self.data = pybamm.FuzzyDict() - - # Add self as sub-solution for compatibility with ProcessedVariable - self._sub_solutions = [self] - - # initialize empty cycles - self._cycles = [] - - # Initialize empty summary variables - self._summary_variables = None - - # Solution now uses CasADi - pybamm.citations.register("Andersson2019") - - def extract_explicit_sensitivities(self): - # if we got here, we haven't set y yet - self.set_y() - - # extract sensitivities from full y solution - self._y, self._sensitivities = self._extract_explicit_sensitivities( - self.all_models[0], self.y, self.t, self.all_inputs[0] - ) - - # make sure we remove all sensitivities from all_ys - for index, (model, ys, ts, inputs) in enumerate( - zip(self.all_models, self.all_ys, self.all_ts, self.all_inputs) - ): - self._all_ys[index], _ = self._extract_explicit_sensitivities( - model, ys, ts, inputs - ) - - def _extract_explicit_sensitivities(self, model, y, t_eval, inputs): - """ - given a model and a solution y, extracts the sensitivities - - Parameters - -------- - model : :class:`pybamm.BaseModel` - A model that has been already setup by this base solver - y: ndarray - The solution of the full explicit sensitivity equations - t_eval: ndarray - The evaluation times - inputs: dict - parameter inputs - - Returns - ------- - y: ndarray - The solution of the ode/dae in model - sensitivities: dict of (string: ndarray) - A dictionary of parameter names, and the corresponding solution of - the sensitivity equations - """ - - n_states = model.len_rhs_and_alg - n_rhs = model.len_rhs - n_alg = model.len_alg - # Get the point where the algebraic equations start - if model.len_rhs != 0: - n_p = model.len_rhs_sens // model.len_rhs +class Solution(pybamm.SolutionBase): + def __new__(cls, *args, **kwargs): + if kwargs.get("output_variables", []): + # Solution contains the full state vector 'y' + return pybamm.SolutionVars(*args, **kwargs) else: - n_p = model.len_alg_sens // model.len_alg - len_rhs_and_sens = model.len_rhs + model.len_rhs_sens - - n_t = len(t_eval) - # y gets the part of the solution vector that correspond to the - # actual ODE/DAE solution - - # save sensitivities as a dictionary - # first save the whole sensitivity matrix - # reshape using Fortran order to get the right array: - # t0_x0_p0, t0_x0_p1, ..., t0_x0_pn - # t0_x1_p0, t0_x1_p1, ..., t0_x1_pn - # ... - # t0_xn_p0, t0_xn_p1, ..., t0_xn_pn - # t1_x0_p0, t1_x0_p1, ..., t1_x0_pn - # t1_x1_p0, t1_x1_p1, ..., t1_x1_pn - # ... - # t1_xn_p0, t1_xn_p1, ..., t1_xn_pn - # ... - # tn_x0_p0, tn_x0_p1, ..., tn_x0_pn - # tn_x1_p0, tn_x1_p1, ..., tn_x1_pn - # ... - # tn_xn_p0, tn_xn_p1, ..., tn_xn_pn - # 1, Extract rhs and alg sensitivities and reshape into 3D matrices - # with shape (n_p, n_states, n_t) - if isinstance(y, casadi.DM): - y_full = y.full() - else: - y_full = y - ode_sens = y_full[n_rhs:len_rhs_and_sens, :].reshape(n_p, n_rhs, n_t) - alg_sens = y_full[len_rhs_and_sens + n_alg :, :].reshape(n_p, n_alg, n_t) - # 2. Concatenate into a single 3D matrix with shape (n_p, n_states, n_t) - # i.e. along first axis - full_sens_matrix = np.concatenate([ode_sens, alg_sens], axis=1) - # Transpose and reshape into a (n_states * n_t, n_p) matrix - full_sens_matrix = full_sens_matrix.transpose(2, 1, 0).reshape( - n_t * n_states, n_p - ) - - # Save the full sensitivity matrix - sensitivity = {"all": full_sens_matrix} - - # also save the sensitivity wrt each parameter (read the columns of the - # sensitivity matrix) - start = 0 - for name in model.calculate_sensitivities: - inp = inputs[name] - input_size = inp.shape[0] - end = start + input_size - sensitivity[name] = full_sens_matrix[:, start:end] - start = end - - y_dae = np.vstack( - [ - y[: model.len_rhs, :], - y[len_rhs_and_sens : len_rhs_and_sens + model.len_alg, :], - ] - ) - return y_dae, sensitivity - - @property - def t(self): - """Times at which the solution is evaluated""" - try: - return self._t - except AttributeError: - self.set_t() - return self._t - - def set_t(self): - self._t = np.concatenate(self.all_ts) - if any(np.diff(self._t) <= 0): - raise ValueError("Solution time vector must be strictly increasing") - - @property - def y(self): - """Values of the solution""" - try: - return self._y - except AttributeError: - self.set_y() - - # if y is evaluated before sensitivities then need to extract them - if isinstance(self._sensitivities, bool) and self._sensitivities: - self.extract_explicit_sensitivities() - - return self._y - - @property - def sensitivities(self): - """Values of the sensitivities. Returns a dict of param_name: np_array""" - if isinstance(self._sensitivities, bool): - if self._sensitivities: - self.extract_explicit_sensitivities() - else: - self._sensitivities = {} - return self._sensitivities - - @sensitivities.setter - def sensitivities(self, value): - """Updates the sensitivity""" - # sensitivities must be a dict or bool - if not isinstance(value, (bool, dict)): - raise TypeError("sensitivities arg needs to be a bool or dict") - self._sensitivities = value - - def set_y(self): - try: - if isinstance(self.all_ys[0], (casadi.DM, casadi.MX)): - self._y = casadi.horzcat(*self.all_ys) - else: - self._y = np.hstack(self.all_ys) - except ValueError: - raise pybamm.SolverError( - "The solution is made up from different models, so `y` cannot be " - "computed explicitly." - ) - - def check_ys_are_not_too_large(self): - # Only check last one so that it doesn't take too long - # We only care about the cases where y is growing too large without any - # restraint, so if y gets large in the middle then comes back down that is ok - y, model = self.all_ys[-1], self.all_models[-1] - y = y[:, -1] - if np.any(y > pybamm.settings.max_y_value): - for var in [*model.rhs.keys(), *model.algebraic.keys()]: - y_var = y[model.variables[var.name].y_slices[0]] - if np.any(y_var > pybamm.settings.max_y_value): - pybamm.logger.error( - f"Solution for '{var}' exceeds the maximum allowed value " - f"of `{pybamm.settings.max_y_value}. This could be due to " - "incorrect scaling, model formulation, or " - "parameter values. The maximum allowed value is set by " - "'pybammm.settings.max_y_value'." - ) - - @property - def all_ts(self): - return self._all_ts - - @property - def all_ys(self): - return self._all_ys - - @property - def all_models(self): - """Model(s) used for solution""" - return self._all_models - - @cached_property - def all_inputs_casadi(self): - return [casadi.vertcat(*inp.values()) for inp in self.all_inputs] - - @property - def t_event(self): - """Time at which the event happens""" - return self._t_event - - @property - def y_event(self): - """Value of the solution at the time of the event""" - return self._y_event - - @property - def termination(self): - """Reason for termination""" - return self._termination - - @termination.setter - def termination(self, value): - """Updates the reason for termination""" - self._termination = value - - @cached_property - def first_state(self): - """ - A Solution object that only contains the first state. This is faster to evaluate - than the full solution when only the first state is needed (e.g. to initialize - a model with the solution) - """ - new_sol = Solution( - self.all_ts[0][:1], - self.all_ys[0][:, :1], - self.all_models[:1], - self.all_inputs[:1], - None, - None, - "final time", - ) - new_sol._all_inputs_casadi = self.all_inputs_casadi[:1] - new_sol._sub_solutions = self.sub_solutions[:1] - - new_sol.solve_time = 0 - new_sol.integration_time = 0 - new_sol.set_up_time = 0 - - return new_sol - - @cached_property - def last_state(self): - """ - A Solution object that only contains the final state. This is faster to evaluate - than the full solution when only the final state is needed (e.g. to initialize - a model with the solution) - """ - new_sol = Solution( - self.all_ts[-1][-1:], - self.all_ys[-1][:, -1:], - self.all_models[-1:], - self.all_inputs[-1:], - self.t_event, - self.y_event, - self.termination, - ) - new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:] - new_sol._sub_solutions = self.sub_solutions[-1:] - - new_sol.solve_time = 0 - new_sol.integration_time = 0 - new_sol.set_up_time = 0 - - return new_sol - - @property - def total_time(self): - return self.set_up_time + self.solve_time - - @property - def cycles(self): - return self._cycles - - @cycles.setter - def cycles(self, cycles): - self._cycles = cycles - - @property - def summary_variables(self): - return self._summary_variables - - def set_summary_variables(self, all_summary_variables): - summary_variables = {var: [] for var in all_summary_variables[0]} - for sum_vars in all_summary_variables: - for name, value in sum_vars.items(): - summary_variables[name].append(value) - - summary_variables["Cycle number"] = range(1, len(all_summary_variables) + 1) - self.all_summary_variables = all_summary_variables - self._summary_variables = pybamm.FuzzyDict( - {name: np.array(value) for name, value in summary_variables.items()} - ) - - def update(self, variables): - """Add ProcessedVariables to the dictionary of variables in the solution""" - # make sure that sensitivities are extracted if required - if isinstance(self._sensitivities, bool) and self._sensitivities: - self.extract_explicit_sensitivities() - - # Convert single entry to list - if isinstance(variables, str): - variables = [variables] - # Process - for key in variables: - cumtrapz_ic = None - pybamm.logger.debug("Post-processing {}".format(key)) - vars_pybamm = [model.variables_and_events[key] for model in self.all_models] - - # Iterate through all models, some may be in the list several times and - # therefore only get set up once - vars_casadi = [] - for i, (model, ys, inputs, var_pybamm) in enumerate( - zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm) - ): - if isinstance(var_pybamm, pybamm.ExplicitTimeIntegral): - cumtrapz_ic = var_pybamm.initial_condition - cumtrapz_ic = cumtrapz_ic.evaluate() - var_pybamm = var_pybamm.child - var_casadi = Solution.process_casadi_var( - var_pybamm, - inputs, - ys.shape, - ) - model._variables_casadi[key] = var_casadi - vars_pybamm[i] = var_pybamm - elif key in model._variables_casadi: - var_casadi = model._variables_casadi[key] - else: - var_casadi = Solution.process_casadi_var( - var_pybamm, - inputs, - ys.shape, - ) - model._variables_casadi[key] = var_casadi - vars_casadi.append(var_casadi) - var = pybamm.ProcessedVariable( - vars_pybamm, vars_casadi, self, cumtrapz_ic=cumtrapz_ic - ) - - # Save variable and data - self._variables[key] = var - self.data[key] = var.data - - @staticmethod - def process_casadi_var(var_pybamm, inputs, ys_shape): - t_MX = casadi.MX.sym("t") - y_MX = casadi.MX.sym("y", ys_shape[0]) - inputs_MX_dict = { - key: casadi.MX.sym("input", value.shape[0]) for key, value in inputs.items() - } - inputs_MX = casadi.vertcat(*[p for p in inputs_MX_dict.values()]) - var_sym = var_pybamm.to_casadi(t_MX, y_MX, inputs=inputs_MX_dict) - var_casadi = casadi.Function("variable", [t_MX, y_MX, inputs_MX], [var_sym]) - return var_casadi - - def __getitem__(self, key): - """Read a variable from the solution. Variables are created 'just in time', i.e. - only when they are called. - - Parameters - ---------- - key : str - The name of the variable - - Returns - ------- - :class:`pybamm.ProcessedVariable` - A variable that can be evaluated at any time or spatial point. The - underlying data for this variable is available in its attribute ".data" - """ - - # return it if it exists - if key in self._variables: - return self._variables[key] - else: - # otherwise create it, save it and then return it - self.update(key) - return self._variables[key] - - def plot(self, output_variables=None, **kwargs): - """ - A method to quickly plot the outputs of the solution. Creates a - :class:`pybamm.QuickPlot` object (with keyword arguments 'kwargs') and - then calls :meth:`pybamm.QuickPlot.dynamic_plot`. - - Parameters - ---------- - output_variables: list, optional - A list of the variables to plot. - **kwargs - Additional keyword arguments passed to - :meth:`pybamm.QuickPlot.dynamic_plot`. - For a list of all possible keyword arguments see :class:`pybamm.QuickPlot`. - """ - return pybamm.dynamic_plot(self, output_variables=output_variables, **kwargs) - - def save(self, filename): - """Save the whole solution using pickle""" - # No warning here if len(self.data)==0 as solution can be loaded - # and used to process new variables - - with open(filename, "wb") as f: - pickle.dump(self, f, pickle.HIGHEST_PROTOCOL) - - def get_data_dict(self, variables=None, short_names=None, cycles_and_steps=True): - """ - Construct a (standard python) dictionary of the solution data containing the - variables in `variables`. If `variables` is None then all variables are - returned. Any variable names in short_names are replaced with the corresponding - short name. - - If the solution has cycles, then the cycle numbers and step numbers are also - returned in the dictionary. - - Parameters - ---------- - variables : list, optional - List of variables to return. If None, returns all variables in solution.data - short_names : dict, optional - Dictionary of shortened names to use when saving. - cycles_and_steps : bool, optional - Whether to include the cycle numbers and step numbers in the dictionary - - Returns - ------- - dict - A dictionary of the solution data - """ - if variables is None: - # variables not explicitly provided -> save all variables that have been - # computed - data_long_names = self.data - else: - if isinstance(variables, str): - variables = [variables] - # otherwise, save only the variables specified - data_long_names = {} - for name in variables: - data_long_names[name] = self[name].data - if len(data_long_names) == 0: - raise ValueError( - """ - Solution does not have any data. Please provide a list of variables - to save. - """ - ) - - # Use any short names if provided - data_short_names = {} - short_names = short_names or {} - for name, var in data_long_names.items(): - name = short_names.get(name, name) # return name if no short name - data_short_names[name] = var - - # Save cycle number and step number if the solution has them - if cycles_and_steps and len(self.cycles) > 0: - data_short_names["Cycle"] = np.array([]) - data_short_names["Step"] = np.array([]) - for i, cycle in enumerate(self.cycles): - data_short_names["Cycle"] = np.concatenate( - [data_short_names["Cycle"], i * np.ones_like(cycle.t)] - ) - for j, step in enumerate(cycle.steps): - data_short_names["Step"] = np.concatenate( - [data_short_names["Step"], j * np.ones_like(step.t)] - ) - - return data_short_names - - def save_data( - self, filename=None, variables=None, to_format="pickle", short_names=None - ): - """ - Save solution data only (raw arrays) - - Parameters - ---------- - filename : str, optional - The name of the file to save data to. If None, then a str is returned - variables : list, optional - List of variables to save. If None, saves all of the variables that have - been created so far - to_format : str, optional - The format to save to. Options are: - - - 'pickle' (default): creates a pickle file with the data dictionary - - 'matlab': creates a .mat file, for loading in matlab - - 'csv': creates a csv file (0D variables only) - - 'json': creates a json file - short_names : dict, optional - Dictionary of shortened names to use when saving. This may be necessary when - saving to MATLAB, since no spaces or special characters are allowed in - MATLAB variable names. Note that not all the variables need to be given - a short name. - - Returns - ------- - data : str, optional - str if 'csv' or 'json' is chosen and filename is None, otherwise None - """ - data = self.get_data_dict(variables=variables, short_names=short_names) - - if to_format == "pickle": - if filename is None: - raise ValueError("pickle format must be written to a file") - with open(filename, "wb") as f: - pickle.dump(data, f, pickle.HIGHEST_PROTOCOL) - elif to_format == "matlab": - if filename is None: - raise ValueError("matlab format must be written to a file") - # Check all the variable names only contain a-z, A-Z or _ or numbers - for name in data.keys(): - # Check the string only contains the following ASCII: - # a-z (97-122) - # A-Z (65-90) - # _ (95) - # 0-9 (48-57) but not in the first position - for i, s in enumerate(name): - if not ( - 97 <= ord(s) <= 122 - or 65 <= ord(s) <= 90 - or ord(s) == 95 - or (i > 0 and 48 <= ord(s) <= 57) - ): - raise ValueError( - "Invalid character '{}' found in '{}'. ".format(s, name) - + "MATLAB variable names must only contain a-z, A-Z, _, " - "or 0-9 (except the first position). " - "Use the 'short_names' argument to pass an alternative " - "variable name, e.g. \n\n" - "\tsolution.save_data(filename, " - "['Electrolyte concentration'], to_format='matlab, " - "short_names={'Electrolyte concentration': 'c_e'})" - ) - savemat(filename, data) - elif to_format == "csv": - for name, var in data.items(): - if var.ndim >= 2: - raise ValueError( - "only 0D variables can be saved to csv, but '{}' is {}D".format( - name, var.ndim - 1 - ) - ) - df = pd.DataFrame(data) - return df.to_csv(filename, index=False) - elif to_format == "json": - if filename is None: - return json.dumps(data, cls=NumpyEncoder) - else: - with open(filename, "w") as outfile: - json.dump(data, outfile, cls=NumpyEncoder) - else: - raise ValueError("format '{}' not recognised".format(to_format)) - - @property - def sub_solutions(self): - """List of sub solutions that have been - concatenated to form the full solution""" - - return self._sub_solutions - - def __add__(self, other): - """Adds two solutions together, e.g. when stepping""" - if other is None or isinstance(other, EmptySolution): - return self.copy() - if not isinstance(other, Solution): - raise pybamm.SolverError( - "Only a Solution or None can be added to a Solution" - ) - # Special case: new solution only has one timestep and it is already in the - # existing solution. In this case, return a copy of the existing solution - if ( - len(other.all_ts) == 1 - and len(other.all_ts[0]) == 1 - and other.all_ts[0][0] == self.all_ts[-1][-1] - ): - new_sol = self.copy() - # Update termination using the latter solution - new_sol._termination = other.termination - new_sol._t_event = other._t_event - new_sol._y_event = other._y_event - return new_sol - - # Update list of sub-solutions - if other.all_ts[0][0] == self.all_ts[-1][-1]: - # Skip first time step if it is repeated - all_ts = self.all_ts + [other.all_ts[0][1:]] + other.all_ts[1:] - all_ys = self.all_ys + [other.all_ys[0][:, 1:]] + other.all_ys[1:] - else: - all_ts = self.all_ts + other.all_ts - all_ys = self.all_ys + other.all_ys - - new_sol = Solution( - all_ts, - all_ys, - self.all_models + other.all_models, - self.all_inputs + other.all_inputs, - other.t_event, - other.y_event, - other.termination, - bool(self.sensitivities), - ) - - new_sol.closest_event_idx = other.closest_event_idx - new_sol._all_inputs_casadi = self.all_inputs_casadi + other.all_inputs_casadi - - # Set solution time - new_sol.solve_time = self.solve_time + other.solve_time - new_sol.integration_time = self.integration_time + other.integration_time - - # Set sub_solutions - new_sol._sub_solutions = self.sub_solutions + other.sub_solutions - - return new_sol - - def __radd__(self, other): - return self.__add__(other) - - def copy(self): - new_sol = self.__class__( - self.all_ts, - self.all_ys, - self.all_models, - self.all_inputs, - self.t_event, - self.y_event, - self.termination, - ) - new_sol._all_inputs_casadi = self.all_inputs_casadi - new_sol._sub_solutions = self.sub_solutions - new_sol.closest_event_idx = self.closest_event_idx - - new_sol.solve_time = self.solve_time - new_sol.integration_time = self.integration_time - new_sol.set_up_time = self.set_up_time - - return new_sol - - -class EmptySolution: - def __init__(self, termination=None, t=None): - self.termination = termination - if t is None: - t = np.array([0]) - elif isinstance(t, numbers.Number): - t = np.array([t]) - - self.t = t - - def __add__(self, other): - if isinstance(other, (EmptySolution, Solution)): - return other.copy() - - def __radd__(self, other): - if other is None: - return self.copy() - - def copy(self): - return EmptySolution(termination=self.termination, t=self.t) - - -def make_cycle_solution(step_solutions, esoh_solver=None, save_this_cycle=True): - """ - Function to create a Solution for an entire cycle, and associated summary variables - - Parameters - ---------- - step_solutions : list of :class:`Solution` - Step solutions that form the entire cycle - esoh_solver : :class:`pybamm.lithium_ion.ElectrodeSOHSolver` - Solver to calculate electrode SOH (eSOH) variables. If `None` (default) - then only summary variables that do not require the eSOH calculation - are calculated. See :footcite:t:`Mohtat2019` for more details on eSOH variables. - save_this_cycle : bool, optional - Whether to save the entire cycle variables or just the summary variables. - Default True - - Returns - ------- - cycle_solution : :class:`pybamm.Solution` or None - The Solution object for this cycle, or None (if save_this_cycle is False) - cycle_summary_variables : dict - Dictionary of summary variables for this cycle - - """ - sum_sols = step_solutions[0].copy() - for step_solution in step_solutions[1:]: - sum_sols = sum_sols + step_solution - - cycle_solution = Solution( - sum_sols.all_ts, - sum_sols.all_ys, - sum_sols.all_models, - sum_sols.all_inputs, - sum_sols.t_event, - sum_sols.y_event, - sum_sols.termination, - ) - cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi - cycle_solution._sub_solutions = sum_sols.sub_solutions - - cycle_solution.solve_time = sum_sols.solve_time - cycle_solution.integration_time = sum_sols.integration_time - cycle_solution.set_up_time = sum_sols.set_up_time - - cycle_solution.steps = step_solutions - - cycle_summary_variables = _get_cycle_summary_variables(cycle_solution, esoh_solver) - - cycle_first_state = cycle_solution.first_state - - if save_this_cycle: - cycle_solution.cycle_summary_variables = cycle_summary_variables - else: - cycle_solution = None - - return cycle_solution, cycle_summary_variables, cycle_first_state - - -def _get_cycle_summary_variables(cycle_solution, esoh_solver): - model = cycle_solution.all_models[0] - cycle_summary_variables = pybamm.FuzzyDict({}) - - # Measured capacity variables - if "Discharge capacity [A.h]" in model.variables: - Q = cycle_solution["Discharge capacity [A.h]"].data - min_Q, max_Q = np.min(Q), np.max(Q) - - cycle_summary_variables.update( - { - "Minimum measured discharge capacity [A.h]": min_Q, - "Maximum measured discharge capacity [A.h]": max_Q, - "Measured capacity [A.h]": max_Q - min_Q, - } - ) - - # Voltage variables - if "Battery voltage [V]" in model.variables: - V = cycle_solution["Battery voltage [V]"].data - min_V, max_V = np.min(V), np.max(V) - - cycle_summary_variables.update( - {"Minimum voltage [V]": min_V, "Maximum voltage [V]": max_V} - ) - - # Degradation variables - degradation_variables = model.summary_variables - first_state = cycle_solution.first_state - last_state = cycle_solution.last_state - for var in degradation_variables: - data_first = first_state[var].data - data_last = last_state[var].data - cycle_summary_variables[var] = data_last[0] - var_lowercase = var[0].lower() + var[1:] - cycle_summary_variables["Change in " + var_lowercase] = ( - data_last[0] - data_first[0] - ) - - # eSOH variables (full-cell lithium-ion model only, for now) - if ( - esoh_solver is not None - and isinstance(model, pybamm.lithium_ion.BaseModel) - and model.options.electrode_types["negative"] == "porous" - ): - Q_n = last_state["Negative electrode capacity [A.h]"].data[0] - Q_p = last_state["Positive electrode capacity [A.h]"].data[0] - Q_Li = last_state["Total lithium capacity in particles [A.h]"].data[0] - - inputs = {"Q_n": Q_n, "Q_p": Q_p, "Q_Li": Q_Li} - - try: - esoh_sol = esoh_solver.solve(inputs) - except pybamm.SolverError: # pragma: no cover - raise pybamm.SolverError( - "Could not solve for summary variables, run " - "`sim.solve(calc_esoh=False)` to skip this step" - ) - - cycle_summary_variables.update(esoh_sol) - - return cycle_summary_variables + # Solution contains only the requested variables + return pybamm.SolutionFull(*args, **kwargs) diff --git a/pybamm/solvers/solution_base.py b/pybamm/solvers/solution_base.py new file mode 100644 index 0000000000..3cdde7dc63 --- /dev/null +++ b/pybamm/solvers/solution_base.py @@ -0,0 +1,922 @@ +# +# SolutionBase class +# +import casadi +import json +import numbers +import numpy as np +import pickle +import pybamm +import pandas as pd +from scipy.io import savemat +from functools import cached_property + + +class NumpyEncoder(json.JSONEncoder): + """ + Numpy serialiser helper class that converts numpy arrays to a list + https://stackoverflow.com/questions/26646362/numpy-array-is-not-json-serializable + """ + + def default(self, obj): + if isinstance(obj, np.ndarray): + return obj.tolist() + # won't be called since we only need to convert numpy arrays + return json.JSONEncoder.default(self, obj) # pragma: no cover + + +class SolutionBase(object): + """ + Class containing the solution of, and various attributes associated with, a PyBaMM + model. + + Parameters + ---------- + all_ts : :class:`numpy.array`, size (n,) (or list of these) + A one-dimensional array containing the times at which the solution is evaluated. + A list of times can be provided instead to initialize a solution with + sub-solutions. + all_ys : :class:`numpy.array`, size (m, n) (or list of these) + A two-dimensional array containing the values of the solution. y[i, :] is the + vector of solutions at time t[i]. + A list of ys can be provided instead to initialize a solution with + sub-solutions. + all_models : :class:`pybamm.BaseModel` + The model that was used to calculate the solution. + A list of models can be provided instead to initialize a solution with + sub-solutions that have been calculated using those models. + all_inputs : dict (or list of these) + The inputs that were used to calculate the solution + A list of inputs can be provided instead to initialize a solution with + sub-solutions. + t_event : :class:`numpy.array`, size (1,) + A zero-dimensional array containing the time at which the event happens. + y_event : :class:`numpy.array`, size (m,) + A one-dimensional array containing the value of the solution at the time when + the event happens. + termination : str + String to indicate why the solution terminated + + sensitivities: bool or dict + True if sensitivities included as the solution of the explicit forwards + equations. False if no sensitivities included/wanted. Dict if sensitivities are + provided as a dict of {parameter: sensitivities} pairs. + + """ + + def __init__( + self, + all_ts, + all_ys, + all_models, + all_inputs, + t_event=None, + y_event=None, + termination="final time", + sensitivities=False, + check_solution=True, + ): + if not isinstance(all_ts, list): + all_ts = [all_ts] + if not isinstance(all_ys, list): + all_ys = [all_ys] + if not isinstance(all_models, list): + all_models = [all_models] + self._all_ts = all_ts + self._all_ys = all_ys + self._all_ys_and_sens = all_ys + self._all_models = all_models + + # Set up inputs + if not isinstance(all_inputs, list): + all_inputs_copy = dict(all_inputs) + for key, value in all_inputs_copy.items(): + if isinstance(value, numbers.Number): + all_inputs_copy[key] = np.array([value]) + self.all_inputs = [all_inputs_copy] + else: + self.all_inputs = all_inputs + + self.sensitivities = sensitivities + + self._t_event = t_event + self._y_event = y_event + self._termination = termination + + # Check no ys are too large + if check_solution: + self.check_ys_are_not_too_large() + + # Events + self._t_event = t_event + self._y_event = y_event + self._termination = termination + self.closest_event_idx = None + + # Initialize times + self.set_up_time = None + self.solve_time = None + self.integration_time = None + + # initialize empty variables and data + self._variables = pybamm.FuzzyDict() + self.data = pybamm.FuzzyDict() + + # Add self as sub-solution for compatibility with ProcessedVariable + self._sub_solutions = [self] + + # initialize empty cycles + self._cycles = [] + + # Initialize empty summary variables + self._summary_variables = None + + # Solution now uses CasADi + pybamm.citations.register("Andersson2019") + + def extract_explicit_sensitivities(self): + # if we got here, we haven't set y yet + self.set_y() + + # extract sensitivities from full y solution + self._y, self._sensitivities = self._extract_explicit_sensitivities( + self.all_models[0], self.y, self.t, self.all_inputs[0] + ) + + # make sure we remove all sensitivities from all_ys + for index, (model, ys, ts, inputs) in enumerate( + zip(self.all_models, self.all_ys, self.all_ts, self.all_inputs) + ): + self._all_ys[index], _ = self._extract_explicit_sensitivities( + model, ys, ts, inputs + ) + + def _extract_explicit_sensitivities(self, model, y, t_eval, inputs): + """ + given a model and a solution y, extracts the sensitivities + + Parameters + -------- + model : :class:`pybamm.BaseModel` + A model that has been already setup by this base solver + y: ndarray + The solution of the full explicit sensitivity equations + t_eval: ndarray + The evaluation times + inputs: dict + parameter inputs + + Returns + ------- + y: ndarray + The solution of the ode/dae in model + sensitivities: dict of (string: ndarray) + A dictionary of parameter names, and the corresponding solution of + the sensitivity equations + """ + + n_states = model.len_rhs_and_alg + n_rhs = model.len_rhs + n_alg = model.len_alg + # Get the point where the algebraic equations start + if model.len_rhs != 0: + n_p = model.len_rhs_sens // model.len_rhs + else: + n_p = model.len_alg_sens // model.len_alg + len_rhs_and_sens = model.len_rhs + model.len_rhs_sens + + n_t = len(t_eval) + # y gets the part of the solution vector that correspond to the + # actual ODE/DAE solution + + # save sensitivities as a dictionary + # first save the whole sensitivity matrix + # reshape using Fortran order to get the right array: + # t0_x0_p0, t0_x0_p1, ..., t0_x0_pn + # t0_x1_p0, t0_x1_p1, ..., t0_x1_pn + # ... + # t0_xn_p0, t0_xn_p1, ..., t0_xn_pn + # t1_x0_p0, t1_x0_p1, ..., t1_x0_pn + # t1_x1_p0, t1_x1_p1, ..., t1_x1_pn + # ... + # t1_xn_p0, t1_xn_p1, ..., t1_xn_pn + # ... + # tn_x0_p0, tn_x0_p1, ..., tn_x0_pn + # tn_x1_p0, tn_x1_p1, ..., tn_x1_pn + # ... + # tn_xn_p0, tn_xn_p1, ..., tn_xn_pn + # 1, Extract rhs and alg sensitivities and reshape into 3D matrices + # with shape (n_p, n_states, n_t) + if isinstance(y, casadi.DM): + y_full = y.full() + else: + y_full = y + ode_sens = y_full[n_rhs:len_rhs_and_sens, :].reshape(n_p, n_rhs, n_t) + alg_sens = y_full[len_rhs_and_sens + n_alg :, :].reshape(n_p, n_alg, n_t) + # 2. Concatenate into a single 3D matrix with shape (n_p, n_states, n_t) + # i.e. along first axis + full_sens_matrix = np.concatenate([ode_sens, alg_sens], axis=1) + # Transpose and reshape into a (n_states * n_t, n_p) matrix + full_sens_matrix = full_sens_matrix.transpose(2, 1, 0).reshape( + n_t * n_states, n_p + ) + + # Save the full sensitivity matrix + sensitivity = {"all": full_sens_matrix} + + # also save the sensitivity wrt each parameter (read the columns of the + # sensitivity matrix) + start = 0 + for name in model.calculate_sensitivities: + inp = inputs[name] + input_size = inp.shape[0] + end = start + input_size + sensitivity[name] = full_sens_matrix[:, start:end] + start = end + + y_dae = np.vstack( + [ + y[: model.len_rhs, :], + y[len_rhs_and_sens : len_rhs_and_sens + model.len_alg, :], + ] + ) + return y_dae, sensitivity + + @property + def t(self): + """Times at which the solution is evaluated""" + try: + return self._t + except AttributeError: + self.set_t() + return self._t + + def set_t(self): + self._t = np.concatenate(self.all_ts) + if any(np.diff(self._t) <= 0): + raise ValueError("Solution time vector must be strictly increasing") + + @property + def y(self): + """Values of the solution""" + try: + return self._y + except AttributeError: + self.set_y() + + # if y is evaluated before sensitivities then need to extract them + if isinstance(self._sensitivities, bool) and self._sensitivities: + self.extract_explicit_sensitivities() + + return self._y + + @property + def sensitivities(self): + """Values of the sensitivities. Returns a dict of param_name: np_array""" + if isinstance(self._sensitivities, bool): + if self._sensitivities: + self.extract_explicit_sensitivities() + else: + self._sensitivities = {} + return self._sensitivities + + @sensitivities.setter + def sensitivities(self, value): + """Updates the sensitivity""" + # sensitivities must be a dict or bool + if not isinstance(value, (bool, dict)): + raise TypeError("sensitivities arg needs to be a bool or dict") + self._sensitivities = value + + def set_y(self): + try: + if isinstance(self.all_ys[0], (casadi.DM, casadi.MX)): + self._y = casadi.horzcat(*self.all_ys) + else: + self._y = np.hstack(self.all_ys) + except ValueError: + raise pybamm.SolverError( + "The solution is made up from different models, so `y` cannot be " + "computed explicitly." + ) + + def check_ys_are_not_too_large(self): + # Only check last one so that it doesn't take too long + # We only care about the cases where y is growing too large without any + # restraint, so if y gets large in the middle then comes back down that is ok + y, model = self.all_ys[-1], self.all_models[-1] + y = y[:, -1] + if np.any(y > pybamm.settings.max_y_value): + for var in [*model.rhs.keys(), *model.algebraic.keys()]: + y_var = y[model.variables[var.name].y_slices[0]] + if np.any(y_var > pybamm.settings.max_y_value): + pybamm.logger.error( + f"Solution for '{var}' exceeds the maximum allowed value " + f"of `{pybamm.settings.max_y_value}. This could be due to " + "incorrect scaling, model formulation, or " + "parameter values. The maximum allowed value is set by " + "'pybammm.settings.max_y_value'." + ) + + @property + def all_ts(self): + return self._all_ts + + @property + def all_ys(self): + return self._all_ys + + @property + def all_models(self): + """Model(s) used for solution""" + return self._all_models + + @cached_property + def all_inputs_casadi(self): + return [casadi.vertcat(*inp.values()) for inp in self.all_inputs] + + @property + def t_event(self): + """Time at which the event happens""" + return self._t_event + + @property + def y_event(self): + """Value of the solution at the time of the event""" + return self._y_event + + @property + def termination(self): + """Reason for termination""" + return self._termination + + @termination.setter + def termination(self, value): + """Updates the reason for termination""" + self._termination = value + + @cached_property + def first_state(self): + """ + A Solution object that only contains the first state. This is faster to evaluate + than the full solution when only the first state is needed (e.g. to initialize + a model with the solution) + """ + new_sol = Solution( + self.all_ts[0][:1], + self.all_ys[0][:, :1], + self.all_models[:1], + self.all_inputs[:1], + None, + None, + "final time", + ) + new_sol._all_inputs_casadi = self.all_inputs_casadi[:1] + new_sol._sub_solutions = self.sub_solutions[:1] + + new_sol.solve_time = 0 + new_sol.integration_time = 0 + new_sol.set_up_time = 0 + + return new_sol + + @cached_property + def last_state(self): + """ + A Solution object that only contains the final state. This is faster to evaluate + than the full solution when only the final state is needed (e.g. to initialize + a model with the solution) + """ + new_sol = Solution( + self.all_ts[-1][-1:], + self.all_ys[-1][:, -1:], + self.all_models[-1:], + self.all_inputs[-1:], + self.t_event, + self.y_event, + self.termination, + ) + new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:] + new_sol._sub_solutions = self.sub_solutions[-1:] + + new_sol.solve_time = 0 + new_sol.integration_time = 0 + new_sol.set_up_time = 0 + + return new_sol + + @property + def total_time(self): + return self.set_up_time + self.solve_time + + @property + def cycles(self): + return self._cycles + + @cycles.setter + def cycles(self, cycles): + self._cycles = cycles + + @property + def summary_variables(self): + return self._summary_variables + + def set_summary_variables(self, all_summary_variables): + summary_variables = {var: [] for var in all_summary_variables[0]} + for sum_vars in all_summary_variables: + for name, value in sum_vars.items(): + summary_variables[name].append(value) + + summary_variables["Cycle number"] = range(1, len(all_summary_variables) + 1) + self.all_summary_variables = all_summary_variables + self._summary_variables = pybamm.FuzzyDict( + {name: np.array(value) for name, value in summary_variables.items()} + ) + + def update(self, variables): + """Add ProcessedVariables to the dictionary of variables in the solution""" + # make sure that sensitivities are extracted if required + if isinstance(self._sensitivities, bool) and self._sensitivities: + self.extract_explicit_sensitivities() + + # Convert single entry to list + if isinstance(variables, str): + variables = [variables] + # Process + for key in variables: + cumtrapz_ic = None + pybamm.logger.debug("Post-processing {}".format(key)) + vars_pybamm = [model.variables_and_events[key] for model in self.all_models] + + # Iterate through all models, some may be in the list several times and + # therefore only get set up once + vars_casadi = [] + for i, (model, ys, inputs, var_pybamm) in enumerate( + zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm) + ): + if isinstance(var_pybamm, pybamm.ExplicitTimeIntegral): + cumtrapz_ic = var_pybamm.initial_condition + cumtrapz_ic = cumtrapz_ic.evaluate() + var_pybamm = var_pybamm.child + var_casadi = self.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) + model._variables_casadi[key] = var_casadi + vars_pybamm[i] = var_pybamm + elif key in model._variables_casadi: + var_casadi = model._variables_casadi[key] + else: + var_casadi = self.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) + model._variables_casadi[key] = var_casadi + vars_casadi.append(var_casadi) + var = pybamm.ProcessedVariable( + vars_pybamm, vars_casadi, self, cumtrapz_ic=cumtrapz_ic + ) + + # Save variable and data + self._variables[key] = var + self.data[key] = var.data + + def process_casadi_var(self, var_pybamm, inputs, ys_shape): + t_MX = casadi.MX.sym("t") + y_MX = casadi.MX.sym("y", ys_shape[0]) + inputs_MX_dict = { + key: casadi.MX.sym("input", value.shape[0]) for key, value in inputs.items() + } + inputs_MX = casadi.vertcat(*[p for p in inputs_MX_dict.values()]) + var_sym = var_pybamm.to_casadi(t_MX, y_MX, inputs=inputs_MX_dict) + var_casadi = casadi.Function("variable", [t_MX, y_MX, inputs_MX], [var_sym]) + return var_casadi + + def __getitem__(self, key): + """Read a variable from the solution. Variables are created 'just in time', i.e. + only when they are called. + + Parameters + ---------- + key : str + The name of the variable + + Returns + ------- + :class:`pybamm.ProcessedVariable` + A variable that can be evaluated at any time or spatial point. The + underlying data for this variable is available in its attribute ".data" + """ + + # return it if it exists + if key in self._variables: + return self._variables[key] + else: + # otherwise create it, save it and then return it + self.update(key) + return self._variables[key] + + def plot(self, output_variables=None, **kwargs): + """ + A method to quickly plot the outputs of the solution. Creates a + :class:`pybamm.QuickPlot` object (with keyword arguments 'kwargs') and + then calls :meth:`pybamm.QuickPlot.dynamic_plot`. + + Parameters + ---------- + output_variables: list, optional + A list of the variables to plot. + **kwargs + Additional keyword arguments passed to + :meth:`pybamm.QuickPlot.dynamic_plot`. + For a list of all possible keyword arguments see :class:`pybamm.QuickPlot`. + """ + return pybamm.dynamic_plot(self, output_variables=output_variables, **kwargs) + + def save(self, filename): + """Save the whole solution using pickle""" + # No warning here if len(self.data)==0 as solution can be loaded + # and used to process new variables + + with open(filename, "wb") as f: + pickle.dump(self, f, pickle.HIGHEST_PROTOCOL) + + def get_data_dict(self, variables=None, short_names=None, cycles_and_steps=True): + """ + Construct a (standard python) dictionary of the solution data containing the + variables in `variables`. If `variables` is None then all variables are + returned. Any variable names in short_names are replaced with the corresponding + short name. + + If the solution has cycles, then the cycle numbers and step numbers are also + returned in the dictionary. + + Parameters + ---------- + variables : list, optional + List of variables to return. If None, returns all variables in solution.data + short_names : dict, optional + Dictionary of shortened names to use when saving. + cycles_and_steps : bool, optional + Whether to include the cycle numbers and step numbers in the dictionary + + Returns + ------- + dict + A dictionary of the solution data + """ + if variables is None: + # variables not explicitly provided -> save all variables that have been + # computed + data_long_names = self.data + else: + if isinstance(variables, str): + variables = [variables] + # otherwise, save only the variables specified + data_long_names = {} + for name in variables: + data_long_names[name] = self[name].data + if len(data_long_names) == 0: + raise ValueError( + """ + Solution does not have any data. Please provide a list of variables + to save. + """ + ) + + # Use any short names if provided + data_short_names = {} + short_names = short_names or {} + for name, var in data_long_names.items(): + name = short_names.get(name, name) # return name if no short name + data_short_names[name] = var + + # Save cycle number and step number if the solution has them + if cycles_and_steps and len(self.cycles) > 0: + data_short_names["Cycle"] = np.array([]) + data_short_names["Step"] = np.array([]) + for i, cycle in enumerate(self.cycles): + data_short_names["Cycle"] = np.concatenate( + [data_short_names["Cycle"], i * np.ones_like(cycle.t)] + ) + for j, step in enumerate(cycle.steps): + data_short_names["Step"] = np.concatenate( + [data_short_names["Step"], j * np.ones_like(step.t)] + ) + + return data_short_names + + def save_data( + self, filename=None, variables=None, to_format="pickle", short_names=None + ): + """ + Save solution data only (raw arrays) + + Parameters + ---------- + filename : str, optional + The name of the file to save data to. If None, then a str is returned + variables : list, optional + List of variables to save. If None, saves all of the variables that have + been created so far + to_format : str, optional + The format to save to. Options are: + + - 'pickle' (default): creates a pickle file with the data dictionary + - 'matlab': creates a .mat file, for loading in matlab + - 'csv': creates a csv file (0D variables only) + - 'json': creates a json file + short_names : dict, optional + Dictionary of shortened names to use when saving. This may be necessary when + saving to MATLAB, since no spaces or special characters are allowed in + MATLAB variable names. Note that not all the variables need to be given + a short name. + + Returns + ------- + data : str, optional + str if 'csv' or 'json' is chosen and filename is None, otherwise None + """ + data = self.get_data_dict(variables=variables, short_names=short_names) + + if to_format == "pickle": + if filename is None: + raise ValueError("pickle format must be written to a file") + with open(filename, "wb") as f: + pickle.dump(data, f, pickle.HIGHEST_PROTOCOL) + elif to_format == "matlab": + if filename is None: + raise ValueError("matlab format must be written to a file") + # Check all the variable names only contain a-z, A-Z or _ or numbers + for name in data.keys(): + # Check the string only contains the following ASCII: + # a-z (97-122) + # A-Z (65-90) + # _ (95) + # 0-9 (48-57) but not in the first position + for i, s in enumerate(name): + if not ( + 97 <= ord(s) <= 122 + or 65 <= ord(s) <= 90 + or ord(s) == 95 + or (i > 0 and 48 <= ord(s) <= 57) + ): + raise ValueError( + "Invalid character '{}' found in '{}'. ".format(s, name) + + "MATLAB variable names must only contain a-z, A-Z, _, " + "or 0-9 (except the first position). " + "Use the 'short_names' argument to pass an alternative " + "variable name, e.g. \n\n" + "\tsolution.save_data(filename, " + "['Electrolyte concentration'], to_format='matlab, " + "short_names={'Electrolyte concentration': 'c_e'})" + ) + savemat(filename, data) + elif to_format == "csv": + for name, var in data.items(): + if var.ndim >= 2: + raise ValueError( + "only 0D variables can be saved to csv, but '{}' is {}D".format( + name, var.ndim - 1 + ) + ) + df = pd.DataFrame(data) + return df.to_csv(filename, index=False) + elif to_format == "json": + if filename is None: + return json.dumps(data, cls=NumpyEncoder) + else: + with open(filename, "w") as outfile: + json.dump(data, outfile, cls=NumpyEncoder) + else: + raise ValueError("format '{}' not recognised".format(to_format)) + + @property + def sub_solutions(self): + """List of sub solutions that have been + concatenated to form the full solution""" + + return self._sub_solutions + + def __add__(self, other): + """Adds two solutions together, e.g. when stepping""" + if other is None or isinstance(other, EmptySolution): + return self.copy() + if not isinstance(other, Solution): + raise pybamm.SolverError( + "Only a Solution or None can be added to a Solution" + ) + # Special case: new solution only has one timestep and it is already in the + # existing solution. In this case, return a copy of the existing solution + if ( + len(other.all_ts) == 1 + and len(other.all_ts[0]) == 1 + and other.all_ts[0][0] == self.all_ts[-1][-1] + ): + new_sol = self.copy() + # Update termination using the latter solution + new_sol._termination = other.termination + new_sol._t_event = other._t_event + new_sol._y_event = other._y_event + return new_sol + + # Update list of sub-solutions + if other.all_ts[0][0] == self.all_ts[-1][-1]: + # Skip first time step if it is repeated + all_ts = self.all_ts + [other.all_ts[0][1:]] + other.all_ts[1:] + all_ys = self.all_ys + [other.all_ys[0][:, 1:]] + other.all_ys[1:] + else: + all_ts = self.all_ts + other.all_ts + all_ys = self.all_ys + other.all_ys + + new_sol = Solution( + all_ts, + all_ys, + self.all_models + other.all_models, + self.all_inputs + other.all_inputs, + other.t_event, + other.y_event, + other.termination, + bool(self.sensitivities), + ) + + new_sol.closest_event_idx = other.closest_event_idx + new_sol._all_inputs_casadi = self.all_inputs_casadi + other.all_inputs_casadi + + # Set solution time + new_sol.solve_time = self.solve_time + other.solve_time + new_sol.integration_time = self.integration_time + other.integration_time + + # Set sub_solutions + new_sol._sub_solutions = self.sub_solutions + other.sub_solutions + + return new_sol + + def __radd__(self, other): + return self.__add__(other) + + def copy(self): + new_sol = self.__class__( + self.all_ts, + self.all_ys, + self.all_models, + self.all_inputs, + self.t_event, + self.y_event, + self.termination, + ) + new_sol._all_inputs_casadi = self.all_inputs_casadi + new_sol._sub_solutions = self.sub_solutions + new_sol.closest_event_idx = self.closest_event_idx + + new_sol.solve_time = self.solve_time + new_sol.integration_time = self.integration_time + new_sol.set_up_time = self.set_up_time + + return new_sol + + +class EmptySolution: + def __init__(self, termination=None, t=None): + self.termination = termination + if t is None: + t = np.array([0]) + elif isinstance(t, numbers.Number): + t = np.array([t]) + + self.t = t + + def __add__(self, other): + if isinstance(other, (EmptySolution, Solution)): + return other.copy() + + def __radd__(self, other): + if other is None: + return self.copy() + + def copy(self): + return EmptySolution(termination=self.termination, t=self.t) + + +def make_cycle_solution(step_solutions, esoh_solver=None, save_this_cycle=True): + """ + Function to create a Solution for an entire cycle, and associated summary variables + + Parameters + ---------- + step_solutions : list of :class:`Solution` + Step solutions that form the entire cycle + esoh_solver : :class:`pybamm.lithium_ion.ElectrodeSOHSolver` + Solver to calculate electrode SOH (eSOH) variables. If `None` (default) + then only summary variables that do not require the eSOH calculation + are calculated. See :footcite:t:`Mohtat2019` for more details on eSOH variables. + save_this_cycle : bool, optional + Whether to save the entire cycle variables or just the summary variables. + Default True + + Returns + ------- + cycle_solution : :class:`pybamm.Solution` or None + The Solution object for this cycle, or None (if save_this_cycle is False) + cycle_summary_variables : dict + Dictionary of summary variables for this cycle + + """ + sum_sols = step_solutions[0].copy() + for step_solution in step_solutions[1:]: + sum_sols = sum_sols + step_solution + + cycle_solution = Solution( + sum_sols.all_ts, + sum_sols.all_ys, + sum_sols.all_models, + sum_sols.all_inputs, + sum_sols.t_event, + sum_sols.y_event, + sum_sols.termination, + ) + cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi + cycle_solution._sub_solutions = sum_sols.sub_solutions + + cycle_solution.solve_time = sum_sols.solve_time + cycle_solution.integration_time = sum_sols.integration_time + cycle_solution.set_up_time = sum_sols.set_up_time + + cycle_solution.steps = step_solutions + + cycle_summary_variables = _get_cycle_summary_variables(cycle_solution, esoh_solver) + + cycle_first_state = cycle_solution.first_state + + if save_this_cycle: + cycle_solution.cycle_summary_variables = cycle_summary_variables + else: + cycle_solution = None + + return cycle_solution, cycle_summary_variables, cycle_first_state + + +def _get_cycle_summary_variables(cycle_solution, esoh_solver): + model = cycle_solution.all_models[0] + cycle_summary_variables = pybamm.FuzzyDict({}) + + # Measured capacity variables + if "Discharge capacity [A.h]" in model.variables: + Q = cycle_solution["Discharge capacity [A.h]"].data + min_Q, max_Q = np.min(Q), np.max(Q) + + cycle_summary_variables.update( + { + "Minimum measured discharge capacity [A.h]": min_Q, + "Maximum measured discharge capacity [A.h]": max_Q, + "Measured capacity [A.h]": max_Q - min_Q, + } + ) + + # Voltage variables + if "Battery voltage [V]" in model.variables: + V = cycle_solution["Battery voltage [V]"].data + min_V, max_V = np.min(V), np.max(V) + + cycle_summary_variables.update( + {"Minimum voltage [V]": min_V, "Maximum voltage [V]": max_V} + ) + + # Degradation variables + degradation_variables = model.summary_variables + first_state = cycle_solution.first_state + last_state = cycle_solution.last_state + for var in degradation_variables: + data_first = first_state[var].data + data_last = last_state[var].data + cycle_summary_variables[var] = data_last[0] + var_lowercase = var[0].lower() + var[1:] + cycle_summary_variables["Change in " + var_lowercase] = ( + data_last[0] - data_first[0] + ) + + # eSOH variables (full-cell lithium-ion model only, for now) + if ( + esoh_solver is not None + and isinstance(model, pybamm.lithium_ion.BaseModel) + and model.options.electrode_types["negative"] == "porous" + ): + Q_n = last_state["Negative electrode capacity [A.h]"].data[0] + Q_p = last_state["Positive electrode capacity [A.h]"].data[0] + Q_Li = last_state["Total lithium capacity in particles [A.h]"].data[0] + + inputs = {"Q_n": Q_n, "Q_p": Q_p, "Q_Li": Q_Li} + + try: + esoh_sol = esoh_solver.solve(inputs) + except pybamm.SolverError: # pragma: no cover + raise pybamm.SolverError( + "Could not solve for summary variables, run " + "`sim.solve(calc_esoh=False)` to skip this step" + ) + + cycle_summary_variables.update(esoh_sol) + + return cycle_summary_variables diff --git a/pybamm/solvers/solution_full.py b/pybamm/solvers/solution_full.py new file mode 100644 index 0000000000..a73363b6ec --- /dev/null +++ b/pybamm/solvers/solution_full.py @@ -0,0 +1,8 @@ +# +# Solution class +# +import pybamm + + +class SolutionFull(pybamm.SolutionBase): + ... diff --git a/pybamm/solvers/solution_vars.py b/pybamm/solvers/solution_vars.py new file mode 100644 index 0000000000..dc0df83833 --- /dev/null +++ b/pybamm/solvers/solution_vars.py @@ -0,0 +1,8 @@ +# +# Solution class +# +import pybamm + + +class SolutionVars(pybamm.SolutionBase): + ... diff --git a/test.py b/test.py index e0f72e83d2..830ad48d6e 100644 --- a/test.py +++ b/test.py @@ -6,7 +6,7 @@ "Voltage [V]", "Current [A]", ] -#output_variables = [] +output_variables = [] import pybamm import numpy as np From df9568d10606eb6282cf25ea584fe0b8f32c2db6 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Tue, 18 Jul 2023 08:57:01 +0000 Subject: [PATCH 06/38] Fix incorrect input structure being passed when building variables' casadi functions --- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 4 +- .../c_solvers/idaklu/casadi_functions.cpp | 4 +- pybamm/solvers/idaklu_solver.py | 28 ++++++--- test.py | 59 ++++++++++++++++--- 4 files changed, 74 insertions(+), 21 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index ce968876c8..f49e21bc7c 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -327,7 +327,7 @@ Solution CasadiSolverOpenMP::solve( // Initial state (t_i=0) t_return[0] = t(0); - realtype *res = new realtype[max_res_size]; + realtype *res = new realtype[number_of_states]; // TODO: Crashes if set to max_res_size if (functions->var_casadi_fcns.size() > 0) { // Evaluate casadi functions for each requested variable and store size_t j = 0; @@ -373,7 +373,6 @@ Solution CasadiSolverOpenMP::solve( // Evaluate and store results for the time step t_return[t_i] = tret; size_t j = 0; - std::cout << "Timestep " << t_return[t_i] << std::endl; if (functions->var_casadi_fcns.size() > 0) { // Evaluate casadi functions for each requested variable and store for (auto& var_fcn : functions->var_casadi_fcns) { @@ -385,7 +384,6 @@ Solution CasadiSolverOpenMP::solve( // store in return vector for (size_t jj=0; jj Date: Wed, 19 Jul 2023 22:33:34 +0000 Subject: [PATCH 07/38] Calculate sensitivities in Idaklu; uses temporary return structure --- pybamm/solvers/c_solvers/idaklu.cpp | 2 + .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 155 ++++++++++++++---- .../c_solvers/idaklu/CasadiSolverOpenMP.hpp | 18 ++ .../c_solvers/idaklu/casadi_functions.cpp | 21 ++- .../c_solvers/idaklu/casadi_functions.hpp | 8 +- .../c_solvers/idaklu/casadi_solver.cpp | 4 + .../c_solvers/idaklu/casadi_solver.hpp | 2 + pybamm/solvers/idaklu_solver.py | 65 ++++++-- pybamm/solvers/processed_variable.py | 3 + test.py | 36 ++-- 10 files changed, 252 insertions(+), 62 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index f5b8b23aed..720dcd3be2 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -81,6 +81,8 @@ PYBIND11_MODULE(idaklu, m) py::arg("rtol"), py::arg("inputs"), py::arg("var_casadi_fcns"), + py::arg("dvar_dy_fcns"), + py::arg("dvar_dp_fcns"), py::arg("options"), py::return_value_policy::take_ownership); diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index f49e21bc7c..f5d8db23ec 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -1,5 +1,8 @@ #include "CasadiSolverOpenMP.hpp" #include "casadi_sundials_functions.hpp" +#include +#include +#include /* * This is an abstract class that implements an OpenMP solution but @@ -218,6 +221,70 @@ CasadiSolverOpenMP::~CasadiSolverOpenMP() SUNContext_Free(&sunctx); } +void CasadiSolverOpenMP::CalcVars( + realtype *y_return, + size_t length_of_return_vector, + size_t t_i, + realtype *tret, + realtype *yval, + std::vector ySval, + realtype *yS_return, + size_t *ySk +) { + // Evaluate casadi functions for each requested variable and store + size_t j = 0; + for (auto& var_fcn : functions->var_casadi_fcns) { + var_fcn({tret, yval, functions->inputs.data()}, {res}); + // store in return vector + for (size_t jj=0; jj ySval, + realtype *yS_return, + size_t *ySk +) { + // Calculate sensitivities + + // Loop over variables + realtype* dens_dvar_dp = new realtype[number_of_parameters]; + for (size_t dvar_k=0; dvar_kdvar_dy_fcns.size(); dvar_k++) { + // Isolate functions + CasadiFunction dvar_dy = functions->dvar_dy_fcns[dvar_k]; + CasadiFunction dvar_dp = functions->dvar_dp_fcns[dvar_k]; + // Calculate dvar/dy + dvar_dy({tret, yval, functions->inputs.data()}, {res_dvar_dy}); + casadi::Sparsity spdy = dvar_dy.m_func.sparsity_out(0); + // Calculate dvar/dp and convert to dense array for indexing + dvar_dp({tret, yval, functions->inputs.data()}, {res_dvar_dp}); + casadi::Sparsity spdp = dvar_dp.m_func.sparsity_out(0); + for(int k=0; kvar_casadi_fcns.size() > 0) { // return only the requested variables list after computation for (auto& var_fcn : functions->var_casadi_fcns) { max_res_size = std::max(max_res_size, var_fcn.m_res.size()); - length_of_return_vector += max_res_size - 1; + if (var_fcn.m_res.size() > 0) + length_of_return_vector += var_fcn.m_res.size() - 1; + for (auto& dvar_fcn : functions->dvar_dy_fcns) + max_res_dvar_dy = std::max(max_res_dvar_dy, dvar_fcn.m_res.size()); + for (auto& dvar_fcn : functions->dvar_dp_fcns) + max_res_dvar_dp = std::max(max_res_dvar_dp, dvar_fcn.m_res.size()); } } else { // Return full y state-vector @@ -302,6 +375,11 @@ Solution CasadiSolverOpenMP::solve( realtype *yS_return = new realtype[number_of_parameters * number_of_timesteps * length_of_return_vector]; + std::cout << number_of_timesteps << " " << length_of_return_vector << std::endl; + + res = new realtype[number_of_states]; // TODO: Crashes if set to max_res_size + res_dvar_dy = new realtype[max_res_dvar_dy]; + res_dvar_dp = new realtype[max_res_dvar_dp]; py::capsule free_t_when_done( t_return, @@ -326,21 +404,13 @@ Solution CasadiSolverOpenMP::solve( ); // Initial state (t_i=0) - t_return[0] = t(0); - realtype *res = new realtype[number_of_states]; // TODO: Crashes if set to max_res_size + int t_i = 0; + size_t ySk = 0; + t_return[t_i] = t(t_i); if (functions->var_casadi_fcns.size() > 0) { // Evaluate casadi functions for each requested variable and store - size_t j = 0; - for (auto& var_fcn : functions->var_casadi_fcns) { - var_fcn.m_arg[0] = &tret; - var_fcn.m_arg[1] = yval; - var_fcn.m_arg[2] = functions->inputs.data(); - var_fcn.m_res[0] = res; - var_fcn(); - // store in return vector - for (size_t jj=0; jj0) int retval; - int t_i = 1; + t_i = 1; while (true) { t_next = t(t_i); @@ -372,19 +442,11 @@ Solution CasadiSolverOpenMP::solve( // Evaluate and store results for the time step t_return[t_i] = tret; - size_t j = 0; if (functions->var_casadi_fcns.size() > 0) { // Evaluate casadi functions for each requested variable and store - for (auto& var_fcn : functions->var_casadi_fcns) { - var_fcn.m_arg[0] = &tret; - var_fcn.m_arg[1] = yval; - var_fcn.m_arg[2] = functions->inputs.data(); - var_fcn.m_res[0] = res; - var_fcn(); - // store in return vector - for (size_t jj=0; jjdvar_dy_fcns.size() << ", " << ySk << ", " + << (length_of_return_vector*number_of_timesteps*number_of_parameters) + << std::endl; + if (ySk != (length_of_return_vector*number_of_timesteps*number_of_parameters)) + // throw std::runtime_error("Sensitivities vector has become misaligned."); + std::cout << "WARNING: Sensitivities vector has become misaligned."; + np_array t_ret = np_array( t_i, &t_return[0], @@ -421,15 +491,30 @@ Solution CasadiSolverOpenMP::solve( &y_return[0], free_y_when_done ); - np_array yS_ret = np_array( - std::vector { - number_of_parameters, - number_of_timesteps, - length_of_return_vector - }, - &yS_return[0], - free_yS_when_done - ); + // Note: Ordering of vector is differnet if computing variables vs returning + // the complete state vector + np_array yS_ret; + if (functions->var_casadi_fcns.size() > 0) { + yS_ret = np_array( + std::vector { + number_of_timesteps, + length_of_return_vector, + number_of_parameters + }, + &yS_return[0], + free_yS_when_done + ); + } else { + yS_ret = np_array( + std::vector { + number_of_parameters, + number_of_timesteps, + length_of_return_vector + }, + &yS_return[0], + free_yS_when_done + ); + } Solution sol(retval, t_ret, y_ret, yS_ret); diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp index 962b156d0a..e4d69a3d78 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -32,6 +32,9 @@ class CasadiSolverOpenMP : public CasadiSolver SUNMatrix J; SUNLinearSolver LS; std::unique_ptr functions; + realtype *res; + realtype *res_dvar_dy; + realtype *res_dvar_dp; Options options; #if SUNDIALS_VERSION_MAJOR >= 6 @@ -51,6 +54,21 @@ class CasadiSolverOpenMP : public CasadiSolver std::unique_ptr functions, const Options& options); ~CasadiSolverOpenMP(); + void CalcVars( + realtype *y_return, + size_t length_of_return_vector, + size_t t_i, + realtype *tret, + realtype *yval, + std::vector ySval, + realtype *yS_return, + size_t *ySk); + void CalcVarsSensitivities( + realtype *tret, + realtype *yval, + std::vector ySval, + realtype *yS_return, + size_t *ySk); Solution solve( np_array t_np, np_array y0_np, diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index c73fc5e643..4c6dfc7362 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -31,6 +31,18 @@ void CasadiFunction::operator()() m_func.release(mem); } +void CasadiFunction::operator()(std::vector inputs, + std::vector results) +{ + // Set-up input arguments, provide result vector, then execute function + // Example call: fcn({in1, in2, in3}, {out1}) + for(size_t k=0; k var_casadi_fcns, + const std::vector dvar_dy_fcns, + const std::vector dvar_dp_fcns, 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), @@ -53,10 +67,15 @@ CasadiFunctions::CasadiFunctions( options(options) { // convert casadi::Function list to CasadiFunction list - this->var_casadi_fcns.clear(); for (auto& var : var_casadi_fcns) { this->var_casadi_fcns.push_back(CasadiFunction(*var)); } + for (auto& var : dvar_dy_fcns) { + this->dvar_dy_fcns.push_back(CasadiFunction(*var)); + } + for (auto& var : dvar_dp_fcns) { + this->dvar_dp_fcns.push_back(CasadiFunction(*var)); + } // copy across numpy array values const int n_row_vals = jac_times_cjmass_rowvals_arg.request().size; diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 956c915e71..1feac1f359 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -42,12 +42,14 @@ class CasadiFunction public: explicit CasadiFunction(const Function &f); void operator()(); + void operator()(std::vector inputs, + std::vector results); public: std::vector m_arg; std::vector m_res; -private: +//private: const Function &m_func; std::vector m_iw; std::vector m_w; @@ -73,6 +75,8 @@ class CasadiFunctions const int n_e, const int n_p, const std::vector var_casadi_fcns, + const std::vector dvar_dy_fcns, + const std::vector dvar_dp_fcns, const Options& options ); @@ -91,6 +95,8 @@ class CasadiFunctions CasadiFunction mass_action; CasadiFunction events; std::vector var_casadi_fcns; + std::vector dvar_dy_fcns; + std::vector dvar_dp_fcns; std::vector jac_times_cjmass_rowvals; std::vector jac_times_cjmass_colptrs; diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp index 82fd9561a2..7040d1398f 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -26,6 +26,8 @@ CasadiSolver *create_casadi_solver( double rel_tol, int inputs_length, const std::vector var_casadi_fcns, + const std::vector dvar_dy_fcns, + const std::vector dvar_dp_fcns, py::dict options ) { auto options_cpp = Options(options); @@ -46,6 +48,8 @@ CasadiSolver *create_casadi_solver( number_of_events, number_of_parameters, var_casadi_fcns, + dvar_dy_fcns, + dvar_dp_fcns, options_cpp ); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index 54cf1e0ea2..023810d2a5 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -23,6 +23,8 @@ CasadiSolver *create_casadi_solver( double rel_tol, int inputs_length, const std::vector var_casadi_fcns, + const std::vector dvar_dy_fcns, + const std::vector dvar_dp_fcns, py::dict options ); diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index c462930565..b4fbb9da54 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -21,7 +21,7 @@ def have_idaklu(): return idaklu_spec is not None -def wrangle_name(name): +def wrangle_name(name: str) -> str: return name.casefold() \ .replace(" ", "_") \ .replace("[", "") \ @@ -285,16 +285,42 @@ def resfn(t, y, inputs, ydot): # if output_variables specified then convert functions to casadi # expressions for evaluation within the idaklu solver + self.var_casadi_fcns = [] + self.dvar_dy_fcns = [] + self.dvar_dp_fcns = [] for key in self.output_variables: - var_casadi = casadi.Function( - wrangle_name(key), + # variable functions + fcn_name = wrangle_name(key) + var_casadi = model.variables_and_events[key].to_casadi( + t_casadi, y_casadi, inputs=p_casadi) + var_casadi_fcn = casadi.Function( + fcn_name, [t_casadi, y_casadi, p_casadi_stacked], - [model.variables_and_events[key].to_casadi( - t_casadi, y_casadi, inputs=p_casadi)], + [var_casadi] ) self.var_casadi_fcns.append( - idaklu.generate_function(var_casadi.serialize()) + idaklu.generate_function(var_casadi_fcn.serialize()) ) + # Generate derivative functions for sensitivities + if len(inputs) > 0: + dvar_dy = casadi.jacobian(var_casadi, y_casadi) + dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) + dvar_dy_fcn = casadi.Function( + f"d{fcn_name}_dy", + [t_casadi, y_casadi, p_casadi_stacked], + [dvar_dy] + ) + dvar_dp_fcn = casadi.Function( + f"d{fcn_name}_dp", + [t_casadi, y_casadi, p_casadi_stacked], + [dvar_dp] + ) + self.dvar_dy_fcns.append( + idaklu.generate_function(dvar_dy_fcn.serialize()) + ) + self.dvar_dp_fcns.append( + idaklu.generate_function(dvar_dp_fcn.serialize()) + ) else: t0 = 0 if t_eval is None else t_eval[0] @@ -468,6 +494,8 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): "number_of_sensitivity_parameters": number_of_sensitivity_parameters, "output_variables": self.output_variables, "var_casadi_fcns": self.var_casadi_fcns, + "dvar_dy_fcns": self.dvar_dy_fcns, + "dvar_dp_fcns": self.dvar_dp_fcns, } solver = idaklu.create_casadi_solver( @@ -490,6 +518,8 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): rtol=rtol, inputs=len(inputs), var_casadi_fcns=self._setup["var_casadi_fcns"], + dvar_dy_fcns=self._setup["dvar_dy_fcns"], + dvar_dp_fcns=self._setup["dvar_dp_fcns"], options=self._options, ) @@ -604,7 +634,11 @@ def _integrate(self, model, t_eval, inputs_dict=None): t = sol.t number_of_timesteps = t.size number_of_states = y0.size - y_out = sol.y.reshape((number_of_timesteps, number_of_states)) + if self.output_variables: + # Substitute empty vectors for state vector 'y' + y_out = np.zeros((number_of_timesteps*number_of_states, 0)) + else: + y_out = sol.y.reshape((number_of_timesteps, number_of_states)) # return sensitivity solution, we need to flatten yS to # (#timesteps * #states (where t is changing the quickest),) @@ -628,7 +662,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): elif sol.flag == 2: termination = "event" - sol = pybamm.Solution( + newsol = pybamm.Solution( sol.t, np.transpose(y_out), model, @@ -638,7 +672,18 @@ def _integrate(self, model, t_eval, inputs_dict=None): termination, sensitivities=yS_out, ) - sol.integration_time = integration_time - return sol + newsol.integration_time = integration_time + if self.output_variables: + # Populate variables and sensititivies dictionaries directly + number_of_samples = sol.y.shape[0] // number_of_timesteps + sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) + newsol._variables = {} + newsol.yS = sol.yS # TODO: Replace + startk = 0 + for name in self.output_variables: + len_of_var = 1 + newsol._variables[name] = sol.y[:, startk:(startk+len_of_var)] + startk += len_of_var + return newsol else: raise pybamm.SolverError("idaklu solver failed") diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py index 75b58e0950..90884019a4 100644 --- a/pybamm/solvers/processed_variable.py +++ b/pybamm/solvers/processed_variable.py @@ -466,9 +466,12 @@ def initialise_sensitivity_explicit_forward(self): dvar_dy_eval = casadi.diagcat(dvar_dy_eval, next_dvar_dy_eval) dvar_dp_eval = casadi.vertcat(dvar_dp_eval, next_dvar_dp_eval) + print(dvar_dy_eval) + print(dvar_dp_eval) # Compute sensitivity dy_dp = self.solution_sensitivities["all"] S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval + print(S_var) sensitivities = {"all": S_var} diff --git a/test.py b/test.py index ea629b5504..c3c9663689 100644 --- a/test.py +++ b/test.py @@ -12,9 +12,11 @@ "Current [A]", ] #output_variables = [] +all_vars = False input_parameters = { "Current function [A]": 0.15652, + "Separator porosity": 0.47, } # check for loading errors @@ -62,7 +64,7 @@ 'num_threads': num_threads, } -if True: +if all_vars: output_variables = [m for m, (k, v) in zip(model.variable_names(), model.variables.items()) if not isinstance(v, pybamm.ExplicitTimeIntegral)] @@ -92,17 +94,21 @@ print(f"Solve time: {sol.solve_time.value*1000} msecs") -if output_variables: - var = output_variables[0] - print(sol[var].data) -else: - var = "Current [A]" -if input_parameters: - param = list(input_parameters.keys())[0] - print(sol[var].sensitivities[param]) - if False: - fig, axs = plt.subplots(1, 2) - axs[0].plot(t_eval, sol[var](t_eval)) - axs[1].plot(t_eval, sol[var].sensitivities[param]) - plt.tight_layout() - plt.show() +if True: + #output_variables = [ + # "Voltage [V]", + # "Time [min]", + # "Current [A]", + #] + fig, axs = plt.subplots(len(output_variables), len(input_parameters)+1) + for k, var in enumerate(output_variables): + if False: + axs[k,0].plot(t_eval, sol[var](t_eval)) + for paramk, param in enumerate(list(input_parameters.keys())): + axs[k,paramk+1].plot(t_eval, sol[var].sensitivities[param]) # time, param, var + else: + axs[k,0].plot(t_eval, sol[var][:,0]) + for paramk, param in enumerate(list(input_parameters.keys())): + axs[k,paramk+1].plot(t_eval, sol.yS[:,k,paramk]) # time, param, var + plt.tight_layout() + plt.show() From c8ad3c466026726797eb07e3025c3721dd482807 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 20 Jul 2023 11:02:07 +0000 Subject: [PATCH 08/38] Provide new ProcessedVariable class when output_variables are specified --- pybamm/__init__.py | 6 +- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 15 - pybamm/solvers/idaklu_solver.py | 17 +- pybamm/solvers/processed_variable.py | 26 +- pybamm/solvers/solution.py | 924 +++++++++++++++++- pybamm/solvers/solution_base.py | 922 ----------------- pybamm/solvers/solution_full.py | 8 - pybamm/solvers/solution_vars.py | 8 - test.py | 55 +- 9 files changed, 957 insertions(+), 1024 deletions(-) delete mode 100644 pybamm/solvers/solution_base.py delete mode 100644 pybamm/solvers/solution_full.py delete mode 100644 pybamm/solvers/solution_vars.py diff --git a/pybamm/__init__.py b/pybamm/__init__.py index cf70c01c7a..f20935a023 100644 --- a/pybamm/__init__.py +++ b/pybamm/__init__.py @@ -200,11 +200,9 @@ # # Solver classes # -from .solvers.solution_base import SolutionBase, EmptySolution, make_cycle_solution -from .solvers.solution_full import SolutionFull -from .solvers.solution_vars import SolutionVars -from .solvers.solution import Solution +from .solvers.solution import Solution, EmptySolution, make_cycle_solution from .solvers.processed_variable import ProcessedVariable +from .solvers.processed_variable_var import ProcessedVariableVar from .solvers.base_solver import BaseSolver from .solvers.dummy_solver import DummySolver from .solvers.algebraic_solver import AlgebraicSolver diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index f5d8db23ec..daf296e3bd 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -243,13 +243,6 @@ void CasadiSolverOpenMP::CalcVars( CalcVarsSensitivities(tret, yval, ySval, yS_return, ySk); } -void print(CasadiFunction dvar, realtype* res_dvar) { - std::cout << dvar.m_func.n_out() << "[" << dvar.m_res.size() << "] ("; - for (int k=0; kdvar_dy_fcns.size() << ", " << ySk << ", " - << (length_of_return_vector*number_of_timesteps*number_of_parameters) - << std::endl; - if (ySk != (length_of_return_vector*number_of_timesteps*number_of_parameters)) - // throw std::runtime_error("Sensitivities vector has become misaligned."); - std::cout << "WARNING: Sensitivities vector has become misaligned."; - np_array t_ret = np_array( t_i, &t_return[0], diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index b4fbb9da54..644ecf3537 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -677,13 +677,20 @@ def _integrate(self, model, t_eval, inputs_dict=None): # Populate variables and sensititivies dictionaries directly number_of_samples = sol.y.shape[0] // number_of_timesteps sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) - newsol._variables = {} - newsol.yS = sol.yS # TODO: Replace startk = 0 - for name in self.output_variables: - len_of_var = 1 - newsol._variables[name] = sol.y[:, startk:(startk+len_of_var)] + for vark, var in enumerate(self.output_variables): + len_of_var = int(model.variables_and_events[var].size) + newsol._variables[var] = pybamm.ProcessedVariableVar( + [model.variables_and_events[var]], + [sol.y[:, startk:(startk+len_of_var)]], + newsol, + ) startk += len_of_var + + # Add sensitivities + newsol[var]._sensitivities = {} + for paramk, param in enumerate(inputs_dict.keys()): + newsol[var]._sensitivities[param] = sol.yS[:, vark, paramk] return newsol else: raise pybamm.SolverError("idaklu solver failed") diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py index 90884019a4..3c920b659e 100644 --- a/pybamm/solvers/processed_variable.py +++ b/pybamm/solvers/processed_variable.py @@ -64,9 +64,8 @@ def __init__( self.t_pts = solution.t # Evaluate base variable at initial time - self.base_eval = self.base_variables_casadi[0]( - self.all_ts[0][0], self.all_ys[0][:, 0], self.all_inputs_casadi[0] - ).full() + self.base_eval_shape = self.base_variables[0].shape + self.base_eval_size = self.base_variables[0].size # handle 2D (in space) finite element variables differently if ( @@ -79,14 +78,13 @@ def __init__( # check variable shape else: if ( - isinstance(self.base_eval, numbers.Number) - or len(self.base_eval.shape) == 0 - or self.base_eval.shape[0] == 1 + len(self.base_eval_shape) == 0 + or self.base_eval_shape[0] == 1 ): self.initialise_0D() else: n = self.mesh.npts - base_shape = self.base_eval.shape[0] + base_shape = self.base_eval_shape[0] # Try some shapes that could make the variable a 1D variable if base_shape in [n, n + 1]: self.initialise_1D() @@ -95,7 +93,7 @@ def __init__( first_dim_nodes = self.mesh.nodes first_dim_edges = self.mesh.edges second_dim_pts = self.base_variables[0].secondary_mesh.nodes - if self.base_eval.size // len(second_dim_pts) in [ + if self.base_eval_size // len(second_dim_pts) in [ len(first_dim_nodes), len(first_dim_edges), ]: @@ -109,9 +107,6 @@ def __init__( def initialise_0D(self): # initialise empty array of the correct size - entries = np.empty(len(self.t_pts)) - idx = 0 - entries = np.empty(len(self.t_pts)) idx = 0 # Evaluate the base_variable index-by-index @@ -137,7 +132,7 @@ def initialise_0D(self): self.dimensions = 0 def initialise_1D(self, fixed_t=False): - len_space = self.base_eval.shape[0] + len_space = self.base_eval_shape[0] entries = np.empty((len_space, len(self.t_pts))) # Evaluate the base_variable index-by-index @@ -213,9 +208,9 @@ def initialise_2D(self): first_dim_edges = self.mesh.edges second_dim_nodes = self.base_variables[0].secondary_mesh.nodes second_dim_edges = self.base_variables[0].secondary_mesh.edges - if self.base_eval.size // len(second_dim_nodes) == len(first_dim_nodes): + if self.base_eval_size // len(second_dim_nodes) == len(first_dim_nodes): first_dim_pts = first_dim_nodes - elif self.base_eval.size // len(second_dim_nodes) == len(first_dim_edges): + elif self.base_eval_size // len(second_dim_nodes) == len(first_dim_edges): first_dim_pts = first_dim_edges second_dim_pts = second_dim_nodes @@ -466,12 +461,9 @@ def initialise_sensitivity_explicit_forward(self): dvar_dy_eval = casadi.diagcat(dvar_dy_eval, next_dvar_dy_eval) dvar_dp_eval = casadi.vertcat(dvar_dp_eval, next_dvar_dp_eval) - print(dvar_dy_eval) - print(dvar_dp_eval) # Compute sensitivity dy_dp = self.solution_sensitivities["all"] S_var = dvar_dy_eval @ dy_dp + dvar_dp_eval - print(S_var) sensitivities = {"all": S_var} diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py index a71c8231ce..f8faa67dca 100644 --- a/pybamm/solvers/solution.py +++ b/pybamm/solvers/solution.py @@ -1,14 +1,922 @@ # -# Solution interface class +# Solution class # +import casadi +import json +import numbers +import numpy as np +import pickle import pybamm +import pandas as pd +from scipy.io import savemat +from functools import cached_property -class Solution(pybamm.SolutionBase): - def __new__(cls, *args, **kwargs): - if kwargs.get("output_variables", []): - # Solution contains the full state vector 'y' - return pybamm.SolutionVars(*args, **kwargs) +class NumpyEncoder(json.JSONEncoder): + """ + Numpy serialiser helper class that converts numpy arrays to a list + https://stackoverflow.com/questions/26646362/numpy-array-is-not-json-serializable + """ + + def default(self, obj): + if isinstance(obj, np.ndarray): + return obj.tolist() + # won't be called since we only need to convert numpy arrays + return json.JSONEncoder.default(self, obj) # pragma: no cover + + +class Solution(object): + """ + Class containing the solution of, and various attributes associated with, a PyBaMM + model. + + Parameters + ---------- + all_ts : :class:`numpy.array`, size (n,) (or list of these) + A one-dimensional array containing the times at which the solution is evaluated. + A list of times can be provided instead to initialize a solution with + sub-solutions. + all_ys : :class:`numpy.array`, size (m, n) (or list of these) + A two-dimensional array containing the values of the solution. y[i, :] is the + vector of solutions at time t[i]. + A list of ys can be provided instead to initialize a solution with + sub-solutions. + all_models : :class:`pybamm.BaseModel` + The model that was used to calculate the solution. + A list of models can be provided instead to initialize a solution with + sub-solutions that have been calculated using those models. + all_inputs : dict (or list of these) + The inputs that were used to calculate the solution + A list of inputs can be provided instead to initialize a solution with + sub-solutions. + t_event : :class:`numpy.array`, size (1,) + A zero-dimensional array containing the time at which the event happens. + y_event : :class:`numpy.array`, size (m,) + A one-dimensional array containing the value of the solution at the time when + the event happens. + termination : str + String to indicate why the solution terminated + + sensitivities: bool or dict + True if sensitivities included as the solution of the explicit forwards + equations. False if no sensitivities included/wanted. Dict if sensitivities are + provided as a dict of {parameter: sensitivities} pairs. + + """ + + def __init__( + self, + all_ts, + all_ys, + all_models, + all_inputs, + t_event=None, + y_event=None, + termination="final time", + sensitivities=False, + check_solution=True, + ): + if not isinstance(all_ts, list): + all_ts = [all_ts] + if not isinstance(all_ys, list): + all_ys = [all_ys] + if not isinstance(all_models, list): + all_models = [all_models] + self._all_ts = all_ts + self._all_ys = all_ys + self._all_ys_and_sens = all_ys + self._all_models = all_models + + # Set up inputs + if not isinstance(all_inputs, list): + all_inputs_copy = dict(all_inputs) + for key, value in all_inputs_copy.items(): + if isinstance(value, numbers.Number): + all_inputs_copy[key] = np.array([value]) + self.all_inputs = [all_inputs_copy] + else: + self.all_inputs = all_inputs + + self.sensitivities = sensitivities + + self._t_event = t_event + self._y_event = y_event + self._termination = termination + + # Check no ys are too large + if check_solution: + self.check_ys_are_not_too_large() + + # Events + self._t_event = t_event + self._y_event = y_event + self._termination = termination + self.closest_event_idx = None + + # Initialize times + self.set_up_time = None + self.solve_time = None + self.integration_time = None + + # initialize empty variables and data + self._variables = pybamm.FuzzyDict() + self.data = pybamm.FuzzyDict() + + # Add self as sub-solution for compatibility with ProcessedVariable + self._sub_solutions = [self] + + # initialize empty cycles + self._cycles = [] + + # Initialize empty summary variables + self._summary_variables = None + + # Solution now uses CasADi + pybamm.citations.register("Andersson2019") + + def extract_explicit_sensitivities(self): + # if we got here, we haven't set y yet + self.set_y() + + # extract sensitivities from full y solution + self._y, self._sensitivities = self._extract_explicit_sensitivities( + self.all_models[0], self.y, self.t, self.all_inputs[0] + ) + + # make sure we remove all sensitivities from all_ys + for index, (model, ys, ts, inputs) in enumerate( + zip(self.all_models, self.all_ys, self.all_ts, self.all_inputs) + ): + self._all_ys[index], _ = self._extract_explicit_sensitivities( + model, ys, ts, inputs + ) + + def _extract_explicit_sensitivities(self, model, y, t_eval, inputs): + """ + given a model and a solution y, extracts the sensitivities + + Parameters + -------- + model : :class:`pybamm.BaseModel` + A model that has been already setup by this base solver + y: ndarray + The solution of the full explicit sensitivity equations + t_eval: ndarray + The evaluation times + inputs: dict + parameter inputs + + Returns + ------- + y: ndarray + The solution of the ode/dae in model + sensitivities: dict of (string: ndarray) + A dictionary of parameter names, and the corresponding solution of + the sensitivity equations + """ + + n_states = model.len_rhs_and_alg + n_rhs = model.len_rhs + n_alg = model.len_alg + # Get the point where the algebraic equations start + if model.len_rhs != 0: + n_p = model.len_rhs_sens // model.len_rhs else: - # Solution contains only the requested variables - return pybamm.SolutionFull(*args, **kwargs) + n_p = model.len_alg_sens // model.len_alg + len_rhs_and_sens = model.len_rhs + model.len_rhs_sens + + n_t = len(t_eval) + # y gets the part of the solution vector that correspond to the + # actual ODE/DAE solution + + # save sensitivities as a dictionary + # first save the whole sensitivity matrix + # reshape using Fortran order to get the right array: + # t0_x0_p0, t0_x0_p1, ..., t0_x0_pn + # t0_x1_p0, t0_x1_p1, ..., t0_x1_pn + # ... + # t0_xn_p0, t0_xn_p1, ..., t0_xn_pn + # t1_x0_p0, t1_x0_p1, ..., t1_x0_pn + # t1_x1_p0, t1_x1_p1, ..., t1_x1_pn + # ... + # t1_xn_p0, t1_xn_p1, ..., t1_xn_pn + # ... + # tn_x0_p0, tn_x0_p1, ..., tn_x0_pn + # tn_x1_p0, tn_x1_p1, ..., tn_x1_pn + # ... + # tn_xn_p0, tn_xn_p1, ..., tn_xn_pn + # 1, Extract rhs and alg sensitivities and reshape into 3D matrices + # with shape (n_p, n_states, n_t) + if isinstance(y, casadi.DM): + y_full = y.full() + else: + y_full = y + ode_sens = y_full[n_rhs:len_rhs_and_sens, :].reshape(n_p, n_rhs, n_t) + alg_sens = y_full[len_rhs_and_sens + n_alg :, :].reshape(n_p, n_alg, n_t) + # 2. Concatenate into a single 3D matrix with shape (n_p, n_states, n_t) + # i.e. along first axis + full_sens_matrix = np.concatenate([ode_sens, alg_sens], axis=1) + # Transpose and reshape into a (n_states * n_t, n_p) matrix + full_sens_matrix = full_sens_matrix.transpose(2, 1, 0).reshape( + n_t * n_states, n_p + ) + + # Save the full sensitivity matrix + sensitivity = {"all": full_sens_matrix} + + # also save the sensitivity wrt each parameter (read the columns of the + # sensitivity matrix) + start = 0 + for name in model.calculate_sensitivities: + inp = inputs[name] + input_size = inp.shape[0] + end = start + input_size + sensitivity[name] = full_sens_matrix[:, start:end] + start = end + + y_dae = np.vstack( + [ + y[: model.len_rhs, :], + y[len_rhs_and_sens : len_rhs_and_sens + model.len_alg, :], + ] + ) + return y_dae, sensitivity + + @property + def t(self): + """Times at which the solution is evaluated""" + try: + return self._t + except AttributeError: + self.set_t() + return self._t + + def set_t(self): + self._t = np.concatenate(self.all_ts) + if any(np.diff(self._t) <= 0): + raise ValueError("Solution time vector must be strictly increasing") + + @property + def y(self): + """Values of the solution""" + try: + return self._y + except AttributeError: + self.set_y() + + # if y is evaluated before sensitivities then need to extract them + if isinstance(self._sensitivities, bool) and self._sensitivities: + self.extract_explicit_sensitivities() + + return self._y + + @property + def sensitivities(self): + """Values of the sensitivities. Returns a dict of param_name: np_array""" + if isinstance(self._sensitivities, bool): + if self._sensitivities: + self.extract_explicit_sensitivities() + else: + self._sensitivities = {} + return self._sensitivities + + @sensitivities.setter + def sensitivities(self, value): + """Updates the sensitivity""" + # sensitivities must be a dict or bool + if not isinstance(value, (bool, dict)): + raise TypeError("sensitivities arg needs to be a bool or dict") + self._sensitivities = value + + def set_y(self): + try: + if isinstance(self.all_ys[0], (casadi.DM, casadi.MX)): + self._y = casadi.horzcat(*self.all_ys) + else: + self._y = np.hstack(self.all_ys) + except ValueError: + raise pybamm.SolverError( + "The solution is made up from different models, so `y` cannot be " + "computed explicitly." + ) + + def check_ys_are_not_too_large(self): + # Only check last one so that it doesn't take too long + # We only care about the cases where y is growing too large without any + # restraint, so if y gets large in the middle then comes back down that is ok + y, model = self.all_ys[-1], self.all_models[-1] + y = y[:, -1] + if np.any(y > pybamm.settings.max_y_value): + for var in [*model.rhs.keys(), *model.algebraic.keys()]: + y_var = y[model.variables[var.name].y_slices[0]] + if np.any(y_var > pybamm.settings.max_y_value): + pybamm.logger.error( + f"Solution for '{var}' exceeds the maximum allowed value " + f"of `{pybamm.settings.max_y_value}. This could be due to " + "incorrect scaling, model formulation, or " + "parameter values. The maximum allowed value is set by " + "'pybammm.settings.max_y_value'." + ) + + @property + def all_ts(self): + return self._all_ts + + @property + def all_ys(self): + return self._all_ys + + @property + def all_models(self): + """Model(s) used for solution""" + return self._all_models + + @cached_property + def all_inputs_casadi(self): + return [casadi.vertcat(*inp.values()) for inp in self.all_inputs] + + @property + def t_event(self): + """Time at which the event happens""" + return self._t_event + + @property + def y_event(self): + """Value of the solution at the time of the event""" + return self._y_event + + @property + def termination(self): + """Reason for termination""" + return self._termination + + @termination.setter + def termination(self, value): + """Updates the reason for termination""" + self._termination = value + + @cached_property + def first_state(self): + """ + A Solution object that only contains the first state. This is faster to evaluate + than the full solution when only the first state is needed (e.g. to initialize + a model with the solution) + """ + new_sol = Solution( + self.all_ts[0][:1], + self.all_ys[0][:, :1], + self.all_models[:1], + self.all_inputs[:1], + None, + None, + "final time", + ) + new_sol._all_inputs_casadi = self.all_inputs_casadi[:1] + new_sol._sub_solutions = self.sub_solutions[:1] + + new_sol.solve_time = 0 + new_sol.integration_time = 0 + new_sol.set_up_time = 0 + + return new_sol + + @cached_property + def last_state(self): + """ + A Solution object that only contains the final state. This is faster to evaluate + than the full solution when only the final state is needed (e.g. to initialize + a model with the solution) + """ + new_sol = Solution( + self.all_ts[-1][-1:], + self.all_ys[-1][:, -1:], + self.all_models[-1:], + self.all_inputs[-1:], + self.t_event, + self.y_event, + self.termination, + ) + new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:] + new_sol._sub_solutions = self.sub_solutions[-1:] + + new_sol.solve_time = 0 + new_sol.integration_time = 0 + new_sol.set_up_time = 0 + + return new_sol + + @property + def total_time(self): + return self.set_up_time + self.solve_time + + @property + def cycles(self): + return self._cycles + + @cycles.setter + def cycles(self, cycles): + self._cycles = cycles + + @property + def summary_variables(self): + return self._summary_variables + + def set_summary_variables(self, all_summary_variables): + summary_variables = {var: [] for var in all_summary_variables[0]} + for sum_vars in all_summary_variables: + for name, value in sum_vars.items(): + summary_variables[name].append(value) + + summary_variables["Cycle number"] = range(1, len(all_summary_variables) + 1) + self.all_summary_variables = all_summary_variables + self._summary_variables = pybamm.FuzzyDict( + {name: np.array(value) for name, value in summary_variables.items()} + ) + + def update(self, variables): + """Add ProcessedVariables to the dictionary of variables in the solution""" + # make sure that sensitivities are extracted if required + if isinstance(self._sensitivities, bool) and self._sensitivities: + self.extract_explicit_sensitivities() + + # Convert single entry to list + if isinstance(variables, str): + variables = [variables] + # Process + for key in variables: + cumtrapz_ic = None + pybamm.logger.debug("Post-processing {}".format(key)) + vars_pybamm = [model.variables_and_events[key] for model in self.all_models] + + # Iterate through all models, some may be in the list several times and + # therefore only get set up once + vars_casadi = [] + for i, (model, ys, inputs, var_pybamm) in enumerate( + zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm) + ): + if isinstance(var_pybamm, pybamm.ExplicitTimeIntegral): + cumtrapz_ic = var_pybamm.initial_condition + cumtrapz_ic = cumtrapz_ic.evaluate() + var_pybamm = var_pybamm.child + var_casadi = self.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) + model._variables_casadi[key] = var_casadi + vars_pybamm[i] = var_pybamm + elif key in model._variables_casadi: + var_casadi = model._variables_casadi[key] + else: + var_casadi = self.process_casadi_var( + var_pybamm, + inputs, + ys.shape, + ) + model._variables_casadi[key] = var_casadi + vars_casadi.append(var_casadi) + var = pybamm.ProcessedVariable( + vars_pybamm, vars_casadi, self, cumtrapz_ic=cumtrapz_ic + ) + + # Save variable and data + self._variables[key] = var + self.data[key] = var.data + + def process_casadi_var(self, var_pybamm, inputs, ys_shape): + t_MX = casadi.MX.sym("t") + y_MX = casadi.MX.sym("y", ys_shape[0]) + inputs_MX_dict = { + key: casadi.MX.sym("input", value.shape[0]) for key, value in inputs.items() + } + inputs_MX = casadi.vertcat(*[p for p in inputs_MX_dict.values()]) + var_sym = var_pybamm.to_casadi(t_MX, y_MX, inputs=inputs_MX_dict) + var_casadi = casadi.Function("variable", [t_MX, y_MX, inputs_MX], [var_sym]) + return var_casadi + + def __getitem__(self, key): + """Read a variable from the solution. Variables are created 'just in time', i.e. + only when they are called. + + Parameters + ---------- + key : str + The name of the variable + + Returns + ------- + :class:`pybamm.ProcessedVariable` + A variable that can be evaluated at any time or spatial point. The + underlying data for this variable is available in its attribute ".data" + """ + + # return it if it exists + if key in self._variables: + return self._variables[key] + else: + # otherwise create it, save it and then return it + self.update(key) + return self._variables[key] + + def plot(self, output_variables=None, **kwargs): + """ + A method to quickly plot the outputs of the solution. Creates a + :class:`pybamm.QuickPlot` object (with keyword arguments 'kwargs') and + then calls :meth:`pybamm.QuickPlot.dynamic_plot`. + + Parameters + ---------- + output_variables: list, optional + A list of the variables to plot. + **kwargs + Additional keyword arguments passed to + :meth:`pybamm.QuickPlot.dynamic_plot`. + For a list of all possible keyword arguments see :class:`pybamm.QuickPlot`. + """ + return pybamm.dynamic_plot(self, output_variables=output_variables, **kwargs) + + def save(self, filename): + """Save the whole solution using pickle""" + # No warning here if len(self.data)==0 as solution can be loaded + # and used to process new variables + + with open(filename, "wb") as f: + pickle.dump(self, f, pickle.HIGHEST_PROTOCOL) + + def get_data_dict(self, variables=None, short_names=None, cycles_and_steps=True): + """ + Construct a (standard python) dictionary of the solution data containing the + variables in `variables`. If `variables` is None then all variables are + returned. Any variable names in short_names are replaced with the corresponding + short name. + + If the solution has cycles, then the cycle numbers and step numbers are also + returned in the dictionary. + + Parameters + ---------- + variables : list, optional + List of variables to return. If None, returns all variables in solution.data + short_names : dict, optional + Dictionary of shortened names to use when saving. + cycles_and_steps : bool, optional + Whether to include the cycle numbers and step numbers in the dictionary + + Returns + ------- + dict + A dictionary of the solution data + """ + if variables is None: + # variables not explicitly provided -> save all variables that have been + # computed + data_long_names = self.data + else: + if isinstance(variables, str): + variables = [variables] + # otherwise, save only the variables specified + data_long_names = {} + for name in variables: + data_long_names[name] = self[name].data + if len(data_long_names) == 0: + raise ValueError( + """ + Solution does not have any data. Please provide a list of variables + to save. + """ + ) + + # Use any short names if provided + data_short_names = {} + short_names = short_names or {} + for name, var in data_long_names.items(): + name = short_names.get(name, name) # return name if no short name + data_short_names[name] = var + + # Save cycle number and step number if the solution has them + if cycles_and_steps and len(self.cycles) > 0: + data_short_names["Cycle"] = np.array([]) + data_short_names["Step"] = np.array([]) + for i, cycle in enumerate(self.cycles): + data_short_names["Cycle"] = np.concatenate( + [data_short_names["Cycle"], i * np.ones_like(cycle.t)] + ) + for j, step in enumerate(cycle.steps): + data_short_names["Step"] = np.concatenate( + [data_short_names["Step"], j * np.ones_like(step.t)] + ) + + return data_short_names + + def save_data( + self, filename=None, variables=None, to_format="pickle", short_names=None + ): + """ + Save solution data only (raw arrays) + + Parameters + ---------- + filename : str, optional + The name of the file to save data to. If None, then a str is returned + variables : list, optional + List of variables to save. If None, saves all of the variables that have + been created so far + to_format : str, optional + The format to save to. Options are: + + - 'pickle' (default): creates a pickle file with the data dictionary + - 'matlab': creates a .mat file, for loading in matlab + - 'csv': creates a csv file (0D variables only) + - 'json': creates a json file + short_names : dict, optional + Dictionary of shortened names to use when saving. This may be necessary when + saving to MATLAB, since no spaces or special characters are allowed in + MATLAB variable names. Note that not all the variables need to be given + a short name. + + Returns + ------- + data : str, optional + str if 'csv' or 'json' is chosen and filename is None, otherwise None + """ + data = self.get_data_dict(variables=variables, short_names=short_names) + + if to_format == "pickle": + if filename is None: + raise ValueError("pickle format must be written to a file") + with open(filename, "wb") as f: + pickle.dump(data, f, pickle.HIGHEST_PROTOCOL) + elif to_format == "matlab": + if filename is None: + raise ValueError("matlab format must be written to a file") + # Check all the variable names only contain a-z, A-Z or _ or numbers + for name in data.keys(): + # Check the string only contains the following ASCII: + # a-z (97-122) + # A-Z (65-90) + # _ (95) + # 0-9 (48-57) but not in the first position + for i, s in enumerate(name): + if not ( + 97 <= ord(s) <= 122 + or 65 <= ord(s) <= 90 + or ord(s) == 95 + or (i > 0 and 48 <= ord(s) <= 57) + ): + raise ValueError( + "Invalid character '{}' found in '{}'. ".format(s, name) + + "MATLAB variable names must only contain a-z, A-Z, _, " + "or 0-9 (except the first position). " + "Use the 'short_names' argument to pass an alternative " + "variable name, e.g. \n\n" + "\tsolution.save_data(filename, " + "['Electrolyte concentration'], to_format='matlab, " + "short_names={'Electrolyte concentration': 'c_e'})" + ) + savemat(filename, data) + elif to_format == "csv": + for name, var in data.items(): + if var.ndim >= 2: + raise ValueError( + "only 0D variables can be saved to csv, but '{}' is {}D".format( + name, var.ndim - 1 + ) + ) + df = pd.DataFrame(data) + return df.to_csv(filename, index=False) + elif to_format == "json": + if filename is None: + return json.dumps(data, cls=NumpyEncoder) + else: + with open(filename, "w") as outfile: + json.dump(data, outfile, cls=NumpyEncoder) + else: + raise ValueError("format '{}' not recognised".format(to_format)) + + @property + def sub_solutions(self): + """List of sub solutions that have been + concatenated to form the full solution""" + + return self._sub_solutions + + def __add__(self, other): + """Adds two solutions together, e.g. when stepping""" + if other is None or isinstance(other, EmptySolution): + return self.copy() + if not isinstance(other, Solution): + raise pybamm.SolverError( + "Only a Solution or None can be added to a Solution" + ) + # Special case: new solution only has one timestep and it is already in the + # existing solution. In this case, return a copy of the existing solution + if ( + len(other.all_ts) == 1 + and len(other.all_ts[0]) == 1 + and other.all_ts[0][0] == self.all_ts[-1][-1] + ): + new_sol = self.copy() + # Update termination using the latter solution + new_sol._termination = other.termination + new_sol._t_event = other._t_event + new_sol._y_event = other._y_event + return new_sol + + # Update list of sub-solutions + if other.all_ts[0][0] == self.all_ts[-1][-1]: + # Skip first time step if it is repeated + all_ts = self.all_ts + [other.all_ts[0][1:]] + other.all_ts[1:] + all_ys = self.all_ys + [other.all_ys[0][:, 1:]] + other.all_ys[1:] + else: + all_ts = self.all_ts + other.all_ts + all_ys = self.all_ys + other.all_ys + + new_sol = Solution( + all_ts, + all_ys, + self.all_models + other.all_models, + self.all_inputs + other.all_inputs, + other.t_event, + other.y_event, + other.termination, + bool(self.sensitivities), + ) + + new_sol.closest_event_idx = other.closest_event_idx + new_sol._all_inputs_casadi = self.all_inputs_casadi + other.all_inputs_casadi + + # Set solution time + new_sol.solve_time = self.solve_time + other.solve_time + new_sol.integration_time = self.integration_time + other.integration_time + + # Set sub_solutions + new_sol._sub_solutions = self.sub_solutions + other.sub_solutions + + return new_sol + + def __radd__(self, other): + return self.__add__(other) + + def copy(self): + new_sol = self.__class__( + self.all_ts, + self.all_ys, + self.all_models, + self.all_inputs, + self.t_event, + self.y_event, + self.termination, + ) + new_sol._all_inputs_casadi = self.all_inputs_casadi + new_sol._sub_solutions = self.sub_solutions + new_sol.closest_event_idx = self.closest_event_idx + + new_sol.solve_time = self.solve_time + new_sol.integration_time = self.integration_time + new_sol.set_up_time = self.set_up_time + + return new_sol + + +class EmptySolution: + def __init__(self, termination=None, t=None): + self.termination = termination + if t is None: + t = np.array([0]) + elif isinstance(t, numbers.Number): + t = np.array([t]) + + self.t = t + + def __add__(self, other): + if isinstance(other, (EmptySolution, Solution)): + return other.copy() + + def __radd__(self, other): + if other is None: + return self.copy() + + def copy(self): + return EmptySolution(termination=self.termination, t=self.t) + + +def make_cycle_solution(step_solutions, esoh_solver=None, save_this_cycle=True): + """ + Function to create a Solution for an entire cycle, and associated summary variables + + Parameters + ---------- + step_solutions : list of :class:`Solution` + Step solutions that form the entire cycle + esoh_solver : :class:`pybamm.lithium_ion.ElectrodeSOHSolver` + Solver to calculate electrode SOH (eSOH) variables. If `None` (default) + then only summary variables that do not require the eSOH calculation + are calculated. See :footcite:t:`Mohtat2019` for more details on eSOH variables. + save_this_cycle : bool, optional + Whether to save the entire cycle variables or just the summary variables. + Default True + + Returns + ------- + cycle_solution : :class:`pybamm.Solution` or None + The Solution object for this cycle, or None (if save_this_cycle is False) + cycle_summary_variables : dict + Dictionary of summary variables for this cycle + + """ + sum_sols = step_solutions[0].copy() + for step_solution in step_solutions[1:]: + sum_sols = sum_sols + step_solution + + cycle_solution = Solution( + sum_sols.all_ts, + sum_sols.all_ys, + sum_sols.all_models, + sum_sols.all_inputs, + sum_sols.t_event, + sum_sols.y_event, + sum_sols.termination, + ) + cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi + cycle_solution._sub_solutions = sum_sols.sub_solutions + + cycle_solution.solve_time = sum_sols.solve_time + cycle_solution.integration_time = sum_sols.integration_time + cycle_solution.set_up_time = sum_sols.set_up_time + + cycle_solution.steps = step_solutions + + cycle_summary_variables = _get_cycle_summary_variables(cycle_solution, esoh_solver) + + cycle_first_state = cycle_solution.first_state + + if save_this_cycle: + cycle_solution.cycle_summary_variables = cycle_summary_variables + else: + cycle_solution = None + + return cycle_solution, cycle_summary_variables, cycle_first_state + + +def _get_cycle_summary_variables(cycle_solution, esoh_solver): + model = cycle_solution.all_models[0] + cycle_summary_variables = pybamm.FuzzyDict({}) + + # Measured capacity variables + if "Discharge capacity [A.h]" in model.variables: + Q = cycle_solution["Discharge capacity [A.h]"].data + min_Q, max_Q = np.min(Q), np.max(Q) + + cycle_summary_variables.update( + { + "Minimum measured discharge capacity [A.h]": min_Q, + "Maximum measured discharge capacity [A.h]": max_Q, + "Measured capacity [A.h]": max_Q - min_Q, + } + ) + + # Voltage variables + if "Battery voltage [V]" in model.variables: + V = cycle_solution["Battery voltage [V]"].data + min_V, max_V = np.min(V), np.max(V) + + cycle_summary_variables.update( + {"Minimum voltage [V]": min_V, "Maximum voltage [V]": max_V} + ) + + # Degradation variables + degradation_variables = model.summary_variables + first_state = cycle_solution.first_state + last_state = cycle_solution.last_state + for var in degradation_variables: + data_first = first_state[var].data + data_last = last_state[var].data + cycle_summary_variables[var] = data_last[0] + var_lowercase = var[0].lower() + var[1:] + cycle_summary_variables["Change in " + var_lowercase] = ( + data_last[0] - data_first[0] + ) + + # eSOH variables (full-cell lithium-ion model only, for now) + if ( + esoh_solver is not None + and isinstance(model, pybamm.lithium_ion.BaseModel) + and model.options.electrode_types["negative"] == "porous" + ): + Q_n = last_state["Negative electrode capacity [A.h]"].data[0] + Q_p = last_state["Positive electrode capacity [A.h]"].data[0] + Q_Li = last_state["Total lithium capacity in particles [A.h]"].data[0] + + inputs = {"Q_n": Q_n, "Q_p": Q_p, "Q_Li": Q_Li} + + try: + esoh_sol = esoh_solver.solve(inputs) + except pybamm.SolverError: # pragma: no cover + raise pybamm.SolverError( + "Could not solve for summary variables, run " + "`sim.solve(calc_esoh=False)` to skip this step" + ) + + cycle_summary_variables.update(esoh_sol) + + return cycle_summary_variables diff --git a/pybamm/solvers/solution_base.py b/pybamm/solvers/solution_base.py deleted file mode 100644 index 3cdde7dc63..0000000000 --- a/pybamm/solvers/solution_base.py +++ /dev/null @@ -1,922 +0,0 @@ -# -# SolutionBase class -# -import casadi -import json -import numbers -import numpy as np -import pickle -import pybamm -import pandas as pd -from scipy.io import savemat -from functools import cached_property - - -class NumpyEncoder(json.JSONEncoder): - """ - Numpy serialiser helper class that converts numpy arrays to a list - https://stackoverflow.com/questions/26646362/numpy-array-is-not-json-serializable - """ - - def default(self, obj): - if isinstance(obj, np.ndarray): - return obj.tolist() - # won't be called since we only need to convert numpy arrays - return json.JSONEncoder.default(self, obj) # pragma: no cover - - -class SolutionBase(object): - """ - Class containing the solution of, and various attributes associated with, a PyBaMM - model. - - Parameters - ---------- - all_ts : :class:`numpy.array`, size (n,) (or list of these) - A one-dimensional array containing the times at which the solution is evaluated. - A list of times can be provided instead to initialize a solution with - sub-solutions. - all_ys : :class:`numpy.array`, size (m, n) (or list of these) - A two-dimensional array containing the values of the solution. y[i, :] is the - vector of solutions at time t[i]. - A list of ys can be provided instead to initialize a solution with - sub-solutions. - all_models : :class:`pybamm.BaseModel` - The model that was used to calculate the solution. - A list of models can be provided instead to initialize a solution with - sub-solutions that have been calculated using those models. - all_inputs : dict (or list of these) - The inputs that were used to calculate the solution - A list of inputs can be provided instead to initialize a solution with - sub-solutions. - t_event : :class:`numpy.array`, size (1,) - A zero-dimensional array containing the time at which the event happens. - y_event : :class:`numpy.array`, size (m,) - A one-dimensional array containing the value of the solution at the time when - the event happens. - termination : str - String to indicate why the solution terminated - - sensitivities: bool or dict - True if sensitivities included as the solution of the explicit forwards - equations. False if no sensitivities included/wanted. Dict if sensitivities are - provided as a dict of {parameter: sensitivities} pairs. - - """ - - def __init__( - self, - all_ts, - all_ys, - all_models, - all_inputs, - t_event=None, - y_event=None, - termination="final time", - sensitivities=False, - check_solution=True, - ): - if not isinstance(all_ts, list): - all_ts = [all_ts] - if not isinstance(all_ys, list): - all_ys = [all_ys] - if not isinstance(all_models, list): - all_models = [all_models] - self._all_ts = all_ts - self._all_ys = all_ys - self._all_ys_and_sens = all_ys - self._all_models = all_models - - # Set up inputs - if not isinstance(all_inputs, list): - all_inputs_copy = dict(all_inputs) - for key, value in all_inputs_copy.items(): - if isinstance(value, numbers.Number): - all_inputs_copy[key] = np.array([value]) - self.all_inputs = [all_inputs_copy] - else: - self.all_inputs = all_inputs - - self.sensitivities = sensitivities - - self._t_event = t_event - self._y_event = y_event - self._termination = termination - - # Check no ys are too large - if check_solution: - self.check_ys_are_not_too_large() - - # Events - self._t_event = t_event - self._y_event = y_event - self._termination = termination - self.closest_event_idx = None - - # Initialize times - self.set_up_time = None - self.solve_time = None - self.integration_time = None - - # initialize empty variables and data - self._variables = pybamm.FuzzyDict() - self.data = pybamm.FuzzyDict() - - # Add self as sub-solution for compatibility with ProcessedVariable - self._sub_solutions = [self] - - # initialize empty cycles - self._cycles = [] - - # Initialize empty summary variables - self._summary_variables = None - - # Solution now uses CasADi - pybamm.citations.register("Andersson2019") - - def extract_explicit_sensitivities(self): - # if we got here, we haven't set y yet - self.set_y() - - # extract sensitivities from full y solution - self._y, self._sensitivities = self._extract_explicit_sensitivities( - self.all_models[0], self.y, self.t, self.all_inputs[0] - ) - - # make sure we remove all sensitivities from all_ys - for index, (model, ys, ts, inputs) in enumerate( - zip(self.all_models, self.all_ys, self.all_ts, self.all_inputs) - ): - self._all_ys[index], _ = self._extract_explicit_sensitivities( - model, ys, ts, inputs - ) - - def _extract_explicit_sensitivities(self, model, y, t_eval, inputs): - """ - given a model and a solution y, extracts the sensitivities - - Parameters - -------- - model : :class:`pybamm.BaseModel` - A model that has been already setup by this base solver - y: ndarray - The solution of the full explicit sensitivity equations - t_eval: ndarray - The evaluation times - inputs: dict - parameter inputs - - Returns - ------- - y: ndarray - The solution of the ode/dae in model - sensitivities: dict of (string: ndarray) - A dictionary of parameter names, and the corresponding solution of - the sensitivity equations - """ - - n_states = model.len_rhs_and_alg - n_rhs = model.len_rhs - n_alg = model.len_alg - # Get the point where the algebraic equations start - if model.len_rhs != 0: - n_p = model.len_rhs_sens // model.len_rhs - else: - n_p = model.len_alg_sens // model.len_alg - len_rhs_and_sens = model.len_rhs + model.len_rhs_sens - - n_t = len(t_eval) - # y gets the part of the solution vector that correspond to the - # actual ODE/DAE solution - - # save sensitivities as a dictionary - # first save the whole sensitivity matrix - # reshape using Fortran order to get the right array: - # t0_x0_p0, t0_x0_p1, ..., t0_x0_pn - # t0_x1_p0, t0_x1_p1, ..., t0_x1_pn - # ... - # t0_xn_p0, t0_xn_p1, ..., t0_xn_pn - # t1_x0_p0, t1_x0_p1, ..., t1_x0_pn - # t1_x1_p0, t1_x1_p1, ..., t1_x1_pn - # ... - # t1_xn_p0, t1_xn_p1, ..., t1_xn_pn - # ... - # tn_x0_p0, tn_x0_p1, ..., tn_x0_pn - # tn_x1_p0, tn_x1_p1, ..., tn_x1_pn - # ... - # tn_xn_p0, tn_xn_p1, ..., tn_xn_pn - # 1, Extract rhs and alg sensitivities and reshape into 3D matrices - # with shape (n_p, n_states, n_t) - if isinstance(y, casadi.DM): - y_full = y.full() - else: - y_full = y - ode_sens = y_full[n_rhs:len_rhs_and_sens, :].reshape(n_p, n_rhs, n_t) - alg_sens = y_full[len_rhs_and_sens + n_alg :, :].reshape(n_p, n_alg, n_t) - # 2. Concatenate into a single 3D matrix with shape (n_p, n_states, n_t) - # i.e. along first axis - full_sens_matrix = np.concatenate([ode_sens, alg_sens], axis=1) - # Transpose and reshape into a (n_states * n_t, n_p) matrix - full_sens_matrix = full_sens_matrix.transpose(2, 1, 0).reshape( - n_t * n_states, n_p - ) - - # Save the full sensitivity matrix - sensitivity = {"all": full_sens_matrix} - - # also save the sensitivity wrt each parameter (read the columns of the - # sensitivity matrix) - start = 0 - for name in model.calculate_sensitivities: - inp = inputs[name] - input_size = inp.shape[0] - end = start + input_size - sensitivity[name] = full_sens_matrix[:, start:end] - start = end - - y_dae = np.vstack( - [ - y[: model.len_rhs, :], - y[len_rhs_and_sens : len_rhs_and_sens + model.len_alg, :], - ] - ) - return y_dae, sensitivity - - @property - def t(self): - """Times at which the solution is evaluated""" - try: - return self._t - except AttributeError: - self.set_t() - return self._t - - def set_t(self): - self._t = np.concatenate(self.all_ts) - if any(np.diff(self._t) <= 0): - raise ValueError("Solution time vector must be strictly increasing") - - @property - def y(self): - """Values of the solution""" - try: - return self._y - except AttributeError: - self.set_y() - - # if y is evaluated before sensitivities then need to extract them - if isinstance(self._sensitivities, bool) and self._sensitivities: - self.extract_explicit_sensitivities() - - return self._y - - @property - def sensitivities(self): - """Values of the sensitivities. Returns a dict of param_name: np_array""" - if isinstance(self._sensitivities, bool): - if self._sensitivities: - self.extract_explicit_sensitivities() - else: - self._sensitivities = {} - return self._sensitivities - - @sensitivities.setter - def sensitivities(self, value): - """Updates the sensitivity""" - # sensitivities must be a dict or bool - if not isinstance(value, (bool, dict)): - raise TypeError("sensitivities arg needs to be a bool or dict") - self._sensitivities = value - - def set_y(self): - try: - if isinstance(self.all_ys[0], (casadi.DM, casadi.MX)): - self._y = casadi.horzcat(*self.all_ys) - else: - self._y = np.hstack(self.all_ys) - except ValueError: - raise pybamm.SolverError( - "The solution is made up from different models, so `y` cannot be " - "computed explicitly." - ) - - def check_ys_are_not_too_large(self): - # Only check last one so that it doesn't take too long - # We only care about the cases where y is growing too large without any - # restraint, so if y gets large in the middle then comes back down that is ok - y, model = self.all_ys[-1], self.all_models[-1] - y = y[:, -1] - if np.any(y > pybamm.settings.max_y_value): - for var in [*model.rhs.keys(), *model.algebraic.keys()]: - y_var = y[model.variables[var.name].y_slices[0]] - if np.any(y_var > pybamm.settings.max_y_value): - pybamm.logger.error( - f"Solution for '{var}' exceeds the maximum allowed value " - f"of `{pybamm.settings.max_y_value}. This could be due to " - "incorrect scaling, model formulation, or " - "parameter values. The maximum allowed value is set by " - "'pybammm.settings.max_y_value'." - ) - - @property - def all_ts(self): - return self._all_ts - - @property - def all_ys(self): - return self._all_ys - - @property - def all_models(self): - """Model(s) used for solution""" - return self._all_models - - @cached_property - def all_inputs_casadi(self): - return [casadi.vertcat(*inp.values()) for inp in self.all_inputs] - - @property - def t_event(self): - """Time at which the event happens""" - return self._t_event - - @property - def y_event(self): - """Value of the solution at the time of the event""" - return self._y_event - - @property - def termination(self): - """Reason for termination""" - return self._termination - - @termination.setter - def termination(self, value): - """Updates the reason for termination""" - self._termination = value - - @cached_property - def first_state(self): - """ - A Solution object that only contains the first state. This is faster to evaluate - than the full solution when only the first state is needed (e.g. to initialize - a model with the solution) - """ - new_sol = Solution( - self.all_ts[0][:1], - self.all_ys[0][:, :1], - self.all_models[:1], - self.all_inputs[:1], - None, - None, - "final time", - ) - new_sol._all_inputs_casadi = self.all_inputs_casadi[:1] - new_sol._sub_solutions = self.sub_solutions[:1] - - new_sol.solve_time = 0 - new_sol.integration_time = 0 - new_sol.set_up_time = 0 - - return new_sol - - @cached_property - def last_state(self): - """ - A Solution object that only contains the final state. This is faster to evaluate - than the full solution when only the final state is needed (e.g. to initialize - a model with the solution) - """ - new_sol = Solution( - self.all_ts[-1][-1:], - self.all_ys[-1][:, -1:], - self.all_models[-1:], - self.all_inputs[-1:], - self.t_event, - self.y_event, - self.termination, - ) - new_sol._all_inputs_casadi = self.all_inputs_casadi[-1:] - new_sol._sub_solutions = self.sub_solutions[-1:] - - new_sol.solve_time = 0 - new_sol.integration_time = 0 - new_sol.set_up_time = 0 - - return new_sol - - @property - def total_time(self): - return self.set_up_time + self.solve_time - - @property - def cycles(self): - return self._cycles - - @cycles.setter - def cycles(self, cycles): - self._cycles = cycles - - @property - def summary_variables(self): - return self._summary_variables - - def set_summary_variables(self, all_summary_variables): - summary_variables = {var: [] for var in all_summary_variables[0]} - for sum_vars in all_summary_variables: - for name, value in sum_vars.items(): - summary_variables[name].append(value) - - summary_variables["Cycle number"] = range(1, len(all_summary_variables) + 1) - self.all_summary_variables = all_summary_variables - self._summary_variables = pybamm.FuzzyDict( - {name: np.array(value) for name, value in summary_variables.items()} - ) - - def update(self, variables): - """Add ProcessedVariables to the dictionary of variables in the solution""" - # make sure that sensitivities are extracted if required - if isinstance(self._sensitivities, bool) and self._sensitivities: - self.extract_explicit_sensitivities() - - # Convert single entry to list - if isinstance(variables, str): - variables = [variables] - # Process - for key in variables: - cumtrapz_ic = None - pybamm.logger.debug("Post-processing {}".format(key)) - vars_pybamm = [model.variables_and_events[key] for model in self.all_models] - - # Iterate through all models, some may be in the list several times and - # therefore only get set up once - vars_casadi = [] - for i, (model, ys, inputs, var_pybamm) in enumerate( - zip(self.all_models, self.all_ys, self.all_inputs, vars_pybamm) - ): - if isinstance(var_pybamm, pybamm.ExplicitTimeIntegral): - cumtrapz_ic = var_pybamm.initial_condition - cumtrapz_ic = cumtrapz_ic.evaluate() - var_pybamm = var_pybamm.child - var_casadi = self.process_casadi_var( - var_pybamm, - inputs, - ys.shape, - ) - model._variables_casadi[key] = var_casadi - vars_pybamm[i] = var_pybamm - elif key in model._variables_casadi: - var_casadi = model._variables_casadi[key] - else: - var_casadi = self.process_casadi_var( - var_pybamm, - inputs, - ys.shape, - ) - model._variables_casadi[key] = var_casadi - vars_casadi.append(var_casadi) - var = pybamm.ProcessedVariable( - vars_pybamm, vars_casadi, self, cumtrapz_ic=cumtrapz_ic - ) - - # Save variable and data - self._variables[key] = var - self.data[key] = var.data - - def process_casadi_var(self, var_pybamm, inputs, ys_shape): - t_MX = casadi.MX.sym("t") - y_MX = casadi.MX.sym("y", ys_shape[0]) - inputs_MX_dict = { - key: casadi.MX.sym("input", value.shape[0]) for key, value in inputs.items() - } - inputs_MX = casadi.vertcat(*[p for p in inputs_MX_dict.values()]) - var_sym = var_pybamm.to_casadi(t_MX, y_MX, inputs=inputs_MX_dict) - var_casadi = casadi.Function("variable", [t_MX, y_MX, inputs_MX], [var_sym]) - return var_casadi - - def __getitem__(self, key): - """Read a variable from the solution. Variables are created 'just in time', i.e. - only when they are called. - - Parameters - ---------- - key : str - The name of the variable - - Returns - ------- - :class:`pybamm.ProcessedVariable` - A variable that can be evaluated at any time or spatial point. The - underlying data for this variable is available in its attribute ".data" - """ - - # return it if it exists - if key in self._variables: - return self._variables[key] - else: - # otherwise create it, save it and then return it - self.update(key) - return self._variables[key] - - def plot(self, output_variables=None, **kwargs): - """ - A method to quickly plot the outputs of the solution. Creates a - :class:`pybamm.QuickPlot` object (with keyword arguments 'kwargs') and - then calls :meth:`pybamm.QuickPlot.dynamic_plot`. - - Parameters - ---------- - output_variables: list, optional - A list of the variables to plot. - **kwargs - Additional keyword arguments passed to - :meth:`pybamm.QuickPlot.dynamic_plot`. - For a list of all possible keyword arguments see :class:`pybamm.QuickPlot`. - """ - return pybamm.dynamic_plot(self, output_variables=output_variables, **kwargs) - - def save(self, filename): - """Save the whole solution using pickle""" - # No warning here if len(self.data)==0 as solution can be loaded - # and used to process new variables - - with open(filename, "wb") as f: - pickle.dump(self, f, pickle.HIGHEST_PROTOCOL) - - def get_data_dict(self, variables=None, short_names=None, cycles_and_steps=True): - """ - Construct a (standard python) dictionary of the solution data containing the - variables in `variables`. If `variables` is None then all variables are - returned. Any variable names in short_names are replaced with the corresponding - short name. - - If the solution has cycles, then the cycle numbers and step numbers are also - returned in the dictionary. - - Parameters - ---------- - variables : list, optional - List of variables to return. If None, returns all variables in solution.data - short_names : dict, optional - Dictionary of shortened names to use when saving. - cycles_and_steps : bool, optional - Whether to include the cycle numbers and step numbers in the dictionary - - Returns - ------- - dict - A dictionary of the solution data - """ - if variables is None: - # variables not explicitly provided -> save all variables that have been - # computed - data_long_names = self.data - else: - if isinstance(variables, str): - variables = [variables] - # otherwise, save only the variables specified - data_long_names = {} - for name in variables: - data_long_names[name] = self[name].data - if len(data_long_names) == 0: - raise ValueError( - """ - Solution does not have any data. Please provide a list of variables - to save. - """ - ) - - # Use any short names if provided - data_short_names = {} - short_names = short_names or {} - for name, var in data_long_names.items(): - name = short_names.get(name, name) # return name if no short name - data_short_names[name] = var - - # Save cycle number and step number if the solution has them - if cycles_and_steps and len(self.cycles) > 0: - data_short_names["Cycle"] = np.array([]) - data_short_names["Step"] = np.array([]) - for i, cycle in enumerate(self.cycles): - data_short_names["Cycle"] = np.concatenate( - [data_short_names["Cycle"], i * np.ones_like(cycle.t)] - ) - for j, step in enumerate(cycle.steps): - data_short_names["Step"] = np.concatenate( - [data_short_names["Step"], j * np.ones_like(step.t)] - ) - - return data_short_names - - def save_data( - self, filename=None, variables=None, to_format="pickle", short_names=None - ): - """ - Save solution data only (raw arrays) - - Parameters - ---------- - filename : str, optional - The name of the file to save data to. If None, then a str is returned - variables : list, optional - List of variables to save. If None, saves all of the variables that have - been created so far - to_format : str, optional - The format to save to. Options are: - - - 'pickle' (default): creates a pickle file with the data dictionary - - 'matlab': creates a .mat file, for loading in matlab - - 'csv': creates a csv file (0D variables only) - - 'json': creates a json file - short_names : dict, optional - Dictionary of shortened names to use when saving. This may be necessary when - saving to MATLAB, since no spaces or special characters are allowed in - MATLAB variable names. Note that not all the variables need to be given - a short name. - - Returns - ------- - data : str, optional - str if 'csv' or 'json' is chosen and filename is None, otherwise None - """ - data = self.get_data_dict(variables=variables, short_names=short_names) - - if to_format == "pickle": - if filename is None: - raise ValueError("pickle format must be written to a file") - with open(filename, "wb") as f: - pickle.dump(data, f, pickle.HIGHEST_PROTOCOL) - elif to_format == "matlab": - if filename is None: - raise ValueError("matlab format must be written to a file") - # Check all the variable names only contain a-z, A-Z or _ or numbers - for name in data.keys(): - # Check the string only contains the following ASCII: - # a-z (97-122) - # A-Z (65-90) - # _ (95) - # 0-9 (48-57) but not in the first position - for i, s in enumerate(name): - if not ( - 97 <= ord(s) <= 122 - or 65 <= ord(s) <= 90 - or ord(s) == 95 - or (i > 0 and 48 <= ord(s) <= 57) - ): - raise ValueError( - "Invalid character '{}' found in '{}'. ".format(s, name) - + "MATLAB variable names must only contain a-z, A-Z, _, " - "or 0-9 (except the first position). " - "Use the 'short_names' argument to pass an alternative " - "variable name, e.g. \n\n" - "\tsolution.save_data(filename, " - "['Electrolyte concentration'], to_format='matlab, " - "short_names={'Electrolyte concentration': 'c_e'})" - ) - savemat(filename, data) - elif to_format == "csv": - for name, var in data.items(): - if var.ndim >= 2: - raise ValueError( - "only 0D variables can be saved to csv, but '{}' is {}D".format( - name, var.ndim - 1 - ) - ) - df = pd.DataFrame(data) - return df.to_csv(filename, index=False) - elif to_format == "json": - if filename is None: - return json.dumps(data, cls=NumpyEncoder) - else: - with open(filename, "w") as outfile: - json.dump(data, outfile, cls=NumpyEncoder) - else: - raise ValueError("format '{}' not recognised".format(to_format)) - - @property - def sub_solutions(self): - """List of sub solutions that have been - concatenated to form the full solution""" - - return self._sub_solutions - - def __add__(self, other): - """Adds two solutions together, e.g. when stepping""" - if other is None or isinstance(other, EmptySolution): - return self.copy() - if not isinstance(other, Solution): - raise pybamm.SolverError( - "Only a Solution or None can be added to a Solution" - ) - # Special case: new solution only has one timestep and it is already in the - # existing solution. In this case, return a copy of the existing solution - if ( - len(other.all_ts) == 1 - and len(other.all_ts[0]) == 1 - and other.all_ts[0][0] == self.all_ts[-1][-1] - ): - new_sol = self.copy() - # Update termination using the latter solution - new_sol._termination = other.termination - new_sol._t_event = other._t_event - new_sol._y_event = other._y_event - return new_sol - - # Update list of sub-solutions - if other.all_ts[0][0] == self.all_ts[-1][-1]: - # Skip first time step if it is repeated - all_ts = self.all_ts + [other.all_ts[0][1:]] + other.all_ts[1:] - all_ys = self.all_ys + [other.all_ys[0][:, 1:]] + other.all_ys[1:] - else: - all_ts = self.all_ts + other.all_ts - all_ys = self.all_ys + other.all_ys - - new_sol = Solution( - all_ts, - all_ys, - self.all_models + other.all_models, - self.all_inputs + other.all_inputs, - other.t_event, - other.y_event, - other.termination, - bool(self.sensitivities), - ) - - new_sol.closest_event_idx = other.closest_event_idx - new_sol._all_inputs_casadi = self.all_inputs_casadi + other.all_inputs_casadi - - # Set solution time - new_sol.solve_time = self.solve_time + other.solve_time - new_sol.integration_time = self.integration_time + other.integration_time - - # Set sub_solutions - new_sol._sub_solutions = self.sub_solutions + other.sub_solutions - - return new_sol - - def __radd__(self, other): - return self.__add__(other) - - def copy(self): - new_sol = self.__class__( - self.all_ts, - self.all_ys, - self.all_models, - self.all_inputs, - self.t_event, - self.y_event, - self.termination, - ) - new_sol._all_inputs_casadi = self.all_inputs_casadi - new_sol._sub_solutions = self.sub_solutions - new_sol.closest_event_idx = self.closest_event_idx - - new_sol.solve_time = self.solve_time - new_sol.integration_time = self.integration_time - new_sol.set_up_time = self.set_up_time - - return new_sol - - -class EmptySolution: - def __init__(self, termination=None, t=None): - self.termination = termination - if t is None: - t = np.array([0]) - elif isinstance(t, numbers.Number): - t = np.array([t]) - - self.t = t - - def __add__(self, other): - if isinstance(other, (EmptySolution, Solution)): - return other.copy() - - def __radd__(self, other): - if other is None: - return self.copy() - - def copy(self): - return EmptySolution(termination=self.termination, t=self.t) - - -def make_cycle_solution(step_solutions, esoh_solver=None, save_this_cycle=True): - """ - Function to create a Solution for an entire cycle, and associated summary variables - - Parameters - ---------- - step_solutions : list of :class:`Solution` - Step solutions that form the entire cycle - esoh_solver : :class:`pybamm.lithium_ion.ElectrodeSOHSolver` - Solver to calculate electrode SOH (eSOH) variables. If `None` (default) - then only summary variables that do not require the eSOH calculation - are calculated. See :footcite:t:`Mohtat2019` for more details on eSOH variables. - save_this_cycle : bool, optional - Whether to save the entire cycle variables or just the summary variables. - Default True - - Returns - ------- - cycle_solution : :class:`pybamm.Solution` or None - The Solution object for this cycle, or None (if save_this_cycle is False) - cycle_summary_variables : dict - Dictionary of summary variables for this cycle - - """ - sum_sols = step_solutions[0].copy() - for step_solution in step_solutions[1:]: - sum_sols = sum_sols + step_solution - - cycle_solution = Solution( - sum_sols.all_ts, - sum_sols.all_ys, - sum_sols.all_models, - sum_sols.all_inputs, - sum_sols.t_event, - sum_sols.y_event, - sum_sols.termination, - ) - cycle_solution._all_inputs_casadi = sum_sols.all_inputs_casadi - cycle_solution._sub_solutions = sum_sols.sub_solutions - - cycle_solution.solve_time = sum_sols.solve_time - cycle_solution.integration_time = sum_sols.integration_time - cycle_solution.set_up_time = sum_sols.set_up_time - - cycle_solution.steps = step_solutions - - cycle_summary_variables = _get_cycle_summary_variables(cycle_solution, esoh_solver) - - cycle_first_state = cycle_solution.first_state - - if save_this_cycle: - cycle_solution.cycle_summary_variables = cycle_summary_variables - else: - cycle_solution = None - - return cycle_solution, cycle_summary_variables, cycle_first_state - - -def _get_cycle_summary_variables(cycle_solution, esoh_solver): - model = cycle_solution.all_models[0] - cycle_summary_variables = pybamm.FuzzyDict({}) - - # Measured capacity variables - if "Discharge capacity [A.h]" in model.variables: - Q = cycle_solution["Discharge capacity [A.h]"].data - min_Q, max_Q = np.min(Q), np.max(Q) - - cycle_summary_variables.update( - { - "Minimum measured discharge capacity [A.h]": min_Q, - "Maximum measured discharge capacity [A.h]": max_Q, - "Measured capacity [A.h]": max_Q - min_Q, - } - ) - - # Voltage variables - if "Battery voltage [V]" in model.variables: - V = cycle_solution["Battery voltage [V]"].data - min_V, max_V = np.min(V), np.max(V) - - cycle_summary_variables.update( - {"Minimum voltage [V]": min_V, "Maximum voltage [V]": max_V} - ) - - # Degradation variables - degradation_variables = model.summary_variables - first_state = cycle_solution.first_state - last_state = cycle_solution.last_state - for var in degradation_variables: - data_first = first_state[var].data - data_last = last_state[var].data - cycle_summary_variables[var] = data_last[0] - var_lowercase = var[0].lower() + var[1:] - cycle_summary_variables["Change in " + var_lowercase] = ( - data_last[0] - data_first[0] - ) - - # eSOH variables (full-cell lithium-ion model only, for now) - if ( - esoh_solver is not None - and isinstance(model, pybamm.lithium_ion.BaseModel) - and model.options.electrode_types["negative"] == "porous" - ): - Q_n = last_state["Negative electrode capacity [A.h]"].data[0] - Q_p = last_state["Positive electrode capacity [A.h]"].data[0] - Q_Li = last_state["Total lithium capacity in particles [A.h]"].data[0] - - inputs = {"Q_n": Q_n, "Q_p": Q_p, "Q_Li": Q_Li} - - try: - esoh_sol = esoh_solver.solve(inputs) - except pybamm.SolverError: # pragma: no cover - raise pybamm.SolverError( - "Could not solve for summary variables, run " - "`sim.solve(calc_esoh=False)` to skip this step" - ) - - cycle_summary_variables.update(esoh_sol) - - return cycle_summary_variables diff --git a/pybamm/solvers/solution_full.py b/pybamm/solvers/solution_full.py deleted file mode 100644 index a73363b6ec..0000000000 --- a/pybamm/solvers/solution_full.py +++ /dev/null @@ -1,8 +0,0 @@ -# -# Solution class -# -import pybamm - - -class SolutionFull(pybamm.SolutionBase): - ... diff --git a/pybamm/solvers/solution_vars.py b/pybamm/solvers/solution_vars.py deleted file mode 100644 index dc0df83833..0000000000 --- a/pybamm/solvers/solution_vars.py +++ /dev/null @@ -1,8 +0,0 @@ -# -# Solution class -# -import pybamm - - -class SolutionVars(pybamm.SolutionBase): - ... diff --git a/test.py b/test.py index c3c9663689..b4080e8c1c 100644 --- a/test.py +++ b/test.py @@ -3,9 +3,6 @@ import matplotlib.pylab as plt import importlib -solver_opt = 2 -jacobian = 'sparse' # sparse, dense, band, none -num_threads = 1 output_variables = [ "Voltage [V]", "Time [min]", @@ -15,8 +12,8 @@ all_vars = False input_parameters = { - "Current function [A]": 0.15652, - "Separator porosity": 0.47, + "Current function [A]": 0.680616, + "Separator porosity": 1.0, } # check for loading errors @@ -40,28 +37,10 @@ disc.process_model(model) t_eval = np.linspace(0, 3600, 100) -if solver_opt == 1: - linear_solver = 'SUNLinSol_Dense' -if solver_opt == 2: - linear_solver = 'SUNLinSol_KLU' -if solver_opt == 3: - linear_solver = 'SUNLinSol_Band' -if solver_opt == 4: - linear_solver = 'SUNLinSol_SPBCGS' -if solver_opt == 5: - linear_solver = 'SUNLinSol_SPFGMR' -if solver_opt == 6: - linear_solver = 'SUNLinSol_SPGMR' -if solver_opt == 7: - linear_solver = 'SUNLinSol_SPTFQMR' -if solver_opt == 8: - linear_solver = 'SUNLinSol_cuSolverSp_batchQR' - jacobian = 'cuSparse_' - options = { - 'linear_solver': linear_solver, - 'jacobian': jacobian, - 'num_threads': num_threads, + 'linear_solver': 'SUNLinSol_KLU', + 'jacobian': 'sparse', + 'num_threads': 4, } if all_vars: @@ -95,20 +74,22 @@ print(f"Solve time: {sol.solve_time.value*1000} msecs") if True: - #output_variables = [ - # "Voltage [V]", - # "Time [min]", - # "Current [A]", - #] + plot_allvars = False + if not output_variables: + plot_allvars = True + output_variables = [ + "Voltage [V]", + "Time [min]", + "Current [A]", + ] fig, axs = plt.subplots(len(output_variables), len(input_parameters)+1) for k, var in enumerate(output_variables): - if False: - axs[k,0].plot(t_eval, sol[var](t_eval)) - for paramk, param in enumerate(list(input_parameters.keys())): - axs[k,paramk+1].plot(t_eval, sol[var].sensitivities[param]) # time, param, var + # Solution variables currently use different classes/calls + if not input_parameters: + axs[k].plot(t_eval, sol[var](t_eval)) else: - axs[k,0].plot(t_eval, sol[var][:,0]) + axs[k,0].plot(t_eval, sol[var](t_eval)) for paramk, param in enumerate(list(input_parameters.keys())): - axs[k,paramk+1].plot(t_eval, sol.yS[:,k,paramk]) # time, param, var + axs[k,paramk+1].plot(t_eval, sol[var].sensitivities[param]) plt.tight_layout() plt.show() From d7de190491a9118d61f96114cb128a827db4f9bc Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 20 Jul 2023 18:00:32 +0000 Subject: [PATCH 09/38] Variables now correctly unroll in the ProcessedVariablesVar class --- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 13 +- .../c_solvers/idaklu/casadi_functions.cpp | 12 +- pybamm/solvers/idaklu_solver.py | 15 +- pybamm/solvers/processed_variable_var.py | 391 ++++++++++++++++++ test.py | 48 ++- 5 files changed, 436 insertions(+), 43 deletions(-) create mode 100644 pybamm/solvers/processed_variable_var.py diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index daf296e3bd..e3e7117896 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -236,7 +236,7 @@ void CasadiSolverOpenMP::CalcVars( for (auto& var_fcn : functions->var_casadi_fcns) { var_fcn({tret, yval, functions->inputs.data()}, {res}); // store in return vector - for (size_t jj=0; jjvar_casadi_fcns.size() > 0) { // return only the requested variables list after computation for (auto& var_fcn : functions->var_casadi_fcns) { - max_res_size = std::max(max_res_size, var_fcn.m_res.size()); - if (var_fcn.m_res.size() > 0) - length_of_return_vector += var_fcn.m_res.size() - 1; + max_res_size = std::max(max_res_size, size_t(var_fcn.m_func.nnz_out())); + length_of_return_vector += var_fcn.m_func.nnz_out(); for (auto& dvar_fcn : functions->dvar_dy_fcns) - max_res_dvar_dy = std::max(max_res_dvar_dy, dvar_fcn.m_res.size()); + max_res_dvar_dy = std::max(max_res_dvar_dy, size_t(dvar_fcn.m_func.nnz_out())); for (auto& dvar_fcn : functions->dvar_dp_fcns) - max_res_dvar_dp = std::max(max_res_dvar_dp, dvar_fcn.m_res.size()); + max_res_dvar_dp = std::max(max_res_dvar_dp, size_t(dvar_fcn.m_func.nnz_out())); } } else { // Return full y state-vector @@ -369,7 +368,7 @@ Solution CasadiSolverOpenMP::solve( number_of_timesteps * length_of_return_vector]; - res = new realtype[number_of_states]; // TODO: Crashes if set to max_res_size + res = new realtype[max_res_size]; res_dvar_dy = new realtype[max_res_dvar_dy]; res_dvar_dp = new realtype[max_res_dvar_dp]; diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index 4c6dfc7362..829e983b42 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -7,16 +7,10 @@ CasadiFunction::CasadiFunction(const Function &f) : m_func(f) size_t sz_iw; size_t sz_w; m_func.sz_work(sz_arg, sz_res, sz_iw, sz_w); + int nnz = (sz_res>0) ? m_func.nnz_out() : 0; //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); - // } + // << sz_res << " iw = " << sz_iw << " w = " << sz_w << " nnz = " << nnz << + // std::endl; m_arg.resize(sz_arg, nullptr); m_res.resize(sz_res, nullptr); m_iw.resize(sz_iw, 0); diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 644ecf3537..82d3831312 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -285,7 +285,8 @@ def resfn(t, y, inputs, ydot): # if output_variables specified then convert functions to casadi # expressions for evaluation within the idaklu solver - self.var_casadi_fcns = [] + self.var_casadi_fcns = {} + self.var_casadi_idaklu_fcns = [] self.dvar_dy_fcns = [] self.dvar_dp_fcns = [] for key in self.output_variables: @@ -293,13 +294,13 @@ def resfn(t, y, inputs, ydot): fcn_name = wrangle_name(key) var_casadi = model.variables_and_events[key].to_casadi( t_casadi, y_casadi, inputs=p_casadi) - var_casadi_fcn = casadi.Function( + self.var_casadi_fcns[key] = casadi.Function( fcn_name, [t_casadi, y_casadi, p_casadi_stacked], [var_casadi] ) - self.var_casadi_fcns.append( - idaklu.generate_function(var_casadi_fcn.serialize()) + self.var_casadi_idaklu_fcns.append( + idaklu.generate_function(self.var_casadi_fcns[key].serialize()) ) # Generate derivative functions for sensitivities if len(inputs) > 0: @@ -494,6 +495,7 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): "number_of_sensitivity_parameters": number_of_sensitivity_parameters, "output_variables": self.output_variables, "var_casadi_fcns": self.var_casadi_fcns, + "var_casadi_idaklu_fcns": self.var_casadi_idaklu_fcns, "dvar_dy_fcns": self.dvar_dy_fcns, "dvar_dp_fcns": self.dvar_dp_fcns, } @@ -517,7 +519,7 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): atol=atol, rtol=rtol, inputs=len(inputs), - var_casadi_fcns=self._setup["var_casadi_fcns"], + var_casadi_fcns=self._setup["var_casadi_idaklu_fcns"], dvar_dy_fcns=self._setup["dvar_dy_fcns"], dvar_dp_fcns=self._setup["dvar_dp_fcns"], options=self._options, @@ -679,9 +681,10 @@ def _integrate(self, model, t_eval, inputs_dict=None): sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) startk = 0 for vark, var in enumerate(self.output_variables): - len_of_var = int(model.variables_and_events[var].size) + len_of_var = self._setup["var_casadi_fcns"][var](0,0,0).sparsity().nnz() newsol._variables[var] = pybamm.ProcessedVariableVar( [model.variables_and_events[var]], + [self._setup["var_casadi_fcns"][var]], [sol.y[:, startk:(startk+len_of_var)]], newsol, ) diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_var.py new file mode 100644 index 0000000000..0ec1c3bfdb --- /dev/null +++ b/pybamm/solvers/processed_variable_var.py @@ -0,0 +1,391 @@ +# +# Processed Variable class +# +import casadi +import numbers +import numpy as np +import pybamm +from scipy.integrate import cumulative_trapezoid +import xarray as xr + + +class ProcessedVariableVar(object): + """ + An object that can be evaluated at arbitrary (scalars or vectors) t and x, and + returns the (interpolated) value of the base variable at that t and x. + + Parameters + ---------- + base_variables : list of :class:`pybamm.Symbol` + A list of base variables with a method `evaluate(t,y)`, each entry of which + returns the value of that variable for that particular sub-solution. + A Solution can be comprised of sub-solutions which are the solutions of + different models. + Note that this can be any kind of node in the expression tree, not + just a :class:`pybamm.Variable`. + When evaluated, returns an array of size (m,n) + base_variable_casadis : list of :numpy:array + A list of numpy arrays; the returns from casadi evaluation. the same thing as + `base_Variable.evaluate` (but more efficient). + solution : :class:`pybamm.Solution` + The solution object to be used to create the processed variables + warn : bool, optional + Whether to raise warnings when trying to evaluate time and length scales. + Default is True. + """ + + def __init__( + self, + base_variables, + base_variables_casadi, + base_variables_data, + solution, + warn=True, + cumtrapz_ic=None, + ): + self.base_variables = base_variables + self.base_variables_casadi = base_variables_casadi + self.base_variables_data = base_variables_data + + self.all_ts = solution.all_ts + self.all_ys = solution.all_ys + self.all_inputs = solution.all_inputs + self.all_inputs_casadi = solution.all_inputs_casadi + + self.mesh = base_variables[0].mesh + self.domain = base_variables[0].domain + self.domains = base_variables[0].domains + self.warn = warn + self.cumtrapz_ic = cumtrapz_ic + + # Sensitivity starts off uninitialized, only set when called + self._sensitivities = None + self.solution_sensitivities = solution.sensitivities + + # Store time + self.t_pts = solution.t + + # Evaluate base variable at initial time + self.base_eval_shape = self.base_variables[0].shape + self.base_eval_size = self.base_variables[0].size + + # handle 2D (in space) finite element variables differently + if ( + self.mesh + and "current collector" in self.domain + and isinstance(self.mesh, pybamm.ScikitSubMesh2D) + ): + self.initialise_2D_scikit_fem() + + # check variable shape + else: + if ( + len(self.base_eval_shape) == 0 + or self.base_eval_shape[0] == 1 + ): + self.initialise_0D() + else: + n = self.mesh.npts + base_shape = self.base_eval_shape[0] + # Try some shapes that could make the variable a 1D variable + if base_shape in [n, n + 1]: + self.initialise_1D() + else: + # Try some shapes that could make the variable a 2D variable + first_dim_nodes = self.mesh.nodes + first_dim_edges = self.mesh.edges + second_dim_pts = self.base_variables[0].secondary_mesh.nodes + if self.base_eval_size // len(second_dim_pts) in [ + len(first_dim_nodes), + len(first_dim_edges), + ]: + self.initialise_2D() + else: + # Raise error for 3D variable + raise NotImplementedError( + "Shape not recognized for {} ".format(base_variables[0]) + + "(note processing of 3D variables is not yet implemented)" + ) + + def initialise_0D(self): + entries = np.concatenate(self.base_variables_data, axis=0).flatten() + + if self.cumtrapz_ic is not None: + entries = cumulative_trapezoid( + entries, self.t_pts, initial=float(self.cumtrapz_ic) + ) + + # set up interpolation + self._xr_data_array = xr.DataArray(entries, coords=[("t", self.t_pts)]) + + self.entries = entries + self.dimensions = 0 + + def _unroll_nnz(self): + # unroll in nnz != numel, otherwise copy + sp = self.base_variables_casadi[0](0,0,0).sparsity() + if sp.nnz() != sp.numel(): + data = [None]*len(self.base_variables_data) + for datak in range(len(self.base_variables_data)): + data[datak] = np.zeros(self.base_eval_shape[0]*len(self.t_pts)) + var_data = self.base_variables_data[0].flatten() + k = 0 + for t_i in range(len(self.t_pts)): + base = t_i*sp.numel() + for r in sp.row(): + data[datak][base + r] = var_data[k] + k = k + 1 + else: + data = self.base_variables_data + return data + + def initialise_1D(self, fixed_t=False): + len_space = self.base_eval_shape[0] + # reshape + entries = np.concatenate(self._unroll_nnz(), axis=0)\ + .flatten()\ + .reshape((len(self.t_pts),len_space))\ + .transpose() + + # Get node and edge values + nodes = self.mesh.nodes + edges = self.mesh.edges + if entries.shape[0] == len(nodes): + space = nodes + elif entries.shape[0] == len(edges): + space = edges + + # add points outside domain for extrapolation to boundaries + extrap_space_left = np.array([2 * space[0] - space[1]]) + extrap_space_right = np.array([2 * space[-1] - space[-2]]) + space = np.concatenate([extrap_space_left, space, extrap_space_right]) + extrap_entries_left = 2 * entries[0] - entries[1] + extrap_entries_right = 2 * entries[-1] - entries[-2] + entries_for_interp = np.vstack( + [extrap_entries_left, entries, extrap_entries_right] + ) + + # assign attributes for reference (either x_sol or r_sol) + self.entries = entries + self.dimensions = 1 + if self.domain[0].endswith("particle"): + self.first_dimension = "r" + self.r_sol = space + elif self.domain[0] in [ + "negative electrode", + "separator", + "positive electrode", + ]: + self.first_dimension = "x" + self.x_sol = space + elif self.domain == ["current collector"]: + self.first_dimension = "z" + self.z_sol = space + elif self.domain[0].endswith("particle size"): + self.first_dimension = "R" + self.R_sol = space + else: + self.first_dimension = "x" + self.x_sol = space + + # assign attributes for reference + pts_for_interp = space + self.internal_boundaries = self.mesh.internal_boundaries + + # Set first_dim_pts to edges for nicer plotting + self.first_dim_pts = edges + + # set up interpolation + self._xr_data_array = xr.DataArray( + entries_for_interp, + coords=[(self.first_dimension, pts_for_interp), ("t", self.t_pts)], + ) + + def initialise_2D(self): + """ + Initialise a 2D object that depends on x and r, x and z, x and R, or R and r. + """ + first_dim_nodes = self.mesh.nodes + first_dim_edges = self.mesh.edges + second_dim_nodes = self.base_variables[0].secondary_mesh.nodes + second_dim_edges = self.base_variables[0].secondary_mesh.edges + if self.base_eval_size // len(second_dim_nodes) == len(first_dim_nodes): + first_dim_pts = first_dim_nodes + elif self.base_eval_size // len(second_dim_nodes) == len(first_dim_edges): + first_dim_pts = first_dim_edges + + second_dim_pts = second_dim_nodes + first_dim_size = len(first_dim_pts) + second_dim_size = len(second_dim_pts) + + len_space = self.base_eval_shape[0] + entries = np.concatenate(self._unroll_nnz(), axis=0)\ + .flatten()\ + .reshape((len(self.t_pts), second_dim_size, first_dim_size)) + entries = np.moveaxis(entries, 0, 2) + entries = np.moveaxis(entries, 0, 1) + + # add points outside first dimension domain for extrapolation to + # boundaries + extrap_space_first_dim_left = np.array( + [2 * first_dim_pts[0] - first_dim_pts[1]] + ) + extrap_space_first_dim_right = np.array( + [2 * first_dim_pts[-1] - first_dim_pts[-2]] + ) + first_dim_pts = np.concatenate( + [extrap_space_first_dim_left, first_dim_pts, extrap_space_first_dim_right] + ) + extrap_entries_left = np.expand_dims(2 * entries[0] - entries[1], axis=0) + extrap_entries_right = np.expand_dims(2 * entries[-1] - entries[-2], axis=0) + entries_for_interp = np.concatenate( + [extrap_entries_left, entries, extrap_entries_right], axis=0 + ) + + # add points outside second dimension domain for extrapolation to + # boundaries + extrap_space_second_dim_left = np.array( + [2 * second_dim_pts[0] - second_dim_pts[1]] + ) + extrap_space_second_dim_right = np.array( + [2 * second_dim_pts[-1] - second_dim_pts[-2]] + ) + second_dim_pts = np.concatenate( + [ + extrap_space_second_dim_left, + second_dim_pts, + extrap_space_second_dim_right, + ] + ) + extrap_entries_second_dim_left = np.expand_dims( + 2 * entries_for_interp[:, 0, :] - entries_for_interp[:, 1, :], axis=1 + ) + extrap_entries_second_dim_right = np.expand_dims( + 2 * entries_for_interp[:, -1, :] - entries_for_interp[:, -2, :], axis=1 + ) + entries_for_interp = np.concatenate( + [ + extrap_entries_second_dim_left, + entries_for_interp, + extrap_entries_second_dim_right, + ], + axis=1, + ) + + # Process r-x, x-z, r-R, R-x, or R-z + if self.domain[0].endswith("particle") and self.domains["secondary"][ + 0 + ].endswith("electrode"): + self.first_dimension = "r" + self.second_dimension = "x" + self.r_sol = first_dim_pts + self.x_sol = second_dim_pts + elif self.domain[0] in [ + "negative electrode", + "separator", + "positive electrode", + ] and self.domains["secondary"] == ["current collector"]: + self.first_dimension = "x" + self.second_dimension = "z" + self.x_sol = first_dim_pts + self.z_sol = second_dim_pts + elif self.domain[0].endswith("particle") and self.domains["secondary"][ + 0 + ].endswith("particle size"): + self.first_dimension = "r" + self.second_dimension = "R" + self.r_sol = first_dim_pts + self.R_sol = second_dim_pts + elif self.domain[0].endswith("particle size") and self.domains["secondary"][ + 0 + ].endswith("electrode"): + self.first_dimension = "R" + self.second_dimension = "x" + self.R_sol = first_dim_pts + self.x_sol = second_dim_pts + elif self.domain[0].endswith("particle size") and self.domains["secondary"] == [ + "current collector" + ]: + self.first_dimension = "R" + self.second_dimension = "z" + self.R_sol = first_dim_pts + self.z_sol = second_dim_pts + else: # pragma: no cover + raise pybamm.DomainError( + f"Cannot process 2D object with domains '{self.domains}'." + ) + + # assign attributes for reference + self.entries = entries + self.dimensions = 2 + first_dim_pts_for_interp = first_dim_pts + second_dim_pts_for_interp = second_dim_pts + + # Set pts to edges for nicer plotting + self.first_dim_pts = first_dim_edges + self.second_dim_pts = second_dim_edges + + # set up interpolation + self._xr_data_array = xr.DataArray( + entries_for_interp, + coords={ + self.first_dimension: first_dim_pts_for_interp, + self.second_dimension: second_dim_pts_for_interp, + "t": self.t_pts, + }, + ) + + def initialise_2D_scikit_fem(self): + y_sol = self.mesh.edges["y"] + len_y = len(y_sol) + z_sol = self.mesh.edges["z"] + len_z = len(z_sol) + entries = np.concatenate(self._unroll_nnz(), axis=0)\ + .flatten()\ + .reshape((len(self.t_pts), len_y, len_z)) + entries = np.moveaxis(entries, 0, 2) + + # assign attributes for reference + self.entries = entries + self.dimensions = 2 + self.y_sol = y_sol + self.z_sol = z_sol + self.first_dimension = "y" + self.second_dimension = "z" + self.first_dim_pts = y_sol + self.second_dim_pts = z_sol + + # set up interpolation + self._xr_data_array = xr.DataArray( + entries, + coords={"y": y_sol, "z": z_sol, "t": self.t_pts}, + ) + + def __call__(self, t=None, x=None, r=None, y=None, z=None, R=None, warn=True): + """ + Evaluate the variable at arbitrary *dimensional* t (and x, r, y, z and/or R), + using interpolation + """ + kwargs = {"t": t, "x": x, "r": r, "y": y, "z": z, "R": R} + # Remove any None arguments + kwargs = {key: value for key, value in kwargs.items() if value is not None} + # Use xarray interpolation, return numpy array + return self._xr_data_array.interp(**kwargs).values + + @property + def data(self): + """Same as entries, but different name""" + return self.entries + + @property + def sensitivities(self): + """ + Returns a dictionary of sensitivities for each input parameter. + The keys are the input parameters, and the value is a matrix of size + (n_x * n_t, n_p), where n_x is the number of states, n_t is the number of time + points, and n_p is the size of the input parameter + """ + # No sensitivities if there are no inputs + if len(self.all_inputs[0]) == 0: + return {} + return self._sensitivities diff --git a/test.py b/test.py index b4080e8c1c..beb612d2de 100644 --- a/test.py +++ b/test.py @@ -7,8 +7,12 @@ "Voltage [V]", "Time [min]", "Current [A]", + "r_n [m]", + "x [m]", + "Gradient of negative electrolyte potential [V.m-1]", + "Negative particle flux [mol.m-2.s-1]", ] -#output_variables = [] +output_variables = [] all_vars = False input_parameters = { @@ -16,15 +20,12 @@ "Separator porosity": 1.0, } -# check for loading errors idaklu_spec = importlib.util.find_spec("pybamm.solvers.idaklu") idaklu = importlib.util.module_from_spec(idaklu_spec) idaklu_spec.loader.exec_module(idaklu) # construct model -# pybamm.set_logging_level("INFO") model = pybamm.lithium_ion.DFN() -# model.convert_to_format = 'jax' geometry = model.default_geometry param = model.default_parameter_values param.update({key: "[input]" for key in input_parameters}) @@ -53,28 +54,33 @@ print("ExplicitTimeIntegral variables:") print(left_out) -print("output_variables:") -print(output_variables) -print("\nInput parameters:") -print(input_parameters) +#var_names = [[]] +var_names = [output_variables] +#var_names = [[m] for m in model.variable_names()] +#var_names = [[m] for m in output_variables] +for output_vars in var_names: + output_variables = output_vars + print("Output variables:") + print(output_variables) + print("\nInput parameters:") + print(input_parameters) -solver = pybamm.IDAKLUSolver( - atol=1e-8, rtol=1e-8, - options=options, - output_variables=output_variables, -) + solver = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + output_variables=output_variables, + ) -sol = solver.solve( - model, - t_eval, - inputs=input_parameters, - calculate_sensitivities=True, -) + sol = solver.solve( + model, + t_eval, + inputs=input_parameters, + calculate_sensitivities=True, + ) -print(f"Solve time: {sol.solve_time.value*1000} msecs") + print(f"Solve time: {sol.solve_time.value*1000} msecs") if True: - plot_allvars = False if not output_variables: plot_allvars = True output_variables = [ From ef90a615bd908be09af0f849ab5265b624330d0f Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 20 Jul 2023 23:47:09 +0000 Subject: [PATCH 10/38] Reshape sensitivities using base variable parameters --- pybamm/solvers/idaklu_solver.py | 55 +++++----- pybamm/solvers/processed_variable.py | 5 +- pybamm/solvers/processed_variable_var.py | 134 ++++++++++++++++------- test.py | 5 +- 4 files changed, 126 insertions(+), 73 deletions(-) diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 82d3831312..a6881c17e5 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -20,19 +20,22 @@ def have_idaklu(): return idaklu_spec is not None - + + def wrangle_name(name: str) -> str: - return name.casefold() \ - .replace(" ", "_") \ - .replace("[", "") \ - .replace("]", "") \ - .replace(".", "_") \ - .replace("-", "_") \ - .replace("(", "") \ - .replace(")", "") \ - .replace("%", "prc") \ - .replace(",", "") \ + return ( + name.casefold() + .replace(" ", "_") + .replace("[", "") + .replace("]", "") + .replace(".", "_") + .replace("-", "_") + .replace("(", "") + .replace(")", "") + .replace("%", "prc") + .replace(",", "") .replace(".", "") + ) class IDAKLUSolver(pybamm.BaseSolver): @@ -293,11 +296,10 @@ def resfn(t, y, inputs, ydot): # variable functions fcn_name = wrangle_name(key) var_casadi = model.variables_and_events[key].to_casadi( - t_casadi, y_casadi, inputs=p_casadi) + t_casadi, y_casadi, inputs=p_casadi + ) self.var_casadi_fcns[key] = casadi.Function( - fcn_name, - [t_casadi, y_casadi, p_casadi_stacked], - [var_casadi] + fcn_name, [t_casadi, y_casadi, p_casadi_stacked], [var_casadi] ) self.var_casadi_idaklu_fcns.append( idaklu.generate_function(self.var_casadi_fcns[key].serialize()) @@ -309,12 +311,12 @@ def resfn(t, y, inputs, ydot): dvar_dy_fcn = casadi.Function( f"d{fcn_name}_dy", [t_casadi, y_casadi, p_casadi_stacked], - [dvar_dy] + [dvar_dy], ) dvar_dp_fcn = casadi.Function( f"d{fcn_name}_dp", [t_casadi, y_casadi, p_casadi_stacked], - [dvar_dp] + [dvar_dp], ) self.dvar_dy_fcns.append( idaklu.generate_function(dvar_dy_fcn.serialize()) @@ -557,8 +559,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): inputs_dict = inputs_dict or {} # stack inputs if inputs_dict: - arrays_to_stack = [np.array(x).reshape(-1, 1) - for x in inputs_dict.values()] + arrays_to_stack = [np.array(x).reshape(-1, 1) for x in inputs_dict.values()] inputs = np.vstack(arrays_to_stack) else: inputs = np.array([[]]) @@ -638,7 +639,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): number_of_states = y0.size if self.output_variables: # Substitute empty vectors for state vector 'y' - y_out = np.zeros((number_of_timesteps*number_of_states, 0)) + y_out = np.zeros((number_of_timesteps * number_of_states, 0)) else: y_out = sol.y.reshape((number_of_timesteps, number_of_states)) @@ -681,19 +682,23 @@ def _integrate(self, model, t_eval, inputs_dict=None): sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) startk = 0 for vark, var in enumerate(self.output_variables): - len_of_var = self._setup["var_casadi_fcns"][var](0,0,0).sparsity().nnz() + len_of_var = ( + self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz() + ) newsol._variables[var] = pybamm.ProcessedVariableVar( [model.variables_and_events[var]], [self._setup["var_casadi_fcns"][var]], - [sol.y[:, startk:(startk+len_of_var)]], + [sol.y[:, startk : (startk + len_of_var)]], newsol, ) - startk += len_of_var - # Add sensitivities newsol[var]._sensitivities = {} for paramk, param in enumerate(inputs_dict.keys()): - newsol[var]._sensitivities[param] = sol.yS[:, vark, paramk] + newsol[var].add_sensitivity( + param, + [sol.yS[:, startk : (startk + len_of_var), paramk]] + ) + startk += len_of_var return newsol else: raise pybamm.SolverError("idaklu solver failed") diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py index 3c920b659e..5c14fd504b 100644 --- a/pybamm/solvers/processed_variable.py +++ b/pybamm/solvers/processed_variable.py @@ -77,10 +77,7 @@ def __init__( # check variable shape else: - if ( - len(self.base_eval_shape) == 0 - or self.base_eval_shape[0] == 1 - ): + if len(self.base_eval_shape) == 0 or self.base_eval_shape[0] == 1: self.initialise_0D() else: n = self.mesh.npts diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_var.py index 0ec1c3bfdb..740470ed8f 100644 --- a/pybamm/solvers/processed_variable_var.py +++ b/pybamm/solvers/processed_variable_var.py @@ -24,9 +24,11 @@ class ProcessedVariableVar(object): Note that this can be any kind of node in the expression tree, not just a :class:`pybamm.Variable`. When evaluated, returns an array of size (m,n) - base_variable_casadis : list of :numpy:array - A list of numpy arrays; the returns from casadi evaluation. the same thing as - `base_Variable.evaluate` (but more efficient). + base_variable_casadis : list of :class:`casadi.Function` + A list of casadi functions. When evaluated, returns the same thing as + `base_Variable.evaluate` (but more efficiently). + base_variable_data : list of :numpy:array + A list of numpy arrays, the returned evaluations. solution : :class:`pybamm.Solution` The solution object to be used to create the processed variables warn : bool, optional @@ -68,6 +70,7 @@ def __init__( # Evaluate base variable at initial time self.base_eval_shape = self.base_variables[0].shape self.base_eval_size = self.base_variables[0].size + self.unroll_params = {} # handle 2D (in space) finite element variables differently if ( @@ -79,10 +82,7 @@ def __init__( # check variable shape else: - if ( - len(self.base_eval_shape) == 0 - or self.base_eval_shape[0] == 1 - ): + if len(self.base_eval_shape) == 0 or self.base_eval_shape[0] == 1: self.initialise_0D() else: n = self.mesh.npts @@ -107,9 +107,80 @@ def __init__( + "(note processing of 3D variables is not yet implemented)" ) + def add_sensitivity(self, param, data): + self._sensitivities[param] = self.unroll(data) + + def _unroll_nnz(self, realdata=None): + # unroll in nnz != numel, otherwise copy + if realdata is None: + realdata = self.base_variables_data + sp = self.base_variables_casadi[0](0, 0, 0).sparsity() + if sp.nnz() != sp.numel(): + data = [None] * len(realdata) + for datak in range(len(realdata)): + data[datak] = np.zeros(self.base_eval_shape[0] * len(self.t_pts)) + var_data = realdata[0].flatten() + k = 0 + for t_i in range(len(self.t_pts)): + base = t_i * sp.numel() + for r in sp.row(): + data[datak][base + r] = var_data[k] + k = k + 1 + else: + data = realdata + return data + + def unroll_0D(self, realdata=None): + if realdata is None: + realdata = self.base_variables_data + return np.concatenate(realdata, axis=0).flatten() + + def unroll_1D(self, realdata=None): + len_space = self.base_eval_shape[0] + return ( + np.concatenate(self._unroll_nnz(realdata), axis=0) + .flatten() + .reshape((len(self.t_pts), len_space)) + .transpose() + ) + + def unroll_2D(self, realdata=None, n_dim1=None, n_dim2=None, axis_swaps=[]): + # initialise settings on first run + if not self.unroll_params: + self.unroll_params["n_dim1"] = n_dim1 + self.unroll_params["n_dim2"]= n_dim2 + self.unroll_params["axis_swaps"] = axis_swaps + # use stored settings on subsequent runs + if not n_dim1: + n_dim1 = self.unroll_params["n_dim1"] + n_dim2 = self.unroll_params["n_dim2"] + axis_swaps = self.unroll_params["axis_swaps"] + len_space = self.base_eval_shape[0] + entries = ( + np.concatenate(self._unroll_nnz(realdata), axis=0) + .flatten() + .reshape((len(self.t_pts), n_dim1, n_dim2)) + ) + for a, b in axis_swaps: + entries = np.moveaxis(entries, a, b) + return entries + + def unroll(self, realdata=None): + if self.dimensions == 0: + return self.unroll_0D(realdata=realdata) + elif self.dimensions == 1: + return self.unroll_1D(realdata=realdata) + elif self.dimensions == 2: + return self.unroll_2D(realdata=realdata) + else: + # Raise error for 3D variable + raise NotImplementedError( + "Unsupported data dimension: {self.dimensions}" + ) + def initialise_0D(self): - entries = np.concatenate(self.base_variables_data, axis=0).flatten() - + entries = self.unroll_0D() + if self.cumtrapz_ic is not None: entries = cumulative_trapezoid( entries, self.t_pts, initial=float(self.cumtrapz_ic) @@ -121,31 +192,9 @@ def initialise_0D(self): self.entries = entries self.dimensions = 0 - def _unroll_nnz(self): - # unroll in nnz != numel, otherwise copy - sp = self.base_variables_casadi[0](0,0,0).sparsity() - if sp.nnz() != sp.numel(): - data = [None]*len(self.base_variables_data) - for datak in range(len(self.base_variables_data)): - data[datak] = np.zeros(self.base_eval_shape[0]*len(self.t_pts)) - var_data = self.base_variables_data[0].flatten() - k = 0 - for t_i in range(len(self.t_pts)): - base = t_i*sp.numel() - for r in sp.row(): - data[datak][base + r] = var_data[k] - k = k + 1 - else: - data = self.base_variables_data - return data - def initialise_1D(self, fixed_t=False): len_space = self.base_eval_shape[0] - # reshape - entries = np.concatenate(self._unroll_nnz(), axis=0)\ - .flatten()\ - .reshape((len(self.t_pts),len_space))\ - .transpose() + entries = self.unroll_1D() # Get node and edge values nodes = self.mesh.nodes @@ -219,11 +268,12 @@ def initialise_2D(self): second_dim_size = len(second_dim_pts) len_space = self.base_eval_shape[0] - entries = np.concatenate(self._unroll_nnz(), axis=0)\ - .flatten()\ - .reshape((len(self.t_pts), second_dim_size, first_dim_size)) - entries = np.moveaxis(entries, 0, 2) - entries = np.moveaxis(entries, 0, 1) + entries = self.unroll_2D( + realdata=None, + n_dim1=second_dim_size, + n_dim2=first_dim_size, + axis_swaps=[(0, 2), (0, 1)], + ) # add points outside first dimension domain for extrapolation to # boundaries @@ -340,10 +390,12 @@ def initialise_2D_scikit_fem(self): len_y = len(y_sol) z_sol = self.mesh.edges["z"] len_z = len(z_sol) - entries = np.concatenate(self._unroll_nnz(), axis=0)\ - .flatten()\ - .reshape((len(self.t_pts), len_y, len_z)) - entries = np.moveaxis(entries, 0, 2) + entries = self.unroll_2D( + realdata=None, + n_dim1=len_y, + n_dim2=len_z, + axis_swaps=[(0, 2)], + ) # assign attributes for reference self.entries = entries diff --git a/test.py b/test.py index beb612d2de..6c17f22bb8 100644 --- a/test.py +++ b/test.py @@ -12,7 +12,7 @@ "Gradient of negative electrolyte potential [V.m-1]", "Negative particle flux [mol.m-2.s-1]", ] -output_variables = [] +#output_variables = [] all_vars = False input_parameters = { @@ -31,8 +31,7 @@ param.update({key: "[input]" for key in input_parameters}) param.process_model(model) param.process_geometry(geometry) -n = 100 # control the complexity of the geometry (increases number of solver states) -var_pts = {"x_n": n, "x_s": n, "x_p": n, "r_n": 10, "r_p": 10} +var_pts = {"x_n": 100, "x_s": 100, "x_p": 100, "r_n": 10, "r_p": 10} mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) disc = pybamm.Discretisation(mesh, model.default_spatial_methods) disc.process_model(model) From c35fbae8b56f58bd79ed0e14c10411b968ca0655 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 21 Jul 2023 10:16:10 +0000 Subject: [PATCH 11/38] Ensure further compatibility with existing tests --- .../idaklu/casadi_sundials_functions.cpp | 3 -- pybamm/solvers/idaklu_solver.py | 17 ++++---- pybamm/solvers/processed_variable_var.py | 5 ++- tests/unit/test_solvers/test_idaklu_solver.py | 41 ++++++++++--------- 4 files changed, 34 insertions(+), 32 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index 5e34bb97a9..0ee97dddce 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -177,9 +177,6 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, if (p_python_functions->options.using_banded_matrix) { - if (SUNSparseMatrix_SparseType(JJ) != CSC_MAT) - throw std::runtime_error("Banded matrix only tested with CSC format"); - // 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(); diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index a6881c17e5..6eddf00a8c 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -24,7 +24,7 @@ def have_idaklu(): def wrangle_name(name: str) -> str: return ( - name.casefold() + "v_" + name.casefold() .replace(" ", "_") .replace("[", "") .replace("]", "") @@ -204,7 +204,7 @@ def inputs_to_dict(inputs): if model.convert_to_format != "casadi": y0S = None if self.output_variables: - raise SolverError( + raise pybamm.SolverError( "output_variables can only be specified " 'with convert_to_format="casadi"' ) @@ -305,7 +305,7 @@ def resfn(t, y, inputs, ydot): idaklu.generate_function(self.var_casadi_fcns[key].serialize()) ) # Generate derivative functions for sensitivities - if len(inputs) > 0: + if (len(inputs) > 0) and (model.calculate_sensitivities): dvar_dy = casadi.jacobian(var_casadi, y_casadi) dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) dvar_dy_fcn = casadi.Function( @@ -693,11 +693,12 @@ def _integrate(self, model, t_eval, inputs_dict=None): ) # Add sensitivities newsol[var]._sensitivities = {} - for paramk, param in enumerate(inputs_dict.keys()): - newsol[var].add_sensitivity( - param, - [sol.yS[:, startk : (startk + len_of_var), paramk]] - ) + if model.calculate_sensitivities: + for paramk, param in enumerate(inputs_dict.keys()): + newsol[var].add_sensitivity( + param, + [sol.yS[:, startk : (startk + len_of_var), paramk]] + ) startk += len_of_var return newsol else: diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_var.py index 740470ed8f..5b5cfc1a9c 100644 --- a/pybamm/solvers/processed_variable_var.py +++ b/pybamm/solvers/processed_variable_var.py @@ -108,7 +108,10 @@ def __init__( ) def add_sensitivity(self, param, data): - self._sensitivities[param] = self.unroll(data) + # unroll from sparse representation into n-d matrix + # Note: then flatten and convert to casadi.DM for consistency with + # full state-vector ProcessedVariable sensitivities + self._sensitivities[param] = casadi.DM(self.unroll(data).flatten()) def _unroll_nnz(self, realdata=None): # unroll in nnz != numel, otherwise copy diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index efd0439f32..12fcf81843 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -182,31 +182,32 @@ def test_input_params(self): np.testing.assert_array_almost_equal(sol.y[1:3], true_solution) def test_sensitivites_initial_condition(self): - model = pybamm.BaseModel() - model.convert_to_format = "casadi" - u = pybamm.Variable("u") - v = pybamm.Variable("v") - a = pybamm.InputParameter("a") - model.rhs = {u: -u} - model.algebraic = {v: a * u - v} - model.initial_conditions = {u: 1, v: 1} - model.variables = {"2v": 2 * v} + for output_variables in [[], ["2v"]]: + model = pybamm.BaseModel() + model.convert_to_format = "casadi" + u = pybamm.Variable("u") + v = pybamm.Variable("v") + a = pybamm.InputParameter("a") + model.rhs = {u: -u} + model.algebraic = {v: a * u - v} + model.initial_conditions = {u: 1, v: 1} + model.variables = {"2v": 2 * v} - disc = pybamm.Discretisation() - disc.process_model(model) + disc = pybamm.Discretisation() + disc.process_model(model) - solver = pybamm.IDAKLUSolver() + solver = pybamm.IDAKLUSolver(output_variables=output_variables) - t_eval = np.linspace(0, 3, 100) - a_value = 0.1 + t_eval = np.linspace(0, 3, 100) + a_value = 0.1 - sol = solver.solve( - model, t_eval, inputs={"a": a_value}, calculate_sensitivities=True - ) + sol = solver.solve( + model, t_eval, inputs={"a": a_value}, calculate_sensitivities=True + ) - np.testing.assert_array_almost_equal( - sol["2v"].sensitivities["a"].full().flatten(), np.exp(-sol.t) * 2, decimal=4 - ) + np.testing.assert_array_almost_equal( + sol["2v"].sensitivities["a"].full().flatten(), np.exp(-sol.t) * 2, decimal=4 + ) def test_ida_roberts_klu_sensitivities(self): # this test implements a python version of the ida Roberts From 80d845ec96fe29bd36941f20771b843811a2e824 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 21 Jul 2023 11:06:59 +0000 Subject: [PATCH 12/38] Add unit tests for ProcessedVariableVar and modify Solution class tests to include output_variables where relevant --- pybamm/solvers/processed_variable_var.py | 4 +- .../test_processed_variable_var.py | 283 ++++++++++++++++++ 2 files changed, 284 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_solvers/test_processed_variable_var.py diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_var.py index 5b5cfc1a9c..57e1614865 100644 --- a/pybamm/solvers/processed_variable_var.py +++ b/pybamm/solvers/processed_variable_var.py @@ -142,7 +142,6 @@ def unroll_1D(self, realdata=None): len_space = self.base_eval_shape[0] return ( np.concatenate(self._unroll_nnz(realdata), axis=0) - .flatten() .reshape((len(self.t_pts), len_space)) .transpose() ) @@ -151,7 +150,7 @@ def unroll_2D(self, realdata=None, n_dim1=None, n_dim2=None, axis_swaps=[]): # initialise settings on first run if not self.unroll_params: self.unroll_params["n_dim1"] = n_dim1 - self.unroll_params["n_dim2"]= n_dim2 + self.unroll_params["n_dim2"]= n_dim2 self.unroll_params["axis_swaps"] = axis_swaps # use stored settings on subsequent runs if not n_dim1: @@ -161,7 +160,6 @@ def unroll_2D(self, realdata=None, n_dim1=None, n_dim2=None, axis_swaps=[]): len_space = self.base_eval_shape[0] entries = ( np.concatenate(self._unroll_nnz(realdata), axis=0) - .flatten() .reshape((len(self.t_pts), n_dim1, n_dim2)) ) for a, b in axis_swaps: diff --git a/tests/unit/test_solvers/test_processed_variable_var.py b/tests/unit/test_solvers/test_processed_variable_var.py new file mode 100644 index 0000000000..03da9ac808 --- /dev/null +++ b/tests/unit/test_solvers/test_processed_variable_var.py @@ -0,0 +1,283 @@ +# +# Tests for the Processed Variable Var class +# +# This class forms a container for variables (and sensitivities) calculted +# by the idaklu solver, and does not possesses any capability to calculate +# values itself since it does not have access to the full state vector +# +from tests import TestCase +import casadi +import pybamm +import tests + +import numpy as np +import unittest + + +def to_casadi(var_pybamm, y, inputs=None): + t_MX = casadi.MX.sym("t") + y_MX = casadi.MX.sym("y", y.shape[0]) + + inputs_MX_dict = {} + inputs = inputs or {} + for key, value in inputs.items(): + inputs_MX_dict[key] = casadi.MX.sym("input", value.shape[0]) + + inputs_MX = casadi.vertcat(*[p for p in inputs_MX_dict.values()]) + + var_sym = var_pybamm.to_casadi(t_MX, y_MX, inputs=inputs_MX_dict) + + var_casadi = casadi.Function("variable", [t_MX, y_MX, inputs_MX], [var_sym]) + return var_casadi + + +def process_and_check_2D_variable( + var, first_spatial_var, second_spatial_var, disc=None +): + # first_spatial_var should be on the "smaller" domain, i.e "r" for an "r-x" variable + if disc is None: + disc = tests.get_discretisation_for_testing() + disc.set_variable_slices([var]) + + first_sol = disc.process_symbol(first_spatial_var).entries[:, 0] + second_sol = disc.process_symbol(second_spatial_var).entries[:, 0] + + # Keep only the first iteration of entries + first_sol = first_sol[: len(first_sol) // len(second_sol)] + var_sol = disc.process_symbol(var) + t_sol = np.linspace(0, 1) + y_sol = np.ones(len(second_sol) * len(first_sol))[:, np.newaxis] * np.linspace(0, 5) + + var_casadi = to_casadi(var_sol, y_sol) + processed_var = pybamm.ProcessedVariableVar( + [var_sol], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + warn=False, + ) + # Ordering from idaklu with output_variables set is different to + # the full solver + y_sol = y_sol.reshape((y_sol.shape[1], y_sol.shape[0])).transpose() + np.testing.assert_array_equal( + processed_var.entries, + np.reshape(y_sol, [len(first_sol), len(second_sol), len(t_sol)]), + ) + return y_sol, first_sol, second_sol, t_sol + + +class TestProcessedVariableVar(TestCase): + def test_processed_variable_0D(self): + # without space + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + var = y + var.mesh = None + t_sol = np.array([0]) + y_sol = np.array([1])[:, np.newaxis] + var_casadi = to_casadi(var, y_sol) + processed_var = pybamm.ProcessedVariableVar( + [var], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + warn=False, + ) + np.testing.assert_array_equal(processed_var.entries, y_sol[0]) + + # check empty sensitivity works + + def test_processed_variable_0D_no_sensitivity(self): + # without space + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + var = t * y + var.mesh = None + t_sol = np.linspace(0, 1) + y_sol = np.array([np.linspace(0, 5)]) + var_casadi = to_casadi(var, y_sol) + processed_var = pybamm.ProcessedVariableVar( + [var], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + warn=False, + ) + + # test no inputs (i.e. no sensitivity) + self.assertDictEqual(processed_var.sensitivities, {}) + + # with parameter + t = pybamm.t + y = pybamm.StateVector(slice(0, 1)) + a = pybamm.InputParameter("a") + var = t * y * a + var.mesh = None + t_sol = np.linspace(0, 1) + y_sol = np.array([np.linspace(0, 5)]) + inputs = {"a": np.array([1.0])} + var_casadi = to_casadi(var, y_sol, inputs=inputs) + processed_var = pybamm.ProcessedVariableVar( + [var], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), inputs), + warn=False, + ) + + # test no sensitivity raises error + assert processed_var.sensitivities is None + + def test_processed_variable_1D(self): + t = pybamm.t + var = pybamm.Variable("var", domain=["negative electrode", "separator"]) + x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"]) + eqn = t * var + x + + # On nodes + disc = tests.get_discretisation_for_testing() + disc.set_variable_slices([var]) + x_sol = disc.process_symbol(x).entries[:, 0] + var_sol = disc.process_symbol(var) + eqn_sol = disc.process_symbol(eqn) + t_sol = np.linspace(0, 1) + y_sol = np.ones_like(x_sol)[:, np.newaxis] * np.linspace(0, 5) + + var_casadi = to_casadi(var_sol, y_sol) + sol = pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}) + processed_var = pybamm.ProcessedVariableVar( + [var_sol], + [var_casadi], + [y_sol], + sol, + warn=False, + ) + + # Ordering from idaklu with output_variables set is different to + # the full solver + y_sol = y_sol.reshape((y_sol.shape[1], y_sol.shape[0])).transpose() + np.testing.assert_array_equal(processed_var.entries, y_sol) + np.testing.assert_array_almost_equal(processed_var(t_sol, x_sol), y_sol) + eqn_casadi = to_casadi(eqn_sol, y_sol) + + + def test_processed_variable_1D_unknown_domain(self): + x = pybamm.SpatialVariable("x", domain="SEI layer", coord_sys="cartesian") + geometry = pybamm.Geometry( + {"SEI layer": {x: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}}} + ) + + submesh_types = {"SEI layer": pybamm.Uniform1DSubMesh} + var_pts = {x: 100} + mesh = pybamm.Mesh(geometry, submesh_types, var_pts) + + nt = 100 + + y_sol = np.zeros((var_pts[x], nt)) + solution = pybamm.Solution( + np.linspace(0, 1, nt), + y_sol, + pybamm.BaseModel(), + {}, + np.linspace(0, 1, 1), + np.zeros((var_pts[x])), + "test", + ) + + c = pybamm.StateVector(slice(0, var_pts[x]), domain=["SEI layer"]) + c.mesh = mesh["SEI layer"] + c_casadi = to_casadi(c, y_sol) + pybamm.ProcessedVariableVar([c], [c_casadi], [y_sol], solution, warn=False) + + def test_processed_variable_2D_space_only(self): + var = pybamm.Variable( + "var", + domain=["negative particle"], + auxiliary_domains={"secondary": ["negative electrode"]}, + ) + x = pybamm.SpatialVariable("x", domain=["negative electrode"]) + r = pybamm.SpatialVariable( + "r", + domain=["negative particle"], + auxiliary_domains={"secondary": ["negative electrode"]}, + ) + + disc = tests.get_p2d_discretisation_for_testing() + disc.set_variable_slices([var]) + x_sol = disc.process_symbol(x).entries[:, 0] + r_sol = disc.process_symbol(r).entries[:, 0] + # Keep only the first iteration of entries + r_sol = r_sol[: len(r_sol) // len(x_sol)] + var_sol = disc.process_symbol(var) + t_sol = np.array([0]) + y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] + + var_casadi = to_casadi(var_sol, y_sol) + processed_var = pybamm.ProcessedVariableVar( + [var_sol], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + warn=False, + ) + np.testing.assert_array_equal( + processed_var.entries, + np.reshape(y_sol, [len(r_sol), len(x_sol), len(t_sol)]), + ) + + def test_processed_variable_2D_fixed_t_scikit(self): + var = pybamm.Variable("var", domain=["current collector"]) + + disc = tests.get_2p1d_discretisation_for_testing() + disc.set_variable_slices([var]) + y = disc.mesh["current collector"].edges["y"] + z = disc.mesh["current collector"].edges["z"] + var_sol = disc.process_symbol(var) + var_sol.mesh = disc.mesh["current collector"] + t_sol = np.array([0]) + u_sol = np.ones(var_sol.shape[0])[:, np.newaxis] + + var_casadi = to_casadi(var_sol, u_sol) + processed_var = pybamm.ProcessedVariableVar( + [var_sol], + [var_casadi], + [u_sol], + pybamm.Solution(t_sol, u_sol, pybamm.BaseModel(), {}), + warn=False, + ) + np.testing.assert_array_equal( + processed_var.entries, np.reshape(u_sol, [len(y), len(z), len(t_sol)]) + ) + + def test_3D_raises_error(self): + var = pybamm.Variable( + "var", + domain=["negative electrode"], + auxiliary_domains={"secondary": ["current collector"]}, + ) + + disc = tests.get_2p1d_discretisation_for_testing() + disc.set_variable_slices([var]) + var_sol = disc.process_symbol(var) + t_sol = np.array([0, 1, 2]) + u_sol = np.ones(var_sol.shape[0] * 3)[:, np.newaxis] + var_casadi = to_casadi(var_sol, u_sol) + + with self.assertRaisesRegex(NotImplementedError, "Shape not recognized"): + pybamm.ProcessedVariableVar( + [var_sol], + [var_casadi], + [u_sol], + pybamm.Solution(t_sol, u_sol, pybamm.BaseModel(), {}), + warn=False, + ) + + +if __name__ == "__main__": + print("Add -v for more debug output") + import sys + + if "-v" in sys.argv: + debug = True + pybamm.settings.debug_mode = True + unittest.main() From 3dfbc15d812d847c33b04a4ab5e5d63faf331313 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Tue, 1 Aug 2023 10:43:12 +0000 Subject: [PATCH 13/38] Account for ExplicitTimeIntegral variables --- pybamm/solvers/idaklu_solver.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 6eddf00a8c..a45755d49d 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -294,6 +294,9 @@ def resfn(t, y, inputs, ydot): self.dvar_dp_fcns = [] for key in self.output_variables: # variable functions + if isinstance(model.variables_and_events[key], + pybamm.ExplicitTimeIntegral): + continue fcn_name = wrangle_name(key) var_casadi = model.variables_and_events[key].to_casadi( t_casadi, y_casadi, inputs=p_casadi @@ -682,6 +685,8 @@ def _integrate(self, model, t_eval, inputs_dict=None): sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) startk = 0 for vark, var in enumerate(self.output_variables): + if isinstance(model.variables_and_events[var], pybamm.ExplicitTimeIntegral): + continue len_of_var = ( self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz() ) From f09ca6d03f7bd4c4e66e4e2f42f3a2d5481db1bb Mon Sep 17 00:00:00 2001 From: John Brittain Date: Tue, 1 Aug 2023 10:45:22 +0000 Subject: [PATCH 14/38] Remove developer quicktest code --- compile | 15 --------------- compile.py | 31 ------------------------------- test.py | 41 +++++++++++++++++++++++++++++++---------- 3 files changed, 31 insertions(+), 56 deletions(-) delete mode 100755 compile delete mode 100644 compile.py diff --git a/compile b/compile deleted file mode 100755 index 2fdcaf2586..0000000000 --- a/compile +++ /dev/null @@ -1,15 +0,0 @@ -#!/usr/bin/env bash - -set -eox pipefail - -source .nox/dev/bin/activate -#rm pybamm/solvers/idaklu.*.so -#rm -rf build/* -python compile.py -cd build -cmake --build . -cp idaklu* ../pybamm/solvers -cd .. - -export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/jsb/.local/lib -python -it test.py diff --git a/compile.py b/compile.py deleted file mode 100644 index edb86b3d52..0000000000 --- a/compile.py +++ /dev/null @@ -1,31 +0,0 @@ -import os -import sys -import subprocess - -# Folder containing hte file -cmake_list_dir = os.path.abspath(os.path.dirname(__file__)) - -build_dir = 'build' - -cmake_args = [ - "-DCMAKE_BUILD_TYPE=RELEASE", - "-DPYTHON_EXECUTABLE={}".format(sys.executable), - "-DUSE_PYTHON_CASADI=TRUE", -] -if True: - cmake_args.append( - "-DSuiteSparse_ROOT={}".format(os.path.abspath("/home/jsb/.local")) - ) -if True: - cmake_args.append( - "-DSUNDIALS_ROOT={}".format(os.path.abspath("/home/jsb/.local")) - ) - -build_env = os.environ -# build_env["vcpkg_root_dir"] = vcpkg_root_dir -# build_env["vcpkg_default_triplet"] = vcpkg_default_triplet -# build_env["vcpkg_feature_flags"] = vcpkg_feature_flags - -subprocess.run( - ["cmake", cmake_list_dir] + cmake_args, cwd=build_dir, env=build_env -) diff --git a/test.py b/test.py index 6c17f22bb8..e419eedc0b 100644 --- a/test.py +++ b/test.py @@ -11,13 +11,15 @@ "x [m]", "Gradient of negative electrolyte potential [V.m-1]", "Negative particle flux [mol.m-2.s-1]", + "Discharge capacity [A.h]", + "Throughput capacity [A.h]", ] #output_variables = [] -all_vars = False +all_vars = True input_parameters = { - "Current function [A]": 0.680616, - "Separator porosity": 1.0, +# "Current function [A]": 0.680616, +# "Separator porosity": 1.0, } idaklu_spec = importlib.util.find_spec("pybamm.solvers.idaklu") @@ -47,16 +49,23 @@ output_variables = [m for m, (k, v) in zip(model.variable_names(), model.variables.items()) if not isinstance(v, pybamm.ExplicitTimeIntegral)] - left_out = [m for m, (k, v) in - zip(model.variable_names(), model.variables.items()) - if isinstance(v, pybamm.ExplicitTimeIntegral)] - print("ExplicitTimeIntegral variables:") - print(left_out) + +solver_all = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, +) +sol_all = solver_all.solve( + model, + t_eval, + inputs=input_parameters, + calculate_sensitivities=True, +) +print(f"Solve time [all]: {sol_all.solve_time.value*1000} msecs") #var_names = [[]] -var_names = [output_variables] +#var_names = [output_variables] #var_names = [[m] for m in model.variable_names()] -#var_names = [[m] for m in output_variables] +var_names = [[m] for m in output_variables] for output_vars in var_names: output_variables = output_vars print("Output variables:") @@ -77,6 +86,18 @@ calculate_sensitivities=True, ) + # Compare output to sol_all + for v in output_variables: + print(f"Comparison variable: {v}:") + print(f" Inidividual size {sol[v].data.shape}") + print(f" All size {sol_all[v].data.shape}") + print(sol_all[v].data) + print(sol[v].data) + rms = np.sqrt(np.mean((sol_all[v].data - sol[v].data)**2)) + print(f" RMS: {rms}") + assert sol[v].data.shape == sol_all[v].data.shape + assert np.allclose(sol[v].data, sol_all[v].data) + print(f"Solve time: {sol.solve_time.value*1000} msecs") if True: From bc2e162402cc838892a2ba0f488a7821ae8cff5a Mon Sep 17 00:00:00 2001 From: John Brittain Date: Tue, 1 Aug 2023 10:57:53 +0000 Subject: [PATCH 15/38] Code formatting --- pybamm/solvers/idaklu_solver.py | 31 +++-- pybamm/solvers/processed_variable.py | 1 - pybamm/solvers/processed_variable_var.py | 19 +-- test.py | 121 ------------------ tests/unit/test_solvers/test_idaklu_solver.py | 4 +- .../test_processed_variable_var.py | 6 - 6 files changed, 24 insertions(+), 158 deletions(-) delete mode 100644 test.py diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index a45755d49d..81f17a8976 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -23,18 +23,14 @@ def have_idaklu(): def wrangle_name(name: str) -> str: - return ( - "v_" + name.casefold() - .replace(" ", "_") - .replace("[", "") - .replace("]", "") - .replace(".", "_") - .replace("-", "_") - .replace("(", "") - .replace(")", "") - .replace("%", "prc") - .replace(",", "") - .replace(".", "") + return "v_" + name.casefold().replace(" ", "_").replace("[", "").replace( + "]", "" + ).replace(".", "_").replace("-", "_").replace("(", "").replace(")", "").replace( + "%", "prc" + ).replace( + ",", "" + ).replace( + ".", "" ) @@ -294,8 +290,9 @@ def resfn(t, y, inputs, ydot): self.dvar_dp_fcns = [] for key in self.output_variables: # variable functions - if isinstance(model.variables_and_events[key], - pybamm.ExplicitTimeIntegral): + if isinstance( + model.variables_and_events[key], pybamm.ExplicitTimeIntegral + ): continue fcn_name = wrangle_name(key) var_casadi = model.variables_and_events[key].to_casadi( @@ -685,7 +682,9 @@ def _integrate(self, model, t_eval, inputs_dict=None): sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) startk = 0 for vark, var in enumerate(self.output_variables): - if isinstance(model.variables_and_events[var], pybamm.ExplicitTimeIntegral): + if isinstance( + model.variables_and_events[var], pybamm.ExplicitTimeIntegral + ): continue len_of_var = ( self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz() @@ -702,7 +701,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): for paramk, param in enumerate(inputs_dict.keys()): newsol[var].add_sensitivity( param, - [sol.yS[:, startk : (startk + len_of_var), paramk]] + [sol.yS[:, startk : (startk + len_of_var), paramk]], ) startk += len_of_var return newsol diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py index 5c14fd504b..f6a0d673da 100644 --- a/pybamm/solvers/processed_variable.py +++ b/pybamm/solvers/processed_variable.py @@ -2,7 +2,6 @@ # Processed Variable class # import casadi -import numbers import numpy as np import pybamm from scipy.integrate import cumulative_trapezoid diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_var.py index 57e1614865..e6cb1c5ea3 100644 --- a/pybamm/solvers/processed_variable_var.py +++ b/pybamm/solvers/processed_variable_var.py @@ -2,7 +2,6 @@ # Processed Variable class # import casadi -import numbers import numpy as np import pybamm from scipy.integrate import cumulative_trapezoid @@ -137,7 +136,7 @@ def unroll_0D(self, realdata=None): if realdata is None: realdata = self.base_variables_data return np.concatenate(realdata, axis=0).flatten() - + def unroll_1D(self, realdata=None): len_space = self.base_eval_shape[0] return ( @@ -150,22 +149,20 @@ def unroll_2D(self, realdata=None, n_dim1=None, n_dim2=None, axis_swaps=[]): # initialise settings on first run if not self.unroll_params: self.unroll_params["n_dim1"] = n_dim1 - self.unroll_params["n_dim2"]= n_dim2 + self.unroll_params["n_dim2"] = n_dim2 self.unroll_params["axis_swaps"] = axis_swaps # use stored settings on subsequent runs if not n_dim1: n_dim1 = self.unroll_params["n_dim1"] n_dim2 = self.unroll_params["n_dim2"] axis_swaps = self.unroll_params["axis_swaps"] - len_space = self.base_eval_shape[0] - entries = ( - np.concatenate(self._unroll_nnz(realdata), axis=0) - .reshape((len(self.t_pts), n_dim1, n_dim2)) + entries = np.concatenate(self._unroll_nnz(realdata), axis=0).reshape( + (len(self.t_pts), n_dim1, n_dim2) ) for a, b in axis_swaps: entries = np.moveaxis(entries, a, b) return entries - + def unroll(self, realdata=None): if self.dimensions == 0: return self.unroll_0D(realdata=realdata) @@ -175,9 +172,7 @@ def unroll(self, realdata=None): return self.unroll_2D(realdata=realdata) else: # Raise error for 3D variable - raise NotImplementedError( - "Unsupported data dimension: {self.dimensions}" - ) + raise NotImplementedError("Unsupported data dimension: {self.dimensions}") def initialise_0D(self): entries = self.unroll_0D() @@ -194,7 +189,6 @@ def initialise_0D(self): self.dimensions = 0 def initialise_1D(self, fixed_t=False): - len_space = self.base_eval_shape[0] entries = self.unroll_1D() # Get node and edge values @@ -268,7 +262,6 @@ def initialise_2D(self): first_dim_size = len(first_dim_pts) second_dim_size = len(second_dim_pts) - len_space = self.base_eval_shape[0] entries = self.unroll_2D( realdata=None, n_dim1=second_dim_size, diff --git a/test.py b/test.py deleted file mode 100644 index e419eedc0b..0000000000 --- a/test.py +++ /dev/null @@ -1,121 +0,0 @@ -import pybamm -import numpy as np -import matplotlib.pylab as plt -import importlib - -output_variables = [ - "Voltage [V]", - "Time [min]", - "Current [A]", - "r_n [m]", - "x [m]", - "Gradient of negative electrolyte potential [V.m-1]", - "Negative particle flux [mol.m-2.s-1]", - "Discharge capacity [A.h]", - "Throughput capacity [A.h]", -] -#output_variables = [] -all_vars = True - -input_parameters = { -# "Current function [A]": 0.680616, -# "Separator porosity": 1.0, -} - -idaklu_spec = importlib.util.find_spec("pybamm.solvers.idaklu") -idaklu = importlib.util.module_from_spec(idaklu_spec) -idaklu_spec.loader.exec_module(idaklu) - -# construct model -model = pybamm.lithium_ion.DFN() -geometry = model.default_geometry -param = model.default_parameter_values -param.update({key: "[input]" for key in input_parameters}) -param.process_model(model) -param.process_geometry(geometry) -var_pts = {"x_n": 100, "x_s": 100, "x_p": 100, "r_n": 10, "r_p": 10} -mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) -disc = pybamm.Discretisation(mesh, model.default_spatial_methods) -disc.process_model(model) -t_eval = np.linspace(0, 3600, 100) - -options = { - 'linear_solver': 'SUNLinSol_KLU', - 'jacobian': 'sparse', - 'num_threads': 4, -} - -if all_vars: - output_variables = [m for m, (k, v) in - zip(model.variable_names(), model.variables.items()) - if not isinstance(v, pybamm.ExplicitTimeIntegral)] - -solver_all = pybamm.IDAKLUSolver( - atol=1e-8, rtol=1e-8, - options=options, -) -sol_all = solver_all.solve( - model, - t_eval, - inputs=input_parameters, - calculate_sensitivities=True, -) -print(f"Solve time [all]: {sol_all.solve_time.value*1000} msecs") - -#var_names = [[]] -#var_names = [output_variables] -#var_names = [[m] for m in model.variable_names()] -var_names = [[m] for m in output_variables] -for output_vars in var_names: - output_variables = output_vars - print("Output variables:") - print(output_variables) - print("\nInput parameters:") - print(input_parameters) - - solver = pybamm.IDAKLUSolver( - atol=1e-8, rtol=1e-8, - options=options, - output_variables=output_variables, - ) - - sol = solver.solve( - model, - t_eval, - inputs=input_parameters, - calculate_sensitivities=True, - ) - - # Compare output to sol_all - for v in output_variables: - print(f"Comparison variable: {v}:") - print(f" Inidividual size {sol[v].data.shape}") - print(f" All size {sol_all[v].data.shape}") - print(sol_all[v].data) - print(sol[v].data) - rms = np.sqrt(np.mean((sol_all[v].data - sol[v].data)**2)) - print(f" RMS: {rms}") - assert sol[v].data.shape == sol_all[v].data.shape - assert np.allclose(sol[v].data, sol_all[v].data) - - print(f"Solve time: {sol.solve_time.value*1000} msecs") - -if True: - if not output_variables: - plot_allvars = True - output_variables = [ - "Voltage [V]", - "Time [min]", - "Current [A]", - ] - fig, axs = plt.subplots(len(output_variables), len(input_parameters)+1) - for k, var in enumerate(output_variables): - # Solution variables currently use different classes/calls - if not input_parameters: - axs[k].plot(t_eval, sol[var](t_eval)) - else: - axs[k,0].plot(t_eval, sol[var](t_eval)) - for paramk, param in enumerate(list(input_parameters.keys())): - axs[k,paramk+1].plot(t_eval, sol[var].sensitivities[param]) - plt.tight_layout() - plt.show() diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index 12fcf81843..9c9be64984 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -206,7 +206,9 @@ def test_sensitivites_initial_condition(self): ) np.testing.assert_array_almost_equal( - sol["2v"].sensitivities["a"].full().flatten(), np.exp(-sol.t) * 2, decimal=4 + sol["2v"].sensitivities["a"].full().flatten(), + np.exp(-sol.t) * 2, + decimal=4, ) def test_ida_roberts_klu_sensitivities(self): diff --git a/tests/unit/test_solvers/test_processed_variable_var.py b/tests/unit/test_solvers/test_processed_variable_var.py index 03da9ac808..65c5b7ec0f 100644 --- a/tests/unit/test_solvers/test_processed_variable_var.py +++ b/tests/unit/test_solvers/test_processed_variable_var.py @@ -69,7 +69,6 @@ def process_and_check_2D_variable( class TestProcessedVariableVar(TestCase): def test_processed_variable_0D(self): # without space - t = pybamm.t y = pybamm.StateVector(slice(0, 1)) var = y var.mesh = None @@ -129,17 +128,14 @@ def test_processed_variable_0D_no_sensitivity(self): assert processed_var.sensitivities is None def test_processed_variable_1D(self): - t = pybamm.t var = pybamm.Variable("var", domain=["negative electrode", "separator"]) x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"]) - eqn = t * var + x # On nodes disc = tests.get_discretisation_for_testing() disc.set_variable_slices([var]) x_sol = disc.process_symbol(x).entries[:, 0] var_sol = disc.process_symbol(var) - eqn_sol = disc.process_symbol(eqn) t_sol = np.linspace(0, 1) y_sol = np.ones_like(x_sol)[:, np.newaxis] * np.linspace(0, 5) @@ -158,8 +154,6 @@ def test_processed_variable_1D(self): y_sol = y_sol.reshape((y_sol.shape[1], y_sol.shape[0])).transpose() np.testing.assert_array_equal(processed_var.entries, y_sol) np.testing.assert_array_almost_equal(processed_var(t_sol, x_sol), y_sol) - eqn_casadi = to_casadi(eqn_sol, y_sol) - def test_processed_variable_1D_unknown_domain(self): x = pybamm.SpatialVariable("x", domain="SEI layer", coord_sys="cartesian") From aadc28714f2a11f65a32f7bf2d9f4543e9c83377 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Tue, 1 Aug 2023 12:10:06 +0100 Subject: [PATCH 16/38] Update changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 187a264223..2d6fcd15bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ ## Features +- Idaklu solver can be given a list of variables to calculate during the solve ([#3217](https://github.com/pybamm-team/PyBaMM/pull/3217)) - Enable multithreading in IDAKLU solver ([#2947](https://github.com/pybamm-team/PyBaMM/pull/2947)) - If a solution contains cycles and steps, the cycle number and step number are now saved when `solution.save_data()` is called ([#2931](https://github.com/pybamm-team/PyBaMM/pull/2931)) - Experiments can now be given a `start_time` to define when each step should be triggered ([#2616](https://github.com/pybamm-team/PyBaMM/pull/2616)) From 58cee45de04766e0f77963f527bcab393438fcf1 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 31 Aug 2023 11:04:05 +0100 Subject: [PATCH 17/38] Pre-commit tidy-up --- pybamm/solvers/c_solvers/idaklu.cpp | 2 +- pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp | 2 +- pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp | 8 ++++---- pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp | 2 +- pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp | 8 ++++---- pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp | 4 ++-- pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp | 4 ++-- .../c_solvers/idaklu/casadi_sundials_functions.cpp | 2 +- .../solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp | 4 ++-- 9 files changed, 18 insertions(+), 18 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index 720dcd3be2..2e5fccfb8e 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -59,7 +59,7 @@ PYBIND11_MODULE(idaklu, m) //py::bind_vector>(m, "VectorFunction"); //py::implicitly_convertible>(); - + m.def("create_casadi_solver", &create_casadi_solver, "Create a casadi idaklu solver object", py::arg("number_of_states"), diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp index f45e667d68..3967449b1f 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp @@ -4,6 +4,6 @@ /* * This is an abstract base class for the Idaklu solver */ - + CasadiSolver::CasadiSolver() {} CasadiSolver::~CasadiSolver() {} diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index e3e7117896..ac695051ce 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -10,7 +10,7 @@ * are intended to be overriden to support alternative solver * approaches, as needed. */ - + /* Skeleton workflow: https://sundials.readthedocs.io/en/latest/ida/Usage/index.html 1. (N/A) Initialize parallel or multi-threaded environment @@ -135,7 +135,7 @@ void CasadiSolverOpenMP::Initialize() { N_VConst(RCONST(0.0), yyS[is]); N_VConst(RCONST(0.0), ypS[is]); } - + // create Matrix objects SetMatrix(); @@ -210,7 +210,7 @@ CasadiSolverOpenMP::~CasadiSolverOpenMP() N_VDestroy(yy); N_VDestroy(yp); N_VDestroy(id); - + if (number_of_parameters > 0) { N_VDestroyVectorArray(yyS, number_of_parameters); @@ -367,7 +367,7 @@ Solution CasadiSolverOpenMP::solve( realtype *yS_return = new realtype[number_of_parameters * number_of_timesteps * length_of_return_vector]; - + res = new realtype[max_res_size]; res_dvar_dy = new realtype[max_res_dvar_dy]; res_dvar_dp = new realtype[max_res_dvar_dp]; diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index 829e983b42..6c5f882b7b 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -70,7 +70,7 @@ CasadiFunctions::CasadiFunctions( for (auto& var : dvar_dp_fcns) { this->dvar_dp_fcns.push_back(CasadiFunction(*var)); } - + // 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>(); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 1feac1f359..848a365842 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -14,7 +14,7 @@ void csc_csr(realtype f[], T1 c[], T1 r[], realtype nf[], T2 nc[], T2 nr[], int int rr[N]; for (int i=0; i var_casadi_fcns; std::vector dvar_dy_fcns; std::vector dvar_dp_fcns; - + std::vector jac_times_cjmass_rowvals; std::vector jac_times_cjmass_colptrs; std::vector inputs; - + Options options; - + realtype *get_tmp_state_vector(); realtype *get_tmp_sparse_jacobian_data(); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp index 7040d1398f..c1c06c0fed 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -29,7 +29,7 @@ CasadiSolver *create_casadi_solver( const std::vector dvar_dy_fcns, const std::vector dvar_dp_fcns, py::dict options -) { +) { auto options_cpp = Options(options); auto functions = std::make_unique( rhs_alg, @@ -54,7 +54,7 @@ CasadiSolver *create_casadi_solver( ); CasadiSolver *casadiSolver = nullptr; - + // Instantiate solver class if (options_cpp.linear_solver == "SUNLinSol_Dense") { diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index 023810d2a5..84e57ed050 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -10,9 +10,9 @@ CasadiSolver *create_casadi_solver( 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 int jac_times_cjmass_nnz, const int jac_bandwidth_lower, - const int jac_bandwidth_upper, + const int jac_bandwidth_upper, const Function &jac_action, const Function &mass_action, const Function &sens, diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index e39d0fce56..0c98ad9d26 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -222,7 +222,7 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, realtype newjac[SUNSparseMatrix_NNZ(JJ)]; sunindextype *jac_ptrs = SUNSparseMatrix_IndexPointers(JJ); sunindextype *jac_vals = SUNSparseMatrix_IndexValues(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(yy); diff --git a/pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp b/pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp index 19bc92dcdf..f4855b1bc4 100644 --- a/pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp +++ b/pybamm/solvers/c_solvers/idaklu/sundials_legacy_wrapper.hpp @@ -26,12 +26,12 @@ { return N_VNew_Serial(vec_length); } - + N_Vector N_VNew_OpenMP(sunindextype vec_length, SUNContext sunctx) { return N_VNew_OpenMP(vec_length); } - + N_Vector N_VNew_Cuda(sunindextype vec_length, SUNContext sunctx) { return N_VNew_Cuda(vec_length); From 2fb4e558962d9d0803225db532a8e8752098b468 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 31 Aug 2023 13:05:50 +0100 Subject: [PATCH 18/38] Improve testing / code-coverage --- pybamm/solvers/processed_variable_var.py | 33 ++++++------- tests/unit/test_solvers/test_idaklu_solver.py | 16 +++++++ .../test_processed_variable_var.py | 48 ++++++++++++++++++- 3 files changed, 80 insertions(+), 17 deletions(-) diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_var.py index e6cb1c5ea3..d318e88067 100644 --- a/pybamm/solvers/processed_variable_var.py +++ b/pybamm/solvers/processed_variable_var.py @@ -116,21 +116,22 @@ def _unroll_nnz(self, realdata=None): # unroll in nnz != numel, otherwise copy if realdata is None: realdata = self.base_variables_data - sp = self.base_variables_casadi[0](0, 0, 0).sparsity() - if sp.nnz() != sp.numel(): - data = [None] * len(realdata) - for datak in range(len(realdata)): - data[datak] = np.zeros(self.base_eval_shape[0] * len(self.t_pts)) - var_data = realdata[0].flatten() - k = 0 - for t_i in range(len(self.t_pts)): - base = t_i * sp.numel() - for r in sp.row(): - data[datak][base + r] = var_data[k] - k = k + 1 - else: - data = realdata - return data + # sp = self.base_variables_casadi[0](0, 0, 0).sparsity() + # if sp.nnz() != sp.numel(): + # data = [None] * len(realdata) + # for datak in range(len(realdata)): + # data[datak] = np.zeros(self.base_eval_shape[0] * len(self.t_pts)) + # var_data = realdata[0].flatten() + # k = 0 + # for t_i in range(len(self.t_pts)): + # base = t_i * sp.numel() + # for r in sp.row(): + # data[datak][base + r] = var_data[k] + # k = k + 1 + # else: + # data = realdata + # return data + return realdata def unroll_0D(self, realdata=None): if realdata is None: @@ -172,7 +173,7 @@ def unroll(self, realdata=None): return self.unroll_2D(realdata=realdata) else: # Raise error for 3D variable - raise NotImplementedError("Unsupported data dimension: {self.dimensions}") + raise NotImplementedError(f"Unsupported data dimension: {self.dimensions}") def initialise_0D(self): entries = self.unroll_0D() diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index 9c9be64984..1cde09c43e 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -76,6 +76,22 @@ def test_model_events(self): # create discretisation disc = pybamm.Discretisation() model_disc = disc.process_model(model, inplace=False) + # Invalid atol (dict) raises error, valid options are float or ndarray + self.assertRaises( + pybamm.SolverError, + pybamm.IDAKLUSolver( + rtol=1e-8, atol={'key': 'value'}, + root_method=root_method) + ) + # output_variables only valid with convert_to_format=='casadi' + if form == "python" or form == "jax": + self.assertRaises( + pybamm.SolverError, + pybamm.IDAKLUSolver( + rtol=1e-8, atol=1e-8, + output_variables=['var'], + root_method=root_method) + ) # Solve solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8, root_method=root_method) t_eval = np.linspace(0, 1, 100) diff --git a/tests/unit/test_solvers/test_processed_variable_var.py b/tests/unit/test_solvers/test_processed_variable_var.py index 65c5b7ec0f..dd9b2f712b 100644 --- a/tests/unit/test_solvers/test_processed_variable_var.py +++ b/tests/unit/test_solvers/test_processed_variable_var.py @@ -82,10 +82,19 @@ def test_processed_variable_0D(self): pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), warn=False, ) + # Assert that the processed variable is the same as the solution np.testing.assert_array_equal(processed_var.entries, y_sol[0]) + # Check that 'data' produces the same output as 'entries' + np.testing.assert_array_equal(processed_var.entries, processed_var.data) - # check empty sensitivity works + # Check unroll function + np.testing.assert_array_equal(processed_var.unroll(), y_sol[0]) + # Check cumtrapz workflow produces no errors + processed_var.cumtrapz_ic = 1 + processed_var.initialise_0D() + + # check empty sensitivity works def test_processed_variable_0D_no_sensitivity(self): # without space t = pybamm.t @@ -153,8 +162,33 @@ def test_processed_variable_1D(self): # the full solver y_sol = y_sol.reshape((y_sol.shape[1], y_sol.shape[0])).transpose() np.testing.assert_array_equal(processed_var.entries, y_sol) + np.testing.assert_array_equal(processed_var.entries, processed_var.data) np.testing.assert_array_almost_equal(processed_var(t_sol, x_sol), y_sol) + # Check unroll function + np.testing.assert_array_equal(processed_var.unroll(), y_sol) + + # Check no error when data dimension is transposed vs node/edge + processed_var.mesh.nodes, processed_var.mesh.edges = \ + processed_var.mesh.edges, processed_var.mesh.nodes + processed_var.initialise_1D() + processed_var.mesh.nodes, processed_var.mesh.edges = \ + processed_var.mesh.edges, processed_var.mesh.nodes + + # Check no errors with domain-specific attributes + # (see ProcessedVariableVar.initialise_2D() for details) + domain_list = [ + ["particle", "electrode"], + ["separator", "current collector"], + ["particle", "particle size"], + ["particle size", "electrode"], + ["particle size", "current collector"] + ] + for domain, secondary in domain_list: + processed_var.domain[0] = domain + processed_var.domains["secondary"] = [secondary] + processed_var.initialise_1D() + def test_processed_variable_1D_unknown_domain(self): x = pybamm.SpatialVariable("x", domain="SEI layer", coord_sys="cartesian") geometry = pybamm.Geometry( @@ -218,6 +252,18 @@ def test_processed_variable_2D_space_only(self): processed_var.entries, np.reshape(y_sol, [len(r_sol), len(x_sol), len(t_sol)]), ) + np.testing.assert_array_equal( + processed_var.entries, + processed_var.data, + ) + + # Check unroll function (2D) + np.testing.assert_array_equal(processed_var.unroll(), y_sol.reshape(10, 40, 1)) + + # Check unroll function (3D) + with self.assertRaises(NotImplementedError): + processed_var.dimensions = 3 + processed_var.unroll() def test_processed_variable_2D_fixed_t_scikit(self): var = pybamm.Variable("var", domain=["current collector"]) From 437a4880bf582e7a87cecd539683f961c32557ab Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 31 Aug 2023 13:27:47 +0100 Subject: [PATCH 19/38] Codacy improvements --- pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp | 5 +++-- pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp | 4 ++-- pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp | 8 ++++---- pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp | 6 +++--- pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp | 6 +++--- pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp | 6 +++--- tests/unit/test_solvers/test_processed_variable_var.py | 2 +- 7 files changed, 19 insertions(+), 18 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index ac695051ce..28818c48f7 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -60,6 +60,7 @@ CasadiSolverOpenMP::CasadiSolverOpenMP( { // Construction code moved to Initialize() which is called from the // (child) CasadiSolver_XXX class constructors. + ida_mem = nullptr; } void CasadiSolverOpenMP::AllocateVectors() { @@ -227,7 +228,7 @@ void CasadiSolverOpenMP::CalcVars( size_t t_i, realtype *tret, realtype *yval, - std::vector ySval, + const std::vector& ySval, realtype *yS_return, size_t *ySk ) { @@ -246,7 +247,7 @@ void CasadiSolverOpenMP::CalcVars( void CasadiSolverOpenMP::CalcVarsSensitivities( realtype *tret, realtype *yval, - std::vector ySval, + const std::vector& ySval, realtype *yS_return, size_t *ySk ) { diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp index e4d69a3d78..f23e4d4e85 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -60,13 +60,13 @@ class CasadiSolverOpenMP : public CasadiSolver size_t t_i, realtype *tret, realtype *yval, - std::vector ySval, + const std::vector& ySval, realtype *yS_return, size_t *ySk); void CalcVarsSensitivities( realtype *tret, realtype *yval, - std::vector ySval, + const std::vector& ySval, realtype *yS_return, size_t *ySk); Solution solve( diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index 6c5f882b7b..8b864d17c4 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -7,7 +7,7 @@ CasadiFunction::CasadiFunction(const Function &f) : m_func(f) size_t sz_iw; size_t sz_w; m_func.sz_work(sz_arg, sz_res, sz_iw, sz_w); - int nnz = (sz_res>0) ? m_func.nnz_out() : 0; + //int nnz = (sz_res>0) ? m_func.nnz_out() : 0; //std::cout << "name = "<< m_func.name() << " arg = " << sz_arg << " res = " // << sz_res << " iw = " << sz_iw << " w = " << sz_w << " nnz = " << nnz << // std::endl; @@ -46,9 +46,9 @@ CasadiFunctions::CasadiFunctions( 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 std::vector var_casadi_fcns, - const std::vector dvar_dy_fcns, - const std::vector dvar_dp_fcns, + const std::vector& var_casadi_fcns, + const std::vector& dvar_dy_fcns, + const std::vector& dvar_dp_fcns, 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), diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 848a365842..2052bf1ee3 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -74,9 +74,9 @@ class CasadiFunctions const int n_s, const int n_e, const int n_p, - const std::vector var_casadi_fcns, - const std::vector dvar_dy_fcns, - const std::vector dvar_dp_fcns, + const std::vector& var_casadi_fcns, + const std::vector& dvar_dy_fcns, + const std::vector& dvar_dp_fcns, const Options& options ); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp index c1c06c0fed..9fcfa06510 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -25,9 +25,9 @@ CasadiSolver *create_casadi_solver( np_array atol_np, double rel_tol, int inputs_length, - const std::vector var_casadi_fcns, - const std::vector dvar_dy_fcns, - const std::vector dvar_dp_fcns, + const std::vector& var_casadi_fcns, + const std::vector& dvar_dy_fcns, + const std::vector& dvar_dp_fcns, py::dict options ) { auto options_cpp = Options(options); diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index 84e57ed050..c358634c93 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -22,9 +22,9 @@ CasadiSolver *create_casadi_solver( np_array atol_np, double rel_tol, int inputs_length, - const std::vector var_casadi_fcns, - const std::vector dvar_dy_fcns, - const std::vector dvar_dp_fcns, + const std::vector& var_casadi_fcns, + const std::vector& dvar_dy_fcns, + const std::vector& dvar_dp_fcns, py::dict options ); diff --git a/tests/unit/test_solvers/test_processed_variable_var.py b/tests/unit/test_solvers/test_processed_variable_var.py index dd9b2f712b..a94cb0875c 100644 --- a/tests/unit/test_solvers/test_processed_variable_var.py +++ b/tests/unit/test_solvers/test_processed_variable_var.py @@ -134,7 +134,7 @@ def test_processed_variable_0D_no_sensitivity(self): ) # test no sensitivity raises error - assert processed_var.sensitivities is None + self.assertIsNone(processed_var.sensitivities) def test_processed_variable_1D(self): var = pybamm.Variable("var", domain=["negative electrode", "separator"]) From 74f06f68cbd3f2992c4366e63f585f068c51ba14 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 31 Aug 2023 13:46:12 +0100 Subject: [PATCH 20/38] Codacy improvements --- pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index 28818c48f7..15d7b60c02 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -56,11 +56,11 @@ CasadiSolverOpenMP::CasadiSolverOpenMP( jac_bandwidth_lower(jac_bandwidth_lower), jac_bandwidth_upper(jac_bandwidth_upper), functions(std::move(functions_arg)), + ida_mem(nullptr), options(options) { // Construction code moved to Initialize() which is called from the // (child) CasadiSolver_XXX class constructors. - ida_mem = nullptr; } void CasadiSolverOpenMP::AllocateVectors() { From 89222a5c099712a8dbe3e67a9e977e2a4fb4354b Mon Sep 17 00:00:00 2001 From: John Brittain Date: Thu, 31 Aug 2023 14:22:21 +0100 Subject: [PATCH 21/38] Additional tests for idaklu solver with output_variables --- pybamm/solvers/idaklu_solver.py | 2 +- pybamm/solvers/processed_variable_var.py | 31 ++-- tests/integration/test_solvers/test_idaklu.py | 72 ++++++++ tests/unit/test_solvers/test_idaklu_solver.py | 21 +-- .../test_processed_variable_var.py | 157 +++++++++++++++--- 5 files changed, 230 insertions(+), 53 deletions(-) diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index a025d56bc3..2abbe9ad96 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -678,7 +678,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): if isinstance( model.variables_and_events[var], pybamm.ExplicitTimeIntegral ): - continue + continue # pragma: no cover len_of_var = ( self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz() ) diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_var.py index d318e88067..11e39f99f1 100644 --- a/pybamm/solvers/processed_variable_var.py +++ b/pybamm/solvers/processed_variable_var.py @@ -116,22 +116,21 @@ def _unroll_nnz(self, realdata=None): # unroll in nnz != numel, otherwise copy if realdata is None: realdata = self.base_variables_data - # sp = self.base_variables_casadi[0](0, 0, 0).sparsity() - # if sp.nnz() != sp.numel(): - # data = [None] * len(realdata) - # for datak in range(len(realdata)): - # data[datak] = np.zeros(self.base_eval_shape[0] * len(self.t_pts)) - # var_data = realdata[0].flatten() - # k = 0 - # for t_i in range(len(self.t_pts)): - # base = t_i * sp.numel() - # for r in sp.row(): - # data[datak][base + r] = var_data[k] - # k = k + 1 - # else: - # data = realdata - # return data - return realdata + sp = self.base_variables_casadi[0](0, 0, 0).sparsity() + if sp.nnz() != sp.numel(): + data = [None] * len(realdata) + for datak in range(len(realdata)): + data[datak] = np.zeros(self.base_eval_shape[0] * len(self.t_pts)) + var_data = realdata[0].flatten() + k = 0 + for t_i in range(len(self.t_pts)): + base = t_i * sp.numel() + for r in sp.row(): + data[datak][base + r] = var_data[k] + k = k + 1 + else: + data = realdata + return data def unroll_0D(self, realdata=None): if realdata is None: diff --git a/tests/integration/test_solvers/test_idaklu.py b/tests/integration/test_solvers/test_idaklu.py index eecc5dfe2b..fedcc78d64 100644 --- a/tests/integration/test_solvers/test_idaklu.py +++ b/tests/integration/test_solvers/test_idaklu.py @@ -90,6 +90,78 @@ def test_changing_grid(self): # solve solver.solve(model_disc, t_eval) + def test_with_output_variables(self): + # Construct a model and solve for all vairables, then test + # the 'output_variables' option for each variable in turn, confirming + # equivalence + + # construct model + model = pybamm.lithium_ion.DFN() + geometry = model.default_geometry + param = model.default_parameter_values + input_parameters = {} + param.update({key: "[input]" for key in input_parameters}) + param.process_model(model) + param.process_geometry(geometry) + var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5} + mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + t_eval = np.linspace(0, 3600, 100) + + options = { + 'linear_solver': 'SUNLinSol_KLU', + 'jacobian': 'sparse', + 'num_threads': 4, + } + + # Select all output_variables (NB: This can be very slow to run all solves) + # output_variables = [m for m, (k, v) in + # zip(model.variable_names(), model.variables.items()) + # if not isinstance(v, pybamm.ExplicitTimeIntegral)] + + # Use a selection of variables (of different types) + output_variables = [ + "Voltage [V]", + "Time [min]", + "Current [A]", + "r_n [m]", + "x [m]", + "Gradient of negative electrolyte potential [V.m-1]", + "Negative particle flux [mol.m-2.s-1]", + "Discharge capacity [A.h]", + "Throughput capacity [A.h]", + ] + + solver_all = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + ) + sol_all = solver_all.solve( + model, + t_eval, + inputs=input_parameters, + calculate_sensitivities=True, + ) + + for varname in output_variables: + print(varname) + solver = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + output_variables=[varname], + ) + + sol = solver.solve( + model, + t_eval, + inputs=input_parameters, + calculate_sensitivities=True, + ) + + # Compare output to sol_all + self.assertTrue(np.allclose(sol[varname].data, sol_all[varname].data)) + if __name__ == "__main__": print("Add -v for more debug output") diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index 1cde09c43e..c6b6b52ddf 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -76,22 +76,6 @@ def test_model_events(self): # create discretisation disc = pybamm.Discretisation() model_disc = disc.process_model(model, inplace=False) - # Invalid atol (dict) raises error, valid options are float or ndarray - self.assertRaises( - pybamm.SolverError, - pybamm.IDAKLUSolver( - rtol=1e-8, atol={'key': 'value'}, - root_method=root_method) - ) - # output_variables only valid with convert_to_format=='casadi' - if form == "python" or form == "jax": - self.assertRaises( - pybamm.SolverError, - pybamm.IDAKLUSolver( - rtol=1e-8, atol=1e-8, - output_variables=['var'], - root_method=root_method) - ) # Solve solver = pybamm.IDAKLUSolver(rtol=1e-8, atol=1e-8, root_method=root_method) t_eval = np.linspace(0, 1, 100) @@ -101,6 +85,10 @@ def test_model_events(self): solution.y[0], np.exp(0.1 * solution.t), decimal=5 ) + # Check invalid atol type raises an error + with self.assertRaises(pybamm.SolverError): + solver._check_atol_type({'key': 'value'}, []) + # enforce events that won't be triggered model.events = [pybamm.Event("an event", var + 1)] model_disc = disc.process_model(model, inplace=False) @@ -211,7 +199,6 @@ def test_sensitivites_initial_condition(self): disc = pybamm.Discretisation() disc.process_model(model) - solver = pybamm.IDAKLUSolver(output_variables=output_variables) t_eval = np.linspace(0, 3, 100) diff --git a/tests/unit/test_solvers/test_processed_variable_var.py b/tests/unit/test_solvers/test_processed_variable_var.py index a94cb0875c..e5b62e8bee 100644 --- a/tests/unit/test_solvers/test_processed_variable_var.py +++ b/tests/unit/test_solvers/test_processed_variable_var.py @@ -32,7 +32,7 @@ def to_casadi(var_pybamm, y, inputs=None): def process_and_check_2D_variable( - var, first_spatial_var, second_spatial_var, disc=None + var, first_spatial_var, second_spatial_var, disc=None, geometry_options={} ): # first_spatial_var should be on the "smaller" domain, i.e "r" for an "r-x" variable if disc is None: @@ -49,20 +49,20 @@ def process_and_check_2D_variable( y_sol = np.ones(len(second_sol) * len(first_sol))[:, np.newaxis] * np.linspace(0, 5) var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariableVar( + model = tests.get_base_model_with_battery_geometry(**geometry_options) + pybamm.ProcessedVariableVar( [var_sol], [var_casadi], [y_sol], - pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + pybamm.Solution(t_sol, y_sol, model, {}), warn=False, ) - # Ordering from idaklu with output_variables set is different to - # the full solver - y_sol = y_sol.reshape((y_sol.shape[1], y_sol.shape[0])).transpose() - np.testing.assert_array_equal( - processed_var.entries, - np.reshape(y_sol, [len(first_sol), len(second_sol), len(t_sol)]), - ) + # NB: ProcessedVariableVar does not interpret y in the same way as + # ProcessedVariable; a better test of equivalence is to check that the + # results are the same between IDAKLUSolver with (and without) + # output_variables. This is implemented in the integration test: + # tests/integration/test_solvers/test_idaklu_solver.py + # ::test_output_variables return y_sol, first_sol, second_sol, t_sol @@ -175,18 +175,17 @@ def test_processed_variable_1D(self): processed_var.mesh.nodes, processed_var.mesh.edges = \ processed_var.mesh.edges, processed_var.mesh.nodes - # Check no errors with domain-specific attributes - # (see ProcessedVariableVar.initialise_2D() for details) + # Check that there are no errors with domain-specific attributes + # (see ProcessedVariableVar.initialise_1D() for details) domain_list = [ - ["particle", "electrode"], - ["separator", "current collector"], - ["particle", "particle size"], - ["particle size", "electrode"], - ["particle size", "current collector"] + "particle", + "separator", + "current collector", + "particle size", + "random-non-specific-domain", ] - for domain, secondary in domain_list: + for domain in domain_list: processed_var.domain[0] = domain - processed_var.domains["secondary"] = [secondary] processed_var.initialise_1D() def test_processed_variable_1D_unknown_domain(self): @@ -217,6 +216,126 @@ def test_processed_variable_1D_unknown_domain(self): c_casadi = to_casadi(c, y_sol) pybamm.ProcessedVariableVar([c], [c_casadi], [y_sol], solution, warn=False) + def test_processed_variable_2D_x_r(self): + var = pybamm.Variable( + "var", + domain=["negative particle"], + auxiliary_domains={"secondary": ["negative electrode"]}, + ) + x = pybamm.SpatialVariable("x", domain=["negative electrode"]) + r = pybamm.SpatialVariable( + "r", + domain=["negative particle"], + auxiliary_domains={"secondary": ["negative electrode"]}, + ) + + disc = tests.get_p2d_discretisation_for_testing() + process_and_check_2D_variable(var, r, x, disc=disc) + + def test_processed_variable_2D_R_x(self): + var = pybamm.Variable( + "var", + domain=["negative particle size"], + auxiliary_domains={"secondary": ["negative electrode"]}, + ) + R = pybamm.SpatialVariable( + "R", + domain=["negative particle size"], + auxiliary_domains={"secondary": ["negative electrode"]}, + ) + x = pybamm.SpatialVariable("x", domain=["negative electrode"]) + + disc = tests.get_size_distribution_disc_for_testing() + process_and_check_2D_variable( + var, + R, + x, + disc=disc, + geometry_options={"options": {"particle size": "distribution"}}, + ) + + def test_processed_variable_2D_R_z(self): + var = pybamm.Variable( + "var", + domain=["negative particle size"], + auxiliary_domains={"secondary": ["current collector"]}, + ) + R = pybamm.SpatialVariable( + "R", + domain=["negative particle size"], + auxiliary_domains={"secondary": ["current collector"]}, + ) + z = pybamm.SpatialVariable("z", domain=["current collector"]) + + disc = tests.get_size_distribution_disc_for_testing() + process_and_check_2D_variable( + var, + R, + z, + disc=disc, + geometry_options={"options": {"particle size": "distribution"}}, + ) + + def test_processed_variable_2D_r_R(self): + var = pybamm.Variable( + "var", + domain=["negative particle"], + auxiliary_domains={"secondary": ["negative particle size"]}, + ) + r = pybamm.SpatialVariable( + "r", + domain=["negative particle"], + auxiliary_domains={"secondary": ["negative particle size"]}, + ) + R = pybamm.SpatialVariable("R", domain=["negative particle size"]) + + disc = tests.get_size_distribution_disc_for_testing() + process_and_check_2D_variable( + var, + r, + R, + disc=disc, + geometry_options={"options": {"particle size": "distribution"}}, + ) + + def test_processed_variable_2D_x_z(self): + var = pybamm.Variable( + "var", + domain=["negative electrode", "separator"], + auxiliary_domains={"secondary": "current collector"}, + ) + x = pybamm.SpatialVariable( + "x", + domain=["negative electrode", "separator"], + auxiliary_domains={"secondary": "current collector"}, + ) + z = pybamm.SpatialVariable("z", domain=["current collector"]) + + disc = tests.get_1p1d_discretisation_for_testing() + y_sol, x_sol, z_sol, t_sol = process_and_check_2D_variable(var, x, z, disc=disc) + del x_sol + + # On edges + x_s_edge = pybamm.Matrix( + np.tile(disc.mesh["separator"].edges, len(z_sol)), + domain="separator", + auxiliary_domains={"secondary": "current collector"}, + ) + x_s_edge.mesh = disc.mesh["separator"] + x_s_edge.secondary_mesh = disc.mesh["current collector"] + x_s_casadi = to_casadi(x_s_edge, y_sol) + processed_x_s_edge = pybamm.ProcessedVariable( + [x_s_edge], + [x_s_casadi], + pybamm.Solution( + t_sol, y_sol, tests.get_base_model_with_battery_geometry(), {} + ), + warn=False, + ) + np.testing.assert_array_equal( + x_s_edge.entries.flatten(), processed_x_s_edge.entries[:, :, 0].T.flatten() + ) + def test_processed_variable_2D_space_only(self): var = pybamm.Variable( "var", From d028668f51a43ac0ca7fe07c8d8c8e4a8c30b659 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 1 Sep 2023 10:22:10 +0100 Subject: [PATCH 22/38] Improve test coverage --- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 1 - .../c_solvers/idaklu/CasadiSolverOpenMP.hpp | 2 +- pybamm/solvers/idaklu_solver.py | 4 +- tests/integration/test_solvers/test_idaklu.py | 72 ------------------- tests/unit/test_solvers/test_idaklu_solver.py | 71 ++++++++++++++++++ 5 files changed, 74 insertions(+), 76 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index 15d7b60c02..93e90c3879 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -56,7 +56,6 @@ CasadiSolverOpenMP::CasadiSolverOpenMP( jac_bandwidth_lower(jac_bandwidth_lower), jac_bandwidth_upper(jac_bandwidth_upper), functions(std::move(functions_arg)), - ida_mem(nullptr), options(options) { // Construction code moved to Initialize() which is called from the diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp index f23e4d4e85..89914e26cd 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -14,7 +14,7 @@ using Function = casadi::Function; class CasadiSolverOpenMP : public CasadiSolver { public: - void *ida_mem; // pointer to memory + void *ida_mem = nullptr; // pointer to memory np_array atol_np; double rel_tol; np_array rhs_alg_id; diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 2abbe9ad96..eec279c968 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -196,7 +196,7 @@ def inputs_to_dict(inputs): raise pybamm.SolverError( "output_variables can only be specified " 'with convert_to_format="casadi"' - ) + ) # pragma: no cover if y0S is not None: if isinstance(y0S, casadi.DM): y0S = (y0S,) @@ -678,7 +678,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): if isinstance( model.variables_and_events[var], pybamm.ExplicitTimeIntegral ): - continue # pragma: no cover + continue len_of_var = ( self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz() ) diff --git a/tests/integration/test_solvers/test_idaklu.py b/tests/integration/test_solvers/test_idaklu.py index fedcc78d64..eecc5dfe2b 100644 --- a/tests/integration/test_solvers/test_idaklu.py +++ b/tests/integration/test_solvers/test_idaklu.py @@ -90,78 +90,6 @@ def test_changing_grid(self): # solve solver.solve(model_disc, t_eval) - def test_with_output_variables(self): - # Construct a model and solve for all vairables, then test - # the 'output_variables' option for each variable in turn, confirming - # equivalence - - # construct model - model = pybamm.lithium_ion.DFN() - geometry = model.default_geometry - param = model.default_parameter_values - input_parameters = {} - param.update({key: "[input]" for key in input_parameters}) - param.process_model(model) - param.process_geometry(geometry) - var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5} - mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) - disc = pybamm.Discretisation(mesh, model.default_spatial_methods) - disc.process_model(model) - t_eval = np.linspace(0, 3600, 100) - - options = { - 'linear_solver': 'SUNLinSol_KLU', - 'jacobian': 'sparse', - 'num_threads': 4, - } - - # Select all output_variables (NB: This can be very slow to run all solves) - # output_variables = [m for m, (k, v) in - # zip(model.variable_names(), model.variables.items()) - # if not isinstance(v, pybamm.ExplicitTimeIntegral)] - - # Use a selection of variables (of different types) - output_variables = [ - "Voltage [V]", - "Time [min]", - "Current [A]", - "r_n [m]", - "x [m]", - "Gradient of negative electrolyte potential [V.m-1]", - "Negative particle flux [mol.m-2.s-1]", - "Discharge capacity [A.h]", - "Throughput capacity [A.h]", - ] - - solver_all = pybamm.IDAKLUSolver( - atol=1e-8, rtol=1e-8, - options=options, - ) - sol_all = solver_all.solve( - model, - t_eval, - inputs=input_parameters, - calculate_sensitivities=True, - ) - - for varname in output_variables: - print(varname) - solver = pybamm.IDAKLUSolver( - atol=1e-8, rtol=1e-8, - options=options, - output_variables=[varname], - ) - - sol = solver.solve( - model, - t_eval, - inputs=input_parameters, - calculate_sensitivities=True, - ) - - # Compare output to sol_all - self.assertTrue(np.allclose(sol[varname].data, sol_all[varname].data)) - if __name__ == "__main__": print("Add -v for more debug output") diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index c6b6b52ddf..8024df7490 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -546,6 +546,77 @@ def test_options(self): with self.assertRaises(ValueError): soln = solver.solve(model, t_eval) + def test_with_output_variables(self): + # Construct a model and solve for all vairables, then test + # the 'output_variables' option for each variable in turn, confirming + # equivalence + + # construct model + model = pybamm.lithium_ion.DFN() + geometry = model.default_geometry + param = model.default_parameter_values + input_parameters = {} + param.update({key: "[input]" for key in input_parameters}) + param.process_model(model) + param.process_geometry(geometry) + var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5} + mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + t_eval = np.linspace(0, 3600, 100) + + options = { + 'linear_solver': 'SUNLinSol_KLU', + 'jacobian': 'sparse', + 'num_threads': 4, + } + + # Use a selection of variables of different types + output_variables = [ + "Voltage [V]", + "Time [min]", + "Current [A]", + "r_n [m]", + "x [m]", + "x_s [m]", + "Gradient of negative electrolyte potential [V.m-1]", + "Negative particle flux [mol.m-2.s-1]", + "Discharge capacity [A.h]", # ExplicitTimeIntegral + "Throughput capacity [A.h]", # ExplicitTimeIntegral + ] + + # Use the full model as comparison (tested separately) + solver_all = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + ) + sol_all = solver_all.solve( + model, + t_eval, + inputs=input_parameters, + calculate_sensitivities=True, + ) + + # Solve for a subset of variables and compare results + solver = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + output_variables=output_variables, + ) + sol = solver.solve( + model, + t_eval, + inputs=input_parameters, + ) + + # Compare output to sol_all + for varname in output_variables: + self.assertTrue(np.allclose(sol[varname].data, sol_all[varname].data)) + + # Mock a 1D current collector and initialise (none in the model) + sol["x_s [m]"].domain = ["current collector"] + sol["x_s [m]"].initialise_1D() + if __name__ == "__main__": print("Add -v for more debug output") From 86b9c01ed8789c5eaddf410d5d1ad0ffab3922f5 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 1 Sep 2023 13:11:09 +0100 Subject: [PATCH 23/38] Generalise idaklu matrix format conversions --- .../solvers/c_solvers/idaklu/casadi_sundials_functions.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index 0c98ad9d26..5b6a3b99f6 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -1,6 +1,7 @@ #include "casadi_sundials_functions.hpp" #include "casadi_functions.hpp" #include "common.hpp" +#include #define NV_DATA NV_DATA_OMP //#define NV_DATA NV_DATA_S @@ -233,7 +234,10 @@ int jacobian_casadi(realtype tt, realtype cj, N_Vector yy, N_Vector yp, p_python_functions->jac_times_cjmass(); // convert (casadi's) CSC format to CSR - csc_csr( + csc_csr< + std::remove_pointer_tjac_times_cjmass_rowvals.data())>, + std::remove_pointer_t + >( newjac, p_python_functions->jac_times_cjmass_rowvals.data(), p_python_functions->jac_times_cjmass_colptrs.data(), From 653dc1d40a7bf0692781af3feafa8ba6bc30e499 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 8 Sep 2023 09:19:22 +0100 Subject: [PATCH 24/38] Suggested changes --- pybamm/solvers/c_solvers/idaklu.cpp | 1 - pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp | 4 ++-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index 2e5fccfb8e..be90955b9c 100644 --- a/pybamm/solvers/c_solvers/idaklu.cpp +++ b/pybamm/solvers/c_solvers/idaklu.cpp @@ -18,7 +18,6 @@ Function generate_function(const std::string &data) namespace py = pybind11; PYBIND11_MAKE_OPAQUE(std::vector); -//PYBIND11_MAKE_OPAQUE(std::vector); PYBIND11_MODULE(idaklu, m) { diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 2052bf1ee3..1fd2ceda2d 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -48,9 +48,9 @@ class CasadiFunction public: std::vector m_arg; std::vector m_res; - -//private: const Function &m_func; + +private: std::vector m_iw; std::vector m_w; }; From f442fb181c8220e76f745c847732da057d3dd45b Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 11 Sep 2023 09:11:26 +0000 Subject: [PATCH 25/38] Improve documentation in idaklu solver --- .../solvers/c_solvers/idaklu/CasadiSolver.cpp | 5 -- .../solvers/c_solvers/idaklu/CasadiSolver.hpp | 23 +++++++ .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 31 --------- .../c_solvers/idaklu/CasadiSolverOpenMP.hpp | 68 ++++++++++++++++++- .../idaklu/CasadiSolverOpenMP_solvers.cpp | 4 -- .../idaklu/CasadiSolverOpenMP_solvers.hpp | 21 ++++++ .../c_solvers/idaklu/casadi_functions.hpp | 35 +++++++++- .../c_solvers/idaklu/casadi_solver.hpp | 5 ++ .../idaklu/casadi_sundials_functions.cpp | 3 +- pybamm/solvers/c_solvers/idaklu/options.hpp | 3 + pybamm/solvers/c_solvers/idaklu/python.hpp | 4 +- pybamm/solvers/c_solvers/idaklu/solution.hpp | 6 ++ 12 files changed, 162 insertions(+), 46 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp index 3967449b1f..cc876edce6 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp @@ -1,9 +1,4 @@ #include "CasadiSolver.hpp" -//#include "casadi_sundials_functions.hpp" - -/* - * This is an abstract base class for the Idaklu solver - */ CasadiSolver::CasadiSolver() {} CasadiSolver::~CasadiSolver() {} diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp index be2336c412..3703e98c99 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp @@ -10,16 +10,39 @@ using Function = casadi::Function; #include "solution.hpp" #include "sundials_legacy_wrapper.hpp" +/** + * Abstract base class for solutions that can use different solvers and vector + * implementations. + * @brief An abstract base class for the Idaklu solver + */ class CasadiSolver { public: + + /** + * @brief Default constructor + */ CasadiSolver(); + + /** + * @brief Default destructor + */ ~CasadiSolver(); + + /** + * @brief Abstract solver method that returns a Solution class + */ virtual Solution solve( np_array t_np, np_array y0_np, np_array yp0_np, np_array_dense inputs) = 0; + + /** + * Abstract method to initialize the solver, once vectors and solver classes + * are set + * @brief Abstract initialization method + */ virtual void Initialize() = 0; }; diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index 93e90c3879..fc4268a636 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -4,37 +4,6 @@ #include #include -/* - * This is an abstract class that implements an OpenMP solution but - * requires a linear solver to create a concrete class. Hook functions - * are intended to be overriden to support alternative solver - * approaches, as needed. - */ - -/* Skeleton workflow: - https://sundials.readthedocs.io/en/latest/ida/Usage/index.html - 1. (N/A) Initialize parallel or multi-threaded environment - 2. Create the SUNDIALS context object - 3. Create the vector of initial values - 4. Create matrix object (if appropriate) - 5. Create linear solver object - 6. (N/A) Create nonlinear solver object - 7. Create IDA object - 8. Initialize IDA solver - 9. Specify integration tolerances - 10. Attach the linear solver - 11. Set linear solver optional inputs - 12. (N/A) Attach nonlinear solver module - 13. (N/A) Set nonlinear solver optional inputs - 14. Specify rootfinding problem (optional) - 15. Set optional inputs - 16. Correct initial values (optional) - 17. Advance solution in time - 18. Get optional outputs - 19. Destroy objects - 20. (N/A) Finalize MPI -*/ - CasadiSolverOpenMP::CasadiSolverOpenMP( np_array atol_np, double rel_tol, diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp index 89914e26cd..260c03ded7 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -11,10 +11,39 @@ using Function = casadi::Function; #include "solution.hpp" #include "sundials_legacy_wrapper.hpp" +/** + * @brief Abstract solver class based on OpenMP vectors + * + * An abstract class that implements a solution based on OpenMP + * vectors but needs to be provided with a suitable linear solver. + * + * This class broadly implements the following skeleton workflow: + * (https://sundials.readthedocs.io/en/latest/ida/Usage/index.html) + * 1. (N/A) Initialize parallel or multi-threaded environment + * 2. Create the SUNDIALS context object + * 3. Create the vector of initial values + * 4. Create matrix object (if appropriate) + * 5. Create linear solver object + * 6. (N/A) Create nonlinear solver object + * 7. Create IDA object + * 8. Initialize IDA solver + * 9. Specify integration tolerances + * 10. Attach the linear solver + * 11. Set linear solver optional inputs + * 12. (N/A) Attach nonlinear solver module + * 13. (N/A) Set nonlinear solver optional inputs + * 14. Specify rootfinding problem (optional) + * 15. Set optional inputs + * 16. Correct initial values (optional) + * 17. Advance solution in time + * 18. Get optional outputs + * 19. Destroy objects + * 20. (N/A) Finalize MPI + */ class CasadiSolverOpenMP : public CasadiSolver { public: - void *ida_mem = nullptr; // pointer to memory + void *ida_mem = nullptr; np_array atol_np; double rel_tol; np_array rhs_alg_id; @@ -42,6 +71,9 @@ class CasadiSolverOpenMP : public CasadiSolver #endif public: + /** + * @brief Constructor + */ CasadiSolverOpenMP( np_array atol_np, double rel_tol, @@ -53,7 +85,17 @@ class CasadiSolverOpenMP : public CasadiSolver int jac_bandwidth_upper, std::unique_ptr functions, const Options& options); + + /** + * @brief Destructor + */ ~CasadiSolverOpenMP(); + + /** + * Evaluate casadi functions (including sensitivies) for each requested + * variable and store + * @brief Evaluate casadi functions + */ void CalcVars( realtype *y_return, size_t length_of_return_vector, @@ -63,20 +105,44 @@ class CasadiSolverOpenMP : public CasadiSolver const std::vector& ySval, realtype *yS_return, size_t *ySk); + + /** + * @brief Evaluate casadi functions for sensitivities + */ void CalcVarsSensitivities( realtype *tret, realtype *yval, const std::vector& ySval, realtype *yS_return, size_t *ySk); + + /** + * @brief The main solve method that solves for each variable and time step + */ Solution solve( np_array t_np, np_array y0_np, np_array yp0_np, np_array_dense inputs) override; + + /** + * @brief Concrete implementation of initialization method + */ void Initialize() override; + + /** + * @brief Allocate memory for OpenMP vectors + */ void AllocateVectors(); + + /** + * @brief Allocate memory for matrices (noting appropriate matrix format/types) + */ void SetMatrix(); + + /** + * @Brief Abstract method to set the linear solver + */ virtual void SetLinearSolver() = 0; }; diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp index 574272eef5..cbb29eb64a 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp @@ -1,9 +1,5 @@ #include "CasadiSolverOpenMP_solvers.hpp" -/* - * CasadiSolver implementations compatible with the OPENMP vector class - */ - void CasadiSolverOpenMP_Dense::SetLinearSolver() { LS = SUNLinSol_Dense(yy, J, sunctx); } diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp index 41413c1a39..5938d2ec15 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp @@ -4,6 +4,9 @@ #include "CasadiSolverOpenMP.hpp" #include "casadi_solver.hpp" +/** + * @brief CasadiSolver Dense implementation with OpenMP class + */ class CasadiSolverOpenMP_Dense : public CasadiSolverOpenMP { public: CasadiSolverOpenMP_Dense( @@ -36,6 +39,9 @@ class CasadiSolverOpenMP_Dense : public CasadiSolverOpenMP { void SetLinearSolver() override; }; +/** + * @brief CasadiSolver KLU implementation with OpenMP class + */ class CasadiSolverOpenMP_KLU : public CasadiSolverOpenMP { public: CasadiSolverOpenMP_KLU( @@ -68,6 +74,9 @@ class CasadiSolverOpenMP_KLU : public CasadiSolverOpenMP { void SetLinearSolver() override; }; +/** + * @brief CasadiSolver Banded implementation with OpenMP class + */ class CasadiSolverOpenMP_Band : public CasadiSolverOpenMP { public: CasadiSolverOpenMP_Band( @@ -100,6 +109,9 @@ class CasadiSolverOpenMP_Band : public CasadiSolverOpenMP { void SetLinearSolver() override; }; +/** + * @brief CasadiSolver SPBCGS implementation with OpenMP class + */ class CasadiSolverOpenMP_SPBCGS : public CasadiSolverOpenMP { public: CasadiSolverOpenMP_SPBCGS( @@ -132,6 +144,9 @@ class CasadiSolverOpenMP_SPBCGS : public CasadiSolverOpenMP { void SetLinearSolver() override; }; +/** + * @brief CasadiSolver SPFGMR implementation with OpenMP class + */ class CasadiSolverOpenMP_SPFGMR : public CasadiSolverOpenMP { public: CasadiSolverOpenMP_SPFGMR( @@ -164,6 +179,9 @@ class CasadiSolverOpenMP_SPFGMR : public CasadiSolverOpenMP { void SetLinearSolver() override; }; +/** + * @brief CasadiSolver SPGMR implementation with OpenMP class + */ class CasadiSolverOpenMP_SPGMR : public CasadiSolverOpenMP { public: CasadiSolverOpenMP_SPGMR( @@ -196,6 +214,9 @@ class CasadiSolverOpenMP_SPGMR : public CasadiSolverOpenMP { void SetLinearSolver() override; }; +/** + * @brief CasadiSolver SPTFQMR implementation with OpenMP class + */ class CasadiSolverOpenMP_SPTFQMR : public CasadiSolverOpenMP { public: CasadiSolverOpenMP_SPTFQMR( diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 1fd2ceda2d..10dd2dfc7d 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -6,8 +6,18 @@ #include #include -// Utility function for compressed-sparse-column (CSC) to/from -// compressed-sparse-row (CSR) matrix representation. +/** + * Utility function to convert compressed-sparse-column (CSC) to/from + * compressed-sparse-row (CSR) matrix representation. Conversion is symmetric / + * invertible using this function. + * @brief Utility function to convert to/from CSC/CSR matrix representations. + * @param f Data vector containing the sparse matrix elements + * @param c Index pointer to column starts + * @param r Array of row indices + * @param nf New data vector that will contain the transformed sparse matrix + * @param nc New array of column indices + * @param nr New index pointer to row starts + */ template void csc_csr(realtype f[], T1 c[], T1 r[], realtype nf[], T2 nc[], T2 nr[], int N, int cols) { int nn[cols+1]; @@ -35,13 +45,28 @@ void csc_csr(realtype f[], T1 c[], T1 r[], realtype nf[], T2 nc[], T2 nr[], int } } + using Function = casadi::Function; +/** + * @brief Class for handling individual casadi functions + */ class CasadiFunction { public: + /** + * @brief Constructor + */ explicit CasadiFunction(const Function &f); + + /** + * @brief Evaluation operator + */ void operator()(); + + /** + * @brief Evaluation operator given data vectors + */ void operator()(std::vector inputs, std::vector results); @@ -55,9 +80,15 @@ class CasadiFunction std::vector m_w; }; +/** + * @brief Class for handling casadi functions + */ class CasadiFunctions { public: + /** + * @brief Create a new CasadiFunctions object + */ CasadiFunctions( const Function &rhs_alg, const Function &jac_times_cjmass, diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp index c358634c93..335907a93a 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -3,6 +3,11 @@ #include "CasadiSolver.hpp" +/** + * Creates a concrete casadi solver given a linear solver, as specified in + * options_cpp.linear_solver. + * @brief Create a concrete casadi solver given a linear solver + */ CasadiSolver *create_casadi_solver( int number_of_states, int number_of_parameters, diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index 5b6a3b99f6..4f31dbe57d 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -3,8 +3,7 @@ #include "common.hpp" #include -#define NV_DATA NV_DATA_OMP -//#define NV_DATA NV_DATA_S +#define NV_DATA NV_DATA_OMP // Serial: NV_DATA_S int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr, void *user_data) diff --git a/pybamm/solvers/c_solvers/idaklu/options.hpp b/pybamm/solvers/c_solvers/idaklu/options.hpp index db5f136a01..b70d0f4a30 100644 --- a/pybamm/solvers/c_solvers/idaklu/options.hpp +++ b/pybamm/solvers/c_solvers/idaklu/options.hpp @@ -3,6 +3,9 @@ #include "common.hpp" +/** + * @brief Options passed to the idaklu solver by pybamm + */ struct Options { bool print_stats; bool using_sparse_matrix; diff --git a/pybamm/solvers/c_solvers/idaklu/python.hpp b/pybamm/solvers/c_solvers/idaklu/python.hpp index 8ae73f2a90..0478d0946f 100644 --- a/pybamm/solvers/c_solvers/idaklu/python.hpp +++ b/pybamm/solvers/c_solvers/idaklu/python.hpp @@ -22,7 +22,9 @@ using event_type = using jac_get_type = std::function; - +/** + * @brief Interface to the python solver + */ Solution solve_python(np_array t_np, np_array y0_np, np_array yp0_np, residual_type res, jacobian_type jac, sensitivities_type sens, diff --git a/pybamm/solvers/c_solvers/idaklu/solution.hpp b/pybamm/solvers/c_solvers/idaklu/solution.hpp index 047ae6ef8e..92e22d02b6 100644 --- a/pybamm/solvers/c_solvers/idaklu/solution.hpp +++ b/pybamm/solvers/c_solvers/idaklu/solution.hpp @@ -3,9 +3,15 @@ #include "common.hpp" +/** + * @brief Solution class + */ class Solution { public: + /** + * @brief Constructor + */ Solution(int retval, np_array t_np, np_array y_np, np_array yS_np) : flag(retval), t(t_np), y(y_np), yS(yS_np) { From 38e4a064d29958a8b6a377d65677c661181874bc Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 11 Sep 2023 09:50:48 +0000 Subject: [PATCH 26/38] Refactor CasadiFunctionsOpenMP concrete implementations to macro generator --- .../idaklu/CasadiSolverOpenMP_solvers.cpp | 48 --- .../idaklu/CasadiSolverOpenMP_solvers.hpp | 302 +++++------------- 2 files changed, 85 insertions(+), 265 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp index cbb29eb64a..868d2b2138 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp @@ -1,49 +1 @@ #include "CasadiSolverOpenMP_solvers.hpp" - -void CasadiSolverOpenMP_Dense::SetLinearSolver() { - LS = SUNLinSol_Dense(yy, J, sunctx); -} - -void CasadiSolverOpenMP_KLU::SetLinearSolver() { - LS = SUNLinSol_KLU(yy, J, sunctx); -} - -void CasadiSolverOpenMP_Band::SetLinearSolver() { - LS = SUNLinSol_Band(yy, J, sunctx); -} - -void CasadiSolverOpenMP_SPBCGS::SetLinearSolver() { - LS = SUNLinSol_SPBCGS( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -} - -void CasadiSolverOpenMP_SPFGMR::SetLinearSolver() { - LS = SUNLinSol_SPFGMR( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -} - -void CasadiSolverOpenMP_SPGMR::SetLinearSolver() { - LS = SUNLinSol_SPGMR( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -} - -void CasadiSolverOpenMP_SPTFQMR::SetLinearSolver() { - LS = SUNLinSol_SPTFQMR( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -} diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp index 5938d2ec15..a1fd8f3ca6 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp @@ -5,248 +5,116 @@ #include "casadi_solver.hpp" /** - * @brief CasadiSolver Dense implementation with OpenMP class + * Macro to generate CasadiSolver OpenMP implementations with specified linear + * solvers */ -class CasadiSolverOpenMP_Dense : public CasadiSolverOpenMP { -public: - CasadiSolverOpenMP_Dense( - 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, - const Options& options - ) : - CasadiSolverOpenMP( - atol_np, - rel_tol, - rhs_alg_id, - number_of_parameters, - number_of_events, - jac_times_cjmass_nnz, - jac_bandwidth_lower, - jac_bandwidth_upper, - std::move(functions), - options - ) - { - Initialize(); - } - void SetLinearSolver() override; +#define CASADISOLVER_NEWCLASS(CLASSNAME, FCN_CALL) \ +class CasadiSolverOpenMP_##CLASSNAME : public CasadiSolverOpenMP { \ +public: \ + CasadiSolverOpenMP_##CLASSNAME( \ + 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, \ + const Options& options \ + ) : \ + CasadiSolverOpenMP( \ + atol_np, \ + rel_tol, \ + rhs_alg_id, \ + number_of_parameters, \ + number_of_events, \ + jac_times_cjmass_nnz, \ + jac_bandwidth_lower, \ + jac_bandwidth_upper, \ + std::move(functions), \ + options \ + ) \ + { \ + Initialize(); \ + } \ + void SetLinearSolver() override { LS = FCN_CALL; }; \ }; +/** + * @brief CasadiSolver Dense implementation with OpenMP class + */ +CASADISOLVER_NEWCLASS( + Dense, + SUNLinSol_Dense(yy, J, sunctx) +) + /** * @brief CasadiSolver KLU implementation with OpenMP class */ -class CasadiSolverOpenMP_KLU : public CasadiSolverOpenMP { -public: - CasadiSolverOpenMP_KLU( - 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, - const Options& options - ) : - CasadiSolverOpenMP( - atol_np, - rel_tol, - rhs_alg_id, - number_of_parameters, - number_of_events, - jac_times_cjmass_nnz, - jac_bandwidth_lower, - jac_bandwidth_upper, - std::move(functions), - options - ) - { - Initialize(); - } - void SetLinearSolver() override; -}; +CASADISOLVER_NEWCLASS( + KLU, + SUNLinSol_KLU(yy, J, sunctx) +) /** * @brief CasadiSolver Banded implementation with OpenMP class */ -class CasadiSolverOpenMP_Band : public CasadiSolverOpenMP { -public: - CasadiSolverOpenMP_Band( - 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, - const Options& options - ) : - CasadiSolverOpenMP( - atol_np, - rel_tol, - rhs_alg_id, - number_of_parameters, - number_of_events, - jac_times_cjmass_nnz, - jac_bandwidth_lower, - jac_bandwidth_upper, - std::move(functions), - options - ) - { - Initialize(); - } - void SetLinearSolver() override; -}; +CASADISOLVER_NEWCLASS( + Band, + SUNLinSol_Band(yy, J, sunctx) +) /** * @brief CasadiSolver SPBCGS implementation with OpenMP class */ -class CasadiSolverOpenMP_SPBCGS : public CasadiSolverOpenMP { -public: - CasadiSolverOpenMP_SPBCGS( - 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, - const Options& options - ) : - CasadiSolverOpenMP( - atol_np, - rel_tol, - rhs_alg_id, - number_of_parameters, - number_of_events, - jac_times_cjmass_nnz, - jac_bandwidth_lower, - jac_bandwidth_upper, - std::move(functions), - options - ) - { - Initialize(); - } - void SetLinearSolver() override; -}; +CASADISOLVER_NEWCLASS( + SPBCGS, + LS = SUNLinSol_SPBCGS( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +) /** * @brief CasadiSolver SPFGMR implementation with OpenMP class */ -class CasadiSolverOpenMP_SPFGMR : public CasadiSolverOpenMP { -public: - CasadiSolverOpenMP_SPFGMR( - 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, - const Options& options - ) : - CasadiSolverOpenMP( - atol_np, - rel_tol, - rhs_alg_id, - number_of_parameters, - number_of_events, - jac_times_cjmass_nnz, - jac_bandwidth_lower, - jac_bandwidth_upper, - std::move(functions), - options - ) - { - Initialize(); - } - void SetLinearSolver() override; -}; +CASADISOLVER_NEWCLASS( + SPFGMR, + LS = SUNLinSol_SPFGMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +) /** * @brief CasadiSolver SPGMR implementation with OpenMP class */ -class CasadiSolverOpenMP_SPGMR : public CasadiSolverOpenMP { -public: - CasadiSolverOpenMP_SPGMR( - 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, - const Options& options - ) : - CasadiSolverOpenMP( - atol_np, - rel_tol, - rhs_alg_id, - number_of_parameters, - number_of_events, - jac_times_cjmass_nnz, - jac_bandwidth_lower, - jac_bandwidth_upper, - std::move(functions), - options - ) - { - Initialize(); - } - void SetLinearSolver() override; -}; +CASADISOLVER_NEWCLASS( + SPGMR, + LS = SUNLinSol_SPGMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +) /** * @brief CasadiSolver SPTFQMR implementation with OpenMP class */ -class CasadiSolverOpenMP_SPTFQMR : public CasadiSolverOpenMP { -public: - CasadiSolverOpenMP_SPTFQMR( - 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, - const Options& options - ) : - CasadiSolverOpenMP( - atol_np, - rel_tol, - rhs_alg_id, - number_of_parameters, - number_of_events, - jac_times_cjmass_nnz, - jac_bandwidth_lower, - jac_bandwidth_upper, - std::move(functions), - options - ) - { - Initialize(); - } - void SetLinearSolver() override; -}; +CASADISOLVER_NEWCLASS( + SPTFQMR, + LS = SUNLinSol_SPTFQMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); +) #endif // PYBAMM_IDAKLU_CASADI_SOLVER_OPENMP_HPP From bb56bea843a99883b1bc10fb347b95bba6bcac4c Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 11 Sep 2023 10:00:10 +0000 Subject: [PATCH 27/38] Reinforce appropriate access rights (provide getters) for function access in idaklu casadi_functions.hpp --- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 16 ++++++++-------- .../c_solvers/idaklu/casadi_functions.cpp | 8 ++++++++ .../c_solvers/idaklu/casadi_functions.hpp | 13 ++++++++++++- 3 files changed, 28 insertions(+), 9 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index fc4268a636..ca682c81ed 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -205,7 +205,7 @@ void CasadiSolverOpenMP::CalcVars( for (auto& var_fcn : functions->var_casadi_fcns) { var_fcn({tret, yval, functions->inputs.data()}, {res}); // store in return vector - for (size_t jj=0; jjdvar_dp_fcns[dvar_k]; // Calculate dvar/dy dvar_dy({tret, yval, functions->inputs.data()}, {res_dvar_dy}); - casadi::Sparsity spdy = dvar_dy.m_func.sparsity_out(0); + casadi::Sparsity spdy = dvar_dy.sparsity_out(0); // Calculate dvar/dp and convert to dense array for indexing dvar_dp({tret, yval, functions->inputs.data()}, {res_dvar_dp}); - casadi::Sparsity spdp = dvar_dp.m_func.sparsity_out(0); + casadi::Sparsity spdp = dvar_dp.sparsity_out(0); for(int k=0; kvar_casadi_fcns.size() > 0) { // return only the requested variables list after computation for (auto& var_fcn : functions->var_casadi_fcns) { - max_res_size = std::max(max_res_size, size_t(var_fcn.m_func.nnz_out())); - length_of_return_vector += var_fcn.m_func.nnz_out(); + max_res_size = std::max(max_res_size, size_t(var_fcn.nnz_out())); + length_of_return_vector += var_fcn.nnz_out(); for (auto& dvar_fcn : functions->dvar_dy_fcns) - max_res_dvar_dy = std::max(max_res_dvar_dy, size_t(dvar_fcn.m_func.nnz_out())); + max_res_dvar_dy = std::max(max_res_dvar_dy, size_t(dvar_fcn.nnz_out())); for (auto& dvar_fcn : functions->dvar_dp_fcns) - max_res_dvar_dp = std::max(max_res_dvar_dp, size_t(dvar_fcn.m_func.nnz_out())); + max_res_dvar_dp = std::max(max_res_dvar_dp, size_t(dvar_fcn.nnz_out())); } } else { // Return full y state-vector diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index 8b864d17c4..ca979d8648 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -25,6 +25,14 @@ void CasadiFunction::operator()() m_func.release(mem); } +casadi_int CasadiFunction::nnz_out() { + return m_func.nnz_out(); +} + +casadi::Sparsity CasadiFunction::sparsity_out(casadi_int ind) { + return m_func.sparsity_out(ind); +} + void CasadiFunction::operator()(std::vector inputs, std::vector results) { diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 10dd2dfc7d..5639e8c89a 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -4,6 +4,7 @@ #include "common.hpp" #include "options.hpp" #include +#include #include /** @@ -70,12 +71,22 @@ class CasadiFunction void operator()(std::vector inputs, std::vector results); + /** + * @brief Return the number of non-zero elements for the function output + */ + casadi_int nnz_out(); + + /** + * @brief Return the number of non-zero elements for the function output + */ + casadi::Sparsity sparsity_out(casadi_int ind); + public: std::vector m_arg; std::vector m_res; - const Function &m_func; private: + const Function &m_func; std::vector m_iw; std::vector m_w; }; From 21d68befa91fbc92d22a723c5a35bc4fa72e18f5 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 11 Sep 2023 10:31:01 +0000 Subject: [PATCH 28/38] Rename ProcessedVariableVar class to ProcessedVariableComputed and update docstring --- pybamm/__init__.py | 2 +- pybamm/solvers/idaklu_solver.py | 2 +- ...essed_variable_var.py => processed_variable_computed.py} | 6 +++++- 3 files changed, 7 insertions(+), 3 deletions(-) rename pybamm/solvers/{processed_variable_var.py => processed_variable_computed.py} (98%) diff --git a/pybamm/__init__.py b/pybamm/__init__.py index f20935a023..9aa1ca79a0 100644 --- a/pybamm/__init__.py +++ b/pybamm/__init__.py @@ -202,7 +202,7 @@ # from .solvers.solution import Solution, EmptySolution, make_cycle_solution from .solvers.processed_variable import ProcessedVariable -from .solvers.processed_variable_var import ProcessedVariableVar +from .solvers.processed_variable_computed import ProcessedVariableComputed from .solvers.base_solver import BaseSolver from .solvers.dummy_solver import DummySolver from .solvers.algebraic_solver import AlgebraicSolver diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index eec279c968..2ed8b106a1 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -682,7 +682,7 @@ def _integrate(self, model, t_eval, inputs_dict=None): len_of_var = ( self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz() ) - newsol._variables[var] = pybamm.ProcessedVariableVar( + newsol._variables[var] = pybamm.ProcessedVariableComputed( [model.variables_and_events[var]], [self._setup["var_casadi_fcns"][var]], [sol.y[:, startk : (startk + len_of_var)]], diff --git a/pybamm/solvers/processed_variable_var.py b/pybamm/solvers/processed_variable_computed.py similarity index 98% rename from pybamm/solvers/processed_variable_var.py rename to pybamm/solvers/processed_variable_computed.py index 11e39f99f1..78d16c27fb 100644 --- a/pybamm/solvers/processed_variable_var.py +++ b/pybamm/solvers/processed_variable_computed.py @@ -8,11 +8,15 @@ import xarray as xr -class ProcessedVariableVar(object): +class ProcessedVariableComputed(object): """ An object that can be evaluated at arbitrary (scalars or vectors) t and x, and returns the (interpolated) value of the base variable at that t and x. + The 'Computed' variant of ProcessedVariable deals with variables that have + been derived at solve time (see the 'output_variables' solver option), + where the full state-vector is not itself propogated and returned. + Parameters ---------- base_variables : list of :class:`pybamm.Symbol` From ab9819037abec5f0753c6f4187da5cfded3a3172 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 11 Sep 2023 11:16:02 +0000 Subject: [PATCH 29/38] Move casadi variable and sensitivity functions to the BaseSolver class --- pybamm/solvers/base_solver.py | 64 +++++++++++++++++ .../c_solvers/idaklu/casadi_functions.hpp | 2 +- pybamm/solvers/idaklu_solver.py | 71 ++++++------------- 3 files changed, 86 insertions(+), 51 deletions(-) diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py index 7740006310..5e65c20d93 100644 --- a/pybamm/solvers/base_solver.py +++ b/pybamm/solvers/base_solver.py @@ -38,6 +38,9 @@ class BaseSolver(object): 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. + output_variables : list[str], optional + List of variables to calculate and return. If none are specified then + the complete state vector is returned (can be very large) (default is []) """ def __init__( @@ -48,6 +51,7 @@ def __init__( root_method=None, root_tol=1e-6, extrap_tol=None, + output_variables=[], ): self.method = method self.rtol = rtol @@ -55,6 +59,7 @@ def __init__( self.root_tol = root_tol self.root_method = root_method self.extrap_tol = extrap_tol or -1e-10 + self.output_variables = output_variables self._model_set_up = {} # Defaults, can be overwritten by specific solver @@ -235,7 +240,9 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): # can use DAE solver to solve model with algebraic equations only if len(model.rhs) > 0: t_casadi = vars_for_processing["t_casadi"] + y_casadi = vars_for_processing["t_casadi"] y_and_S = vars_for_processing["y_and_S"] + p_casadi = vars_for_processing["p_casadi"] p_casadi_stacked = vars_for_processing["p_casadi_stacked"] mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries) explicit_rhs = mass_matrix_inv @ rhs( @@ -250,8 +257,65 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): model.casadi_sensitivities_rhs = jacp_rhs model.casadi_sensitivities_algebraic = jacp_algebraic + # if output_variables specified then convert functions to casadi + # expressions for evaluation within the respective solver + self.var_casadi_fcns = {} + self.dvar_dy_casadi_fcns = {} + self.dvar_dp_casadi_fcns = {} + for key in self.output_variables: + # ExplicitTimeIntegral's are not computed as part of the solver and + # do not need to be converted + if isinstance( + model.variables_and_events[key], pybamm.ExplicitTimeIntegral + ): + continue + # Generate Casadi function to calculate variable + fcn_name = BaseSolver._wrangle_name(key) + var_casadi = model.variables_and_events[key].to_casadi( + t_casadi, y_casadi, inputs=p_casadi + ) + self.var_casadi_fcns[key] = casadi.Function( + fcn_name, [t_casadi, y_casadi, p_casadi_stacked], [var_casadi] + ) + # Generate derivative functions for sensitivities + if (len(inputs) > 0) and (model.calculate_sensitivities): + dvar_dy = casadi.jacobian(var_casadi, y_casadi) + dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) + self.dvar_dy_casadi_fcns[key] = casadi.Function( + f"d{fcn_name}_dy", + [t_casadi, y_casadi, p_casadi_stacked], + [dvar_dy], + ) + self.dvar_dp_casadi_fcns[key] = casadi.Function( + f"d{fcn_name}_dp", + [t_casadi, y_casadi, p_casadi_stacked], + [dvar_dp], + ) + pybamm.logger.info("Finish solver set-up") + @classmethod + def _wrangle_name(cls, name: str) -> str: + """ + Wrangle a function name to replace special characters + """ + replacements = [ + (" ", "_"), + ("[", ""), + ("]", ""), + (".", "_"), + ("-", "_"), + ("(", ""), + (")", ""), + ("%", "prc"), + (",", ""), + (".", ""), + ] + name = "v_" + name.casefold() + for string, replacement in replacements: + name = name.replace(string, replacement) + return name + def _check_and_prepare_model_inplace(self, model, inputs, ics_only): """ Performs checks on the model and prepares it for solving. diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 5639e8c89a..3aa891a16a 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -75,7 +75,7 @@ class CasadiFunction * @brief Return the number of non-zero elements for the function output */ casadi_int nnz_out(); - + /** * @brief Return the number of non-zero elements for the function output */ diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 2ed8b106a1..c8ff7c844d 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -22,18 +22,6 @@ def have_idaklu(): return idaklu_spec is not None -def wrangle_name(name: str) -> str: - return "v_" + name.casefold().replace(" ", "_").replace("[", "").replace( - "]", "" - ).replace(".", "_").replace("-", "_").replace("(", "").replace(")", "").replace( - "%", "prc" - ).replace( - ",", "" - ).replace( - ".", "" - ) - - class IDAKLUSolver(pybamm.BaseSolver): """ Solve a discretised model, using sundials with the KLU sparse linear solver. @@ -275,47 +263,28 @@ def resfn(t, y, inputs, ydot): "mass_action", [v_casadi], [casadi.densify(mass_matrix @ v_casadi)] ) - # if output_variables specified then convert functions to casadi - # expressions for evaluation within the idaklu solver - self.var_casadi_fcns = {} - self.var_casadi_idaklu_fcns = [] - self.dvar_dy_fcns = [] - self.dvar_dp_fcns = [] + # if output_variables specified then convert 'variable' casadi + # function expressions to idaklu-compatible functions + self.var_idaklu_fcns = [] + self.dvar_dy_idaklu_fcns = [] + self.dvar_dp_idaklu_fcns = [] for key in self.output_variables: - # variable functions + # ExplicitTimeIntegral's are not computed as part of the solver and + # do not need to be converted if isinstance( model.variables_and_events[key], pybamm.ExplicitTimeIntegral ): continue - fcn_name = wrangle_name(key) - var_casadi = model.variables_and_events[key].to_casadi( - t_casadi, y_casadi, inputs=p_casadi - ) - self.var_casadi_fcns[key] = casadi.Function( - fcn_name, [t_casadi, y_casadi, p_casadi_stacked], [var_casadi] - ) - self.var_casadi_idaklu_fcns.append( + self.var_idaklu_fcns.append( idaklu.generate_function(self.var_casadi_fcns[key].serialize()) ) - # Generate derivative functions for sensitivities + # Convert derivative functions for sensitivities if (len(inputs) > 0) and (model.calculate_sensitivities): - dvar_dy = casadi.jacobian(var_casadi, y_casadi) - dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) - dvar_dy_fcn = casadi.Function( - f"d{fcn_name}_dy", - [t_casadi, y_casadi, p_casadi_stacked], - [dvar_dy], - ) - dvar_dp_fcn = casadi.Function( - f"d{fcn_name}_dp", - [t_casadi, y_casadi, p_casadi_stacked], - [dvar_dp], - ) - self.dvar_dy_fcns.append( - idaklu.generate_function(dvar_dy_fcn.serialize()) + self.dvar_dy_idaklu_fcns.append( + idaklu.generate_function(self.dvar_dy_casadi_fcns[key].serialize()) ) - self.dvar_dp_fcns.append( - idaklu.generate_function(dvar_dp_fcn.serialize()) + self.dvar_dp_idaklu_fcns.append( + idaklu.generate_function(self.dvar_dp_casadi_fcns[key].serialize()) ) else: @@ -490,9 +459,9 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): "number_of_sensitivity_parameters": number_of_sensitivity_parameters, "output_variables": self.output_variables, "var_casadi_fcns": self.var_casadi_fcns, - "var_casadi_idaklu_fcns": self.var_casadi_idaklu_fcns, - "dvar_dy_fcns": self.dvar_dy_fcns, - "dvar_dp_fcns": self.dvar_dp_fcns, + "var_idaklu_fcns": self.var_idaklu_fcns, + "dvar_dy_idaklu_fcns": self.dvar_dy_idaklu_fcns, + "dvar_dp_idaklu_fcns": self.dvar_dp_idaklu_fcns, } solver = idaklu.create_casadi_solver( @@ -514,9 +483,9 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): atol=atol, rtol=rtol, inputs=len(inputs), - var_casadi_fcns=self._setup["var_casadi_idaklu_fcns"], - dvar_dy_fcns=self._setup["dvar_dy_fcns"], - dvar_dp_fcns=self._setup["dvar_dp_fcns"], + var_casadi_fcns=self._setup["var_idaklu_fcns"], + dvar_dy_fcns=self._setup["dvar_dy_idaklu_fcns"], + dvar_dp_fcns=self._setup["dvar_dp_idaklu_fcns"], options=self._options, ) @@ -675,6 +644,8 @@ def _integrate(self, model, t_eval, inputs_dict=None): sol.y = sol.y.reshape((number_of_timesteps, number_of_samples)) startk = 0 for vark, var in enumerate(self.output_variables): + # ExplicitTimeIntegral's are not computed as part of the solver and + # do not need to be converted if isinstance( model.variables_and_events[var], pybamm.ExplicitTimeIntegral ): From f9b709a723f18c680337133ff0387958f5c824b5 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 11 Sep 2023 12:39:59 +0100 Subject: [PATCH 30/38] Update test functions for ProcessedVariableComputed class --- ...py => test_processed_variable_computed.py} | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) rename tests/unit/test_solvers/{test_processed_variable_var.py => test_processed_variable_computed.py} (95%) diff --git a/tests/unit/test_solvers/test_processed_variable_var.py b/tests/unit/test_solvers/test_processed_variable_computed.py similarity index 95% rename from tests/unit/test_solvers/test_processed_variable_var.py rename to tests/unit/test_solvers/test_processed_variable_computed.py index e5b62e8bee..e31f51ab1e 100644 --- a/tests/unit/test_solvers/test_processed_variable_var.py +++ b/tests/unit/test_solvers/test_processed_variable_computed.py @@ -1,5 +1,5 @@ # -# Tests for the Processed Variable Var class +# Tests for the Processed Variable Computed class # # This class forms a container for variables (and sensitivities) calculted # by the idaklu solver, and does not possesses any capability to calculate @@ -50,14 +50,14 @@ def process_and_check_2D_variable( var_casadi = to_casadi(var_sol, y_sol) model = tests.get_base_model_with_battery_geometry(**geometry_options) - pybamm.ProcessedVariableVar( + pybamm.ProcessedVariableComputed( [var_sol], [var_casadi], [y_sol], pybamm.Solution(t_sol, y_sol, model, {}), warn=False, ) - # NB: ProcessedVariableVar does not interpret y in the same way as + # NB: ProcessedVariableComputed does not interpret y in the same way as # ProcessedVariable; a better test of equivalence is to check that the # results are the same between IDAKLUSolver with (and without) # output_variables. This is implemented in the integration test: @@ -66,7 +66,7 @@ def process_and_check_2D_variable( return y_sol, first_sol, second_sol, t_sol -class TestProcessedVariableVar(TestCase): +class TestProcessedVariableComputed(TestCase): def test_processed_variable_0D(self): # without space y = pybamm.StateVector(slice(0, 1)) @@ -75,7 +75,7 @@ def test_processed_variable_0D(self): t_sol = np.array([0]) y_sol = np.array([1])[:, np.newaxis] var_casadi = to_casadi(var, y_sol) - processed_var = pybamm.ProcessedVariableVar( + processed_var = pybamm.ProcessedVariableComputed( [var], [var_casadi], [y_sol], @@ -104,7 +104,7 @@ def test_processed_variable_0D_no_sensitivity(self): t_sol = np.linspace(0, 1) y_sol = np.array([np.linspace(0, 5)]) var_casadi = to_casadi(var, y_sol) - processed_var = pybamm.ProcessedVariableVar( + processed_var = pybamm.ProcessedVariableComputed( [var], [var_casadi], [y_sol], @@ -125,7 +125,7 @@ def test_processed_variable_0D_no_sensitivity(self): y_sol = np.array([np.linspace(0, 5)]) inputs = {"a": np.array([1.0])} var_casadi = to_casadi(var, y_sol, inputs=inputs) - processed_var = pybamm.ProcessedVariableVar( + processed_var = pybamm.ProcessedVariableComputed( [var], [var_casadi], [y_sol], @@ -150,7 +150,7 @@ def test_processed_variable_1D(self): var_casadi = to_casadi(var_sol, y_sol) sol = pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}) - processed_var = pybamm.ProcessedVariableVar( + processed_var = pybamm.ProcessedVariableComputed( [var_sol], [var_casadi], [y_sol], @@ -176,7 +176,7 @@ def test_processed_variable_1D(self): processed_var.mesh.edges, processed_var.mesh.nodes # Check that there are no errors with domain-specific attributes - # (see ProcessedVariableVar.initialise_1D() for details) + # (see ProcessedVariableComputed.initialise_1D() for details) domain_list = [ "particle", "separator", @@ -214,7 +214,7 @@ def test_processed_variable_1D_unknown_domain(self): c = pybamm.StateVector(slice(0, var_pts[x]), domain=["SEI layer"]) c.mesh = mesh["SEI layer"] c_casadi = to_casadi(c, y_sol) - pybamm.ProcessedVariableVar([c], [c_casadi], [y_sol], solution, warn=False) + pybamm.ProcessedVariableComputed([c], [c_casadi], [y_sol], solution, warn=False) def test_processed_variable_2D_x_r(self): var = pybamm.Variable( @@ -360,7 +360,7 @@ def test_processed_variable_2D_space_only(self): y_sol = np.ones(len(x_sol) * len(r_sol))[:, np.newaxis] var_casadi = to_casadi(var_sol, y_sol) - processed_var = pybamm.ProcessedVariableVar( + processed_var = pybamm.ProcessedVariableComputed( [var_sol], [var_casadi], [y_sol], @@ -397,7 +397,7 @@ def test_processed_variable_2D_fixed_t_scikit(self): u_sol = np.ones(var_sol.shape[0])[:, np.newaxis] var_casadi = to_casadi(var_sol, u_sol) - processed_var = pybamm.ProcessedVariableVar( + processed_var = pybamm.ProcessedVariableComputed( [var_sol], [var_casadi], [u_sol], @@ -423,7 +423,7 @@ def test_3D_raises_error(self): var_casadi = to_casadi(var_sol, u_sol) with self.assertRaisesRegex(NotImplementedError, "Shape not recognized"): - pybamm.ProcessedVariableVar( + pybamm.ProcessedVariableComputed( [var_sol], [var_casadi], [u_sol], From c3e421870203f33768f4c3c6bb0559ab96a408e2 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 11 Sep 2023 15:26:59 +0100 Subject: [PATCH 31/38] Fix IdakluSolver issue overwriting casadi functions and add additional sensitivity tests when output_variables are specified --- pybamm/solvers/base_solver.py | 3 +- pybamm/solvers/idaklu_solver.py | 2 +- tests/unit/test_solvers/test_idaklu_solver.py | 72 ++++++++++++++++++- 3 files changed, 74 insertions(+), 3 deletions(-) diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py index 5e65c20d93..3d42b4e4cf 100644 --- a/pybamm/solvers/base_solver.py +++ b/pybamm/solvers/base_solver.py @@ -67,6 +67,7 @@ def __init__( self.ode_solver = False self.algebraic_solver = False self._on_extrapolation = "warn" + self.var_casadi_fcns = {} @property def root_method(self): @@ -240,7 +241,7 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): # can use DAE solver to solve model with algebraic equations only if len(model.rhs) > 0: t_casadi = vars_for_processing["t_casadi"] - y_casadi = vars_for_processing["t_casadi"] + y_casadi = vars_for_processing["y_casadi"] y_and_S = vars_for_processing["y_and_S"] p_casadi = vars_for_processing["p_casadi"] p_casadi_stacked = vars_for_processing["p_casadi_stacked"] diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index c8ff7c844d..b9e474d2af 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -122,6 +122,7 @@ def __init__( root_method, root_tol, extrap_tol, + output_variables, ) self.name = "IDA KLU solver" @@ -221,7 +222,6 @@ def resfn(t, y, inputs, ydot): raise pybamm.SolverError("KLU requires the Jacobian") # need to provide jacobian_rhs_alg - cj * mass_matrix - self.var_casadi_fcns = [] if model.convert_to_format == "casadi": t_casadi = casadi.MX.sym("t") y_casadi = casadi.MX.sym("y", model.len_rhs_and_alg) diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index 8024df7490..c1b7d42bd0 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -555,7 +555,7 @@ def test_with_output_variables(self): model = pybamm.lithium_ion.DFN() geometry = model.default_geometry param = model.default_parameter_values - input_parameters = {} + input_parameters = {} # Sensitivities dictionary param.update({key: "[input]" for key in input_parameters}) param.process_model(model) param.process_geometry(geometry) @@ -616,6 +616,76 @@ def test_with_output_variables(self): # Mock a 1D current collector and initialise (none in the model) sol["x_s [m]"].domain = ["current collector"] sol["x_s [m]"].initialise_1D() + + def test_with_output_variables_and_sensitivities(self): + # Construct a model and solve for all vairables, then test + # the 'output_variables' option for each variable in turn, confirming + # equivalence + + # construct model + model = pybamm.lithium_ion.DFN() + geometry = model.default_geometry + param = model.default_parameter_values + input_parameters = { # Sensitivities dictionary + "Current function [A]": 0.680616, + "Separator porosity": 1.0, + } + param.update({key: "[input]" for key in input_parameters}) + param.process_model(model) + param.process_geometry(geometry) + var_pts = {"x_n": 50, "x_s": 50, "x_p": 50, "r_n": 5, "r_p": 5} + mesh = pybamm.Mesh(geometry, model.default_submesh_types, var_pts) + disc = pybamm.Discretisation(mesh, model.default_spatial_methods) + disc.process_model(model) + t_eval = np.linspace(0, 3600, 100) + + options = { + 'linear_solver': 'SUNLinSol_KLU', + 'jacobian': 'sparse', + 'num_threads': 4, + } + + # Use a selection of variables of different types + output_variables = [ + "Voltage [V]", + "Time [min]", + "x [m]", + "Negative particle flux [mol.m-2.s-1]", + "Throughput capacity [A.h]", # ExplicitTimeIntegral + ] + + # Use the full model as comparison (tested separately) + solver_all = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + ) + sol_all = solver_all.solve( + model, + t_eval, + inputs=input_parameters, + calculate_sensitivities=True, + ) + + # Solve for a subset of variables and compare results + solver = pybamm.IDAKLUSolver( + atol=1e-8, rtol=1e-8, + options=options, + output_variables=output_variables, + ) + sol = solver.solve( + model, + t_eval, + inputs=input_parameters, + calculate_sensitivities=True, + ) + + # Compare output to sol_all + for varname in output_variables: + self.assertTrue(np.allclose(sol[varname].data, sol_all[varname].data)) + + # Mock a 1D current collector and initialise (none in the model) + sol["x_s [m]"].domain = ["current collector"] + sol["x_s [m]"].initialise_1D() if __name__ == "__main__": From 5c5222089153ae99ab112186d86688dc88c33337 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 11 Sep 2023 14:29:02 +0000 Subject: [PATCH 32/38] style: pre-commit fixes --- tests/unit/test_solvers/test_idaklu_solver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index c1b7d42bd0..cc54f3dfd5 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -616,7 +616,7 @@ def test_with_output_variables(self): # Mock a 1D current collector and initialise (none in the model) sol["x_s [m]"].domain = ["current collector"] sol["x_s [m]"].initialise_1D() - + def test_with_output_variables_and_sensitivities(self): # Construct a model and solve for all vairables, then test # the 'output_variables' option for each variable in turn, confirming From 2187b895c10eb6ac96565a5d84e7ce3f5cfc86e3 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Tue, 12 Sep 2023 14:33:51 +0100 Subject: [PATCH 33/38] Fix IDAKLUSolver crash on successive runs --- .../solvers/c_solvers/idaklu/CasadiSolver.cpp | 3 - .../solvers/c_solvers/idaklu/CasadiSolver.hpp | 4 +- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 96 +++++++++---------- .../c_solvers/idaklu/CasadiSolverOpenMP.hpp | 5 - .../idaklu/CasadiSolverOpenMP_solvers.hpp | 2 +- 5 files changed, 50 insertions(+), 60 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp index cc876edce6..16a04f8eb9 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp @@ -1,4 +1 @@ #include "CasadiSolver.hpp" - -CasadiSolver::CasadiSolver() {} -CasadiSolver::~CasadiSolver() {} diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp index 3703e98c99..dac94579f3 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp @@ -22,12 +22,12 @@ class CasadiSolver /** * @brief Default constructor */ - CasadiSolver(); + CasadiSolver() = default; /** * @brief Default destructor */ - ~CasadiSolver(); + ~CasadiSolver() = default; /** * @brief Abstract solver method that returns a Solution class diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index ca682c81ed..8319076310 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -29,6 +29,52 @@ CasadiSolverOpenMP::CasadiSolverOpenMP( { // Construction code moved to Initialize() which is called from the // (child) CasadiSolver_XXX class constructors. + DEBUG("CasadiSolverOpenMP::CasadiSolverOpenMP"); + auto atol = atol_np.unchecked<1>(); + + // create SUNDIALS context object + SUNContext_Create(NULL, &sunctx); // calls null-wrapper if Sundials Ver<6 + + // allocate memory for solver + ida_mem = IDACreate(sunctx); + + // create the vector of initial values + AllocateVectors(); + if (number_of_parameters > 0) + { + yyS = N_VCloneVectorArray(number_of_parameters, yy); + ypS = N_VCloneVectorArray(number_of_parameters, yp); + } + // set initial values + realtype *atval = N_VGetArrayPointer(avtol); + for (int i = 0; i < number_of_states; i++) + 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]); + } + + // create Matrix objects + SetMatrix(); + + // 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); + + // specify preconditioner type + precon_type = SUN_PREC_NONE; + if (options.preconditioner != "none") { + precon_type = SUN_PREC_LEFT; + } } void CasadiSolverOpenMP::AllocateVectors() { @@ -79,55 +125,7 @@ void CasadiSolverOpenMP::SetMatrix() { } void CasadiSolverOpenMP::Initialize() { - DEBUG("CasadiSolverOpenMP::Initialize"); - auto atol = atol_np.unchecked<1>(); - - // create SUNDIALS context object - SUNContext_Create(NULL, &sunctx); // calls null-wrapper if Sundials Ver<6 - - // allocate memory for solver - ida_mem = IDACreate(sunctx); - - // create the vector of initial values - AllocateVectors(); - if (number_of_parameters > 0) - { - yyS = N_VCloneVectorArray(number_of_parameters, yy); - ypS = N_VCloneVectorArray(number_of_parameters, yp); - } - // set initial values - realtype *atval = N_VGetArrayPointer(avtol); - for (int i = 0; i < number_of_states; i++) - 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]); - } - - // create Matrix objects - SetMatrix(); - - // 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); - - // specify preconditioner type - precon_type = SUN_PREC_NONE; - if (options.preconditioner != "none") { - precon_type = SUN_PREC_LEFT; - } - - // create linear solver object - SetLinearSolver(); + // Call after setting the solver // attach the linear solver IDASetLinearSolver(ida_mem, LS, J); diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp index 260c03ded7..0129caa9e7 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -139,11 +139,6 @@ class CasadiSolverOpenMP : public CasadiSolver * @brief Allocate memory for matrices (noting appropriate matrix format/types) */ void SetMatrix(); - - /** - * @Brief Abstract method to set the linear solver - */ - virtual void SetLinearSolver() = 0; }; #endif // PYBAMM_IDAKLU_CASADISOLVEROPENMP_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp index a1fd8f3ca6..1729f3ada3 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp @@ -36,9 +36,9 @@ public: \ options \ ) \ { \ + LS = FCN_CALL; \ Initialize(); \ } \ - void SetLinearSolver() override { LS = FCN_CALL; }; \ }; /** From 7f9aa24c9c6310788334518a4dd8eb2708eb6d95 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Wed, 13 Sep 2023 13:26:56 +0100 Subject: [PATCH 34/38] Refactor variable generation through the standard process() function --- pybamm/solvers/base_solver.py | 56 ++++++++++++++++----------------- pybamm/solvers/idaklu_solver.py | 8 ++--- 2 files changed, 32 insertions(+), 32 deletions(-) diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py index 3d42b4e4cf..b6ab7f1e2b 100644 --- a/pybamm/solvers/base_solver.py +++ b/pybamm/solvers/base_solver.py @@ -67,7 +67,7 @@ def __init__( self.ode_solver = False self.algebraic_solver = False self._on_extrapolation = "warn" - self.var_casadi_fcns = {} + self.computed_var_fcns = {} @property def root_method(self): @@ -241,9 +241,7 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): # can use DAE solver to solve model with algebraic equations only if len(model.rhs) > 0: t_casadi = vars_for_processing["t_casadi"] - y_casadi = vars_for_processing["y_casadi"] y_and_S = vars_for_processing["y_and_S"] - p_casadi = vars_for_processing["p_casadi"] p_casadi_stacked = vars_for_processing["p_casadi_stacked"] mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries) explicit_rhs = mass_matrix_inv @ rhs( @@ -260,9 +258,9 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): # if output_variables specified then convert functions to casadi # expressions for evaluation within the respective solver - self.var_casadi_fcns = {} - self.dvar_dy_casadi_fcns = {} - self.dvar_dp_casadi_fcns = {} + self.computed_var_fcns = {} + self.computed_dvar_dy_fcns = {} + self.computed_dvar_dp_fcns = {} for key in self.output_variables: # ExplicitTimeIntegral's are not computed as part of the solver and # do not need to be converted @@ -270,28 +268,20 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False): model.variables_and_events[key], pybamm.ExplicitTimeIntegral ): continue - # Generate Casadi function to calculate variable - fcn_name = BaseSolver._wrangle_name(key) - var_casadi = model.variables_and_events[key].to_casadi( - t_casadi, y_casadi, inputs=p_casadi - ) - self.var_casadi_fcns[key] = casadi.Function( - fcn_name, [t_casadi, y_casadi, p_casadi_stacked], [var_casadi] + # Generate Casadi function to calculate variable and derivates + # to enable sensitivites to be computed within the solver + ( + self.computed_var_fcns[key], + self.computed_dvar_dy_fcns[key], + self.computed_dvar_dp_fcns[key], + _, + ) = process( + model.variables_and_events[key], + BaseSolver._wrangle_name(key), + vars_for_processing, + use_jacobian=True, + return_jacp_stacked=True, ) - # Generate derivative functions for sensitivities - if (len(inputs) > 0) and (model.calculate_sensitivities): - dvar_dy = casadi.jacobian(var_casadi, y_casadi) - dvar_dp = casadi.jacobian(var_casadi, p_casadi_stacked) - self.dvar_dy_casadi_fcns[key] = casadi.Function( - f"d{fcn_name}_dy", - [t_casadi, y_casadi, p_casadi_stacked], - [dvar_dy], - ) - self.dvar_dp_casadi_fcns[key] = casadi.Function( - f"d{fcn_name}_dp", - [t_casadi, y_casadi, p_casadi_stacked], - [dvar_dp], - ) pybamm.logger.info("Finish solver set-up") @@ -1430,7 +1420,7 @@ def _set_up_model_inputs(self, model, inputs): return ordered_inputs -def process(symbol, name, vars_for_processing, use_jacobian=None): +def process(symbol, name, vars_for_processing, use_jacobian=None, return_jacp_stacked=None): """ Parameters ---------- @@ -1440,6 +1430,8 @@ def process(symbol, name, vars_for_processing, use_jacobian=None): function evaluators created will have this base name use_jacobian: bool, optional whether to return Jacobian functions + return_jacp_stacked: bool, optional + returns Jacobian function wrt stacked parameters instead of jacp Returns ------- @@ -1650,6 +1642,14 @@ def jacp(*args, **kwargs): [t_casadi, y_and_S, p_casadi_stacked, v], [jac_action_casadi], ) + # Compute derivate wrt p-stacked (can be passed to solver to + # compute sensitivities online) + if return_jacp_stacked: + jacp = casadi.Function( + f"d{name}_dp", + [t_casadi, y_casadi, p_casadi_stacked], + [casadi.jacobian(casadi_expression, p_casadi_stacked)], + ) else: jac = None jac_action = None diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index b9e474d2af..d9819f1608 100644 --- a/pybamm/solvers/idaklu_solver.py +++ b/pybamm/solvers/idaklu_solver.py @@ -276,15 +276,15 @@ def resfn(t, y, inputs, ydot): ): continue self.var_idaklu_fcns.append( - idaklu.generate_function(self.var_casadi_fcns[key].serialize()) + idaklu.generate_function(self.computed_var_fcns[key].serialize()) ) # Convert derivative functions for sensitivities if (len(inputs) > 0) and (model.calculate_sensitivities): self.dvar_dy_idaklu_fcns.append( - idaklu.generate_function(self.dvar_dy_casadi_fcns[key].serialize()) + idaklu.generate_function(self.computed_dvar_dy_fcns[key].serialize()) ) self.dvar_dp_idaklu_fcns.append( - idaklu.generate_function(self.dvar_dp_casadi_fcns[key].serialize()) + idaklu.generate_function(self.computed_dvar_dp_fcns[key].serialize()) ) else: @@ -458,7 +458,7 @@ def sensfn(resvalS, t, y, inputs, yp, yS, ypS): "sensitivity_names": sensitivity_names, "number_of_sensitivity_parameters": number_of_sensitivity_parameters, "output_variables": self.output_variables, - "var_casadi_fcns": self.var_casadi_fcns, + "var_casadi_fcns": self.computed_var_fcns, "var_idaklu_fcns": self.var_idaklu_fcns, "dvar_dy_idaklu_fcns": self.dvar_dy_idaklu_fcns, "dvar_dp_idaklu_fcns": self.dvar_dp_idaklu_fcns, From f381e31c23ac02d5d6f4b6eb32294d45fe4c84f2 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Wed, 13 Sep 2023 16:04:03 +0100 Subject: [PATCH 35/38] Pre-commit fixes --- pybamm/solvers/base_solver.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py index b6ab7f1e2b..c0ee4c32fc 100644 --- a/pybamm/solvers/base_solver.py +++ b/pybamm/solvers/base_solver.py @@ -1420,7 +1420,13 @@ def _set_up_model_inputs(self, model, inputs): return ordered_inputs -def process(symbol, name, vars_for_processing, use_jacobian=None, return_jacp_stacked=None): +def process( + symbol, + name, + vars_for_processing, + use_jacobian=None, + return_jacp_stacked=None +): """ Parameters ---------- From 174bec41f3152b95065e8ca598cc751ca524d57c Mon Sep 17 00:00:00 2001 From: John Brittain Date: Mon, 18 Sep 2023 09:31:38 +0100 Subject: [PATCH 36/38] Streamline base_solver process() --- pybamm/solvers/base_solver.py | 47 +++++++++++++++++------------------ 1 file changed, 23 insertions(+), 24 deletions(-) diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py index c0ee4c32fc..0a412f6a06 100644 --- a/pybamm/solvers/base_solver.py +++ b/pybamm/solvers/base_solver.py @@ -1421,11 +1421,7 @@ def _set_up_model_inputs(self, model, inputs): def process( - symbol, - name, - vars_for_processing, - use_jacobian=None, - return_jacp_stacked=None + symbol, name, vars_for_processing, use_jacobian=None, return_jacp_stacked=None ): """ Parameters @@ -1615,17 +1611,28 @@ def jacp(*args, **kwargs): "CasADi" ) ) - # WARNING, jacp for convert_to_format=casadi does not return a dict - # instead it returns multiple return values, one for each param - # TODO: would it be faster to do the jacobian wrt pS_casadi_stacked? - jacp = casadi.Function( - name + "_jacp", - [t_casadi, y_and_S, p_casadi_stacked], - [ - casadi.densify(casadi.jacobian(casadi_expression, p_casadi[pname])) - for pname in model.calculate_sensitivities - ], - ) + # Compute derivate wrt p-stacked (can be passed to solver to + # compute sensitivities online) + if return_jacp_stacked: + jacp = casadi.Function( + f"d{name}_dp", + [t_casadi, y_casadi, p_casadi_stacked], + [casadi.jacobian(casadi_expression, p_casadi_stacked)], + ) + else: + # WARNING, jacp for convert_to_format=casadi does not return a dict + # instead it returns multiple return values, one for each param + # TODO: would it be faster to do the jacobian wrt pS_casadi_stacked? + jacp = casadi.Function( + name + "_jacp", + [t_casadi, y_and_S, p_casadi_stacked], + [ + casadi.densify( + casadi.jacobian(casadi_expression, p_casadi[pname]) + ) + for pname in model.calculate_sensitivities + ], + ) if use_jacobian: report(f"Calculating jacobian for {name} using CasADi") @@ -1648,14 +1655,6 @@ def jacp(*args, **kwargs): [t_casadi, y_and_S, p_casadi_stacked, v], [jac_action_casadi], ) - # Compute derivate wrt p-stacked (can be passed to solver to - # compute sensitivities online) - if return_jacp_stacked: - jacp = casadi.Function( - f"d{name}_dp", - [t_casadi, y_casadi, p_casadi_stacked], - [casadi.jacobian(casadi_expression, p_casadi_stacked)], - ) else: jac = None jac_action = None From e8ff644babae52d7aabd904970f1d3fd189a3464 Mon Sep 17 00:00:00 2001 From: John Brittain Date: Tue, 19 Sep 2023 22:41:35 +0100 Subject: [PATCH 37/38] Separate idaklu implementation classes --- .../idaklu/CasadiSolverOpenMP_solvers.hpp | 175 +++++++++--------- .../c_solvers/idaklu/casadi_functions.cpp | 4 +- .../c_solvers/idaklu/casadi_functions.hpp | 4 +- 3 files changed, 94 insertions(+), 89 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp index 1729f3ada3..3e39e5a303 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp @@ -4,117 +4,122 @@ #include "CasadiSolverOpenMP.hpp" #include "casadi_solver.hpp" -/** - * Macro to generate CasadiSolver OpenMP implementations with specified linear - * solvers - */ -#define CASADISOLVER_NEWCLASS(CLASSNAME, FCN_CALL) \ -class CasadiSolverOpenMP_##CLASSNAME : public CasadiSolverOpenMP { \ -public: \ - CasadiSolverOpenMP_##CLASSNAME( \ - 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, \ - const Options& options \ - ) : \ - CasadiSolverOpenMP( \ - atol_np, \ - rel_tol, \ - rhs_alg_id, \ - number_of_parameters, \ - number_of_events, \ - jac_times_cjmass_nnz, \ - jac_bandwidth_lower, \ - jac_bandwidth_upper, \ - std::move(functions), \ - options \ - ) \ - { \ - LS = FCN_CALL; \ - Initialize(); \ - } \ -}; - /** * @brief CasadiSolver Dense implementation with OpenMP class */ -CASADISOLVER_NEWCLASS( - Dense, - SUNLinSol_Dense(yy, J, sunctx) -) +class CasadiSolverOpenMP_Dense : public CasadiSolverOpenMP { +public: + template + CasadiSolverOpenMP_Dense(Args&& ... args) + : CasadiSolverOpenMP(std::forward(args) ...) + { + LS = SUNLinSol_Dense(yy, J, sunctx); + Initialize(); + } +}; /** * @brief CasadiSolver KLU implementation with OpenMP class */ -CASADISOLVER_NEWCLASS( - KLU, - SUNLinSol_KLU(yy, J, sunctx) -) +class CasadiSolverOpenMP_KLU : public CasadiSolverOpenMP { +public: + template + CasadiSolverOpenMP_KLU(Args&& ... args) + : CasadiSolverOpenMP(std::forward(args) ...) + { + LS = SUNLinSol_KLU(yy, J, sunctx); + Initialize(); + } +}; /** * @brief CasadiSolver Banded implementation with OpenMP class */ -CASADISOLVER_NEWCLASS( - Band, - SUNLinSol_Band(yy, J, sunctx) -) +class CasadiSolverOpenMP_Band : public CasadiSolverOpenMP { +public: + template + CasadiSolverOpenMP_Band(Args&& ... args) + : CasadiSolverOpenMP(std::forward(args) ...) + { + LS = SUNLinSol_Band(yy, J, sunctx); + Initialize(); + } +}; /** * @brief CasadiSolver SPBCGS implementation with OpenMP class */ -CASADISOLVER_NEWCLASS( - SPBCGS, - LS = SUNLinSol_SPBCGS( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -) +class CasadiSolverOpenMP_SPBCGS : public CasadiSolverOpenMP { +public: + template + CasadiSolverOpenMP_SPBCGS(Args&& ... args) + : CasadiSolverOpenMP(std::forward(args) ...) + { + LS = SUNLinSol_SPBCGS( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); + Initialize(); + } +}; /** * @brief CasadiSolver SPFGMR implementation with OpenMP class */ -CASADISOLVER_NEWCLASS( - SPFGMR, - LS = SUNLinSol_SPFGMR( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -) +class CasadiSolverOpenMP_SPFGMR : public CasadiSolverOpenMP { +public: + template + CasadiSolverOpenMP_SPFGMR(Args&& ... args) + : CasadiSolverOpenMP(std::forward(args) ...) + { + LS = SUNLinSol_SPFGMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); + Initialize(); + } +}; /** * @brief CasadiSolver SPGMR implementation with OpenMP class */ -CASADISOLVER_NEWCLASS( - SPGMR, - LS = SUNLinSol_SPGMR( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -) +class CasadiSolverOpenMP_SPGMR : public CasadiSolverOpenMP { +public: + template + CasadiSolverOpenMP_SPGMR(Args&& ... args) + : CasadiSolverOpenMP(std::forward(args) ...) + { + LS = SUNLinSol_SPGMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); + Initialize(); + } +}; /** * @brief CasadiSolver SPTFQMR implementation with OpenMP class */ -CASADISOLVER_NEWCLASS( - SPTFQMR, - LS = SUNLinSol_SPTFQMR( - yy, - precon_type, - options.linsol_max_iterations, - sunctx - ); -) +class CasadiSolverOpenMP_SPTFQMR : public CasadiSolverOpenMP { +public: + template + CasadiSolverOpenMP_SPTFQMR(Args&& ... args) + : CasadiSolverOpenMP(std::forward(args) ...) + { + LS = SUNLinSol_SPTFQMR( + yy, + precon_type, + options.linsol_max_iterations, + sunctx + ); + Initialize(); + } +}; #endif // PYBAMM_IDAKLU_CASADI_SOLVER_OPENMP_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp index ca979d8648..ddad4612c9 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -33,8 +33,8 @@ casadi::Sparsity CasadiFunction::sparsity_out(casadi_int ind) { return m_func.sparsity_out(ind); } -void CasadiFunction::operator()(std::vector inputs, - std::vector results) +void CasadiFunction::operator()(const std::vector& inputs, + const std::vector& results) { // Set-up input arguments, provide result vector, then execute function // Example call: fcn({in1, in2, in3}, {out1}) diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 3aa891a16a..4b16639882 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -68,8 +68,8 @@ class CasadiFunction /** * @brief Evaluation operator given data vectors */ - void operator()(std::vector inputs, - std::vector results); + void operator()(const std::vector& inputs, + const std::vector& results); /** * @brief Return the number of non-zero elements for the function output From 9ac838321f71d43c2a5a81d4a8fa525ce82d634b Mon Sep 17 00:00:00 2001 From: John Brittain Date: Fri, 22 Sep 2023 11:27:56 +0100 Subject: [PATCH 38/38] Codacy improvements --- .../c_solvers/idaklu/CasadiSolverOpenMP.cpp | 5 +-- .../c_solvers/idaklu/CasadiSolverOpenMP.hpp | 31 ++++++++++--------- .../c_solvers/idaklu/casadi_functions.hpp | 13 +++++--- 3 files changed, 28 insertions(+), 21 deletions(-) diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp index 8319076310..c1c71a967d 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -128,6 +128,8 @@ void CasadiSolverOpenMP::Initialize() { // Call after setting the solver // attach the linear solver + if (LS == nullptr) + throw std::invalid_argument("Linear solver not set"); IDASetLinearSolver(ida_mem, LS, J); if (options.preconditioner != "none") @@ -307,7 +309,6 @@ Solution CasadiSolverOpenMP::solve( IDAGetSens(ida_mem, &t0, yyS); realtype tret; - realtype t_next; realtype t_final = t(number_of_timesteps - 1); // set return vectors @@ -386,7 +387,7 @@ Solution CasadiSolverOpenMP::solve( t_i = 1; while (true) { - t_next = t(t_i); + realtype t_next = t(t_i); IDASetStopTime(ida_mem, t_next); DEBUG("IDASolve"); retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp index 0129caa9e7..2312f9cf8f 100644 --- a/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -42,28 +42,31 @@ using Function = casadi::Function; */ class CasadiSolverOpenMP : public CasadiSolver { + // NB: cppcheck-suppress unusedStructMember is used because codacy reports + // these members as unused even though they are important in child + // classes, but are passed by variadic arguments (and are therefore unnamed) public: void *ida_mem = nullptr; np_array atol_np; - double rel_tol; np_array rhs_alg_id; - int number_of_states; - int number_of_parameters; - int number_of_events; - int precon_type; - N_Vector yy, yp, avtol; // y, y', and absolute tolerance - N_Vector *yyS, *ypS; // y, y' for sensitivities + int number_of_states; // cppcheck-suppress unusedStructMember + int number_of_parameters; // cppcheck-suppress unusedStructMember + int number_of_events; // cppcheck-suppress unusedStructMember + int precon_type; // cppcheck-suppress unusedStructMember + N_Vector yy, yp, avtol; // y, y', and absolute tolerance + N_Vector *yyS; // cppcheck-suppress unusedStructMember + N_Vector *ypS; // cppcheck-suppress unusedStructMember N_Vector id; // rhs_alg_id realtype rtol; - const int jac_times_cjmass_nnz; - int jac_bandwidth_lower; - int jac_bandwidth_upper; + const int jac_times_cjmass_nnz; // cppcheck-suppress unusedStructMember + int jac_bandwidth_lower; // cppcheck-suppress unusedStructMember + int jac_bandwidth_upper; // cppcheck-suppress unusedStructMember SUNMatrix J; - SUNLinearSolver LS; + SUNLinearSolver LS = nullptr; std::unique_ptr functions; - realtype *res; - realtype *res_dvar_dy; - realtype *res_dvar_dp; + realtype *res = nullptr; + realtype *res_dvar_dy = nullptr; + realtype *res_dvar_dp = nullptr; Options options; #if SUNDIALS_VERSION_MAJOR >= 6 diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp index 4b16639882..d29cb7c961 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -20,9 +20,9 @@ * @param nr New index pointer to row starts */ template -void csc_csr(realtype f[], T1 c[], T1 r[], realtype nf[], T2 nc[], T2 nr[], int N, int cols) { +void csc_csr(const realtype f[], const T1 c[], const T1 r[], realtype nf[], T2 nc[], T2 nr[], int N, int cols) { int nn[cols+1]; - int rr[N]; + std::vector rr(N); for (int i=0; i var_casadi_fcns; - std::vector dvar_dy_fcns; - std::vector dvar_dp_fcns; + + // NB: cppcheck-suppress unusedStructMember is used because codacy reports + // these members as unused even though they are important + std::vector var_casadi_fcns; // cppcheck-suppress unusedStructMember + std::vector dvar_dy_fcns; // cppcheck-suppress unusedStructMember + std::vector dvar_dp_fcns; // cppcheck-suppress unusedStructMember std::vector jac_times_cjmass_rowvals; std::vector jac_times_cjmass_colptrs;