Merge pull request #2002 from pybamm-team/i1863-idklu-casadi

I1863 idklu casadi

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Features
 
+-   Added a casadi version of the IDKLU solver, which is used for `model.convert_to_format = "casadi"` ([#2002](https://github.com/pybamm-team/PyBaMM/pull/2002))
+
 ## Bug fixes
 
 -   Remove old deprecation errors, including those in `parameter_values.py` that caused the simulation if, for example, the reaction rate is re-introduced manually ([#2022](https://github.com/pybamm-team/PyBaMM/pull/2022))

CMakeLists.txt

@@ -17,21 +17,45 @@ endif()
 
 project(idaklu)
 
-set (CMAKE_CXX_STANDARD 11)
+set(CMAKE_CXX_STANDARD 14)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CXX_EXTENSIONS OFF)
+set(CMAKE_EXPORT_COMPILE_COMMANDS 1)
 set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 
+# casadi seems to compile without the newer versions of std::string
+add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=0)
+
 if(NOT PYBIND11_DIR)
   set(PYBIND11_DIR pybind11)
 endif()
 
 add_subdirectory(${PYBIND11_DIR})
-pybind11_add_module(idaklu pybamm/solvers/c_solvers/idaklu.cpp)
+pybind11_add_module(idaklu 
+  pybamm/solvers/c_solvers/idaklu_python.cpp
+  pybamm/solvers/c_solvers/idaklu_python.hpp
+  pybamm/solvers/c_solvers/idaklu.cpp
+  pybamm/solvers/c_solvers/idaklu.hpp
+  pybamm/solvers/c_solvers/idaklu_casadi.cpp
+  pybamm/solvers/c_solvers/idaklu_casadi.hpp
+  pybamm/solvers/c_solvers/solution.cpp
+  pybamm/solvers/c_solvers/solution.hpp
+)
+
+if(NOT CASADI_DIR)
+  execute_process(
+    COMMAND "${PYTHON_EXECUTABLE}" -c 
+    "import casadi as _; print(_.__path__[0])"
+    OUTPUT_VARIABLE CASADI_DIR
+    OUTPUT_STRIP_TRAILING_WHITESPACE)
+endif()
+find_package(casadi CONFIG PATHS ${CASADI_DIR} REQUIRED)
 
 set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${PROJECT_SOURCE_DIR})
 # Sundials
 find_package(SUNDIALS)
 target_include_directories(idaklu PRIVATE ${SUNDIALS_INCLUDE_DIR})
-target_link_libraries(idaklu PRIVATE ${SUNDIALS_LIBRARIES})
+target_link_libraries(idaklu PRIVATE ${SUNDIALS_LIBRARIES} casadi)
 
 # link suitesparse
 # if using vcpkg, use config mode to

docs/install/install-from-source.rst

@@ -85,7 +85,7 @@ If you are using MacOS, an alternative to the above is to get the required SUNDI
 
 	  brew install sundials
 
-Next, clone the pybind11 repository:
+Next, clone the pybind11 and casadi-headers repositories:
 
 .. code:: bash
 
@@ -103,6 +103,7 @@ If you'd rather do things yourself,
 2. Compile and install SuiteSparse (PyBaMM only requires the ``KLU`` component).
 3. Compile and install SUNDIALS.
 4. Clone the pybind11 repository in the ``PyBaMM/`` directory (make sure the directory is named ``pybind11``).
+
 
 PyBaMM ships with a python script that automates points 2. and 3. You can run it with
 

pybamm/discretisations/discretisation.py

@@ -1036,14 +1036,13 @@ def _process_symbol(self, symbol):
             return new_symbol
 
         elif isinstance(symbol, pybamm.InputParameter):
-            # Return a new copy of the input parameter, but set the expected size
-            # according to the domain of the input parameter
-            expected_size = self._get_variable_size(symbol)
-            new_input_parameter = pybamm.InputParameter(
-                symbol.name, symbol.domain, expected_size
-            )
-            return new_input_parameter
-
+            if symbol.domain != []:
+                expected_size = self._get_variable_size(symbol)
+            else:
+                expected_size = None
+            if symbol._expected_size is None:
+                symbol._expected_size = expected_size
+            return symbol.create_copy()
         else:
             # Backup option: return the object
             return symbol

pybamm/expression_tree/input_parameter.py

@@ -3,6 +3,7 @@
 #
 import numbers
 import numpy as np
+import scipy.sparse
 import pybamm
 
 
@@ -55,7 +56,16 @@ def _evaluate_for_shape(self):
 
     def _jac(self, variable):
         """See :meth:`pybamm.Symbol._jac()`."""
-        return pybamm.Scalar(0)
+        n_variable = variable.evaluation_array.count(True)
+        nan_vector = self._evaluate_for_shape()
+        if isinstance(nan_vector, numbers.Number):
+            n_self = 1
+        else:
+            n_self = nan_vector.shape[0]
+        zero_matrix = scipy.sparse.csr_matrix(
+            (n_self, n_variable)
+        )
+        return pybamm.Matrix(zero_matrix)
 
     def _base_evaluate(self, t=None, y=None, y_dot=None, inputs=None):
         # inputs should be a dictionary

pybamm/expression_tree/operations/evaluate_python.py

@@ -369,7 +369,7 @@ def find_symbols(symbol, constant_symbols, variable_symbols, output_jax=False):
         symbol_str = "t"
 
     elif isinstance(symbol, pybamm.InputParameter):
-        symbol_str = "inputs['{}']".format(symbol.name)
+        symbol_str = 'inputs["{}"]'.format(symbol.name)
 
     else:
         raise NotImplementedError(
@@ -621,6 +621,9 @@ def get_jacobian(self):
 
         return EvaluatorJaxJacobian(self._jac_evaluate, self._constants)
 
+    def get_jacobian_action(self):
+        return self.jvp
+
     def get_sensitivities(self):
         n = len(self._arg_list)
 
@@ -662,6 +665,22 @@ def __call__(self, t=None, y=None, inputs=None):
 
         return result
 
+    def jvp(self, t=None, y=None, v=None, inputs=None):
+        """
+        evaluate jacobian vector product of function
+        """
+
+        # generated code assumes y is a column vector
+        if y is not None and y.ndim == 1:
+            y = y.reshape(-1, 1)
+        if v is not None and v.ndim == 1:
+            v = v.reshape(-1, 1)
+
+        def bind_t_and_inputs(the_y):
+            return self._jit_evaluate(*self._constants, t, the_y, inputs)
+
+        return jax.jvp(bind_t_and_inputs, (y,), (v,))[1]
+
 
 class EvaluatorJaxJacobian:
     def __init__(self, jac_evaluate, constants):
@@ -690,13 +709,17 @@ def __init__(self, jac_evaluate, constants):
 
     def __call__(self, t=None, y=None, inputs=None):
         """
-        Acts as a drop-in replacement for :func:`pybamm.Symbol.evaluate`
+        evaluate function
         """
         # generated code assumes y is a column vector
         if y is not None and y.ndim == 1:
             y = y.reshape(-1, 1)
 
         # execute code
         result = self._jac_evaluate(*self._constants, t, y, inputs)
+        result = {
+            key: value.reshape(value.shape[0], -1)
+            for key, value in result.items()
+        }
 
         return result

pybamm/solvers/base_solver.py

@@ -297,6 +297,43 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
                 )
 
         def process(symbol, name, use_jacobian=None):
+            """
+            Parameters
+            ----------
+            symbol: :class:`pybamm.Symbol`
+                expression tree to convert
+            name: str
+                function evaluators created will have this base name
+            use_jacobian: bool, optional
+                wether to return jacobian functions
+
+            Returns
+            -------
+            func: :class:`pybamm.EvaluatorPython` or
+                  :class:`pybamm.EvaluatorJax` or
+                  :class:`casadi.Function`
+                evaluator for the function $f(y, t, p)$ given by `symbol`
+
+            jac: :class:`pybamm.EvaluatorPython` or
+                  :class:`pybamm.EvaluatorJaxJacobian` or
+                  :class:`casadi.Function`
+                evaluator for the jacobian $\frac{\partial f}{\partial y}$
+                of the function given by `symbol`
+
+            jacp: :class:`pybamm.EvaluatorPython` or
+                  :class:`pybamm.EvaluatorJaxSensitivities` or
+                  :class:`casadi.Function`
+                evaluator for the parameter sensitivities
+                $\frac{\partial f}{\partial p}$
+                of the function given by `symbol`
+
+            jac_action: :class:`pybamm.EvaluatorPython` or
+                  :class:`pybamm.EvaluatorJax` or
+                  :class:`casadi.Function`
+                evaluator for product of the jacobian with a vector $v$,
+                i.e. $\frac{\partial f}{\partial y} * v$
+            """
+
             def report(string):
                 # don't log event conversion
                 if "event" not in string:
@@ -320,8 +357,10 @@ def report(string):
                 if use_jacobian:
                     report(f"Calculating jacobian for {name} using jax")
                     jac = func.get_jacobian()
+                    jac_action = func.get_jacobian_action()
                 else:
                     jac = None
+                    jac_action = None
 
             elif model.convert_to_format != "casadi":
                 # Process with pybamm functions, converting
@@ -356,11 +395,13 @@ def jacp(*args, **kwargs):
                 if use_jacobian:
                     report(f"Calculating jacobian for {name}")
                     jac = jacobian.jac(symbol, y)
-                    if model.convert_to_format == "python":
-                        report(f"Converting jacobian for {name} to python")
-                        jac = pybamm.EvaluatorPython(jac)
+                    report(f"Converting jacobian for {name} to python")
+                    jac = pybamm.EvaluatorPython(jac)
+                    # cannot do jacobian action efficiently for now
+                    jac_action = None
                 else:
                     jac = None
+                    jac_action = None
 
                 report(f"Converting {name} to python")
                 func = pybamm.EvaluatorPython(symbol)
@@ -444,32 +485,44 @@ def jacp(*args, **kwargs):
                             "CasADi"
                         )
                     )
