Updating arrow size and fixing issue with color of arrows during path (#75)

* updating arrow width and fixing issue with color of arrows during path

* put continue while rendering path step by step

* simplfying interface

* make implementation a bit more clear and robust

Author: weinbe58

src/bloqade/shuttle/visualizer/impl/path.py

@@ -24,7 +24,6 @@ class PathVisualizerMethods(interp.MethodTable):
 
     @interp.impl(Play)
     def play(self, _interp: PathVisualizer[Renderer], frame: interp.Frame, stmt: Play):
-        _interp.renderer.clear_paths()
 
         path = frame.get(stmt.path)
         # path is generated by "main" interpreter so
@@ -40,7 +39,6 @@ def play(self, _interp: PathVisualizer[Renderer], frame: interp.Frame, stmt: Pla
                 f"Expected a Path or tuple of Paths, got {path}"
             )
 
-        _interp.renderer.show()
         return ()
 
     @interp.impl(Parallel)

src/bloqade/shuttle/visualizer/renderers/matplotlib.py

@@ -37,6 +37,7 @@ class MatplotlibRenderer(RendererInterface):
     ax: Axes = field(default_factory=default_ax, repr=False)
 
     gate_display_options: GateDisplayOptions = field(default_factory=GateDisplayOptions)
+    arrow_rescale: float = field(default=1.0, kw_only=True)
 
     active_x_tones: set[int] = field(default_factory=set, repr=False, init=False)
     active_y_tones: set[int] = field(default_factory=set, repr=False, init=False)
@@ -59,6 +60,14 @@ def __post_init__(self) -> None:
         self.exit_button = Button(exit_ax, "Exit")
         self.exit_button.on_clicked(lambda event: exit())
 
+    @property
+    def fov_size(self) -> float:
+        return np.sqrt((self.xmax - self.xmin) ** 2 + (self.ymax - self.ymin) ** 2)
+
+    @property
+    def arrow_scale(self) -> float:
+        return self.fov_size / 400.0 * self.arrow_rescale
+
     def update_x_bounds(self, y: float) -> None:
         xmin = min(curr_xmin := getattr(self, "xmin", float("inf")), y - 3)
         xmax = max(curr_xmax := getattr(self, "xmax", float("-inf")), y + 3)
@@ -185,29 +194,34 @@ def render_path(self, pth: path.Path) -> None:
             for way_point in path_action.way_points
         ]
 
-        if len(all_waypoints) == 0:
+        num_unique_waypoints = len(set(all_waypoints))
+        if num_unique_waypoints < 2:
             return
 
+        num_arrows = num_unique_waypoints - 1
         first_waypoint = all_waypoints[0]
         curr_x = first_waypoint.x_positions
         curr_y = first_waypoint.y_positions
 
         color_map = plt.get_cmap("viridis")
 
-        num_steps = len(all_waypoints) - 1
         step = 0
 
         x_tones = np.array(pth.x_tones)
         y_tones = np.array(pth.y_tones)
 
         x = all_waypoints[0].x_positions
         y = all_waypoints[0].y_positions
+        self.clear_paths()
+        self.show()
 
         for action in pth.path:
             if isinstance(action, taskgen.WayPointsAction):
-                for way_point in action.way_points:
-                    x = way_point.x_positions
-                    y = way_point.y_positions
+                for start, end in zip(action.way_points[:-1], action.way_points[1:]):
+                    x = end.x_positions
+                    y = end.y_positions
+                    curr_x = start.x_positions
+                    curr_y = start.y_positions
 
                     for (x_tone, x_start, x_end), (y_tone, y_start, y_end) in product(
                         zip(pth.x_tones, curr_x, x), zip(pth.y_tones, curr_y, y)
@@ -226,14 +240,14 @@ def render_path(self, pth: path.Path) -> None:
                             x_tone in self.active_x_tones
                             and y_tone in self.active_y_tones
                         )
-
+                        p = step / (num_arrows - 1) if num_arrows > 1 else 0.0
                         line = self.ax.arrow(
                             x_start,
                             y_start,
                             dx,
                             dy,
-                            width=0.1,
-                            color=color_map(step / num_steps),
+                            width=self.arrow_scale,
+                            color=color_map(p),
                             length_includes_head=True,
                             linestyle="-" if is_on else (0, (5, 10)),
                             alpha=1.0 if is_on else 0.5,
@@ -243,9 +257,9 @@ def render_path(self, pth: path.Path) -> None:
                         line.set_edgecolor(line.get_facecolor())
                         self.curr_path_lines.append(line)
 
-                    curr_x = x
-                    curr_y = y
-                    step += 1
+                    if curr_x != x or curr_y != y:
+                        step += 1
+                        self.show()
 
             elif isinstance(action, taskgen.TurnOnAction):
                 self.active_x_tones.update(x_tones[action.x_tone_indices])