Belief propagation for purified iPEPO

src/PEPSKit.jl

@@ -121,7 +121,9 @@ export InfinitePEPS, InfiniteTransferPEPS
 export SUWeight
 export InfinitePEPO, InfiniteTransferPEPO
 
-export BPEnv, BeliefPropagation, BPGauge
+export BPEnv, BeliefPropagation
+export BPGauge, SUGauge
+export gauge_fix
 
 export initialize_mps, initializePEPS
 export ReflectDepth, ReflectWidth, Rotate, RotateReflect

src/algorithms/bp/gaugefix.jl

@@ -8,9 +8,11 @@ Algorithm for gauging PEPS with belief propagation fixed point messages.
 end
 
 """
-$(SIGNATURES)
+    gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::BPGauge, env::BPEnv)
 
-Fix the gauge of `psi` using fixed point environment `env` of belief propagation.
+Fix the gauge of `psi` (which can be an [`InfinitePEPS`](@ref), or
+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)
     psi′ = copy(psi)
@@ -20,6 +22,24 @@ function gauge_fix(psi::InfinitePEPS, alg::BPGauge, env::BPEnv)
     end
     return psi′, XXinv
 end
+function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv)
+    # convert to iPEPS, and store physical leg fusers
+    Nr, Nc, Nh = size(psi)
+    @assert Nh == 1
+    psi_Fs = map(psi.A) do t
+        return fuse_physicalspaces(t)
+    end
+    psi_Fs = reshape(psi_Fs, (Nr, Nc))
+    psi′ = map(Base.Fix2(getindex, 1), psi_Fs)
+    Fs = map(Base.Fix2(getindex, 2), psi_Fs)
+    psi′, XXinv = gauge_fix(InfinitePEPS(psi′), alg, env)
+    # convert back to iPEPO
+    psi′ = map(zip(psi′.A, Fs)) do (t, F)
+        return F' * t
+    end
+    psi′ = reshape(psi′, (Nr, Nc, 1))
+    return InfinitePEPO(psi′), XXinv
+end
 
 function _sqrt_bp_messages(I::CartesianIndex{3}, env::BPEnv)
     dir, row, col = Tuple(I)

src/algorithms/ctmrg/gaugefix.jl

@@ -6,7 +6,7 @@ This assumes that the `envfinal` is the result of one CTMRG iteration on `envpre
 Given that the CTMRG run is converged, the returned environment will be
 element-wise converged to `envprev`.
 """
-function gauge_fix(envprev::CTMRGEnv{C, T}, envfinal::CTMRGEnv{C, T}) where {C, T}
+function env_gauge_fix(envprev::CTMRGEnv{C, T}, envfinal::CTMRGEnv{C, T}) where {C, T}
     # Check if spaces in envprev and envfinal are the same
     same_spaces = map(eachcoordinate(envfinal, 1:4)) do (dir, r, c)
         space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) &&

src/algorithms/optimization/fixed_point_differentiation.jl