-                    jacp_dict = {}
-                    for pname in model.calculate_sensitivities:
-                        p_diff = casadi.jacobian(casadi_expression, p_casadi[pname])
-                        jacp_dict[pname] = casadi.Function(
-                            name, [t_casadi, y_and_S, p_casadi_stacked], [p_diff]
-                        )
-
-                    # jacp should be a casadi_expressiontion that returns
-                    # a dict of sensitivities
-                    def jacp(*args, **kwargs):
-                        return {k: v(*args, **kwargs) for k, v in jacp_dict.items()}
+                    # 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")
                     jac_casadi = casadi.jacobian(casadi_expression, y_and_S)
                     jac = casadi.Function(
-                        name, [t_casadi, y_and_S, p_casadi_stacked], [jac_casadi]
+                        name + "_jac", [t_casadi, y_and_S, p_casadi_stacked],
+                        [jac_casadi]
+                    )
+
+                    v = casadi.MX.sym("v", model.len_rhs_and_alg + model.len_rhs_sens +
+                                      model.len_alg_sens)
+                    jac_action_casadi = casadi.densify(
+                        casadi.jtimes(casadi_expression, y_and_S, v)
+                    )
+                    jac_action = casadi.Function(
+                        name + "_jac_action", [t_casadi, y_and_S, p_casadi_stacked, v],
+                        [jac_action_casadi]
                     )
                 else:
                     jac = None
