From 79297a0617120546981d1a80cb1a49aa09cc5879 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 12 Dec 2025 16:49:26 +0800
Subject: [PATCH 01/13] Make SUWeight axis order definite

---
 src/environments/suweight.jl | 184 +++++++++++++++--------------------
 1 file changed, 78 insertions(+), 106 deletions(-)

diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index 75e3180bc..3055cd86c 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -9,31 +9,22 @@ const PEPSWeight{T, S} = AbstractTensorMap{T, S, 1, 1}
     struct SUWeight{E<:PEPSWeight}
 
 Schmidt bond weights used in simple/cluster update. 
-Weight elements are always real and non-negative.
-The domain and codomain of each weight matrix must be an un-dualed `ElementarySpace`.
+Each weight is a real and semi-positive definite
+`DiagonalTensorMap`, with the same codomain and domain.
 
-For a square lattice InfinitePEPS, the weights are placed as
-```
-        |
-    -T[r-1,c]-
-        |
-    wt[2,r,c]
-        |                   |
-    --T[r,c]--wt[1,r,c]--T[r,c+1]--
-        |                   |
-```
+On the square lattice,
+- `wt[1,r,c]` is on the x-bond between `[r,c]` and `[r,c+1]`;
+- `wt[2,r,c]` is on the y-bond between `[r,c]` and `[r-1,c]`.
 
 Axis order of each weight matrix is
 ```
-    x-weights:
-        1 ← x ← 2   or   2 → x → 1
-    
-    y-weights:
-        2           1
-        ↓           ↑  
-        y    or     y
-        ↓           ↑
-        1           2
+    x-weights:      y-weights:
+
+    1 - x - 2           2
+                        |
+                        y
+                        |
+                        1
 ```
 
 ## Fields
@@ -54,10 +45,6 @@ function SUWeight(data::Array{E, 3}) where {E <: PEPSWeight}
     for wt in data
         isa(wt, DiagonalTensorMap) ||
             error("Each weight matrix should be a DiagonalTensorMap")
-        domain(wt, 1) == codomain(wt, 1) ||
-            error("Domain and codomain of each weight matrix must be the same.")
-        !isdual(codomain(wt, 1)) ||
-            error("Domain and codomain of each weight matrix cannot be a dual space.")
         all(wt.data .>= 0) || error("Weight elements must be non-negative.")
     end
     return SUWeight{E}(data)
@@ -76,15 +63,12 @@ end
 """
     SUWeight(Nspaces::M, [Espaces::M]) where {M<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
 
-Create an SUWeight by specifying the vertical (north) or horizontal (east) virtual bond spaces.
+Create a trivial SUWeight by specifying the vertical (north) or horizontal (east) virtual bond spaces.
 Each individual space can be specified as either an `Int` or an `ElementarySpace`.
-The weights are initialized as identity matrices of element type `Float64`.
 """
 function SUWeight(
         Nspaces::M, Espaces::M = Nspaces
     ) where {M <: AbstractMatrix{<:Union{Int, ElementarySpace}}}
-    @assert all(!isdual, Nspaces)
-    @assert all(!isdual, Espaces)
     @assert size(Nspaces) == size(Espaces)
     Nr, Nc = size(Nspaces)
     weights = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
