Add test cases

theano/d3viz/d3viz.py

 """Dynamic visualization of Theano graphs.
 
-Author: Christof Angermueller <cangermueller@gmail.com
+Author: Christof Angermueller <cangermueller@gmail.com>
 """
 
 import os
@@ -19,19 +19,34 @@ def replace_patterns(x, replace):
 
 
 def d3write(fct, path, *args, **kwargs):
-    """Convert theano graph to pydot graph and write to file."""
+    """Convert Theano graph to pydot graph and write to file."""
     gf = PyDotFormatter(*args, **kwargs)
     g = gf(fct)
     g.write_dot(path)
 
 
 def d3viz(fct, outfile,  *args, **kwargs):
-    """Create dynamic graph visualization using d3.js javascript library.
+    """Create HTML file with dynamic visualizing of a Theano function graph.
 
     :param fct: A compiled Theano function, variable, apply or a list of
                 variables.
-    :param outfile: The output file
-    :param args, kwargs: Additional parameters passed to PyDotFromatter.
+    :param outfile: The output HTML file
+    :param compact=False: if True, will remove intermediate variables without
+                          name.
+
+    In the HTML file, the whole graph or single nodes can be moved by drag and
+    drop. Zooming is possible via the mouse wheel. Detailed information about
+    nodes and edges are displayed via mouse-over events. Node labels can be
+    edited by selecting Edit from the context menu.
+
+    Input nodes are colored in green, output nodes in blue. Apply nodes are
+    ellipses, and colored depending on the type of operation they perform. Red
+    ellipses are transfers from/to the GPU (ops with names GpuFromHost,
+    HostFromGpu).
+
+    Edges are black by default. If a node returns a view of an
+    input, the input edge will be blue. If it returns a destroyed input, the
+    edge will be red.
     """
 
     outdir = pt.dirname(outfile)

theano/d3viz/formatting.py

 """Functions for formatting Theano compute graphs.
 
