#664 convert scalars and operations to casadi

Author: valentinsulzer

.requirements-docs.txt

@@ -5,5 +5,6 @@ pandas>=0.23
 anytree>=2.4.3
 autograd>=1.2
 scikit-fem>=0.2.0
+casadi>=3.5.0
 guzzle-sphinx-theme
 sphinx>=1.5

docs/source/expression_tree/index.rst

@@ -15,7 +15,6 @@ Expression Tree
   unary_operator
   concatenations
   broadcasts
-  simplify
   functions
   interpolant
-  evaluate
+  operations/index

docs/source/expression_tree/operations/convert_to_casadi.rst

@@ -0,0 +1,5 @@
+Convert to CasADi
+=================
+
+.. autoclass:: pybamm.CasadiConverter
+  :members:

docs/source/expression_tree/evaluate.rst docs/source/expression_tree/operations/evaluate.rst

@@ -0,0 +1,10 @@
+Operations on expression trees
+==============================
+
+Classes and functions that operate on the expression tree
+
+.. toctree::
+
+  simplify
+  evaluate
+  convert_to_casadi

docs/source/expression_tree/operations/index.rst

@@ -0,0 +1,5 @@
+Casadi Solver
+=============
+
+.. autoclass:: pybamm.CasadiSolver
+  :members:

docs/source/expression_tree/simplify.rst docs/source/expression_tree/operations/simplify.rst

@@ -7,4 +7,5 @@ Solvers
   base_solvers
   scipy_solver
   scikits_solvers
+  casadi_solver
   solution

docs/source/solvers/casadi_solver.rst

@@ -150,18 +150,21 @@ def version(formatted=False):
     UndefinedOperationError,
     GeometryError,
 )