@@ -97,9 +81,8 @@ end
 """
     SUWeight(Nspace::S, Espace::S=Nspace; unitcell::Tuple{Int,Int}=(1, 1)) where {S<:ElementarySpace}
 
-Create an SUWeight by specifying its vertical (north) and horizontal (east) 
+Create an trivial SUWeight by specifying its vertical (north) and horizontal (east) 
 as `ElementarySpace`s) and unit cell size.
-The weights are initialized as identity matrices of element type `Float64`.
 """
 function SUWeight(
         Nspace::S, Espace::S = Nspace; unitcell::Tuple{Int, Int} = (1, 1)
@@ -110,18 +93,11 @@ end
 """
     SUWeight(peps::InfinitePEPS)
 
-Create an SUWeight for a given InfinitePEPS. 
-The weights are initialized as identity matrices of element type `Float64`.
+Create a trivial SUWeight for a given InfinitePEPS. 
 """
 function SUWeight(peps::InfinitePEPS)
-    Nspaces = map(peps.A) do t
-        V = domain(t, NORTH)
-        isdual(V) ? V' : V
-    end
-    Espaces = map(peps.A) do t
-        V = domain(t, EAST)
-        isdual(V) ? V' : V
-    end
+    Nspaces = collect(domain(t, NORTH) for t in peps.A)
+    Espaces = collect(domain(t, EAST) for t in peps.A)
     return SUWeight(Nspaces, Espaces)
 end
 
@@ -133,14 +109,8 @@ The weights are initialized as identity matrices of element type `Float64`.
 """
 function SUWeight(pepo::InfinitePEPO)
     @assert size(pepo, 3) == 1
-    Nspaces = map(@view(pepo.A[:, :, 1])) do t
-        V = north_virtualspace(t)
-        isdual(V) ? V' : V
-    end
-    Espaces = map(@view(pepo.A[:, :, 1])) do t
-        V = east_virtualspace(t)
-        isdual(V) ? V' : V
-    end
+    Nspaces = collect(domain(t, NORTH) for t in @view(pepo.A[:, :, 1]))
+    Espaces = collect(domain(t, EAST) for t in @view(pepo.A[:, :, 1]))
     return SUWeight(Nspaces, Espaces)
 end
 
@@ -182,7 +152,8 @@ end
 function Base.show(io::IO, ::MIME"text/plain", wts::SUWeight)
     println(io, typeof(wts))
     for idx in CartesianIndices(wts.data)
-        println(io, Tuple(idx), ":")
+        print(io, Tuple(idx), ": ")
+        println(space(wts.data[idx]))
         for (k, b) in blocks(wts.data[idx])
             println(io, k, " = ", diag(b))
         end
@@ -194,9 +165,9 @@ end
     absorb_weight(t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight, row::Int, col::Int, ax::Int; inv::Bool = false)
     absorb_weight(t::Union{PEPSTensor, PEPOTensor}, weights::SUWeight, row::Int, col::Int, ax::NTuple{N, Int}; inv::Bool = false)
 
-Absorb or remove environment weight on an axis of tensor `t` known to be located at
-position (`row`, `col`) in the unit cell of an InfinitePEPS or InfinitePEPO. 
-Weights around the tensor at `(row, col)` are
+Absorb or remove (in a twist-free way) the square root of environment weight 
+on an axis of the PEPS/PEPO tensor `t` known to be at position (`row`, `col`)
+in the unit cell of an InfinitePEPS/InfinitePEPO. The involved weights are
 ```
                     |
                 [2,r,c]
@@ -209,7 +180,7 @@ Weights around the tensor at `(row, col)` are
 
 ## Arguments
 
-- `t::PT` : PEPSTensor or PEPOTensor to which the weight will be absorbed. 
+- `t::Union{PEPSTensor, PEPOTensor}` : PEPSTensor or PEPOTensor to which the weight will be absorbed. 
 - `weights::SUWeight` : All simple update weights.
 - `row::Int` : The row index specifying the position in the tensor network.
 - `col::Int` : The column index specifying the position in the tensor network.
@@ -252,11 +223,9 @@ function absorb_weight(
     )
     t_idx = [(n - nout == ax) ? 1 : -n for n in 1:ntol]
     ax′ = ax + nout
-    wt_idx = if isdual(space(t, ax′))
-        [1, -ax′] # t ← wt
-    else
-        [-ax′, 1] # t → wt
-    end
+    wt_idx = (ax == NORTH || ax == EAST) ? [1, -ax′] : [-ax′, 1]
+    # make absorption/removal twist-free
+    twistdual!(wt, 1)
     return permute(ncon((t, wt), (t_idx, wt_idx)), (Tuple(1:nout), Tuple((nout + 1):ntol)))
 end
 function absorb_weight(
@@ -278,13 +247,13 @@ end
             y₁₁       y₁₂       y₁₃
             |         |         |
     ..x₁₃...┼---x₁₁---┼---x₁₂---┼---x₁₃---
-            |         |         |
+            |         |         |           2
             y₂₁       y₂₂       y₂₃         |
             |         |         |           y
     ..x₂₃...┼---x₂₁---┼---x₂₂---┼---x₂₃---  |
-            |         |         |
-            y₃₁       y₃₂       y₃₃         -- x --
-            |         |         |
+            |         |         |           1
+            y₃₁       y₃₂       y₃₃
+            |         |         |           1 -- x -- 2
     ..x₃₃...┼---x₃₁---┼---x₃₂---┼---x₃₃---
             :         :         :
             y₁₁       y₁₂       y₁₃
@@ -297,19 +266,20 @@ end
             x₁₃       x₂₃       x₃₃
             |         |         |
     --y₁₃---┼---y₂₃---┼---y₃₃---┼...y₁₃...
-            |         |         |
+            |         |         |                   2
             x₁₂       x₂₂       x₃₂                 |
             |         |         |                   x
     --y₁₂---┼---y₂₂---┼---y₃₂---┼...y₁₂...          |
-            |         |         |
-            x₁₁       x₂₁       x₃₁         -- y --
-            |         |         |
+            |         |         |                   1
+            x₁₁       x₂₁       x₃₁
+            |         |         |           2 -- y -- 1
     --y₁₁---┼---y₂₁---┼---y₃₁---┼...y₁₁...
             :         :         :
             x₁₃       x₂₃       x₃₃
             :         :         :
     ```
     - x/y-weights are exchanged.
+    - need to further transpose x-weights.
     - need to further move 1st column of x-weights to the last column.
 
 - `rotr90`:
@@ -319,18 +289,19 @@ end
             :         :         :
     ..y₁₁...┼---y₃₁---┼---y₂₁---┼---y₁₁---
             |         |         |
-            x₃₁       x₂₁       x₁₁         -- y --
-            |         |         |
-    ..y₁₂...┼---y₃₂---┼---y₂₂---┼---y₁₂---  |
-            |         |         |           x
-            x₃₂       x₂₂       x₁₂         |
+            x₃₁       x₂₁       x₁₁         1 -- y -- 2
             |         |         |
-    ..y₁₃...┼---y₃₃---┼---y₂₃---┼---y₁₃---
+    ..y₁₂...┼---y₃₂---┼---y₂₂---┼---y₁₂---  1
+            |         |         |           |
+            x₃₂       x₂₂       x₁₂         x
+            |         |         |           |
+    ..y₁₃...┼---y₃₃---┼---y₂₃---┼---y₁₃---  2
             |         |         |
             x₃₃       x₂₃       x₁₃
             |         |         |
     ```
     - x/y-weights are exchanged.
+    - need to further transpose y-weights.
     - need to further move last row of y-weights to the 1st row.
 
 - `rot180`:
@@ -340,40 +311,60 @@ end
             :         :         :
     --x₃₃---┼---x₃₂---┼---x₃₁---┼...x₃₃...
             |         |         |
-            y₃₃       y₃₂       y₃₁         -- x --
+            y₃₃       y₃₂       y₃₁        2 -- x -- 1
             |         |         |
-    --x₂₃---┼---x₂₂---┼---x₂₁---┼...x₂₃...          |
-            |         |         |                   y
-            y₂₃       y₂₂       y₂₁                 |
-            |         |         |
-    --x₁₃---┼---x₁₂---┼---x₁₁---┼...x₁₃...
+    --x₂₃---┼---x₂₂---┼---x₂₁---┼...x₂₃...          1
+            |         |         |                   |
+            y₂₃       y₂₂       y₂₁                 y
+            |         |         |                   |
+    --x₁₃---┼---x₁₂---┼---x₁₁---┼...x₁₃...          2
             |         |         |
             y₁₃       y₁₂       y₁₁
             |         |         |
     ```
+    - need to transpose all weights.
     - need to move 1st column of x-weights to the last column.
     - need to move last row of y-weights to the 1st row.
 =#
 
 function Base.rotl90(wts::SUWeight)
     wts_x = circshift(rotl90(wts[2, :, :]), (0, -1))
+    for (i, wt) in enumerate(wts_x)
+        wts_x[i] = DiagonalTensorMap(transpose(wt; copy = true))
+    end
     wts_y = rotl90(wts[1, :, :])
     return SUWeight(wts_x, wts_y)
 end
 function Base.rotr90(wts::SUWeight)
     wts_x = rotr90(wts[2, :, :])
     wts_y = circshift(rotr90(wts[1, :, :]), (1, 0))
+    for (i, wt) in enumerate(wts_y)
+        wts_y[i] = DiagonalTensorMap(transpose(wt; copy = true))
+    end
     return SUWeight(wts_x, wts_y)
 end
 function Base.rot180(wts::SUWeight)
     wts_x = circshift(rot180(wts[1, :, :]), (0, -1))
     wts_y = circshift(rot180(wts[2, :, :]), (1, 0))
+    for (i, wt) in enumerate(wts_x)
+        wts_x[i] = DiagonalTensorMap(transpose(wt; copy = true))
+    end
+    for (i, wt) in enumerate(wts_y)
+        wts_y[i] = DiagonalTensorMap(transpose(wt; copy = true))
+    end
     return SUWeight(wts_x, wts_y)
 end
 
-function _CTMRGEnv(wts::SUWeight, flips::Array{Bool, 3})
-    @assert size(wts) == size(flips)
+"""
+    CTMRGEnv(wts::SUWeight)
+
+Construct a CTMRG environment with bond dimension χ = 1 from SUWeight `wts`.
+The scalartype of the returned environment is always `Float64`.
+"""
+function CTMRGEnv(wts::SUWeight)
     _, Nr, Nc = size(wts)
+    S = sectortype(wts)
+    V_env = Vect[S](one(S) => 1)
     edges = map(Iterators.product(1:4, 1:Nr, 1:Nc)) do (d, r, c)
         wt_idx = if d == NORTH
             CartesianIndex(2, _next(r, Nr), c)
@@ -384,35 +375,16 @@ function _CTMRGEnv(wts::SUWeight, flips::Array{Bool, 3})
         else # WEST
             CartesianIndex(1, r, c)
         end
+        # temporarily make wt axis order ([bra], [ket])
         wt = deepcopy(wts[wt_idx])
-        if (flips[wt_idx] && d in (NORTH, EAST)) || (!flips[wt_idx] && d in (SOUTH, WEST))
-            wt = permute(wt, ((2,), (1,)))
+        if d in (NORTH, EAST)
+            wt = transpose(wt)
         end
-        twistdual!(wt, 1)
-        wt = insertrightunit(insertleftunit(wt, 1))
-        return permute(wt, ((1, 2, 3), (4,)))
+        # attach identity on environment space
+        return permute(wt ⊗ TensorKit.id(Float64, V_env), ((2, 3, 1), (4,)))
     end
     corners = map(CartesianIndices(edges)) do idx
-        return TensorKit.id(Float64, codomain(edges[idx], 1))
+        return TensorKit.id(Float64, V_env)
     end
     return CTMRGEnv(corners, edges)
 end
-
-"""
-    CTMRGEnv(wts::SUWeight, peps::InfinitePEPS)
-
-Construct a CTMRG environment with bond dimension χ = 1 
-for an InfinitePEPS `peps` from the accompanying SUWeight `wts`.
-The scalartype of the returned environment is always `Float64`.
-"""
-function CTMRGEnv(wts::SUWeight, peps::InfinitePEPS)
-    _, Nr, Nc = size(wts)
-    flips = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
-        return if d == 1
-            isdual(domain(peps[r, c], EAST))
-        else
-            isdual(domain(peps[r, c], NORTH))
-        end
-    end
-    return _CTMRGEnv(wts, flips)
-end

From 638e9d5d388b0ded7d6491aaf544b1fb137a7683 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 12 Dec 2025 16:57:12 +0800
Subject: [PATCH 02/13] Update related tests

---
 test/ctmrg/suweight.jl | 37 ++++++++++++++++++-------
 test/types/suweight.jl | 61 +++++++++++++++---------------------------
 2 files changed, 50 insertions(+), 48 deletions(-)

diff --git a/test/ctmrg/suweight.jl b/test/ctmrg/suweight.jl
index 1945bd3b2..4a195997b 100644
--- a/test/ctmrg/suweight.jl
+++ b/test/ctmrg/suweight.jl
@@ -2,7 +2,18 @@ using Test
 using Random
 using TensorKit
 using PEPSKit
-using PEPSKit: str
+using PEPSKit: str, twistdual, _prev, _next
+
+Vps = Dict(
+    Z2Irrep => Vect[Z2Irrep](0 => 1, 1 => 2),
+    U1Irrep => Vect[U1Irrep](0 => 2, 1 => 2, -1 => 1),
+    FermionParity => Vect[FermionParity](0 => 1, 1 => 2),
+)
+Vvs = Dict(
+    Z2Irrep => Vect[Z2Irrep](0 => 2, 1 => 2),
+    U1Irrep => Vect[U1Irrep](0 => 3, 1 => 1, -1 => 2),
+    FermionParity => Vect[FermionParity](0 => 2, 1 => 2),
+)
 
 function rand_wts(peps::InfinitePEPS)
     wts = SUWeight(peps)
@@ -13,27 +24,35 @@ function rand_wts(peps::InfinitePEPS)
     return wts
 end
 
-function su_rdm_1x1(row::Int, col::Int, peps::InfinitePEPS, wts::SUWeight)
+function su_rdm_1x1(
+        row::Int, col::Int, peps::InfinitePEPS, wts::Union{Nothing, SUWeight} = nothing
+    )
     Nr, Nc = size(peps)
     @assert 1 <= row <= Nr && 1 <= col <= Nc
     t = peps.A[row, col]
-    t = absorb_weight(t, wts, row, col, Tuple(1:4))
-    ρ = t * t'
+    if !(wts === nothing)
+        t = absorb_weight(t, wts, row, col, Tuple(1:4))
+    end
+    # contract local ⟨t|t⟩ without virtual twists
+    @tensor ρ[k; b] := conj(t[b; n e s w]) * twistdual(t, 2:5)[k; n e s w]
     return ρ / str(ρ)
 end
 
-Vps = [Vect[U1Irrep](0 => 2, 1 => 2, -1 => 1), Vect[FermionParity](0 => 1, 1 => 2)]
-Vvs = [Vect[U1Irrep](0 => 3, 1 => 1, -1 => 2), Vect[FermionParity](0 => 2, 1 => 2)]
+@testset "SUWeight ($(init) init, $(sect))" for (init, sect) in
+    Iterators.product([:trivial, :random], keys(Vps))
 
-for (Vp, Vv) in zip(Vps, Vvs)
+    Vp, Vv = Vps[sect], Vvs[sect]
     Nspaces = [Vv Vv' Vv; Vv' Vv Vv']
     Espaces = [Vv Vv Vv'; Vv Vv' Vv']
     Pspaces = fill(Vp, size(Nspaces))
     peps = InfinitePEPS(randn, ComplexF64, Pspaces, Nspaces, Espaces)
-    wts = rand_wts(peps)
-    env = CTMRGEnv(wts, peps)
+    wts = (init == :trivial) ? SUWeight(peps) : rand_wts(peps)
+    env = CTMRGEnv(wts)
     for (r, c) in Tuple.(CartesianIndices(peps.A))
         ρ1 = su_rdm_1x1(r, c, peps, wts)
+        if init == :trivial
+            @test ρ1 ≈ su_rdm_1x1(r, c, peps, nothing)
+        end
         ρ2 = reduced_densitymatrix(((r, c),), peps, env)
         @test ρ1 ≈ ρ2
     end
diff --git a/test/types/suweight.jl b/test/types/suweight.jl
index e5b142485..b9ee3d368 100644
--- a/test/types/suweight.jl
+++ b/test/types/suweight.jl
@@ -2,6 +2,15 @@ using Test
 using TensorKit
 using PEPSKit
 
+function rand_wts(state::Union{InfinitePEPS, InfinitePEPO})
+    wts = SUWeight(state)
+    for idx in CartesianIndices(wts.data)
+        n = length(wts.data[idx].data)
+        wts.data[idx].data[:] = sort(rand(Float64, n); lt = !isless)
+    end
+    return wts
+end
+
 function compose_n(f, n)
     n == 0 && return identity
     return f ∘ compose_n(f, n - 1)
@@ -18,43 +27,20 @@ function test_rotation(wts::SUWeight)
     return nothing
 end
 
-function test_rotation(peps::InfinitePEPS, wts::SUWeight)
-    for n in 1:4
-        rot = compose_n(rotl90, n)
-        A1 = InfinitePEPS(
-            collect(
-                absorb_weight(peps.A[idx], wts, idx[1], idx[2], Tuple(1:4)) for
-                    idx in CartesianIndices(peps.A)
-            ),
-        )
-        peps2, wts2 = rot(peps), rot(wts)
-        A2 = InfinitePEPS(
-            collect(
-                absorb_weight(peps2.A[idx], wts2, idx[1], idx[2], Tuple(1:4)) for
-                    idx in CartesianIndices(peps2.A)
-            ),
-        )
-        @test A2 ≈ rot(A1)
-    end
-    return
-end
-
-function test_rotation(pepo::InfinitePEPO, wts::SUWeight)
-    @assert size(pepo, 3) == 1
+function test_rotation(state::Union{InfinitePEPS, InfinitePEPO}, wts::SUWeight)
+    InfiniteState = (state isa InfinitePEPS) ? InfinitePEPS : InfinitePEPO
+    A1 = InfiniteState(
+        map(CartesianIndices(state.A)) do idx
+            return absorb_weight(state.A[idx], wts, idx[1], idx[2], Tuple(1:4))
+        end
+    )
     for n in 1:4
         rot = compose_n(rotl90, n)
-        A1 = InfinitePEPO(
-            collect(
-                absorb_weight(pepo.A[idx], wts, idx[1], idx[2], Tuple(1:4)) for
-                    idx in CartesianIndices(pepo.A)
-            ),
-        )
-        pepo2, wts2 = rot(pepo), rot(wts)
-        A2 = InfinitePEPO(
-            collect(
-                absorb_weight(pepo2.A[idx], wts2, idx[1], idx[2], Tuple(1:4)) for
-                    idx in CartesianIndices(pepo2.A)
-            ),
+        state2, wts2 = rot(state), rot(wts)
+        A2 = InfiniteState(
+            map(CartesianIndices(state2.A)) do idx
+                return absorb_weight(state2.A[idx], wts2, idx[1], idx[2], Tuple(1:4))
+            end
         )
         @test A2 ≈ rot(A1)
     end
@@ -71,10 +57,7 @@ Vs = (
 Nr, Nc = 2, 3
 peps = InfinitePEPS(rand, Float64, Vphy, Vs[2], Vs[1]'; unitcell = (Nr, Nc))
 pepo = InfinitePEPO(rand, Float64, Vphy, Vs[2], Vs[1]'; unitcell = (Nr, Nc, 1))
-wts = collect(
-    svd_compact(rand(Float64, Vs[dir] ← Vs[dir]))[2] for dir in 1:2, r in 1:Nr, c in 1:Nc
-)
-wts = SUWeight(wts)
+wts = rand_wts(peps)
 
 @test sectortype(wts) === sectortype(Vs[1])
 @test spacetype(wts) === spacetype(Vs[1])

From 259c97f4b66087374726136b9aa8ebcaa178d79b Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 12 Dec 2025 17:06:10 +0800
Subject: [PATCH 03/13] Update docs on CTMRGEnv(wts)

---
 docs/src/examples/hubbard_su/index.md   |   2 +-
 docs/src/examples/hubbard_su/main.ipynb | 114 ++++++++++++------------
 examples/hubbard_su/main.jl             |   2 +-
 3 files changed, 59 insertions(+), 59 deletions(-)

diff --git a/docs/src/examples/hubbard_su/index.md b/docs/src/examples/hubbard_su/index.md
index 6b347c15d..5af3f5045 100644
--- a/docs/src/examples/hubbard_su/index.md
+++ b/docs/src/examples/hubbard_su/index.md
@@ -130,7 +130,7 @@ is achieved by using `trunc=truncrank(χ)` with different `χ`s in the CTMRG run
 χenv₀, χenv = 6, 16
 env_space = Vect[fℤ₂](0 => χenv₀ / 2, 1 => χenv₀ / 2)
 normalize!.(peps.A, Inf)
-env = CTMRGEnv(wts, peps)
+env = CTMRGEnv(wts)
 for χ in [χenv₀, χenv]
     global env, = leading_boundary(
         env, peps; alg = :sequential, tol = 1.0e-8, maxiter = 50, trunc = truncrank(χ)
diff --git a/docs/src/examples/hubbard_su/main.ipynb b/docs/src/examples/hubbard_su/main.ipynb
index a15f8199b..0e4fb5280 100644
--- a/docs/src/examples/hubbard_su/main.ipynb
+++ b/docs/src/examples/hubbard_su/main.ipynb
@@ -1,16 +1,17 @@
 {
  "cells": [
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "using Markdown #hide"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "# Simple update for the Fermi-Hubbard model at half-filling\n",
     "\n",
@@ -26,46 +27,46 @@
     "with $\\sigma \\in \\{\\uparrow,\\downarrow\\}$ and $n_{i,\\sigma} = c_{i,\\sigma}^+ c_{i,\\sigma}^-$.\n",
     "\n",
     "Let's get started by seeding the RNG and importing the required modules:"
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "using Random\n",
     "using TensorKit, PEPSKit\n",
     "Random.seed!(12329348592498);"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## Defining the Hamiltonian\n",
     "\n",
     "First, we define the Hubbard model at $t=1$ hopping and $U=6$ using `Trivial` sectors for\n",
     "the particle and spin symmetries, and set $\\mu = U/2$ for half-filling. The model will be\n",
     "constructed on a $2 \\times 2$ unit cell, so we have:"
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "t = 1\n",
     "U = 6\n",
     "Nr, Nc = 2, 2\n",
     "H = hubbard_model(Float64, Trivial, Trivial, InfiniteSquare(Nr, Nc); t, U, mu = U / 2);\n",
     "physical_space = Vect[fℤ₂](0 => 2, 1 => 2);"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## Running the simple update algorithm\n",
     "\n",
@@ -78,33 +79,34 @@
     "First, we shall use a small D for the random PEPS initialization, which is chosen as 4 here.\n",
     "For convenience, here we work with real tensors with `Float64` entries.\n",
     "The bond weights are still initialized as identity matrices."
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "virtual_space = Vect[fℤ₂](0 => 2, 1 => 2)\n",
     "peps = InfinitePEPS(rand, Float64, physical_space, virtual_space; unitcell = (Nr, Nc));\n",
     "wts = SUWeight(peps);"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "Starting from the random state, we first use a relatively large evolution time step\n",
     "`dt = 1e-2`. After convergence at D = 4, to avoid stucking at some bad local minimum,\n",
     "we first increase D to 12, and drop it back to D = 8 after a while.\n",
     "Afterwards, we keep D = 8 and gradually decrease `dt` to `1e-4` to improve convergence."
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "dts = [1.0e-2, 1.0e-2, 1.0e-3, 4.0e-4, 1.0e-4]\n",
     "tols = [1.0e-7, 1.0e-7, 1.0e-8, 1.0e-8, 1.0e-8]\n",
@@ -116,12 +118,11 @@
     "    alg = SimpleUpdate(; trunc, bipartite = false)\n",
     "    global peps, wts, = time_evolve(peps, H, dt, maxiter, alg, wts; tol, check_interval = 2000)\n",
     "end"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## Computing the ground-state energy\n",
     "\n",
@@ -130,87 +131,86 @@
     "which is initialized using the simple update bond weights. Next we use it to initialize\n",
     "another run with bigger environment dimension. The dynamic adjustment of environment dimension\n",
     "is achieved by using `trunc=truncrank(χ)` with different `χ`s in the CTMRG runs:"
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "χenv₀, χenv = 6, 16\n",
     "env_space = Vect[fℤ₂](0 => χenv₀ / 2, 1 => χenv₀ / 2)\n",
     "normalize!.(peps.A, Inf)\n",
-    "env = CTMRGEnv(wts, peps)\n",
+    "env = CTMRGEnv(wts)\n",
     "for χ in [χenv₀, χenv]\n",
     "    global env, = leading_boundary(\n",
     "        env, peps; alg = :sequential, tol = 1.0e-8, maxiter = 50, trunc = truncrank(χ)\n",
     "    )\n",
     "end"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "We measure the energy by computing the `H` expectation value, where we have to make sure to\n",
     "normalize with respect to the unit cell to obtain the energy per site:"
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "E = expectation_value(peps, H, env) / (Nr * Nc)\n",
     "@show E;"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "Finally, we can compare the obtained ground-state energy against the literature, namely the\n",
     "QMC estimates from [Qin et al.](@cite qin_benchmark_2016). We find that the results generally\n",
     "agree:"
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "outputs": [],
    "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "Es_exact = Dict(0 => -1.62, 2 => -0.176, 4 => 0.8603, 6 => -0.6567, 8 => -0.5243)\n",
     "E_exact = Es_exact[U] - U / 2\n",
     "@show (E - E_exact) / abs(E_exact);"
-   ],
-   "metadata": {},
-   "execution_count": null
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "---\n",
     "\n",
     "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
-   ],
-   "metadata": {}
+   ]
   }
  ],
- "nbformat_minor": 3,
  "metadata": {
+  "kernelspec": {
+   "display_name": "Julia 1.11.7",
+   "language": "julia",
+   "name": "julia-1.11"
+  },
   "language_info": {
    "file_extension": ".jl",
    "mimetype": "application/julia",
    "name": "julia",
    "version": "1.11.7"
-  },
-  "kernelspec": {
-   "name": "julia-1.11",
-   "display_name": "Julia 1.11.7",
-   "language": "julia"
   }
  },
- "nbformat": 4
-}
\ No newline at end of file
+ "nbformat": 4,
+ "nbformat_minor": 3
+}
diff --git a/examples/hubbard_su/main.jl b/examples/hubbard_su/main.jl
index 786dab680..2bb0f411c 100644
--- a/examples/hubbard_su/main.jl
+++ b/examples/hubbard_su/main.jl
@@ -83,7 +83,7 @@ is achieved by using `trunc=truncrank(χ)` with different `χ`s in the CTMRG run
 χenv₀, χenv = 6, 16
 env_space = Vect[fℤ₂](0 => χenv₀ / 2, 1 => χenv₀ / 2)
 normalize!.(peps.A, Inf)
-env = CTMRGEnv(wts, peps)
+env = CTMRGEnv(wts)
 for χ in [χenv₀, χenv]
     global env, = leading_boundary(
         env, peps; alg = :sequential, tol = 1.0e-8, maxiter = 50, trunc = truncrank(χ)

From 2b3a5322540d74c45b40359f9c6ff24538ad76f0 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 12 Dec 2025 17:37:48 +0800
Subject: [PATCH 04/13] Update SU and bond truncation

---
 src/algorithms/time_evolution/evoltools.jl     |  2 +-
 .../time_evolution/simpleupdate3site.jl        |  2 +-
 src/algorithms/truncation/bond_truncation.jl   |  4 ++--
 src/utility/util.jl                            | 18 +-----------------
 4 files changed, 5 insertions(+), 21 deletions(-)

diff --git a/src/algorithms/time_evolution/evoltools.jl b/src/algorithms/time_evolution/evoltools.jl
index 40ceee032..146ede367 100644
--- a/src/algorithms/time_evolution/evoltools.jl
+++ b/src/algorithms/time_evolution/evoltools.jl
@@ -233,7 +233,7 @@ function _apply_gate(
 
     a, b = absorb_s(a, s, b)
     if need_flip
-        a, s, b = flip_svd(a, s, b)
+        a, s, b = flip(a, numind(a)), _fliptwist_s(s), flip(b, 1)
     end
     return a, s, b, ϵ
 end
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index d72d61b8a..895df8dca 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -479,7 +479,7 @@ function _su3site_se!(
             Ms = [first(_unfuse_physicalspace(M, Vphy)) for (M, Vphy) in zip(Ms, Vphys)]
         end
         for (wt, wt_idx, flip) in zip(wts, wt_idxs, flips)
-            env[CartesianIndex(wt_idx)] = normalize(flip ? _flip_s(wt) : wt, Inf)
+            env[CartesianIndex(wt_idx)] = normalize(flip ? _fliptwist_s(wt) : wt, Inf)
         end
     end
     # restore virtual arrows in `Ms`
diff --git a/src/algorithms/truncation/bond_truncation.jl b/src/algorithms/truncation/bond_truncation.jl
index d4d1fde07..9d70947b1 100644
--- a/src/algorithms/truncation/bond_truncation.jl
+++ b/src/algorithms/truncation/bond_truncation.jl
@@ -139,7 +139,7 @@ function bond_truncate(
     a, s, b = svd_trunc!(permute(_combine_ab(a, b), perm_ab); trunc = alg.trunc)
     a, b = absorb_s(a, s, b)
     if need_flip
-        a, s, b = flip_svd(a, s, b)
+        a, s, b = flip(a, numind(a)), _fliptwist_s(s), flip(b, 1)
     end
     return a, s, b, (; fid, Δfid, Δs)
 end
@@ -186,7 +186,7 @@ function bond_truncate(
     @tensor a[-1 -2; -3] := Qa[-1 -2 3] * u[3 -3]
     @tensor b[-1; -2 -3] := vh[-1 1] * Qb[1 -2 -3]
     if need_flip
-        a, s, b = flip_svd(a, s, b)
+        a, s, b = flip(a, numind(a)), _fliptwist_s(s), flip(b, 1)
     end
     return a, s, b, info
 end
diff --git a/src/utility/util.jl b/src/utility/util.jl
index f5c5c31f3..a63a2a97f 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -100,23 +100,7 @@ function absorb_s(U::AbstractTensorMap, S::DiagonalTensorMap, V::AbstractTensorM
     return U * sqrt_S, sqrt_S * V
 end
 
-"""
-    flip_svd(u::AbstractTensorMap, s::DiagonalTensorMap, vh::AbstractTensorMap)
-
-Given SVD result `u ← s ← vh`, flip the arrow between the three tensors 
-to `u2 → s2 → vh2` such that
-```
-    u * s * vh = (@tensor t2[-1, ...; -2, ...] := u2[-1, ...; 2] * s2[1; 2] * vh2[1; -2, ...])
-```
-The axis orders for `s`, `s2` are
-```
-    1 ← s ← 2   2 → s2 → 1
-```
-"""
-function flip_svd(u::AbstractTensorMap, s::DiagonalTensorMap, vh::AbstractTensorMap)
-    return flip(u, numind(u)), _flip_s(s), flip(vh, 1)
-end
-_flip_s(s::DiagonalTensorMap) = permute(DiagonalTensorMap(flip(s, (1, 2))), ((2,), (1,)))
+_fliptwist_s(s::DiagonalTensorMap) = twist!(DiagonalTensorMap(flip(s, 1:2)), 1)
 
 """
     twistdual(t::AbstractTensorMap, i)

From 6ec6239ec6a890c8d5b70bbaf83f8f058c5e6f22 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 12 Dec 2025 17:43:00 +0800
Subject: [PATCH 05/13] Test SU with non-standard virtual dualness

---
 test/timeevol/cluster_projectors.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index f232c5df4..035ef3826 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -111,7 +111,7 @@ end
     Espace = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 8, (1, 1 // 2) => 4, (1, -1 // 2) => 4)
     trunc_env0 = truncerror(; atol = 1.0e-12) & truncrank(4)
     trunc_env = truncerror(; atol = 1.0e-12) & truncrank(16)
-    peps = InfinitePEPS(rand, Float64, Pspace, Vspace; unitcell = (Nr, Nc))
+    peps = InfinitePEPS(rand, Float64, Pspace, Vspace, Vspace'; unitcell = (Nr, Nc))
     wts = SUWeight(peps)
     ham = real(
         hubbard_model(
@@ -127,7 +127,7 @@ end
         peps, wts, = time_evolve(peps, ham, dt, 10000, alg, wts; tol, check_interval = 1000)
     end
     normalize!.(peps.A, Inf)
-    env = CTMRGEnv(wts, peps)
+    env = CTMRGEnv(wts)
     env, = leading_boundary(env, peps; tol = ctmrg_tol, trunc = trunc_env0)
     env, = leading_boundary(env, peps; tol = ctmrg_tol, trunc = trunc_env)
     e_site = cost_function(peps, env, ham) / (Nr * Nc)

From f230f2fe6cc83aa057e81f4ff71915b04b26dfa4 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 12 Dec 2025 17:49:07 +0800
Subject: [PATCH 06/13] Undo metadata changes

---
 docs/src/examples/hubbard_su/main.ipynb | 112 ++++++++++++------------
 1 file changed, 56 insertions(+), 56 deletions(-)

diff --git a/docs/src/examples/hubbard_su/main.ipynb b/docs/src/examples/hubbard_su/main.ipynb
index 0e4fb5280..4dbda8971 100644
--- a/docs/src/examples/hubbard_su/main.ipynb
+++ b/docs/src/examples/hubbard_su/main.ipynb
@@ -1,17 +1,16 @@
 {
  "cells": [
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "using Markdown #hide"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "# Simple update for the Fermi-Hubbard model at half-filling\n",
     "\n",
@@ -27,46 +26,46 @@
     "with $\\sigma \\in \\{\\uparrow,\\downarrow\\}$ and $n_{i,\\sigma} = c_{i,\\sigma}^+ c_{i,\\sigma}^-$.\n",
     "\n",
     "Let's get started by seeding the RNG and importing the required modules:"
-   ]
+   ],
+   "metadata": {}
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "using Random\n",
     "using TensorKit, PEPSKit\n",
     "Random.seed!(12329348592498);"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "## Defining the Hamiltonian\n",
     "\n",
     "First, we define the Hubbard model at $t=1$ hopping and $U=6$ using `Trivial` sectors for\n",
     "the particle and spin symmetries, and set $\\mu = U/2$ for half-filling. The model will be\n",
     "constructed on a $2 \\times 2$ unit cell, so we have:"
-   ]
+   ],
+   "metadata": {}
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "t = 1\n",
     "U = 6\n",
     "Nr, Nc = 2, 2\n",
     "H = hubbard_model(Float64, Trivial, Trivial, InfiniteSquare(Nr, Nc); t, U, mu = U / 2);\n",
     "physical_space = Vect[fℤ₂](0 => 2, 1 => 2);"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "## Running the simple update algorithm\n",
     "\n",
@@ -79,34 +78,33 @@
     "First, we shall use a small D for the random PEPS initialization, which is chosen as 4 here.\n",
     "For convenience, here we work with real tensors with `Float64` entries.\n",
     "The bond weights are still initialized as identity matrices."
-   ]
+   ],
+   "metadata": {}
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "virtual_space = Vect[fℤ₂](0 => 2, 1 => 2)\n",
     "peps = InfinitePEPS(rand, Float64, physical_space, virtual_space; unitcell = (Nr, Nc));\n",
     "wts = SUWeight(peps);"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "Starting from the random state, we first use a relatively large evolution time step\n",
     "`dt = 1e-2`. After convergence at D = 4, to avoid stucking at some bad local minimum,\n",
     "we first increase D to 12, and drop it back to D = 8 after a while.\n",
     "Afterwards, we keep D = 8 and gradually decrease `dt` to `1e-4` to improve convergence."
-   ]
+   ],
+   "metadata": {}
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "dts = [1.0e-2, 1.0e-2, 1.0e-3, 4.0e-4, 1.0e-4]\n",
     "tols = [1.0e-7, 1.0e-7, 1.0e-8, 1.0e-8, 1.0e-8]\n",
