Quantum programming baby steps

from cirq import NamedQubit, TOFFOLI, X, CNOT, H
import math

nqubits = 4
state_space_size = 2 ** nqubits
num_grover_iterations = round(math.pi / (4. * math.asin(1. / math.sqrt(state_space_size))) - 0.5)

print(f'State space size=#{state_space_size}, num iters=#{num_grover_iterations}')


x1 = NamedQubit('x1')
x2 = NamedQubit('x2')
x3 = NamedQubit('x3')
x4 = NamedQubit('x4')

w1 = NamedQubit('w1')
w2 = NamedQubit('w2')
w3 = NamedQubit('w3')
w4 = NamedQubit('w4')
w5 = NamedQubit('w5')

# I am not sure if reusing ancilla across iterations is ok, so creating
# one per iteration.
ancilla = [cirq.NamedQubit(f'a#{i}') for i in range(num_grover_iterations)]

def make_ancilla():
  """Returns a sequence of operations that prepare each ancilla bit in the
  |-⟩ = H|1⟩ state."""
  prepare_ancilla = []
  for q in ancilla:
    prepare_ancilla.append(X(q))
    prepare_ancilla.append(H(q))
  return prepare_ancilla

def make_equal_superposition_state():
  return [H(x1), H(x2), H(x3), H(x4)]

def make_oracle(solution, z):
  """Returns a sequence of gates that takes |x>|z> to -|x>|z> if
  |x> == |solution>, else leaves the state alone. Note that the solution
  is known here, but in general this oracle will check if there is a solution
  in |x>."""
  assert len(solution) == 4, "I assumed a 4 qubit solution"
  gates = []
  uncompute_gates = []

  # Convenience fn to add uncomputation operations as we add the main
  # operations.
  def _add(gate, *operands):
    gates.append(gate(*operands))
    uncompute_gates.append(cirq.inverse(gate(*operands)))

  input_qubits = [x1, x2, x3, x4]
  for i, x in enumerate(solution):
    assert x in (1, 0), "Expected a binary solution vector"
    if x == 0:
      _add(X, input_qubits[i])
  
  _add(TOFFOLI, x1, x2, w1)
  _add(TOFFOLI, w1, x3, w2)
  _add(TOFFOLI, w2, x4, w3)

  # Copy the output to the ancilla qubit since uncomputation will reset w3
  # to its input state.
  gates.append(CNOT(w3, z))

  uncompute_gates.reverse()
  gates.extend(uncompute_gates)
  return gates
  

def make_search_step(z):
  """Returns a sequence of ops that runs the diffusion step of Grover search
  by reflecting |x> about |E> where |E> is the equal superposition state.
  z is an ancilla qubit. |x>|z> should be first passed through an oracle step.
  """
  return [   
    H(x1), H(x2), H(x3), H(x4),
    X(x1), X(x2), X(x3), X(x4),
    TOFFOLI(x1, x2, w1),
    TOFFOLI(w1, x3, w2),
    TOFFOLI(w2, x4, w3),

    # Leave the result in the ancilla.
    CNOT(w3, z),

    # uncompute
    TOFFOLI(w2, x4, w3),
    TOFFOLI(w1, x3, w2),
    TOFFOLI(x1, x2, w1),
    X(x1), X(x2), X(x3), X(x4),
    H(x1), H(x2), H(x3), H(x4),
  ]

def make_ckt(solution):
  """Creates a circuit to search for the given solution, which should be a
  4-element binary (0/1) array."""
  circuit = cirq.Circuit()
  circuit.append(make_equal_superposition_state())
  circuit.append(make_ancilla())
  for iter in range(num_grover_iterations):
    circuit.append(make_oracle(solution, ancilla[iter]))
    circuit.append(make_search_step(ancilla[iter]))
  circuit.append([cirq.measure(q) for q in (x1, x2, x3, x4)])
  return circuit

circuit = make_ckt([1, 0, 0, 1])
print(circuit)

sim = cirq.Simulator()
result = sim.simulate(circuit, initial_state=0)
print(result)