-from .expression_tree.simplify import (
+
+# Operations
+from .expression_tree.operations.simplify import (
     Simplification,
     simplify_if_constant,
     simplify_addition_subtraction,
     simplify_multiplication_division,
 )
-from .expression_tree.evaluate import (
+from .expression_tree.operations.evaluate import (
     find_symbols,
     id_to_python_variable,
     to_python,
     EvaluatorPython,
 )
+from .expression_tree.operations.convert_to_casadi import CasadiConverter
 
 #
 # Model classes

docs/source/solvers/index.rst

@@ -0,0 +1,75 @@
+#
+# Convert a PyBaMM expression tree to a CasADi expression tree
+#
+import pybamm
+import casadi
+
+
+class CasadiConverter(object):
+    def __init__(self, casadi_symbols=None):
+        self._casadi_symbols = casadi_symbols or {}
+
+    def convert(self, symbol):
+        """
+        This function recurses down the tree, applying any simplifications defined in
+        classes derived from pybamm.Symbol. E.g. any expression multiplied by a
+        pybamm.Scalar(0) will be simplified to a pybamm.Scalar(0).
+        If a symbol has already been simplified, the stored value is returned.
+
+        Parameters
+        ----------
+        symbol : :class:`pybamm.Symbol`
+            The symbol to convert
+
+        Returns
+        -------
+        CasADi symbol
+            The convert symbol
+        """
+
+        try:
+            return self._casadi_symbols[symbol.id]
+        except KeyError:
+            casadi_symbol = self._convert(symbol)
+            self._casadi_symbols[symbol.id] = casadi_symbol
+
+            return casadi_symbol
+
+    def _convert(self, symbol):
+        """ See :meth:`Simplification.convert()`. """
+        if isinstance(symbol, pybamm.Scalar):
+            return casadi.SX(symbol.evaluate())
+
+        if isinstance(symbol, pybamm.BinaryOperator):
+            left, right = symbol.children
+            # process children
+            converted_left = self.convert(left)
+            converted_right = self.convert(right)
+            # _binary_evaluate defined in derived classes for specific rules
+            return symbol._binary_evaluate(converted_left, converted_right)
+
+        elif isinstance(symbol, pybamm.UnaryOperator):
+            converted_child = self.convert(symbol.child)
+            if isinstance(symbol, pybamm.AbsoluteValue):
+                return casadi.fabs(converted_child)
+            return symbol._unary_evaluate(converted_child)
+
+        elif isinstance(symbol, pybamm.Function):
+            converted_children = [None] * len(symbol.children)
+            for i, child in enumerate(symbol.children):
+                converted_children[i] = self.convert(child)
+            return symbol._function_evaluate(converted_children)
+
+        elif isinstance(symbol, pybamm.Concatenation):
+            converted_children = [self.convert(child) for child in symbol.children]
+            return symbol._concatenation_evaluate(converted_children)
+
+        else:
+            raise TypeError(
+                """
+                Cannot convert symbol of type '{}' to CasADi. Symbols must all be
+                'linear algebra' at this stage.
+                """.format(
+                    type(symbol)
+                )
+            )

pybamm/__init__.py

@@ -584,6 +584,13 @@ def simplify(self, simplified_symbols=None):
         """ Simplify the expression tree. See :class:`pybamm.Simplification`. """
         return pybamm.Simplification(simplified_symbols).simplify(self)
 
+    def to_casadi(self, casadi_symbols=None):
+        """
+        Convert the expression tree to a CasADi expression tree.
+        See :class:`pybamm.CasadiConverter`.
+        """
+        return pybamm.CasadiConverter(casadi_symbols).convert(self)
+
     def new_copy(self):
         """
         Make a new copy of a symbol, to avoid Tree corruption errors while bypassing

pybamm/expression_tree/operations/__init__.py

@@ -0,0 +1,9 @@
+#
+# Wrap CasADi
+#
+import pybamm
+import casadi
+
+
+class CasadiSolver(pybamm.DaeSolver):
+    pass

pybamm/expression_tree/operations/convert_to_casadi.py

@@ -45,6 +45,7 @@ def load_version():
         "anytree>=2.4.3",
         "autograd>=1.2",
         "scikit-fem>=0.2.0",
+        "casadi>=3.5.0",
         # Note: Matplotlib is loaded for debug plots, but to ensure pybamm runs
         # on systems without an attached display, it should never be imported
         # outside of plot() methods.

pybamm/expression_tree/evaluate.py pybamm/expression_tree/operations/evaluate.py

@@ -0,0 +1,62 @@
+#
+# Test for the Simplify class
+#
+import casadi
+import math
+import numpy as np
+import pybamm
+import unittest
+from tests import get_discretisation_for_testing
+
+
+class TestCasadiConverter(unittest.TestCase):
+    def test_convert_scalar_symbols(self):
+        a = pybamm.Scalar(0)
+        b = pybamm.Scalar(1)
+        c = pybamm.Scalar(-1)
+        d = pybamm.Scalar(2)
+
+        self.assertEqual(a.to_casadi(), casadi.SX(0))
+        self.assertEqual(d.to_casadi(), casadi.SX(2))
+
+        # negate
+        self.assertEqual((-b).to_casadi(), casadi.SX(-1))
+        # absolute value
+        self.assertEqual(abs(c).to_casadi(), casadi.SX(1))
+
+        # function
+        def sin(x):
+            return np.sin(x)
+
+        f = pybamm.Function(sin, b)
+        self.assertEqual((f).to_casadi(), casadi.SX(np.sin(1)))
+
+        def myfunction(x, y):
+            return x + y
+
+        f = pybamm.Function(myfunction, b, d)
+        self.assertEqual((f).to_casadi(), casadi.SX(3))
+
+        # addition
+        self.assertEqual((a + b).to_casadi(), casadi.SX(1))
+        # subtraction
+        self.assertEqual((c - d).to_casadi(), casadi.SX(-3))
+        # multiplication
+        self.assertEqual((c * d).to_casadi(), casadi.SX(-2))
+        # power
+        self.assertEqual((c ** d).to_casadi(), casadi.SX(1))
+        # division
+        self.assertEqual((b / d).to_casadi(), casadi.SX(1 / 2))
+
+    def test_convert_array_symbols(self):
+        pass
+
+
+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()