@@ -118,11 +116,12 @@
     "    alg = SimpleUpdate(; trunc, bipartite = false)\n",
     "    global peps, wts, = time_evolve(peps, H, dt, maxiter, alg, wts; tol, check_interval = 2000)\n",
     "end"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "## Computing the ground-state energy\n",
     "\n",
@@ -131,13 +130,12 @@
     "which is initialized using the simple update bond weights. Next we use it to initialize\n",
     "another run with bigger environment dimension. The dynamic adjustment of environment dimension\n",
     "is achieved by using `trunc=truncrank(χ)` with different `χ`s in the CTMRG runs:"
-   ]
+   ],
+   "metadata": {}
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "χenv₀, χenv = 6, 16\n",
     "env_space = Vect[fℤ₂](0 => χenv₀ / 2, 1 => χenv₀ / 2)\n",
@@ -148,69 +146,71 @@
     "        env, peps; alg = :sequential, tol = 1.0e-8, maxiter = 50, trunc = truncrank(χ)\n",
     "    )\n",
     "end"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "We measure the energy by computing the `H` expectation value, where we have to make sure to\n",
     "normalize with respect to the unit cell to obtain the energy per site:"
-   ]
+   ],
+   "metadata": {}
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "E = expectation_value(peps, H, env) / (Nr * Nc)\n",
     "@show E;"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "Finally, we can compare the obtained ground-state energy against the literature, namely the\n",
     "QMC estimates from [Qin et al.](@cite qin_benchmark_2016). We find that the results generally\n",
     "agree:"
