Merge pull request #17 from CodingThrust/dev

BP gauge for graph tensor network state

Author: xuanzhaogao

README.md

@@ -1,9 +1,7 @@
 # BPGauge
 
-<!-- [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://ArrogantGao.github.io/BPGauge.jl/stable/) -->
-<!-- [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://ArrogantGao.github.io/BPGauge.jl/dev/) -->
-<!-- [![Build Status](https://github.com/ArrogantGao/BPGauge.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/ArrogantGao/BPGauge.jl/actions/workflows/CI.yml?query=branch%3Amain) -->
-<!-- [![Coverage](https://codecov.io/gh/ArrogantGao/BPGauge.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/ArrogantGao/BPGauge.jl) -->
+[![Build Status](https://github.com/CodingThrust/BPGauge.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/CodingThrust/BPGauge.jl/actions/workflows/CI.yml?query=branch%3Amain)
+<!-- [![Coverage](https://codecov.io/gh/CodingThrust/BPGauge.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/CodingThrust/BPGauge.jl) -->
 
 ## Introduction
 
@@ -13,7 +11,7 @@ A simple example is to consider a cycle graph:
 ```julia
 using BPGauge, Graphs
 
-g = chain(100)
+g = path_graph(100)
 add_edge!(g, 1, 100)
 
 # generate a random state
@@ -23,7 +21,7 @@ normalize_state!(tn)
 # generate a BP state, and do BP until convergence
 bp_state = BPState(tn)
 bp_path = BPPath(tn)
-bp!(bp_state, bp_path, tn, err_bound = 1e-8)
+bp!(bp_state, bp_path, tn, atol = 1e-8)
 
 # gauge transform the state
 gauge!(tn, bp_state)

exm/Note.md

@@ -12,7 +12,7 @@ More complex geometry and property will reduce more the dimension of system by i
 
 $D=2^N \xrightarrow{\text{Constraint}} \alpha^N  \xrightarrow{\text{Particle Number}}\alpha^{N-1}/k\xrightarrow{\text{translation}}\alpha^{N-1}/kN\xrightarrow{\text{spin flip/inversion}}\alpha^{N-2}/kN$.
 
-where $\alpha \sim 1.618, 1.502$ for Fibonacci chain, hard core square lattice, $k \sim 5$. So for fully ED, $D \sim 10^6$. When doing dynamics, Krylov subspace $K_m​=\text{span}\{v,Hv,H^2v,⋯,H^{m−1}v\}$ is usually incorporated, which can scale up to $N=32$ even without anyon symmetry ($D \sim 10^{10}, m \sim 10^4$,  as long as the locality of interaction ensuring sparse matrix). In such basis, the Its complexity decreases to $O(m^2D)$
+where $\alpha \sim 1.618, 1.502$ for Fibonacci path_graph, hard core square lattice, $k \sim 5$. So for fully ED, $D \sim 10^6$. When doing dynamics, Krylov subspace $K_m​=\text{span}\{v,Hv,H^2v,⋯,H^{m−1}v\}$ is usually incorporated, which can scale up to $N=32$ even without anyon symmetry ($D \sim 10^{10}, m \sim 10^4$,  as long as the locality of interaction ensuring sparse matrix). In such basis, the Its complexity decreases to $O(m^2D)$
 
 
 ## *Tensor Networks, TN*
@@ -95,7 +95,7 @@ Supplemental Material](https://arxiv.org/pdf/1911.04882), which maybe helpful?
 <img width="872" alt="Image" src="https://github.com/user-attachments/assets/df2cf6ed-da28-4320-9f28-335b383a5cbb" />
 
 [^Leshouches]: Les houch notes for Exact Diagonalizaton. https://indico.ictp.it/event/a14246/session/31/contribution/51/material/0/0.pdf
-[^QMBSPRB]: Quantum scarred eigenstates in a Rydberg atom chain: Entanglement, breakdown of thermalization, and stability to perturbations. https://journals.aps.org/prb/pdf/10.1103/PhysRevB.98.155134
+[^QMBSPRB]: Quantum scarred eigenstates in a Rydberg atom path_graph: Entanglement, breakdown of thermalization, and stability to perturbations. https://journals.aps.org/prb/pdf/10.1103/PhysRevB.98.155134
 [^LauchliKagome]: S = 1/2 kagome Heisenberg antiferromagnet revisited https://journals.aps.org/prb/pdf/10.1103/PhysRevB.100.155142
 [^Dowling]: Monte Carlo techniques for real-time
 quantum dynamics https://arxiv.org/pdf/quant-ph/0507003

src/BPGauge.jl

@@ -0,0 +1,30 @@
+module BPGauge
+
+using LinearAlgebra, Graphs, Random
+using OMEinsum
+
+export TensorNetworkAnsatz
+export BPState, BPPath, BPStep
+
+export zero_state, random_state, inner_product, normalize_state!
+export bp!, bp_update!
+export absorb!
+export apply_gauge!, gauge!
+
+export square_lattice
+
+include("utils.jl")
+
+# define and construct the network
+include("ansatz.jl")
+
+# bp on the tensor network ansatz
+include("bp.jl")
+
+# operations about gauging
+include("gauge.jl")
+
+# simple update
+include("su.jl")
+
+end

src/ansatz.jl

@@ -0,0 +1,83 @@
+# the tensor network ansatz, |ket> 
+struct TensorNetworkAnsatz{TA <: AbstractArray, TB <: AbstractArray}
+    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{TB} # corresponding to edges of the graph, listed in the order of edges(g), and enforce e.src < e.dst, real valued diagonal matrices
+    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; type::Type = Float64)
+        TA = Array{Complex{type}}
+        TB = Array{type}
+
+        site_tensors = Vector{TA}()
+        gauge_tensors = Vector{TB}()
+
+        for v in vertices(g)
+            t = zeros(Complex{type}, [d_virtual for _ in 1:length(neighbors(g, v))]..., d_open)
+            push!(site_tensors, t)
+        end
+
+        for e in edges(g)
+            t = (one(type) * 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{TA, TB}(g, site_tensors, gauge_tensors, map)
+    end
+
+    function TensorNetworkAnsatz(g::SimpleGraph{Int}, site_tensors::Vector{TA}, gauge_tensors::Vector{TB}, map::Dict{Tuple{Int, Int}, Int}) where {TA, TB}
+        @assert length(site_tensors) == nv(g)
+        @assert length(gauge_tensors) == ne(g)
+        new{TA, TB}(g, site_tensors, gauge_tensors, map)
+    end
+end
+
+Base.show(io::IO, ansatz::TensorNetworkAnsatz{TA, TB}) where {TA, TB} = print(io, "TensorNetworkAnsatz{$(TA), $(TB)} with $(nv(ansatz.g)) sites and $(ne(ansatz.g)) edges")
+
+Base.adjoint(ansatz::TensorNetworkAnsatz{TA, TB}) where {TA, TB} = TensorNetworkAnsatz(ansatz.g, TA[conj(t) for t in ansatz.site_tensors], ansatz.gauge_tensors, ansatz.gauge_tensors_map)
+
+function inner_product(bra::TensorNetworkAnsatz{TA, TB}, ket::TensorNetworkAnsatz{TA}; optimizer = GreedyMethod()) where {TA, TB}
+    @assert bra.g == ket.g
+
+    raw_code = inner_product_eins(bra.g)
+
+    xs = [bra.site_tensors..., bra.gauge_tensors..., ket.site_tensors..., ket.gauge_tensors...]
+    size_dict = OMEinsum.get_size_dict!(raw_code.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
+
+# 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
+
+# generate non-normalized random state
+function random_state(g::SimpleGraph{Int}; d_virtual::Int = 1)
+    phi = TensorNetworkAnsatz(g, d_virtual, 2)
+    for tensor in phi.site_tensors
+        tensor .= rand(ComplexF64, size(tensor))
+    end
+    return phi
+end
+
+function normalize_state!(phi::TensorNetworkAnsatz{TA, TB}) where {TA, TB}
+    n = inner_product(phi, adjoint(phi))
+    phi.site_tensors .= phi.site_tensors ./ (sqrt(real(n))^(1 / nv(phi.g)))
+    return phi
+end

src/bp.jl

@@ -0,0 +1,99 @@
+# bp on the TensorNetworkAnsatz
+
+# each message tensor have two legs, as follows
+# --Ta --   |--1--Tb --
+#          Mab
+# --Ta*--   |--2--Tb*--
+struct BPState{TA}
+    messages::Dict{Tuple{Int, Int}, TA}
+    function BPState(ansatz::TensorNetworkAnsatz{TA, TB}) where {TA, TB}
+        messages = Dict{Tuple{Int, Int}, TA}()
+        TE = eltype(TA)
+
+        # initialize the messages from neighbors to sites
+        for dst in vertices(ansatz.g)
+            for (i, src) in enumerate(neighbors(ansatz.g, dst))
+                d_virtual = size(ansatz.site_tensors[dst])[i]
+                messages[(src, dst)] = ones(TE, d_virtual, d_virtual)
+                normalize_message!(messages[(src, dst)])
+            end
+        end
+
+        new{TA}(messages)
+    end
+end
+
+# a single bp update step: \sum_{T_neighbors} T * T* * M_inputs... -> M_output
+struct BPStep
+    eincode::AbstractEinsum # Eincode([T..., T*..., M_inputs...], [M_output])
+    Tid::Int # the id of the tensor to interact with
+    inputs::Vector{Tuple{Int, Int}} # input messages 
+    output::Tuple{Int, Int} # output message
+end
+
+struct BPPath
+    bp_steps::Vector{BPStep} # steps of bp update, each step is a single bp update
+    function BPPath(ansatz::TensorNetworkAnsatz{TA, TB}; random_order::Bool = true, seed::Int = 1234) where {TA, TB}
+        g = ansatz.g
+        Random.seed!(seed)
+        bp_steps = BPStep[]
+        # the different step are taken in random order
+        outputs = Tuple{Int, Int}[]
+        for e in edges(g)
+            push!(outputs, (src(e), dst(e)))
+            push!(outputs, (dst(e), src(e)))
+        end
+
+        # the steps are taken in random order
+        random_order && shuffle!(outputs)
+
+        for (s, d) in outputs
+            inputs = Tuple{Int, Int}[]
+            Tid = s
+            for v in neighbors(g, s)
+                if v != d
+                    push!(inputs, (v, s))
+                end
+            end
+
+            eincode = bp_eins(neighbors(g, s), inputs, (s, d))
+            push!(bp_steps, BPStep(eincode, Tid, inputs, (s, d)))
+        end
+
+        new(bp_steps)
+    end
+end
+
+function bp!(state::BPState, path::BPPath, ansatz::TensorNetworkAnsatz; max_iter::Int = 1000, atol::Float64 = 1e-6, damping::Float64 = 0.2, verbose::Bool = false)
+    Ts = ansatz.site_tensors
+    Tcs = conj.(ansatz.site_tensors)
+    for iter in 1:max_iter
+        error = bp_update!(state, path, Ts, Tcs, damping)
+        verbose && @info "BP iteration $iter, error: $error"
+        error < atol && break
+    end
+    nothing
+end
+
+function bp_update!(state::BPState, path::BPPath, Ts::Vector{TA}, Tcs::Vector{TB}, damping::Float64) where {TA, TB}
+
+    error = 0.0
+    for step in path.bp_steps
+        eincode = step.eincode
+        Tid = step.Tid
+        inputs = step.inputs
+        output = step.output
+
+        T = Ts[Tid]
+        Tc = Tcs[Tid]
+
+        M_inputs = [state.messages[input] for input in inputs]
+
+        # update the message tensor
+        new_message = normalize_message!(eincode(T, Tc, M_inputs...))
+        error = max(error, maximum(abs.(state.messages[output] - new_message)))
+        state.messages[output] .*= damping
+        state.messages[output] .+= (1 - damping) * new_message
+    end
+    return error
+end

src/gauge.jl

@@ -0,0 +1,87 @@
+# absorb the gauge tensors (diagonal matrices) into the site tensors, for bp to work
+# after the gauge is absorbed, set them to identity
+function absorb!(ansatz::TensorNetworkAnsatz{TA, TB}) where {TA, TB}
+    for e in edges(ansatz.g)
+        s, d = minmax(e.src, e.dst)
+        Gamma = ansatz.gauge_tensors[ansatz.gauge_tensors_map[(s, d)]]
+        sqrt_Gamma = sqrt.(Gamma)
+
+        eincode_s, eincode_d = absorb_eins(s, d, neighbors(ansatz.g, s), neighbors(ansatz.g, d))
+
+        s_xs = (ansatz.site_tensors[s], sqrt_Gamma)
+        sds = OMEinsum.get_size_dict(eincode_s.ixs, s_xs)
+        einsum!(eincode_s.ixs, eincode_s.iy, s_xs, ansatz.site_tensors[s], 1, 0, sds)
+
+        d_xs = (ansatz.site_tensors[d], sqrt_Gamma)
+        sdd = OMEinsum.get_size_dict(eincode_d.ixs, d_xs)
+        einsum!(eincode_d.ixs, eincode_d.iy, d_xs, ansatz.site_tensors[d], 1, 0, sdd)
+
+        # reset the gauge tensor to identity
+        Gamma .= I(size(Gamma, 1))
+    end
+    nothing
+end
+
+# assume that the existing gauge tensors are already absorbed into the site tensors
+# the bp converge, using the bp messages to generate the new gauge tensors, which are real valued diagonal matrices
+# according to Eqs.12~17 of https://scipost.org/SciPostPhys.15.6.222
+# site s and site d
+# Ts -i- sqrt_Msd_inv (j, i) -j- sqrt_Msd (k, j) -k- sqrt_Mds (k, l) -l- sqrt_Mds_inv (l, m) -m- Td
+# Ts -i- sqrt_Msd_inv (j, i) -j- U(j, n) -n- S(n, o) -o- Vt(o, l) -l- sqrt_Mds_inv (l, m) -m- Td
+# Ts -i- As (i, n) -n- Gamma -o- Ad (m, o) -m- Td
+# Ts -i- Gamma -m- Td
+function gauge!(ansatz::TensorNetworkAnsatz{TA, TB}, state::BPState{TA}) where {TA, TB}
+    for e in edges(ansatz.g)
+        s, d = minmax(e.src, e.dst)
+        M_sd = state.messages[(s, d)]
+        M_ds = state.messages[(d, s)]
+
+        apply_gauge!(ansatz, s, d, M_sd, M_ds)
+    end
+    nothing
+end
+
+# given the message matrix M_sd and M_ds, generate the new gauge tensor Gamma, update T_s and T_d
+function apply_gauge!(ansatz::TensorNetworkAnsatz{TA, TB}, s::Int, d::Int, M_sd::Matrix{T}, M_ds::Matrix{T}) where {TA, TB, T}
+    # square root the message matrix
+    sqrt_Msd, sqrt_Msd_inv = square_root(M_sd)
+    sqrt_Mds, sqrt_Mds_inv = square_root(M_ds)
+
+    # the eigen decomposition is only correct when the message matrix is positive definite
+    # @show ein"ji, kj -> ik"(sqrt_Msd_inv, sqrt_Msd) ≈ I(size(M_sd, 1))
+    # @show ein"kl, lm -> km"(sqrt_Mds, sqrt_Mds_inv) ≈ I(size(M_sd, 1))
+    # @show ein"ji, kj, kl, lm -> im"(sqrt_Msd_inv, sqrt_Msd, sqrt_Mds, sqrt_Mds_inv) ≈ I(size(M_sd, 1))
+
+    # generate the new gauge tensor
+    M_mid = ein"kj, kl -> jl"(sqrt_Msd, sqrt_Mds)
+
+    # @show ein"ji, jl, lm -> im"(sqrt_Msd_inv, M_mid, sqrt_Mds_inv) ≈ I(size(M_mid, 1))
+
+    res = svd(M_mid)
+    U = res.U # jn
+    S = diagm(res.S) # no
+    Vt = res.Vt # ol
+
+    # @show ein"jn, no, ol -> jl"(U, S, Vt) ≈ M_mid
+    # @show maximum(abs.(ein"ji, jn, no, ol, lm -> im"(sqrt_Msd_inv, U, S, Vt, sqrt_Mds_inv) - I(size(M_mid, 1))))
+
+    As = ein"ji, jn -> in"(sqrt_Msd_inv, U)
+    Ad = ein"lm, ol -> mo"(sqrt_Mds_inv, Vt)
+
+    # @show maximum(abs.(ein"in, no, mo -> im"(As, S, Ad) - I(size(As, 1))))
+
+    #update the gauge tensor
+    ansatz.gauge_tensors[ansatz.gauge_tensors_map[(s, d)]] .= S
+
+    # absorb As and Ad into the site tensors
+    eincode_s, eincode_d = absorb_eins(s, d, neighbors(ansatz.g, s), neighbors(ansatz.g, d))
+    s_xs = (ansatz.site_tensors[s], As)
+    sds = OMEinsum.get_size_dict(eincode_s.ixs, s_xs)
+    einsum!(eincode_s.ixs, eincode_s.iy, s_xs, ansatz.site_tensors[s], 1, 0, sds)
+
+    d_xs = (ansatz.site_tensors[d], Ad)
+    sdd = OMEinsum.get_size_dict(eincode_d.ixs, d_xs)
+    einsum!(eincode_d.ixs, eincode_d.iy, d_xs, ansatz.site_tensors[d], 1, 0, sdd)
+
+    nothing
+end

src/su.jl

@@ -0,0 +1,4 @@
+# TODO
+# 1. single site update
+# 2. two site update
+# 3. truncation