Formalize gauge transformation for iPEPS, iPEPO and CTMRGEnv

ext/PEPSKitAdaptExt.jl

@@ -35,4 +35,18 @@ function Adapt.adapt_structure(to, x::PEPSKit.C4vCTMRGEnv{C, T}) where {C, T}
     return C4vCTMRGEnv{eltype(C′), eltype(T′)}(C′, T′)
 end
 
+function Adapt.adapt_structure(to, gauge::VirtualGaugeTransform)
+    adapt_tensor = Base.Fix1(adapt, to)
+    return VirtualGaugeTransform(
+        map(adapt_tensor, gauge.matrices), map(adapt_tensor, gauge.inverses)
+    )
+end
+
+function Adapt.adapt_structure(to, gauge::CTMRGEnvGaugeTransform)
+    adapt_tensor = Base.Fix1(adapt, to)
+    return CTMRGEnvGaugeTransform(
+        map(adapt_tensor, gauge.matrices), map(adapt_tensor, gauge.inverses)
+    )
+end
+
 end

src/PEPSKit.jl

@@ -75,6 +75,7 @@ include("utility/initialization.jl")
 include("networks/tensors.jl")
 include("networks/local_sandwich.jl")
 include("networks/infinitesquarenetwork.jl")
+include("networks/gaugetrans.jl")
 
 include("states/infinitepeps.jl")
 include("states/infinitepartitionfunction.jl")
@@ -108,7 +109,6 @@ include("algorithms/contractions/ctmrg/fullinf_env.jl")
 include("algorithms/contractions/ctmrg/renormalize_corner.jl")
 include("algorithms/contractions/ctmrg/renormalize_edge.jl")
 include("algorithms/contractions/ctmrg/network_value.jl")
-include("algorithms/contractions/ctmrg/gaugefix.jl")
 include("algorithms/contractions/ctmrg/characteristic_equations.jl")
 
 include("algorithms/contractions/absorb.jl")
@@ -142,6 +142,7 @@ include("algorithms/ctmrg/sequential.jl")
 include("algorithms/ctmrg/gaugefix.jl")
 include("algorithms/ctmrg/c4v.jl")
 include("algorithms/ctmrg/initialization.jl")
+include("algorithms/ctmrg/gaugetrans.jl")
 
 include("algorithms/truncation/truncationschemes.jl")
 include("algorithms/truncation/fullenv_truncation.jl")
@@ -212,7 +213,8 @@ export InfinitePEPO, InfiniteTransferPEPO
 
 export BPEnv, BeliefPropagation
 export BPGauge, SUGauge
-export gauge_fix
+export gauge_fix, gauge_transform
+export VirtualGaugeTransform, CTMRGEnvGaugeTransform
 
 export initialize_mps, initializePEPS
 export ReflectDepth, ReflectWidth, Rotate, RotateReflect

src/algorithms/bp/gaugefix.jl

