Add bipartite specialization of belief propagation and gauging (#319)

* Add bipartite specialization of belief propagation

* Fix formatting

* Fix bipartite SU test

* Rename variable

* Avoid double `map` calls

Author: Yue-Zhengyuan

src/algorithms/bp/beliefpropagation.jl

@@ -20,6 +20,9 @@ $(TYPEDFIELDS)
     "Toggle for projecting messages onto the hermitian subspace immediately after update through BP equation"
     project_hermitian::Bool = true
 
+    "When true, preserve bipartite structure of BPEnv inherited from input network"
+    bipartite::Bool = false
+
     "Output verbosity level"
     verbosity::Int = 2
 end
@@ -54,21 +57,43 @@ function leading_boundary(env₀::BPEnv, network::InfiniteSquareNetwork, alg::Be
         return env, ϵ
     end
 end
-function leading_boundary(env₀::BPEnv, state, args...)
-    return leading_boundary(env₀, InfiniteSquareNetwork(state), args...)
+function leading_boundary(env₀::BPEnv, state::InfiniteState, alg::BeliefPropagation)
+    if alg.bipartite
+        @assert _state_bipartite_check(state)
+    end
+    return leading_boundary(env₀, InfiniteSquareNetwork(state), alg)
 end
 
 """
 One iteration to update the BP environment.
 """
 function bp_iteration(network::InfiniteSquareNetwork, env::BPEnv, alg::BeliefPropagation)
-    messages = map(eachindex(env)) do I
-        M = update_message(I, network, env)
-        normalize!(M)
-        alg.project_hermitian && (M = project_hermitian!!(M))
-        return M
+    if alg.bipartite
+        @assert size(network, 1) == size(network, 2) == 2
+        messages = similar(env.messages)
+        for (d, r) in Iterators.product(1:4, 1:2)
+            # update BP env around 1st column of state
+            # [N/S, 1:2, 1], [E/W, 1:2, 2]
+            c = (d == NORTH || d == SOUTH) ? 1 : 2
+            I = CartesianIndex(d, r, c)
+            M = update_message(I, network, env)
+            normalize!(M)
+            alg.project_hermitian && (M = project_hermitian!!(M))
+            messages[I] = M
+            # copy to the other column
+            I′ = CartesianIndex(d, _next(r, 2), _next(c, 2))
+            messages[I′] = copy(M)
+        end
+        return BPEnv(messages)
+    else
+        messages = map(eachindex(env)) do I
+            M = update_message(I, network, env)
+            normalize!(M)
+            alg.project_hermitian && (M = project_hermitian!!(M))
+            return M
+        end
+        return BPEnv(messages)
     end
-    return BPEnv(messages)
 end
 
 """

src/algorithms/bp/gaugefix.jl

@@ -7,6 +7,16 @@ Algorithm for gauging PEPS with belief propagation fixed point messages.
     # TODO: add options
 end
 
+function _bpenv_bipartite_check(env::BPEnv)
+    for (r, c) in Iterators.product(1:2, 1:2)
+        r′, c′ = _next(r, 2), _next(c, 2)
+        if !all(env[:, r, c] .== env[:, r′, c′])
+            return false
+        end
+    end
+    return true
+end
+
 """
     gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::BPGauge, env::BPEnv)
 
@@ -15,11 +25,19 @@ an [`InfinitePEPO`](@ref) interpreted as purified state with two physical legs)
 using fixed point environment `env` of belief propagation.
 """
 function gauge_fix(psi::InfinitePEPS, alg::BPGauge, env::BPEnv)
+    bipartite = _state_bipartite_check(psi) && _bpenv_bipartite_check(env)
     psi′ = copy(psi)
     XXinv = map(eachcoordinate(psi, 1:2)) do I
         _, X, Xinv = _bp_gauge_fix!(CartesianIndex(I), psi′, env)
         return X, Xinv
     end
+    if bipartite
+        # copy 1st column to 2nd column to eliminate differences
+        # caused by order of applying gauge transformations
+        for r in 1:2
+            psi′[_next(r, 2), 2] = copy(psi′[r, 1])
+        end
+    end
     return psi′, XXinv
 end
 function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv)

src/algorithms/time_evolution/gaugefix_su.jl

@@ -36,13 +36,13 @@ function _trivial_gates(elt::Type{<:Number}, lattice::Matrix{S}) where {S <: Ele
 end
 
 """
