Belief propagation gauge fixing

.github/workflows/Tests.yml

@@ -26,6 +26,7 @@ jobs:
           - 'lts' # minimal supported version
           - '1'
         group:
+          - bp
           - types
           - ctmrg
           - boundarymps

src/PEPSKit.jl

@@ -51,11 +51,13 @@ include("operators/models.jl")
 include("environments/ctmrg_environments.jl")
 include("environments/vumps_environments.jl")
 include("environments/suweight.jl")
+include("environments/bp_environments.jl")
 
 include("algorithms/contractions/ctmrg_contractions.jl")
 include("algorithms/contractions/transfer.jl")
 include("algorithms/contractions/localoperator.jl")
 include("algorithms/contractions/vumps_contractions.jl")
+include("algorithms/contractions/bp_contractions.jl")
 include("algorithms/contractions/bondenv/benv_tools.jl")
 include("algorithms/contractions/bondenv/gaugefix.jl")
 include("algorithms/contractions/bondenv/als_solve.jl")
@@ -78,6 +80,10 @@ include("algorithms/time_evolution/evoltools.jl")
 include("algorithms/time_evolution/time_evolve.jl")
 include("algorithms/time_evolution/simpleupdate.jl")
 include("algorithms/time_evolution/simpleupdate3site.jl")
+include("algorithms/time_evolution/gaugefix_su.jl")
+
+include("algorithms/bp/beliefpropagation.jl")
+include("algorithms/bp/gaugefix.jl")
 
 include("algorithms/transfermatrix.jl")
 include("algorithms/toolbox.jl")
@@ -114,6 +120,9 @@ export InfinitePartitionFunction
 export InfinitePEPS, InfiniteTransferPEPS
 export SUWeight
 export InfinitePEPO, InfiniteTransferPEPO
+
+export BPEnv, BeliefPropagation, BPGauge
+
 export initialize_mps, initializePEPS
 export ReflectDepth, ReflectWidth, Rotate, RotateReflect
 export symmetrize!, symmetrize_retract_and_finalize!

src/algorithms/bp/beliefpropagation.jl

