Merge pull request #7 from QuantumBFS/better-boundary-specification

boundary condition specification

Project.toml

@@ -17,6 +17,7 @@ julia = "1"
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 YaoExtensions = "7a06699c-c960-11e9-3c98-9f78548b5f0f"
+SymEngine = "123dc426-2d89-5057-bbad-38513e3affd8"
 
 [targets]
-test = ["Test", "OMEinsumContractionOrders", "YaoExtensions"]
+test = ["Test", "OMEinsumContractionOrders", "YaoExtensions", "SymEngine"]

README.md

@@ -32,18 +32,16 @@ pkg> add YaoToEinsum#master
 If you have problem to install the package, please [file us an issue](https://github.com/QuantumBFS/YaoToEinsum.jl/issues/new).
 
 ## Example
+This package contains one main function `yao2einsum(circuit; initial_state=Dict(), final_state=Dict())`.
+It transform a [`Yao`](https://github.com/QuantumBFS/Yao.jl) circuit to a generalized tensor network (einsum) notation. 
+This function returns a 2-tuple of (einsum code, input tensors). 
+`initial_state` and `final_state` specifies the initial state and final state.
+They can specified as a dictionary with integer keys, with value either integer or a single qubit register.
+If a qubit of initial state or final state is not specified, the circuit will have open edges.
+
 ```julia
 julia> import Yao, YaoToEinsum
 
-help?> YaoToEinsum.yao2einsum
-  yao2einsum(circuit; initial_state=nothing, final_state=nothing)
-
-  Transform a Yao circuit to a generalized tensor network (einsum) notation. 
-  This function returns a 2-tuple of (einsum code, input tensors). 
-  initial_state and final_state specifies the initial state and final state
-  as product states, e.g. a vector [1, 1, 0, 1] specifies a product state |1⟩⊗|1⟩⊗|0⟩⊗|1⟩. 
-  If initial state or final state is not specified, the circuit will have open edges.
-
 julia> using YaoExtensions: qft_circuit
 
 julia> using OMEinsumContractionOrders: optimize_code, TreeSA, uniformsize
@@ -62,7 +60,8 @@ julia> reshape(optcode(tensors...; size_info=uniformsize(code, 2)), 1<<n, 1<<n)
 true
 
 # convert circuit (applied on product state `initial_state` and projected to output state `final_state`)
-julia> code, tensors = YaoToEinsum.yao2einsum(circuit; initial_state=zeros(Bool, n), final_state=zeros(Bool, n));
+julia> code, tensors = YaoToEinsum.yao2einsum(circuit;
+        initial_state=Dict([i=>0 for i=1:n]), final_state=Dict([i=>0 for i=1:n]));
 
 julia> optcode = optimize_code(code, uniformsize(code, 2), TreeSA(ntrials=1));
 

src/YaoToEinsum.jl

@@ -13,19 +13,16 @@ struct EinBuilder{T}
 end
 
 Yao.nqubits(eb::EinBuilder) = length(eb.slots)
-
-function EinBuilder(n::Int)
-    EinBuilder(collect(1:n), Vector{Int}[], AbstractArray{ComplexF64}[], Ref(n))
+function add_tensor!(eb::EinBuilder{T}, tensor::AbstractArray{T,N}, labels::Vector{Int}) where {N,T}
+    @assert N == length(labels)
+    push!(eb.tensors, tensor)
+    push!(eb.labels, labels)
 end
-newlabel!(eb::EinBuilder) = (eb.maxlabel[] += 1; eb.maxlabel[])
 
-function add_product_state!(eb::EinBuilder{T}, bitstring) where T
-    for i=1:nqubits(eb)
-        push!(eb.tensors, bitstring[i] == 0 ? T[1, 0] : T[0, 1])
-        push!(eb.labels, [eb.slots[i]])
-    end
-    return eb
+function EinBuilder(::Type{T}, n::Int) where T
+    EinBuilder(collect(1:n), Vector{Int}[], AbstractArray{T}[], Ref(n))
 end
+newlabel!(eb::EinBuilder) = (eb.maxlabel[] += 1; eb.maxlabel[])
 
 function add_gate!(eb::EinBuilder{T}, b::PutBlock{N,C}) where {T, N,C}
     return add_matrix!(eb, C, mat(T, b.content), collect(b.locs))
@@ -34,12 +31,10 @@ end
 function add_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector) where T
     if !isdiag(m)
         nlabels = [newlabel!(eb) for _=1:k]
-        push!(eb.tensors, reshape(Matrix{T}(m), fill(2, 2k)...))  # need to check
-        push!(eb.labels, [nlabels..., eb.slots[locs]...])
+        add_tensor!(eb, reshape(Matrix{T}(m), fill(2, 2k)...), [nlabels..., eb.slots[locs]...])
         eb.slots[locs] .= nlabels
     else
-        push!(eb.tensors, reshape(Vector{T}(diag(m)), fill(2, k)...))  # need to check
-        push!(eb.labels, eb.slots[locs])
+        add_tensor!(eb, reshape(Vector{T}(diag(m)), fill(2, k)...), eb.slots[locs])
     end
     return eb
 end
@@ -63,8 +58,7 @@ function add_controlled_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, lo
     val = control_vals[1]
     for i=1:length(control_locs)-1
         newsig = newlabel!(eb)
-        push!(eb.labels, [newsig,eb.slots[control_locs[i+1]],sig])
-        push!(eb.tensors, and_gate(T, control_vals[i+1], val))
+        add_tensor!(eb, and_gate(T, control_vals[i+1], val), [newsig,eb.slots[control_locs[i+1]],sig])
         sig = newsig
         val = 1
     end
@@ -74,18 +68,16 @@ function add_controlled_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, lo
         if val == 1
             t1, t2 = t2, t1
         end
-        push!(eb.tensors, cat(t1, t2; dims=2k+1))  # need to check
         nlabels = [newlabel!(eb) for _=1:k]
-        push!(eb.labels, [nlabels..., eb.slots[locs]..., sig])
+        add_tensor!(eb, cat(t1, t2; dims=2k+1), [nlabels..., eb.slots[locs]..., sig])
         eb.slots[locs] .= nlabels
     else
         t1 = reshape(Vector{T}(diag(m)), fill(2, k)...)
         t2 = reshape(ones(T, 1<<k), fill(2, k)...)
         if val == 1
             t1, t2 = t2, t1
         end
-        push!(eb.tensors, cat(t1, t2; dims=k+1))  # need to check
-        push!(eb.labels, [eb.slots[locs]..., sig])
+        add_tensor!(eb, cat(t1, t2; dims=k+1), [eb.slots[locs]..., sig])
     end
     return eb
 end
@@ -109,7 +101,7 @@ function add_gate!(eb::EinBuilder, b::ChainBlock)
 end
 
 function add_gate!(eb::EinBuilder, b::AbstractBlock)
-    B = to_basic_types(b)
+    B = Optimise.to_basictypes(b)
     if typeof(B) == typeof(b)
         throw("block of type `$(typeof(b))` can not be converted to tensor network representation!")
     else
@@ -119,32 +111,66 @@ function add_gate!(eb::EinBuilder, b::AbstractBlock)
 end
 
 """
-    yao2einsum(circuit; initial_state=nothing, final_state=nothing)
+    yao2einsum(circuit, initial_state::Dict, final_state::Dict)
+    yao2einsum(circuit; initial_state=Dict(), final_state=Dict())
 
 Transform a Yao `circuit` to a generalized tensor network (einsum) notation.
 This function returns a 2-tuple of (einsum code, input tensors).
-`initial_state` and `final_state` specifies the initial state and final state as product states,
-e.g. a vector `[1, 1, 0, 1]` specifies a product state `|1⟩⊗|1⟩⊗|0⟩⊗|1⟩`.
-If initial state or final state is not specified, the circuit will have open edges.
+`initial_state` and `final_state` are dictionaries that specify the initial state and final state as product states,
+e.g. a vector `Dict(1=>1, 2=>1, 3=>0, 4=>1)` specifies a product state `|1⟩⊗|1⟩⊗|0⟩⊗|1⟩`.
+If an qubit in initial state or final state is not specified, it will be treated as an open edge.
+
+```jldoctest
+julia> using YaoToEinsum, Yao
+
+julia> c = chain(3, put(3, 2=>X), put(3, 1=>Y), control(3, 1, 3=>Y))
+nqubits: 3
+chain
+├─ put on (2)
+│  └─ X
+├─ put on (1)
+│  └─ Y
+└─ control(1)
+   └─ (3,) Y
+
+
+julia> yao2einsum(c; initial_state=Dict(1=>0, 2=>1), final_state=Dict(1=>ArrayReg([0.6, 0.8im]), 2=>1))
+(1, 2, 4∘2, 5∘1, 6∘3∘5, 5, 4 -> 6∘3, AbstractArray{ComplexF64}[[1.0 + 0.0im, 0.0 + 0.0im], [0.0 + 0.0im, 1.0 + 0.0im], [0.0 + 0.0im 1.0 + 0.0im; 1.0 + 0.0im 0.0 + 0.0im], [0.0 + 0.0im 0.0 - 1.0im; 0.0 + 1.0im 0.0 + 0.0im], [1.0 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 1.0 + 0.0im;;; 0.0 + 0.0im 0.0 - 1.0im; 0.0 + 1.0im 0.0 + 0.0im], [0.6 + 0.0im, 0.0 + 0.8im], [0.0 + 0.0im, 1.0 + 0.0im]])
+```
 """
-function yao2einsum(circuit::AbstractBlock; initial_state=nothing, final_state=nothing)
+function yao2einsum(circuit::AbstractBlock; initial_state::Dict=Dict{Int,Int}(), final_state::Dict=Dict{Int,Int}())
+    T = promote_type(ComplexF64, dict_regtype(initial_state), dict_regtype(final_state), Yao.parameters_eltype(circuit))
+    vec_initial_state = Dict{Int,ArrayReg{1,T}}([k=>render_single_qubit_state(T, v) for (k, v) in initial_state])
+    vec_final_state = Dict{Int,ArrayReg{1,T}}([k=>render_single_qubit_state(T, v) for (k, v) in final_state])
+    yao2einsum(circuit, vec_initial_state, vec_final_state)
+end
+dict_regtype(d::Dict) = promote_type(_regtype.(values(d))...)
+_regtype(::ArrayReg{1,VT}) where {VT} = VT
+_regtype(::Int) = ComplexF64
+render_single_qubit_state(::Type{T}, x::Int) where T = x == 0 ? zero_state(T, 1) : product_state(T, bit"1")
+render_single_qubit_state(::Type{T}, x::ArrayReg{1}) where T = ArrayReg(collect(T, statevec(x)))
+
+function yao2einsum(circuit::AbstractBlock, initial_state::Dict{Int,<:ArrayReg{1,T}}, final_state::Dict{Int,<:ArrayReg{1,T}}) where T
     n = nqubits(circuit)
-    eb = EinBuilder(nqubits(circuit))
+    eb = EinBuilder(T, n)
     openindices = Int[]
-    if initial_state===nothing
-        append!(openindices, eb.slots)
-    else
-        @assert n == length(initial_state)
-        add_product_state!(eb, initial_state)
+    for k=1:n
+        if haskey(initial_state, k)
+            add_tensor!(eb, statevec(initial_state[k]), [eb.slots[k]])
+        else
+            push!(openindices, eb.slots[k])
+        end
     end
     add_gate!(eb, circuit)
-    if final_state===nothing
-        prepend!(openindices, eb.slots)
-    else
-        @assert n == length(final_state)
-        add_product_state!(eb, final_state)
+    openindices2 = Int[]
+    for k=1:n
+        if haskey(final_state, k)
+            add_tensor!(eb, statevec(final_state[k]), [eb.slots[k]])
+        else
+            push!(openindices2, eb.slots[k])
+        end
     end
-    return build_einsum(eb, openindices)
+    return build_einsum(eb, vcat(openindices2, openindices))
 end
 
 function build_einsum(eb::EinBuilder, openindices)

test/runtests.jl

@@ -2,6 +2,7 @@ using YaoToEinsum
 using Test, OMEinsum, OMEinsumContractionOrders
 using Yao
 using YaoExtensions: qft_circuit, variational_circuit, rand_google53
+using SymEngine
 
 @testset "YaoToEinsum.jl" begin
     n = 5
@@ -10,16 +11,45 @@ using YaoExtensions: qft_circuit, variational_circuit, rand_google53
         @show c
         C = chain([put(n, i=>Rx(rand()*2π)) for i=1:n]..., c)
         code, xs = yao2einsum(C)
-        optcode = optimize_code(code, uniformsize(code, 2), GreedyMethod())
-        @test reshape(optcode(xs...; size_info=uniformsize(code, 2)), 1<<n, 1<<n) ≈ mat(C)
+        optcode = optimize_code(code, OMEinsumContractionOrders.uniformsize(code, 2), GreedyMethod())
+        @test reshape(optcode(xs...; size_info=OMEinsumContractionOrders.uniformsize(code, 2)), 1<<n, 1<<n) ≈ mat(C)
     end
 end
 
 @testset "Yao Extensions" begin
     n = 5
     for c in [qft_circuit(n), variational_circuit(n, 2), rand_google53(5; nbits=n)]
         code, xs = yao2einsum(c)
-        optcode = optimize_code(code, uniformsize(code, 2), TreeSA(nslices=3))
-        @test reshape(optcode(xs...; size_info=uniformsize(code, 2)), 1<<n, 1<<n) ≈ mat(c)
+        optcode = optimize_code(code, OMEinsumContractionOrders.uniformsize(code, 2), TreeSA(nslices=3))
+        @test reshape(optcode(xs...; size_info=OMEinsumContractionOrders.uniformsize(code, 2)), 1<<n, 1<<n) ≈ mat(c)
     end
 end
+
+@testset "boundary conditions" begin
+    n = 5
+    c = qft_circuit(n)
+    initial_state = Dict([i=>rand_state(1) for i=1:n])
+    reg = join([initial_state[i] for i=n:-1:1]...)
+    reg |> c
+    inner = (2,3)
+    focus!(reg, inner)
+    for final_state in [Dict([i=>rand_state(1) for i in inner]), Dict([i=>1 for i in inner])]
+        freg = join(YaoToEinsum.render_single_qubit_state(ComplexF64, final_state[3]), YaoToEinsum.render_single_qubit_state(ComplexF64, final_state[2]))
+        code, xs = yao2einsum(c; initial_state=initial_state, final_state=final_state)
+        optcode = optimize_code(code, OMEinsumContractionOrders.uniformsize(code, 2), TreeSA(nslices=3))
+        @test vec(optcode(xs...; size_info=OMEinsumContractionOrders.uniformsize(code, 2))) ≈ vec(transpose(statevec(freg)) * state(reg))
+    end
+end
+
+@testset "symbolic" begin
+    n = 5
+    c = qft_circuit(n)
+    initial_state = Dict([i=>zero_state(Basic, 1) for i=1:n])
+    code, xs = yao2einsum(c; initial_state=initial_state)
+    @test eltype(xs) == AbstractArray{Basic}
+end
+
+@testset "fix to basic type" begin
+    c = chain(kron(X,X))
+    @test yao2einsum(c)[1] isa EinCode
+end