Skip to content
This repository has been archived by the owner on Aug 11, 2024. It is now read-only.

Commit

Permalink
add type TensorNetwork
Browse files Browse the repository at this point in the history
  • Loading branch information
GiggleLiu committed Oct 5, 2023
1 parent 7087660 commit 549344b
Show file tree
Hide file tree
Showing 7 changed files with 301 additions and 266 deletions.
2 changes: 1 addition & 1 deletion Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "YaoToEinsum"
uuid = "9b173c7b-dc24-4dc5-a0e1-adab2f7b6ba9"
authors = ["GiggleLiu <cacate0129@gmail.com> and contributors"]
version = "0.1.5"
version = "0.2.0"

[deps]
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Expand Down
63 changes: 24 additions & 39 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,66 +6,51 @@ Convert Yao circuit to OMEinsum notation for tensor network based simulation.

## Installation

<p>
YaoToEinsum is a &nbsp;
<a href="https://julialang.org">
<img src="https://raw.githubusercontent.com/JuliaLang/julia-logo-graphics/master/images/julia.ico" width="16em">
Julia Language
</a>
&nbsp; package. To install YaoToEinsum,
please <a href="https://docs.julialang.org/en/v1/manual/getting-started/">open
Julia's interactive session (known as REPL)</a> and press <kbd>]</kbd> key in the REPL to use the package mode, then type the following command
</p>

For stable release
`YaoToEinsum` is a [Julia language](https://julialang.org/) package. To install `YaoToEinsum`, please [open Julia's interactive session (known as REPL)](https://docs.julialang.org/en/v1/manual/getting-started/) and press <kbd>]</kbd> key in the REPL to use the package mode, then type the following command

```julia
pkg> add YaoToEinsum
```

For current master

```julia
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.
This package contains one main function `yao2einsum(circuit; initial_state=Dict(), final_state=Dict(), optimizer=TreeSA())`.
It transform a [`Yao`](https://github.com/QuantumBFS/Yao.jl) circuit to a generalized tensor network (einsum notation). The return value is a `TensorNetwork` object.

* `initial_state` and `final_state` are for specifying the initial state and final state.
If any of them is not specified, the function will return a tensor network with open legs.
* `optimizer` is for optimizing the contraction order of the tensor network. The default value is `TreeSA()`. Please check the README of [OMEinsumContractors.jl](https://github.com/TensorBFS/OMEinsumContractionOrders.jl) for more information.

```julia
julia> import Yao, YaoToEinsum

julia> using Yao.EasyBuild: qft_circuit

julia> using YaoToEinsum: uniformsize, TreeSA, optimize_code
julia> using YaoToEinsum: TreeSA

julia> n = 10;

julia> circuit = qft_circuit(n);

# convert circuit (open in both left and right) to einsum notation (code) and tensors.
julia> code, tensors = YaoToEinsum.yao2einsum(circuit);
# convert this circuit to tensor network
julia> network = YaoToEinsum.yao2einsum(circuit)
TensorNetwork
Time complexity: 2^20.03816881914695
Space complexity: 2^20.0
Read-write complexity: 2^20.07564105083201

# optimize contraction order, for more algorithms, please check `OMEinsumContractionOrders`.
julia> optcode = optimize_code(code, uniformsize(code, 2), TreeSA(ntrials=1));

julia> reshape(optcode(tensors...; size_info=uniformsize(code, 2)), 1<<n, 1<<n) Yao.mat(circuit)
julia> reshape(contract(network), 1<<n, 1<<n) Yao.mat(circuit)
true

# convert circuit (applied on product state `initial_state` and projected to output state `final_state`)
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));
# convert circuit sandwiched by zero states
julia> network = YaoToEinsum.yao2einsum(circuit;
initial_state=Dict([i=>0 for i=1:n]), final_state=Dict([i=>0 for i=1:n]),
optimizer=TreeSA(; nslices=3)) # slicing technique
TensorNetwork
Time complexity: 2^12.224001674198101
Space complexity: 2^5.0
Read-write complexity: 2^13.036173612553485