@@ -0,0 +1,111 @@
+"""
+    struct BeliefPropagation
+
+Algorithm for computing the belief propagation fixed point messages.
+
+## Fields
+
+$(TYPEDFIELDS)
+"""
+@kwdef struct BeliefPropagation
+    "Stopping criterion for the BP iterations in relative trace norm difference"
+    tol::Float64 = 1.0e-6
+
+    "Minimal number of BP iterations"
+    miniter::Int = 2
+
+    "Maximal number of BP iterations"
+    maxiter::Int = 50
+
+    "Toggle for projecting messages onto the hermitian subspace immediately after update through BP equation"
+    project_hermitian::Bool = true
+
+    "Output verbosity level"
+    verbosity::Int = 2
+end
+
+"""
+    leading_boundary(env₀::BPEnv, network, alg::BeliefPropagation)
+
+Contract `network` in the BP approximation and return the corresponding messages.
+"""
+function leading_boundary(env₀::BPEnv, network::InfiniteSquareNetwork, alg::BeliefPropagation)
+    return LoggingExtras.withlevel(; alg.verbosity) do
+        env = deepcopy(env₀)
+        log = MPSKit.IterLog("BP")
+        ϵ = Inf
+        @infov 1 loginit!(log, ϵ)
+        for iter in 1:(alg.maxiter)
+            env′ = bp_iteration(network, env, alg)
+            ϵ = oftype(ϵ, tr_distance(env, env′))
+            env = env′
+
+            if ϵ <= alg.tol && iter >= alg.miniter
+                @infov 2 logfinish!(log, iter, ϵ)
+                break
+            end
+            if iter ≥ alg.maxiter
+                @warnv 1 logcancel!(log, iter, ϵ)
+            else
+                @infov 3 logiter!(log, iter, ϵ)
+            end
+        end
+
+        return env, ϵ
+    end
+end
+function leading_boundary(env₀::BPEnv, state, args...)
+    return leading_boundary(env₀, InfiniteSquareNetwork(state), args...)
+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
+    end
+    return BPEnv(messages)
+end
+
+"""
+Update the BP message in `env.messages[I]`.
+"""
+function update_message(I::CartesianIndex{3}, network::InfiniteSquareNetwork, env::BPEnv)
+    dir, row, col = Tuple(I)
+    (1 <= dir <= 4) || throw(ArgumentError("Invalid direction $dir"))
+
+    A = network[row, col]
+    dir == SOUTH || (M_north = env[NORTH, _prev(row, end), col])
+    dir == WEST || (M_east = env[EAST, row, _next(col, end)])
+    dir == NORTH || (M_south = env[SOUTH, _next(row, end), col])
+    dir == EAST || (M_west = env[WEST, row, _prev(col, end)])
+
+    return if dir == NORTH
+        contract_north_message(A, M_west, M_north, M_east)
+    elseif dir == EAST
+        contract_east_message(A, M_north, M_east, M_south)
+    elseif dir == SOUTH
+        contract_south_message(A, M_east, M_south, M_west)
+    else # dir == WEST
+        contract_west_message(A, M_south, M_west, M_north)
+    end
+end
+
+function tr_distance(A::BPEnv, B::BPEnv)
+    return sum(zip(A.messages, B.messages)) do (a, b)
+        return trnorm(add(a, b, -inv(tr(b)), inv(tr(a))))
+    end / length(A.messages)
+end
+
+function trnorm(M::AbstractTensorMap, p::Real = 1)
+    return TensorKit._norm(svdvals(M), p, zero(real(scalartype(M))))
+end
+function trnorm!(M::AbstractTensorMap, p::Real = 1)
+    return TensorKit._norm(svdvals!(M), p, zero(real(scalartype(M))))
+end
+
+project_hermitian!!(t) = add(t, t', 1 / 2, 1 / 2)

src/algorithms/bp/gaugefix.jl

@@ -0,0 +1,127 @@
+"""
+    struct BPGauge
+
+Algorithm for gauging PEPS with belief propagation fixed point messages.
+"""
+@kwdef struct BPGauge
+    # TODO: add options
+end
+
+"""
+$(SIGNATURES)
+
+Fix the gauge of `psi` using fixed point environment `env` of belief propagation.
+"""
+function gauge_fix(psi::InfinitePEPS, alg::BPGauge, env::BPEnv)
+    psi′ = copy(psi)
+    XXinv = map(eachcoordinate(psi, 1:2)) do I
+        _, X, Xinv = _bp_gauge_fix!(CartesianIndex(I), psi′, env)
+        return X, Xinv
+    end
+    return psi′, XXinv
+end
+
+function _sqrt_bp_messages(I::CartesianIndex{3}, env::BPEnv)
+    dir, row, col = Tuple(I)
+    @assert dir == NORTH || dir == EAST
+    M12 = env[dir, dir == NORTH ? _prev(row, end) : row, dir == EAST ? _next(col, end) : col]
+    sqrtM12, isqrtM12 = sqrt_invsqrt(twist(M12, 1))
+    M21 = env[dir + 2, row, col]
+    sqrtM21, isqrtM21 = sqrt_invsqrt(M21)
+    return sqrtM12, isqrtM12, sqrtM21, isqrtM21
+end
+
+"""
+    _bp_gauge_fix!(I, psi::InfinitePEPS, env::BPEnv) -> psi, X, X⁻¹
+
+For the bond at direction `I[1]` (which can be `NORTH` or `EAST`)
+from site `I[2], I[3]`, we identify the following gauge matrices,
+along the canonical direction of the PEPS arrows (`SOUTH ← NORTH` or `WEST ← EAST`):
+
+```math
+    I = √M₁₂⁻¹ √M₁₂ √M₂₁ √M₂₁⁻¹
+      = √M₁₂⁻¹ (U Λ Vᴴ) √M₂₁⁻¹
+      = (√M₁₂⁻¹ U √Λ) (√Λ Vᴴ √M₂₁⁻¹)
+      = X X⁻¹
+```
+
+Which are then used to update the gauge of `psi`. Thus, by convention `X` is attached to the `SOUTH`/`WEST` directions
+and `X⁻¹` is attached to the `NORTH`/`EAST` directions.
+"""
+function _bp_gauge_fix!(I::CartesianIndex{3}, psi::InfinitePEPS, env::BPEnv)
+    dir, row, col = Tuple(I)
+    @assert dir == NORTH || dir == EAST
+
+    sqrtM12, isqrtM12, sqrtM21, isqrtM21 = _sqrt_bp_messages(I, env)
+    U, Λ, Vᴴ = svd_compact!(sqrtM12 * sqrtM21)
+    sqrtΛ = sdiag_pow(Λ, 1 / 2)
+    X = isqrtM12 * U * sqrtΛ
+    invX = sqrtΛ * Vᴴ * isqrtM21
+    if isdual(space(sqrtM12, 1))
+        X, invX = twist(flip(X, 2), 1), flip(invX, 1)
+    end
+    if dir == NORTH
+        psi[row, col] = absorb_north_message(psi[row, col], X)
+        psi[_prev(row, end), col] = absorb_south_message(psi[_prev(row, end), col], invX)
+    elseif dir == EAST
+        psi[row, col] = absorb_east_message(psi[row, col], X)
+        psi[row, _next(col, end)] = absorb_west_message(psi[row, _next(col, end)], invX)
+    end
+    return psi, X, invX
+end
+
+"""
+    SUWeight(env::BPEnv)
+
+Construct `SUWeight` from belief propagation fixed point environment `env`.
+"""
+function SUWeight(env::BPEnv)
+    wts = map(Iterators.product(1:2, axes(env, 2), axes(env, 3))) do (dir′, row, col)
+        I = CartesianIndex(mod1(dir′ + 1, 2), row, col)
+        sqrtM12, _, sqrtM21, _ = _sqrt_bp_messages(I, env)
+        Λ = svd_vals!(sqrtM12 * sqrtM21)
+        return isdual(space(sqrtM12, 1)) ? _fliptwist_s(Λ) : Λ
+    end
+    return SUWeight(wts)
+end
+
+"""
+    BPEnv(wts::SUWeight)
+
+Convert fixed point weights `wts` of trivial simple update
+to a belief propagation environment.
+"""
+function BPEnv(wts::SUWeight)
+    messages = map(Iterators.product(1:4, axes(wts, 2), axes(wts, 3))) do (d, r, c)
+        wt = if d == NORTH
+            twist(wts[2, _next(r, end), c], 1)
+        elseif d == EAST
+            twist(wts[1, r, _prev(c, end)], 1)
+        elseif d == SOUTH
+            copy(wts[2, r, c])
+        else # WEST
+            copy(wts[1, r, c])
+        end
+        return TensorMap(wt)
+    end
+    return BPEnv(messages)
+end
+
+function sqrt_invsqrt(A::PEPSMessage)
+    if isposdef(A)
+        D, V = eigh_full(A)
+        sqrtA = V * sdiag_pow(D, 1 / 2) * V'
+        isqrtA = V * sdiag_pow(D, -1 / 2) * V'
+    else
+        D, V = eig_full(A)
+        V⁻¹ = inv(V)
+        sqrtA = V * sdiag_pow(D, 1 / 2) * V⁻¹
+        isqrtA = V * sdiag_pow(D, -1 / 2) * V⁻¹
+        if scalartype(A) <: Real
+            # TODO: is this valid?
+            sqrtA = real(sqrtA)
+            isqrtA = real(isqrtA)
+        end
+    end
+    return sqrtA, isqrtA
+end