-   ]
+   ],
+   "metadata": {}
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
    "outputs": [],
+   "cell_type": "code",
    "source": [
     "Es_exact = Dict(0 => -1.62, 2 => -0.176, 4 => 0.8603, 6 => -0.6567, 8 => -0.5243)\n",
     "E_exact = Es_exact[U] - U / 2\n",
     "@show (E - E_exact) / abs(E_exact);"
-   ]
+   ],
+   "metadata": {},
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
    "source": [
     "---\n",
     "\n",
     "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
-   ]
+   ],
+   "metadata": {}
   }
  ],
+ "nbformat_minor": 3,
  "metadata": {
-  "kernelspec": {
-   "display_name": "Julia 1.11.7",
-   "language": "julia",
-   "name": "julia-1.11"
-  },
   "language_info": {
    "file_extension": ".jl",
    "mimetype": "application/julia",
    "name": "julia",
    "version": "1.11.7"
+  },
+  "kernelspec": {
+   "name": "julia-1.11",
+   "display_name": "Julia 1.11.7",
+   "language": "julia"
   }
  },
- "nbformat": 4,
- "nbformat_minor": 3
-}
+ "nbformat": 4
+}
\ No newline at end of file

From d64219778fe25507e895861b3f5120128ea9971b Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 12 Dec 2025 17:56:49 +0800
Subject: [PATCH 07/13] Update docstring

