add ansatz

Author: xuanzhaogao

Project.toml

@@ -3,7 +3,18 @@ uuid = "1c4bdfd3-3b3e-42fc-86f5-70ab4d610eac"
 authors = ["ArrogantGao <xz.gao@connect.ust.hk> and contributors"]
 version = "1.0.0-DEV"
 
+[deps]
+GenericMessagePassing = "09522ef3-fe1e-4698-a056-61d75781c616"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
+Yao = "5872b779-8223-5990-8dd0-5abbb0748c8c"
+
 [compat]
+Graphs = "1.12.1"
+LinearAlgebra = "1.11.0"
+OMEinsum = "0.8.7"
+Yao = "0.9.1"
 julia = "1.6.7"
 
 [extras]

src/BPGauge.jl

@@ -1,5 +1,19 @@
 module BPGauge
 
-# Write your package code here.
+using LinearAlgebra, Graphs
+using OMEinsum
+using Yao
+using GenericMessagePassing # provide the `bp` backend
+
+export TensorNetworkAnsatz
+
+include("eins.jl")
+
+# define and construct the network
+include("ansatz.jl")
+include("rydberg.jl")
+
+# bp on the tensor network ansatz
+include("bp.jl")
 
 end

src/ansatz.jl

@@ -0,0 +1,55 @@
+# the tensor network ansatz, |ket> 
+struct TensorNetworkAnsatz{TA}
+    g::SimpleGraph{Int} # g is used to store the connections between the tensors
+    site_tensors::Vector{TA} # the tensors are stored as vector of arrays, corresponding to vertices of the graph. The virtual bounds are in order of its neighbors, the last dimension is the open dimension
+    gauge_tensors::Vector{TA} # corresponding to edges of the graph, listed in the order of edges(g), and enforce e.src < e.dst
+    gauge_tensors_map::Dict{Tuple{Int, Int}, Int} # maps between the edge and the index of the gauge tensor
+
+    function TensorNetworkAnsatz(g::SimpleGraph{Int}, d_virtual::Int, d_open::Int)
+        TA = Array{ComplexF64}
+
+        site_tensors = Vector{TA}()
+        gauge_tensors = Vector{TA}()
+
+        for v in vertices(g)
+            t = zeros(ComplexF64, [d_virtual for _ in 1:length(neighbors(g, v))]..., d_open)
+            push!(site_tensors, t)
+        end
+
+        for e in edges(g)
+            t = (one(ComplexF64) * I)(d_virtual)
+            push!(gauge_tensors, t)
+        end
+
+        map = Dict{Tuple{Int, Int}, Int}()
+        for (i, e) in enumerate(edges(g))
+            src, dst = minmax(e.src, e.dst)
+            map[(src, dst)] = i
+        end
+
+        new{Array{ComplexF64}}(g, site_tensors, gauge_tensors, map)
+    end
+end
+
+Base.show(io::IO, ansatz::TensorNetworkAnsatz{TA}) where {TA} = print(io, "TensorNetworkAnsatz{$(TA)} with $(nv(ansatz.g)) sites and $(ne(ansatz.g)) edges")
+
+function inner_product(bra::TensorNetworkAnsatz{TA}, ket::TensorNetworkAnsatz{TA}; optimizer = GreedyMethod()) where {TA}
+    @assert bra.g == ket.g
+
+    # 1:nv(g) are the indices of open indices
+    count = nv(bra.g) + 1
+    ixs_bra, count = all_eins(bra.g, count)
+    ixs_ket, count = all_eins(ket.g, count)
+
+    ixs = vcat(ixs_bra, ixs_ket)
+    raw_code = EinCode(ixs, Int[])
+
+    xs = [bra.site_tensors..., bra.gauge_tensors..., ket.site_tensors..., ket.gauge_tensors...]
+    size_dict = OMEinsum.get_size_dict!(ixs, xs, Dict{Int, Int}())
+
+    opt_code = optimize_code(raw_code, size_dict, optimizer)
+    @info "contraction complexity: $(contraction_complexity(opt_code, size_dict))"
+
+    res = opt_code(xs...)
+    return res[]
+end

src/bp.jl

@@ -0,0 +1,6 @@
+# bp on the TensorNetworkAnsatz
+
+struct BPState{TA}
+    message_e2v::Dict{Tuple{Int, Int}, TA}
+    message_v2e::Dict{Tuple{Int, Int}, TA}
+end

src/eins.jl

@@ -0,0 +1,26 @@
+# codes about constructing the einsum expression
+
+function all_eins(g::SimpleGraph{Int}, count::Int)
+    indices_dict = Dict{Tuple{Int, Int}, Int}()
+    ixs = Vector{Vector{Int}}()
+
+    # site tensors
+    for v in vertices(g)
+        ix = Int[]
+        for n in neighbors(g, v)
+            push!(ix, count)
+            indices_dict[(v, n)] = count
+            count += 1
+        end
+        push!(ix, v)
+        push!(ixs, ix)
+    end
+
+    # gauge tensors
+    for e in edges(g)
+        src, dst = minmax(e.src, e.dst)
+        push!(ixs, [indices_dict[(src, dst)], indices_dict[(dst, src)]])
+    end
+
+    return ixs, count
+end

src/guage.jl

@@ -0,0 +1,9 @@
+# initialize the state for rydberg system, all sites are in the |0> state
+function zero_state(g::SimpleGraph{Int}; d_virtual::Int = 1)
+    phi = TensorNetworkAnsatz(g, d_virtual, 2)
+    # set all to the |0> state
+    for tensor in phi.site_tensors
+        tensor[1] = ComplexF64(1.0)
+    end
+    return phi
+end