From 399a6000955707c7fc6655b725374cf003936004 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 1 Nov 2024 21:12:57 +0800
Subject: [PATCH 01/20] remove CTMRG redundant logging

---
 src/algorithms/ctmrg/ctmrg.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl
index 79d3f033..00fc4597 100644
--- a/src/algorithms/ctmrg/ctmrg.jl
+++ b/src/algorithms/ctmrg/ctmrg.jl
@@ -115,10 +115,8 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
         ctmrg_loginit!(log, η, N)
         for iter in 1:(alg.maxiter)
             env, = ctmrg_iter(state, env, alg)  # Grow and renormalize in all 4 directions
-
             η, CS, TS = calc_convergence(env, CS, TS)
             N = norm(state, env)
-            ctmrg_logiter!(log, iter, η, N)
 
             if η ≤ alg.tol && iter ≥ alg.miniter
                 ctmrg_logfinish!(log, iter, η, N)

From e6f8f060a972e62fa47732eccbaf06d23a76a329 Mon Sep 17 00:00:00 2001
From: Paul Brehmer <paul.brehmer@univie.ac.at>
Date: Tue, 5 Nov 2024 16:19:02 +0100
Subject: [PATCH 02/20] Make :sequential act column-wise

---
 .../contractions/ctmrg_contractions.jl        |  8 +-
 src/algorithms/ctmrg/ctmrg.jl                 | 73 ++++++++-----------
 test/heisenberg.jl                            |  6 ++
 3 files changed, 41 insertions(+), 46 deletions(-)

diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl
index e53bf2ac..cb920d76 100644
--- a/src/algorithms/contractions/ctmrg_contractions.jl
+++ b/src/algorithms/contractions/ctmrg_contractions.jl
@@ -530,7 +530,7 @@ Apply bottom projector to southwest corner and south edge.
 function renormalize_bottom_corner((row, col), envs::CTMRGEnv, projectors)
     C_southwest = envs.corners[SOUTHWEST, row, _prev(col, end)]
     E_south = envs.edges[SOUTH, row, col]
-    P_bottom = projectors[1][row, col]
+    P_bottom = projectors[1][row]
     return @autoopt @tensor corner[χ_in; χ_out] :=
         E_south[χ_in D1 D2; χ1] * C_southwest[χ1; χ2] * P_bottom[χ2 D1 D2; χ_out]
 end
@@ -549,7 +549,7 @@ Apply top projector to northwest corner and north edge.
 function renormalize_top_corner((row, col), envs::CTMRGEnv, projectors)
     C_northwest = envs.corners[NORTHWEST, row, _prev(col, end)]
     E_north = envs.edges[NORTH, row, col]
-    P_top = projectors[2][_next(row, end), col]
+    P_top = projectors[2][_next(row, end)]
     return @autoopt @tensor corner[χ_in; χ_out] :=
         P_top[χ_in; χ1 D1 D2] * C_northwest[χ1; χ2] * E_north[χ2 D1 D2; χ_out]
 end
@@ -705,8 +705,8 @@ function renormalize_west_edge(  # For sequential CTMRG scheme
 )
     return renormalize_west_edge(
         envs.edges[WEST, row, _prev(col, end)],
-        projectors[1][row, col],
-        projectors[2][_next(row, end), col],
+        projectors[1][row],
+        projectors[2][_next(row, end)],
         ket[row, col],
         bra[row, col],
     )
diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl
index 04b8651b..27f8745a 100644
--- a/src/algorithms/ctmrg/ctmrg.jl
+++ b/src/algorithms/ctmrg/ctmrg.jl
@@ -141,22 +141,23 @@ and also returns the truncation error.
 """
 function ctmrg_iter(state, envs::CTMRGEnv, alg::SequentialCTMRG)
     ϵ = zero(real(scalartype(state)))
-    for _ in 1:4
-        # left move
-        enlarged_envs = ctmrg_expand(state, envs, alg)
-        projectors, info = ctmrg_projectors(enlarged_envs, envs, alg)
-        envs = ctmrg_renormalize(projectors, state, envs, alg)
-
-        # rotate
+    for _ in 1:4 # rotate
+        for col in 1:size(state, 2) # left move column-wise
+            enlarged_envs = ctmrg_expand(
+                eachcoordinate(envs, [4, 1])[:, :, col], state, envs
+            )
+            projectors, info = ctmrg_projectors(col, enlarged_envs, envs, alg)
+            envs = ctmrg_renormalize(col, projectors, state, envs, alg)
+            ϵ = max(ϵ, info.err)
+        end
         state = rotate_north(state, EAST)
         envs = rotate_north(envs, EAST)
-        ϵ = max(ϵ, info.err)
     end
 
     return envs, (; err=ϵ)
 end
 function ctmrg_iter(state, envs::CTMRGEnv, alg::SimultaneousCTMRG)
-    enlarged_envs = ctmrg_expand(state, envs, alg)
+    enlarged_envs = ctmrg_expand(eachcoordinate(state, 1:4), state, envs)
     projectors, info = ctmrg_projectors(enlarged_envs, envs, alg)
     envs′ = ctmrg_renormalize(enlarged_envs, projectors, state, envs, alg)
     return envs′, info
@@ -184,14 +185,7 @@ There are two modes of expansion: `M = :sequential` and `M = :simultaneous`.
 The first mode expands the environment in one direction at a time, for convenience towards
 the left. The second mode expands the environment in all four directions simultaneously.
 """
-function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG)
-    return ctmrg_expand([4, 1], state, envs)
-end
-function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG)
-    return ctmrg_expand(1:4, state, envs)
-end
-function ctmrg_expand(dirs, state, envs::CTMRGEnv)
-    coordinates = eachcoordinate(state, dirs)
+function ctmrg_expand(coordinates, state, envs::CTMRGEnv)
     return dtmap(idx -> TensorMap(EnlargedCorner(state, envs, idx), idx[1]), coordinates)
 end
 
@@ -206,17 +200,16 @@ Compute the CTMRG projectors based from enlarged environments.
 In the `:simultaneous` mode, the environment SVD is run in parallel.
 """
 function ctmrg_projectors(
-    enlarged_envs, envs::CTMRGEnv{C,E}, alg::SequentialCTMRG
+    col::Int, enlarged_envs, envs::CTMRGEnv{C,E}, alg::SequentialCTMRG
 ) where {C,E}
     projector_alg = alg.projector_alg
     ϵ = zero(real(scalartype(envs)))
 
-    coordinates = eachcoordinate(envs)
+    # SVD half-infinite environment
+    coordinates = eachcoordinate(envs)[:, col]
     projectors = dtmap(coordinates) do (r, c)
-        # SVD half-infinite environment
         r′ = _prev(r, size(envs.corners, 2))
-        QQ = halfinfinite_environment(enlarged_envs[1, r, c], enlarged_envs[2, r′, c])
-
+        QQ = halfinfinite_environment(enlarged_envs[1, r], enlarged_envs[2, r′])
         trscheme = truncation_scheme(projector_alg, envs.edges[WEST, r′, c])
         svd_alg = svd_algorithm(projector_alg, (WEST, r, c))
         U, S, V, ϵ_local = PEPSKit.tsvd!(QQ, svd_alg; trunc=trscheme)
@@ -231,9 +224,8 @@ function ctmrg_projectors(
         end
 
         # Compute projectors
-        return build_projectors(U, S, V, enlarged_envs[1, r, c], enlarged_envs[2, r′, c])
+        return build_projectors(U, S, V, enlarged_envs[1, r], enlarged_envs[2, r′])
     end
-
     return (map(first, projectors), map(last, projectors)), (; err=ϵ)
 end
 function ctmrg_projectors(
@@ -339,32 +331,29 @@ end
 
 Apply projectors to renormalize corners and edges.
 """
-function ctmrg_renormalize(projectors, state, envs, ::SequentialCTMRG)
+function ctmrg_renormalize(col::Int, projectors, state, envs, ::SequentialCTMRG)
     corners = Zygote.Buffer(envs.corners)
     edges = Zygote.Buffer(envs.edges)
 
-    # copy environments that do not participate
-    for dir in (NORTHEAST, SOUTHEAST)
-        for r in axes(envs.corners, 2), c in axes(envs.corners, 3)
-            corners[dir, r, c] = envs.corners[dir, r, c]
-        end
+    for (dir, r, c) in eachcoordinate(state, 1:4)
+        (c == col && dir in [SOUTHWEST, NORTHWEST]) && continue
+        corners[dir, r, c] = envs.corners[dir, r, c]
     end
-    for dir in (NORTH, EAST, SOUTH)
-        for r in axes(envs.corners, 2), c in axes(envs.corners, 3)
-            edges[dir, r, c] = envs.edges[dir, r, c]
-        end
+    for (dir, r, c) in eachcoordinate(state, 1:4)
+        (c == col && dir == WEST) && continue
+        edges[dir, r, c] = envs.edges[dir, r, c]
     end
 
-    # Apply projectors to renormalize corners and edges
-    for (r, c) in eachcoordinate(state)
-        C_southwest = renormalize_bottom_corner((r, c), envs, projectors)
-        corners[SOUTHWEST, r, c] = C_southwest / norm(C_southwest)
+    # Apply projectors to renormalize corners and edge
+    for row in axes(envs.corners, 2)
+        C_southwest = renormalize_bottom_corner((row, col), envs, projectors)
+        corners[SOUTHWEST, row, col] = C_southwest / norm(C_southwest)
 
-        C_northwest = renormalize_top_corner((r, c), envs, projectors)
-        corners[NORTHWEST, r, c] = C_northwest / norm(C_northwest)
+        C_northwest = renormalize_top_corner((row, col), envs, projectors)
+        corners[NORTHWEST, row, col] = C_northwest / norm(C_northwest)
 
-        E_west = renormalize_west_edge((r, c), envs, projectors, state)
-        edges[WEST, r, c] = E_west / norm(E_west)
+        E_west = renormalize_west_edge((row, col), envs, projectors, state)
+        edges[WEST, row, col] = E_west / norm(E_west)
     end
 
     return CTMRGEnv(copy(corners), copy(edges))
diff --git a/test/heisenberg.jl b/test/heisenberg.jl
index 1c05b084..c4941c40 100644
--- a/test/heisenberg.jl
+++ b/test/heisenberg.jl
@@ -12,6 +12,12 @@ ctm_alg = CTMRG()
 opt_alg = PEPSOptimize(;
     boundary_alg=ctm_alg, optimizer=LBFGS(4; gradtol=1e-3, verbosity=2)
 )
+ctm_alg = CTMRG(; ctmrgscheme=:sequential)
+opt_alg = PEPSOptimize(;
+    boundary_alg=ctm_alg,
+    optimizer=LBFGS(4; gradtol=1e-3, verbosity=2),
+    gradient_alg=LinSolver(; iterscheme=:diffgauge),
+)
 
 # initialize states
 Random.seed!(91283219347)

From d9b08bbe3709d0cdbf2e253e94d645af424cd6ec Mon Sep 17 00:00:00 2001
From: Paul Brehmer <paul.brehmer@univie.ac.at>
Date: Tue, 5 Nov 2024 16:27:32 +0100
Subject: [PATCH 03/20] Update docstrings

---
 src/algorithms/ctmrg/ctmrg.jl | 22 ++++++++++++----------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl
index 27f8745a..8e74436f 100644
--- a/src/algorithms/ctmrg/ctmrg.jl
+++ b/src/algorithms/ctmrg/ctmrg.jl
@@ -137,7 +137,7 @@ end
     ctmrg_iter(state, envs::CTMRGEnv, alg::CTMRG) -> envs′, info
 
 Perform one iteration of CTMRG that maps the `state` and `envs` to a new environment,
-and also returns the truncation error.
+and also returns the `info` `NamedTuple`.
 """
 function ctmrg_iter(state, envs::CTMRGEnv, alg::SequentialCTMRG)
     ϵ = zero(real(scalartype(state)))
@@ -178,12 +178,9 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N)
 # ======================================================================================== #
 
 """
-    ctmrg_expand(state, envs, alg::CTMRG{M})
+    ctmrg_expand(coordinates, state, envs)
 
-Expand the environment by absorbing a new PEPS tensor.
-There are two modes of expansion: `M = :sequential` and `M = :simultaneous`.
-The first mode expands the environment in one direction at a time, for convenience towards
-the left. The second mode expands the environment in all four directions simultaneously.
+Expand the environment by absorbing a new PEPS tensor on the given coordinates.
 """
 function ctmrg_expand(coordinates, state, envs::CTMRGEnv)
     return dtmap(idx -> TensorMap(EnlargedCorner(state, envs, idx), idx[1]), coordinates)
@@ -194,10 +191,12 @@ end
 # ======================================================================================== #
 
 """
-    ctmrg_projectors(enlarged_envs, env, alg::CTMRG{M})
+    ctmrg_projectors(col::Int, enlarged_envs, env, alg::CTMRG{:sequential})
+    ctmrg_projectors(enlarged_envs, env, alg::CTMRG{:simultaneous})
 
-Compute the CTMRG projectors based from enlarged environments.
-In the `:simultaneous` mode, the environment SVD is run in parallel.
+Compute the CTMRG projectors based on enlarged environments.
+In the `:sequential` mode the projectors are computed for the column `col`, whereas
+in the `:simultaneous` mode, all projectors (and corresponding SVDs) are computed in parallel.
 """
 function ctmrg_projectors(
     col::Int, enlarged_envs, envs::CTMRGEnv{C,E}, alg::SequentialCTMRG
@@ -327,9 +326,12 @@ end
 # ======================================================================================== #
 
 """
-    ctmrg_renormalize(enlarged_envs, projectors, state, envs, alg::CTMRG{M})
+    ctmrg_renormalize(col::Int, projectors, state, envs, ::CTMRG{:sequential})
+    ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::CTMRG{:simultaneous})
 
 Apply projectors to renormalize corners and edges.
