Grover's search in qutip

grover_qutip.py

from random import randint
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from qutip import tensor, basis, hadamard_transform, identity

# initialize the |0> and |1> qubit states
zero = basis(2, 0)
one = basis(2, 1)

# how many times have we applied the oracle and diffuser?
rounds = 0


def measure(state):
    # get the probability amplitudes of each qubit in the series, and plot it
    probs = []
    for i in range(n):
        pa = one.trans()*state.ptrace(i)*one
        probs.append(abs(pa[0][0]))
    fig, ax = plt.subplots(1, 1)
    ax.plot([i for i in range(n)], probs)
    ax.set_ylim(0, 1)
    ax.xaxis.set_major_formatter(FormatStrFormatter('%d'))
    ax.set_title("Solution is %s, round %s" % (solution, rounds))
    ax.set_xlabel("Qubit Number")
    ax.set_ylabel("Measurement Probability (%)")
    fig.savefig("round%s.png" % rounds)

# the number of qubits
n = 10

# choose a solution at random from the number of qubits
solution = randint(0, n-1)

n_hadamard = hadamard_transform(n)
one_hadamard = hadamard_transform(1)
# create the oracle
solution_state = tensor(tensor(*[one if i == solution else zero for i in range(0, n)]), one)
oracle = tensor(*[identity(2) for i in range(0, n+1)]) - 2*solution_state*solution_state.trans()

# create the grover diffusion operator
n_zeros = tensor(*[zero for i in range(0, n)])
input_state = tensor(n_hadamard*n_zeros, one_hadamard*one)
grover_diffusion = n_hadamard*(2*n_zeros*n_zeros.trans() - tensor(*[identity(2) for i in range(0, n)]))*n_hadamard

measure(input_state)
# apply the oracle ~ n**0.5 times
for i in range(int(n**0.5)+1):
    # apply the oracle
    input_state = tensor(grover_diffusion, identity(2))*oracle*input_state
    rounds += 1
    measure(input_state)