From 97f832031a40eeffa477ef753263000c7c6a23ac Mon Sep 17 00:00:00 2001
From: Brendan Smithyman <brendan@bitsmithy.net>
Date: Tue, 5 May 2015 17:03:45 -0400
Subject: [PATCH] Created SystemGraph subclass of the networkx.DiGraph class
 that knows how to render itself in the IPython Notebook using d3.js.

---
 SimPEG/Parallel.py                         | 50 +++++++++++++++-
 SimPEG/Resources/Parallel/SystemGraph.html | 66 ++++++++++++++++++++++
 2 files changed, 115 insertions(+), 1 deletion(-)
 create mode 100644 SimPEG/Resources/Parallel/SystemGraph.html

diff --git a/SimPEG/Parallel.py b/SimPEG/Parallel.py
index 7056da05..73c02408 100644
--- a/SimPEG/Parallel.py
+++ b/SimPEG/Parallel.py
@@ -53,6 +53,54 @@ def adjustMKLVectorization(nt=1):
     finally:
         mkl.set_num_threads(nt)
 
+class SystemGraph(networkx.DiGraph):
+
+    def _codeStatus(self, data):
+        
+        status = 0
+        
+        if 'jobs' in data:
+            status = 1 * data['jobs'][-1].ready() + 1
+            if status > 1:
+                status += 1 * (not data['jobs'][-1].successful())
+        
+        return status 
+
+    def _codeGraph(self):
+        from networkx.readwrite import json_graph
+        
+        G = networkx.DiGraph()
+        
+        for e in self.edges_iter():
+            G.add_edge(e[0], e[1])
+        
+        for n, data in self.nodes_iter(data=True):
+            G.add_node(n, status=self._codeStatus(data))
+        
+        return json_graph.node_link_data(G)
+
+    def RenderHTML(self):
+        import pkg_resources
+        from IPython.core import display
+        import time
+
+        data = str(self._codeGraph())
+        uniqueID = hash(time.time())
+
+        formatstr = {
+            'uniqueID': 'Graph%s'%uniqueID,
+            'JSONData': data,
+        }
+
+        code = pkg_resources.resource_string('SimPEG', 'Resources/Parallel/SystemGraph.html')%formatstr
+
+        return display.HTML(data=code)._repr_html_()
+
+try:
+    get_ipython().display_formatter.formatters['text/html'].for_type(SystemGraph, SystemGraph.RenderHTML)
+except NameError:
+    pass
+
 class SystemSolver(object):
 
     def __init__(self, dispatcher, schedule):
@@ -71,7 +119,7 @@ class SystemSolver(object):
 
         chunksPerWorker = getattr(self.dispatcher, 'chunksPerWorker', 1)
 
-        G = networkx.DiGraph()
+        G = SystemGraph()
 
         mainNode = 'Beginning'
         G.add_node(mainNode)
diff --git a/SimPEG/Resources/Parallel/SystemGraph.html b/SimPEG/Resources/Parallel/SystemGraph.html
new file mode 100644
index 00000000..fa9e17b1
--- /dev/null
+++ b/SimPEG/Resources/Parallel/SystemGraph.html
@@ -0,0 +1,66 @@
+<div id="%(uniqueID)s"></div>
+<style>
+.node {stroke: #fff; stroke-width: 1.5px;}
+.link {stroke: #999; stroke-opacity: .6;}
+</style>
+
+<script>
+var True = true;
+var False = false;
+var graph = %(JSONData)s;
+
+require.config({paths: {d3: "http://d3js.org/d3.v3.min"}});
+require(["d3"], function(d3) {
+
+    var width = 800, height = 300;
+
+    var color = d3.scale.category10();
+    var domain = [0, 1, 2, 3];
+    color.domain(domain);
+
+    var force = d3.layout.force()
+        .charge(-120)
+        .linkDistance(30)
+        .size([width, height]);
+
+    var svg = d3.select("#%(uniqueID)s").select("svg");
+    if (svg.empty()) {
+        svg = d3.select("#%(uniqueID)s").append("svg")
+                    .attr("width", width)
+                    .attr("height", height);
+    };
+
+    force.nodes(graph.nodes)
+        .links(graph.links)
+        .start();
+
+    var link = svg.selectAll(".link")
+        .data(graph.links)
+        .enter().append("line")
+        .attr("class", "link");
+
+    var node = svg.selectAll(".node")
+        .data(graph.nodes)
+        .enter().append("circle")
+        .attr("class", "node")
+        .attr("r", 5)
+        .style("fill", function(d) {
+            return color(d.status); 
+        })
+        .call(force.drag);
+
+    node.append("title")
+        .text(function(d) { return d.id; });
+
+    force.on("tick", function() {
+        link.attr("x1", function(d) { return d.source.x; })
+            .attr("y1", function(d) { return d.source.y; })
+            .attr("x2", function(d) { return d.target.x; })
+            .attr("y2", function(d) { return d.target.y; });
+
+        node.attr("cx", function(d) { return d.x; })
+            .attr("cy", function(d) { return d.y; });
+    });
+
+});
+</script>
\ No newline at end of file