+The `:sequential` mode renormalizes the environment on the column `col`, where as the
+`:simultaneous` mode renormalizes all environment tensors simultaneously.
 """
 function ctmrg_renormalize(col::Int, projectors, state, envs, ::SequentialCTMRG)
     corners = Zygote.Buffer(envs.corners)

From 8d814db32e69c4e44d25096b29b2316b20023e25 Mon Sep 17 00:00:00 2001
From: Paul Brehmer <paul.brehmer@univie.ac.at>
Date: Tue, 5 Nov 2024 16:42:38 +0100
Subject: [PATCH 04/20] Remove duplicate lines in Heisenberg test

---
 test/heisenberg.jl | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/test/heisenberg.jl b/test/heisenberg.jl
index c4941c40..1c05b084 100644
--- a/test/heisenberg.jl
+++ b/test/heisenberg.jl
@@ -12,12 +12,6 @@ ctm_alg = CTMRG()
 opt_alg = PEPSOptimize(;
     boundary_alg=ctm_alg, optimizer=LBFGS(4; gradtol=1e-3, verbosity=2)
 )
-ctm_alg = CTMRG(; ctmrgscheme=:sequential)
-opt_alg = PEPSOptimize(;
-    boundary_alg=ctm_alg,
-    optimizer=LBFGS(4; gradtol=1e-3, verbosity=2),
-    gradient_alg=LinSolver(; iterscheme=:diffgauge),
-)
 
 # initialize states
 Random.seed!(91283219347)

From 55f1e0085db0d17526ebfe3950f4826b11979dad Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 6 Nov 2024 16:15:32 +0800
Subject: [PATCH 05/20] add simple update algorithm

---
 src/PEPSKit.jl                          |   8 +
 src/algorithms/timeevol/simpleupdate.jl | 186 ++++++++++++++++++++++++
 src/states/suweight.jl                  |  46 ++++++
 src/utility/mirror.jl                   |  14 ++
 src/utility/svd.jl                      |  13 ++
 src/utility/util.jl                     |  25 ++++
 6 files changed, 292 insertions(+)
 create mode 100644 src/algorithms/timeevol/simpleupdate.jl
 create mode 100644 src/states/suweight.jl
 create mode 100644 src/utility/mirror.jl

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index c8a6b032..957c3f51 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -17,12 +17,14 @@ include("utility/util.jl")
 include("utility/diffable_threads.jl")
 include("utility/svd.jl")
 include("utility/rotations.jl")
+include("utility/mirror.jl")
 include("utility/diffset.jl")
 include("utility/hook_pullback.jl")
 include("utility/autoopt.jl")
 
 include("states/abstractpeps.jl")
 include("states/infinitepeps.jl")
+include("states/suweight.jl")
 
 include("operators/transferpeps.jl")
 include("operators/infinitepepo.jl")
@@ -43,6 +45,9 @@ include("algorithms/ctmrg/sparse_environments.jl")
 include("algorithms/ctmrg/ctmrg.jl")
 include("algorithms/ctmrg/gaugefix.jl")
 
+include("algorithms/timeevol/simpleupdate.jl")
+include("algorithms/timeevol/fullupdate.jl")
+
 include("algorithms/toolbox.jl")
 
 include("algorithms/peps_opt.jl")
@@ -167,6 +172,9 @@ export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolver
 export fixedpoint
 
+export simpleupdate!, absorb_wt
+
+export SUWeight
 export InfinitePEPS, InfiniteTransferPEPS
 export InfinitePEPO, InfiniteTransferPEPO
 export initializeMPS, initializePEPS
diff --git a/src/algorithms/timeevol/simpleupdate.jl b/src/algorithms/timeevol/simpleupdate.jl
new file mode 100644
index 00000000..52400ae9
--- /dev/null
+++ b/src/algorithms/timeevol/simpleupdate.jl
@@ -0,0 +1,186 @@
+"""
+Mirror the unit cell of an iPEPS by its anti-diagonal line
+"""
+function mirror_antidiag!(peps::InfinitePEPS)
+    peps.A[:] = mirror_antidiag(peps.A)
+    for (i, t) in enumerate(peps.A)
+        peps.A[i] = permute(t, (1,), (3,2,5,4))
+    end
+end
+
+"""
+Mirror the unit cell of an iPEPS with weights by its anti-diagonal line
+"""
+function mirror_antidiag!(wts::SUWeight)
+    wts.x[:], wts.y[:] = mirror_antidiag(wts.y), mirror_antidiag(wts.x)
+end
+
+"""
+Absorb environment weight on axis `ax` into tensor `t` at position `(row,col)`
+```
+                ↓
+                y[r,c]
+                ↓
+    ←x[r,c-1] ← T[r,c] ← x[r,c] ←
+                ↓
+                y[r+1,c]
+                ↓
+```
+"""
+function absorb_wt(
+    t::AbstractTensorMap, row::Int, col::Int, 
+    ax::Int, wts::SUWeight; 
+    sqrtwt::Bool=false, invwt::Bool=false
+)
+    Nr, Nc = size(wts)
+    @assert 1 <= row <= Nr && 1 <= col <= Nc
+    @assert 2 <= ax <= 5
+    pow = (sqrtwt ? 1/2 : 1) * (invwt ? -1 : 1)
+    if ax == 2 # north
+        wt = wts.y[row, col]
+    elseif ax == 3 # east
+        wt = wts.x[row, col]
+    elseif ax == 4 # south
+        wt = wts.y[_next(row,Nr), col]
+    else # west
+        wt = wts.x[row, _prev(col,Nc)]
+    end
+    wt2 = sdiag_pow(wt, pow)
+    indices_t = collect(-1:-1:-5)
+    indices_t[ax] = 1
+    indices_wt = (ax in (2,3) ? [1,-ax] : [-ax,1])
+    t2 = ncon((t, wt2), (indices_t, indices_wt))
+    # restore codomain and domain
+    t2 = permute(t2, (1,), Tuple(2:5))
+    return t2
+end
+
+
+"""
+Simple update of bond `wts.x[r,c]`
+```
+                y[r,c]              y[r,c+1]
+                ↓                   ↓
+    x[r,c-1] ←- T[r,c] ←- x[r,c] ←- T[r,c+1] ← x[r,c+1]
+                ↓                   ↓
+                y[r+1,c]            y[r+1,c+1]
+```
+"""
+function _su_bondx!(
+    row::Int, col::Int, gate::AbstractTensorMap, 
+    peps::InfinitePEPS, wts::SUWeight,
+    Dcut::Int, svderr::Float64=1e-10
+)
+    Nr, Nc = size(peps)
+    @assert 1 <= row <= Nr && 1 <= col <= Nc
+    row2, col2 = row, _next(col,Nc)
+    T1, T2 = peps[row,col], peps[row2,col2]
+    # absorb environment weights
+    for ax in (2,4,5)
+        T1 = absorb_wt(T1, row, col, ax, wts)
+    end
+    for ax in (2,3,4)
+        T2 = absorb_wt(T2, row2, col2, ax, wts)
+    end
+    # absorb bond weight
+    T1 = absorb_wt(T1, row, col, 3, wts; sqrtwt=true)
+    T2 = absorb_wt(T2, row2, col2, 5, wts; sqrtwt=true)
+    # QR and LQ decomposition
+    """
+        2   1               1             2
+        ↓ ↗                 ↓            ↗
+    5 ← T ← 3   ====>   3 ← X ← 4 ← 1 ← aR ← 3
+        ↓                   ↓
+        4                   2
+    """
+    X, aR = leftorth(T1, ((2,4,5), (1,3)), alg=QRpos())
+    """
+        2   1                 2         2
+        ↓ ↗                 ↗           ↓
+    5 ← T ← 3   ====>  1 ← bL ← 3 ← 1 ← Y ← 3
+        ↓                               ↓
+        4                               4
+    """
+    bL, Y = rightorth(T2, ((5,1), (2,3,4)), alg=LQpos())
+    # apply gate
+    """
+            -2          -3
+            ↑           ↑
+            |----gate---|
+            ↑           ↑
+            1           2
+            ↑           ↑
+        -1← aR -← 3 -← bL ← -4
+    """
+    tmp = ncon((gate, aR, bL), ([-2,-3,1,2], [-1,1,3], [3,2,-4]))
+    # SVD
+    truncscheme = truncerr(svderr) & truncdim(Dcut)
+    aR, s, bL, ϵ = tsvd(tmp, ((1,2), (3,4)); trunc=truncscheme)
+    """
+            -2         -1              -1    -2
+            |         ↗               ↗       |
+        -5- X ← 1 ← aR - -3     -5 - bL ← 1 ← Y - -3
+            |                                 |
+            -4                               -4
+    """
+    T1 = ncon((X, aR), ([-2,-4,-5,1], [1,-1,-3]))
+    T2 = ncon((bL, Y), ([-5,-1,1], [1,-2,-3,-4]))
+    # remove environment weights
+    for ax in (2,4,5)
+        T1 = absorb_wt(T1, row, col, ax, wts; invwt=true)
+    end
+    for ax in (2,3,4)
+        T2 = absorb_wt(T2, row2, col2, ax, wts; invwt=true)
+    end
+    # update tensor dict and weight on current bond (with normalization)
+    peps.A[row,col], peps.A[row2,col2] = T1, T2
+    wts.x[row,col] = s / maxabs(s)
+    return ϵ
+end
+
+
+"""
+One round of simple update on the input InfinitePEPS `peps` 
+and SUWeight `wts` with the nearest neighbor gate `gate`
+
+When `bipartite === true` (for square lattice), the unit cell size should be 2 x 2, 
+and the tensor and x/y weight at `(row, col)` is the same as `(row+1, col+1)`
+"""
+function simpleupdate!(
+    gate::AbstractTensorMap, peps::InfinitePEPS, wts::SUWeight, 
+    Dcut::Int, svderr::Float64=1e-10; bipartite::Bool=false,
+)
+    Nr, Nc = size(peps)
+    if bipartite
+        @assert Nr == Nc == 2
+    end
+    # TODO: make algorithm independent on the choice of dual in the network
+    for (r, c) in Iterators.product(1:Nr, 1:Nc)
+        @assert [isdual(space(peps.A[r, c], ax)) for ax in 1:5] == [0,1,1,0,0]
+        @assert [isdual(space(wts.x[r, c], ax)) for ax in 1:2] == [0,1]
+        @assert [isdual(space(wts.y[r, c], ax)) for ax in 1:2] == [0,1]
+    end
+    for direction in 1:2
+        # mirror the y-weights to x-direction 
+        # to update them using code for x-weights
+        if direction == 2
+            mirror_antidiag!(peps); mirror_antidiag!(wts)
+        end
+        if bipartite
+            ϵ = _su_bondx!(1, 1, gate, peps, wts, Dcut, svderr)
+            (peps.A[2,2], peps.A[2,1], wts.x[2,2]) = deepcopy.((peps.A[1,1], peps.A[1,2], wts.x[1,1]))
+            ϵ = _su_bondx!(2, 1, gate, peps, wts, Dcut, svderr)
+            (peps.A[1,2], peps.A[1,1], wts.x[1,2]) = deepcopy.((peps.A[2,1], peps.A[2,2], wts.x[2,1]))
+        else
+            for site in CartesianIndices(peps.A)
+                row, col = Tuple(site)
+                ϵ = _su_bondx!(row, col, gate, peps, wts, Dcut)
+            end
+        end
+        if direction == 2
+            mirror_antidiag!(peps); mirror_antidiag!(wts)
+        end
+    end
+    return nothing
+end
+
diff --git a/src/states/suweight.jl b/src/states/suweight.jl
new file mode 100644
index 00000000..97b75832
--- /dev/null
+++ b/src/states/suweight.jl
@@ -0,0 +1,46 @@
+"""
+Schmidt bond weight used in simple/cluster update
+"""
+struct SUWeight{T<:AbstractTensorMap}
+    x::Matrix{T}
+    y::Matrix{T}
+
+    function SUWeight(wxs::Matrix{T}, wys::Matrix{T}) where {T}
+        new{T}(wxs, wys)
+    end
+end
+
+function Base.size(wts::SUWeight)
+    @assert size(wts.x) == size(wts.y)
+    return size(wts.x)
+end
+
+function Base.:(==)(wts1::SUWeight, wts2::SUWeight)
+    return wts1.x == wts2.x && wts1.y == wts2.y
+end
+
+function Base.:(+)(wts1::SUWeight, wts2::SUWeight)
+    return SUWeight(wts1.x + wts2.x, wts1.y + wts2.y)
+end
+
+function Base.:(-)(wts1::SUWeight, wts2::SUWeight)
+    return SUWeight(wts1.x - wts2.x, wts1.y - wts2.y)
+end
+
+function Base.iterate(wts::SUWeight, state=1)
+    nx = prod(size(wts.x))
+    if 1 <= state <= nx
+        return wts.x[state], state+1
+    elseif nx+1 <= state <= 2*nx
+        return wts.y[state-nx], state+1
+    else
+        return nothing
+    end
+end
+
+function Base.isapprox(wts1::SUWeight, wts2::SUWeight; atol=0.0, rtol=1e-5)
+    return (
+        isapprox(wts1.x, wts2.x; atol=atol, rtol=rtol) &&
+        isapprox(wts1.y, wts2.y; atol=atol, rtol=rtol)
+    )
+end
diff --git a/src/utility/mirror.jl b/src/utility/mirror.jl
new file mode 100644
index 00000000..f9bc74ae
--- /dev/null
+++ b/src/utility/mirror.jl
@@ -0,0 +1,14 @@
+"""
+Mirror a matrix by its anti-diagonal line
+(the 45 degree line through the lower-left corner)
+
+The element originally at [r, c] is moved [Nc-c+1, Nr-r+1], 
+i.e. the element now at [r, c] was originally at [Nr-c+1, Nc-r+1]
+"""
+function mirror_antidiag(arr::AbstractMatrix)
+    Nr, Nc = size(arr)
+    return collect(
+        arr[Nr-c+1, Nc-r+1]
+        for (r, c) in Iterators.product(1:Nc, 1:Nr)
+    )
+end
diff --git a/src/utility/svd.jl b/src/utility/svd.jl
index a590c87a..20950ac9 100644
--- a/src/utility/svd.jl
+++ b/src/utility/svd.jl
@@ -291,3 +291,16 @@ function _lorentz_broaden(x::Real, ε=1e-12)
     x′ = 1 / x
     return x′ / (x′^2 + ε)
 end
+
+
+"""
+Given `tsvd` result `u`, `s` and `vh`, 
+absorb singular values `s` into `u` and `vh` by
+```
+    u -> u * sqrt(s), vh -> sqrt(s) * vh
+```
+"""
+function absorb_s(u::AbstractTensorMap, s::AbstractTensorMap, vh::AbstractTensorMap)
+    sqrt_s = sdiag_pow(s, 0.5)
+    return u * sqrt_s, sqrt_s * vh
+end
diff --git a/src/utility/util.jl b/src/utility/util.jl
index 0b9fd4c4..181eb3a9 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -21,6 +21,31 @@ function _elementwise_mult(a::AbstractTensorMap, b::AbstractTensorMap)
     return dst
 end
 
+"""
+Return the maximum absolute value of tensor elements
+"""
+function maxabs(t::AbstractTensorMap)
+    maxel = 0.0
+    for (k, b) in blocks(t)
+        maxelb = maximum(abs.(b))
+        if maxelb > maxel
+            maxel = maxelb
+        end
+    end
+    return maxel
+end
+
+"""
+Compute S^(pow) for diagonal matrices `S`
+"""
+function sdiag_pow(S::AbstractTensorMap, pow::Real)
+    S2 = similar(S)
+    for (k, b) in blocks(S)
+        copyto!(blocks(S2)[k], diagm(diag(b).^pow))
+    end
+    return S2
+end
+
 # Compute √S⁻¹ for diagonal TensorMaps
 _safe_inv(a, tol) = abs(a) < tol ? zero(a) : inv(a)
 function sdiag_inv_sqrt(S::AbstractTensorMap; tol::Real=eps(eltype(S))^(3 / 4))

From 8c03cc220ff25586f9afc20ffaffe5d1438cc9ee Mon Sep 17 00:00:00 2001
From: Paul Brehmer <paul.brehmer@univie.ac.at>
Date: Wed, 6 Nov 2024 11:53:13 +0100
Subject: [PATCH 06/20] Stabilize large unit cell energy test

---
 test/ctmrg/ctmrgschemes.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/ctmrg/ctmrgschemes.jl b/test/ctmrg/ctmrgschemes.jl
index cf556be2..406d4c1f 100644
--- a/test/ctmrg/ctmrgschemes.jl
+++ b/test/ctmrg/ctmrgschemes.jl
@@ -48,7 +48,7 @@ unitcells = [(1, 1), (3, 4)]
     H = heisenberg_XYZ(InfiniteSquare(unitcell...))
     E_sequential = costfun(psi, env_sequential, H)
     E_simultaneous = costfun(psi, env_simultaneous, H)
-    @test E_sequential ≈ E_simultaneous rtol = 1e-4
+    @test E_sequential ≈ E_simultaneous rtol = 1e-3
 end
 
 # test fixedspace actually fixes space

From f9d19fd04950ed6082c535228b1f42b1c1139e30 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 6 Nov 2024 20:39:13 +0800
Subject: [PATCH 07/20] add full update core algorithm

---
 src/PEPSKit.jl                          |   3 +-
 src/algorithms/ctmrg/ctmrg.jl           |  29 ++-
 src/algorithms/timeevol/fu_gaugefix.jl  |  92 +++++++
 src/algorithms/timeevol/fu_optimize.jl  | 325 ++++++++++++++++++++++++
 src/algorithms/timeevol/fullupdate.jl   | 195 ++++++++++++++
 src/algorithms/timeevol/simpleupdate.jl |  27 +-
 src/environments/ctmrg_environments.jl  |  38 +++
 src/states/infinitepeps.jl              |  12 +
 8 files changed, 701 insertions(+), 20 deletions(-)
 create mode 100644 src/algorithms/timeevol/fu_gaugefix.jl
 create mode 100644 src/algorithms/timeevol/fu_optimize.jl
 create mode 100644 src/algorithms/timeevol/fullupdate.jl

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index 957c3f51..bd71e728 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -64,7 +64,7 @@ include("utility/symmetrization.jl")
         const ctmrgscheme = :simultaneous
         const reuse_env = true
         const trscheme = FixedSpaceTruncation()
-        const fwd_alg = TensorKit.SVD()
+        const fwd_alg = TensorKit.SDD()
         const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
         const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
         const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=2)
@@ -173,6 +173,7 @@ export PEPSOptimize, GeomSum, ManualIter, LinSolver
 export fixedpoint
 
 export simpleupdate!, absorb_wt
+export fullupdate!
 
 export SUWeight
 export InfinitePEPS, InfiniteTransferPEPS
diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl
index 8e74436f..79fb5317 100644
--- a/src/algorithms/ctmrg/ctmrg.jl
+++ b/src/algorithms/ctmrg/ctmrg.jl
@@ -133,6 +133,28 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
     end
 end
 
+"""
+Perform CTMRG left move on the `col`-th column
+"""
+function ctmrg_leftmove(
+    col::Int, state, envs::CTMRGEnv, alg::SequentialCTMRG
+)
+    """
+        ----> left move
+        C1 ← T1 ←   r-1
+        ↓    ‖
+        T4 = M ==   r
+        ↓    ‖
+        C4 → T3 →   r+1
+        c-1  c 
+    """
+    enlarged_envs = ctmrg_expand(
+        eachcoordinate(envs, [4, 1])[:, :, col], state, envs
+    )
+    projectors, info = ctmrg_projectors(col, enlarged_envs, envs, alg)
+    envs = ctmrg_renormalize(col, projectors, state, envs, alg)
+    return envs, info
+end
 """
     ctmrg_iter(state, envs::CTMRGEnv, alg::CTMRG) -> envs′, info
 
@@ -143,17 +165,12 @@ function ctmrg_iter(state, envs::CTMRGEnv, alg::SequentialCTMRG)
     ϵ = zero(real(scalartype(state)))
     for _ in 1:4 # rotate
         for col in 1:size(state, 2) # left move column-wise
-            enlarged_envs = ctmrg_expand(
-                eachcoordinate(envs, [4, 1])[:, :, col], state, envs
-            )
-            projectors, info = ctmrg_projectors(col, enlarged_envs, envs, alg)
-            envs = ctmrg_renormalize(col, projectors, state, envs, alg)
+            envs, info = ctmrg_leftmove(col, state, envs, alg)
             ϵ = max(ϵ, info.err)
         end
         state = rotate_north(state, EAST)
         envs = rotate_north(envs, EAST)
     end
-
     return envs, (; err=ϵ)
 end
 function ctmrg_iter(state, envs::CTMRGEnv, alg::SimultaneousCTMRG)