@@ -29,27 +29,26 @@ here. The available algorithms are:
 end
 
 """
-    gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::BPGauge, env::BPEnv)
+    gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::BPGauge, env::BPEnv) -> psi′, gauge
 
 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)
     bipartite = _is_bipartite(psi) && _is_bipartite(env)
-    psi′ = copy(psi)
     XXinv = map(eachcoordinate(psi, 1:2)) do I
-        _, X, Xinv = _bp_gauge_fix!(CartesianIndex(I), psi′, env, alg)
-        return X, Xinv
+        return _bp_gauge(CartesianIndex(I), env, alg)
     end
+    gauge = VirtualGaugeTransform(XXinv)
+    psi′ = gauge_transform(psi, gauge)
     if bipartite
-        # copy 1st column to 2nd column to eliminate differences
-        # caused by order of applying gauge transformations
+        # Enforce exact equality between equivalent sites.
         for r in 1:2
             psi′[r + 1, 2] = copy(psi′[r, 1])
         end
     end
-    return psi′, XXinv
+    return psi′, gauge
 end
 function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv)
     # convert to iPEPS, and store physical leg fusers
@@ -59,15 +58,19 @@ function gauge_fix(psi::InfinitePEPO, alg::BPGauge, env::BPEnv)
         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)
+    psi′ = first.(psi_Fs)
+    Fs = last.(psi_Fs)
+    psi′, gauge = gauge_fix(InfinitePEPS(psi′), alg, env)
     # convert back to iPEPO
     psi′ = map(psi′.A, Fs) do t, F
         return F' * t
     end
     psi′ = reshape(psi′, (Nr, Nc, 1))
-    return InfinitePEPO(psi′), XXinv
+    gauge_dims = (2, Nr, Nc, 1)
+    gauge = VirtualGaugeTransform(
+        reshape(gauge.matrices, gauge_dims), reshape(gauge.inverses, gauge_dims)
+    )
+    return InfinitePEPO(psi′), gauge
 end
 
 function _sqrt_bp_messages(I::CartesianIndex{3}, env::BPEnv)
@@ -81,7 +84,7 @@ function _sqrt_bp_messages(I::CartesianIndex{3}, env::BPEnv)
 end
 
 """
-    _bp_gauge_fix!(I, psi::InfinitePEPS, env::BPEnv) -> psi, X, X⁻¹
+    _bp_gauge(I, env::BPEnv, alg::BPGauge) -> 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,
@@ -94,13 +97,10 @@ along the canonical direction of the PEPS arrows (`SOUTH ← NORTH` or `WEST ←
       = 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.
+These matrices update the gauge of `psi` using the [`VirtualGaugeTransform`](@ref) convention.
+`X` is attached to the `NORTH`/`EAST` leg at the indexed site, and `X⁻¹` to the neighboring `SOUTH`/`WEST` leg.
 """
-function _bp_gauge_fix!(I::CartesianIndex{3}, psi::InfinitePEPS, env::BPEnv, alg::BPGauge)
-    dir, row, col = Tuple(I)
-    @assert dir == NORTH || dir == EAST
-
+function _bp_gauge(I::CartesianIndex{3}, env::BPEnv, alg::BPGauge)
     sqrtM12, isqrtM12, sqrtM21, isqrtM21 = _sqrt_bp_messages(I, env)
     U, Λ, Vᴴ = svd_compact!(sqrtM12 * sqrtM21; alg = alg.svd_alg)
     sqrtΛ = sdiag_pow(Λ, 1 / 2)
@@ -109,14 +109,7 @@ function _bp_gauge_fix!(I::CartesianIndex{3}, psi::InfinitePEPS, env::BPEnv, alg
     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[row - 1, col] = absorb_south_message(psi[row - 1, col], invX)
-    elseif dir == EAST
-        psi[row, col] = absorb_east_message(psi[row, col], X)
-        psi[row, col + 1] = absorb_west_message(psi[row, col + 1], invX)
-    end
-    return psi, X, invX
+    return X, invX
 end
 
 """

src/algorithms/contractions/bp_messages.jl

@@ -38,12 +38,3 @@ function contract_west_message(
             M_south[DSb; DSt] * M_west[DWb; DWt] * M_north[DNt; DNb]
     end
 end
-
-absorb_north_message(A::PEPSTensor, M::PEPSMessage) =
-    @tensor A′[d; N' E S W] := A[d; N E S W] * M[N; N']
-absorb_east_message(A::PEPSTensor, M::PEPSMessage) =
-    @tensor A′[d; N E' S W] := A[d; N E S W] * M[E; E']
-absorb_south_message(A::PEPSTensor, M::PEPSMessage) =
-    @tensor A′[d; N E S' W] := A[d; N E S W] * M[S'; S]
-absorb_west_message(A::PEPSTensor, M::PEPSMessage) =
-    @tensor A′[d; N E S W'] := A[d; N E S W] * M[W'; W]