diff --git a/CHANGELOG.md b/CHANGELOG.md index 8a93addf65..21cfb2ae87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,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)) diff --git a/CMakeLists.txt b/CMakeLists.txt index c3c5141d4f..38e3d4976c 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/__init__.py b/pybamm/__init__.py index 6c2636ba51..9aa1ca79a0 100644 --- a/pybamm/__init__.py +++ b/pybamm/__init__.py @@ -202,6 +202,7 @@ # from .solvers.solution import Solution, EmptySolution, make_cycle_solution from .solvers.processed_variable import ProcessedVariable +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/base_solver.py b/pybamm/solvers/base_solver.py index a92a9309d4..c2b81c1568 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 @@ -62,6 +67,7 @@ def __init__( self.ode_solver = False self.algebraic_solver = False self._on_extrapolation = "warn" + self.computed_var_fcns = {} @property def root_method(self): @@ -250,8 +256,57 @@ 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.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 + if isinstance( + model.variables_and_events[key], pybamm.ExplicitTimeIntegral + ): + continue + # 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, + ) + 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. @@ -1366,7 +1421,9 @@ 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 ---------- @@ -1376,6 +1433,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 ------- @@ -1553,17 +1612,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") diff --git a/pybamm/solvers/c_solvers/idaklu.cpp b/pybamm/solvers/c_solvers/idaklu.cpp index 132e8883f4..be90955b9c 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 @@ -25,39 +26,75 @@ 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("var_casadi_fcns"), + py::arg("dvar_dy_fcns"), + py::arg("dvar_dp_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/CasadiSolver.cpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp new file mode 100644 index 0000000000..16a04f8eb9 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.cpp @@ -0,0 +1 @@ +#include "CasadiSolver.hpp" diff --git a/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp new file mode 100644 index 0000000000..dac94579f3 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolver.hpp @@ -0,0 +1,49 @@ +#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" + +/** + * 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() = default; + + /** + * @brief Default destructor + */ + ~CasadiSolver() = default; + + /** + * @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; +}; + +#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..c1c71a967d --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.cpp @@ -0,0 +1,518 @@ +#include "CasadiSolverOpenMP.hpp" +#include "casadi_sundials_functions.hpp" +#include +#include +#include + +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. + 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() { + // 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() { + // 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") + { + 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); +} + +void CasadiSolverOpenMP::CalcVars( + realtype *y_return, + size_t length_of_return_vector, + size_t t_i, + realtype *tret, + realtype *yval, + const 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.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.sparsity_out(0); + for(int k=0; k(); + 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]; + } + + 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); + + realtype tret; + 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) + size_t max_res_dvar_dy = 0, max_res_dvar_dp = 0; + 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, 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.nnz_out())); + for (auto& dvar_fcn : functions->dvar_dp_fcns) + max_res_dvar_dp = std::max(max_res_dvar_dp, size_t(dvar_fcn.nnz_out())); + } + } 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 * + length_of_return_vector]; + 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]; + + 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; + } + ); + + // Initial state (t_i=0) + 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 + CalcVars(y_return, length_of_return_vector, t_i, + &tret, yval, ySval, yS_return, &ySk); + } else { + // Retain complete copy of the state vector + for (int j = 0; j < number_of_states; j++) + y_return[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]; + } + } + + // Subsequent states (t_i>0) + int retval; + t_i = 1; + while (true) + { + realtype t_next = t(t_i); + IDASetStopTime(ida_mem, t_next); + DEBUG("IDASolve"); + retval = IDASolve(ida_mem, t_final, &tret, yy, yp, IDA_NORMAL); + + if (retval == IDA_TSTOP_RETURN || + retval == IDA_SUCCESS || + retval == IDA_ROOT_RETURN) + { + if (number_of_parameters > 0) + IDAGetSens(ida_mem, &tret, yyS); + + // Evaluate and store results for the time step + t_return[t_i] = tret; + if (functions->var_casadi_fcns.size() > 0) { + // Evaluate casadi functions for each requested variable and store + // NOTE: Indexing of yS_return is (time:var:param) + CalcVars(y_return, length_of_return_vector, t_i, + &tret, yval, ySval, yS_return, &ySk); + } else { + // Retain complete copy of the state vector + for (int j = 0; j < number_of_states; j++) + y_return[t_i * number_of_states + j] = yval[j]; + for (int j = 0; j < number_of_parameters; j++) + { + const int base_index = + j * number_of_timesteps * number_of_states + + t_i * number_of_states; + for (int k = 0; k < number_of_states; k++) + // NOTE: Indexing of yS_return is (time:param:yvec) + yS_return[base_index + k] = ySval[j][k]; + } + } + t_i += 1; + + if (retval == IDA_SUCCESS || + retval == IDA_ROOT_RETURN) + break; + } + else + { + // failed + break; + } + } + + np_array t_ret = np_array( + t_i, + &t_return[0], + free_t_when_done + ); + np_array y_ret = np_array( + t_i * length_of_return_vector, + &y_return[0], + free_y_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); + + 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..2312f9cf8f --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP.hpp @@ -0,0 +1,147 @@ +#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" + +/** + * @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 +{ + // 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; + np_array rhs_alg_id; + 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; // cppcheck-suppress unusedStructMember + int jac_bandwidth_lower; // cppcheck-suppress unusedStructMember + int jac_bandwidth_upper; // cppcheck-suppress unusedStructMember + SUNMatrix J; + SUNLinearSolver LS = nullptr; + std::unique_ptr functions; + realtype *res = nullptr; + realtype *res_dvar_dy = nullptr; + realtype *res_dvar_dp = nullptr; + Options options; + +#if SUNDIALS_VERSION_MAJOR >= 6 + SUNContext sunctx; +#endif + +public: + /** + * @brief Constructor + */ + 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); + + /** + * @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, + size_t t_i, + realtype *tret, + realtype *yval, + 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(); +}; + +#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..868d2b2138 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.cpp @@ -0,0 +1 @@ +#include "CasadiSolverOpenMP_solvers.hpp" 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..3e39e5a303 --- /dev/null +++ b/pybamm/solvers/c_solvers/idaklu/CasadiSolverOpenMP_solvers.hpp @@ -0,0 +1,125 @@ +#ifndef PYBAMM_IDAKLU_CASADI_SOLVER_OPENMP_HPP +#define PYBAMM_IDAKLU_CASADI_SOLVER_OPENMP_HPP + +#include "CasadiSolverOpenMP.hpp" +#include "casadi_solver.hpp" + +/** + * @brief CasadiSolver Dense implementation with OpenMP class + */ +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 + */ +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 + */ +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 + */ +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 + */ +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 + */ +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 + */ +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 310575742d..ddad4612c9 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.cpp @@ -7,23 +7,17 @@ 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); - // 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); + //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; + 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(); @@ -31,25 +25,59 @@ 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()(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}) + 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), + 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) { + // convert casadi::Function list to CasadiFunction list + 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; @@ -67,8 +95,11 @@ CasadiFunctions::CasadiFunctions( } 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..d29cb7c961 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_functions.hpp @@ -3,20 +3,87 @@ #include "common.hpp" #include "options.hpp" -#include "solution.hpp" #include +#include +#include + +/** + * 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(const realtype f[], const T1 c[], const T1 r[], realtype nf[], T2 nc[], T2 nr[], int N, int cols) { + int nn[cols+1]; + std::vector rr(N); + for (int i=0; i& inputs, + const 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; - void operator()(); private: const Function &m_func; @@ -24,8 +91,37 @@ 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, + 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 std::vector& dvar_dy_fcns, + const std::vector& dvar_dp_fcns, + const Options& options + ); + public: int number_of_states; int number_of_parameters; @@ -33,26 +129,25 @@ 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; - 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); + // 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; + 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 67bb2793ae..9fcfa06510 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp @@ -1,494 +1,177 @@ #include "casadi_solver.hpp" +#include "CasadiSolver.hpp" +#include "CasadiSolverOpenMP_solvers.hpp" #include "casadi_sundials_functions.hpp" #include "common.hpp" #include #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, + 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); 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, + var_casadi_fcns, + dvar_dy_fcns, + dvar_dp_fcns, + 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") + 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"); -#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_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 + 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 + ); } - - - 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]; - } - - 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 75bc73e9d3..335907a93a 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_solver.hpp @@ -1,59 +1,36 @@ -#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" + +/** + * 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, + 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, + const std::vector& var_casadi_fcns, + const std::vector& dvar_dy_fcns, + const std::vector& dvar_dp_fcns, + py::dict options +); + +#endif // PYBAMM_IDAKLU_CREATE_CASADI_SOLVER_HPP diff --git a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp index 8a3d96966f..4f31dbe57d 100644 --- a/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp +++ b/pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp @@ -1,6 +1,9 @@ #include "casadi_sundials_functions.hpp" #include "casadi_functions.hpp" #include "common.hpp" +#include + +#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) @@ -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; @@ -190,30 +193,61 @@ 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(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< + 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(), + 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); @@ -227,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(); @@ -271,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(); @@ -285,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/c_solvers/idaklu/common.hpp b/pybamm/solvers/c_solvers/idaklu/common.hpp index d6ddb5c16b..55fd4b1c5d 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 33998470ed..efad4d5de0 100644 --- a/pybamm/solvers/c_solvers/idaklu/options.cpp +++ b/pybamm/solvers/c_solvers/idaklu/options.cpp @@ -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) { - throw std::domain_error( - "Unknown preconditioner \""s + preconditioner + - "\", use one of \"BBDP\" or \"none\""s - ); - } - } - else - { - preconditioner = "none"; - } + if (preconditioner != "none" && preconditioner != "BBDP") + { + throw std::domain_error( + "Unknown preconditioner \""s + preconditioner + + "\", use one of \"BBDP\" or \"none\""s + ); + } + } + else + { + preconditioner = "none"; + } } diff --git a/pybamm/solvers/c_solvers/idaklu/options.hpp b/pybamm/solvers/c_solvers/idaklu/options.hpp 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.cpp b/pybamm/solvers/c_solvers/idaklu/python.cpp index 9ec018109e..03090c9850 100644 --- a/pybamm/solvers/c_solvers/idaklu/python.cpp +++ b/pybamm/solvers/c_solvers/idaklu/python.cpp @@ -5,203 +5,211 @@ 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) { + 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]. @@ -223,255 +231,255 @@ int sensitivities(int Ns, realtype t, N_Vector yy, N_Vector yp, // 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/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) { 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..f4855b1bc4 --- /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 diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py index 5ccff7ed14..d9819f1608 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: @@ -84,6 +87,7 @@ def __init__( root_method="casadi", root_tol=1e-6, extrap_tol=None, + output_variables=[], options=None, ): # set default options, @@ -106,6 +110,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") @@ -116,6 +122,7 @@ def __init__( root_method, root_tol, extrap_tol, + output_variables, ) self.name = "IDA KLU solver" @@ -174,6 +181,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 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,) @@ -251,6 +263,30 @@ def resfn(t, y, inputs, ydot): "mass_action", [v_casadi], [casadi.densify(mass_matrix @ v_casadi)] ) + # 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: + # 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 + self.var_idaklu_fcns.append( + 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.computed_dvar_dy_fcns[key].serialize()) + ) + self.dvar_dp_idaklu_fcns.append( + idaklu.generate_function(self.computed_dvar_dp_fcns[key].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) @@ -421,28 +457,36 @@ 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.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, } 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), + 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, ) self._setup["solver"] = solver @@ -555,7 +599,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),) @@ -579,7 +627,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, @@ -589,7 +637,37 @@ 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)) + 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 + ): + continue + len_of_var = ( + self._setup["var_casadi_fcns"][var](0, 0, 0).sparsity().nnz() + ) + newsol._variables[var] = pybamm.ProcessedVariableComputed( + [model.variables_and_events[var]], + [self._setup["var_casadi_fcns"][var]], + [sol.y[:, startk : (startk + len_of_var)]], + newsol, + ) + # Add sensitivities + newsol[var]._sensitivities = {} + if model.calculate_sensitivities: + for paramk, param in enumerate(inputs_dict.keys()): + newsol[var].add_sensitivity( + param, + [sol.yS[:, startk : (startk + len_of_var), paramk]], + ) + startk += len_of_var + return newsol else: raise pybamm.SolverError("idaklu solver failed") diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py index 9c404b72a2..f9d967c4b0 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 @@ -73,9 +72,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 ( @@ -87,15 +85,11 @@ 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 - ): + 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] + 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() @@ -104,7 +98,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), ]: @@ -118,9 +112,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 @@ -146,7 +137,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 @@ -208,9 +199,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 diff --git a/pybamm/solvers/processed_variable_computed.py b/pybamm/solvers/processed_variable_computed.py new file mode 100644 index 0000000000..78d16c27fb --- /dev/null +++ b/pybamm/solvers/processed_variable_computed.py @@ -0,0 +1,441 @@ +# +# Processed Variable class +# +import casadi +import numpy as np +import pybamm +from scipy.integrate import cumulative_trapezoid +import xarray as xr + + +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` + 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 :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 + 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 + self.unroll_params = {} + + # 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 add_sensitivity(self, param, 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 + 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) + .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"] + 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) + 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(f"Unsupported data dimension: {self.dimensions}") + + def initialise_0D(self): + entries = self.unroll_0D() + + 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 initialise_1D(self, fixed_t=False): + entries = self.unroll_1D() + + # 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) + + 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 + 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 = self.unroll_2D( + realdata=None, + n_dim1=len_y, + n_dim2=len_z, + axis_swaps=[(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/pybamm/solvers/solution.py b/pybamm/solvers/solution.py index 4c9ccb993d..d7a27f142c 100644 --- a/pybamm/solvers/solution.py +++ b/pybamm/solvers/solution.py @@ -483,13 +483,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 = 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) + 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( @@ -500,9 +508,9 @@ def update(self, variables): self._variables[key] = var self.data[key] = var.data - def process_casadi_var(self, var_pybamm, inputs, ys): + 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]) + 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/tests/unit/test_solvers/test_idaklu_solver.py b/tests/unit/test_solvers/test_idaklu_solver.py index efd0439f32..cc54f3dfd5 100644 --- a/tests/unit/test_solvers/test_idaklu_solver.py +++ b/tests/unit/test_solvers/test_idaklu_solver.py @@ -85,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) @@ -182,31 +186,33 @@ 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} - - disc = pybamm.Discretisation() - disc.process_model(model) + 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} - solver = pybamm.IDAKLUSolver() + disc = pybamm.Discretisation() + disc.process_model(model) + 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 @@ -540,6 +546,147 @@ 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 = {} # Sensitivities dictionary + 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() + + 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__": print("Add -v for more debug output") diff --git a/tests/unit/test_solvers/test_processed_variable_computed.py b/tests/unit/test_solvers/test_processed_variable_computed.py new file mode 100644 index 0000000000..e31f51ab1e --- /dev/null +++ b/tests/unit/test_solvers/test_processed_variable_computed.py @@ -0,0 +1,442 @@ +# +# 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 +# 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, geometry_options={} +): + # 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) + model = tests.get_base_model_with_battery_geometry(**geometry_options) + pybamm.ProcessedVariableComputed( + [var_sol], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, model, {}), + warn=False, + ) + # NB: ProcessedVariableComputed does not interpret y in the same way as + # ProcessedVariable; a better test of equivalence is to check that the + # 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 + + +class TestProcessedVariableComputed(TestCase): + def test_processed_variable_0D(self): + # without space + 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.ProcessedVariableComputed( + [var], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), {}), + warn=False, + ) + # Assert that the processed variable is the same as the solution + np.testing.assert_array_equal(processed_var.entries, y_sol[0]) + # Check that 'data' produces the same output as 'entries' + np.testing.assert_array_equal(processed_var.entries, processed_var.data) + + # 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 + 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.ProcessedVariableComputed( + [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.ProcessedVariableComputed( + [var], + [var_casadi], + [y_sol], + pybamm.Solution(t_sol, y_sol, pybamm.BaseModel(), inputs), + warn=False, + ) + + # test no sensitivity raises error + self.assertIsNone(processed_var.sensitivities) + + def test_processed_variable_1D(self): + var = pybamm.Variable("var", domain=["negative electrode", "separator"]) + x = pybamm.SpatialVariable("x", domain=["negative electrode", "separator"]) + + # 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) + 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.ProcessedVariableComputed( + [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_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 that there are no errors with domain-specific attributes + # (see ProcessedVariableComputed.initialise_1D() for details) + domain_list = [ + "particle", + "separator", + "current collector", + "particle size", + "random-non-specific-domain", + ] + for domain in domain_list: + processed_var.domain[0] = domain + 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( + {"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.ProcessedVariableComputed([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", + 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.ProcessedVariableComputed( + [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)]), + ) + 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"]) + + 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.ProcessedVariableComputed( + [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.ProcessedVariableComputed( + [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()