diff --git a/src/algorithms/timeevol/fu_gaugefix.jl b/src/algorithms/timeevol/fu_gaugefix.jl
new file mode 100644
index 00000000..760d781d
--- /dev/null
+++ b/src/algorithms/timeevol/fu_gaugefix.jl
@@ -0,0 +1,92 @@
+"""
+Replace `env` by its positive/negative approximant `± Z Z†`
+(returns the sign and Z†)
+```
+                        |-→ 1   2 ←-|
+                        |           |
+    |----env----|       |←--- Z ---→|
+    |→ 1     2 ←|   =         ↑
+    |← 3     4 →|       |---→ Z† ←--|
+    |-----------|       |           |
+                        |←- 3   4 -→|
+```
+"""
+function positive_approx(env::AbstractTensorMap)
+    @assert [isdual(space(env, ax)) for ax in 1:4] == [0,0,1,1]
+    # hermitize env, and perform eigen-decomposition
+    # env = U D U'
+    D, U = eigh((env + env')/2)
+    # determine env is (mostly) positive or negative
+    sgn = sign(mean(vcat((diag(b) for (k,b) in blocks(D))...)))
+    if sgn == -1
+        D *= -1
+    end
+    # set negative eigenvalues to 0
+    for (k, b) in blocks(D)
+        for i in diagind(b)
+            if b[i] < 0
+                b[i] = 0.0
+            end
+        end
+    end
+    Zdg = sdiag_pow(D, 1/2) * U'
+    return sgn, Zdg
+end
+
+"""
+Fix local gauge of the env tensor around a bond
+"""
+function fu_fixgauge(
+    Zdg::AbstractTensorMap, 
+    X::AbstractTensorMap, Y::AbstractTensorMap, 
+    aR::AbstractTensorMap, bL::AbstractTensorMap
+)
+    #= 
+            1               1
+            ↑               ↑
+        2 → Z† ← 3  =   2 → QR ← 3  1 ← R ← 2
+
+                                        1
+                                        ↑
+                    =   2 → L → 1   3 → QL ← 2
+    =#
+    QR, R = leftorth(Zdg, ((1,2), (3,)), alg=QRpos())
+    QL, L = leftorth(Zdg, ((1,3), (2,)), alg=QRpos())
+    @assert !isdual(codomain(R)[1]) && !isdual(domain(R)[1])
+    @assert !isdual(codomain(L)[1]) && !isdual(domain(L)[1])
+    Rinv, Linv = inv(R), inv(L)
+    #= fix gauge of aR, bL, Z†
+
+                    ↑
+        |→-(Linv -→ Z† ← Rinv)←-|
+        |                       |
+        ↑                       ↑
+        |        ↑     ↑        |
+        |← (L ← aR) ← (bL → R) →|
+        |-----------------------|
+
+                     -2              -2
+                      ↑               ↑        
+        -1 ← L ← 1 ← aR2 ← -3   -1 ← bL2 → 1 → R → -3
+
+                        -1
+                        ↑
+        -2 → Linv → 1 → Z† ← 2 ← Rinv ← -3
+    =#
+    aR = ncon([L, aR], [[-1,1], [1,-2,-3]])
+    bL = ncon([bL, R], [[-1,-2,1], [-3,1]])
+    Zdg = permute(
+        ncon([Zdg, Linv, Rinv], [[-1,1,2], [1,-2], [2,-3]]),
+        (1,), (2,3)
+    )
+    #= fix gauge of X, Y
+            -1                                      -1
+             |                                      |
+        -4 - X ← 1 ← Linv ← -2      -4 → Rinv → 1 → Y - -2
+             |                                      |
+            -3                                      -3
+    =#
+    X = ncon([X, Linv], [[-1,1,-3,-4], [1,-2]])
+    Y = ncon([Y, Rinv], [[-1,-2,-3,1], [1,-4]])
+    return Zdg, X, Y, aR, bL
+end
diff --git a/src/algorithms/timeevol/fu_optimize.jl b/src/algorithms/timeevol/fu_optimize.jl
new file mode 100644
index 00000000..a89f8d3e
--- /dev/null
+++ b/src/algorithms/timeevol/fu_optimize.jl
@@ -0,0 +1,325 @@
+"""
+Construct the environment (norm) tensor
+```
+    left half                       right half
+    C1 -χ4 - T1 ------- χ6 ------- T1 - χ8 - C2     r-1
+    |        ‖                      ‖        |
+    χ2      DNX                    DNY      χ10
+    |        ‖                      ‖        |
+    T4 =DWX= XX = DX =       = DY = YY =DEY= T2     r
+    |        ‖                      ‖        |
+    χ1      DSX                    DSY       χ9
+    |        ‖                      ‖        |
+    C4 -χ3 - T3 ------- χ5 ------- T3 - χ7 - C3     r+1
+    c-1      c                      c+1     c+2
+```
+which can be more simply denoted as
+```
+    |------------|
+    |→ DX1  DY1 ←|   axis order
+    |← DX0  DX1 →|   (DX1, DY1, DX0, DY0)
+    |------------|
+```
+The axes 1, 2 (or 3, 4) come from X†, Y† (or X, Y)
+"""
+function tensor_env(
+    row::Int, col::Int, X::AbstractTensorMap, 
+    Y::AbstractTensorMap, envs::CTMRGEnv
+)
+    Nr, Nc = size(envs.corners)[[2,3]]
+    cm1 = _prev(col, Nc);
+    cp1 = _next(col, Nc); cp2 = _next(cp1, Nc)
+    rm1 = _prev(row, Nr); rp1 = _next(row, Nr)
+    c1 = envs.corners[1, rm1, cm1]
+    c2 = envs.corners[2, rm1, cp2]
+    c3 = envs.corners[3, rp1, cp2]
+    c4 = envs.corners[4, rp1, cm1]
+    t1X, t1Y = envs.edges[1, rm1, col], envs.edges[1, rm1, cp1]
+    t2 = envs.edges[2, row, cp2]
+    t3X, t3Y = envs.edges[3, rp1, col], envs.edges[3, rp1, cp1]
+    t4 = envs.edges[4, row, cm1]
+    # left half
+    @autoopt @tensor lhalf[DX1, DX0, χ5, χ6] := (
+        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]
+    )
+    # right half
+    @autoopt @tensor rhalf[DY1, DY0, χ5, χ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]
+    )
+    # combine
+    @autoopt @tensor env[DX1, DY1; DX0, DY0] := (
+        lhalf[DX1, DX0, χ5, χ6] * rhalf[DY1, DY0, χ5, χ6]
+    )
+    return env
+end
+
+
+"""
+Construct the tensor
+```
+    |------------env------------|
+    |→ DX1     Db1 → bL† ← DY1 ←|
+    |                ↑          |
+    |                db         |
+    |                ↑          |
+    |← DX0     Db0 ← bL -→ DY0 →|
+    |---------------------------|
+```
+"""
+function tensor_Ra(env::AbstractTensorMap, bL::AbstractTensorMap)
+    @autoopt @tensor Ra[DX1, Db1, DX0, Db0] := (
+        env[DX1, DY1, DX0, DY0] *
+        bL[Db0, db, DY0] * conj(bL[Db1, db, DY1])
+    )
+    return Ra
+end
+
+
+"""
+Construct the tensor
+```
+    |--------------env--------------|
+    |→ DX1         Db1 → bL† ← DY1 ←|
+    |                    ↑          |
+    |         da         db         |
+    |         ↑          ↑          |
+    |← DX0 ←- aR2 ←- D ← bL2 → DY0 →|
+    |-------------------------------|
+```
+"""
+function tensor_Sa(
+    env::AbstractTensorMap, aR2::AbstractTensorMap, 
+    bL::AbstractTensorMap, bL2::AbstractTensorMap
+)
+    @autoopt @tensor Sa[DX1, Db1, da] := (
+        env[DX1, DY1, DX0, DY0] * conj(bL[Db1, db, DY1]) *
+        bL2[D, db, DY0] * aR2[DX0, da, D]
+    )
+    return Sa
+end
+
+
+"""
+Construct the tensor
+```
+    |------------env------------|
+    |→ DX1 → aR† → Da1     DY1 ←|
+    |        ↑                  |
+    |        da                 |
+    |        ↑                  |
+    |← DX0 ← aR ←- Da0     DY0 →|
+    |---------------------------|
+```
+"""
+function tensor_Rb(env::AbstractTensorMap, aR::AbstractTensorMap)
+    @autoopt @tensor Rb[Da1, DY1, Da0, DY0] := (
+        env[DX1, DY1, DX0, DY0] *
+        aR[DX0, da, Da0] * conj(aR[DX1, da, Da1])
+    )
+    return Rb
+end
+
+
+"""
+Construct the tensor
+```
+    |--------------env--------------|
+    |→ DX1 → aR† → Da1         DY1 ←|
+    |        ↑                      |
+    |        da         db          |
+    |        ↑          ↑           |
+    |← DX0 ← aR2 ← D ←- bL2 -→ DY0 →|
+    |-------------------------------|
+```
+"""
+function tensor_Sb(
+    env::AbstractTensorMap, aR::AbstractTensorMap, 
+    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+)
+    @autoopt @tensor Sb[Da1, DY1, db] := (
+        env[DX1, DY1, DX0, DY0] * conj(aR[DX1, da, Da1]) *
+        aR2[DX0, da, D] * bL2[D, db, DY0]
+    )
+    return Sb
+end
+
+
+"""
+Calculate the norm <Psi(a1,b1)|Psi(a2,b2)>
+```
+    |--------------env--------------|
+    |→ DX1 → aR1†→ D1 → bL1† ← DY1 ←|
+    |        ↑           ↑          |
+    |        da          db         |
+    |        ↑           ↑          |
+    |← DX0 ← aR2 ← D0 ← bL2 → DY0 -→|
+    |-------------------------------|
+```
+"""
+function inner_prod(
+    env::AbstractTensorMap, 
+    aR1::AbstractTensorMap, bL1::AbstractTensorMap,
+    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+)
+    @autoopt @tensor t[:] := (
+        env[DX1, DY1, DX0, DY0] *
+        conj(aR1[DX1, da, D1]) * conj(bL1[D1, db, DY1]) *
+        aR2[DX0, da, D0] * bL2[D0, db, DY0]
+    )
+    return first(blocks(t))[2][1]
+end
+
+"""
+Calculate the cost function
+```
+    f(a,b)  = | |Psi(a,b)> - |Psi(a2,b2)> |^2
+    = <Psi(a,b)|Psi(a,b)> + <Psi(a2,b2)|Psi(a2,b2)>
+        - 2 Re<Psi(a,b)|Psi(a2,b2)>
+```
+"""
+function cost_func(
+    env::AbstractTensorMap, 
+    aR::AbstractTensorMap, bL::AbstractTensorMap,
+    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+)
+    t1 = inner_prod(env, aR, bL, aR, bL)
+    t2 = inner_prod(env, aR2, bL2, aR2, bL2)
+    t3 = inner_prod(env, aR, bL, aR2, bL2)
+    return real(t1) + real(t2) - 2 * real(t3)
+end
+
+
+"""
+Calculate the approximate local inner product
+`<aR1 bL1|aR2 bL2>`
+```
+    |→ aR1† → D1 → bL1† ←|
+    |   ↑          ↑     |
+    DW  da         db    DE
+    |   ↑          ↑     |
+    |← aR2 ←- D0 ← bL2 -→|
+```
+"""
+function inner_prod_local(
+    aR1::AbstractTensorMap, bL1::AbstractTensorMap,
+    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+)
+    @autoopt @tensor t[:] := (
+        conj(aR1[DW, da, D1]) * conj(bL1[D1, db, DE]) *
+        aR2[DW, da, D0] * bL2[D0, db, DE]
+    )
+    return first(blocks(t))[2][1]
+end
+
+"""
+Calculate the fidelity using aR, bL
+between two evolution steps
+```
+                |<aR1 bL1 | aR2 bL2>|
+    ---------------------------------------------
+    sqrt(<aR1 bL1 | aR1 bL1> <aR2 bL2 | aR2 bL2>)
+```
+"""
+function local_fidelity(
+    aR1::AbstractTensorMap, bL1::AbstractTensorMap, 
+    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+)
+    b12 = inner_prod_local(aR1, bL1, aR2, bL2)
+    b11 = inner_prod_local(aR1, bL1, aR1, bL1)
+    b22 = inner_prod_local(aR2, bL2, aR2, bL2)
+    return abs(b12) / sqrt(abs(b11*b22))
+end
+
+"""
+Solving the equations
+```
+    Ra aR = Sa, Rb bL = Sb
+```
+"""
+function solve_ab(
+    R::AbstractTensorMap, S::AbstractTensorMap, 
+    ab0::AbstractTensorMap
+)
+    f(x) = ncon((R, x), ([-1,-2,1,2], [1,2,-3]))
+    ab, info = linsolve(f, S, permute(ab0, (1,3,2)), 0, 1)
+    return permute(ab, (1,3,2)), info
+end
+
+"""
+Minimize the cost function
+```
+    fix bL:
+    d(aR,aR†) = aR† Ra aR - aR† Sa - Sa† aR + T
+    minimized by Ra aR = Sa
+
+    fix aR:
+    d(bL,bL†) = bL† Rb bL - bL† Sb - Sb† bL + T
+    minimized by Rb bL = Sb
+```
+`aR0`, `bL0` are initial values of `aR`, `bL`
+"""
+function fu_optimize(
+    aR0::AbstractTensorMap, bL0::AbstractTensorMap, 
+    aR2::AbstractTensorMap, bL2::AbstractTensorMap, 
+    env::AbstractTensorMap;
+    maxiter::Int=50, maxdiff::Float64=1e-15, 
+    check_int::Int=1, verbose::Bool=false
+)
+    if verbose
+        println("---- Iterative optimization ----")
+        @printf(
+            "%-6s%12s%12s%12s %10s\n", 
+            "Step", "Cost", "ϵ_d", "ϵ_ab", "Time/s"
+        )
+    end
+    aR, bL = deepcopy(aR0), deepcopy(bL0)
+    time0 = time()
+    cost00 = cost_func(env, aR, bL, aR2, bL2)
+    fid00 = local_fidelity(aR, bL, aR2, bL2)
+    cost0, fid0 = cost00, fid00
+    # no need to further optimize
+    if abs(cost0) < 5e-15
+        if verbose
+            time1 = time()
+            println(@sprintf(
+                "%-6d%12.3e%12.3e%12.3e %10.3f\n", 
+                0, cost0, NaN, NaN, time1 - time0
+            ))
+        end
+        return aR, bL, cost0
+    end
+    for count in 1:maxiter
+        time0 = time()
+        Ra = tensor_Ra(env, bL)
+        Sa = tensor_Sa(env, aR2, bL, bL2)
+        aR, info_a = solve_ab(Ra, Sa, aR)
+        Rb = tensor_Rb(env, aR)
+        Sb = tensor_Sb(env, aR, aR2, bL2)
+        bL, info_b = solve_ab(Rb, Sb, bL)
+        cost = cost_func(env, aR, bL, aR2, bL2)
+        fid = local_fidelity(aR, bL, aR2, bL2)
+        diff_d = abs(cost - cost0) / cost00
+        diff_ab = abs(fid - fid0) / fid00
+        time1 = time()
+        if verbose && (count == 1 || count % check_int == 0)
+            @printf(
+                "%-6d%12.3e%12.3e%12.3e %10.3f\n", 
+                count, cost, diff_d, diff_ab, time1 - time0
+            )
+        end
+        if diff_ab < maxdiff
+            break
+        end
+        aR0, bL0 = deepcopy(aR), deepcopy(bL)
+        cost0, fid0 = cost, fid
+        if count == maxiter
+            println("Warning: max iter $maxiter reached for optimization")
+        end
+    end
+    return aR, bL, cost0
+end
+
diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/timeevol/fullupdate.jl
new file mode 100644
index 00000000..fdede082
--- /dev/null
+++ b/src/algorithms/timeevol/fullupdate.jl
@@ -0,0 +1,195 @@
+include("fu_gaugefix.jl")
+include("fu_optimize.jl")
+
+"""
+CTMRG left-move to update CTMRGEnv in the c-th column 
+```
+    ---> absorb
+    C1 ← T1 ←   r-1
+    ↓    ‖
+    T4 = M' =   r
+    ↓    ‖
+    C4 → T3 →   r+1
+    c-1  c 
+```
+"""
+function ctmrg_leftmove!(
+    col::Int, peps::InfinitePEPS, envs::CTMRGEnv,
+    chi::Int, svderr::Float64=1e-9
+)
+    trscheme = truncerr(svderr) & truncdim(chi)
+    alg = CTMRG(
+        verbosity=0, miniter=1, maxiter=10, 
+        trscheme=trscheme, ctmrgscheme=:sequential
+    )
+    envs2, info = ctmrg_leftmove(col, peps, envs, alg)
+    envs.corners[:, :, col] = envs2.corners[:, :, col]
+    envs.edges[:, :, col] = envs2.edges[:, :, col]
+    return info
+end
+
+"""
+CTMRG right-move to update CTMRGEnv in the c-th column
+```
+    absorb <---
+    ←-- T1 ← C2     r-1
+        ‖    ↑
+    === M' = T2     r
+        ‖    ↑
+    --→ T3 → C3     r+1
+        c   c+1
+```
+"""
+function ctmrg_rightmove!(
+    col::Int, peps::InfinitePEPS, envs::CTMRGEnv,
+    chi::Int, svderr::Float64=1e-9
+)
+    Nr, Nc = size(peps)
+    @assert 1 <= col <= Nc
+    PEPSKit.rot180!(envs)
+    ctmrg_leftmove!(Nc + 1 - col, rot180(peps), envs, chi, svderr)
+    PEPSKit.rot180!(envs)
+    return nothing
+end
+
+"""
+Update all horizontal bonds in the c-th column
+(i.e. `(r,c) (r,c+1)` for all `r = 1, ..., Nr`).
+To update rows, rotate the network clockwise by 90 degrees.
+"""
+function update_column!(
+    col::Int, gate::AbstractTensorMap,
+    peps::InfinitePEPS, envs::CTMRGEnv,
+    Dcut::Int, chi::Int;
+    svderr::Float64=1e-9, maxiter::Int=50,
+    maxdiff::Float64=1e-15, gaugefix::Bool=true,
+)
+    Nr, Nc = size(peps)
+    @assert 1 <= col <= Nc
+    localfid = 0.0
+    costs = zeros(Nr)
+    truncscheme = truncerr(svderr) & truncdim(Dcut)
+    #= Axis order of X, aR, Y, bL
+
+            1             2            2         1
+            |            ↗           ↗           |
+        4 - X ← 2   1 ← aR ← 3  1 ← bL → 3   4 → Y - 2
+            |                                    |
+            3                                    3
+    =#
+    for row in 1:Nr
+        cp1 = _next(col, Nc)
+        A, B = peps[row, col], peps[row, cp1]
+        # TODO: relax dual requirement on the bonds
+        @assert !isdual(domain(A)[2])
+        #= QR and LQ decomposition
+
+            2   1               1             2
+            | ↗                 |            ↗
+        5 - A ← 3   ====>   4 - X ← 2   1 ← aR ← 3
+            |                   |
+            4                   3
+        =#
+        X, aR0 = leftorth(A, ((2, 4, 5), (1, 3)); alg=QRpos())
+        X = permute(X, (1, 4, 2, 3))
+        #=
+            2   1                 2         2
+            | ↗                 ↗           |
+        5 → B - 3   ====>  1 ← bL → 3   1 → Y - 3
+            |                               |
+            4                               4
+        =#
+        Y, bL0 = leftorth(B, ((2, 3, 4), (1, 5)); alg=QRpos())
+        bL0 = permute(bL0, (3, 2, 1))
+        env = tensor_env(row, col, X, Y, envs)
+        # positive/negative-definite approximant: env = ± Z Z†
+        sgn, Zdg = positive_approx(env)
+        # fix gauge
+        if gaugefix
+            Zdg, X, Y, aR0, bL0 = fu_fixgauge(Zdg, X, Y, aR0, bL0)
+        end
+        env = sgn * Zdg' * Zdg
+        #= apply gate
+        
+                -2          -3
+                ↑           ↑
+                |----gate---|
+                ↑           ↑
+                1           2
+                ↑           ↑
+            -1← aR -← 3 -← bL → -4
+        =#
+        tmp = ncon((gate, aR0, bL0), ([-2, -3, 1, 2], [-1, 1, 3], [3, 2, -4]))
+        # initialize truncated tensors using simple SVD truncation
+        aR2, s, bL2, ϵ = tsvd(tmp, ((1, 2), (3, 4)); trunc=truncerr(1e-15))
+        aR, s_cut, bL, ϵ = tsvd(tmp, ((1, 2), (3, 4)); trunc=truncscheme)
+        aR2, bL2 = absorb_s(aR2, s, bL2)
+        aR, bL = absorb_s(aR, s_cut, bL)
+        # optimize aR, bL
+        aR, bL, cost = fu_optimize(
+            aR, bL, aR2, bL2, env; maxiter=maxiter, maxdiff=maxdiff, verbose=false
+        )
+        costs[row] = cost
+        aR /= maxabs(aR)
+        bL /= maxabs(bL)
+        localfid += local_fidelity(aR, bL, aR0, bL0)
+        #= update and normalize peps, ms
+        
+                -2        -1               -1     -2
+                |        ↗                ↗       |
+            -5- X ← 1 ← aR ← -3     -5 ← bL → 1 → Y - -3
+                |                                 |
+                -4                                -4
+        =#
+        peps.A[row, col] = permute(
+            ncon([X, aR], [[-2, 1, -4, -5], [1, -1, -3]]), (1,), Tuple(2:5)
+        )
+        peps.A[row, cp1] = permute(
+            ncon([bL, Y], [[-5, -1, 1], [-2, -3, -4, 1]]), (1,), Tuple(2:5)
+        )
+        # normalize
+        for c_ in [col, cp1]
+            peps.A[row, c_] /= maxabs(peps.A[row, c_])
+        end
+    end
+    # update CTMRGEnv
+    ctmrg_leftmove!(col, peps, envs, chi, svderr)
+    ctmrg_rightmove!(_next(col, Nc), peps, envs, chi, svderr)
+    return localfid, costs
+end
+
+"""
+One round of full update on the input InfinitePEPS `peps` and its CTMRGEnv `envs`
+
+When `cheap === true`, use half-infinite environment to construct CTMRG projectors.
+Otherwise, use full-infinite environment instead.
+
+Reference: Physical Review B 92, 035142 (2015)
+"""
+function fullupdate!(
+    gate::AbstractTensorMap, peps::InfinitePEPS, envs::CTMRGEnv,
+    Dcut::Int, chi::Int, svderr::Float64=1e-9
+)
+    Nr, Nc = size(peps)
+    fid, maxcost = 0.0, 0.0
+    for col in 1:Nc
+        tmpfid, costs = update_column!(
+            col, gate, peps, envs, Dcut, chi; svderr=svderr
+        )
+        fid += tmpfid
+        maxcost = max(maxcost, maximum(costs))
+    end
+    rotr90!(peps)
+    rotr90!(envs)
+    for row in 1:Nr
+        tmpfid, costs = update_column!(
+            row, gate, peps, envs, Dcut, chi; svderr=svderr
+        )
+        fid += tmpfid
+        maxcost = max(maxcost, maximum(costs))
+    end
+    rotl90!(peps)
+    rotl90!(envs)
+    fid /= (2 * Nr * Nc)
+    return fid, maxcost
+end
diff --git a/src/algorithms/timeevol/simpleupdate.jl b/src/algorithms/timeevol/simpleupdate.jl
index 52400ae9..2e464cde 100644
--- a/src/algorithms/timeevol/simpleupdate.jl
+++ b/src/algorithms/timeevol/simpleupdate.jl
@@ -17,6 +17,8 @@ end
 
 """
 Absorb environment weight on axis `ax` into tensor `t` at position `(row,col)`
+
+Weights around the tensor at `(row, col)` are
 ```
                 ↓
                 y[r,c]
@@ -50,7 +52,6 @@ function absorb_wt(
     indices_t[ax] = 1
     indices_wt = (ax in (2,3) ? [1,-ax] : [-ax,1])
     t2 = ncon((t, wt2), (indices_t, indices_wt))
-    # restore codomain and domain
     t2 = permute(t2, (1,), Tuple(2:5))
     return t2
 end
@@ -85,25 +86,24 @@ function _su_bondx!(
     # absorb bond weight
     T1 = absorb_wt(T1, row, col, 3, wts; sqrtwt=true)
     T2 = absorb_wt(T2, row2, col2, 5, wts; sqrtwt=true)
-    # QR and LQ decomposition
-    """
+    #= QR and LQ decomposition
+    
         2   1               1             2
         ↓ ↗                 ↓            ↗
     5 ← T ← 3   ====>   3 ← X ← 4 ← 1 ← aR ← 3
         ↓                   ↓
         4                   2
-    """
-    X, aR = leftorth(T1, ((2,4,5), (1,3)), alg=QRpos())
-    """
+
         2   1                 2         2
         ↓ ↗                 ↗           ↓
     5 ← T ← 3   ====>  1 ← bL ← 3 ← 1 ← Y ← 3
         ↓                               ↓
         4                               4
-    """
+    =#
+    X, aR = leftorth(T1, ((2,4,5), (1,3)), alg=QRpos())
     bL, Y = rightorth(T2, ((5,1), (2,3,4)), alg=LQpos())
-    # apply gate
-    """
+    #= apply gate
+    
             -2          -3
             ↑           ↑
             |----gate---|