---
 src/environments/suweight.jl | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index 3055cd86c..abd7b68fe 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -104,8 +104,7 @@ end
 """
     SUWeight(pepo::InfinitePEPO)
 
-Create an SUWeight for a given one-layer InfinitePEPO.
-The weights are initialized as identity matrices of element type `Float64`.
+Create a trivial SUWeight for a given one-layer InfinitePEPO.
 """
 function SUWeight(pepo::InfinitePEPO)
     @assert size(pepo, 3) == 1

From dcbe1ac66877414fe982d30e4def2b5c9f8deec3 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Mon, 15 Dec 2025 20:20:50 +0800
Subject: [PATCH 08/13] Refactor SUWeight rotation

---
 src/environments/suweight.jl | 58 ++++++++++++++++++++++++++----------
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index abd7b68fe..3f35339ac 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -326,31 +326,57 @@ end
     - need to move last row of y-weights to the 1st row.
 =#
 
-function Base.rotl90(wts::SUWeight)
-    wts_x = circshift(rotl90(wts[2, :, :]), (0, -1))
-    for (i, wt) in enumerate(wts_x)
-        wts_x[i] = DiagonalTensorMap(transpose(wt; copy = true))
-    end
-    wts_y = rotl90(wts[1, :, :])
-    return SUWeight(wts_x, wts_y)
+function _rotl90_wts_x(wts_x::AbstractMatrix{<:PEPSWeight})
+    wts_y = rotl90(wts_x)
+    return wts_y
 end
-function Base.rotr90(wts::SUWeight)
-    wts_x = rotr90(wts[2, :, :])
-    wts_y = circshift(rotr90(wts[1, :, :]), (1, 0))
+function _rotr90_wts_x(wts_x::AbstractMatrix{<:PEPSWeight})
+    wts_y = circshift(rotr90(wts_x), (1, 0))
     for (i, wt) in enumerate(wts_y)
         wts_y[i] = DiagonalTensorMap(transpose(wt; copy = true))
     end
