Add layout

abp/graphstate.py

@@ -6,6 +6,10 @@ from collections import defaultdict
 import itertools as it
 import clifford
 import json
+try:
+    import networkx as nx
+except ImportError:
+    print "Could not import networkx: layout will not work"
 
 
 class GraphState(object):
@@ -13,6 +17,7 @@ class GraphState(object):
     def __init__(self):
         self.ngbh = defaultdict(set)
         self.vops = defaultdict(int)
+        self.meta = defaultdict(dict)
 
     def add_vertex(self, v):
         """ Add a vertex if it doesn't already exist """
@@ -84,7 +89,6 @@ class GraphState(object):
         if self.ngbh[a] - {b}:
             self.remove_vop(a, b)
         edge = self.has_edge(a, b)
-        # TODO put this in a new function for diff hook
         new_edge, self.vops[a], self.vops[b] = clifford.cz_table[edge, self.vops[a], self.vops[b]]
         if new_edge != edge:
             self.toggle_edge(a, b)
@@ -97,22 +101,24 @@ class GraphState(object):
     def to_json(self):
         """ Convert the graph to JSON form """
         ngbh = {key: tuple(value) for key, value in self.ngbh.items()}
-        return json.dumps({"vops": self.vops, "ngbh": ngbh})
-
-
-class DiffedGraphState(GraphState):
-    """ Just like a graph state, but tracks changes for rendering purposes """
-
-    def __init__(self):
-        GraphState.__init__(self)
-        self.diff = []
-
-    def add_vertex(self, v):
-        GraphState.add_vertex(self, v)
-        self.diff.append("add_vertex {}".format(v))
-
-    def add_edge(self, v1, v2):
-        GraphState.add_edge(self, v1, v2)
-        self.diff.append("add_edge {} {}".format(v1, v2))
+        meta = {key: value for key, value in self.meta.items()}
+        return json.dumps({"vops": self.vops, "ngbh": ngbh, "meta": meta})
+
+    def to_networkx(self):
+        """ Convert the graph to a networkx graph """
+        g = nx.Graph()
+        g.edge = {node: {neighbour: {} for neighbour in neighbours} 
+                for node, neighbours in self.ngbh.items()}
+        g.node = {node: {"vop": vop} for node, vop in self.vops.items()}
+        for node, metadata in self.meta.items():
+            g.node[node].update(metadata)
+        return g
+
+    def layout(self):
+        """ Automatically lay out the graph """
+        g = self.to_networkx()
+        pos = nx.spring_layout(g, dim=3, scale=10)
+        for key, (x, y, z) in pos.items():
+            self.meta[key]["pos"] = {"x": round(x, 0), "y": round(y, 0), "z": round(z, 0)}
 
         

tests/test_graph.py

@@ -1,4 +1,4 @@
-from abp.graphstate import GraphState, DiffedGraphState
+from abp.graphstate import GraphState
 from abp import clifford
 import time
 
@@ -81,13 +81,3 @@ def test_cz():
     assert g.has_edge(0, 1)
 
 
-def test_diff():
-    """ Test diffing """
-    g = DiffedGraphState()
-    g.add_vertex(0)
-    g.add_vertex(1)
-    g.act_local_rotation(0, clifford.by_name["hadamard"])
-    g.act_local_rotation(1, clifford.by_name["hadamard"])
-    g.act_local_rotation(1, clifford.by_name["py"])
-    g.act_cz(0, 1)
-    assert len(g.diff) == 3

tests/test_json.py

@@ -1,4 +1,4 @@
-from abp.graphstate import GraphState, DiffedGraphState
+from abp.graphstate import GraphState
 from abp import clifford
 import time
 import json

tests/test_layout.py

@@ -0,0 +1,26 @@
+from abp.graphstate import GraphState
+
+def demograph():
+    """ A graph for testing with """
+    g = GraphState()
+    g.add_edge(0, 1)
+    g.add_edge(1, 2)
+    g.add_edge(2, 0)
+    g.add_edge(0, 3)
+    g.add_edge(100, 200)
+    return g
+
+
+def test_nx_convert():
+    g = demograph()
+    n = g.to_networkx()
+    assert len(g.ngbh) == len(n.edge)
+    assert len(g.vops) == len(n.node)
+
+def test_layout():
+    g = demograph()
+    g.layout()
+    assert len(g.meta) == len(g.vops)
+    assert "pos" in g.meta[0]
+    assert "x" in g.meta[0]["pos"]
+