@@ -111,18 +111,18 @@ function _su_bondx!(
             1           2
             ↑           ↑
         -1← aR -← 3 -← bL ← -4
-    """
+    =#
     tmp = ncon((gate, aR, bL), ([-2,-3,1,2], [-1,1,3], [3,2,-4]))
     # SVD
     truncscheme = truncerr(svderr) & truncdim(Dcut)
     aR, s, bL, ϵ = tsvd(tmp, ((1,2), (3,4)); trunc=truncscheme)
-    """
+    #=
             -2         -1              -1    -2
             |         ↗               ↗       |
         -5- X ← 1 ← aR - -3     -5 - bL ← 1 ← Y - -3
             |                                 |
             -4                               -4
-    """
+    =#
     T1 = ncon((X, aR), ([-2,-4,-5,1], [1,-1,-3]))
     T2 = ncon((bL, Y), ([-5,-1,1], [1,-2,-3,-4]))
     # remove environment weights
@@ -132,7 +132,8 @@ function _su_bondx!(
     for ax in (2,3,4)
         T2 = absorb_wt(T2, row2, col2, ax, wts; invwt=true)
     end
-    # update tensor dict and weight on current bond (with normalization)
+    # update tensor dict and weight on current bond 
+    # (max element of weight is normalized to 1)
     peps.A[row,col], peps.A[row2,col2] = T1, T2
     wts.x[row,col] = s / maxabs(s)
     return ϵ
diff --git a/src/environments/ctmrg_environments.jl b/src/environments/ctmrg_environments.jl
index 460faf5b..65cd05c1 100644
--- a/src/environments/ctmrg_environments.jl
+++ b/src/environments/ctmrg_environments.jl
@@ -371,6 +371,44 @@ function Base.rotl90(env::CTMRGEnv{C,T}) where {C,T}
     return CTMRGEnv(copy(corners′), copy(edges′))
 end
 
+# in-place rotations (incompatible with autodiff)
+"""
+Rotate the CTMRGEnv `envs` left 90 degrees (anti-clockwise) in place
+"""
+function rotl90!(envs::CTMRGEnv)
+    envs2 = deepcopy(envs)
+    for dir in 1:4
+        dir2 = _prev(dir, 4)
+        envs.corners[dir2, :, :] = rotl90(envs2.corners[dir, :, :])
+        envs.edges[dir2, :, :] = rotl90(envs2.edges[dir, :, :])
+    end
+    return nothing
+end
+"""
+Rotate the CTMRGEnv `envs` right 90 degrees (clockwise) in place
+"""
+function rotr90!(envs::CTMRGEnv)
+    envs2 = deepcopy(envs)
+    for dir in 1:4
+        dir2 = _next(dir, 4)
+        envs.corners[dir2, :, :] = rotr90(envs2.corners[dir, :, :])
+        envs.edges[dir2, :, :] = rotr90(envs2.edges[dir, :, :])
+    end
+    return nothing
+end
+"""
+Rotate the CTMRGEnv `envs` 180 degrees in place
+"""
+function rot180!(envs::CTMRGEnv)
+    envs2 = deepcopy(envs)
+    for dir in 1:4
+        dir2 = _next(_next(dir, 4), 4)
+        envs.corners[dir2, :, :] = rot180(envs2.corners[dir, :, :])
+        envs.edges[dir2, :, :] = rot180(envs2.edges[dir, :, :])
+    end
+    return nothing
+end
+
 Base.eltype(env::CTMRGEnv) = eltype(env.corners[1])
 Base.axes(x::CTMRGEnv, args...) = axes(x.corners, args...)
 function eachcoordinate(x::CTMRGEnv)
diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl
index 51b353d4..827244e9 100644
--- a/src/states/infinitepeps.jl
+++ b/src/states/infinitepeps.jl
@@ -172,6 +172,18 @@ Base.rotl90(t::InfinitePEPS) = InfinitePEPS(rotl90(rotl90.(t.A)))
 Base.rotr90(t::InfinitePEPS) = InfinitePEPS(rotr90(rotr90.(t.A)))
 Base.rot180(t::InfinitePEPS) = InfinitePEPS(rot180(rot180.(t.A)))
 
+# In-place rotations
+function rotl90!(peps::InfinitePEPS)
+    peps.A[:] = rotl90(rotl90.(peps.A))
+end
+function rotr90!(peps::InfinitePEPS)
+    peps.A[:] = rotr90(rotr90.(peps.A))
+end
+function rot180!(peps::InfinitePEPS)
+    peps.A[:] = rot180(rot180.(peps.A))
+end
+
+
 # Chainrules
 function ChainRulesCore.rrule(
     ::typeof(Base.getindex), state::InfinitePEPS, row::Int, col::Int

From 0119c5559f8c0c8367164ee02e753016b0b1bbaa Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Thu, 7 Nov 2024 09:06:16 +0800
Subject: [PATCH 08/20] improve formatting

---
 src/algorithms/ctmrg/ctmrg.jl           |   8 +-
 src/algorithms/timeevol/fu_gaugefix.jl  |  33 +++---
 src/algorithms/timeevol/fu_optimize.jl  | 148 +++++++++++++-----------
 src/algorithms/timeevol/fullupdate.jl   |  46 ++++----
 src/algorithms/timeevol/simpleupdate.jl |  93 ++++++++-------
 src/states/infinitepeps.jl              |   4 +-
 src/states/suweight.jl                  |   8 +-
 src/utility/mirror.jl                   |   5 +-
 src/utility/svd.jl                      |   1 -
 src/utility/util.jl                     |   2 +-
 10 files changed, 185 insertions(+), 163 deletions(-)

diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl
index 79fb5317..1ef2f5c1 100644
--- a/src/algorithms/ctmrg/ctmrg.jl
+++ b/src/algorithms/ctmrg/ctmrg.jl
@@ -136,9 +136,7 @@ end
 """
 Perform CTMRG left move on the `col`-th column
 """
-function ctmrg_leftmove(
-    col::Int, state, envs::CTMRGEnv, alg::SequentialCTMRG
-)
+function ctmrg_leftmove(col::Int, state, envs::CTMRGEnv, alg::SequentialCTMRG)
     """
         ----> left move
         C1 ← T1 ←   r-1
@@ -148,9 +146,7 @@ function ctmrg_leftmove(
         C4 → T3 →   r+1
         c-1  c 
     """
-    enlarged_envs = ctmrg_expand(
-        eachcoordinate(envs, [4, 1])[:, :, col], state, envs
-    )
+    enlarged_envs = ctmrg_expand(eachcoordinate(envs, [4, 1])[:, :, col], state, envs)
     projectors, info = ctmrg_projectors(col, enlarged_envs, envs, alg)
     envs = ctmrg_renormalize(col, projectors, state, envs, alg)
     return envs, info
diff --git a/src/algorithms/timeevol/fu_gaugefix.jl b/src/algorithms/timeevol/fu_gaugefix.jl
index 760d781d..3ee68069 100644
--- a/src/algorithms/timeevol/fu_gaugefix.jl
+++ b/src/algorithms/timeevol/fu_gaugefix.jl
@@ -12,12 +12,12 @@ Replace `env` by its positive/negative approximant `± Z Z†`
 ```
 """
 function positive_approx(env::AbstractTensorMap)
-    @assert [isdual(space(env, ax)) for ax in 1:4] == [0,0,1,1]
+    @assert [isdual(space(env, ax)) for ax in 1:4] == [0, 0, 1, 1]
     # hermitize env, and perform eigen-decomposition
     # env = U D U'
-    D, U = eigh((env + env')/2)
+    D, U = eigh((env + env') / 2)
     # determine env is (mostly) positive or negative
-    sgn = sign(mean(vcat((diag(b) for (k,b) in blocks(D))...)))
+    sgn = sign(mean(vcat((diag(b) for (k, b) in blocks(D))...)))
     if sgn == -1
         D *= -1
     end
@@ -29,7 +29,7 @@ function positive_approx(env::AbstractTensorMap)
             end
         end
     end
-    Zdg = sdiag_pow(D, 1/2) * U'
+    Zdg = sdiag_pow(D, 1 / 2) * U'
     return sgn, Zdg
 end
 
@@ -37,9 +37,11 @@ end
 Fix local gauge of the env tensor around a bond
 """
 function fu_fixgauge(
-    Zdg::AbstractTensorMap, 
-    X::AbstractTensorMap, Y::AbstractTensorMap, 
-    aR::AbstractTensorMap, bL::AbstractTensorMap
+    Zdg::AbstractTensorMap,
+    X::AbstractTensorMap,
+    Y::AbstractTensorMap,
+    aR::AbstractTensorMap,
+    bL::AbstractTensorMap,
 )
     #= 
             1               1
@@ -50,8 +52,8 @@ function fu_fixgauge(
                                         ↑
                     =   2 → L → 1   3 → QL ← 2
     =#
-    QR, R = leftorth(Zdg, ((1,2), (3,)), alg=QRpos())
-    QL, L = leftorth(Zdg, ((1,3), (2,)), alg=QRpos())
+    QR, R = leftorth(Zdg, ((1, 2), (3,)); alg=QRpos())
+    QL, L = leftorth(Zdg, ((1, 3), (2,)); alg=QRpos())
     @assert !isdual(codomain(R)[1]) && !isdual(domain(R)[1])
     @assert !isdual(codomain(L)[1]) && !isdual(domain(L)[1])
     Rinv, Linv = inv(R), inv(L)
@@ -73,12 +75,9 @@ function fu_fixgauge(
                         ↑
         -2 → Linv → 1 → Z† ← 2 ← Rinv ← -3
     =#
-    aR = ncon([L, aR], [[-1,1], [1,-2,-3]])
-    bL = ncon([bL, R], [[-1,-2,1], [-3,1]])
-    Zdg = permute(
-        ncon([Zdg, Linv, Rinv], [[-1,1,2], [1,-2], [2,-3]]),
-        (1,), (2,3)
-    )
+    aR = ncon([L, aR], [[-1, 1], [1, -2, -3]])
+    bL = ncon([bL, R], [[-1, -2, 1], [-3, 1]])
+    Zdg = permute(ncon([Zdg, Linv, Rinv], [[-1, 1, 2], [1, -2], [2, -3]]), (1,), (2, 3))
     #= fix gauge of X, Y
             -1                                      -1
              |                                      |
@@ -86,7 +85,7 @@ function fu_fixgauge(
              |                                      |
             -3                                      -3
     =#
-    X = ncon([X, Linv], [[-1,1,-3,-4], [1,-2]])
-    Y = ncon([Y, Rinv], [[-1,-2,-3,1], [1,-4]])
+    X = ncon([X, Linv], [[-1, 1, -3, -4], [1, -2]])
+    Y = ncon([Y, Rinv], [[-1, -2, -3, 1], [1, -4]])
     return Zdg, X, Y, aR, bL
 end
diff --git a/src/algorithms/timeevol/fu_optimize.jl b/src/algorithms/timeevol/fu_optimize.jl
index a89f8d3e..e90d4d1a 100644
--- a/src/algorithms/timeevol/fu_optimize.jl
+++ b/src/algorithms/timeevol/fu_optimize.jl
@@ -23,13 +23,14 @@ which can be more simply denoted as
 The axes 1, 2 (or 3, 4) come from X†, Y† (or X, Y)
 """
 function tensor_env(
-    row::Int, col::Int, X::AbstractTensorMap, 
-    Y::AbstractTensorMap, envs::CTMRGEnv
+    row::Int, col::Int, X::AbstractTensorMap, Y::AbstractTensorMap, envs::CTMRGEnv
 )
-    Nr, Nc = size(envs.corners)[[2,3]]
-    cm1 = _prev(col, Nc);
-    cp1 = _next(col, Nc); cp2 = _next(cp1, Nc)
-    rm1 = _prev(row, Nr); rp1 = _next(row, Nr)
+    Nr, Nc = size(envs.corners)[[2, 3]]
+    cm1 = _prev(col, Nc)
+    cp1 = _next(col, Nc)
+    cp2 = _next(cp1, Nc)
+    rm1 = _prev(row, Nr)
+    rp1 = _next(row, Nr)
     c1 = envs.corners[1, rm1, cm1]
     c2 = envs.corners[2, rm1, cp2]
     c3 = envs.corners[3, rp1, cp2]
@@ -40,15 +41,23 @@ function tensor_env(
     t4 = envs.edges[4, row, cm1]
     # left half
     @autoopt @tensor lhalf[DX1, DX0, χ5, χ6] := (
-        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]
+        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]
     )
     # right half
     @autoopt @tensor rhalf[DY1, DY0, χ5, χ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]
+        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]
     )
     # combine
     @autoopt @tensor env[DX1, DY1; DX0, DY0] := (
@@ -57,7 +66,6 @@ function tensor_env(
     return env
 end
 
-
 """
 Construct the tensor
 ```
@@ -72,13 +80,11 @@ Construct the tensor
 """
 function tensor_Ra(env::AbstractTensorMap, bL::AbstractTensorMap)
     @autoopt @tensor Ra[DX1, Db1, DX0, Db0] := (
-        env[DX1, DY1, DX0, DY0] *
-        bL[Db0, db, DY0] * conj(bL[Db1, db, DY1])
+        env[DX1, DY1, DX0, DY0] * bL[Db0, db, DY0] * conj(bL[Db1, db, DY1])
     )
     return Ra
 end
 
-
 """
 Construct the tensor
 ```
@@ -92,17 +98,17 @@ Construct the tensor
 ```
 """
 function tensor_Sa(
-    env::AbstractTensorMap, aR2::AbstractTensorMap, 
-    bL::AbstractTensorMap, bL2::AbstractTensorMap
+    env::AbstractTensorMap,
+    aR2::AbstractTensorMap,
+    bL::AbstractTensorMap,
+    bL2::AbstractTensorMap,
 )
     @autoopt @tensor Sa[DX1, Db1, da] := (
-        env[DX1, DY1, DX0, DY0] * conj(bL[Db1, db, DY1]) *
-        bL2[D, db, DY0] * aR2[DX0, da, D]
+        env[DX1, DY1, DX0, DY0] * conj(bL[Db1, db, DY1]) * bL2[D, db, DY0] * aR2[DX0, da, D]
     )
     return Sa
 end
 
-
 """
 Construct the tensor
 ```
@@ -117,13 +123,11 @@ Construct the tensor
 """
 function tensor_Rb(env::AbstractTensorMap, aR::AbstractTensorMap)
     @autoopt @tensor Rb[Da1, DY1, Da0, DY0] := (
-        env[DX1, DY1, DX0, DY0] *
-        aR[DX0, da, Da0] * conj(aR[DX1, da, Da1])
+        env[DX1, DY1, DX0, DY0] * aR[DX0, da, Da0] * conj(aR[DX1, da, Da1])
     )
     return Rb
 end
 
-
 """
 Construct the tensor
 ```
@@ -137,17 +141,17 @@ Construct the tensor
 ```
 """
 function tensor_Sb(
-    env::AbstractTensorMap, aR::AbstractTensorMap, 
-    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+    env::AbstractTensorMap,
+    aR::AbstractTensorMap,
+    aR2::AbstractTensorMap,
+    bL2::AbstractTensorMap,
 )
     @autoopt @tensor Sb[Da1, DY1, db] := (
-        env[DX1, DY1, DX0, DY0] * conj(aR[DX1, da, Da1]) *
-        aR2[DX0, da, D] * bL2[D, db, DY0]
+        env[DX1, DY1, DX0, DY0] * conj(aR[DX1, da, Da1]) * aR2[DX0, da, D] * bL2[D, db, DY0]
     )
     return Sb
 end
 
-
 """
 Calculate the norm <Psi(a1,b1)|Psi(a2,b2)>
 ```
@@ -161,14 +165,18 @@ Calculate the norm <Psi(a1,b1)|Psi(a2,b2)>
 ```
 """
 function inner_prod(
-    env::AbstractTensorMap, 
-    aR1::AbstractTensorMap, bL1::AbstractTensorMap,
-    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+    env::AbstractTensorMap,
+    aR1::AbstractTensorMap,
+    bL1::AbstractTensorMap,
+    aR2::AbstractTensorMap,
+    bL2::AbstractTensorMap,
 )
     @autoopt @tensor t[:] := (
         env[DX1, DY1, DX0, DY0] *
-        conj(aR1[DX1, da, D1]) * conj(bL1[D1, db, DY1]) *
-        aR2[DX0, da, D0] * bL2[D0, db, DY0]
+        conj(aR1[DX1, da, D1]) *
+        conj(bL1[D1, db, DY1]) *
+        aR2[DX0, da, D0] *
+        bL2[D0, db, DY0]
     )
     return first(blocks(t))[2][1]
 end
@@ -182,9 +190,11 @@ Calculate the cost function
 ```
 """
 function cost_func(
-    env::AbstractTensorMap, 
-    aR::AbstractTensorMap, bL::AbstractTensorMap,
-    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+    env::AbstractTensorMap,
+    aR::AbstractTensorMap,
+    bL::AbstractTensorMap,
+    aR2::AbstractTensorMap,
+    bL2::AbstractTensorMap,
 )
     t1 = inner_prod(env, aR, bL, aR, bL)
     t2 = inner_prod(env, aR2, bL2, aR2, bL2)
@@ -192,7 +202,6 @@ function cost_func(
     return real(t1) + real(t2) - 2 * real(t3)
 end
 
-
 """
 Calculate the approximate local inner product
 `<aR1 bL1|aR2 bL2>`
@@ -205,12 +214,13 @@ Calculate the approximate local inner product
 ```
 """
 function inner_prod_local(
-    aR1::AbstractTensorMap, bL1::AbstractTensorMap,
-    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+    aR1::AbstractTensorMap,
+    bL1::AbstractTensorMap,
+    aR2::AbstractTensorMap,
+    bL2::AbstractTensorMap,
 )
     @autoopt @tensor t[:] := (
-        conj(aR1[DW, da, D1]) * conj(bL1[D1, db, DE]) *
-        aR2[DW, da, D0] * bL2[D0, db, DE]
+        conj(aR1[DW, da, D1]) * conj(bL1[D1, db, DE]) * aR2[DW, da, D0] * bL2[D0, db, DE]
     )
     return first(blocks(t))[2][1]
 end
@@ -225,13 +235,15 @@ between two evolution steps
 ```
 """
 function local_fidelity(
-    aR1::AbstractTensorMap, bL1::AbstractTensorMap, 
-    aR2::AbstractTensorMap, bL2::AbstractTensorMap
+    aR1::AbstractTensorMap,
+    bL1::AbstractTensorMap,
+    aR2::AbstractTensorMap,
+    bL2::AbstractTensorMap,
 )
     b12 = inner_prod_local(aR1, bL1, aR2, bL2)
     b11 = inner_prod_local(aR1, bL1, aR1, bL1)
     b22 = inner_prod_local(aR2, bL2, aR2, bL2)
-    return abs(b12) / sqrt(abs(b11*b22))
+    return abs(b12) / sqrt(abs(b11 * b22))
 end
 
 """
@@ -240,13 +252,10 @@ Solving the equations
     Ra aR = Sa, Rb bL = Sb
 ```
 """
-function solve_ab(
-    R::AbstractTensorMap, S::AbstractTensorMap, 
-    ab0::AbstractTensorMap
-)
-    f(x) = ncon((R, x), ([-1,-2,1,2], [1,2,-3]))
-    ab, info = linsolve(f, S, permute(ab0, (1,3,2)), 0, 1)
-    return permute(ab, (1,3,2)), info
+function solve_ab(R::AbstractTensorMap, S::AbstractTensorMap, ab0::AbstractTensorMap)
+    f(x) = ncon((R, x), ([-1, -2, 1, 2], [1, 2, -3]))
+    ab, info = linsolve(f, S, permute(ab0, (1, 3, 2)), 0, 1)
+    return permute(ab, (1, 3, 2)), info
 end
 
 """
@@ -263,18 +272,19 @@ Minimize the cost function
 `aR0`, `bL0` are initial values of `aR`, `bL`
 """
 function fu_optimize(
-    aR0::AbstractTensorMap, bL0::AbstractTensorMap, 
-    aR2::AbstractTensorMap, bL2::AbstractTensorMap, 
+    aR0::AbstractTensorMap,
+    bL0::AbstractTensorMap,
+    aR2::AbstractTensorMap,
+    bL2::AbstractTensorMap,
     env::AbstractTensorMap;
-    maxiter::Int=50, maxdiff::Float64=1e-15, 
-    check_int::Int=1, verbose::Bool=false
+    maxiter::Int=50,
+    maxdiff::Float64=1e-15,
+    check_int::Int=1,
+    verbose::Bool=false,
 )
     if verbose
         println("---- Iterative optimization ----")
-        @printf(
-            "%-6s%12s%12s%12s %10s\n", 
-            "Step", "Cost", "ϵ_d", "ϵ_ab", "Time/s"
-        )
+        @printf("%-6s%12s%12s%12s %10s\n", "Step", "Cost", "ϵ_d", "ϵ_ab", "Time/s")
     end
     aR, bL = deepcopy(aR0), deepcopy(bL0)
     time0 = time()
@@ -285,10 +295,11 @@ function fu_optimize(
     if abs(cost0) < 5e-15
         if verbose
             time1 = time()
-            println(@sprintf(
-                "%-6d%12.3e%12.3e%12.3e %10.3f\n", 
-                0, cost0, NaN, NaN, time1 - time0
-            ))
+            println(
+                @sprintf(
+                    "%-6d%12.3e%12.3e%12.3e %10.3f\n", 0, cost0, NaN, NaN, time1 - time0
+                )
+            )
         end
         return aR, bL, cost0
     end
@@ -307,8 +318,12 @@ function fu_optimize(
         time1 = time()
         if verbose && (count == 1 || count % check_int == 0)
             @printf(
-                "%-6d%12.3e%12.3e%12.3e %10.3f\n", 
-                count, cost, diff_d, diff_ab, time1 - time0
+                "%-6d%12.3e%12.3e%12.3e %10.3f\n",
+                count,
+                cost,
+                diff_d,
+                diff_ab,
+                time1 - time0
             )
         end
         if diff_ab < maxdiff
@@ -322,4 +337,3 @@ function fu_optimize(
     end
     return aR, bL, cost0
 end
-
diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/timeevol/fullupdate.jl
index fdede082..e0a16efd 100644
--- a/src/algorithms/timeevol/fullupdate.jl
+++ b/src/algorithms/timeevol/fullupdate.jl
@@ -14,13 +14,11 @@ CTMRG left-move to update CTMRGEnv in the c-th column
 ```
 """
 function ctmrg_leftmove!(
-    col::Int, peps::InfinitePEPS, envs::CTMRGEnv,
-    chi::Int, svderr::Float64=1e-9
+    col::Int, peps::InfinitePEPS, envs::CTMRGEnv, chi::Int, svderr::Float64=1e-9
 )
     trscheme = truncerr(svderr) & truncdim(chi)
-    alg = CTMRG(
-        verbosity=0, miniter=1, maxiter=10, 
-        trscheme=trscheme, ctmrgscheme=:sequential
+    alg = CTMRG(;
+        verbosity=0, miniter=1, maxiter=10, trscheme=trscheme, ctmrgscheme=:sequential
     )
     envs2, info = ctmrg_leftmove(col, peps, envs, alg)
     envs.corners[:, :, col] = envs2.corners[:, :, col]
@@ -41,8 +39,7 @@ CTMRG right-move to update CTMRGEnv in the c-th column
 ```
 """
 function ctmrg_rightmove!(
-    col::Int, peps::InfinitePEPS, envs::CTMRGEnv,
-    chi::Int, svderr::Float64=1e-9
+    col::Int, peps::InfinitePEPS, envs::CTMRGEnv, chi::Int, svderr::Float64=1e-9
 )
     Nr, Nc = size(peps)
     @assert 1 <= col <= Nc
@@ -58,11 +55,16 @@ Update all horizontal bonds in the c-th column
 To update rows, rotate the network clockwise by 90 degrees.
 """
 function update_column!(
-    col::Int, gate::AbstractTensorMap,
-    peps::InfinitePEPS, envs::CTMRGEnv,
-    Dcut::Int, chi::Int;
-    svderr::Float64=1e-9, maxiter::Int=50,
-    maxdiff::Float64=1e-15, gaugefix::Bool=true,
+    col::Int,
+    gate::AbstractTensorMap,
+    peps::InfinitePEPS,
+    envs::CTMRGEnv,
+    Dcut::Int,
+    chi::Int;
+    svderr::Float64=1e-9,
+    maxiter::Int=50,
+    maxdiff::Float64=1e-15,
+    gaugefix::Bool=true,
 )
     Nr, Nc = size(peps)
     @assert 1 <= col <= Nc
@@ -110,7 +112,7 @@ function update_column!(
         end
         env = sgn * Zdg' * Zdg
         #= apply gate
-        
+
                 -2          -3
                 ↑           ↑
                 |----gate---|
@@ -134,7 +136,7 @@ function update_column!(
         bL /= maxabs(bL)
         localfid += local_fidelity(aR, bL, aR0, bL0)
         #= update and normalize peps, ms
-        
+
                 -2        -1               -1     -2
                 |        ↗                ↗       |
             -5- X ← 1 ← aR ← -3     -5 ← bL → 1 → Y - -3
@@ -167,24 +169,24 @@ Otherwise, use full-infinite environment instead.
 Reference: Physical Review B 92, 035142 (2015)
 """
 function fullupdate!(
-    gate::AbstractTensorMap, peps::InfinitePEPS, envs::CTMRGEnv,
-    Dcut::Int, chi::Int, svderr::Float64=1e-9
+    gate::AbstractTensorMap,
+    peps::InfinitePEPS,
+    envs::CTMRGEnv,
+    Dcut::Int,
+    chi::Int,
+    svderr::Float64=1e-9,
 )
     Nr, Nc = size(peps)
     fid, maxcost = 0.0, 0.0
     for col in 1:Nc
-        tmpfid, costs = update_column!(
-            col, gate, peps, envs, Dcut, chi; svderr=svderr
-        )
+        tmpfid, costs = update_column!(col, gate, peps, envs, Dcut, chi; svderr=svderr)
         fid += tmpfid
         maxcost = max(maxcost, maximum(costs))
     end
     rotr90!(peps)
     rotr90!(envs)
     for row in 1:Nr
-        tmpfid, costs = update_column!(
-            row, gate, peps, envs, Dcut, chi; svderr=svderr
-        )
+        tmpfid, costs = update_column!(row, gate, peps, envs, Dcut, chi; svderr=svderr)
         fid += tmpfid
         maxcost = max(maxcost, maximum(costs))
     end
diff --git a/src/algorithms/timeevol/simpleupdate.jl b/src/algorithms/timeevol/simpleupdate.jl
index 2e464cde..a57795be 100644
--- a/src/algorithms/timeevol/simpleupdate.jl
+++ b/src/algorithms/timeevol/simpleupdate.jl
@@ -4,7 +4,7 @@ Mirror the unit cell of an iPEPS by its anti-diagonal line
 function mirror_antidiag!(peps::InfinitePEPS)
     peps.A[:] = mirror_antidiag(peps.A)
     for (i, t) in enumerate(peps.A)
-        peps.A[i] = permute(t, (1,), (3,2,5,4))
+        peps.A[i] = permute(t, (1,), (3, 2, 5, 4))
     end
 end
 
@@ -12,7 +12,7 @@ end
 Mirror the unit cell of an iPEPS with weights by its anti-diagonal line
 """
 function mirror_antidiag!(wts::SUWeight)
-    wts.x[:], wts.y[:] = mirror_antidiag(wts.y), mirror_antidiag(wts.x)
+    return wts.x[:], wts.y[:] = mirror_antidiag(wts.y), mirror_antidiag(wts.x)
 end
 
 """
@@ -30,33 +30,36 @@ Weights around the tensor at `(row, col)` are
 ```
 """
 function absorb_wt(
-    t::AbstractTensorMap, row::Int, col::Int, 
-    ax::Int, wts::SUWeight; 
-    sqrtwt::Bool=false, invwt::Bool=false
+    t::AbstractTensorMap,
+    row::Int,
+    col::Int,
+    ax::Int,
+    wts::SUWeight;
+    sqrtwt::Bool=false,
+    invwt::Bool=false,
 )
     Nr, Nc = size(wts)
     @assert 1 <= row <= Nr && 1 <= col <= Nc
     @assert 2 <= ax <= 5
-    pow = (sqrtwt ? 1/2 : 1) * (invwt ? -1 : 1)
+    pow = (sqrtwt ? 1 / 2 : 1) * (invwt ? -1 : 1)
     if ax == 2 # north
         wt = wts.y[row, col]
     elseif ax == 3 # east
         wt = wts.x[row, col]
     elseif ax == 4 # south
-        wt = wts.y[_next(row,Nr), col]
+        wt = wts.y[_next(row, Nr), col]
     else # west
-        wt = wts.x[row, _prev(col,Nc)]
+        wt = wts.x[row, _prev(col, Nc)]
     end
     wt2 = sdiag_pow(wt, pow)
     indices_t = collect(-1:-1:-5)
     indices_t[ax] = 1
-    indices_wt = (ax in (2,3) ? [1,-ax] : [-ax,1])
+    indices_wt = (ax in (2, 3) ? [1, -ax] : [-ax, 1])
     t2 = ncon((t, wt2), (indices_t, indices_wt))
     t2 = permute(t2, (1,), Tuple(2:5))
     return t2
 end
 
-
 """
 Simple update of bond `wts.x[r,c]`
 ```
@@ -68,26 +71,30 @@ Simple update of bond `wts.x[r,c]`
 ```
 """
 function _su_bondx!(
-    row::Int, col::Int, gate::AbstractTensorMap, 
-    peps::InfinitePEPS, wts::SUWeight,
-    Dcut::Int, svderr::Float64=1e-10
+    row::Int,
+    col::Int,
+    gate::AbstractTensorMap,
+    peps::InfinitePEPS,
+    wts::SUWeight,
+    Dcut::Int,
+    svderr::Float64=1e-10,
 )
     Nr, Nc = size(peps)
     @assert 1 <= row <= Nr && 1 <= col <= Nc
-    row2, col2 = row, _next(col,Nc)
-    T1, T2 = peps[row,col], peps[row2,col2]
+    row2, col2 = row, _next(col, Nc)
+    T1, T2 = peps[row, col], peps[row2, col2]
     # absorb environment weights
-    for ax in (2,4,5)
+    for ax in (2, 4, 5)
         T1 = absorb_wt(T1, row, col, ax, wts)
     end
-    for ax in (2,3,4)
+    for ax in (2, 3, 4)
         T2 = absorb_wt(T2, row2, col2, ax, wts)
     end
     # absorb bond weight
     T1 = absorb_wt(T1, row, col, 3, wts; sqrtwt=true)
     T2 = absorb_wt(T2, row2, col2, 5, wts; sqrtwt=true)
     #= QR and LQ decomposition
-    
+
         2   1               1             2
         ↓ ↗                 ↓            ↗
     5 ← T ← 3   ====>   3 ← X ← 4 ← 1 ← aR ← 3
@@ -100,10 +107,10 @@ function _su_bondx!(
         ↓                               ↓
         4                               4
     =#
-    X, aR = leftorth(T1, ((2,4,5), (1,3)), alg=QRpos())
-    bL, Y = rightorth(T2, ((5,1), (2,3,4)), alg=LQpos())
+    X, aR = leftorth(T1, ((2, 4, 5), (1, 3)); alg=QRpos())
+    bL, Y = rightorth(T2, ((5, 1), (2, 3, 4)); alg=LQpos())
     #= apply gate
-    
+
             -2          -3
             ↑           ↑
             |----gate---|
@@ -112,10 +119,10 @@ function _su_bondx!(
             ↑           ↑
         -1← aR -← 3 -← bL ← -4
     =#
-    tmp = ncon((gate, aR, bL), ([-2,-3,1,2], [-1,1,3], [3,2,-4]))
+    tmp = ncon((gate, aR, bL), ([-2, -3, 1, 2], [-1, 1, 3], [3, 2, -4]))
     # SVD
     truncscheme = truncerr(svderr) & truncdim(Dcut)
-    aR, s, bL, ϵ = tsvd(tmp, ((1,2), (3,4)); trunc=truncscheme)
+    aR, s, bL, ϵ = tsvd(tmp, ((1, 2), (3, 4)); trunc=truncscheme)
     #=
             -2         -1              -1    -2
             |         ↗               ↗       |
@@ -123,23 +130,22 @@ function _su_bondx!(
             |                                 |
             -4                               -4
     =#
-    T1 = ncon((X, aR), ([-2,-4,-5,1], [1,-1,-3]))
-    T2 = ncon((bL, Y), ([-5,-1,1], [1,-2,-3,-4]))
+    T1 = ncon((X, aR), ([-2, -4, -5, 1], [1, -1, -3]))
+    T2 = ncon((bL, Y), ([-5, -1, 1], [1, -2, -3, -4]))
     # remove environment weights
-    for ax in (2,4,5)
+    for ax in (2, 4, 5)
         T1 = absorb_wt(T1, row, col, ax, wts; invwt=true)
     end
-    for ax in (2,3,4)
+    for ax in (2, 3, 4)
         T2 = absorb_wt(T2, row2, col2, ax, wts; invwt=true)
     end
     # update tensor dict and weight on current bond 
     # (max element of weight is normalized to 1)
-    peps.A[row,col], peps.A[row2,col2] = T1, T2
-    wts.x[row,col] = s / maxabs(s)
+    peps.A[row, col], peps.A[row2, col2] = T1, T2
+    wts.x[row, col] = s / maxabs(s)
     return ϵ
 end
 
-
 """
 One round of simple update on the input InfinitePEPS `peps` 
 and SUWeight `wts` with the nearest neighbor gate `gate`
@@ -148,8 +154,12 @@ When `bipartite === true` (for square lattice), the unit cell size should be 2 x
 and the tensor and x/y weight at `(row, col)` is the same as `(row+1, col+1)`
 """
 function simpleupdate!(
-    gate::AbstractTensorMap, peps::InfinitePEPS, wts::SUWeight, 
-    Dcut::Int, svderr::Float64=1e-10; bipartite::Bool=false,
+    gate::AbstractTensorMap,
+    peps::InfinitePEPS,
+    wts::SUWeight,
+    Dcut::Int,
+    svderr::Float64=1e-10;
+    bipartite::Bool=false,
 )
     Nr, Nc = size(peps)
     if bipartite
@@ -157,21 +167,24 @@ function simpleupdate!(
     end
     # TODO: make algorithm independent on the choice of dual in the network
     for (r, c) in Iterators.product(1:Nr, 1:Nc)
-        @assert [isdual(space(peps.A[r, c], ax)) for ax in 1:5] == [0,1,1,0,0]
-        @assert [isdual(space(wts.x[r, c], ax)) for ax in 1:2] == [0,1]
-        @assert [isdual(space(wts.y[r, c], ax)) for ax in 1:2] == [0,1]
+        @assert [isdual(space(peps.A[r, c], ax)) for ax in 1:5] == [0, 1, 1, 0, 0]
+        @assert [isdual(space(wts.x[r, c], ax)) for ax in 1:2] == [0, 1]
+        @assert [isdual(space(wts.y[r, c], ax)) for ax in 1:2] == [0, 1]
     end
     for direction in 1:2
         # mirror the y-weights to x-direction 
         # to update them using code for x-weights
         if direction == 2
-            mirror_antidiag!(peps); mirror_antidiag!(wts)
+            mirror_antidiag!(peps)
+            mirror_antidiag!(wts)
         end
         if bipartite
             ϵ = _su_bondx!(1, 1, gate, peps, wts, Dcut, svderr)
-            (peps.A[2,2], peps.A[2,1], wts.x[2,2]) = deepcopy.((peps.A[1,1], peps.A[1,2], wts.x[1,1]))
+            (peps.A[2, 2], peps.A[2, 1], wts.x[2, 2]) =
+                deepcopy.((peps.A[1, 1], peps.A[1, 2], wts.x[1, 1]))
             ϵ = _su_bondx!(2, 1, gate, peps, wts, Dcut, svderr)
-            (peps.A[1,2], peps.A[1,1], wts.x[1,2]) = deepcopy.((peps.A[2,1], peps.A[2,2], wts.x[2,1]))
+            (peps.A[1, 2], peps.A[1, 1], wts.x[1, 2]) =
+                deepcopy.((peps.A[2, 1], peps.A[2, 2], wts.x[2, 1]))
         else
             for site in CartesianIndices(peps.A)
                 row, col = Tuple(site)
@@ -179,9 +192,9 @@ function simpleupdate!(
             end
         end
         if direction == 2
-            mirror_antidiag!(peps); mirror_antidiag!(wts)
+            mirror_antidiag!(peps)
+            mirror_antidiag!(wts)
         end
     end
     return nothing
 end
-
diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl
index 827244e9..423c6f08 100644
--- a/src/states/infinitepeps.jl
+++ b/src/states/infinitepeps.jl
@@ -175,15 +175,17 @@ Base.rot180(t::InfinitePEPS) = InfinitePEPS(rot180(rot180.(t.A)))
 # In-place rotations
 function rotl90!(peps::InfinitePEPS)
     peps.A[:] = rotl90(rotl90.(peps.A))
+    return nothing
 end
 function rotr90!(peps::InfinitePEPS)
     peps.A[:] = rotr90(rotr90.(peps.A))
+    return nothing
 end
 function rot180!(peps::InfinitePEPS)
     peps.A[:] = rot180(rot180.(peps.A))
+    return nothing
 end
 
-
 # Chainrules
 function ChainRulesCore.rrule(
     ::typeof(Base.getindex), state::InfinitePEPS, row::Int, col::Int
diff --git a/src/states/suweight.jl b/src/states/suweight.jl
index 97b75832..19421d33 100644
--- a/src/states/suweight.jl
+++ b/src/states/suweight.jl
@@ -6,7 +6,7 @@ struct SUWeight{T<:AbstractTensorMap}
     y::Matrix{T}
 
     function SUWeight(wxs::Matrix{T}, wys::Matrix{T}) where {T}
-        new{T}(wxs, wys)
+        return new{T}(wxs, wys)
     end
 end
 
@@ -30,9 +30,9 @@ end
 function Base.iterate(wts::SUWeight, state=1)
     nx = prod(size(wts.x))
     if 1 <= state <= nx
-        return wts.x[state], state+1
-    elseif nx+1 <= state <= 2*nx
-        return wts.y[state-nx], state+1
+        return wts.x[state], state + 1
+    elseif nx + 1 <= state <= 2 * nx
+        return wts.y[state - nx], state + 1
     else
         return nothing
     end
diff --git a/src/utility/mirror.jl b/src/utility/mirror.jl
index f9bc74ae..8028d0b3 100644
--- a/src/utility/mirror.jl
+++ b/src/utility/mirror.jl
@@ -7,8 +7,5 @@ i.e. the element now at [r, c] was originally at [Nr-c+1, Nc-r+1]
 """
 function mirror_antidiag(arr::AbstractMatrix)
     Nr, Nc = size(arr)
-    return collect(
-        arr[Nr-c+1, Nc-r+1]
-        for (r, c) in Iterators.product(1:Nc, 1:Nr)
-    )
+    return collect(arr[Nr - c + 1, Nc - r + 1] for (r, c) in Iterators.product(1:Nc, 1:Nr))
 end
diff --git a/src/utility/svd.jl b/src/utility/svd.jl
index 20950ac9..7d974109 100644
--- a/src/utility/svd.jl
+++ b/src/utility/svd.jl
@@ -292,7 +292,6 @@ function _lorentz_broaden(x::Real, ε=1e-12)
     return x′ / (x′^2 + ε)
 end
 
-
 """
 Given `tsvd` result `u`, `s` and `vh`, 
 absorb singular values `s` into `u` and `vh` by
diff --git a/src/utility/util.jl b/src/utility/util.jl
index 181eb3a9..5b61dd2b 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -41,7 +41,7 @@ Compute S^(pow) for diagonal matrices `S`
 function sdiag_pow(S::AbstractTensorMap, pow::Real)
     S2 = similar(S)
     for (k, b) in blocks(S)
-        copyto!(blocks(S2)[k], diagm(diag(b).^pow))
+        copyto!(blocks(S2)[k], diagm(diag(b) .^ pow))
     end
     return S2
 end

From 043a7feb5d6f6d6aa63c6dac51b6d66142044791 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Thu, 7 Nov 2024 11:00:33 +0800
Subject: [PATCH 09/20] add TODO for svd initialization of ALS optimization

---
 src/algorithms/timeevol/fullupdate.jl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/timeevol/fullupdate.jl
index e0a16efd..b63edc28 100644
--- a/src/algorithms/timeevol/fullupdate.jl
+++ b/src/algorithms/timeevol/fullupdate.jl
@@ -123,6 +123,7 @@ function update_column!(
         =#
         tmp = ncon((gate, aR0, bL0), ([-2, -3, 1, 2], [-1, 1, 3], [3, 2, -4]))
         # initialize truncated tensors using simple SVD truncation
+        # TODO: return truncated and untruncated SVD result at once, without repeated calculation
         aR2, s, bL2, ϵ = tsvd(tmp, ((1, 2), (3, 4)); trunc=truncerr(1e-15))
         aR, s_cut, bL, ϵ = tsvd(tmp, ((1, 2), (3, 4)); trunc=truncscheme)
         aR2, bL2 = absorb_s(aR2, s, bL2)

From 09270827af2a69b8ead12c2dc971cc223ec28dbf Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Thu, 7 Nov 2024 16:30:32 +0800
Subject: [PATCH 10/20] remove a redundant SVD in full update

---
 src/algorithms/timeevol/fu_optimize.jl | 143 +++++++++++--------------
 src/algorithms/timeevol/fullupdate.jl  |  13 +--
 2 files changed, 69 insertions(+), 87 deletions(-)

diff --git a/src/algorithms/timeevol/fu_optimize.jl b/src/algorithms/timeevol/fu_optimize.jl
index e90d4d1a..de52d3b5 100644
--- a/src/algorithms/timeevol/fu_optimize.jl
+++ b/src/algorithms/timeevol/fu_optimize.jl
@@ -88,23 +88,18 @@ end
 """
 Construct the tensor
 ```
-    |--------------env--------------|
-    |→ DX1         Db1 → bL† ← DY1 ←|
-    |                    ↑          |
-    |         da         db         |
-    |         ↑          ↑          |
-    |← DX0 ←- aR2 ←- D ← bL2 → DY0 →|
-    |-------------------------------|
+    |-----------env-----------|
+    |→ DX1   Db1 → bL† ← DY1 ←|
+    |               ↑         |
+    |         da    db        |
+    |         ↑     ↑         |
+    |← DX0 ←- aR2 bL2 -→ DY0 →|
+    |-------------------------|
 ```
 """
-function tensor_Sa(
-    env::AbstractTensorMap,
-    aR2::AbstractTensorMap,
-    bL::AbstractTensorMap,
-    bL2::AbstractTensorMap,
-)
+function tensor_Sa(env::AbstractTensorMap, bL::AbstractTensorMap, aR2bL2::AbstractTensorMap)
     @autoopt @tensor Sa[DX1, Db1, da] := (
-        env[DX1, DY1, DX0, DY0] * conj(bL[Db1, db, DY1]) * bL2[D, db, DY0] * aR2[DX0, da, D]
+        env[DX1, DY1, DX0, DY0] * conj(bL[Db1, db, DY1]) * aR2bL2[DX0, da, db, DY0]
     )
     return Sa
 end
@@ -131,23 +126,18 @@ end
 """
 Construct the tensor
 ```
-    |--------------env--------------|
-    |→ DX1 → aR† → Da1         DY1 ←|
-    |        ↑                      |
-    |        da         db          |
-    |        ↑          ↑           |
-    |← DX0 ← aR2 ← D ←- bL2 -→ DY0 →|
-    |-------------------------------|
+    |-----------env-----------|
+    |→ DX1 → aR† → Da1   DY1 ←|
+    |         ↑               |
+    |         da   db         |
+    |         ↑     ↑         |
+    |← DX0 ←- aR2 bL2 -→ DY0 →|
+    |-------------------------|
 ```
 """
-function tensor_Sb(
-    env::AbstractTensorMap,
-    aR::AbstractTensorMap,
-    aR2::AbstractTensorMap,
-    bL2::AbstractTensorMap,
-)
+function tensor_Sb(env::AbstractTensorMap, aR::AbstractTensorMap, aR2bL2::AbstractTensorMap)
     @autoopt @tensor Sb[Da1, DY1, db] := (
-        env[DX1, DY1, DX0, DY0] * conj(aR[DX1, da, Da1]) * aR2[DX0, da, D] * bL2[D, db, DY0]
+        env[DX1, DY1, DX0, DY0] * conj(aR[DX1, da, Da1]) * aR2bL2[DX0, da, db, DY0]
     )
     return Sb
 end
@@ -155,32 +145,37 @@ end
 """
 Calculate the norm <Psi(a1,b1)|Psi(a2,b2)>
 ```
-    |--------------env--------------|
-    |→ DX1 → aR1†→ D1 → bL1† ← DY1 ←|
-    |        ↑           ↑          |
-    |        da          db         |
-    |        ↑           ↑          |
-    |← DX0 ← aR2 ← D0 ← bL2 → DY0 -→|
-    |-------------------------------|
+    |----------env----------|
+    |→ DX1 → aR1bL1† ← DY1 ←|
+    |        ↑    ↑         |
+    |        da   db        |
+    |        ↑    ↑         |
+    |← DX0 ← aR2bL2 → DY0 -→|
+    |-----------------------|
 ```
 """
 function inner_prod(
-    env::AbstractTensorMap,
-    aR1::AbstractTensorMap,
-    bL1::AbstractTensorMap,
-    aR2::AbstractTensorMap,
-    bL2::AbstractTensorMap,
+    env::AbstractTensorMap, aR1bL1::AbstractTensorMap, aR2bL2::AbstractTensorMap
 )
     @autoopt @tensor t[:] := (
-        env[DX1, DY1, DX0, DY0] *
-        conj(aR1[DX1, da, D1]) *
-        conj(bL1[D1, db, DY1]) *
-        aR2[DX0, da, D0] *
-        bL2[D0, db, DY0]
+        env[DX1, DY1, DX0, DY0] * conj(aR1bL1[DX1, da, db, DY1]) * aR2bL2[DX0, da, db, DY0]
     )
     return first(blocks(t))[2][1]
 end
 
+"""
+Contract the axis between `aR` and `bL` tensors
+"""
+function _combine_aRbL(aR::AbstractTensorMap, bL::AbstractTensorMap)
+    #= 
+            da      db
+            ↑       ↑
+    ← DX ← aR ← D ← bL → DY →
+    =#
+    @tensor aRbL[DX, da, db, DY] := aR[DX, da, D] * bL[D, db, DY]
+    return aRbL
+end
+
 """
 Calculate the cost function
 ```
@@ -193,12 +188,12 @@ function cost_func(
     env::AbstractTensorMap,
     aR::AbstractTensorMap,
     bL::AbstractTensorMap,
-    aR2::AbstractTensorMap,
-    bL2::AbstractTensorMap,
+    aR2bL2::AbstractTensorMap,
 )
-    t1 = inner_prod(env, aR, bL, aR, bL)
-    t2 = inner_prod(env, aR2, bL2, aR2, bL2)
-    t3 = inner_prod(env, aR, bL, aR2, bL2)
+    aRbL = _combine_aRbL(aR, bL)
+    t1 = inner_prod(env, aRbL, aRbL)
+    t2 = inner_prod(env, aR2bL2, aR2bL2)
+    t3 = inner_prod(env, aRbL, aR2bL2)
     return real(t1) + real(t2) - 2 * real(t3)
 end
 
@@ -206,22 +201,15 @@ end
 Calculate the approximate local inner product
 `<aR1 bL1|aR2 bL2>`
 ```
-    |→ aR1† → D1 → bL1† ←|
-    |   ↑          ↑     |
-    DW  da         db    DE
-    |   ↑          ↑     |
-    |← aR2 ←- D0 ← bL2 -→|
+    |→ aR1bL1† ←|
+    |   ↑   ↑   |
+    DW  da  db  DE
+    |   ↑   ↑   |
+    |← aR2 bL2 →|
 ```
 """
-function inner_prod_local(
-    aR1::AbstractTensorMap,
-    bL1::AbstractTensorMap,
-    aR2::AbstractTensorMap,
-    bL2::AbstractTensorMap,
-)
-    @autoopt @tensor t[:] := (
-        conj(aR1[DW, da, D1]) * conj(bL1[D1, db, DE]) * aR2[DW, da, D0] * bL2[D0, db, DE]
-    )
+function inner_prod_local(aR1bL1::AbstractTensorMap, aR2bL2::AbstractTensorMap)
+    @autoopt @tensor t[:] := (conj(aR1bL1[DW, da, db, DE]) * aR2bL2[DW, da, db, DE])
     return first(blocks(t))[2][1]
 end
 
@@ -235,14 +223,12 @@ between two evolution steps
 ```
 """
 function local_fidelity(
-    aR1::AbstractTensorMap,
-    bL1::AbstractTensorMap,
-    aR2::AbstractTensorMap,
-    bL2::AbstractTensorMap,
+    aR1::AbstractTensorMap, bL1::AbstractTensorMap, aR2bL2::AbstractTensorMap
 )
-    b12 = inner_prod_local(aR1, bL1, aR2, bL2)
-    b11 = inner_prod_local(aR1, bL1, aR1, bL1)
-    b22 = inner_prod_local(aR2, bL2, aR2, bL2)
+    aR1bL1 = _combine_aRbL(aR1, bL1)
+    b12 = inner_prod_local(aR1bL1, aR2bL2)
+    b11 = inner_prod_local(aR1bL1, aR1bL1)
+    b22 = inner_prod_local(aR2bL2, aR2bL2)
     return abs(b12) / sqrt(abs(b11 * b22))
 end
 
@@ -274,8 +260,7 @@ Minimize the cost function
 function fu_optimize(
     aR0::AbstractTensorMap,
     bL0::AbstractTensorMap,
-    aR2::AbstractTensorMap,
-    bL2::AbstractTensorMap,
+    aR2bL2::AbstractTensorMap,
     env::AbstractTensorMap;
     maxiter::Int=50,
     maxdiff::Float64=1e-15,
@@ -288,8 +273,8 @@ function fu_optimize(
     end
     aR, bL = deepcopy(aR0), deepcopy(bL0)
     time0 = time()
-    cost00 = cost_func(env, aR, bL, aR2, bL2)
-    fid00 = local_fidelity(aR, bL, aR2, bL2)
+    cost00 = cost_func(env, aR, bL, aR2bL2)
+    fid00 = local_fidelity(aR, bL, aR2bL2)
     cost0, fid0 = cost00, fid00
     # no need to further optimize
     if abs(cost0) < 5e-15
@@ -306,13 +291,13 @@ function fu_optimize(
     for count in 1:maxiter
         time0 = time()
         Ra = tensor_Ra(env, bL)
-        Sa = tensor_Sa(env, aR2, bL, bL2)
+        Sa = tensor_Sa(env, bL, aR2bL2)
         aR, info_a = solve_ab(Ra, Sa, aR)
         Rb = tensor_Rb(env, aR)
-        Sb = tensor_Sb(env, aR, aR2, bL2)
+        Sb = tensor_Sb(env, aR, aR2bL2)
         bL, info_b = solve_ab(Rb, Sb, bL)
-        cost = cost_func(env, aR, bL, aR2, bL2)
-        fid = local_fidelity(aR, bL, aR2, bL2)
+        cost = cost_func(env, aR, bL, aR2bL2)
+        fid = local_fidelity(aR, bL, aR2bL2)
         diff_d = abs(cost - cost0) / cost00
         diff_ab = abs(fid - fid0) / fid00
         time1 = time()
diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/timeevol/fullupdate.jl
index b63edc28..0b7ee9ad 100644
--- a/src/algorithms/timeevol/fullupdate.jl
+++ b/src/algorithms/timeevol/fullupdate.jl
@@ -121,21 +121,18 @@ function update_column!(
                 ↑           ↑
             -1← aR -← 3 -← bL → -4
         =#
-        tmp = ncon((gate, aR0, bL0), ([-2, -3, 1, 2], [-1, 1, 3], [3, 2, -4]))
-        # initialize truncated tensors using simple SVD truncation
-        # TODO: return truncated and untruncated SVD result at once, without repeated calculation
-        aR2, s, bL2, ϵ = tsvd(tmp, ((1, 2), (3, 4)); trunc=truncerr(1e-15))
-        aR, s_cut, bL, ϵ = tsvd(tmp, ((1, 2), (3, 4)); trunc=truncscheme)
-        aR2, bL2 = absorb_s(aR2, s, bL2)
+        aR2bL2 = ncon((gate, aR0, bL0), ([-2, -3, 1, 2], [-1, 1, 3], [3, 2, -4]))
+        # initialize truncated tensors using SVD truncation
+        aR, s_cut, bL, ϵ = tsvd(aR2bL2, ((1, 2), (3, 4)); trunc=truncscheme)
         aR, bL = absorb_s(aR, s_cut, bL)
         # optimize aR, bL
         aR, bL, cost = fu_optimize(
-            aR, bL, aR2, bL2, env; maxiter=maxiter, maxdiff=maxdiff, verbose=false
+            aR, bL, aR2bL2, env; maxiter=maxiter, maxdiff=maxdiff, verbose=false
         )
         costs[row] = cost
         aR /= maxabs(aR)
         bL /= maxabs(bL)
-        localfid += local_fidelity(aR, bL, aR0, bL0)
+        localfid += local_fidelity(aR, bL, _combine_aRbL(aR0, bL0))
         #= update and normalize peps, ms
 
                 -2        -1               -1     -2

From 6735adeb33badfcee4d2274a12ff2f56bf20ebf7 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sun, 10 Nov 2024 16:37:03 +0800
Subject: [PATCH 11/20] prepare for addition of full-infinite env CTMRG

---
 src/algorithms/timeevol/fullupdate.jl | 33 +++++++++++++++++++++------
 1 file changed, 26 insertions(+), 7 deletions(-)

diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/timeevol/fullupdate.jl
index 0b7ee9ad..e8265779 100644
--- a/src/algorithms/timeevol/fullupdate.jl
+++ b/src/algorithms/timeevol/fullupdate.jl
@@ -1,6 +1,9 @@
 include("fu_gaugefix.jl")
 include("fu_optimize.jl")
 
+# TODO: add option to use full-infinite environment 
+# for CTMRG moves when it is implemented in PEPSKit
+
 """
 CTMRG left-move to update CTMRGEnv in the c-th column 
 ```
@@ -14,7 +17,12 @@ CTMRG left-move to update CTMRGEnv in the c-th column
 ```
 """
 function ctmrg_leftmove!(
-    col::Int, peps::InfinitePEPS, envs::CTMRGEnv, chi::Int, svderr::Float64=1e-9
+    col::Int,
+    peps::InfinitePEPS,
+    envs::CTMRGEnv,
+    chi::Int,
+    svderr::Float64=1e-9;
+    cheap::Bool=true,
 )
     trscheme = truncerr(svderr) & truncdim(chi)
     alg = CTMRG(;
@@ -39,7 +47,12 @@ CTMRG right-move to update CTMRGEnv in the c-th column
 ```
 """
 function ctmrg_rightmove!(
-    col::Int, peps::InfinitePEPS, envs::CTMRGEnv, chi::Int, svderr::Float64=1e-9
+    col::Int,
+    peps::InfinitePEPS,
+    envs::CTMRGEnv,
+    chi::Int,
+    svderr::Float64=1e-9;
+    cheap::Bool=true,
 )
     Nr, Nc = size(peps)
     @assert 1 <= col <= Nc
@@ -64,6 +77,7 @@ function update_column!(
     svderr::Float64=1e-9,
     maxiter::Int=50,
     maxdiff::Float64=1e-15,
+    cheap=true,
     gaugefix::Bool=true,
 )
     Nr, Nc = size(peps)
@@ -153,8 +167,8 @@ function update_column!(
         end
     end
     # update CTMRGEnv
-    ctmrg_leftmove!(col, peps, envs, chi, svderr)
-    ctmrg_rightmove!(_next(col, Nc), peps, envs, chi, svderr)
+    ctmrg_leftmove!(col, peps, envs, chi, svderr; cheap=cheap)
+    ctmrg_rightmove!(_next(col, Nc), peps, envs, chi, svderr; cheap=cheap)
     return localfid, costs
 end
 
@@ -172,19 +186,24 @@ function fullupdate!(
     envs::CTMRGEnv,
     Dcut::Int,
     chi::Int,
-    svderr::Float64=1e-9,
+    svderr::Float64=1e-9;
+    cheap=false,
 )
     Nr, Nc = size(peps)
     fid, maxcost = 0.0, 0.0
     for col in 1:Nc
-        tmpfid, costs = update_column!(col, gate, peps, envs, Dcut, chi; svderr=svderr)
+        tmpfid, costs = update_column!(
+            col, gate, peps, envs, Dcut, chi; svderr=svderr, cheap=cheap
+        )
         fid += tmpfid
         maxcost = max(maxcost, maximum(costs))
     end
     rotr90!(peps)
     rotr90!(envs)
     for row in 1:Nr
-        tmpfid, costs = update_column!(row, gate, peps, envs, Dcut, chi; svderr=svderr)
+        tmpfid, costs = update_column!(
+            row, gate, peps, envs, Dcut, chi; svderr=svderr, cheap=cheap
+        )
         fid += tmpfid
         maxcost = max(maxcost, maximum(costs))
     end

From 3a82f06528aa9ddaa74cb1c858617732f6345fdb Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Thu, 14 Nov 2024 10:00:20 +0800
Subject: [PATCH 12/20] Add `length` for SUWeight

---
 src/states/suweight.jl | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/states/suweight.jl b/src/states/suweight.jl
index 19421d33..037ac322 100644
--- a/src/states/suweight.jl
+++ b/src/states/suweight.jl
@@ -38,6 +38,11 @@ function Base.iterate(wts::SUWeight, state=1)
     end
 end
 
+function Base.length(wts::SUWeight)
+    @assert size(wts.x) == size(wts.y)
+    return 2 * prod(size(wts.x))
+end
+
 function Base.isapprox(wts1::SUWeight, wts2::SUWeight; atol=0.0, rtol=1e-5)
     return (
         isapprox(wts1.x, wts2.x; atol=atol, rtol=rtol) &&

From 9577febf02548a89e04af27522bdba44da1cca10 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 15 Nov 2024 15:14:19 +0800
Subject: [PATCH 13/20] Define `Base.show` for `SUWeight`

---
 src/states/suweight.jl | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/states/suweight.jl b/src/states/suweight.jl
index 037ac322..6995ab9a 100644
--- a/src/states/suweight.jl
+++ b/src/states/suweight.jl
@@ -27,6 +27,17 @@ function Base.:(-)(wts1::SUWeight, wts2::SUWeight)
     return SUWeight(wts1.x - wts2.x, wts1.y - wts2.y)
 end
 
+function Base.show(io::IO, wts::SUWeight)
+    N1, N2 = size(wts)
+    for (direction, r, c) in Iterators.product("xy", 1:N1, 1:N2)
+        println(io, "$direction[$r,$c]: ")
+        wt = (direction == 'x' ? wts.x[r,c] : wts.y[r,c])
+        for (k, b) in blocks(wt)
+            println(io, k, " = ", diag(b))
+        end
+    end
+end
+
 function Base.iterate(wts::SUWeight, state=1)
     nx = prod(size(wts.x))
     if 1 <= state <= nx

From a1c3e062c690a32384c73e3bd7a6acd93e77b89b Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sun, 17 Nov 2024 11:29:49 +0800
Subject: [PATCH 14/20] add test for simple update

---
 src/PEPSKit.jl                              |   9 +-
 src/algorithms/contractions/ctmrg_rhos.jl   | 189 ++++++++++++++++++++
 src/algorithms/contractions/measure_rhos.jl |  39 ++++
 src/algorithms/ctmrg/gaugefix.jl            |  13 ++
 src/algorithms/timeevol/fullupdate.jl       | 112 +++++++++++-
 src/algorithms/timeevol/simpleupdate.jl     |  67 ++++++-
 test/heisenberg_sufu.jl                     |  59 ++++++
 test/utility/heis.jl                        | 106 +++++++++++
 8 files changed, 590 insertions(+), 4 deletions(-)
 create mode 100644 src/algorithms/contractions/ctmrg_rhos.jl
 create mode 100644 src/algorithms/contractions/measure_rhos.jl
 create mode 100644 test/heisenberg_sufu.jl
 create mode 100644 test/utility/heis.jl

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index bd71e728..0726ca29 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -40,6 +40,8 @@ include("environments/transferpepo_environments.jl")
 
 include("algorithms/contractions/localoperator.jl")
 include("algorithms/contractions/ctmrg_contractions.jl")
+include("algorithms/contractions/ctmrg_rhos.jl")
+include("algorithms/contractions/measure_rhos.jl")
 
 include("algorithms/ctmrg/sparse_environments.jl")
 include("algorithms/ctmrg/ctmrg.jl")
@@ -172,8 +174,11 @@ export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolver
 export fixedpoint
 
-export simpleupdate!, absorb_wt
-export fullupdate!
+export absorb_wt, absorb_wt!
+export su_iter!, simpleupdate!
+export fu_iter!, fullupdate!
+export meas_site, meas_bond
+export calrho_site, calrho_bondx, calrho_bondy, calrho_all
 
 export SUWeight
 export InfinitePEPS, InfiniteTransferPEPS
diff --git a/src/algorithms/contractions/ctmrg_rhos.jl b/src/algorithms/contractions/ctmrg_rhos.jl
new file mode 100644
index 00000000..13585fc4
--- /dev/null
+++ b/src/algorithms/contractions/ctmrg_rhos.jl
@@ -0,0 +1,189 @@
+"""
+Calculate 1-site rho at site `(r,c)`
+```
+    C1 - χ4 - T1 - χ6 - C2  r-1
+    |         ‖         |
+    χ2        DN        χ8
+    |         ‖         |
+    T4 = DW =k/b = DE = T2  r
+    |         ‖         |
+    χ1        DS        χ7
+    |         ‖         |
+    C4 - χ3 - T3 - χ5 - C3  r+1
+    c-1       c        c+1
+```
+Indices d0, d1 are physical indices of ket, bra
+"""
+function calrho_site(
+    row::Int, col::Int, envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket
+)
+    N1, N2 = size(ket)
+    @assert 1 <= row <= N1 && 1 <= col <= N2
+    rp1, rm1 = _next(row, N1), _prev(row, N1)
+    cp1, cm1 = _next(col, N2), _prev(col, N2)
+    tket, tbra = ket[row, col], bra[row, col]
+    c1 = envs.corners[1, rm1, cm1]
+    t1 = envs.edges[1, rm1, col]
+    c2 = envs.corners[2, rm1, cp1]
+    t2 = envs.edges[2, row, cp1]
+    c3 = envs.corners[3, rp1, cp1]
+    t3 = envs.edges[3, rp1, col]
+    c4 = envs.corners[4, rp1, cm1]
+    t4 = envs.edges[4, row, cm1]
+    PEPSKit.@autoopt @tensor rho1[d1; d0] := (
+        c4[χ3, χ1] *
+        t4[χ1, DW0, DW1, χ2] *
+        c1[χ2, χ4] *
+        t3[χ5, DS0, DS1, χ3] *
+        tket[d0, DN0, DE0, DS0, DW0] *
+        conj(tbra[d1, DN1, DE1, DS1, DW1]) *
+        t1[χ4, DN0, DN1, χ6] *
+        c3[χ7, χ5] *
+        t2[χ8, DE0, DE1, χ7] *
+        c2[χ6, χ8]
+    )
+    return rho1
+end
+
+"""
+Calculate 2-site rho on sites `(r,c)(r,c+1)`
+```
+    C1 - χ4 - T1 - χ6 - T1 - χ8 - C2    r-1
+    |         ‖         ‖         |
+    χ2       DN1       DN2        χ10
+    |         ‖         ‖         |
+    T4 = DW =k/b = DM =k/b = DE = T2    r
+    |         ‖         ‖         |
+    χ1       DS1       DS2        χ9
+    |         ‖         ‖         |
+    C4 - χ3 - T3 - χ5 - T3 - χ7 - C3    r+1
+    c-1       c        c+1       c+2
+```
+Indices d0, d1 are physical indices of ket, bra
+"""
+function calrho_bondx(
+    row::Int, col::Int, envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket
+)
+    N1, N2 = size(ket)
+    @assert 1 <= row <= N1 && 1 <= col <= N2
+    rp1, rm1 = _next(row, N1), _prev(row, N1)
+    cp1, cm1 = _next(col, N2), _prev(col, N2)
+    cp2 = _next(cp1, N2)
+    tket1, tbra1 = ket[row, col], bra[row, col]
+    tket2, tbra2 = ket[row, cp1], bra[row, cp1]
+    c1 = envs.corners[1, rm1, cm1]
+    t11, t12 = envs.edges[1, rm1, col], envs.edges[1, rm1, cp1]
+    c2 = envs.corners[2, rm1, cp2]
+    t2 = envs.edges[2, row, cp2]
+    c3 = envs.corners[3, rp1, cp2]
+    t31, t32 = envs.edges[3, rp1, col], envs.edges[3, rp1, cp1]
+    c4 = envs.corners[4, rp1, cm1]
+    t4 = envs.edges[4, row, cm1]
+    PEPSKit.@autoopt @tensor rho2[d11, d21; d10, d20] := (
+        c4[χ3, χ1] *
+        t4[χ1, DW0, DW1, χ2] *
+        c1[χ2, χ4] *
+        t31[χ5, DS10, DS11, χ3] *
+        tket1[d10, DN10, DM0, DS10, DW0] *
+        conj(tbra1[d11, DN11, DM1, DS11, DW1]) *
+        t11[χ4, DN10, DN11, χ6] *
+        t32[χ7, DS20, DS21, χ5] *
+        tket2[d20, DN20, DE0, DS20, DM0] *
+        conj(tbra2[d21, DN21, DE1, DS21, DM1]) *
+        t12[χ6, DN20, DN21, χ8] *
+        c3[χ9, χ7] *
+        t2[χ10, DE0, DE1, χ9] *
+        c2[χ8, χ10]
+    )
+    return rho2
+end
+
+"""
+Calculate 2-site rho on sites `(r,c)(r-1,c)`
+```
+    C1 - χ9 - T1 -χ10 - C2  r-2
+    |         ‖         |
+    χ7        DN        χ8
+    |         ‖         |
+    T4 = DW2=k/b =DE2 = T2  r-1
+    |         ‖         |
+    χ5        DM        χ6
+    |         ‖         |
+    T4 = DW1=k/b =DE1 = T2  r
+    |         ‖         |
+    χ3        DS        χ4
+    |         ‖         |
+    C4 - χ1 - T3 - χ2 - C3  r+1
+    c-1       c        c+1
+```
+Indices d0, d1 are physical indices of ket, bra
+"""
+function calrho_bondy(
+    row::Int, col::Int, envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket
+)
+    N1, N2 = size(ket)
+    @assert 1 <= row <= N1 && 1 <= col <= N2
+    rp1, rm1 = _next(row, N1), _prev(row, N1)
+    cp1, cm1 = _next(col, N2), _prev(col, N2)
+    rm2 = _prev(rm1, N1)
+    tket1, tbra1 = ket[row, col], bra[row, col]
+    tket2, tbra2 = ket[rm1, col], bra[rm1, col]
+    c1 = envs.corners[1, rm2, cm1]
+    t1 = envs.edges[1, rm2, col]
+    c2 = envs.corners[2, rm2, cp1]
+    t21, t22 = envs.edges[2, row, cp1], envs.edges[2, rm1, cp1]
+    c3 = envs.corners[3, rp1, cp1]
+    t3 = envs.edges[3, rp1, col]
+    c4 = envs.corners[4, rp1, cm1]
+    t41, t42 = envs.edges[4, row, cm1], envs.edges[4, rm1, cm1]
+    PEPSKit.@autoopt @tensor rho2[d11, d21; d10, d20] := (
+        c4[χ1, χ3] *
+        t3[χ2, DS0, DS1, χ1] *
+        c3[χ4, χ2] *
+        t41[χ3, DW10, DW11, χ5] *
+        tket1[d10, DM0, DE10, DS0, DW10] *
+        conj(tbra1[d11, DM1, DE11, DS1, DW11]) *
+        t21[χ6, DE10, DE11, χ4] *
+        t42[χ5, DW20, DW21, χ7] *
+        tket2[d20, DN0, DE20, DM0, DW20] *
+        conj(tbra2[d21, DN1, DE21, DM1, DW21]) *
+        t22[χ8, DE20, DE21, χ6] *
+        c1[χ7, χ9] *
+        t1[χ9, DN0, DN1, χ10] *
+        c2[χ10, χ8]
+    )
+    return rho2
+end
+
+"""
+Calculate rho for all sites
+"""
+function calrho_allsites(envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket)
+    Nr, Nc = size(ket)
+    return collect(
+        calrho_site(r, c, envs, ket, bra) for (r, c) in Iterators.product(1:Nr, 1:Nc)
+    )
+end
+
+"""
+Calculate rho for all nearest-neighbor bonds
+"""
+function calrho_allnbs(envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket)
+    Nr, Nc = size(ket)
+    rhoxss = collect(
+        calrho_bondx(r, c, envs, ket, bra) for (r, c) in Iterators.product(1:Nr, 1:Nc)
+    )
+    rhoyss = collect(
+        calrho_bondy(r, c, envs, ket, bra) for (r, c) in Iterators.product(1:Nr, 1:Nc)
+    )
+    return [rhoxss, rhoyss]
+end
+
+"""
+Calculate rho for all sites and nearest-neighbor bonds
+"""
+function calrho_all(envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket)
+    rho1ss = calrho_allsites(envs, ket, bra)
+    rho2sss = calrho_allnbs(envs, ket, bra)
+    return rho1ss, rho2sss
+end
diff --git a/src/algorithms/contractions/measure_rhos.jl b/src/algorithms/contractions/measure_rhos.jl
new file mode 100644
index 00000000..a075a077
--- /dev/null
+++ b/src/algorithms/contractions/measure_rhos.jl
@@ -0,0 +1,39 @@
+"""
+Get identity operator on the physical space
+"""
+function _getid_from_rho(rho::AbstractTensorMap)
+    Pspace = codomain(rho)[1]
+    if isdual(Pspace)
+        Pspace = adjoint(Pspace)
+    end
+    return TensorKit.id(Pspace)
+end
+
+"""
+Measure `<op>` using 1-site rho
+"""
+function meas_site(op::AbstractTensorMap, rho1::AbstractTensorMap)
+    Id = _getid_from_rho(rho1)
+    val = ncon((rho1, op), ([1, 2], [1, 2]))
+    nrm = ncon((rho1, Id), ([1, 2], [1, 2]))
+    meas = first(blocks(val / nrm))[2][1]
+    return meas
+end
+
+"""
+Measure `<op1 op2>` using 2-site rho
+"""
+function meas_bond(op1::AbstractTensorMap, op2::AbstractTensorMap, rho2::AbstractTensorMap)
+    return meas_bond(op1 ⊗ op2, rho2)
+end
+
+"""
+Measure `<gate>` using 2-site rho
+"""
+function meas_bond(gate::AbstractTensorMap, rho2::AbstractTensorMap)
+    Id = _getid_from_rho(rho2)
+    val = ncon((rho2, gate), ([1, 2, 3, 4], [1, 2, 3, 4]))
+    nrm = ncon((rho2, Id ⊗ Id), ([1, 2, 3, 4], [1, 2, 3, 4]))
+    meas = first(blocks(val / nrm))[2][1]
+    return meas
+end
diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl
index 5e63d5ef..060db633 100644
--- a/src/algorithms/ctmrg/gaugefix.jl
+++ b/src/algorithms/ctmrg/gaugefix.jl
@@ -173,6 +173,19 @@ function calc_convergence(envs, CSold, TSold)
     return max(ΔCS, ΔTS), CSnew, TSnew
 end
 
+"""
+Calculate convergence of CTMRG by comparing the singular values of CTM tensors
+"""
+function calc_convergence(envsNew::CTMRGEnv, envsOld::CTMRGEnv)
+    CSNew = map(x -> tsvd(x)[2], envsNew.corners)
+    TSNew = map(x -> tsvd(x)[2], envsNew.edges)
+    CSOld = map(x -> tsvd(x)[2], envsOld.corners)
+    TSOld = map(x -> tsvd(x)[2], envsOld.edges)
+    ΔCS = maximum(_singular_value_distance, zip(CSOld, CSNew))
+    ΔTS = maximum(_singular_value_distance, zip(TSOld, TSNew))
+    return max(ΔCS, ΔTS)
+end
+
 @non_differentiable calc_convergence(args...)
 
 """
diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/timeevol/fullupdate.jl
index e8265779..f25e346f 100644
--- a/src/algorithms/timeevol/fullupdate.jl
+++ b/src/algorithms/timeevol/fullupdate.jl
@@ -180,7 +180,7 @@ Otherwise, use full-infinite environment instead.
 
 Reference: Physical Review B 92, 035142 (2015)
 """
-function fullupdate!(
+function fu_iter!(
     gate::AbstractTensorMap,
     peps::InfinitePEPS,
     envs::CTMRGEnv,
@@ -212,3 +212,113 @@ function fullupdate!(
     fid /= (2 * Nr * Nc)
     return fid, maxcost
 end
+
+# TODO: pass Hamiltonian gate as `LocalOperator` 
+"""
+Perform full update
+"""
+function fullupdate!(
+    peps::InfinitePEPS,
+    envs::CTMRGEnv,
+    ham::AbstractTensorMap,
+    dt::Float64,
+    Dcut::Int,
+    chi::Int;
+    evolstep::Int=5000,
+    svderr::Float64=1e-9,
+    rgint::Int=10,
+    rgtol::Float64=1e-6,
+    rgmaxiter::Int=10,
+    ctmrgscheme=:sequential,
+    cheap=false,
+)
+    time_start = time()
+    N1, N2 = size(peps)
+    @assert endswith(folder, "/")
+    # CTMRG algorithm to reconverge environment
+    ctm_alg = CTMRG(;
+        tol=rgtol,
+        maxiter=rgmaxiter,
+        miniter=1,
+        verbosity=2,
+        trscheme=truncerr(svderr) & truncdim(chi),
+        svd_alg=SVDAdjoint(; fwd_alg=TensorKit.SDD()),
+        ctmrgscheme=ctmrgscheme,
+    )
+    @printf(
+        "%-4s %7s%10s%12s%11s  %s/%s\n",
+        "step",
+        "dt",
+        "energy",
+        "Δe",
+        "svd_diff",
+        "speed",
+        "meas(s)"
+    )
+    flush(stdout)
+    gate = exp(-dt * ham)
+    esite0, peps0, envs0 = Inf, deepcopy(peps), deepcopy(envs)
+    diff_energy = 0.0
+    for count in 1:evolstep
+        time0 = time()
+        fid, cost = fu_iter!(gate, peps, envs, Dcut, chi, svderr; cheap=cheap)
+        time1 = time()
+        if count == 1 || count % rgint == 0
+            meast0 = time()
+            # reconverge `env` (in place)
+            println(stderr, "---- FU step $count: reconverging envs ----")
+            envs2 = leading_boundary(envs, peps, ctm_alg)
+            envs.edges[:], envs.corners[:] = envs2.edges, envs2.corners
+            # TODO: monitor energy with costfun
+            # esite = costfun(peps, envs, ham)
+            rho2sss = calrho_allnbs(envs, peps)
+            ebonds = [collect(meas_bond(ham, rho2) for rho2 in rho2sss[n]) for n in 1:2]
+            esite = sum(sum(ebonds)) / (N1 * N2)
+            meast1 = time()
+            # monitor change of CTMRGEnv by its singular values
+            diff_energy = esite - esite0
+            diff_ctm = calc_convergence(envs, envs0)
+            @printf(
+                "%-4d %7.0e%10.5f%12.3e%11.3e  %.3f/%.3f\n",
+                count,
+                dt,
+                esite,
+                diff_energy,
+                diff_ctm,
+                time1 - time0,
+                meast1 - meast0
+            )
+            if diff_energy > 0
+                @printf("Energy starts to increase. Abort evolution.\n")
+                # restore peps and envs at last checking
+                peps.A[:] = peps0.A
+                envs.corners[:], envs.edges[:] = envs0.corners, envs0.edges
+                break
+            end
+            esite0, peps0, envs0 = esite, deepcopy(peps), deepcopy(envs)
+        end
+    end
+    # reconverge the environment tensors
+    for io in (stdout, stderr)
+        @printf(io, "Reconverging final envs ... \n")
+        flush(io)
+    end
+    envs2 = leading_boundary(
+        envs,
+        peps,
+        CTMRG(;
+            tol=1e-10,
+            maxiter=50,
+            miniter=1,
+            verbosity=2,
+            trscheme=truncerr(svderr) & truncdim(chi),
+            svd_alg=SVDAdjoint(; fwd_alg=TensorKit.SDD()),
+            ctmrgscheme=ctmrgscheme,
+        ),
+    )
+    envs.edges[:], envs.corners[:] = envs2.edges, envs2.corners
+    time_end = time()
+    @printf("Evolution time: %.3f s\n\n", time_end - time_start)
+    print(stderr, "\n----------\n\n")
+    return esite0, diff_energy
+end
diff --git a/src/algorithms/timeevol/simpleupdate.jl b/src/algorithms/timeevol/simpleupdate.jl
index a57795be..451ff8c9 100644
--- a/src/algorithms/timeevol/simpleupdate.jl
+++ b/src/algorithms/timeevol/simpleupdate.jl
@@ -60,6 +60,19 @@ function absorb_wt(
     return t2
 end
 
+"""
+Absorb bond weights into iPEPS site tensors
+"""
+function absorb_wt!(peps::InfinitePEPS, wts::SUWeight)
+    N1, N2 = size(peps)
+    for (r, c) in Iterators.product(1:N1, 1:N2)
+        for ax in 2:5
+            peps.A[r, c] = absorb_wt(peps.A[r, c], r, c, ax, wts; sqrtwt=true)
+        end
+    end
+    return nothing
+end
+
 """
 Simple update of bond `wts.x[r,c]`
 ```
@@ -153,7 +166,7 @@ and SUWeight `wts` with the nearest neighbor gate `gate`
 When `bipartite === true` (for square lattice), the unit cell size should be 2 x 2, 
 and the tensor and x/y weight at `(row, col)` is the same as `(row+1, col+1)`
 """
-function simpleupdate!(
+function su_iter!(
     gate::AbstractTensorMap,
     peps::InfinitePEPS,
     wts::SUWeight,
@@ -198,3 +211,55 @@ function simpleupdate!(
     end
     return nothing
 end
+
+function compare_weights(wts1::SUWeight, wts2::SUWeight)
+    wtdiff = sum(_singular_value_distance((wt1, wt2)) for (wt1, wt2) in zip(wts1, wts2))
+    wtdiff /= 2 * prod(size(wts1))
+    return wtdiff
+end
+
+"""
+Perform simple update (maximum `evolstep` iterations)
+with nearest neighbor Hamiltonian `ham` and time step `dt` 
+until the change of bond weights is smaller than `wtdiff_tol` 
+"""
+function simpleupdate!(
+    peps::InfinitePEPS,
+    wts::SUWeight,
+    ham::AbstractTensorMap,
+    dt::Float64,
+    Dcut::Int;
+    evolstep::Int=400000,
+    svderr::Float64=1e-10,
+    wtdiff_tol::Float64=1e-10,
+    bipartite::Bool=false,
+    check_int::Int=500,
+)
+    time_start = time()
+    N1, N2 = size(peps)
+    if bipartite
+        @assert N1 == N2 == 2
+    end
+    @printf("%-9s%6s%12s  %s\n", "Step", "dt", "wt_diff", "speed/s")
+    # exponentiating the 2-site Hamiltonian gate
+    gate = exp(-dt * ham)
+    wtdiff = 1e+3
+    wts0 = deepcopy(wts)
+    for count in 1:evolstep
+        time0 = time()
+        su_iter!(gate, peps, wts, Dcut, svderr; bipartite=bipartite)
+        wtdiff = compare_weights(wts, wts0)
+        stop = (wtdiff < wtdiff_tol) || (count == evolstep)
+        wts0 = deepcopy(wts)
+        time1 = time()
+        if ((count == 1) || (count % check_int == 0) || stop)
+            @printf("%-9d%6.0e%12.3e  %.3f\n", count, dt, wtdiff, time1 - time0)
+        end
+        if stop
+            break
+        end
+    end
+    time_end = time()
+    @printf("Evolution time: %.2f s\n\n", time_end - time_start)
+    return wtdiff
+end
diff --git a/test/heisenberg_sufu.jl b/test/heisenberg_sufu.jl
new file mode 100644
index 00000000..e0a9e28c
--- /dev/null
+++ b/test/heisenberg_sufu.jl
@@ -0,0 +1,59 @@
+using Test
+using Random
+using PEPSKit
+using TensorKit
+import Statistics: mean
+include("utility/heis.jl")
+using .OpsHeis, .RhoMeasureHeis
+
+# benchmark data for D = 3 is from
+# Phys. Rev. B 94, 035133 (2016)
+
+# random initialization of 2x2 iPEPS and CTMRGEnv
+# (using real numbersf)
+Dcut = 3
+χenv = 24
+N1, N2 = 2, 2
+Pspace = ℂ^2
+Vspace = ℂ^Dcut
+Random.seed!(0)
+peps = InfinitePEPS(
+    collect(
+        TensorMap(rand, Float64, Pspace, Vspace ⊗ Vspace ⊗ Vspace' ⊗ Vspace') for
+        (row, col) in Iterators.product(1:N1, 1:N2)
+    ),
+)
+wts = SUWeight(
+    collect(id(Vspace) for (row, col) in Iterators.product(1:N1, 1:N2)),
+    collect(id(Vspace) for (row, col) in Iterators.product(1:N1, 1:N2)),
+)
+# normalize peps
+for ind in CartesianIndices(peps.A)
+    peps.A[ind] /= PEPSKit.maxabs(peps.A[ind])
+end
+# Heisenberg model Hamiltonian
+ham = gen_gate()
+
+# simple update energy and magnetization
+dts = [1e-2, 1e-3, 4e-4, 1e-4]
+tols = [1e-6, 1e-7, 1e-8, 1e-9]
+for (dt, tol) in zip(dts, tols)
+    simpleupdate!(peps, wts, ham, dt, Dcut; evolstep=30000, wtdiff_tol=tol)
+end
+# absort weight into site tensors
+absorb_wt!(peps, wts)
+# CTMRG
+envs = CTMRGEnv(rand, Float64, peps, ℂ^χenv)
+trscheme = truncerr(1e-9) & truncdim(χenv)
+ctm_alg = CTMRG(; tol=1e-10, verbosity=2, trscheme=trscheme, ctmrgscheme=:sequential)
+envs = leading_boundary(envs, peps, ctm_alg)
+# measure physical quantities
+rho1ss, rho2sss = calrho_all(envs, peps)
+result = measrho_all(rho1ss, rho2sss)
+display(result)
+@test isapprox(result["e_site"], -0.6633; atol=1e-3)
+@test isapprox(mean(result["mag_norm"]), 0.3972; atol=1e-3)
+
+# full update energy and magnetization
+# e_fu = -0.6654
+# mag_fu = 0.3634
diff --git a/test/utility/heis.jl b/test/utility/heis.jl
new file mode 100644
index 00000000..4dbd6827
--- /dev/null
+++ b/test/utility/heis.jl
@@ -0,0 +1,106 @@
+module OpsHeis
+
+export gen_gate, gen_siteop, gen_bondop
+using TensorKit
+
+"""
+Create 1-site operators for Heisenberg model
+"""
+function gen_siteop(name::String)
+    Pspace = ℂ^2
+    if name == "Id"
+        return id(Pspace)
+    end
+    op = TensorMap(zeros, Pspace, Pspace)
+    if name == "Nud"
+        block(op, Trivial())[:] = [1.0 0.0; 0.0 1.0]
+    elseif name == "Sp"
+        block(op, Trivial())[:] = [0.0 1.0; 0.0 0.0]
+    elseif name == "Sm"
+        block(op, Trivial())[:] = [0.0 0.0; 1.0 0.0]
+    elseif name == "Sz"
+        block(op, Trivial())[:] = [1.0 0.0; 0.0 -1.0] / 2
+    elseif name == "Sx"
+        block(op, Trivial())[:] = [0.0 1.0; 1.0 0.0] / 2
+    elseif name == "iSy"
+        block(op, Trivial())[:] = [0.0 1.0; -1.0 0.0] / 2
+    else
+        throw(ArgumentError("Invalid 1-site spin operator"))
+    end
+    return op
+end
+
+"""
+Create 2-site operators for Heisenberg model
+"""
+function gen_bondop(name1::String, name2::String)
+    op1 = gen_siteop(name1)
+    op2 = gen_siteop(name2)
+    op = op1 ⊗ op2
+    return op
+end
+
+"""
+Create nearest neighbor gate for Heisenberg model
+"""
+function gen_gate(J::Float64=1.0; dens_shift::Bool=false)
+    heis =
+        J * (
+            (1 / 2) * gen_bondop("Sp", "Sm") +
+            (1 / 2) * gen_bondop("Sm", "Sp") +
+            gen_bondop("Sz", "Sz")
+        )
+    if dens_shift
+        heis = heis - (J / 4) * gen_bondop("Nud", "Nud")
+    end
+    return heis
+end
+
+end
+
+module RhoMeasureHeis
+
+export measrho_all, cal_Esite
+
+using TensorKit
+using PEPSKit
+using Statistics: mean
+using ..OpsHeis
+
+function cal_mags(rho1ss::Matrix{<:AbstractTensorMap})
+    Pspace = codomain(rho1ss[1, 1])[1]'
+    Sas = [gen_siteop(name) for name in ("Sx", "iSy", "Sz")]
+    return [collect(meas_site(Sa, rho1) for rho1 in rho1ss) for Sa in Sas]
+end
+
+function cal_spincor(rho2ss::Matrix{<:AbstractTensorMap})
+    SpSm = gen_bondop("Sp", "Sm")
+    SzSz = gen_bondop("Sz", "Sz")
+    return collect(meas_bond(SpSm, rho2) + meas_bond(SzSz, rho2) for rho2 in rho2ss)
+end
+
+function cal_Esite(rho2sss::Vector{<:Matrix{<:AbstractTensorMap}})
+    N1, N2 = size(rho2sss[1])
+    gate1 = gen_gate(; dens_shift=false)
+    # 1st neighbor bond energy
+    ebond1s = [collect(meas_bond(gate1, rho2) for rho2 in rho2sss[n]) for n in 1:2]
+    esite = sum(sum(ebond1s)) / (N1 * N2)
+    return esite, ebond1s
+end
+
+function measrho_all(
+    rho1ss::Matrix{<:AbstractTensorMap}, rho2sss::Vector{<:Matrix{<:AbstractTensorMap}}
+)
+    results = Dict{String,Any}()
+    N1, N2 = size(rho1ss)
+    results["e_site"], results["energy"] = cal_Esite(rho2sss)
+    results["mag"] = cal_mags(rho1ss)
+    results["mag_norm"] = collect(
+        norm([results["mag"][n][r, c] for n in 1:3]) for
+        (r, c) in Iterators.product(1:N1, 1:N2)
+    )
+    results["spincor"] = [cal_spincor(rho2ss) for rho2ss in rho2sss]
+    return results
+end
+
+end

From 448176ceade38d593c192083e21ad62f0ffe75eb Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sun, 17 Nov 2024 17:30:18 +0800
Subject: [PATCH 15/20] add test for full update

---
 src/algorithms/contractions/ctmrg_rhos.jl | 41 +++++++++++++++++++
 src/algorithms/timeevol/fullupdate.jl     |  3 --
 test/heisenberg_sufu.jl                   | 49 +++++++++++++----------
 test/runtests.jl                          |  2 +-
 4 files changed, 69 insertions(+), 26 deletions(-)

diff --git a/src/algorithms/contractions/ctmrg_rhos.jl b/src/algorithms/contractions/ctmrg_rhos.jl
index 13585fc4..7898d25f 100644
--- a/src/algorithms/contractions/ctmrg_rhos.jl
+++ b/src/algorithms/contractions/ctmrg_rhos.jl
@@ -155,6 +155,47 @@ function calrho_bondy(
     return rho2
 end
 
+# TODO: add rhos on next nearest neighbor bonds
+
+"""
+Calculate 2-site rho on 2nd nearest neighbor sites `(r,c)(r-1,c+1)`
+```
+    C1 -χ10 - T1 -χ11 - T1 -χ12 - C2    r-2
+    |         ‖         ‖         |
+    χ8       DN1       DN2        χ9
+    |         ‖         ‖         |
+    T4 =DW2= k/b =DH2= k/b =DE2== T2    r-1
+    |         ‖         ‖         |
+    χ6       DV1       DV2        χ7
+    |         ‖         ‖         |
+    T4 =DW1= k/b =DH1= k/b =DE1== T2    r
+    |         ‖         ‖         |
+    χ4       DS1       DS2        χ5
+    |         ‖         ‖         |
+    C4 - χ1 - T3 - χ2 - T3 - χ3 - C4    r+1
+    c-1       c        c+1       c+2
+```
+Indices d0, d1 are physical indices of ket, bra
+"""
+function calrho_bondd1(
+    row::Int, col::Int, envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket
+)
+    N1, N2 = size(ket)
+    @assert 1 <= row <= N1 && 1 <= col <= N2
+    throw("not implemented")
+end
+
+"""
+Calculate 2-site rho on 2nd nearest neighbor sites `(r,c+1)(r-1,c)`
+"""
+function calrho_bondd2(
+    row::Int, col::Int, envs::CTMRGEnv, ket::InfinitePEPS, bra::InfinitePEPS=ket
+)
+    N1, N2 = size(ket)
+    @assert 1 <= row <= N1 && 1 <= col <= N2
+    throw("not implemented")
+end
+
 """
 Calculate rho for all sites
 """
diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/timeevol/fullupdate.jl
index f25e346f..3eb85a4b 100644
--- a/src/algorithms/timeevol/fullupdate.jl
+++ b/src/algorithms/timeevol/fullupdate.jl
@@ -234,7 +234,6 @@ function fullupdate!(
 )
     time_start = time()
     N1, N2 = size(peps)
-    @assert endswith(folder, "/")
     # CTMRG algorithm to reconverge environment
     ctm_alg = CTMRG(;
         tol=rgtol,
@@ -255,7 +254,6 @@ function fullupdate!(
         "speed",
         "meas(s)"
     )
-    flush(stdout)
     gate = exp(-dt * ham)
     esite0, peps0, envs0 = Inf, deepcopy(peps), deepcopy(envs)
     diff_energy = 0.0
@@ -301,7 +299,6 @@ function fullupdate!(
     # reconverge the environment tensors
     for io in (stdout, stderr)
         @printf(io, "Reconverging final envs ... \n")
-        flush(io)
     end
     envs2 = leading_boundary(
         envs,
diff --git a/test/heisenberg_sufu.jl b/test/heisenberg_sufu.jl
index e0a9e28c..86372e67 100644
--- a/test/heisenberg_sufu.jl
+++ b/test/heisenberg_sufu.jl
@@ -1,28 +1,22 @@
 using Test
+using Printf
 using Random
 using PEPSKit
 using TensorKit
 import Statistics: mean
 include("utility/heis.jl")
-using .OpsHeis, .RhoMeasureHeis
+import .OpsHeis: gen_gate
+import .RhoMeasureHeis: measrho_all
 
 # benchmark data for D = 3 is from
 # Phys. Rev. B 94, 035133 (2016)
 
-# random initialization of 2x2 iPEPS and CTMRGEnv
-# (using real numbersf)
-Dcut = 3
-χenv = 24
+# random initialization of 2x2 iPEPS and CTMRGEnv (using real numbers)
+Dcut, χenv = 4, 16
 N1, N2 = 2, 2
-Pspace = ℂ^2
-Vspace = ℂ^Dcut
+Pspace, Vspace = ℂ^2, ℂ^Dcut
 Random.seed!(0)
-peps = InfinitePEPS(
-    collect(
-        TensorMap(rand, Float64, Pspace, Vspace ⊗ Vspace ⊗ Vspace' ⊗ Vspace') for
-        (row, col) in Iterators.product(1:N1, 1:N2)
-    ),
-)
+peps = InfinitePEPS(rand, Float64, 2, Dcut; unitcell=(N1, N2))
 wts = SUWeight(
     collect(id(Vspace) for (row, col) in Iterators.product(1:N1, 1:N2)),
     collect(id(Vspace) for (row, col) in Iterators.product(1:N1, 1:N2)),
@@ -34,11 +28,12 @@ end
 # Heisenberg model Hamiltonian
 ham = gen_gate()
 
-# simple update energy and magnetization
+# simple update
 dts = [1e-2, 1e-3, 4e-4, 1e-4]
 tols = [1e-6, 1e-7, 1e-8, 1e-9]
-for (dt, tol) in zip(dts, tols)
-    simpleupdate!(peps, wts, ham, dt, Dcut; evolstep=30000, wtdiff_tol=tol)
+for (n, (dt, tol)) in enumerate(zip(dts, tols))
+    Dcut2 = (n == 1 ? Dcut + 1 : Dcut)
+    simpleupdate!(peps, wts, ham, dt, Dcut2; bipartite=true, evolstep=30000, wtdiff_tol=tol)
 end
 # absort weight into site tensors
 absorb_wt!(peps, wts)
@@ -50,10 +45,20 @@ envs = leading_boundary(envs, peps, ctm_alg)
 # measure physical quantities
 rho1ss, rho2sss = calrho_all(envs, peps)
 result = measrho_all(rho1ss, rho2sss)
-display(result)
-@test isapprox(result["e_site"], -0.6633; atol=1e-3)
-@test isapprox(mean(result["mag_norm"]), 0.3972; atol=1e-3)
+@printf("Energy = %.8f\n", result["e_site"])
+@printf("Staggered magnetization = %.8f\n", mean(result["mag_norm"]))
+@test isapprox(result["e_site"], -0.6675; atol=1e-3)
+@test isapprox(mean(result["mag_norm"]), 0.3767; atol=1e-3)
 
-# full update energy and magnetization
-# e_fu = -0.6654
-# mag_fu = 0.3634
+# continue with full update
+dts = [2e-2, 1e-2, 5e-3, 1e-3, 5e-4]
+for dt in dts
+    fullupdate!(peps, envs, ham, dt, Dcut, χenv; rgmaxiter=5, cheap=true)
+end
+# measure physical quantities
+rho1ss, rho2sss = calrho_all(envs, peps)
+result = measrho_all(rho1ss, rho2sss)
+@printf("Energy = %.8f\n", result["e_site"])
+@printf("Staggered magnetization = %.8f\n", mean(result["mag_norm"]))
+@test isapprox(result["e_site"], -0.66875; atol=1e-4)
+@test isapprox(mean(result["mag_norm"]), 0.3510; atol=1e-3)
diff --git a/test/runtests.jl b/test/runtests.jl
index 07a4cd2f..b15ad1c6 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -53,7 +53,7 @@ end
         @time @safetestset "Heisenberg model" begin
             include("heisenberg.jl")
         end
-        @time @safetestset "Heisenberg model" begin
+        @time @safetestset "J1-J2 model" begin
             include("j1j2_model.jl")
         end
         @time @safetestset "P-wave superconductor" begin

From 18948737c5efd4d36a24a9f8248041af5824a827 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sun, 17 Nov 2024 17:32:29 +0800
Subject: [PATCH 16/20] update formatting

---
 src/states/suweight.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/states/suweight.jl b/src/states/suweight.jl
index 6995ab9a..a232b1c1 100644
--- a/src/states/suweight.jl
+++ b/src/states/suweight.jl
@@ -31,7 +31,7 @@ function Base.show(io::IO, wts::SUWeight)
     N1, N2 = size(wts)
     for (direction, r, c) in Iterators.product("xy", 1:N1, 1:N2)
         println(io, "$direction[$r,$c]: ")
-        wt = (direction == 'x' ? wts.x[r,c] : wts.y[r,c])
+        wt = (direction == 'x' ? wts.x[r, c] : wts.y[r, c])
         for (k, b) in blocks(wt)
             println(io, k, " = ", diag(b))
         end

From 98992980fb3cace482d440573bce50ccc4d869ff Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Mon, 18 Nov 2024 11:14:56 +0800
Subject: [PATCH 17/20] Refactor calc_convergence

Co-authored-by: Lukas Devos <ldevos98@gmail.com>
---
 src/algorithms/ctmrg/gaugefix.jl | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/algorithms/ctmrg/gaugefix.jl b/src/algorithms/ctmrg/gaugefix.jl
index 060db633..b22aa256 100644
--- a/src/algorithms/ctmrg/gaugefix.jl
+++ b/src/algorithms/ctmrg/gaugefix.jl
@@ -177,13 +177,9 @@ end
 Calculate convergence of CTMRG by comparing the singular values of CTM tensors
 """
 function calc_convergence(envsNew::CTMRGEnv, envsOld::CTMRGEnv)
-    CSNew = map(x -> tsvd(x)[2], envsNew.corners)
-    TSNew = map(x -> tsvd(x)[2], envsNew.edges)
     CSOld = map(x -> tsvd(x)[2], envsOld.corners)
     TSOld = map(x -> tsvd(x)[2], envsOld.edges)
-    ΔCS = maximum(_singular_value_distance, zip(CSOld, CSNew))
-    ΔTS = maximum(_singular_value_distance, zip(TSOld, TSNew))
-    return max(ΔCS, ΔTS)
+    return calc_convergence(envsNew, CSOld, TSOld)
 end
 
 @non_differentiable calc_convergence(args...)

From b4d43b5f8887d6c9d050de0313a355dcc9e60457 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Mon, 18 Nov 2024 11:19:50 +0800
Subject: [PATCH 18/20] Update sdiag_pow for latest TensorKit

Co-authored-by: Lukas Devos <ldevos98@gmail.com>
---
 src/utility/util.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/utility/util.jl b/src/utility/util.jl
index 5b61dd2b..8c240903 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -41,7 +41,7 @@ Compute S^(pow) for diagonal matrices `S`
 function sdiag_pow(S::AbstractTensorMap, pow::Real)
     S2 = similar(S)
     for (k, b) in blocks(S)
-        copyto!(blocks(S2)[k], diagm(diag(b) .^ pow))
+        copyto!(block(S2, k), diagm(diag(b) .^ pow))
     end
     return S2
 end

From 4bc9e46ba4034db3f512fc4fed3cab0b8cfb38ca Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Mon, 18 Nov 2024 12:35:07 +0800
Subject: [PATCH 19/20] Rename folder "timeevol" to "time_evolution"

---
 src/PEPSKit.jl                                              | 4 ++--
 src/algorithms/{timeevol => time_evolution}/fu_gaugefix.jl  | 0
 src/algorithms/{timeevol => time_evolution}/fu_optimize.jl  | 0
 src/algorithms/{timeevol => time_evolution}/fullupdate.jl   | 0
 src/algorithms/{timeevol => time_evolution}/simpleupdate.jl | 0
 5 files changed, 2 insertions(+), 2 deletions(-)
 rename src/algorithms/{timeevol => time_evolution}/fu_gaugefix.jl (100%)
 rename src/algorithms/{timeevol => time_evolution}/fu_optimize.jl (100%)
 rename src/algorithms/{timeevol => time_evolution}/fullupdate.jl (100%)
 rename src/algorithms/{timeevol => time_evolution}/simpleupdate.jl (100%)

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index 0726ca29..32f66b60 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -47,8 +47,8 @@ include("algorithms/ctmrg/sparse_environments.jl")
 include("algorithms/ctmrg/ctmrg.jl")
 include("algorithms/ctmrg/gaugefix.jl")
 
-include("algorithms/timeevol/simpleupdate.jl")
-include("algorithms/timeevol/fullupdate.jl")
+include("algorithms/time_evolution/simpleupdate.jl")
+include("algorithms/time_evolution/fullupdate.jl")
 
 include("algorithms/toolbox.jl")
 
diff --git a/src/algorithms/timeevol/fu_gaugefix.jl b/src/algorithms/time_evolution/fu_gaugefix.jl
similarity index 100%
rename from src/algorithms/timeevol/fu_gaugefix.jl
rename to src/algorithms/time_evolution/fu_gaugefix.jl
diff --git a/src/algorithms/timeevol/fu_optimize.jl b/src/algorithms/time_evolution/fu_optimize.jl
similarity index 100%
rename from src/algorithms/timeevol/fu_optimize.jl
rename to src/algorithms/time_evolution/fu_optimize.jl
diff --git a/src/algorithms/timeevol/fullupdate.jl b/src/algorithms/time_evolution/fullupdate.jl
similarity index 100%
rename from src/algorithms/timeevol/fullupdate.jl
rename to src/algorithms/time_evolution/fullupdate.jl
diff --git a/src/algorithms/timeevol/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
similarity index 100%
rename from src/algorithms/timeevol/simpleupdate.jl
rename to src/algorithms/time_evolution/simpleupdate.jl

From 1e97ec042fb64481211871bfb7575aff0b217ccb Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Mon, 18 Nov 2024 15:05:13 +0800
Subject: [PATCH 20/20] Replace `maxabs` by infinity norm

---
 src/algorithms/time_evolution/fullupdate.jl   | 10 +++++-----
 src/algorithms/time_evolution/simpleupdate.jl |  2 +-
 src/utility/util.jl                           | 14 --------------
 test/heisenberg_sufu.jl                       |  2 +-
 4 files changed, 7 insertions(+), 21 deletions(-)

diff --git a/src/algorithms/time_evolution/fullupdate.jl b/src/algorithms/time_evolution/fullupdate.jl
index 3eb85a4b..a8d01eb3 100644
--- a/src/algorithms/time_evolution/fullupdate.jl
+++ b/src/algorithms/time_evolution/fullupdate.jl
@@ -144,8 +144,8 @@ function update_column!(
             aR, bL, aR2bL2, env; maxiter=maxiter, maxdiff=maxdiff, verbose=false
         )
         costs[row] = cost
-        aR /= maxabs(aR)
-        bL /= maxabs(bL)
+        aR /= norm(aR, Inf)
+        bL /= norm(bL, Inf)
         localfid += local_fidelity(aR, bL, _combine_aRbL(aR0, bL0))
         #= update and normalize peps, ms
 
@@ -163,7 +163,7 @@ function update_column!(
         )
         # normalize
         for c_ in [col, cp1]
-            peps.A[row, c_] /= maxabs(peps.A[row, c_])
+            peps.A[row, c_] /= norm(peps.A[row, c_], Inf)
         end
     end
     # update CTMRGEnv
@@ -187,7 +187,7 @@ function fu_iter!(
     Dcut::Int,
     chi::Int,
     svderr::Float64=1e-9;
-    cheap=false,
+    cheap=true,
 )
     Nr, Nc = size(peps)
     fid, maxcost = 0.0, 0.0
@@ -230,7 +230,7 @@ function fullupdate!(
     rgtol::Float64=1e-6,
     rgmaxiter::Int=10,
     ctmrgscheme=:sequential,
-    cheap=false,
+    cheap=true,
 )
     time_start = time()
     N1, N2 = size(peps)
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 451ff8c9..6b5cc9b7 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -155,7 +155,7 @@ function _su_bondx!(
     # update tensor dict and weight on current bond 
     # (max element of weight is normalized to 1)
     peps.A[row, col], peps.A[row2, col2] = T1, T2
-    wts.x[row, col] = s / maxabs(s)
+    wts.x[row, col] = s / norm(s, Inf)
     return ϵ
 end
 
diff --git a/src/utility/util.jl b/src/utility/util.jl
index 8c240903..5b894bf9 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -21,20 +21,6 @@ function _elementwise_mult(a::AbstractTensorMap, b::AbstractTensorMap)
     return dst
 end
 
-"""
-Return the maximum absolute value of tensor elements
-"""
-function maxabs(t::AbstractTensorMap)
-    maxel = 0.0
-    for (k, b) in blocks(t)
-        maxelb = maximum(abs.(b))
-        if maxelb > maxel
-            maxel = maxelb
-        end
-    end
-    return maxel
-end
-
 """
 Compute S^(pow) for diagonal matrices `S`
 """
diff --git a/test/heisenberg_sufu.jl b/test/heisenberg_sufu.jl
index 86372e67..8786e8bc 100644
--- a/test/heisenberg_sufu.jl
+++ b/test/heisenberg_sufu.jl
@@ -23,7 +23,7 @@ wts = SUWeight(
 )
 # normalize peps
 for ind in CartesianIndices(peps.A)
-    peps.A[ind] /= PEPSKit.maxabs(peps.A[ind])
+    peps.A[ind] /= norm(peps.A[ind], Inf)
 end
 # Heisenberg model Hamiltonian
 ham = gen_gate()