From fc92ce374735cc4a54a6de339806548b41e3b10b Mon Sep 17 00:00:00 2001
From: Pete Shadbolt <hello@peteshadbolt.co.uk>
Date: Sat, 19 Nov 2016 18:41:23 -0800
Subject: [PATCH] Add shorthand function to run multiple CZs.

---
 abp/graphstate.py                             |  5 +++
 examples/mix_graph_and_networkx.py            |  5 +--
 examples/visualization/auto_layout.py         | 13 ++++----
 .../issues/unpositioned_nodes.py              | 32 ++++++++++---------
 tests/test_graphstate.py                      |  7 ++++
 5 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/abp/graphstate.py b/abp/graphstate.py
index b3255e2..25d1b09 100755
--- a/abp/graphstate.py
+++ b/abp/graphstate.py
@@ -112,6 +112,11 @@ class GraphState(object):
             else:
                 self.act_local_rotation(node, operation)
 
+    def act_czs(self, *pairs):
+        """ Shorthand to act many CZs """
+        for a, b in pairs:
+            self.act_cz(a, b)
+
     def _add_edge(self, v1, v2, data={}):
         """ Add an edge between two vertices """
         self.adj[v1][v2] = data
diff --git a/examples/mix_graph_and_networkx.py b/examples/mix_graph_and_networkx.py
index c89af56..8d92b8d 100644
--- a/examples/mix_graph_and_networkx.py
+++ b/examples/mix_graph_and_networkx.py
@@ -1,9 +1,10 @@
-from abp.fancy import GraphState
+from abp import NXGraphState
 from abp.util import xyz
 import networkx as nx
 
 n = 10
-g = GraphState(range(n))
+g = NXGraphState(range(n))
 nx.set_node_attributes(g, "color", "red")
 g.add_edges_from([i, i+1] for i in range(n-1))
 print g.node[0]["color"]
+
diff --git a/examples/visualization/auto_layout.py b/examples/visualization/auto_layout.py
index 8367e51..f6aed5b 100644
--- a/examples/visualization/auto_layout.py
+++ b/examples/visualization/auto_layout.py
@@ -1,4 +1,4 @@
-from abp.fancy import GraphState
+from abp import GraphState, VizClient
 from abp.util import xyz
 import numpy as np
 import time
@@ -43,11 +43,10 @@ def lattice(unit_cell, size):
 
 nodes, edges = lattice(threedee_unit_cell, (3, 3, 3))
 
-psi = GraphState()
-for node in nodes:
-    psi.add_node(str(node))
-    psi.act_hadamard(str(node))
+psi = GraphState(nodes)
 
-for edge in edges:
-    psi.act_cz(str(edge[0]), str(edge[1]))
+for a, b in edges:
+    psi.act_cz(a, b)
 
+v = VizClient()
+v.update(psi)
diff --git a/examples/visualization/issues/unpositioned_nodes.py b/examples/visualization/issues/unpositioned_nodes.py
index 476fe7a..0309b56 100644
--- a/examples/visualization/issues/unpositioned_nodes.py
+++ b/examples/visualization/issues/unpositioned_nodes.py
@@ -1,18 +1,20 @@
-from abp.fancy import GraphState
-import networkx as nx
+from abp import GraphState, VizClient
+from abp.util import xyz
 
-edges = [(0,1),(1,2),(2,3),(3,4)]
-nodes = [(i, {'x': i, 'y': 0, 'z':0}) for i in range(5)]
-gs = GraphState()
+# Prepare to visualize
+v = VizClient()
 
-for node, position in nodes:
-    gs.add_qubit(node, position=position)
-    gs.act_hadamard(node)
+# Make a graph state with position attributes
+g = GraphState()
+for i in range(5):
+    g.add_qubit(i, position=xyz(i, 0, 0), vop="identity")
+g.act_czs((0,1),(1,2),(2,3),(3,4))
 
-for edge in edges:
-    gs.act_cz(*edge)
-gs.update(3)
-# a single line of qubits are created along the x axis
-gs.add_qubit('start')
-gs.update(0)
-# a curved 5-qubit cluster and single qubit is depicted
+# Show it
+v.update(g, 3)
+
+# Add a qubit with no position
+g.add_qubit('start')
+
+# Show it
+v.update(g)
diff --git a/tests/test_graphstate.py b/tests/test_graphstate.py
index 5670507..9689321 100644
--- a/tests/test_graphstate.py
+++ b/tests/test_graphstate.py
@@ -83,6 +83,13 @@ def test_cz():
     assert g.has_edge(0, 1)
 
 
+def test_czs():
+    """ Test multiple CZ shorthand """
+    g = GraphState([0, 1, 2])
+    g.act_czs((0, 1), (1, 2))
+    assert len(g.edgelist()) == 2
+
+
 def test_local_complementation():
     """ Test that local complementation works okay """
     pairs = (0, 1), (0, 3), (1, 3), (1, 2),