Repository: tinosulzer/PyBaMM

Issue:
Improve tree visualisation
**Summary**
Add a method to visualize the tree within an svg file. 

**Reason** 
Is useful for demonstrating what the code is doing and for a user to make manual checks on inputted models.


Discussion:

Patch:
diff --git a/pybamm/expression_tree/symbol.py b/pybamm/expression_tree/symbol.py
index 6f6acfc6..99d5f617 100644
--- a/pybamm/expression_tree/symbol.py
+++ b/pybamm/expression_tree/symbol.py
@@ -9,6 +9,8 @@ import anytree
 import numbers
 import copy
 
+from anytree.exporter import DotExporter
+
 
 class Symbol(anytree.NodeMixin):
     """Base node class for the expression tree
@@ -89,6 +91,57 @@ class Symbol(anytree.NodeMixin):
         for pre, _, node in anytree.RenderTree(self):
             print("%s%s" % (pre, str(node)))
 
+    def visualise(self, filename, test=False):
+        """Produces a .png file of the tree (this node and its children) with the
+        name filename"""
+
+        new_node, counter = self.relabel_tree(self, 0)
+
+        # check that filename ends in .png.
+        filename = "view_tree/" + filename + ".png"
+
+        if test is False:
+            DotExporter(
+                new_node, nodeattrfunc=lambda node: 'label="{}"'.format(node.label)
+            ).to_picture(filename)
+
+    def relabel_tree(self, symbol, counter):
+        """ Finds all children of a symbol and assigns them a new id so that they can be
+                visualised properly using the graphviz output
+        """
+        name = symbol.name
+        if name == "div":
+            name = "&nabla;&sdot;"
+        elif name == "grad":
+            name = "&nabla;"
+        elif name == "/":
+            name = "&divide;"
+        elif name == "*":
+            name = "&times;"
+        elif name == "-":
+            name = "&minus;"
+        elif name == "+":
+            name = "&#43;"
+        elif name == "**":
+            name = "^"
+        elif name == "epsilon_s":
+            name = "&#603;"
+
+        new_node = anytree.Node(str(counter), label=name)
+        counter += 1
+
+        if isinstance(symbol, pybamm.BinaryOperator):
+            left, right = symbol.children
+            new_left, counter = self.relabel_tree(left, counter)
+            new_right, counter = self.relabel_tree(right, counter)
+            new_node.children = [new_left, new_right]
+
+        elif isinstance(symbol, pybamm.UnaryOperator):
+            new_child, counter = self.relabel_tree(symbol.children[0], counter)
+            new_node.children = [new_child]
+
+        return new_node, counter
+
     def pre_order(self):
         """returns an iterable that steps through the tree in pre-order
         fashion
diff --git a/pybamm/models/submodels/electrolyte.py b/pybamm/models/submodels/electrolyte.py
index b41b5e50..acad18a1 100644
--- a/pybamm/models/submodels/electrolyte.py
+++ b/pybamm/models/submodels/electrolyte.py
@@ -44,11 +44,9 @@ class StefanMaxwellDiffusion(pybamm.BaseModel):
 
         c_e = pybamm.Variable("c_e", domain=electrolyte_domain)
 
-        N_e = -(epsilon ** b) * pybamm.Gradient(c_e)
+        N_e = -(epsilon ** b) * pybamm.grad(c_e)
 
-        self.rhs = {
-            c_e: -pybamm.Divergence(N_e) / delta / epsilon + nu * (1 - t_plus) * G
-        }
+        self.rhs = {c_e: -pybamm.div(N_e) / delta / epsilon + nu * (1 - t_plus) * G}
         self.initial_conditions = {c_e: ce0}
         self.boundary_conditions = {
             N_e: {"left": pybamm.Scalar(0), "right": pybamm.Scalar(0)}


Test patch:
diff --git a/tests/test_base_model.py b/tests/test_base_model.py
new file mode 100644
index 00000000..94dddd06
--- /dev/null
+++ b/tests/test_base_model.py
@@ -0,0 +1,73 @@
+#
+# Tests for the base model class
+#
+import pybamm
+
+import unittest
+
+
+class TestBaseModel(unittest.TestCase):
+    def test_rhs_set_get(self):
+        model = pybamm.BaseModel()
+        rhs = {
+            pybamm.Symbol("c"): pybamm.Symbol("alpha"),
+            pybamm.Symbol("d"): pybamm.Symbol("beta"),
+        }
+        model.rhs = rhs
+        self.assertEqual(rhs, model.rhs)
+        # test domains
+        rhs = {
+            pybamm.Symbol("c", domain=["negative electrode"]): pybamm.Symbol(
+                "alpha", domain=["negative electrode"]
+            ),
+            pybamm.Symbol("d", domain=["positive electrode"]): pybamm.Symbol(
+                "beta", domain=["positive electrode"]
+            ),
+        }
+        model.rhs = rhs
+        self.assertEqual(rhs, model.rhs)
+        # non-matching domains should fail
+        with self.assertRaises(pybamm.DomainError):
+            model.rhs = {
+                pybamm.Symbol("c", domain=["positive electrode"]): pybamm.Symbol(
+                    "alpha", domain=["negative electrode"]
+                )
+            }
+
+    def test_initial_conditions_set_get(self):
+        model = pybamm.BaseModel()
+        initial_conditions = {
+            pybamm.Symbol("c0"): pybamm.Symbol("gamma"),
+            pybamm.Symbol("d0"): pybamm.Symbol("delta"),
+        }
+        model.initial_conditions = initial_conditions
+        self.assertEqual(initial_conditions, model.initial_conditions)
+
+    def test_boundary_conditions_set_get(self):
+        model = pybamm.BaseModel()
+        boundary_conditions = {"c_left": "epsilon", "c_right": "eta"}
+        model.boundary_conditions = boundary_conditions
+        self.assertEqual(boundary_conditions, model.boundary_conditions)
+
+    def test_variables_set_get(self):
+        model = pybamm.BaseModel()
+        variables = {"c": "alpha", "d": "beta"}
+        model.variables = variables
+        self.assertEqual(variables, model.variables)
+
+    def test_model_dict_behaviour(self):
+        model = pybamm.BaseModel()
+        key = pybamm.Symbol("c")
+        rhs = {key: pybamm.Symbol("alpha")}
+        model.rhs = rhs
+        self.assertEqual(model[key], rhs[key])
+        self.assertEqual(model[key], model.rhs[key])
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    unittest.main()
diff --git a/tests/test_expression_tree/test_symbol.py b/tests/test_expression_tree/test_symbol.py
index b787d7f7..147a8477 100644
--- a/tests/test_expression_tree/test_symbol.py
+++ b/tests/test_expression_tree/test_symbol.py
@@ -136,6 +136,13 @@ class TestSymbol(unittest.TestCase):
             r"Symbol\(" + hex_regex + ", a, grad\(a\)\)",
         )
 
+    def test_symbol_visualise(self):
+        G = pybamm.Symbol("G")
+        model = pybamm.electrolyte.StefanMaxwellDiffusion(G)
+        c_e = list(model.rhs.keys())[0]
+        rhs = model.rhs[c_e]
+        rhs.visualise("StefanMaxwell_test", test=True)
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")