-    return SUWeight(wts_x, wts_y)
+    return wts_y
 end
-function Base.rot180(wts::SUWeight)
-    wts_x = circshift(rot180(wts[1, :, :]), (0, -1))
-    wts_y = circshift(rot180(wts[2, :, :]), (1, 0))
+function _rot180_wts_x(wts_x::AbstractMatrix{<:PEPSWeight})
+    wts_x_ = circshift(rot180(wts_x), (0, -1))
+    for (i, wt) in enumerate(wts_x_)
+        wts_x_[i] = DiagonalTensorMap(transpose(wt; copy = true))
+    end
+    return wts_x_
+end
+
+function _rotl90_wts_y(wts_y::AbstractMatrix{<:PEPSWeight})
+    wts_x = circshift(rotl90(wts_y), (0, -1))
     for (i, wt) in enumerate(wts_x)
         wts_x[i] = DiagonalTensorMap(transpose(wt; copy = true))
     end
-    for (i, wt) in enumerate(wts_y)
-        wts_y[i] = DiagonalTensorMap(transpose(wt; copy = true))
+    return wts_x
+end
+function _rotr90_wts_y(wts_y::AbstractMatrix{<:PEPSWeight})
+    wts_x = rotr90(wts_y)
+    return wts_x
+end
+function _rot180_wts_y(wts_y::AbstractMatrix{<:PEPSWeight})
+    wts_y_ = circshift(rot180(wts_y), (1, 0))
+    for (i, wt) in enumerate(wts_y_)
+        wts_y_[i] = DiagonalTensorMap(transpose(wt; copy = true))
     end
+    return wts_y_
+end
+
+function Base.rotl90(wts::SUWeight)
+    wts_y = _rotl90_wts_x(wts[1, :, :])
+    wts_x = _rotl90_wts_y(wts[2, :, :])
+    return SUWeight(wts_x, wts_y)
+end
+function Base.rotr90(wts::SUWeight)
+    wts_y = _rotr90_wts_x(wts[1, :, :])
+    wts_x = _rotr90_wts_y(wts[2, :, :])
+    return SUWeight(wts_x, wts_y)
+end
+function Base.rot180(wts::SUWeight)
+    wts_x = _rot180_wts_x(wts[1, :, :])
+    wts_y = _rot180_wts_y(wts[2, :, :])
     return SUWeight(wts_x, wts_y)
 end
 

From 07f3fc692f28d77e5f32eeffff6813173e4538b3 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 16 Dec 2025 09:31:29 +0800
Subject: [PATCH 09/13] Apply suggestions and add SUWeight random init

---
 src/environments/suweight.jl | 75 +++++++++++++++++++++---------------
 test/ctmrg/suweight.jl       | 11 +-----
 test/types/suweight.jl       | 13 +------
 3 files changed, 46 insertions(+), 53 deletions(-)

diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index 3f35339ac..e8ba335b4 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -8,7 +8,7 @@ const PEPSWeight{T, S} = AbstractTensorMap{T, S, 1, 1}
 """
     struct SUWeight{E<:PEPSWeight}
 
-Schmidt bond weights used in simple/cluster update. 
+Schmidt bond weights used in simple/cluster update.
 Each weight is a real and semi-positive definite
 `DiagonalTensorMap`, with the same codomain and domain.
 
@@ -45,6 +45,9 @@ function SUWeight(data::Array{E, 3}) where {E <: PEPSWeight}
     for wt in data
         isa(wt, DiagonalTensorMap) ||
             error("Each weight matrix should be a DiagonalTensorMap")
+        # in case TensorKit drops this requirement
+        domain(wt, 1) == codomain(wt, 1) ||
+            error("Domain and codomain of each weight matrix must be the same.")
         all(wt.data .>= 0) || error("Weight elements must be non-negative.")
     end
     return SUWeight{E}(data)
@@ -54,63 +57,71 @@ function SUWeight(wts_mats::AbstractMatrix{E}...) where {E <: PEPSWeight}
     n_mat = length(wts_mats)
     Nr, Nc = size(wts_mats[1])
     @assert all((Nr, Nc) == size(wts_mat) for wts_mat in wts_mats)
-    weights = collect(
-        wts_mats[d][r, c] for (d, r, c) in Iterators.product(1:n_mat, 1:Nr, 1:Nc)
-    )
+    weights = map(Iterators.product(1:n_mat, 1:Nr, 1:Nc)) do (d, r, c)
+        return wts_mats[d][r, c]
+    end
     return SUWeight(weights)
 end
 
 """
-    SUWeight(Nspaces::M, [Espaces::M]) where {M<:AbstractMatrix{<:Union{Int,ElementarySpace}}}
+    SUWeight(Nspaces::M, [Espaces::M]; random::Bool=false) where {M<:AbstractMatrix{<:ElementarySpace}}
 
-Create a trivial SUWeight by specifying the vertical (north) or horizontal (east) virtual bond spaces.
-Each individual space can be specified as either an `Int` or an `ElementarySpace`.
+Create an `SUWeight` by specifying the vertical (north) or horizontal (east) virtual bond spaces.
+When `random = false`, a trivial `SUWeight` is returned.
+Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
 """
 function SUWeight(
-        Nspaces::M, Espaces::M = Nspaces
-    ) where {M <: AbstractMatrix{<:Union{Int, ElementarySpace}}}
+        Nspaces::M, Espaces::M = Nspaces; random::Bool = false
+    ) where {M <: AbstractMatrix{<:ElementarySpace}}
     @assert size(Nspaces) == size(Espaces)
     Nr, Nc = size(Nspaces)
     weights = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
         V = (d == 1 ? Espaces[r, c] : Nspaces[r, c])
-        DiagonalTensorMap(ones(reduceddim(V)), V)
+        data = random ? sort!(normalize!(rand(reduceddim(V)), Inf)) : ones(reduceddim(V))
+        DiagonalTensorMap(data, V)
     end
     return SUWeight(weights)
 end
 
 """
-    SUWeight(Nspace::S, Espace::S=Nspace; unitcell::Tuple{Int,Int}=(1, 1)) where {S<:ElementarySpace}
+    SUWeight(Nspace::S, Espace::S=Nspace; unitcell::Tuple{Int,Int}=(1, 1), random::Bool=false) where {S<:ElementarySpace}
 
-Create an trivial SUWeight by specifying its vertical (north) and horizontal (east) 
+Create an `SUWeight` by specifying its vertical (north) and horizontal (east) 
 as `ElementarySpace`s) and unit cell size.
+When `random = false`, a trivial `SUWeight` is returned.
+Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
 """
 function SUWeight(
-        Nspace::S, Espace::S = Nspace; unitcell::Tuple{Int, Int} = (1, 1)
+        Nspace::S, Espace::S = Nspace; unitcell::Tuple{Int, Int} = (1, 1), random::Bool = false
     ) where {S <: ElementarySpace}
-    return SUWeight(fill(Nspace, unitcell), fill(Espace, unitcell))
+    return SUWeight(fill(Nspace, unitcell), fill(Espace, unitcell); random)
 end
 
 """
-    SUWeight(peps::InfinitePEPS)
+    SUWeight(peps::InfinitePEPS; random::Bool = false)
 
-Create a trivial SUWeight for a given InfinitePEPS. 
+Create an `SUWeight` for a given InfinitePEPS.
+When `random = false`, a trivial `SUWeight` is returned.
+Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
 """
-function SUWeight(peps::InfinitePEPS)
-    Nspaces = collect(domain(t, NORTH) for t in peps.A)
-    Espaces = collect(domain(t, EAST) for t in peps.A)
-    return SUWeight(Nspaces, Espaces)
+function SUWeight(peps::InfinitePEPS; random::Bool = false)
+    Nspaces = map(Base.Fix2(domain, NORTH), unitcell(peps))
+    Espaces = map(Base.Fix2(domain, EAST), unitcell(peps))
+    return SUWeight(Nspaces, Espaces; random)
 end
 
 """
-    SUWeight(pepo::InfinitePEPO)
+    SUWeight(pepo::InfinitePEPO; random::Bool = false)
 
