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/optimization/fixed_point_differentiation.jl

@@ -221,7 +221,7 @@ 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]
         end

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,6 @@ using PEPSKit
 using PEPSKit:
     FixedSVD,
     ctmrg_iteration,
-    gauge_fix,
     fix_relative_phases,
     fix_global_phases,
     calc_elementwise_convergence,

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, calc_elementwise_convergence
 
 spacetypes = [ComplexSpace, Z2Space]
 scalartypes = [Float64, ComplexF64]

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, fix_relative_phases, fix_global_phases, _fix_svd_algorithm
 
 algs = [
     (:fixed, SimultaneousCTMRG(; projector_alg = :halfinfinite)),

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, _fix_svd_algorithm
 using TensorKit
 
 # settings

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)"