-$(SIGNATURES)
+    gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::SUGauge)
 
 Fix the gauge of `psi` using trivial simple update.
 """
-function gauge_fix(psi::InfinitePEPS, alg::SUGauge)
+function gauge_fix(psi::InfiniteState, alg::SUGauge)
     gates = _trivial_gates(scalartype(psi), physicalspace(psi))
-    su_alg = SimpleUpdate(; trunc = FixedSpaceTruncation())
+    su_alg = SimpleUpdate(; trunc = FixedSpaceTruncation(), bipartite = _state_bipartite_check(psi))
     wts0 = SUWeight(psi)
     # use default constructor to avoid calculation of exp(-H * 0)
     evolver = TimeEvolver(su_alg, 0.0, alg.maxiter, gates, SUState(0, 0.0, psi, wts0))

src/algorithms/time_evolution/simpleupdate.jl

@@ -14,7 +14,7 @@ $(TYPEDFIELDS)
     imaginary_time::Bool = true
     "When true, force the usage of 3-site simple update"
     force_3site::Bool = false
-    "(Only applicable to InfinitePEPS) When true, assume bipartite unit cell structure"
+    "When true, assume bipartite unit cell structure"
     bipartite::Bool = false
     "(Only applicable to InfinitePEPO) 
     When true, the PEPO is regarded as a purified PEPS, and updated as

src/algorithms/time_evolution/time_evolve.jl

@@ -30,6 +30,22 @@ end
 
 Base.iterate(it::TimeEvolver) = iterate(it, it.state)
 
+function _state_bipartite_check(psi::InfiniteState)
+    if isa(psi, InfinitePEPO)
+        @assert size(psi, 3) == 1 "Input InfinitePEPO is expected to have only one layer."
+    end
+    if !(size(psi, 1) == size(psi, 2) == 2)
+        return false
+    end
+    for (r, c) in Iterators.product(1:2, 1:2)
+        r′, c′ = _next(r, 2), _next(c, 2)
+        if psi[r, c] != psi[r′, c′]
+            return false
+        end
+    end
+    return true
+end
+
 function _timeevol_sanity_check(
         ψ₀::InfiniteState, Pspaces::M, alg::A
     ) where {A <: TimeEvolution, M <: AbstractMatrix{<:ElementarySpace}}
@@ -40,7 +56,7 @@ function _timeevol_sanity_check(
         @assert ψ₀ isa InfinitePEPO "alg.purified = false is only applicable to PEPO."
     end
     if hasfield(typeof(alg), :bipartite) && alg.bipartite
-        @assert Nr == Nc == 2 "`bipartite = true` requires 2 x 2 unit cell size."
+        @assert _state_bipartite_check(ψ₀) "Input state is not bipartite with 2 x 2 unit cell."
     end
     return nothing
 end

test/timeevol/cluster_projectors.jl

@@ -4,7 +4,7 @@ using PEPSKit
 using LinearAlgebra
 using Random
 import MPSKitModels: hubbard_space
-using PEPSKit: sdiag_pow, _cluster_truncate!, _flip_virtuals!
+using PEPSKit: sdiag_pow, _cluster_truncate!, _flip_virtuals!, _next
 using MPSKit: GenericMPSTensor, MPSBondTensor
 include("cluster_tools.jl")
 
@@ -112,6 +112,10 @@ end
     trunc_env0 = truncerror(; atol = 1.0e-12) & truncrank(4)
     trunc_env = truncerror(; atol = 1.0e-12) & truncrank(16)
     peps = InfinitePEPS(rand, Float64, Pspace, Vspace, Vspace'; unitcell = (Nr, Nc))
+    # make state bipartite
+    for r in 1:2
+        peps.A[_next(r, 2), 2] = copy(peps.A[r, 1])
+    end
     wts = SUWeight(peps)
     ham = real(
         hubbard_model(

test/timeevol/sitedep_truncation.jl

@@ -3,7 +3,7 @@ using LinearAlgebra
 using Random
 using TensorKit
 using PEPSKit
-using PEPSKit: NORTH, EAST
+using PEPSKit: NORTH, EAST, _next
 
 function get_bonddims(peps::InfinitePEPS)
     xdims = collect(dim(domain(t, EAST)) for t in peps.A)
@@ -22,6 +22,10 @@ end
     ham = real(heisenberg_XYZ(InfiniteSquare(Nr, Nc); Jx = 1.0, Jy = 1.0, Jz = 1.0))
     Random.seed!(100)
     peps0 = InfinitePEPS(rand, Float64, ℂ^2, ℂ^10; unitcell = (Nr, Nc))
+    # make state bipartite
+    for r in 1:2
+        peps0.A[_next(r, 2), 2] = copy(peps0.A[r, 1])
+    end
     env0 = SUWeight(peps0)
     normalize!.(peps0.A, Inf)
     # set trunc to be compatible with bipartite structure

test/timeevol/tf_ising_finiteT.jl

@@ -59,6 +59,12 @@ dt, nstep = 1.0e-3, 400
     # PEPO approach: results at β, or T = 2.5
     alg = SimpleUpdate(; trunc = trunc_pepo, purified = false, bipartite)
     pepo, wts, info = time_evolve(pepo0, ham, dt, nstep, alg, wts0)
+
+    ## BP gauge fixing
+    bp_alg = BeliefPropagation(; maxiter = 100, tol = 1.0e-9, bipartite)
+    bp_env, = leading_boundary(BPEnv(ones, Float64, pepo), pepo, bp_alg)
+    pepo, = gauge_fix(pepo, BPGauge(), bp_env)
+
     env = converge_env(InfinitePartitionFunction(pepo), 16)
     result_β = measure_mag(pepo, env)
     @info "tr(σ(x,z)ρ) at T = $(1 / β): $(result_β)."