-Create a trivial SUWeight for a given one-layer InfinitePEPO.
+Create an `SUWeight` for a given one-layer InfinitePEPO.
+When `random = false`, a trivial `SUWeight` is returned.
+Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
 """
-function SUWeight(pepo::InfinitePEPO)
+function SUWeight(pepo::InfinitePEPO; random::Bool = false)
     @assert size(pepo, 3) == 1
-    Nspaces = collect(domain(t, NORTH) for t in @view(pepo.A[:, :, 1]))
-    Espaces = collect(domain(t, EAST) for t in @view(pepo.A[:, :, 1]))
-    return SUWeight(Nspaces, Espaces)
+    Nspaces = map(Base.Fix2(domain, NORTH), @view(unitcell(pepo)[:, :, 1]))
+    Espaces = map(Base.Fix2(domain, EAST), @view(unitcell(pepo)[:, :, 1]))
+    return SUWeight(Nspaces, Espaces; random)
 end
 
 ## Shape and size
@@ -179,7 +190,7 @@ in the unit cell of an InfinitePEPS/InfinitePEPO. The involved weights are
 
 ## Arguments
 
-- `t::Union{PEPSTensor, PEPOTensor}` : PEPSTensor or PEPOTensor to which the weight will be absorbed. 
+- `t::Union{PEPSTensor, PEPOTensor}` : PEPSTensor or PEPOTensor to which the weight will be absorbed.
 - `weights::SUWeight` : All simple update weights.
 - `row::Int` : The row index specifying the position in the tensor network.
 - `col::Int` : The column index specifying the position in the tensor network.
@@ -383,12 +394,12 @@ end
 """
     CTMRGEnv(wts::SUWeight)
 
-Construct a CTMRG environment with bond dimension χ = 1 from SUWeight `wts`.
-The scalartype of the returned environment is always `Float64`.
+Construct a CTMRG environment with bond dimension χ = 1 from SUWeight `wts`,
+with a real scalartype the same as `scalartype(wts)`.
 """
 function CTMRGEnv(wts::SUWeight)
     _, Nr, Nc = size(wts)
-    S = sectortype(wts)
+    elt, S = scalartype(wts), sectortype(wts)
     V_env = Vect[S](one(S) => 1)
     edges = map(Iterators.product(1:4, 1:Nr, 1:Nc)) do (d, r, c)
         wt_idx = if d == NORTH
@@ -406,10 +417,10 @@ function CTMRGEnv(wts::SUWeight)
             wt = transpose(wt)
         end
         # attach identity on environment space
-        return permute(wt ⊗ TensorKit.id(Float64, V_env), ((2, 3, 1), (4,)))
+        return @tensor edge[l t b; r] := wt[b; t] * TensorKit.id(elt, V_env)[l; r]
     end
     corners = map(CartesianIndices(edges)) do idx
-        return TensorKit.id(Float64, V_env)
+        return TensorKit.id(elt, V_env)
     end
     return CTMRGEnv(corners, edges)
 end
diff --git a/test/ctmrg/suweight.jl b/test/ctmrg/suweight.jl
index 4a195997b..05f8255e1 100644
--- a/test/ctmrg/suweight.jl
+++ b/test/ctmrg/suweight.jl
@@ -15,15 +15,6 @@ Vvs = Dict(
     FermionParity => Vect[FermionParity](0 => 2, 1 => 2),
 )
 
-function rand_wts(peps::InfinitePEPS)
-    wts = SUWeight(peps)
-    for idx in CartesianIndices(wts.data)
-        n = length(wts.data[idx].data)
-        wts.data[idx].data[:] = sort(rand(Float64, n); lt = !isless)
-    end
-    return wts
-end
-
 function su_rdm_1x1(
         row::Int, col::Int, peps::InfinitePEPS, wts::Union{Nothing, SUWeight} = nothing
     )
@@ -46,7 +37,7 @@ end
     Espaces = [Vv Vv Vv'; Vv Vv' Vv']
     Pspaces = fill(Vp, size(Nspaces))
     peps = InfinitePEPS(randn, ComplexF64, Pspaces, Nspaces, Espaces)
-    wts = (init == :trivial) ? SUWeight(peps) : rand_wts(peps)
+    wts = SUWeight(peps; random = (init != :trivial))
     env = CTMRGEnv(wts)
     for (r, c) in Tuple.(CartesianIndices(peps.A))
         ρ1 = su_rdm_1x1(r, c, peps, wts)
diff --git a/test/types/suweight.jl b/test/types/suweight.jl
index b9ee3d368..672a9d6fe 100644
--- a/test/types/suweight.jl
+++ b/test/types/suweight.jl
@@ -2,15 +2,6 @@ using Test
 using TensorKit
 using PEPSKit
 
-function rand_wts(state::Union{InfinitePEPS, InfinitePEPO})
-    wts = SUWeight(state)
-    for idx in CartesianIndices(wts.data)
-        n = length(wts.data[idx].data)
-        wts.data[idx].data[:] = sort(rand(Float64, n); lt = !isless)
-    end
-    return wts
-end
-
 function compose_n(f, n)
     n == 0 && return identity
     return f ∘ compose_n(f, n - 1)
@@ -57,11 +48,11 @@ Vs = (
 Nr, Nc = 2, 3
 peps = InfinitePEPS(rand, Float64, Vphy, Vs[2], Vs[1]'; unitcell = (Nr, Nc))
 pepo = InfinitePEPO(rand, Float64, Vphy, Vs[2], Vs[1]'; unitcell = (Nr, Nc, 1))
-wts = rand_wts(peps)
+wts = SUWeight(peps; random = true)
 
 @test sectortype(wts) === sectortype(Vs[1])
 @test spacetype(wts) === spacetype(Vs[1])
 
 test_rotation(wts)
 test_rotation(peps, wts)
-test_rotation(pepo, wts)
+test_rotation(pepo, SUWeight(pepo; random = true))

From b8f668603c5ebfa3eca6d134a2f22dab60d5b3f2 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 16 Dec 2025 10:07:11 +0800
Subject: [PATCH 10/13] Streamline CTMRGEnv(wts)

---
 src/environments/suweight.jl | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index e8ba335b4..ac4d9c0f0 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -394,8 +394,9 @@ end
 """
     CTMRGEnv(wts::SUWeight)
 
