#664 add casadi set up to solver

examples/scripts/compare_lithium_ion.py

@@ -18,9 +18,9 @@
 # load models
 options = {"thermal": "isothermal"}
 models = [
-    # pybamm.lithium_ion.SPM(options),
-    # pybamm.lithium_ion.SPMe(options),
-    pybamm.lithium_ion.DFN(options)
+    pybamm.lithium_ion.SPM(options),
+    pybamm.lithium_ion.SPMe(options),
+    pybamm.lithium_ion.DFN(options),
 ]
 
 

pybamm/discretisations/discretisation.py

@@ -127,7 +127,7 @@ def process_model(self, model, inplace=True):
             model_disc.options = model.options
             model_disc.use_jacobian = model.use_jacobian
             model_disc.use_simplify = model.use_simplify
-            model_disc.use_to_python = model.use_to_python
+            model_disc.convert_to_format = model.convert_to_format
 
         model_disc.bcs = self.bcs
 

pybamm/models/base_model.py

@@ -59,11 +59,17 @@ class BaseModel(object):
         Whether to simplify the expression tress representing the rhs and
         algebraic equations, Jacobain (if using) and events, before solving the
         model (default is True)
-    use_to_python : bool
-        Whether to convert the expression tress representing the rhs and
-        algebraic equations, Jacobain (if using) and events into pure python code
-        that will calculate the result of calling `evaluate(t, y)` on the given
-        expression tree (default is True)
+    convert_to_format : str
+        Whether to convert the expression trees representing the rhs and
+        algebraic equations, Jacobain (if using) and events into a different format:
+
+        - None: keep PyBaMM expression tree structure.
+        - "python": convert into pure python code that will calculate the result of
+        calling `evaluate(t, y)` on the given expression treeself.
+        - "casadi": convert into CasADi expression tree, which then uses CasADi's
+        algorithm to calculate the Jacobian.
+
+        Default is "python".
     """
 
     def __init__(self, name="Unnamed model"):
@@ -86,7 +92,7 @@ def __init__(self, name="Unnamed model"):
         # Default behaviour is to use the jacobian and simplify
         self.use_jacobian = True
         self.use_simplify = True
-        self.use_to_python = True
+        self.convert_to_format = "python"
 
     def _set_dictionary(self, dict, name):
         """

pybamm/solvers/algebraic_solver.py

@@ -199,14 +199,14 @@ def set_up(self, model):
                 pybamm.logger.info("Simplifying jacobian")
                 jac = jac.simplify()
 
-            if model.use_to_python:
+            if model.convert_to_format == "python":
                 pybamm.logger.info("Converting jacobian to python")
                 jac = pybamm.EvaluatorPython(jac)
 
         else:
             jac = None
 
-        if model.use_to_python:
+        if model.convert_to_format == "python":
             pybamm.logger.info("Converting algebraic to python")
             concatenated_algebraic = pybamm.EvaluatorPython(concatenated_algebraic)
 

pybamm/solvers/base_solver.py

@@ -73,7 +73,10 @@ def solve(self, model, t_eval):
         # Set up
         timer = pybamm.Timer()
         start_time = timer.time()
-        self.set_up(model)
+        if model.convert_to_format == "casadi":
+            self.set_up_casadi(model)
+        else:
+            self.set_up(model)
         set_up_time = timer.time() - start_time
 
         # Solve
@@ -127,7 +130,10 @@ def step(self, model, dt, npts=2):
         # Run set up on first step
         if not hasattr(self, "y0"):
             start_time = timer.time()
-            self.set_up(model)
+            if model.convert_to_format == "casadi":
+                self.set_up_casadi(model)
+            else:
+                self.set_up(model)
             self.t = 0.0
             set_up_time = timer.time() - start_time
         else:
@@ -187,11 +193,17 @@ def set_up(self, model):
             The model whose solution to calculate. Must have attributes rhs and
             initial_conditions
 
