Updating GHZ Demo. (#93)

* using more spec api in demo

* fill gate locations from middle index

Author: weinbe58

demo/log_depth_ghz.py

@@ -1,39 +1,42 @@
 import math
-from typing import Any, Literal, TypeVar
+from typing import Any, TypeVar
 
 from bloqade.geometry.dialects import grid
 from kirin.dialects import ilist
 
-from bloqade.shuttle import action, gate, init, measure, schedule, spec
+from bloqade.shuttle import action, gate, init, schedule, spec
 from bloqade.shuttle.prelude import move, tweezer
-from bloqade.shuttle.visualizer import MatplotlibRenderer, PathVisualizer
+from bloqade.shuttle.visualizer import PathVisualizer
 
 NMove = TypeVar("NMove")
 
 
 @tweezer
 def entangle_move(
-    mem_zone: grid.Grid[Any, Literal[1]],
-    gate_zone: grid.Grid[NMove, Literal[2]],
     ctrl_ids: ilist.IList[int, NMove],
     qarg_ids: ilist.IList[int, NMove],
     gate_ids: ilist.IList[int, NMove],
 ):
 
+    mem_zone = spec.get_static_trap(zone_id="mem")
+    gate_zone = spec.get_special_grid(grid_id="gate")
+
     mem_y = grid.get_ypos(mem_zone)[0]
-    ctrl_start = grid.get_xpos(grid.sub_grid(mem_zone, ctrl_ids, [0]))
-    qarg_start = grid.get_xpos(grid.sub_grid(mem_zone, qarg_ids, [0]))
+    ctrl_start = grid.get_xpos(mem_zone[ctrl_ids, 0])
+    qarg_start = grid.get_xpos(mem_zone[qarg_ids, 0])
 
-    pos_1 = grid.from_positions(ctrl_start, [mem_y, mem_y])
-    pos_2 = grid.shift(grid.from_positions(ctrl_start, [mem_y - 4.0, mem_y]), 2.0, 0.0)
-    pos_3 = grid.from_positions(qarg_start, [mem_y - 4.0, mem_y])
-    gate_pos = grid.sub_grid(gate_zone, gate_ids, [0, 1])
+    pos_1 = grid.from_positions(ctrl_start, [mem_y, mem_y + 4.0])
+    pos_2 = grid.shift(pos_1, 0.0, -4.0)
+    pos_3 = grid.from_positions(qarg_start, grid.get_ypos(pos_2))
+    pos_4 = grid.shift(pos_3, 0.0, -4.0)
+    gate_pos = gate_zone[gate_ids, :]
 
     action.set_loc(pos_1)
     action.turn_on(action.ALL, [0])
     action.move(pos_2)
     action.move(pos_3)
     action.turn_on([], [1])
+    action.move(pos_4)
     action.move(gate_pos)
 
 
@@ -45,33 +48,28 @@ def apply_h(zone: grid.Grid[Any, Any]):
     gate.global_r(0.0, -math.pi / 2.0)
 
 
-@move
-def apply_cx(gate_zone: grid.Grid[Any, Any]):
-    # TODO: This is wrong and needs to be fixed.
-    gate.global_r(0.0, math.pi)
-    gate.top_hat_cz(gate_zone, upper_buffer=3.0, lower_buffer=3.0)
-    gate.global_r(0.0, -math.pi)
-
-
 @move
 def run_entangle_move(
-    mem_zone: grid.Grid[Any, Literal[1]],
-    gate_zone: grid.Grid[Any, Literal[2]],
     ctrl_ids: ilist.IList[int, NMove],
     qarg_ids: ilist.IList[int, NMove],
     gate_ids: ilist.IList[int, NMove],
 ):
 
+    gate_zone = spec.get_special_grid(grid_id="gate")
+    mem_zone = spec.get_static_trap(zone_id="mem")
+
     num = len(ctrl_ids)
     xtones = ilist.range(num)
     ytones = [0, 1]
 
     dtask = schedule.device_fn(entangle_move, xtones, ytones)
     rev_dtask = schedule.reverse(dtask)
 
-    dtask(mem_zone, gate_zone, ctrl_ids, qarg_ids, gate_ids)
-    apply_cx(gate_zone)
-    rev_dtask(mem_zone, gate_zone, ctrl_ids, qarg_ids, gate_ids)
+    gate.local_r(0.0, math.pi, mem_zone[qarg_ids, 0])
+    dtask(ctrl_ids, qarg_ids, gate_ids)
+    gate.top_hat_cz(gate_zone)
+    rev_dtask(ctrl_ids, qarg_ids, gate_ids)
+    gate.local_r(0.0, -math.pi, mem_zone[qarg_ids, 0])
 
 
 N = TypeVar("N")
@@ -106,7 +104,7 @@ def ghz_prep_steps(
 ):
     jobs = []
     layers = get_layers(qubit_ids)
-
+    mid = gate_width // 2
     for i in range(len(layers)):
         ctrl_ids = layers[i][0]
         qarg_ids = layers[i][1]
@@ -119,8 +117,12 @@ def ghz_prep_steps(
             if end > n_gates:
                 end = n_gates
 
+            layer_size = end - start
+            gate_start = mid - (layer_size // 2)
+            gate_end = gate_start + layer_size
+
             jobs = jobs + [
-                (ctrl_ids[start:end], qarg_ids[start:end], range(end - start))
+                (ctrl_ids[start:end], qarg_ids[start:end], range(gate_start, gate_end))
             ]
 
     return jobs
@@ -129,20 +131,20 @@ def ghz_prep_steps(
 @move
 def run_prep_steps(
     qubit_ids: ilist.IList[int, N],
-    gate: grid.Grid[Any, Literal[2]],
-    mem: grid.Grid[Any, Literal[1]],
 ):
 
-    gate_shape = grid.shape(gate)
+    gate_zone = spec.get_special_grid(grid_id="gate")
+    mem_zone = spec.get_static_trap(zone_id="mem")
+    gate_shape = grid.shape(gate_zone)
     gate_width = gate_shape[0]
 
     jobs = ghz_prep_steps(qubit_ids, gate_width)
-    apply_h(grid.sub_grid(mem, [qubit_ids[0]], [0]))
+    apply_h(mem_zone[qubit_ids[0], 0])
     for i in range(len(jobs)):
         ctrl_ids = jobs[i][0]
         qarg_ids = jobs[i][1]
         gate_ids = jobs[i][2]
-        run_entangle_move(mem, gate, ctrl_ids, qarg_ids, gate_ids)
+        run_entangle_move(ctrl_ids, qarg_ids, gate_ids)
 
 
 def run_ghz():
@@ -167,7 +169,6 @@ def run_ghz():
     gate_shift = mid_mem - mid_gate
 
     gate = gate.shift(gate_shift, 0.0)
-
     spec_value = spec.ArchSpec(
         layout=spec.Layout(
             static_traps={
@@ -185,17 +186,16 @@ def run_ghz():
     @move
     def main():
         init.fill([spec.get_static_trap(zone_id="mem")])
-        run_prep_steps([1, 2, 4, 5, 7, 9, 14, 28, 29, 32], gate, mem)  # type: ignore
-        return measure.measure((mem,))
+        run_prep_steps([1, 2, 4, 5, 7, 9, 14, 28, 29, 32])  # type: ignore
+        # return measure.measure((mem,))
 
     return main, spec_value
 
 
 def run_plotter():
     main, spec_value = run_ghz()
 
-    renderer = MatplotlibRenderer()
-    PathVisualizer(main.dialects, renderer=renderer, arch_spec=spec_value).run(main, ())
+    PathVisualizer(main.dialects, arch_spec=spec_value).run(main, args=())
 
 
 if __name__ == "__main__":