-Author: Christof Angermueller <cangermueller@gmail.com
+Author: Christof Angermueller <cangermueller@gmail.com>
 """
 
 import numpy as np
@@ -48,6 +48,9 @@ class PyDotFormatter(object):
                              'Elemwise': '#FFAABB',  # dark pink
                              'Subtensor': '#FFAAFF',  # purple
                              'Alloc': '#FFAA22'}  # orange
+        self.shapes = {'input': 'box',
+                       'output': 'box',
+                       'apply': 'ellipse'}
         self.__node_prefix = 'n'
 
     def __add_node(self, node):
@@ -117,6 +120,7 @@ class PyDotFormatter(object):
             nparams['profile'] = apply_profile(node, profile)
             nparams['node_type'] = 'apply'
             nparams['apply_op'] = apply_op(node)
+            nparams['shape'] = self.shapes['apply']
 
             use_color = None
             for opName, color in self.apply_colors.items():
@@ -148,7 +152,7 @@ class PyDotFormatter(object):
                     vparams['style'] = 'filled'
                     vparams['fillcolor'] = self.node_colors[
                         vparams['node_type']]
-                    vparams['shape'] = 'ellipse'
+                    vparams['shape'] = self.shapes['input']
                     pd_var = dict_to_pdnode(vparams)
                     graph.add_node(pd_var)
 
@@ -185,7 +189,7 @@ class PyDotFormatter(object):
                         vparams['fillcolor'] = self.node_colors['unused']
                     else:
                         vparams['fillcolor'] = self.node_colors['output']
-                    vparams['shape'] = 'box'
+                    vparams['shape'] = self.shapes['output']
                     pd_var = dict_to_pdnode(vparams)
                     graph.add_node(pd_var)
 

theano/d3viz/js/d3viz.js

 /*
- * Theano javascript library for interactive visualiztion.
+ * Theano javascript library for interactive visualization.
  *
- * Author: Christof Angermueller <cangermueller@gmail.com
+ * Author: Christof Angermueller <cangermueller@gmail.com>
 */
 
 

theano/d3viz/tests/models.py

+import numpy as np
+
+import theano as th
+import theano.tensor as T
+
+
+class Mlp(object):
+
+    def __init__(self, nfeatures=100, noutputs=10, nhiddens=50, rng=None):
+        self.rng = rng
+        if self.rng is None:
+            self.rng = np.random.RandomState(0)
+        self.nfeatures = nfeatures
+        self.noutputs = noutputs
+        self.nhiddens = nhiddens
+
+        x = T.dmatrix('x')
+        wh = th.shared(self.rng.normal(0, 1, (nfeatures, nhiddens)),
+                       borrow=True)
+        bh = th.shared(np.zeros(nhiddens), borrow=True)
+        h = T.nnet.sigmoid(T.dot(x, wh) + bh)
+
+        wy = th.shared(self.rng.normal(0, 1, (nhiddens, noutputs)))
+        by = th.shared(np.zeros(noutputs), borrow=True)
+        y = T.nnet.softmax(T.dot(h, wy) + by)
+
+        self.inputs = [x]
+        self.outputs = [y]
+
+
+class OfgNested(object):
+
+    def __init__(self):
+        x, y, z = T.scalars('xyz')
+        e = x * y
+        op = th.OpFromGraph([x, y], [e])
+        e2 = op(x, y) + z
+        op2 = th.OpFromGraph([x, y, z], [e2])
+        e3 = op2(x, y, z) + z
+
+        self.inputs = [x, y, z]
+        self.outputs = [e3]
+
+
+class Ofg(object):
+
+    def __init__(self):
+        x, y, z = T.scalars('xyz')
+        e = T.nnet.sigmoid((x + y + z)**2)
+        op = th.OpFromGraph([x, y, z], [e])
+        e2 = op(x, y, z) + op(z, y, x)
+
+        self.inputs = [x, y, z]
+        self.outputs = [e2]

theano/d3viz/tests/test_d3viz.py

+import numpy as np
+import os.path as pt
+from tempfile import mkstemp
+import unittest
+
+import theano as th
+import theano.d3viz as d3v
+from theano.d3viz.tests import models
+
+
+class TestD3Viz(unittest.TestCase):
+
+    def setUp(self):
+        self.rng = np.random.RandomState(0)
+
+    def check(self, f):
+        _, html_file = mkstemp('.html')
+        dot_file = html_file.replace('.html', '.dot')
+        d3v.d3viz(f, html_file)
+        assert pt.getsize(html_file) > 0
+        assert pt.getsize(dot_file) > 0
+
+    def test_mlp(self):
+        m = models.Mlp(rng=self.rng)
+        f = th.function(m.inputs, m.outputs)
+        self.check(f)
+
+    def test_mlp_profiled(self):
+        m = models.Mlp(rng=self.rng)
+        f = th.function(m.inputs, m.outputs, profile=True)
+        x_val = self.rng.normal(0, 1, (1000, m.nfeatures))
+        f(x_val)
+        self.check(f)
+
+    def test_ofg(self):
+        m = models.Ofg()
+        f = th.function(m.inputs, m.outputs)
+        self.check(f)
+
+    def test_ofg_nested(self):
+        m = models.OfgNested()
+        f = th.function(m.inputs, m.outputs)
+        self.check(f)

theano/d3viz/tests/test_formatting.py

+import numpy as np
+import unittest
+
+import theano as th
+from theano.d3viz.formatting import PyDotFormatter
+from theano.d3viz.tests import models
+
+
+class TestPyDotFormatter(unittest.TestCase):
+
+    def setUp(self):
+        self.rng = np.random.RandomState(0)
+
+    def node_counts(self, graph):
+        node_types = [node.get_attributes()['node_type']
+                      for node in graph.get_nodes()]
+        a, b  = np.unique(node_types, return_counts=True)
+        nc = dict(zip(a, b))
+        return nc
+
+    def test_mlp(self):
+        m = models.Mlp()
+        f = th.function(m.inputs, m.outputs)
+        pdf = PyDotFormatter()
+        graph = pdf(f)
+        self.assertEqual(len(graph.get_nodes()), 11)
+        nc = self.node_counts(graph)
+        assert nc['apply'] == 5
+        assert nc['output'] == 1
+
+    def test_ofg(self):
+        m = models.Ofg()
+        f = th.function(m.inputs, m.outputs)
+        pdf = PyDotFormatter()
+        graph = pdf(f)
+        assert len(graph.get_nodes()) == 10
+        sub_graphs = graph.get_subgraph_list()
+        assert len(sub_graphs) == 2
+        ofg1, ofg2 = sub_graphs
+        assert len(ofg1.get_nodes()) == 5
+        assert len(ofg1.get_nodes()) == len(ofg2.get_nodes())
+
+    def test_ofg_nested(self):
+        m = models.OfgNested()
+        f = th.function(m.inputs, m.outputs)
+        pdf = PyDotFormatter()
+        graph = pdf(f)
+        assert len(graph.get_nodes()) == 7
+        assert len(graph.get_subgraph_list()) == 1
+        ofg1 = graph.get_subgraph_list()[0]
+        assert len(ofg1.get_nodes()) == 6
+        assert len(ofg1.get_subgraph_list()) == 1
+        ofg2 = ofg1.get_subgraph_list()[0]
+        assert len(ofg2.get_nodes()) == 4