-Construct a CTMRG environment with bond dimension χ = 1 from SUWeight `wts`,
-with a real scalartype the same as `scalartype(wts)`.
+Construct a CTMRG environment with a trivial environment space 
+(bond dimension χ = 1) from SUWeight `wts`,
+which has the same real scalartype as ``wts`.
 """
 function CTMRGEnv(wts::SUWeight)
     _, Nr, Nc = size(wts)
@@ -412,12 +413,13 @@ function CTMRGEnv(wts::SUWeight)
             CartesianIndex(1, r, c)
         end
         # temporarily make wt axis order ([bra], [ket])
-        wt = deepcopy(wts[wt_idx])
-        if d in (NORTH, EAST)
-            wt = transpose(wt)
+        wt = if d in (NORTH, EAST)
+            twist!(repartition(wts[wt_idx], 2, 0; copy = true), 1)
+        else
+            permute(wts[wt_idx], ((2, 1), ()); copy = true)
         end
         # attach identity on environment space
-        return @tensor edge[l t b; r] := wt[b; t] * TensorKit.id(elt, V_env)[l; r]
+        return insertleftunit(insertleftunit(wt), 1)
     end
     corners = map(CartesianIndices(edges)) do idx
         return TensorKit.id(elt, V_env)

From 1a3e8169a64dcb8c952bed69c91e91335dea499c Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 16 Dec 2025 11:02:06 +0800
Subject: [PATCH 11/13] Correct a comment

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

diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index ac4d9c0f0..35b2c0273 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -412,7 +412,7 @@ function CTMRGEnv(wts::SUWeight)
         else # WEST
             CartesianIndex(1, r, c)
         end
-        # temporarily make wt axis order ([bra], [ket])
+        # make wt axis order ([ket], [bra])
         wt = if d in (NORTH, EAST)
             twist!(repartition(wts[wt_idx], 2, 0; copy = true), 1)
         else

From 12f9dc38027b69f2a1078e87bb1cff00257f49c5 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 2 Jan 2026 18:13:27 +0800
Subject: [PATCH 12/13] Apply suggestions

---
 src/PEPSKit.jl               |  1 +
 src/environments/suweight.jl | 62 ++++++++++++++++--------------------
 test/ctmrg/suweight.jl       |  6 +++-
 test/types/suweight.jl       | 10 ++++--
 4 files changed, 41 insertions(+), 38 deletions(-)

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index 4a39d6bd3..57d359cd3 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -1,6 +1,7 @@
 module PEPSKit
 
 using LinearAlgebra, Statistics, Base.Threads, Base.Iterators, Printf
+using Random
 using Compat
 using Accessors: @set, @reset
 using VectorInterface
diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index 35b2c0273..51301ae5a 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -54,74 +54,67 @@ function SUWeight(data::Array{E, 3}) where {E <: PEPSWeight}
 end
 
 function SUWeight(wts_mats::AbstractMatrix{E}...) where {E <: PEPSWeight}
-    n_mat = length(wts_mats)
-    Nr, Nc = size(wts_mats[1])
-    @assert all((Nr, Nc) == size(wts_mat) for wts_mat in wts_mats)
-    weights = map(Iterators.product(1:n_mat, 1:Nr, 1:Nc)) do (d, r, c)
-        return wts_mats[d][r, c]
-    end
-    return SUWeight(weights)
+    return SUWeight(stack(wts_mats; dims = 1))
 end
 
 """
-    SUWeight(Nspaces::M, [Espaces::M]; random::Bool=false) where {M<:AbstractMatrix{<:ElementarySpace}}
+    SUWeight(Nspaces::M, [Espaces::M]) where {M<:AbstractMatrix{<:ElementarySpace}}
 
-Create an `SUWeight` by specifying the vertical (north) or horizontal (east) virtual bond spaces.
-When `random = false`, a trivial `SUWeight` is returned.
-Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
+Create a trivial `SUWeight` by specifying the vertical (north) or horizontal (east) virtual bond spaces.
 """
 function SUWeight(
-        Nspaces::M, Espaces::M = Nspaces; random::Bool = false
+        Nspaces::M, Espaces::M = Nspaces
     ) where {M <: AbstractMatrix{<:ElementarySpace}}
     @assert size(Nspaces) == size(Espaces)
     Nr, Nc = size(Nspaces)
     weights = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
         V = (d == 1 ? Espaces[r, c] : Nspaces[r, c])
-        data = random ? sort!(normalize!(rand(reduceddim(V)), Inf)) : ones(reduceddim(V))
-        DiagonalTensorMap(data, V)
+        DiagonalTensorMap(ones(reduceddim(V)), V)
     end
     return SUWeight(weights)
 end
 
 """
-    SUWeight(Nspace::S, Espace::S=Nspace; unitcell::Tuple{Int,Int}=(1, 1), random::Bool=false) where {S<:ElementarySpace}
+    SUWeight(Nspace::S, Espace::S=Nspace; unitcell::Tuple{Int,Int}=(1, 1)) where {S<:ElementarySpace}
 
-Create an `SUWeight` by specifying its vertical (north) and horizontal (east) 
+Create a trivial `SUWeight` by specifying its vertical (north) and horizontal (east) 
 as `ElementarySpace`s) and unit cell size.
-When `random = false`, a trivial `SUWeight` is returned.
-Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
 """
 function SUWeight(
-        Nspace::S, Espace::S = Nspace; unitcell::Tuple{Int, Int} = (1, 1), random::Bool = false
+        Nspace::S, Espace::S = Nspace; unitcell::Tuple{Int, Int} = (1, 1)
     ) where {S <: ElementarySpace}
-    return SUWeight(fill(Nspace, unitcell), fill(Espace, unitcell); random)
+    return SUWeight(fill(Nspace, unitcell), fill(Espace, unitcell))
 end
 
 """
-    SUWeight(peps::InfinitePEPS; random::Bool = false)
+    SUWeight(peps::InfinitePEPS)
 
-Create an `SUWeight` for a given InfinitePEPS.
-When `random = false`, a trivial `SUWeight` is returned.
-Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
+Create a trivial `SUWeight` for a given InfinitePEPS.
 """
-function SUWeight(peps::InfinitePEPS; random::Bool = false)
+function SUWeight(peps::InfinitePEPS)
     Nspaces = map(Base.Fix2(domain, NORTH), unitcell(peps))
     Espaces = map(Base.Fix2(domain, EAST), unitcell(peps))
-    return SUWeight(Nspaces, Espaces; random)
+    return SUWeight(Nspaces, Espaces)
 end
 
 """
-    SUWeight(pepo::InfinitePEPO; random::Bool = false)
+    SUWeight(pepo::InfinitePEPO)
 
-Create an `SUWeight` for a given one-layer InfinitePEPO.
-When `random = false`, a trivial `SUWeight` is returned.
-Otherwise each weight is randomly generated, and normalized with its `Inf`-norm.
+Create a trivial `SUWeight` for a given one-layer InfinitePEPO.
 """
-function SUWeight(pepo::InfinitePEPO; random::Bool = false)
+function SUWeight(pepo::InfinitePEPO)
     @assert size(pepo, 3) == 1
     Nspaces = map(Base.Fix2(domain, NORTH), @view(unitcell(pepo)[:, :, 1]))
     Espaces = map(Base.Fix2(domain, EAST), @view(unitcell(pepo)[:, :, 1]))
-    return SUWeight(Nspaces, Espaces; random)
+    return SUWeight(Nspaces, Espaces)
+end
+
+function Random.rand!(wts::SUWeight)
+    for idx in CartesianIndices(wts.data)
+        newdata = rand(length(wts.data[idx].data))
+        wts.data[idx].data[:] = sort!(newdata; lt = !isless)
+    end
+    return wts
 end
 
 ## Shape and size
@@ -400,8 +393,8 @@ which has the same real scalartype as ``wts`.
 """
 function CTMRGEnv(wts::SUWeight)
     _, Nr, Nc = size(wts)
-    elt, S = scalartype(wts), sectortype(wts)
-    V_env = Vect[S](one(S) => 1)
+    elt = scalartype(wts)
+    V_env = oneunit(spacetype(wts))
     edges = map(Iterators.product(1:4, 1:Nr, 1:Nc)) do (d, r, c)
         wt_idx = if d == NORTH
             CartesianIndex(2, _next(r, Nr), c)
@@ -412,7 +405,6 @@ function CTMRGEnv(wts::SUWeight)
         else # WEST
             CartesianIndex(1, r, c)
         end
-        # make wt axis order ([ket], [bra])
         wt = if d in (NORTH, EAST)
             twist!(repartition(wts[wt_idx], 2, 0; copy = true), 1)
         else
diff --git a/test/ctmrg/suweight.jl b/test/ctmrg/suweight.jl
index 05f8255e1..3e90e28cc 100644
--- a/test/ctmrg/suweight.jl
+++ b/test/ctmrg/suweight.jl
@@ -37,7 +37,11 @@ end
     Espaces = [Vv Vv Vv'; Vv Vv' Vv']
     Pspaces = fill(Vp, size(Nspaces))
     peps = InfinitePEPS(randn, ComplexF64, Pspaces, Nspaces, Espaces)
-    wts = SUWeight(peps; random = (init != :trivial))
+    wts = SUWeight(peps)
+    if init != :trivial
+        rand!(wts)
+        normalize!.(wts.data, Inf)
+    end
     env = CTMRGEnv(wts)
     for (r, c) in Tuple.(CartesianIndices(peps.A))
         ρ1 = su_rdm_1x1(r, c, peps, wts)
diff --git a/test/types/suweight.jl b/test/types/suweight.jl
index 672a9d6fe..dee69599a 100644
--- a/test/types/suweight.jl
+++ b/test/types/suweight.jl
@@ -48,11 +48,17 @@ Vs = (
 Nr, Nc = 2, 3
 peps = InfinitePEPS(rand, Float64, Vphy, Vs[2], Vs[1]'; unitcell = (Nr, Nc))
 pepo = InfinitePEPO(rand, Float64, Vphy, Vs[2], Vs[1]'; unitcell = (Nr, Nc, 1))
-wts = SUWeight(peps; random = true)
+wts = SUWeight(peps)
+rand!(wts)
+normalize!.(wts.data, Inf)
 
 @test sectortype(wts) === sectortype(Vs[1])
 @test spacetype(wts) === spacetype(Vs[1])
 
 test_rotation(wts)
 test_rotation(peps, wts)
-test_rotation(pepo, SUWeight(pepo; random = true))
+
+wts = SUWeight(pepo)
+rand!(wts)
+normalize!.(wts.data, Inf)
+test_rotation(pepo, wts)

From a8aae2b90c09367a260d4901341f3a323480d847 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sat, 3 Jan 2026 10:11:11 +0800
Subject: [PATCH 13/13] Improve `rand!(wts)`

---
 src/environments/suweight.jl | 10 ++++++----
 test/types/suweight.jl       |  2 ++
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/environments/suweight.jl b/src/environments/suweight.jl
index 51301ae5a..ea4491aec 100644
--- a/src/environments/suweight.jl
+++ b/src/environments/suweight.jl
@@ -109,10 +109,12 @@ function SUWeight(pepo::InfinitePEPO)
     return SUWeight(Nspaces, Espaces)
 end
 
-function Random.rand!(wts::SUWeight)
-    for idx in CartesianIndices(wts.data)
-        newdata = rand(length(wts.data[idx].data))
-        wts.data[idx].data[:] = sort!(newdata; lt = !isless)
+Random.rand!(wts::SUWeight) = rand!(Random.default_rng(), wts)
+function Random.rand!(rng::Random.AbstractRNG, wts::SUWeight)
+    foreach(wts.data) do wt
+        for (_, b) in blocks(wt)
+            sort!(rand!(rng, b.diag); rev = true)
+        end
     end
     return wts
 end
diff --git a/test/types/suweight.jl b/test/types/suweight.jl
index dee69599a..cc7440920 100644
--- a/test/types/suweight.jl
+++ b/test/types/suweight.jl
@@ -51,6 +51,8 @@ pepo = InfinitePEPO(rand, Float64, Vphy, Vs[2], Vs[1]'; unitcell = (Nr, Nc, 1))
 wts = SUWeight(peps)
 rand!(wts)
 normalize!.(wts.data, Inf)
+# check that elements of wts are successfully randomized
+@test !(wts ≈ SUWeight(peps))
 
 @test sectortype(wts) === sectortype(Vs[1])
 @test spacetype(wts) === spacetype(Vs[1])