-        Raises
-        ------
-        :class:`pybamm.SolverError`
-            If the model contains any algebraic equations (in which case a DAE solver
-            should be used instead)
+        """
+        raise NotImplementedError
+
+    def set_up_casadi(self, model):
+        """Convert model to casadi format and use their inbuilt functionalities.
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel`
+            The model whose solution to calculate. Must have attributes rhs and
+            initial_conditions
 
         """
         raise NotImplementedError

pybamm/solvers/dae_solver.py

@@ -140,7 +140,7 @@ def set_up(self, model):
                 jac_algebraic = simp.simplify(jac_algebraic)
                 jac = simp.simplify(jac)
 
-            if model.use_to_python:
+            if model.convert_to_format == "python":
                 pybamm.logger.info("Converting jacobian to python")
                 jac_algebraic = pybamm.EvaluatorPython(jac_algebraic)
                 jac = pybamm.EvaluatorPython(jac)
@@ -152,7 +152,7 @@ def jac_alg_fn(t, y):
             jac = None
             jac_alg_fn = None
 
-        if model.use_to_python:
+        if model.convert_to_format == "python":
             pybamm.logger.info("Converting RHS to python")
             concatenated_rhs = pybamm.EvaluatorPython(concatenated_rhs)
             pybamm.logger.info("Converting algebraic to python")

pybamm/solvers/ode_solver.py

@@ -99,13 +99,13 @@ def set_up(self, model):
                 pybamm.logger.info("Simplifying jacobian")
                 jac_rhs = simp.simplify(jac_rhs)
 
-            if model.use_to_python:
+            if model.convert_to_format == "python":
                 pybamm.logger.info("Converting jacobian to python")
                 jac_rhs = pybamm.EvaluatorPython(jac_rhs)
         else:
             jac_rhs = None
 
-        if model.use_to_python:
+        if model.convert_to_format == "python":
             pybamm.logger.info("Converting RHS to python")
             concatenated_rhs = pybamm.EvaluatorPython(concatenated_rhs)
             pybamm.logger.info("Converting events to python")
@@ -145,6 +145,83 @@ def jacobian(t, y):
         self.event_funs = event_funs
         self.jacobian = jacobian
 
+    def set_up_casadi(self, model):
+        """Convert model to casadi format and use their inbuilt functionalities.
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel`
+            The model whose solution to calculate. Must have attributes rhs and
+            initial_conditions
+
+        Raises
+        ------
+        :class:`pybamm.SolverError`
+            If the model contains any algebraic equations (in which case a DAE solver
+            should be used instead)
+
+        """
+        # Check for algebraic equations
+        if len(model.algebraic) > 0:
+            raise pybamm.SolverError(
+                """Cannot use ODE solver to solve model with DAEs"""
+            )
+
+        # create simplified rhs and event expressions
+        concatenated_rhs = model.concatenated_rhs
+        events = model.events
+
+        y0 = model.concatenated_initial_conditions[:, 0]
+
+        t_casadi = casadi.SX.sym("t")
+        y_casadi = casadi.SX.sym("y", len(y0))
+        pybamm.logger.info("Converting RHS to CasADi")
+        concatenated_rhs = concatenated_rhs.to_casadi(t_casadi, y_casadi)
+        pybamm.logger.info("Converting events to CasADi")
+        events = {
+            name: event.to_casadi(t_casadi, y_casadi) for name, event in events.items()
+        }
+
+        # Create function to evaluate rhs
+        concatenated_rhs_fn = casadi.Function(
+            "rhs", [t_casadi, y_casadi], [concatenated_rhs]
+        )
+
+        def dydt(t, y):
+            pybamm.logger.debug("Evaluating RHS for {} at t={}".format(model.name, t))
+            dy = concatenated_rhs_fn(t, y).full()
+            return dy[:, 0]
+
+        # Create event-dependent function to evaluate events
+        def event_fun(event):
+            event_fn = casadi.Function("event", [t_casadi, y_casadi], [event])
+
+            def eval_event(t, y):
+                return event_fun(t, y)
+
+            return eval_event
+
+        event_funs = [event_fun(event) for event in events.values()]
+
+        # Create function to evaluate jacobian
+        if model.use_jacobian:
+
+            casadi_jac = casadi.jacobian()
+            casadi_jac_fn = casadi.Function("jac", [t_casadi, y_casadi], [casadi_jac])
+
+            def jacobian(t, y):
+                return casadi_jac_fn(t, y)
+
+        else:
+            jacobian = None
+
+        # Add the solver attributes
+        self.y0 = y0
+        self.dydt = dydt
+        self.events = events
+        self.event_funs = event_funs
+        self.jacobian = jacobian
+
     def integrate(
         self, derivs, y0, t_eval, events=None, mass_matrix=None, jacobian=None
     ):

tests/integration/test_models/standard_model_tests.py

@@ -172,6 +172,7 @@ def evaluate_model(self, simplify=False, use_known_evals=False, to_python=False)
 
     def set_up_model(self, simplify=False, to_python=False):
         self.model.use_simplify = simplify
-        self.model.use_to_python = to_python
+        if to_python is True:
+            self.model.convert_to_format = "python"
         self.model.default_solver.set_up(self.model)
         return None