Demo the discrete quantum random walk with a Hadamard coin.

quantum_random_walk.py

import numpy as np
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.animation as animation


def hadamard_walk(
    state: dict[tuple[int, int], np.complex128]
) -> dict[tuple[int, int], np.complex128]:
    new_state = defaultdict(lambda: np.complex128(0))
    for (spin, position), coefficient in state.items():
        new_state[(1, position + 1)] += coefficient / np.sqrt(2)
        new_state[(-1, position - 1)] += spin * coefficient / np.sqrt(2)
    return new_state


def plot_probability(state: dict[tuple[int, int], np.complex128]):
    # Compute the probability distribution by summing over the coin states (spin)
    probabilities = defaultdict(float)
    for (_, position), coefficient in state.items():
        probabilities[position] += np.abs(coefficient)**2
    xs = sorted(probabilities.keys())
    ys = [probabilities[x] for x in xs]

    plt.figure(figsize=(10, 5))
    plt.bar(xs, ys)
    plt.xlabel("Position")
    plt.ylabel("Probability")
    plt.title("Probability Distribution after Quantum Walk")
    plt.show()
    plt.close()


def animate_quantum_walk(initial_state: dict[tuple[int, int], np.complex128],
                         steps: int):
    # Pre-calculate all states from step 0 to the desired number of steps.
    states = []
    state = initial_state
    states.append(state)
    for _ in range(steps):
        state = hadamard_walk(state)
        states.append(state)

    fig, ax = plt.subplots(figsize=(10, 5))

    def update(frame: int):
        ax.clear()
        # Compute probabilities for the current frame
        probabilities = defaultdict(float)
        for (_, position), coefficient in states[frame].items():
            probabilities[position] += np.abs(coefficient)**2
        xs = probabilities.keys()
        ys = probabilities.values()
        ax.bar(xs, ys)
        ax.set_xlabel("Position")
        ax.set_ylabel("Probability")
        ax.set_title(f"Quantum Walk - Step {frame}")
        ax.set_xlim(-steps, steps)
        ax.set_ylim(0, 1)

    ani = animation.FuncAnimation(
        fig, update, frames=len(states), interval=40, repeat=False)
    plt.show()
    return ani


# Initialize the state with np.complex128 coefficients.
initial_state = defaultdict(lambda: np.complex128(0))
initial_state[(1, 0)] = np.complex128(np.cos(np.pi / 4))
initial_state[(-1, 0)] = np.complex128(0, np.sin(np.pi / 4))

step = 100

anim = animate_quantum_walk(initial_state, step)

state = initial_state
for _ in range(step):
    state = hadamard_walk(state)
plot_probability(state)