Grover Search and Quantum Inference

using Yao
using Yao.Blocks
using Compat
using Compat.Test
using StatsBase

"""
A way to construct oracle, e.g. inference_oracle([1,2,-3,5]) will
invert the sign when a qubit configuration matches: 1=>1, 2=>1, 3=>0, 5=>1.
"""
function inference_oracle(locs::Vector{Int})
    control(locs[1:end-1], abs(locs[end]) => (locs[end]>0 ? Z : chain(phase(π), Z)))
end

function reflectblock(A::MatrixBlock{N}) where N
    chain(N, A |> adjoint, inference_oracle(-collect(1:N)), A)
end

nbit = 12
A = repeat(nbit, H)
ref = reflectblock(A)

@testset "test reflect" begin
    reg = rand_state(nbit)
    ref_vec = apply!(zero_state(nbit), A) |> statevec
    v0 = reg |> statevec
    @test -2*(ref_vec'*v0)*ref_vec + v0 ≈ apply!(copy(reg), ref) |> statevec
end

Test Summary: | Pass  Total
test reflect  |    1      1
Base.Test.DefaultTestSet("test reflect", Any[], 1, false)

# first, construct the oracle with desired state in the range 100-105.
oracle!(reg::DefaultRegister) = (reg.state[100:105,:]*=-1; reg)

# transform it into a function block, so it can be put inside a `Sequential`.
fb_oracle = FunctionBlock{:Oracle}(reg->oracle!(reg))

"""
ratio of components in a wavefunction that flip sign under oracle.
"""
function prob_match_oracle(psi::DefaultRegister, oracle)
    fliped_reg = apply!(register(ones(Complex128, 1<<nqubits(psi))), oracle)
    match_mask = fliped_reg |> statevec |> real .< 0
    norm(statevec(psi)[match_mask])^2
end

# uniform state as initial state
psi0 = apply!(zero_state(nbit), A)

# the number of grover steps that can make it reach first maximum overlap.
num_grover_step(prob::Real) = Int(round(pi/4/sqrt(prob)))-1
niter = num_grover_step(prob_match_oracle(psi0, fb_oracle))

# construct the whole circuit
gb = sequence(sequence(fb_oracle, ref) for i = 1:niter);

for (i, blk) in enumerate(gb)
    apply!(psi0, blk)
    overlap = prob_match_oracle(psi0, fb_oracle)
    println("step $i, overlap = $overlap")
end

step 1, overlap = 0.01313214562833301
step 2, overlap = 0.03619369756224663
step 3, overlap = 0.07010978618384153
step 4, overlap = 0.11408666758254363
step 5, overlap = 0.16709514342697387
step 6, overlap = 0.22789464746611884
step 7, overlap = 0.29506227870937746
step 8, overlap = 0.3670261018170988
step 9, overlap = 0.44210193536698855
step 10, overlap = 0.518532767034414
step 11, overlap = 0.5945298732465308
step 12, overlap = 0.6683146809800716
step 13, overlap = 0.7381603920041205
step 14, overlap = 0.8024323954287373
step 15, overlap = 0.8596265227777784
step 16, overlap = 0.9084042502960318
step 17, overlap = 0.947624024645163
step 18, overlap = 0.9763679788679578
step 19, overlap = 0.9939634133826709
step 20, overlap = 0.9999985392841676

function rand_circuit(nbit::Int, ngate::Int)
    circuit = chain(nbit)
    gate_list = [X, H, Ry(0.3), CNOT]
    for i = 1:ngate
        gate = rand(gate_list)
        push!(circuit, put(nbit, (sample(1:nbit, nqubits(gate),replace=false)...,)=>gate))
    end
    circuit
end
A = rand_circuit(nbit, 200)
psi0 = apply!(zero_state(nbit), A)

# now we want to search the subspace with [1,3,5,8,9,11,12]
# fixed to 1 and [4,6] fixed to 0.
evidense = [1, 3, -4, 5, -6, 8, 9, 11, 12]

"""
Doing Inference, psi is the initial state,
the target is to search target space with specific evidense.
e.g. evidense [1, -3, 6] means the [1, 3, 6]-th bits take value [1, 0, 1].
"""
oracle_infer = inference_oracle(evidense)(nqubits(psi0))

niter = num_grover_step(prob_match_oracle(psi0, oracle_infer))
gb_infer = chain(nbit, chain(oracle_infer, reflectblock(A)) for i = 1:niter);

for (i, blk) in enumerate(gb_infer)
    apply!(psi0, blk)
    p_target = prob_match_oracle(psi0, oracle_infer)
    println("step $i, overlap^2 = $p_target")
end

step 1, overlap^2 = 0.032262547409078794
step 2, overlap^2 = 0.08789567787386308
step 3, overlap^2 = 0.16730873626365186
step 4, overlap^2 = 0.26591929841198225
step 5, overlap^2 = 0.3780371739116427
step 6, overlap^2 = 0.4971927509305307
step 7, overlap^2 = 0.6165103164968588
step 8, overlap^2 = 0.7291048103066807
step 9, overlap^2 = 0.8284791179067036
step 10, overlap^2 = 0.9088989779814067
step 11, overlap^2 = 0.9657238702216389
step 12, overlap^2 = 0.9956747903334211

using Yao.Intrinsics

x1 = [0 1 0; 0 1 0; 0 1 0; 0 1 0; 0 1 0]
x2 = [1 1 1; 0 0 1; 1 1 1; 1 0 0; 1 1 1]
x0 = [1 1 1; 1 0 1; 1 0 1; 1 0 1; 1 1 1]

nbit = 15
v = zeros(1<<nbit)

# they occur with different probabilities.
for (x, p) in [(x0, 0.7), (x1, 0.29), (x2,0.01)]
    v[(x |> vec |> BitArray |> packbits)+1] = sqrt(p)
end

rb = reflect(copy(v))
psi0 = register(v)

# we want to find the digits with the first 5 qubits [1, 0, 1, 1, 1].
evidense = [1, -2, 3, 4, 5]
oracle_infer = inference_oracle(evidense)(nbit)

niter = num_grover_step(prob_match_oracle(psi0, oracle_infer))
gb_infer = chain(nbit, chain(oracle_infer, rb) for i = 1:niter)

Total: 15, DataType: Complex{Float64}
chain
├─ chain
│  ├─ control(1, 2, 3, 4)
│  │  └─ (5,)=>Z gate
│  └─
├─ chain
│  ├─ control(1, 2, 3, 4)
│  │  └─ (5,)=>Z gate
│  └─
├─ chain
│  ├─ control(1, 2, 3, 4)
│  │  └─ (5,)=>Z gate
│  └─
├─ chain
│  ├─ control(1, 2, 3, 4)
│  │  └─ (5,)=>Z gate
│  └─
├─ chain
│  ├─ control(1, 2, 3, 4)
│  │  └─ (5,)=>Z gate
│  └─
├─ chain
│  ├─ control(1, 2, 3, 4)
│  │  └─ (5,)=>Z gate
│  └─
└─ chain
   ├─ control(1, 2, 3, 4)
   │  └─ (5,)=>Z gate
   └─

for (i, blk) in enumerate(gb_infer)
    apply!(psi0, blk)
    p_target = prob_match_oracle(psi0, oracle_infer)
    println("step $i, overlap^2 = $p_target")
end

step 1, overlap^2 = 0.08761600000000003
step 2, overlap^2 = 0.23055362560000003
step 3, overlap^2 = 0.4161715569049601
step 4, overlap^2 = 0.6150679135961745
step 5, overlap^2 = 0.7957375127737549
step 6, overlap^2 = 0.9295622899279725
step 7, overlap^2 = 0.9953444003575993

pl = psi0 |> probs
config = findn(pl.>0.5)[] - 1 |> bitarray(nbit)
res = reshape(config, 5,3)

5×3 BitArray{2}:
  true   true   true
 false  false   true
  true   true   true
  true  false  false
  true   true   true