feat(circle): add triangulation circle

src/grovers_visualizer/main.py

@@ -8,20 +8,24 @@ using matplotlib.
 
 from collections.abc import Iterator
 from itertools import product
-from math import floor, pi, sqrt
-from typing import Callable
+from math import asin, cos, floor, pi, sin, sqrt
+from typing import TYPE_CHECKING, Callable
 
 import matplotlib.pyplot as plt
 import numpy as np
 import numpy.typing as npt
 from matplotlib.axes import Axes
 from matplotlib.container import BarContainer
+from matplotlib.patches import Circle
 from qiskit import QuantumCircuit
 from qiskit.quantum_info import Statevector
 
 from grovers_visualizer.gates import apply_phase_inversion, encode_target_state
 from grovers_visualizer.state import QubitState
 
+if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+
 
 def oracle(qc: QuantumCircuit, target_state: QubitState) -> None:
     """Oracle that flips the sign of the target state."""
@@ -84,11 +88,15 @@ def optimal_grover_iterations(n_qubits: int) -> int:
     return floor(pi / 4 * sqrt(2**n_qubits))
 
 
+def is_optimal_iteration(iteration: int, optimal_iteration: int) -> bool:
+    return iteration % optimal_iteration == 0 and iteration != 0
+
+
 def get_bar_color(state: str, target_state: QubitState | None, iteration: int, optimal_iteration: int | None) -> str:
     """Return the color for a bar based on state and iteration."""
     if state != target_state:
         return "skyblue"
-    if optimal_iteration and iteration % optimal_iteration == 0 and iteration != 0:
+    if optimal_iteration and is_optimal_iteration(iteration, optimal_iteration):
         return "green"
     return "orange"
 
@@ -98,7 +106,7 @@ def plot_amplitudes_live(
     bars: BarContainer,
     statevector: Statevector,
     basis_states: list[str],
-    step_label: str,
+    iteration_label: str,
     iteration: int,
     target_state: QubitState | None = None,
     optimal_iteration: int | None = None,
@@ -110,7 +118,7 @@ def plot_amplitudes_live(
         bar.set_height(amp)
         bar.set_color(get_bar_color(state, target_state, iteration, optimal_iteration))
 
-    ax.set_title(f"Iteration {iteration}: {step_label}")
+    ax.set_title(f"Iteration {iteration}: {iteration_label}")
     ax.set_ylim(-1, 1)
 
     for l in ax.lines:  # Remove previous mean line(s)
@@ -120,7 +128,50 @@ def plot_amplitudes_live(
 
     if not ax.get_legend():
         ax.legend(loc="upper right")
-    plt.pause(1)
+
+
+def draw_grover_circle(
+    ax: Axes,
+    iteration: int,
+    optimal_iterations: int,
+    theta: float,
+    state_angle: float,
+) -> None:
+    ax.clear()
+    ax.set_aspect("equal")
+    ax.set_xlim(-1.1, 1.1)
+    ax.set_ylim(-1.1, 1.1)
+    ax.set_xlabel("Unmarked amplitude")
+    ax.set_ylabel("Target amplitude")
+    ax.set_title("Grover State Vector Rotation")
+
+    # Draw unit circle
+    circle = Circle((0, 0), 1, color="gray", fill=False)
+    ax.add_artist(circle)
+
+    # Draw axes
+    ax.axhline(0, color="black", linewidth=0.5)
+    ax.axvline(0, color="black", linewidth=0.5)
+
+    # Draw labels
+    ax.text(1.05, 0, "", va="center", ha="left", fontsize=10)
+    ax.text(0, 1.05, "1", va="bottom", ha="center", fontsize=10)
+    ax.text(-1.05, 0, "", va="center", ha="right", fontsize=10)
+    ax.text(0, -1.05, "-1", va="top", ha="center", fontsize=10)
+
+    angle = state_angle + iteration * theta
+    x, y = cos(angle), sin(angle)
+    is_optimal = optimal_iterations and is_optimal_iteration(iteration, optimal_iterations)
+
+    # Arrow color: green at optimal, blue otherwise
+    color = "green" if is_optimal else "blue"
+    ax.arrow(0, 0, x, y, head_width=0.07, head_length=0.1, fc=color, ec=color, length_includes_head=True)
+
+    # Probability of target state is y^2
+    prob = y**2
+    ax.set_title(
+        f"Grover State Vector Rotation\nIteration {iteration} | Probability of target: {prob:.2f}{' (optimal)' if is_optimal else ''}"
+    )
 
 
 def main() -> None:
@@ -128,16 +179,19 @@ def main() -> None:
     n_qubits = len(target_state)
     basis_states = [str(bit) for bit in all_states(n_qubits)]
     optimal_iterations = optimal_grover_iterations(n_qubits)
+    theta = 2 * asin(1 / sqrt(2**n_qubits))
+    state_angle = 0.5 * theta
 
     plt.ion()
-    fig, ax = plt.subplots(figsize=(8, 3))
-    bars = ax.bar(basis_states, [0] * len(basis_states), color="skyblue")
-    ax.set_xlabel("Basis State")
-    ax.set_ylabel("Real Amplitude")
-    ax.set_ylim(-1, 1)
-    ax.set_title("Grover Amplitudes")
+    fig: Figure
+    ax_bar: Axes
+    ax_circle: Axes
+    fig, (ax_bar, ax_circle) = plt.subplots(1, 2, figsize=(12, 4))
+    bars = ax_bar.bar(basis_states, [0] * len(basis_states), color="skyblue")
+    ax_bar.set_ylim(-1, 1)
+    ax_bar.set_title("Amplitudes (example)")
 
-    def step_and_plot(
+    def iterate_and_plot(
         operation: Callable[[QuantumCircuit], None] | None,
         step_label: str,
         iteration: int,
@@ -145,17 +199,21 @@ def main() -> None:
         if operation is not None:
             operation(qc)
         sv = Statevector.from_instruction(qc)
-        plot_amplitudes_live(ax, bars, sv, basis_states, step_label, iteration, target_state, optimal_iterations)
+        plot_amplitudes_live(ax_bar, bars, sv, basis_states, step_label, iteration, target_state, optimal_iterations)
+        draw_grover_circle(ax_circle, iteration, optimal_iterations, theta, state_angle)
+
+        plt.pause(1)
 
     # Start with Hadamard
     qc = QuantumCircuit(n_qubits)
     qc.h(range(n_qubits))
-    step_and_plot(None, "Hadamard (Initialization)", 0)
+    iterate_and_plot(None, "Hadamard (Initialization)", 0)
 
     iteration = 1
     while plt.fignum_exists(fig.number):
-        step_and_plot(lambda qc: oracle(qc, target_state), "Oracle (Query Phase)", iteration)
-        step_and_plot(lambda qc: diffusion(qc, n_qubits), "Diffusion (Inversion Phase)", iteration)
+        iterate_and_plot(lambda qc: oracle(qc, target_state), "Oracle (Query Phase)", iteration)
+        iterate_and_plot(lambda qc: diffusion(qc, n_qubits), "Diffusion (Inversion Phase)", iteration)
+
         iteration += 1
 
     plt.ioff()