julia> optcode(tensors...; size_info=uniformsize(code, 2))[] Yao.zero_state(n)' * (Yao.zero_state(n) |> circuit)
julia> contract(network)[] Yao.zero_state(n)' * (Yao.zero_state(n) |> circuit)
true
```

Expand Down
31 changes: 31 additions & 0 deletions src/Core.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
struct TensorNetwork
code::AbstractEinsum
tensors::Vector
end
function Base.show(io::IO, c::TensorNetwork)
print(io, "TensorNetwork")
print(io, "\n")
print(io, contraction_complexity(c))
end
function Base.show(io::IO, ::MIME"text/plain", c::TensorNetwork)
print(io, c)
end
function Base.iterate(c::TensorNetwork, state=1)
if state > 2
return nothing
elseif state == 1
return (c.code, 2)
else
return (c.tensors, 3)
end
end
function contract(c::TensorNetwork)
return c.code(c.tensors...; size_info=uniformsize(c.code, 2))
end
function OMEinsum.optimize_code(c::TensorNetwork, args...)
optcode = optimize_code(c.code, uniformsize(c.code, 2), args...)
return TensorNetwork(optcode, c.tensors)
end
function OMEinsum.contraction_complexity(c::TensorNetwork)
return contraction_complexity(c.code, uniformsize(c.code, 2))
end
176 changes: 4 additions & 172 deletions src/YaoToEinsum.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,178 +3,10 @@ module YaoToEinsum
using Yao, OMEinsum
using LinearAlgebra

export yao2einsum, add_gate!
export yao2einsum
export TensorNetwork, optimize_code, contraction_complexity, contract

struct EinBuilder{T}
slots::Vector{Int}
labels::Vector{Vector{Int}}
tensors::Vector{AbstractArray{T}}
maxlabel::Base.RefValue{Int}
end

Yao.nqubits(eb::EinBuilder) = length(eb.slots)
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

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{D,C}) where {T,D,C}
return add_matrix!(eb, C, mat(T, b.content), collect(b.locs))
end
# general and diagonal gates
function add_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector) where T
if !isdiag(m)
nlabels = [newlabel!(eb) for _=1:k]
add_tensor!(eb, reshape(Matrix{T}(m), fill(2, 2k)...), [nlabels..., eb.slots[locs]...])
eb.slots[locs] .= nlabels
else
add_tensor!(eb, reshape(Vector{T}(diag(m)), fill(2, k)...), eb.slots[locs])
end
return eb
end
# swap gate
function add_gate!(eb::EinBuilder{T}, b::PutBlock{2,2,ConstGate.SWAPGate}) where {T,N}
lj = eb.slots[b.locs[2]]
eb.slots[b.locs[2]] = eb.slots[b.locs[1]]
eb.slots[b.locs[1]] = lj
return eb
end

# control gates
function add_gate!(eb::EinBuilder{T}, b::ControlBlock{BT,C,M}) where {T, BT,C,M}
return add_controlled_matrix!(eb, M, mat(T, b.content), collect(b.locs), collect(b.ctrl_locs), collect(b.ctrl_config))
end
function add_controlled_matrix!(eb::EinBuilder{T}, k::Int, m::AbstractMatrix, locs::Vector, control_locs, control_vals) where T
if length(control_locs) == 0
return add_matrix!(eb, k, m, locs)
end
sig = eb.slots[control_locs[1]]
val = control_vals[1]
for i=1:length(control_locs)-1
newsig = newlabel!(eb)
add_tensor!(eb, and_gate(T, control_vals[i+1], val), [newsig,eb.slots[control_locs[i+1]],sig])
sig = newsig
val = 1
end
if !isdiag(m)
t1 = reshape(Matrix{T}(m), fill(2, 2k)...)
t2 = reshape(Matrix{T}(I, 1<<k, 1<<k), fill(2, 2k)...)
if val == 1
t1, t2 = t2, t1
end
nlabels = [newlabel!(eb) for _=1:k]
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
add_tensor!(eb, cat(t1, t2; dims=k+1), [eb.slots[locs]..., sig])
end
return eb
end

function and_gate(::Type{T}, a::Int, b::Int) where T
m = zeros(T, 2, 2, 2)
for v1 in (0, 1)
for v2 in (0, 1)
# the first is output
m[(v1==a && v2==b)+1, v1+1,v2+1] = 1
end
end
return m
end

function add_gate!(eb::EinBuilder, b::ChainBlock)
for ib in subblocks(b)
add_gate!(eb, ib)
end
return eb
end

function add_gate!(eb::EinBuilder, b::AbstractBlock)
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
add_gate!(eb, B)
end
return eb
end

"""
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` 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{D}; initial_state::Dict=Dict{Int,Int}(), final_state::Dict=Dict{Int,Int}()) where {D}
T = promote_type(ComplexF64, dict_regtype(initial_state), dict_regtype(final_state), Yao.parameters_eltype(circuit))
vec_initial_state = Dict{Int,ArrayReg{D,T}}([k=>render_single_qubit_state(T, v) for (k, v) in initial_state])
vec_final_state = Dict{Int,ArrayReg{D,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{D,VT}) where {D,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) where T = ArrayReg(collect(T, statevec(x)))

function yao2einsum(circuit::AbstractBlock{D}, initial_state::Dict{Int,<:ArrayReg{D,T}}, final_state::Dict{Int,<:ArrayReg{D,T}}) where {D,T}
n = nqubits(circuit)
eb = EinBuilder(T, n)
openindices = Int[]
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)
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, vcat(openindices2, openindices))
end

function build_einsum(eb::EinBuilder, openindices)
return EinCode(eb.labels, openindices), eb.tensors
end
include("Core.jl")
include("circuitmap.jl")

end
Loading

0 comments on commit 549344b

Please sign in to comment.