Add full update bond environment and gauge fixing (#203)

* Add full update bond environment and gauge fixing

* There's still room to improve gauge fixing by avoiding `inv` of `R`, `L`

Author: Yue-Zhengyuan

src/PEPSKit.jl

@@ -45,7 +45,9 @@ include("algorithms/contractions/ctmrg_contractions.jl")
 include("algorithms/contractions/localoperator.jl")
 include("algorithms/contractions/vumps_contractions.jl")
 include("algorithms/contractions/bondenv/benv_tools.jl")
+include("algorithms/contractions/bondenv/gaugefix.jl")
 include("algorithms/contractions/bondenv/als_solve.jl")
+include("algorithms/contractions/bondenv/benv_ctm.jl")
 
 include("algorithms/ctmrg/sparse_environments.jl")
 include("algorithms/ctmrg/ctmrg.jl")

src/algorithms/contractions/bondenv/benv_ctm.jl

@@ -0,0 +1,70 @@
+"""
+Construct the environment (norm) tensor
+```
+    C1---T1---------T1---C2   r-1
+    |    ‖          ‖    |
+    T4===XX==     ==YY===T2    r
+    |    ‖          ‖    |
+    C4---T3---------T3---C3   r+1
+    c-1  c         c+1  c+2
+```
+where `XX = X' X` and `YY = Y' Y` (stacked together).
+
+Axis order: `[DX1 DY1; DX0 DY0]`, as in
+```
+    ┌---------------------┐
+    | ┌----┐              |
+    └-|    |---DX0  DY0---┘
+      |benv|
+    ┌-|    |---DX1  DY1---┐
+    | └----┘              |
+    └---------------------┘
+```
+"""
+function bondenv_fu(row::Int, col::Int, X::PEPSOrth, Y::PEPSOrth, env::CTMRGEnv)
+    Nr, Nc = size(env.corners)[[2, 3]]
+    cm1 = _prev(col, Nc)
+    cp1 = _next(col, Nc)
+    cp2 = _next(cp1, Nc)
+    rm1 = _prev(row, Nr)
+    rp1 = _next(row, Nr)
+    c1 = env.corners[1, rm1, cm1]
+    c2 = env.corners[2, rm1, cp2]
+    c3 = env.corners[3, rp1, cp2]
+    c4 = env.corners[4, rp1, cm1]
+    t1X, t1Y = env.edges[1, rm1, col], env.edges[1, rm1, cp1]
+    t2 = env.edges[2, row, cp2]
+    t3X, t3Y = env.edges[3, rp1, col], env.edges[3, rp1, cp1]
+    t4 = env.edges[4, row, cm1]
+    #= index labels
+
+    C1--χ4--T1X---------χ6---------T1Y--χ8---C2     r-1
+    |        ‖                      ‖        |
+    χ2      DNX                    DNY      χ10
+    |        ‖                      ‖        |
+    T4==DWX==XX===DX==       ==DY===YY==DEY==T2     r
+    |        ‖                      ‖        |
+    χ1      DSX                    DSY       χ9
+    |        ‖                      ‖        |
+    C4--χ3--T3X---------χ5---------T3Y--χ7---C3     r+1
+    c-1      c                      c+1     c+2
+    =#
+    @autoopt @tensor benv[DX1 DY1; DX0 DY0] := (
+        c4[χ3 χ1] *
+        t4[χ1 DWX0 DWX1 χ2] *
+        c1[χ2 χ4] *
+        t3X[χ5 DSX0 DSX1 χ3] *
+        X[DNX0 DX0 DSX0 DWX0] *
+        conj(X[DNX1 DX1 DSX1 DWX1]) *
+        t1X[χ4 DNX0 DNX1 χ6] *
+        c3[χ9 χ7] *
+        t2[χ10 DEY0 DEY1 χ9] *
+        c2[χ8 χ10] *
+        t3Y[χ7 DSY0 DSY1 χ5] *
+        Y[DNY0 DEY0 DSY0 DY0] *
+        conj(Y[DNY1 DEY1 DSY1 DY1]) *
+        t1Y[χ6 DNY0 DNY1 χ8]
+    )
+    normalize!(benv, Inf)
+    return benv
+end

src/algorithms/contractions/bondenv/gaugefix.jl

@@ -0,0 +1,120 @@
+"""
+Replace bond environment `benv` by its positive approximant `Z† Z`
+(returns the "half environment" `Z`)
+```
+    ┌-----------------┐     ┌---------------┐
+    | ┌----┐          |     |               |
+    └-|    |-- 3  4 --┘     └-- Z -- 3  4 --┘
+      |benv|            =       ↓
+    ┌-|    |-- 1  2 --┐     ┌-- Z†-- 1  2 --┐
+    | └----┘          |     |               |
+    └-----------------┘     └---------------┘
+```
+"""
+function positive_approx(benv::BondEnv)
+    # eigen-decomposition: benv = U D U'
+    D, U = eigh((benv + benv') / 2)
+    # determine if `env` is (mostly) positive or negative
+    sgn = sign(sum(D.data))
+    # When optimizing the truncation of a bond, 
+    # its environment can always be multiplied by a number.
+    # If `benv` is negative (e.g. obtained approximately from CTMRG), 
+    # we can multiply it by (-1).
+    data = D.data
+    @inbounds for i in eachindex(data)
+        d = (sgn == -1) ? -data[i] : data[i]
+        data[i] = (d > 0) ? sqrt(d) : zero(d)
+    end
+    Z = D * U'
+    return Z
+end
+
+"""
+Use QR decomposition to fix gauge of the half bond environment `Z`.
+The reduced bond tensors `a`, `b` and `Z` are arranged as
+```
+    ┌---------------┐
+    |               |
+    └---Z---a---b---┘
+        |   ↓   ↓
+        ↓
+```
+Reference: 
+- Physical Review B 90, 064425 (2014)
+- Physical Review B 92, 035142 (2015)
+"""
+function fixgauge_benv(
+    Z::AbstractTensorMap{T,S,1,2},
+    a::AbstractTensorMap{T,S,1,2},
+    b::AbstractTensorMap{T,S,2,1},
+) where {T<:Number,S<:ElementarySpace}
+    @assert !isdual(space(Z, 1))
+    @assert !isdual(space(a, 2))
+    @assert !isdual(space(b, 2))
+    #= QR/LQ decomposition of Z 
+
+        3 - Z - 2   =   2 - L - 1   3 - QL - 1
+            ↓                           ↓
+            1                           2
+
+                    =   1 - QR - 3  1 - R - 2
+                            ↓
+                            2
+    =#
+    QL, L = leftorth(Z, ((2, 1), (3,)))
+    QR, R = leftorth(Z, ((3, 1), (2,)))
+    @debug "cond(L) = $(LinearAlgebra.cond(L)); cond(R) = $(LinearAlgebra.cond(R))"
+    # TODO: find a better way to fix gauge that avoids `inv`
+    Linv, Rinv = inv(L), inv(R)
+    #= fix gauge of Z, a, b
+        ┌---------------------------------------┐
+        |                                       |
+        └---Z---Rinv)---(R--a)--(b--L)---(Linv--┘
+            |               ↓    ↓
+            ↓
+
+        -1 - R - 1 - a - -3   -1 - b - 1 - L - -3
+                     ↓             ↓        
+                    -2            -2
+
+        ┌-----------------------------------------┐
+        |                                         |
+        └---Z-- 1 --Rinv-- -2      -3 --Linv-- 2 -┘
+            ↓
+            -1
+    =#
+    @plansor a[-1; -2 -3] := R[-1; 1] * a[1; -2 -3]
+    @plansor b[-1 -2; -3] := b[-1 -2; 1] * L[-3; 1]
+    @plansor Z[-1; -2 -3] := Z[-1; 1 2] * Rinv[1; -2] * Linv[2; -3]
+    (isdual(space(R, 1)) == isdual(space(R, 2))) && twist!(a, 1)
+    (isdual(space(L, 1)) == isdual(space(L, 2))) && twist!(b, 3)
+    return Z, a, b, (Linv, Rinv)
+end
+
+"""
+When the (half) bond environment `Z` consists of two `PEPSOrth` tensors `X`, `Y` as
+```
+    ┌---------------┐   ┌-------------------┐
+    |               | = |                   | ,
+    └---Z--       --┘   └--Z0---X--    --Y--┘
+        ↓                  ↓
+```
+apply the gauge transformation `Linv`, `Rinv` for `Z` to `X`, `Y`:
+```
+        -1                                     -1
+         |                                      |
+    -4 - X - 1 - Rinv - -2      -4 - Linv - 1 - Y - -2
+         |                                      |
+        -3                                     -3
+```
+"""
+function _fixgauge_benvXY(
+    X::PEPSOrth{T,S},
+    Y::PEPSOrth{T,S},
+    Linv::AbstractTensorMap{T,S,1,1},
+    Rinv::AbstractTensorMap{T,S,1,1},
+) where {T<:Number,S<:ElementarySpace}
+    @plansor X[-1 -2 -3 -4] := X[-1 1 -3 -4] * Rinv[1; -2]
+    @plansor Y[-1 -2 -3 -4] := Y[-1 -2 -3 1] * Linv[1; -4]
+    return X, Y
+end

test/bondenv/benv_fu.jl

@@ -0,0 +1,48 @@
+using Test
+using TensorKit
+using PEPSKit
+using LinearAlgebra
+using KrylovKit
+using Random
+
+Random.seed!(100)
+Nr, Nc = 2, 2
+# create Hubbard iPEPS using simple update
+function get_hubbard_state(t::Float64=1.0, U::Float64=8.0)
+    H = hubbard_model(ComplexF64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t, U, mu=U / 2)
+    Vphy = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2, (1, 1//2) => 1, (1, -1//2) => 1)
+    wpeps = InfiniteWeightPEPS(rand, ComplexF64, Vphy, Vphy; unitcell=(Nr, Nc))
+    alg = SimpleUpdate(1e-2, 1e-8, 10000, truncerr(1e-10) & truncdim(4))
+    wpeps, = simpleupdate(wpeps, H, alg; bipartite=false, check_interval=2000)
+    peps = InfinitePEPS(wpeps)
+    normalize!.(peps.A, Inf)
+    return peps
+end
+
+peps = get_hubbard_state()
+# calculate CTMRG environment
+Envspace = Vect[FermionParity ⊠ U1Irrep](
+    (0, 0) => 4, (1, 1//2) => 1, (1, -1//2) => 1, (0, 1) => 1, (0, -1) => 1
+)
+ctm_alg = SequentialCTMRG(; tol=1e-10, verbosity=2, trscheme=truncerr(1e-10) & truncdim(8))
+env, = leading_boundary(CTMRGEnv(rand, ComplexF64, peps, Envspace), peps, ctm_alg)
+for row in 1:Nr, col in 1:Nc
+    cp1 = PEPSKit._next(col, Nc)
+    A, B = peps.A[row, col], peps.A[row, cp1]
+    X, a, b, Y = PEPSKit._qr_bond(A, B)
+    benv = PEPSKit.bondenv_fu(row, col, X, Y, env)
+    @assert [isdual(space(benv, ax)) for ax in 1:numind(benv)] == [0, 0, 1, 1]
+    Z = PEPSKit.positive_approx(benv)
+    # verify that gauge fixing can greatly reduce 
+    # condition number for physical state bond envs
+    cond1 = cond(Z' * Z)
+    Z2, a2, b2, (Linv, Rinv) = PEPSKit.fixgauge_benv(Z, a, b)
+    benv2 = Z2' * Z2
+    cond2 = cond(benv2)
+    @test 1 <= cond2 < cond1
+    @info "benv cond number: (gauge-fixed) $(cond2) ≤ $(cond1) (initial)"
+    # verify gauge fixing is done correctly
+    @tensor half[:] := Z[-1; 1 3] * a[1; -2 2] * b[2 -3; 3]
+    @tensor half2[:] := Z2[-1; 1 3] * a2[1; -2 2] * b2[2 -3; 3]
+    @test half ≈ half2
+end

test/bondenv/benv_gaugefix.jl

@@ -0,0 +1,43 @@
+using Test
+using TensorKit
+using PEPSKit
+using LinearAlgebra
+using KrylovKit
+using Random
+
+Vphy = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 1, (1, 1) => 1, (1, -1) => 2)
+Vin = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 1, (1, 1) => 3, (1, -1) => 2)
+V = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 1, (1, 1) => 2, (1, -1) => 3)
+Vs = (V, V')
+for V1 in Vs, V2 in Vs, V3 in Vs
+    #=
+        ┌---┬---------------┬---┐
+        |   |               |   |    ┌--------------┐
+        ├---X--- -2   -3 ---Y---┤  = |              |
+        |   |               |   |    └--Z-- -2  -3 -┘
+        └---┴-------Z0------┴---┘       ↓
+                    ↓                   -1
+                    -1
+    =#
+    X = rand(ComplexF64, Vin ⊗ V1' ⊗ Vin' ⊗ Vin)
+    Y = rand(ComplexF64, Vin ⊗ Vin ⊗ Vin' ⊗ V3)
+    Z0 = randn(ComplexF64, Vphy ← Vin ⊗ Vin' ⊗ Vin ⊗ Vin ⊗ Vin ⊗ Vin')
+    @tensor Z[p; Xe Yw] := Z0[p; Xn Xs Xw Yn Ye Ys] * X[Xn Xe Xs Xw] * Y[Yn Ye Ys Yw]
+    #= 
+        ┌---------------------------┐
+        |                           |
+        └---Z-- 1 --a-- 2 --b-- 3 --┘
+            ↓       ↓       ↓
+            -1      -2      -3
+    =#
+    a = randn(ComplexF64, V1 ← Vphy' ⊗ V2)
+    b = randn(ComplexF64, V2 ⊗ Vphy ← V3)
+    @tensor half[:] := Z[-1; 1 3] * a[1; -2 2] * b[2 -3; 3]
+    Z2, a2, b2, (Linv, Rinv) = PEPSKit.fixgauge_benv(Z, a, b)
+    @tensor half2[:] := Z2[-1; 1 3] * a2[1; -2 2] * b2[2 -3; 3]
+    @test half ≈ half2
+    # test gauge transformation of X, Y
+    X2, Y2 = PEPSKit._fixgauge_benvXY(X, Y, Linv, Rinv)
+    @tensor Z2_[p; Xe Yw] := Z0[p; Xn Xs Xw Yn Ye Ys] * X2[Xn Xe Xs Xw] * Y2[Yn Ye Ys Yw]
+    @test Z2 ≈ Z2_
+end

test/runtests.jl

@@ -50,6 +50,12 @@ end
         @time @safetestset "Iterative optimization after truncation" begin
             include("bondenv/bond_truncate.jl")
         end
+        @time @safetestset "Gauge fixing" begin
+            include("bondenv/benv_gaugefix.jl")
+        end
+        @time @safetestset "Full update bond environment" begin
+            include("bondenv/benv_fu.jl")
+        end
     end
     if GROUP == "ALL" || GROUP == "TIMEEVOL"
         @time @safetestset "Cluster truncation with projectors" begin