|
| 1 | +import numpy as np |
| 2 | +import matplotlib.pyplot as plt |
| 3 | +from matplotlib.animation import FuncAnimation |
| 4 | +import os |
| 5 | + |
| 6 | +class QuantumFourierTransform: |
| 7 | + def __init__(self, n_qubits): |
| 8 | + self.n_qubits = n_qubits |
| 9 | + self.N = 2 ** n_qubits |
| 10 | + self.state = np.zeros(self.N, dtype=complex) |
| 11 | + |
| 12 | + def initialize_basis_state(self, index): |
| 13 | + """Set state to basis state |index⟩""" |
| 14 | + self.state = np.zeros(self.N, dtype=complex) |
| 15 | + self.state[index] = 1.0 |
| 16 | + |
| 17 | + def initialize_custom_state(self, state_vector): |
| 18 | + """Set state to custom normalized state vector""" |
| 19 | + assert len(state_vector) == self.N, "State vector length mismatch." |
| 20 | + norm = np.linalg.norm(state_vector) |
| 21 | + assert np.isclose(norm, 1.0), "State vector must be normalized." |
| 22 | + self.state = np.array(state_vector, dtype=complex) |
| 23 | + |
| 24 | + def initialize_superposition(self): |
| 25 | + """Create |+⟩^n superposition state""" |
| 26 | + self.state = np.ones(self.N, dtype=complex) / np.sqrt(self.N) |
| 27 | + def initialize_bell_pair(self): |
| 28 | + """2-qubit Bell state: (|00⟩ + |11⟩)/√2""" |
| 29 | + if self.n_qubits != 2: |
| 30 | + raise ValueError("Bell pair requires exactly 2 qubits.") |
| 31 | + self.state = np.zeros(self.N, dtype=complex) |
| 32 | + self.state[0] = 1 / np.sqrt(2) |
| 33 | + self.state[3] = 1 / np.sqrt(2) |
| 34 | + |
| 35 | + def initialize_ghz_state(self): |
| 36 | + """n-qubit GHZ state: (|00...0⟩ + |11...1⟩)/√2""" |
| 37 | + self.state = np.zeros(self.N, dtype=complex) |
| 38 | + self.state[0] = 1 / np.sqrt(2) |
| 39 | + self.state[-1] = 1 / np.sqrt(2) |
| 40 | + |
| 41 | + def hadamard(self): |
| 42 | + return np.array([[1, 1], [1, -1]]) / np.sqrt(2) |
| 43 | + |
| 44 | + def apply_single_qubit_gate(self, gate, target): |
| 45 | + I = np.eye(2) |
| 46 | + ops = [I] * self.n_qubits |
| 47 | + ops[target] = gate |
| 48 | + U = ops[0] |
| 49 | + for op in ops[1:]: |
| 50 | + U = np.kron(U, op) |
| 51 | + self.state = U @ self.state |
| 52 | + def apply_controlled_phase(self, control, target, theta): |
| 53 | + U = np.eye(self.N, dtype=complex) |
| 54 | + for i in range(self.N): |
| 55 | + b = format(i, f'0{self.n_qubits}b') |
| 56 | + if b[self.n_qubits - 1 - control] == '1' and b[self.n_qubits - 1 - target] == '1': |
| 57 | + U[i, i] *= np.exp(1j * theta) |
| 58 | + self.state = U @ self.state |
| 59 | + |
| 60 | + def swap_registers(self): |
| 61 | + perm = [int(format(i, f'0{self.n_qubits}b')[::-1], 2) for i in range(self.N)] |
| 62 | + self.state = self.state[perm] |
| 63 | + def apply_qft(self, inverse=False): |
| 64 | + for target in range(self.n_qubits): |
| 65 | + idx = self.n_qubits - 1 - target |
| 66 | + for offset in range(1, self.n_qubits - target): |
| 67 | + control = self.n_qubits - 1 - (target + offset) |
| 68 | + angle = np.pi / (2 ** offset) |
| 69 | + if inverse: |
| 70 | + angle *= -1 |
| 71 | + self.apply_controlled_phase(control, idx, angle) |
| 72 | + self.apply_single_qubit_gate(self.hadamard(), idx) |
| 73 | + self.swap_registers() |
| 74 | + def measure(self, shots=1024): |
| 75 | + probs = np.abs(self.state) ** 2 |
| 76 | + outcomes = np.random.choice(self.N, size=shots, p=probs) |
| 77 | + counts = {} |
| 78 | + for o in outcomes: |
| 79 | + bitstring = format(o, f'0{self.n_qubits}b') |
| 80 | + counts[bitstring] = counts.get(bitstring, 0) + 1 |
| 81 | + return counts |
| 82 | + |
| 83 | + def print_amplitudes(self, title="Quantum State"): |
| 84 | + print(f"\n{title}:") |
| 85 | + for i, amp in enumerate(self.state): |
| 86 | + print(f"|{i:0{self.n_qubits}b}>: {amp:.4f}") |
| 87 | + |
| 88 | + |
| 89 | +# Inside the class |
| 90 | + def reset_animation_log(self): |
| 91 | + self.animation_frames = [] |
| 92 | + |
| 93 | + def record_probability_frame(self): |
| 94 | + """Store current probabilities for animation.""" |
| 95 | + probs = np.abs(self.state) ** 2 |
| 96 | + self.animation_frames.append(probs.copy()) |
| 97 | + |
| 98 | + def apply_qft_with_recording(self, inverse=False): |
| 99 | + self.reset_animation_log() |
| 100 | + self.record_probability_frame() # Record initial state |
| 101 | + |
| 102 | + for target in range(self.n_qubits): |
| 103 | + idx = self.n_qubits - 1 - target |
| 104 | + for offset in range(1, self.n_qubits - target): |
| 105 | + control = self.n_qubits - 1 - (target + offset) |
| 106 | + angle = np.pi / (2 ** offset) |
| 107 | + if inverse: |
| 108 | + angle *= -1 |
| 109 | + self.apply_controlled_phase(control, idx, angle) |
| 110 | + self.record_probability_frame() |
| 111 | + self.apply_single_qubit_gate(self.hadamard(), idx) |
| 112 | + self.record_probability_frame() |
| 113 | + |
| 114 | + self.swap_registers() |
| 115 | + self.record_probability_frame() |
| 116 | + |
| 117 | + def animate_probability_evolution(self, save_path="qft_animation.gif", interval=600): |
| 118 | + labels = [format(i, f'0{self.n_qubits}b') for i in range(self.N)] |
| 119 | + fig, ax = plt.subplots() |
| 120 | + bar = ax.bar(labels, [0]*self.N) |
| 121 | + ax.set_ylim(0, 1) |
| 122 | + ax.set_ylabel("Probability") |
| 123 | + ax.set_title("QFT Probability Evolution") |
| 124 | + |
| 125 | + def update(frame_probs): |
| 126 | + for rect, prob in zip(bar, frame_probs): |
| 127 | + rect.set_height(prob) |
| 128 | + return bar |
| 129 | + |
| 130 | + ani = FuncAnimation(fig, update, frames=self.animation_frames, |
| 131 | + interval=interval, blit=False, repeat=False) |
| 132 | + |
| 133 | + if save_path.endswith(".gif"): |
| 134 | + ani.save(save_path, writer='pillow') |
| 135 | + elif save_path.endswith(".mp4"): |
| 136 | + ani.save(save_path, writer='ffmpeg') |
| 137 | + else: |
| 138 | + raise ValueError("Unsupported format. Use .gif or .mp4") |
| 139 | + |
| 140 | + print(f"Animation saved to {save_path}") |
| 141 | + |
| 142 | +# --------------------------- |
| 143 | +# Example usage |
| 144 | +# --------------------------- |
| 145 | +if __name__ == "__main__": |
| 146 | + qft = QuantumFourierTransform(n_qubits=3) |
| 147 | + qft.initialize_basis_state(5) |
| 148 | + |
| 149 | +# Apply QFT with intermediate state recording |
| 150 | + qft.apply_qft_with_recording() |
| 151 | + |
| 152 | +# Animate the probability distribution over steps |
| 153 | + qft.animate_probability_evolution(save_path="qft_probs.gif", interval=800) |
| 154 | + |
| 155 | + |
| 156 | + |
| 157 | +""" |
| 158 | +Quick Tips on Usage |
| 159 | +
|
| 160 | +
|
| 161 | +Superposition inputs (e.g. Hadamards to create |+\rangle^{\otimes n}) |
| 162 | +Custom state initialization (e.g. any normalized complex vector) |
| 163 | +Entangled state preparation (e.g. Bell states, GHZ states, etc.) |
| 164 | +
|
| 165 | +
|
| 166 | +
|
| 167 | +initialize_superposition() gives |+\rangle^{\otimes n} |
| 168 | +initialize_bell_pair() is for 2-qubit entangled states |
| 169 | +initialize_ghz_state() creates a GHZ state |
| 170 | +initialize_custom_state(vec) lets you pass any normalized state |
| 171 | +""" |
0 commit comments