+                    jac_action = None
 
                 func = casadi.Function(
                     name, [t_casadi, y_and_S, p_casadi_stacked], [casadi_expression]
                 )
 
-            return func, jac, jacp
+            return func, jac, jacp, jac_action
 
         # Process initial conditions
         initial_conditions = process(
@@ -558,9 +611,9 @@ def jacp(*args, **kwargs):
 
         # Process rhs, algebraic, residual and event expressions
         # and wrap in callables
-        rhs, jac_rhs, jacp_rhs = process(model.concatenated_rhs, "RHS")
+        rhs, jac_rhs, jacp_rhs, jac_rhs_action = process(model.concatenated_rhs, "RHS")
 
-        algebraic, jac_algebraic, jacp_algebraic = process(
+        algebraic, jac_algebraic, jacp_algebraic, jac_algebraic_action = process(
             model.concatenated_algebraic, "algebraic"
         )
 
@@ -573,7 +626,10 @@ def jacp(*args, **kwargs):
             rhs_algebraic = pybamm.NumpyConcatenation(
                 model.concatenated_rhs, model.concatenated_algebraic
             )
-        rhs_algebraic, jac_rhs_algebraic, jacp_rhs_algebraic = process(
+
+        (rhs_algebraic,
+         jac_rhs_algebraic,
+         jacp_rhs_algebraic, jac_rhs_algebraic_action) = process(
             rhs_algebraic, "rhs_algebraic"
         )
 
@@ -633,12 +689,15 @@ def jacp(*args, **kwargs):
         model.interpolant_extrapolation_events_eval = interpolant_extrapolation_events
 
         model.jac_rhs_eval = jac_rhs
+        model.jac_rhs_action_eval = jac_rhs_action
         model.jacp_rhs_eval = jacp_rhs
 
         model.jac_algebraic_eval = jac_algebraic
+        model.jac_algebraic_action_eval = jac_algebraic_action
         model.jacp_algebraic_eval = jacp_algebraic
 
         model.jac_rhs_algebraic_eval = jac_rhs_algebraic
+        model.jac_rhs_algebraic_action_eval = jac_rhs_algebraic_action
         model.jacp_rhs_algebraic_eval = jacp_rhs_algebraic
 
         # Save CasADi functions for the CasADi solver