@@ -221,9 +221,9 @@ function _rrule(
     function leading_boundary_diffgauge_pullback((Δenv′, Δinfo))
         Δenv = unthunk(Δenv′)
 
-        # find partial gradients of gauge_fixed single CTMRG iteration
+        # find partial gradients of gauge-fixed single CTMRG iteration
         function f(A, x)
-            return gauge_fix(x, ctmrg_iteration(InfiniteSquareNetwork(A), x, alg_fixed)[1])[1]
+            return env_gauge_fix(x, ctmrg_iteration(InfiniteSquareNetwork(A), x, alg_fixed)[1])[1]
         end
         _, env_vjp = rrule_via_ad(config, f, state, env)
 
@@ -250,7 +250,7 @@ function _rrule(
     env, = leading_boundary(envinit, state, alg)
     alg_fixed = @set alg.projector_alg.trunc = FixedSpaceTruncation() # fix spaces during differentiation
     env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed)
-    env_fixed, signs = gauge_fix(env, env_conv)
+    env_fixed, signs = env_gauge_fix(env, env_conv)
 
     # Fix SVD
     svd_alg_fixed = _fix_svd_algorithm(alg.projector_alg.svd_alg, signs, info)

src/environments/bp_environments.jl

@@ -134,12 +134,12 @@ end
 function BPEnv(state::InfinitePartitionFunction, args...; kwargs...)
     return BPEnv(InfiniteSquareNetwork(state), args...; kwargs...)
 end
-function BPEnv(state::InfinitePEPS, args...; kwargs...)
+function BPEnv(state::Union{InfinitePEPS, InfinitePEPO}, args...; kwargs...)
     bp_env = BPEnv(InfiniteSquareNetwork(state), args...; kwargs...)
     TensorKit.id!.(bp_env.messages)
     return bp_env
 end
-function BPEnv(f, T, state::Union{InfinitePartitionFunction, InfinitePEPS}, args...; kwargs...)
+function BPEnv(f, T, state::Union{InfinitePartitionFunction, InfinitePEPS, InfinitePEPO}, args...; kwargs...)
     return BPEnv(f, T, InfiniteSquareNetwork(state), args...; kwargs...)
 end
 

src/operators/infinitepepo.jl

@@ -182,6 +182,10 @@ function InfiniteSquareNetwork(top::InfinitePEPS, mid::InfinitePEPO, bot::Infini
     )
 end
 
+function InfiniteSquareNetwork(top::InfinitePEPO, bot::InfinitePEPO = top)
+    return InfiniteSquareNetwork(InfinitePEPS(top), InfinitePEPS(bot))
+end
+
 ## Conversion to states
 
 function InfinitePartitionFunction(ρ::InfinitePEPO)

test/bp/expvals.jl

@@ -2,7 +2,7 @@ using Test
 using Random
 using TensorKit
 using PEPSKit
-using PEPSKit: random_dual!, SUGauge, gauge_fix
+using PEPSKit: random_dual!
 
 ds = Dict(
     U1Irrep => U1Space(i => d for (i, d) in zip(-1:1, (1, 1, 2))),

test/bp/gaugefix.jl

@@ -2,7 +2,7 @@ using Test, TestExtras
 using Random
 using TensorKit
 using PEPSKit
-using PEPSKit: SUGauge, gauge_fix, compare_weights, random_dual!, twistdual
+using PEPSKit: compare_weights, random_dual!, twistdual
 
 @testset "Compare BP and SU ($S, posdef msgs = $h)" for (S, h) in
     Iterators.product([U1Irrep, FermionParity], [true, false])

test/ctmrg/fixed_iterscheme.jl

@@ -9,7 +9,7 @@ using PEPSKit
 using PEPSKit:
     FixedSVD,
     ctmrg_iteration,
-    gauge_fix,
+    env_gauge_fix,
     fix_relative_phases,
     fix_global_phases,
     calc_elementwise_convergence,
@@ -40,7 +40,7 @@ atol = 1.0e-5
 
     # do extra iteration to get SVD
     env_conv2, info = @constinferred ctmrg_iteration(n, env_conv1, ctm_alg)
-    env_fix, signs = gauge_fix(env_conv1, env_conv2)
+    env_fix, signs = env_gauge_fix(env_conv1, env_conv2)
     @test calc_elementwise_convergence(env_conv1, env_fix) ≈ 0 atol = atol
 
     # fix gauge of SVD
@@ -70,11 +70,11 @@ end
 
     # do extra iteration to get SVD
     env_conv2_iter, info_iter = ctmrg_iteration(n, env_conv1, ctm_alg_iter)
-    env_fix_iter, signs_iter = gauge_fix(env_conv1, env_conv2_iter)
+    env_fix_iter, signs_iter = env_gauge_fix(env_conv1, env_conv2_iter)
     @test calc_elementwise_convergence(env_conv1, env_fix_iter) ≈ 0 atol = atol
 
     env_conv2_full, info_full = ctmrg_iteration(n, env_conv1, ctm_alg_full)
-    env_fix_full, signs_full = gauge_fix(env_conv1, env_conv2_full)
+    env_fix_full, signs_full = env_gauge_fix(env_conv1, env_conv2_full)
     @test calc_elementwise_convergence(env_conv1, env_fix_full) ≈ 0 atol = atol
 
     # fix gauge of SVD

test/ctmrg/gaugefix.jl

@@ -3,7 +3,7 @@ using Random
 using PEPSKit
 using TensorKit
 
-using PEPSKit: ctmrg_iteration, gauge_fix, calc_elementwise_convergence
+using PEPSKit: ctmrg_iteration, env_gauge_fix, calc_elementwise_convergence
 
 spacetypes = [ComplexSpace, Z2Space]
 scalartypes = [Float64, ComplexF64]
@@ -48,6 +48,6 @@ end
     env_pre
     env, = leading_boundary(env_pre, psi, alg)
     env′, = ctmrg_iteration(n, env, alg)
-    env_fixed, = gauge_fix(env, env′)
+    env_fixed, = env_gauge_fix(env, env′)
     @test calc_elementwise_convergence(env, env_fixed) ≈ 0 atol = atol
 end

test/ctmrg/jacobian_real_linear.jl

@@ -4,7 +4,7 @@ using Accessors
 using Zygote
 using TensorKit, KrylovKit, PEPSKit
 using PEPSKit:
-    ctmrg_iteration, gauge_fix, fix_relative_phases, fix_global_phases, _fix_svd_algorithm
+    ctmrg_iteration, env_gauge_fix, fix_relative_phases, fix_global_phases, _fix_svd_algorithm
 
 algs = [
     (:fixed, SimultaneousCTMRG(; projector_alg = :halfinfinite)),
@@ -25,7 +25,7 @@ Dbond, χenv = 2, 16
     # follow code of _rrule
     if iterscheme == :fixed
         env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, ctm_alg)
-        env_fixed, signs = gauge_fix(env, env_conv)
+        env_fixed, signs = env_gauge_fix(env, env_conv)
         svd_alg_fixed = _fix_svd_algorithm(ctm_alg.projector_alg.svd_alg, signs, info)
         alg_fixed = @set ctm_alg.projector_alg.svd_alg = svd_alg_fixed
         alg_fixed = @set alg_fixed.projector_alg.trunc = notrunc()
@@ -36,7 +36,7 @@ Dbond, χenv = 2, 16
         end
     elseif iterscheme == :diffgauge
         _, env_vjp = pullback(state, env) do A, x
-            return gauge_fix(x, ctmrg_iteration(InfiniteSquareNetwork(A), x, ctm_alg)[1])[1]
+            return env_gauge_fix(x, ctmrg_iteration(InfiniteSquareNetwork(A), x, ctm_alg)[1])[1]
         end
     end
 

test/ctmrg/unitcell.jl

@@ -1,7 +1,7 @@
 using Test
 using Random
 using PEPSKit
-using PEPSKit: _prev, _next, ctmrg_iteration, gauge_fix, _fix_svd_algorithm
+using PEPSKit: _prev, _next, ctmrg_iteration, env_gauge_fix, _fix_svd_algorithm
 using TensorKit
 
 # settings
@@ -39,7 +39,7 @@ function test_unitcell(
     @test expectation_value(peps, random_op, env′) isa Number
 
     # test if gauge fixing routines run through
-    _, signs = gauge_fix(env′, env″)
+    _, signs = env_gauge_fix(env′, env″)
     @test signs isa Array
     return if ctm_alg isa SimultaneousCTMRG # also test :fixed mode gauge fixing for simultaneous CTMRG
         svd_alg_fixed_full = _fix_svd_algorithm(SVDAdjoint(; fwd_alg = (; alg = :sdd)), signs, info)

test/timeevol/j1j2_finiteT.jl

@@ -47,6 +47,11 @@ energy = expectation_value(pepo, ham, env) / (Nr * Nc)
 @info "β = $(dt * nstep) / 2: tr(ρH) = $(energy)"
 @test energy ≈ bm[1] atol = 5.0e-3
 
+# test BP gauge fixing for purified iPEPO
+bp_alg = BeliefPropagation(; maxiter = 100, tol = 1.0e-9)
+bp_env, = leading_boundary(BPEnv(ones, Float64, pepo), pepo, bp_alg)
+pepo, = gauge_fix(pepo, BPGauge(), bp_env)
+
 env = converge_env(InfinitePEPS(pepo), 16)
 energy = expectation_value(pepo, ham, pepo, env) / (Nr * Nc)
 @info "β = $(dt * nstep): ⟨ρ|H|ρ⟩ = $(energy)"