From 30cc848ecd761568df279aa9d570f8533f2bc279 Mon Sep 17 00:00:00 2001
From: Kristofers Solo <dev@kristofers.xyz>
Date: Thu, 8 May 2025 18:49:39 +0300
Subject: [PATCH] feat(cli): add phase arg

---
 src/grovers_visualizer/args.py       | 14 ++++++++++++++
 src/grovers_visualizer/circuit.py    | 28 +++++++++++++++-------------
 src/grovers_visualizer/cli.py        |  2 +-
 src/grovers_visualizer/simulation.py | 12 +++++++++---
 4 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/src/grovers_visualizer/args.py b/src/grovers_visualizer/args.py
index f94287e..68f85fa 100644
--- a/src/grovers_visualizer/args.py
+++ b/src/grovers_visualizer/args.py
@@ -1,3 +1,4 @@
+import math
 from argparse import ArgumentParser
 from dataclasses import dataclass
 
@@ -11,6 +12,7 @@ class Args:
     iterations: int
     speed: float
     ui: bool
+    phase: float
 
 
 def parse_args() -> Args:
@@ -26,6 +28,7 @@ def parse_args() -> Args:
         iterations=ns.iterations,
         speed=ns.speed,
         ui=ns.ui,
+        phase=ns.phase,
     )
 
 
@@ -61,4 +64,15 @@ def parse_cli(base_parser: ArgumentParser) -> None:
         default=0.5,
         help="Pause duration (seconds) between steps (deafult: 0.5)",
     )
+    parser.add_argument(
+        "-p",
+        "--phase",
+        type=float,
+        default=math.pi,
+        help=(
+            "The phase φ (in radians) used for the oracle and diffusion steps. "
+            "Defaults to π, which implements the usual sign flip e^(iπ) = -1."
+        ),
+    )
+
     parser.add_argument("--ui", action="store_true", help="Run with DearPyGui UI")
diff --git a/src/grovers_visualizer/circuit.py b/src/grovers_visualizer/circuit.py
index e93d831..76313da 100644
--- a/src/grovers_visualizer/circuit.py
+++ b/src/grovers_visualizer/circuit.py
@@ -1,35 +1,38 @@
+import math
+
 from qiskit import QuantumCircuit
+from qiskit.circuit.library import PhaseGate
 
 from .state import QubitState
 
 
-def oracle(qc: QuantumCircuit, target_state: QubitState) -> None:
+def oracle(qc: QuantumCircuit, target_state: QubitState, /, *, phase: float = math.pi) -> None:
     """Oracle that flips the sign of the target state."""
     n = len(target_state)
     encode_target_state(qc, target_state)
-    apply_phase_inversion(qc, n)
+    apply_phase_inversion(qc, n, phase=phase)
     encode_target_state(qc, target_state)  # Undo
 
 
-def oracle_circuit(target: QubitState) -> QuantumCircuit:
+def oracle_circuit(target: QubitState, /, *, phase: float = math.pi) -> QuantumCircuit:
     n = len(target)
     qc = QuantumCircuit(n)
-    oracle(qc, target)
+    oracle(qc, target, phase=phase)
     return qc
 
 
-def diffusion(qc: QuantumCircuit, n: int) -> None:
+def diffusion(qc: QuantumCircuit, n: int, /, *, phase: float = math.pi) -> None:
     """Apply the Grovers diffusion operator."""
     qc.h(range(n))
     qc.x(range(n))
-    apply_phase_inversion(qc, n)
+    apply_phase_inversion(qc, n, phase=phase)
     qc.x(range(n))
     qc.h(range(n))
 
 
-def diffusion_circuit(n: int) -> QuantumCircuit:
+def diffusion_circuit(n: int, /, *, phase: float = math.pi) -> QuantumCircuit:
     qc = QuantumCircuit(n)
-    diffusion(qc, n)
+    diffusion(qc, n, phase=phase)
     return qc
 
 
@@ -40,11 +43,10 @@ def encode_target_state(qc: QuantumCircuit, target_state: QubitState) -> None:
             qc.x(i)
 
 
-def apply_phase_inversion(qc: QuantumCircuit, n: int) -> None:
+def apply_phase_inversion(qc: QuantumCircuit, n: int, /, *, phase: float = math.pi) -> None:
     """Apply a multi-controlled phase inversion (Z) to the marked state."""
     if n == 1:
-        qc.z(0)
+        qc.p(phase, 0)
         return
-    qc.h(n - 1)
-    qc.mcx(list(range(n - 1)), n - 1)
-    qc.h(n - 1)
+    mc_phase = PhaseGate(phase).control(n - 1)
+    qc.append(mc_phase, list(range(n)))
diff --git a/src/grovers_visualizer/cli.py b/src/grovers_visualizer/cli.py
index 8cc1b0e..531219f 100644
--- a/src/grovers_visualizer/cli.py
+++ b/src/grovers_visualizer/cli.py
@@ -7,7 +7,7 @@ from .args import Args
 def run_cli(args: Args) -> None:
     vis = GroverVisualizer(args.target, pause=args.speed)
 
-    for it, sv in grover_evolver(vis.target, args.iterations):
+    for it, sv in grover_evolver(vis.target, args.iterations, phase=args.phase):
         if not vis.is_running:
             break
         vis.update(it, sv)
diff --git a/src/grovers_visualizer/simulation.py b/src/grovers_visualizer/simulation.py
index f3ca0bc..5d86281 100644
--- a/src/grovers_visualizer/simulation.py
+++ b/src/grovers_visualizer/simulation.py
@@ -1,3 +1,4 @@
+import math
 from collections.abc import Iterator
 from itertools import count
 
@@ -9,7 +10,12 @@ from grovers_visualizer.circuit import diffusion_circuit, oracle_circuit
 from grovers_visualizer.state import QubitState
 
 
-def grover_evolver(target: QubitState, max_iterations: int = 0) -> Iterator[tuple[int, Statevector]]:
+def grover_evolver(
+    target: QubitState,
+    max_iterations: int = 0,
+    *,
+    phase: float = math.pi,
+) -> Iterator[tuple[int, Statevector]]:
     """Yields (iteration, statevector) pairs.
 
     - iteration=0 is the uniform-Hadamard initialization
@@ -24,8 +30,8 @@ def grover_evolver(target: QubitState, max_iterations: int = 0) -> Iterator[tupl
     sv = Statevector.from_instruction(qc)
     yield 0, sv
 
-    oracle_op = Operator(oracle_circuit(target))
-    diffusion_op = Operator(diffusion_circuit(n_qubits))
+    oracle_op = Operator(oracle_circuit(target, phase=phase))
+    diffusion_op = Operator(diffusion_circuit(n_qubits, phase=phase))
 
     iters = range(1, max_iterations + 1) if max_iterations > 0 else count(1)
     for i in iters: