From 87511dd749ba0ad017a2d69adf9767b7c6473b00 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Mon, 2 Mar 2026 22:35:55 +0800
Subject: [PATCH 01/27] Introduce `TrotterMPOs`

---
 src/algorithms/time_evolution/simpleupdate.jl | 23 ++----
 .../time_evolution/simpleupdate3site.jl       |  7 +-
 src/algorithms/time_evolution/time_evolve.jl  | 11 +++
 src/algorithms/time_evolution/trotter_gate.jl | 61 ++++++++++++----
 src/algorithms/time_evolution/trotter_mpo.jl  | 70 ++++++++++++-------
 5 files changed, 115 insertions(+), 57 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 9df0b8034..a0ca7de08 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -55,22 +55,11 @@ function TimeEvolver(
     )
     # sanity checks
     _timeevol_sanity_check(psi0, physicalspace(H), alg)
-    # create Trotter gates
-    nnonly = is_nearest_neighbour(H)
-    use_3site = alg.force_3site || !nnonly
-    if alg.bipartite
-        @assert !use_3site "3-site simple update is incompatible with bipartite lattice."
-    end
     dt′ = _get_dt(psi0, dt, alg.imaginary_time)
-    gate = if use_3site
-        [
-            _get_gatempos_se(H, dt′),
-            _get_gatempos_se(rotl90(H), dt′),
-            _get_gatempos_se(rot180(H), dt′),
-            _get_gatempos_se(rotr90(H), dt′),
-        ]
-    else
-        get_expham(H, dt′)
+    # create Trotter gates
+    gate = trotterize(H, dt′; force_3site = alg.force_3site)
+    if isa(gate, TrotterMPOs)
+        @assert !alg.bipartite "Trotter MPOs are incompatible with bipartite lattice structure."
     end
     state = SUState(0, t0, psi0, env0)
     return TimeEvolver(alg, dt, nstep, gate, state)
@@ -166,7 +155,7 @@ function su_iter(
     for r in 1:Nr, c in 1:Nc
         (alg.bipartite && r > 1) && continue
         # update x-bonds
-        term = get_gateterm(gate, (CartesianIndex(r, c), CartesianIndex(r, c + 1)))
+        term = _get_bond_term(gate, (CartesianIndex(r, c), CartesianIndex(r, c + 1)))
         trunc = truncation_strategy(alg.trunc, 1, r, c)
         for gate_ax in gate_axs
             ϵ′ = _su_xbond!(state2, term, env2, r, c, trunc; gate_ax)
@@ -179,7 +168,7 @@ function su_iter(
             env2.data[1, rp1, cp1] = deepcopy(env2.data[1, r, c])
         end
         # update y-bonds
-        term = get_gateterm(gate, (CartesianIndex(r, c), CartesianIndex(r - 1, c)))
+        term = _get_bond_term(gate, (CartesianIndex(r, c), CartesianIndex(r - 1, c)))
         trunc = truncation_strategy(alg.trunc, 2, r, c)
         for gate_ax in gate_axs
             ϵ′ = _su_ybond!(state2, term, env2, r, c, trunc; gate_ax)
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index ea3d0e2e5..e1e639413 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -97,8 +97,9 @@ function _su3site_se!(
 end
 
 function su_iter(
-        state::InfiniteState, gatempos::Vector{G}, alg::SimpleUpdate, env::SUWeight
-    ) where {G <: AbstractMatrix}
+        state::InfiniteState, gatempos::TrotterMPOs2ndNeighbor,
+        alg::SimpleUpdate, env::SUWeight
+    )
     if state isa InfinitePEPO
         @assert size(state, 3) == 1
     end
@@ -113,7 +114,7 @@ function su_iter(
     for i in 1:4
         Nr, Nc = size(state2)[1:2]
         for r in 1:Nr, c in 1:Nc
-            gs = gatempos[i][r, c]
+            gs = gatempos[i, r, c]
             truncs = [
                 truncation_strategy(trunc, 1, r, c)
                 truncation_strategy(trunc, 2, r, _next(c, Nc))
diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index d2b208542..b4fd1c153 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -91,3 +91,14 @@ function MPSKit.infinite_temperature_density_matrix(H::LocalOperator)
     end
     return InfinitePEPO(cat(A; dims = 3))
 end
+
+function trotterize(H::LocalOperator, dt::Number; force_3site::Bool = false)
+    nnonly = is_nearest_neighbour(H)
+    use_3site = force_3site || !nnonly
+    gate = if use_3site
+        TrotterMPOs2ndNeighbor(H, dt)
+    else
+        get_expham(H, dt)
+    end
+    return gate
+end
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index dfe151e42..09d463c89 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -39,27 +39,26 @@ function is_equivalent_bond(
 end
 
 """
-    get_gateterm(gate::LocalOperator, bond::NTuple{2,CartesianIndex{2}})
+    _get_bond_term(ham::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
 
-Get the term of a 2-site gate acting on a certain bond.
-Input `gate` should only include one term for each nearest neighbor bond.
+Get the 2-site term on `bond` in `ham`.
 """
-function get_gateterm(gate::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
-    bonds = findall(p -> is_equivalent_bond(p.first, bond, size(gate.lattice)), gate.terms)
+function _get_bond_term(ham::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
+    bonds = findall(p -> is_equivalent_bond(p.first, bond, size(ham.lattice)), ham.terms)
     if length(bonds) == 0
         # try reversed site order
         bonds = findall(
-            p -> is_equivalent_bond(p.first, reverse(bond), size(gate.lattice)), gate.terms
+            p -> is_equivalent_bond(p.first, reverse(bond), size(ham.lattice)), ham.terms
         )
         if length(bonds) == 1
-            return permute(gate.terms[bonds[1]].second, ((2, 1), (4, 3)))
+            return permute(ham.terms[bonds[1]].second, ((2, 1), (4, 3)))
         elseif length(bonds) == 0
             # if term not found, return the zero operator on this bond
-            dtype = scalartype(gate)
-            r1, c1 = (mod1(bond[1][i], n) for (i, n) in zip(1:2, size(gate)))
-            r2, c2 = (mod1(bond[2][i], n) for (i, n) in zip(1:2, size(gate)))
-            V1 = physicalspace(gate)[r1, c1]
-            V2 = physicalspace(gate)[r2, c2]
+            dtype = scalartype(ham)
+            r1, c1 = (mod1(bond[1][i], n) for (i, n) in zip(1:2, size(ham)))
+            r2, c2 = (mod1(bond[2][i], n) for (i, n) in zip(1:2, size(ham)))
+            V1 = physicalspace(ham)[r1, c1]
+            V2 = physicalspace(ham)[r2, c2]
             return zeros(dtype, V1 ⊗ V2 ← V1 ⊗ V2)
         else
             error("There are multiple terms in `gate` corresponding to the bond $(bond).")
@@ -67,6 +66,42 @@ function get_gateterm(gate::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
     else
         (length(bonds) == 1) ||
             error("There are multiple terms in `gate` corresponding to the bond $(bond).")
-        return gate.terms[bonds[1]].second
+        return ham.terms[bonds[1]].second
     end
 end
+
+"""
+    _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
+
+Construct the term acting on the southeast 3-site cluster in `ham`.
+```
+    r-1        g3
+                |
+                ↓
+    r   g1 -←- g2
+        c      c+1
+```
+"""
+function _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
+    Nr, Nc = size(ham)
+    @assert 1 <= row <= Nr && 1 <= col <= Nc
+    sites = [
+        CartesianIndex(row, col),
+        CartesianIndex(row, col + 1),
+        CartesianIndex(row - 1, col + 1),
+    ]
+    nb1x = _get_bond_term(ham, (sites[1], sites[2]))
+    nb1y = _get_bond_term(ham, (sites[2], sites[3]))
+    nb2 = _get_bond_term(ham, (sites[1], sites[3]))
+    # identity operator at each site
+    units = map(sites) do site
+        site_ = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
+        return id(physicalspace(ham)[site_])
+    end
+    # when iterating through ┘, └, ┌, ┐ clusters in the unit cell,
+    # NN / NNN bonds are counted 4 / 2 times, respectively.
+    @tensor term[i' j' k'; i j k] :=
+        (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4 +
+        (nb2[i' k'; i k] * units[2][j'; j]) / 2
+    return term
+end
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
index d8f1d3e87..720079a71 100644
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ b/src/algorithms/time_evolution/trotter_mpo.jl
@@ -1,3 +1,47 @@
+"""
+$(TYPEDEF)
+
+Abstract super type for the collection of
+Trotter evolution MPOs acting on 3 or more sites.
+"""
+abstract type TrotterMPOs end
+
+Base.getindex(gate::TrotterMPOs, args...) = Base.getindex(gate.mpos, args...)
+
+"""
+    struct TrotterMPOs2ndNeighbor{T}
+
+Collection of all Trotter evolution MPOs obtained from a Hamiltonian
+containing up to 2nd nearest neighbor terms.
+
+Before exponentiating, terms in the Hamiltonian are organized as
+```
+    H = ∑ᵢⱼ(┘ᵢⱼ + ┐ᵢⱼ + ┌ᵢⱼ + └ᵢⱼ)
+```
+where `┘`, `┐`, `┌`, `└` refer to the following 3-site clusters 
+```
+        3   3---2   2---1   1
+        |       |   |       |
+    1---2       1   3       2---3
+```
+Then each Trotter MPO is `exp(-dt * ┘ᵢⱼ)`, etc.
+"""
+struct TrotterMPOs2ndNeighbor{T} <: TrotterMPOs
+    mpos::T
+end
+
+function TrotterMPOs2ndNeighbor(H::LocalOperator, dt::Number)
+    mpos = stack(
+        [
+            _get_gatempos_se(H, dt),
+            _get_gatempos_se(rotl90(H), dt),
+            _get_gatempos_se(rot180(H), dt),
+            _get_gatempos_se(rotr90(H), dt),
+        ]; dims = 1
+    )
+    return TrotterMPOs2ndNeighbor(mpos)
+end
+
 """
 Convert a 3-site gate to MPO form by SVD, 
 in which the axes are ordered as
@@ -30,30 +74,8 @@ Obtain the 3-site gate MPO on the southeast cluster at position `[row, col]`
 ```
 """
 function _get_gatempo_se(ham::LocalOperator, dt::Number, row::Int, col::Int)
-    Nr, Nc = size(ham)
-    @assert 1 <= row <= Nr && 1 <= col <= Nc
-    sites = [
-        CartesianIndex(row, col),
-        CartesianIndex(row, col + 1),
-        CartesianIndex(row - 1, col + 1),
-    ]
-    nb1x = get_gateterm(ham, (sites[1], sites[2]))
-    nb1y = get_gateterm(ham, (sites[2], sites[3]))
-    nb2 = get_gateterm(ham, (sites[1], sites[3]))
-    # identity operator at each site
-    units = map(sites) do site
-        site_ = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
-        return id(physicalspace(ham)[site_])
-    end
-    # when iterating through ┘, └, ┌, ┐ clusters in the unit cell,
-    # NN / NNN bonds are counted 4 / 2 times, respectively.
-    @tensor Odt[i' j' k'; i j k] :=
-        -dt * (
-        (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4 +
-            (nb2[i' k'; i k] * units[2][j'; j]) / 2
-    )
-    op = exp(Odt)
-    return gate_to_mpo3(op)
+    term = _get_se3site_term(ham, row, col)
+    return gate_to_mpo3(exp(-dt * term))
 end
 
 """

From ab60954541e47b99591839757d4f6b504da964ee Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 3 Mar 2026 09:36:03 +0800
Subject: [PATCH 02/27] Remove `force_3site`

---
 src/algorithms/time_evolution/simpleupdate.jl |  4 +-
 .../time_evolution/simpleupdate3site.jl       |  2 +-
 src/algorithms/time_evolution/time_evolve.jl  |  9 ++--
 src/algorithms/time_evolution/trotter_mpo.jl  | 14 +++---
 test/timeevol/cluster_projectors.jl           | 50 -------------------
 test/timeevol/sitedep_truncation.jl           |  6 +--
 6 files changed, 15 insertions(+), 70 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index a0ca7de08..66ff2613b 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -12,8 +12,6 @@ $(TYPEDFIELDS)
     trunc::TruncationStrategy
     "When true (or false), the Trotter gate is `exp(-H dt)` (or `exp(-iH dt)`)"
     imaginary_time::Bool = true
-    "When true, force the usage of 3-site simple update"
-    force_3site::Bool = false
     "When true, assume bipartite unit cell structure"
     bipartite::Bool = false
     "(Only applicable to InfinitePEPO) 
@@ -57,7 +55,7 @@ function TimeEvolver(
     _timeevol_sanity_check(psi0, physicalspace(H), alg)
     dt′ = _get_dt(psi0, dt, alg.imaginary_time)
     # create Trotter gates
-    gate = trotterize(H, dt′; force_3site = alg.force_3site)
+    gate = trotterize(H, dt′)
     if isa(gate, TrotterMPOs)
         @assert !alg.bipartite "Trotter MPOs are incompatible with bipartite lattice structure."
     end
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index e1e639413..35a31206e 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -114,7 +114,7 @@ function su_iter(
     for i in 1:4
         Nr, Nc = size(state2)[1:2]
         for r in 1:Nr, c in 1:Nc
-            gs = gatempos[i, r, c]
+            gs = gatempos[i][r, c]
             truncs = [
                 truncation_strategy(trunc, 1, r, c)
                 truncation_strategy(trunc, 2, r, _next(c, Nc))
diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index b4fd1c153..9220ce9e8 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -92,13 +92,12 @@ function MPSKit.infinite_temperature_density_matrix(H::LocalOperator)
     return InfinitePEPO(cat(A; dims = 3))
 end
 
-function trotterize(H::LocalOperator, dt::Number; force_3site::Bool = false)
+function trotterize(H::LocalOperator, dt::Number)
     nnonly = is_nearest_neighbour(H)
-    use_3site = force_3site || !nnonly
-    gate = if use_3site
-        TrotterMPOs2ndNeighbor(H, dt)
-    else
+    gate = if nnonly
         get_expham(H, dt)
+    else
+        TrotterMPOs2ndNeighbor(H, dt)
     end
     return gate
 end
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
index 720079a71..861b90904 100644
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ b/src/algorithms/time_evolution/trotter_mpo.jl
@@ -31,14 +31,12 @@ struct TrotterMPOs2ndNeighbor{T} <: TrotterMPOs
 end
 
 function TrotterMPOs2ndNeighbor(H::LocalOperator, dt::Number)
-    mpos = stack(
-        [
-            _get_gatempos_se(H, dt),
-            _get_gatempos_se(rotl90(H), dt),
-            _get_gatempos_se(rot180(H), dt),
-            _get_gatempos_se(rotr90(H), dt),
-        ]; dims = 1
-    )
+    mpos = [
+        _get_gatempos_se(H, dt),
+        _get_gatempos_se(rotl90(H), dt),
+        _get_gatempos_se(rot180(H), dt),
+        _get_gatempos_se(rotr90(H), dt)
+    ]
     return TrotterMPOs2ndNeighbor(mpos)
 end
 
diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index 0ad3ae2ce..faff68774 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -100,53 +100,3 @@ end
         end
     end
 end
-
-@testset "Hubbard model with 2-site and 3-site SU" begin
-    Nr, Nc = 2, 2
-    ctmrg_tol = 1.0e-9
-    Random.seed!(1459)
-    # with U(1) spin rotation symmetry
-    Pspace = hubbard_space(Trivial, U1Irrep)
-    Vspace = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2, (1, 1 // 2) => 1, (1, -1 // 2) => 1)
-    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, Vspace'; unitcell = (Nr, Nc))
-    # make state bipartite
-    for r in 1:2
-        peps.A[_next(r, 2), 2] = copy(peps.A[r, 1])
-    end
-    wts = SUWeight(peps)
-    ham = real(
-        hubbard_model(
-            ComplexF64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t = 1.0, U = 8.0, mu = 0.0
-        ),
-    )
-    # usual 2-site simple update, and measure energy
-    dts = [1.0e-2, 1.0e-2]
-    tols = [1.0e-8, 1.0e-8]
-    for (n, (dt, tol)) in enumerate(zip(dts, tols))
-        trunc = truncerror(; atol = 1.0e-10) & truncrank(n == 1 ? 4 : 2)
-        alg = SimpleUpdate(; trunc, bipartite = true)
-        peps, wts, = time_evolve(peps, ham, dt, 10000, alg, wts; tol, check_interval = 1000)
-    end
-    normalize!.(peps.A, Inf)
-    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)
-    @info "2-site simple update energy = $e_site"
-    # continue with 3-site simple update; energy should not change much
-    dts = [1.0e-2]
-    tols = [1.0e-8]
-    trunc = truncerror(; atol = 1.0e-10) & truncrank(2)
-    alg = SimpleUpdate(; trunc, force_3site = true)
-    for (n, (dt, tol)) in enumerate(zip(dts, tols))
-        peps, wts, = time_evolve(peps, ham, dt, 5000, alg, wts; tol, check_interval = 1000)
-    end
-    normalize!.(peps.A, Inf)
-    env, = leading_boundary(env, peps; tol = ctmrg_tol, trunc = trunc_env)
-    e_site2 = cost_function(peps, env, ham) / (Nr * Nc)
-    @info "3-site simple update energy = $e_site2"
-    @test e_site ≈ e_site2 atol = 5.0e-4
-end
diff --git a/test/timeevol/sitedep_truncation.jl b/test/timeevol/sitedep_truncation.jl
index 355fa085c..a36f3de75 100644
--- a/test/timeevol/sitedep_truncation.jl
+++ b/test/timeevol/sitedep_truncation.jl
@@ -48,7 +48,6 @@ end
 
 @testset "Simple update: generic 2-site and 3-site" begin
     Nr, Nc = 3, 4
-    ham = real(heisenberg_XYZ(InfiniteSquare(Nr, Nc); Jx = 1.0, Jy = 1.0, Jz = 1.0))
     Random.seed!(100)
     peps0 = InfinitePEPS(rand, Float64, ℂ^2, ℂ^10; unitcell = (Nr, Nc))
     normalize!.(peps0.A, Inf)
@@ -57,13 +56,14 @@ end
     bonddims = rand(2:8, 2, Nr, Nc)
     @show bonddims
     trunc = SiteDependentTruncation(collect(truncrank(d) for d in bonddims))
-    # 2-site SU
     alg = SimpleUpdate(; trunc)
+    # 2-site SU
+    ham = j1_j2_model(Float64, Trivial, InfiniteSquare(Nr, Nc); J2 = 0.0, sublattice = false)
     peps, env, = time_evolve(peps0, ham, 1.0e-2, 4, alg, env0)
     @test get_bonddims(peps) == bonddims
     @test get_bonddims(env) == bonddims
     # 3-site SU
-    alg = SimpleUpdate(; trunc, force_3site = true)
+    ham = j1_j2_model(Float64, Trivial, InfiniteSquare(Nr, Nc); J2 = 0.2, sublattice = false)
     peps, env, = time_evolve(peps0, ham, 1.0e-2, 4, alg, env0)
     @test get_bonddims(peps) == bonddims
     @test get_bonddims(env) == bonddims

From 0836f750ae246063547120d5631c2c50b2278693 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 3 Mar 2026 10:07:01 +0800
Subject: [PATCH 03/27] Introduce `TrotterGates`

---
 src/algorithms/time_evolution/simpleupdate.jl |  9 ++---
 src/algorithms/time_evolution/time_evolve.jl  |  2 +-
 src/algorithms/time_evolution/trotter_gate.jl | 38 +++++++++++++++----
 src/algorithms/time_evolution/trotter_mpo.jl  |  9 ++---
 4 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 66ff2613b..1ef38b566 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -145,7 +145,8 @@ end
 One iteration of simple update
 """
 function su_iter(
-        state::InfiniteState, gate::LocalOperator, alg::SimpleUpdate, env::SUWeight
+        state::InfiniteState, gate::TrotterGates1stNeighbor,
+        alg::SimpleUpdate, env::SUWeight
     )
     Nr, Nc, = size(state)
     state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
@@ -153,10 +154,9 @@ function su_iter(
     for r in 1:Nr, c in 1:Nc
         (alg.bipartite && r > 1) && continue
         # update x-bonds
-        term = _get_bond_term(gate, (CartesianIndex(r, c), CartesianIndex(r, c + 1)))
         trunc = truncation_strategy(alg.trunc, 1, r, c)
         for gate_ax in gate_axs
-            ϵ′ = _su_xbond!(state2, term, env2, r, c, trunc; gate_ax)
+            ϵ′ = _su_xbond!(state2, gate[1, r, c], env2, r, c, trunc; gate_ax)
             ϵ = max(ϵ, ϵ′)
         end
         if alg.bipartite
@@ -166,10 +166,9 @@ function su_iter(
             env2.data[1, rp1, cp1] = deepcopy(env2.data[1, r, c])
         end
         # update y-bonds
-        term = _get_bond_term(gate, (CartesianIndex(r, c), CartesianIndex(r - 1, c)))
         trunc = truncation_strategy(alg.trunc, 2, r, c)
         for gate_ax in gate_axs
-            ϵ′ = _su_ybond!(state2, term, env2, r, c, trunc; gate_ax)
+            ϵ′ = _su_ybond!(state2, gate[2, r, c], env2, r, c, trunc; gate_ax)
             ϵ = max(ϵ, ϵ′)
         end
         if alg.bipartite
diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index 9220ce9e8..a6e2dbd5a 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -95,7 +95,7 @@ end
 function trotterize(H::LocalOperator, dt::Number)
     nnonly = is_nearest_neighbour(H)
     gate = if nnonly
-        get_expham(H, dt)
+        TrotterGates1stNeighbor(H, dt)
     else
         TrotterMPOs2ndNeighbor(H, dt)
     end
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 09d463c89..245580c76 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -1,14 +1,36 @@
 """
-    get_expham(H::LocalOperator, dt::Number)
+$(TYPEDEF)
 
-Compute `exp(-dt * op)` for each term `op` in `H`,
-and combine them into a new `LocalOperator`.
-Each `op` in `H` must be a single `TensorMap`.
+Abstract super type for the collection of 2-body Trotter evolution gates.
 """
-function get_expham(H::LocalOperator, dt::Number)
-    return LocalOperator(
-        physicalspace(H), (sites => exp(-dt * op) for (sites, op) in H.terms)...
-    )
+abstract type TrotterGates end
+
+Base.getindex(gates::TrotterGates, args...) = Base.getindex(gates.gates, args...)
+
+"""
+Collection of 1st (nearest) neighbor 2-body Trotter gates.
+
+Before exponentiating, terms in the Hamiltonian are organized as
+```
+    H = ∑ᵢⱼ (Xᵢⱼ + Yᵢⱼ)
+```
+where each `Xᵢⱼ` (or `Yᵢⱼ`) acts on a horizontal (or vertical) bond.
+The Trotter gates are `exp(-dt * Xᵢⱼ)`, `exp(-dt * Yᵢⱼ)`.
+"""
+struct TrotterGates1stNeighbor{G} <: TrotterGates
+    gates::G
+end
+
+function TrotterGates1stNeighbor(H::LocalOperator, dt::Number)
+    Nr, Nc = size(H)
+    gates = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
+        # d = 1: horizontal bond; d = 2: vertical bond
+        site1 = CartesianIndex(r, c)
+        site2 = (d == 1) ?  CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
+        term = _get_bond_term(H, (site1, site2))
+        return exp(-dt * term)
+    end
+    return TrotterGates1stNeighbor(gates)
 end
 
 """
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
index 861b90904..c6d1d8265 100644
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ b/src/algorithms/time_evolution/trotter_mpo.jl
@@ -11,10 +11,8 @@ Base.getindex(gate::TrotterMPOs, args...) = Base.getindex(gate.mpos, args...)
 """
     struct TrotterMPOs2ndNeighbor{T}
 
-Collection of all Trotter evolution MPOs obtained from a Hamiltonian
-containing up to 2nd nearest neighbor terms.
-
-Before exponentiating, terms in the Hamiltonian are organized as
+Collection of all Trotter evolution MPOs obtained from
+a Hamiltonian containing up to 2nd neighbor terms
 ```
     H = ∑ᵢⱼ(┘ᵢⱼ + ┐ᵢⱼ + ┌ᵢⱼ + └ᵢⱼ)
 ```
@@ -24,7 +22,8 @@ where `┘`, `┐`, `┌`, `└` refer to the following 3-site clusters
         |       |   |       |
     1---2       1   3       2---3
 ```
-Then each Trotter MPO is `exp(-dt * ┘ᵢⱼ)`, etc.
+`mpos[d][i, j]` is the `┘ᵢⱼ` MPO acting on the `[i, j]` southeast
+cluster after the network is left-rotated by `90 x (d - 1)` degrees.
 """
 struct TrotterMPOs2ndNeighbor{T} <: TrotterMPOs
     mpos::T

From a0046a3452ee17a47ccfe6cbc1bd17869c51ef3e Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 3 Mar 2026 10:52:16 +0800
Subject: [PATCH 04/27] Improve _get_bond_term; introduce _get_site_term

---
 src/algorithms/time_evolution/trotter_gate.jl | 81 +++++++++++++------
 src/algorithms/time_evolution/trotter_mpo.jl  |  2 +-
 2 files changed, 57 insertions(+), 26 deletions(-)

diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 245580c76..1b5730f78 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -26,7 +26,7 @@ function TrotterGates1stNeighbor(H::LocalOperator, dt::Number)
     gates = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
         # d = 1: horizontal bond; d = 2: vertical bond
         site1 = CartesianIndex(r, c)
-        site2 = (d == 1) ?  CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
+        site2 = (d == 1) ? CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
         term = _get_bond_term(H, (site1, site2))
         return exp(-dt * term)
     end
@@ -45,7 +45,45 @@ function is_nearest_neighbour(H::LocalOperator)
 end
 
 """
-    is_equivalent_bond(bond1::NTuple{2,CartesianIndex{2}}, bond2::NTuple{2,CartesianIndex{2}}, (Nrow, Ncol)::NTuple{2,Int})
+    is_equivalent_site(
+        site1::CartesianIndex{2}, site2::CartesianIndex{2},
+        (Nrow, Ncol)::NTuple{2, Int}
+    )
+
+Check if two lattice sites are related by a (periodic) lattice translation.
+"""
+function is_equivalent_site(
+        site1::CartesianIndex{2}, site2::CartesianIndex{2},
+        (Nrow, Ncol)::NTuple{2, Int}
+    )
+    shift = site1 - site2
+    return mod(shift[1], Nrow) == 0 && mod(shift[2], Ncol) == 0
+end
+
+"""
+    _get_site_term(ham::LocalOperator, site::CartesianIndex{2})
+
+Get the sum of all 1-site terms at `site` in `ham`.
+If there are no such terms, return the zero operator at `site`.
+"""
+function _get_site_term(ham::LocalOperator, site::CartesianIndex{2})
+    r, c = mod1.(Tuple(site), size(ham))
+    V = physicalspace(ham)[r, c]
+    term = zeros(scalartype(ham), V ← V)
+    for (sites, op) in ham.terms
+        length(sites) != 1 && continue
+        if is_equivalent_site(sites[1], site, size(ham))
+            term = term + op
+        end
+    end
+    return term
+end
+
+"""
+    is_equivalent_bond(
+        bond1::NTuple{2, CartesianIndex{2}}, bond2::NTuple{2, CartesianIndex{2}},
+        (Nrow, Ncol)::NTuple{2, Int},
+    )
 
 Check if two 2-site bonds are related by a (periodic) lattice translation.
 """
@@ -63,33 +101,26 @@ end
 """
     _get_bond_term(ham::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
 
-Get the 2-site term on `bond` in `ham`.
+Get the sum of all 2-site terms on `bond` in `ham`.
+If there are no such terms, return the zero operator on `bond`.
 """
 function _get_bond_term(ham::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
-    bonds = findall(p -> is_equivalent_bond(p.first, bond, size(ham.lattice)), ham.terms)
-    if length(bonds) == 0
-        # try reversed site order
-        bonds = findall(
-            p -> is_equivalent_bond(p.first, reverse(bond), size(ham.lattice)), ham.terms
-        )
-        if length(bonds) == 1
-            return permute(ham.terms[bonds[1]].second, ((2, 1), (4, 3)))
-        elseif length(bonds) == 0
-            # if term not found, return the zero operator on this bond
-            dtype = scalartype(ham)
-            r1, c1 = (mod1(bond[1][i], n) for (i, n) in zip(1:2, size(ham)))
-            r2, c2 = (mod1(bond[2][i], n) for (i, n) in zip(1:2, size(ham)))
-            V1 = physicalspace(ham)[r1, c1]
-            V2 = physicalspace(ham)[r2, c2]
-            return zeros(dtype, V1 ⊗ V2 ← V1 ⊗ V2)
-        else
-            error("There are multiple terms in `gate` corresponding to the bond $(bond).")
+    # create zero operator
+    r1, c1 = mod1.(Tuple(bond[1]), size(ham))
+    r2, c2 = mod1.(Tuple(bond[2]), size(ham))
+    V1 = physicalspace(ham)[r1, c1]
+    V2 = physicalspace(ham)[r2, c2]
+    term = zeros(scalartype(ham), V1 ⊗ V2 ← V1 ⊗ V2)
+    for (sites, op) in ham.terms
+        length(sites) != 2 && continue
+        if is_equivalent_bond(sites, bond, size(ham))
+            term += op
+        elseif is_equivalent_bond(sites, reverse(bond), size(ham))
+            op′ = permute(op, ((2, 1), (4, 3)); copy = true)
+            term += op′
         end
-    else
-        (length(bonds) == 1) ||
-            error("There are multiple terms in `gate` corresponding to the bond $(bond).")
-        return ham.terms[bonds[1]].second
     end
+    return term
 end
 
 """
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
index c6d1d8265..8a63aa922 100644
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ b/src/algorithms/time_evolution/trotter_mpo.jl
@@ -34,7 +34,7 @@ function TrotterMPOs2ndNeighbor(H::LocalOperator, dt::Number)
         _get_gatempos_se(H, dt),
         _get_gatempos_se(rotl90(H), dt),
         _get_gatempos_se(rot180(H), dt),
-        _get_gatempos_se(rotr90(H), dt)
+        _get_gatempos_se(rotr90(H), dt),
     ]
     return TrotterMPOs2ndNeighbor(mpos)
 end

From 6345f38c88a0fa0d58ce4c3678f9092f0e41e528 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 3 Mar 2026 19:54:53 +0800
Subject: [PATCH 05/27] Allow one-site terms in Hamiltonian for time evolution

---
 src/algorithms/time_evolution/time_evolve.jl  | 36 +++++++++++--
 src/algorithms/time_evolution/trotter_gate.jl | 53 ++++++++++---------
 test/timeevol/tf_ising_finiteT.jl             | 13 +----
 3 files changed, 63 insertions(+), 39 deletions(-)

diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index a6e2dbd5a..939bfb609 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -92,11 +92,41 @@ function MPSKit.infinite_temperature_density_matrix(H::LocalOperator)
     return InfinitePEPO(cat(A; dims = 3))
 end
 
+"""
+    _check_hamiltonian_for_trotter(H::LocalOperator)
+
+Assert that operator `H` contains only one-site and two-site terms,
+so that it is suited for Trotter time evolution algorithms.
+Returns the maximum squared distance covered by a two-site term in `H`.
+On the square lattice, the neighbor distances are
+```
+    1st nb.     2nd nb.     3rd nb.
+
+                    o
+                    |
+    o---o       o---o       o---o---o
+
+    dist² = 1   dist² = 2   dist² = 4
+```
+"""
+function _check_hamiltonian_for_trotter(H::LocalOperator)
+    dist = 0
+    for (sites, op) in H.terms
+        @assert numin(op) <= 2 "Hamiltonians containing multi-site (> 2) terms are not supported."
+        if numin(op) == 2
+            dist = max(dist, sum(Tuple(sites[1] - sites[2]) .^ 2))
+        end
+    end
+    @assert dist > 0 "Hamiltonians with only one-site terms are not suited for current implementation of Trotter time evolution focusing on 2-site gates."
+    @assert dist <= 2 "Hamiltonians with 2-site terms on beyond 2nd-neighbor bonds are not supported."
+    return dist
+end
+
 function trotterize(H::LocalOperator, dt::Number)
-    nnonly = is_nearest_neighbour(H)
-    gate = if nnonly
+    dist = _check_hamiltonian_for_trotter(H)
+    gate = if dist == 1
         TrotterGates1stNeighbor(H, dt)
-    else
+    elseif dist == 2
         TrotterMPOs2ndNeighbor(H, dt)
     end
     return gate
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 1b5730f78..1e9662db4 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -23,27 +23,24 @@ end
 
 function TrotterGates1stNeighbor(H::LocalOperator, dt::Number)
     Nr, Nc = size(H)
+    T = scalartype(H)
     gates = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
         # d = 1: horizontal bond; d = 2: vertical bond
         site1 = CartesianIndex(r, c)
         site2 = (d == 1) ? CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
-        term = _get_bond_term(H, (site1, site2))
+        s1term = _get_site_term(H, site1)
+        unit1 = TensorKit.id(T, space(s1term, 1))
+        s2term = _get_site_term(H, site2)
+        unit2 = TensorKit.id(T, space(s2term, 1))
+        b_term = _get_bond_term(H, (site1, site2))
+        # When iterating through horizontal and vertical bonds in the unit cell,
+        # each site / NN-bond is counted 4 / 1 times, respectively.
+        term = b_term + (s1term ⊗ unit2 + unit1 ⊗ s2term) / 4
         return exp(-dt * term)
     end
     return TrotterGates1stNeighbor(gates)
 end
 
-"""
-    is_nearest_neighbour(H::LocalOperator)
-
-Check if an operator `H` contains only nearest neighbor terms.
-"""
-function is_nearest_neighbour(H::LocalOperator)
-    return all(H.terms) do (sites, op)
-        return numin(op) == 2 && sum(abs, Tuple(sites[2] - sites[1])) == 1
-    end
-end
-
 """
     is_equivalent_site(
         site1::CartesianIndex{2}, site2::CartesianIndex{2},
@@ -128,33 +125,41 @@ end
 
 Construct the term acting on the southeast 3-site cluster in `ham`.
 ```
-    r-1        g3
-                |
-                ↓
-    r   g1 -←- g2
-        c      c+1
+    r-1     3
+            ↓
+    r   1-←-2
+        c   c+1
 ```
 """
 function _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
     Nr, Nc = size(ham)
+    T = scalartype(ham)
     @assert 1 <= row <= Nr && 1 <= col <= Nc
     sites = [
         CartesianIndex(row, col),
         CartesianIndex(row, col + 1),
         CartesianIndex(row - 1, col + 1),
     ]
+    ss = map(sites) do site
+        _get_site_term(ham, site)
+    end
     nb1x = _get_bond_term(ham, (sites[1], sites[2]))
     nb1y = _get_bond_term(ham, (sites[2], sites[3]))
     nb2 = _get_bond_term(ham, (sites[1], sites[3]))
     # identity operator at each site
     units = map(sites) do site
         site_ = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
-        return id(physicalspace(ham)[site_])
+        return id(T, physicalspace(ham)[site_])
     end
-    # when iterating through ┘, └, ┌, ┐ clusters in the unit cell,
-    # NN / NNN bonds are counted 4 / 2 times, respectively.
-    @tensor term[i' j' k'; i j k] :=
-        (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4 +
-        (nb2[i' k'; i k] * units[2][j'; j]) / 2
-    return term
+    # When iterating through ┘, └, ┌, ┐ clusters in the unit cell,
+    # each site / NN-bond / NNN-bond is counted 12 / 4 / 2 times, respectively.
+    term_site = (
+        ss[1] ⊗ units[2] ⊗ units[3] +
+            units[1] ⊗ ss[2] ⊗ units[3] +
+            units[1] ⊗ units[2] ⊗ ss[3]
+    ) / 12
+    @tensor term_nb1[i' j' k'; i j k] :=
+        (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4
+    @tensor term_nb2[i' j' k'; i j k] := (nb2[i' k'; i k] * units[2][j'; j]) / 2
+    return term_site + term_nb1 + term_nb2
 end
diff --git a/test/timeevol/tf_ising_finiteT.jl b/test/timeevol/tf_ising_finiteT.jl
index 8381d3213..087fe7d46 100644
--- a/test/timeevol/tf_ising_finiteT.jl
+++ b/test/timeevol/tf_ising_finiteT.jl
@@ -9,17 +9,6 @@ using PEPSKit
 bm_β = [0.5632, 0.0]
 bm_2β = [0.5297, 0.8265]
 
-# only contains 2-site gates, convenient for time evolution
-function tfising_model(T::Type{<:Number}, lattice::InfiniteSquare; J = 1.0, g = 1.0)
-    pspace, S = ℂ^2, Trivial
-    ZZ = rmul!(σᶻ(T, S) ⊗ σᶻ(T, S), -J)
-    X = rmul!(σˣ(T, S), g * -J)
-    unit = TensorKit.id(pspace)
-    spaces = fill(pspace, (lattice.Nrows, lattice.Ncols))
-    gate = ZZ + (unit ⊗ X + X ⊗ unit) / 4
-    return PEPSKit.nearest_neighbour_hamiltonian(spaces, gate)
-end
-
 function converge_env(state, χ::Int)
     trunc1 = truncrank(4) & truncerror(; atol = 1.0e-12)
     env0 = CTMRGEnv(ones, Float64, state, ℂ^1)
@@ -45,7 +34,7 @@ function measure_mag(pepo::InfinitePEPO, env::CTMRGEnv; purified::Bool = false)
 end
 
 Nr, Nc = 2, 2
-ham = tfising_model(Float64, InfiniteSquare(Nr, Nc); J = 1.0, g = 2.0)
+ham = transverse_field_ising(Float64, Trivial, InfiniteSquare(Nr, Nc); J = 1.0, g = 2.0)
 pepo0 = PEPSKit.infinite_temperature_density_matrix(ham)
 wts0 = SUWeight(pepo0)
 

From 975985ca9c40c9e78bdc55f56859d0fbdc832156 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 4 Mar 2026 12:29:36 +0800
Subject: [PATCH 06/27] Cleanups

---
 src/algorithms/time_evolution/simpleupdate.jl |  2 +-
 src/algorithms/time_evolution/time_evolve.jl  | 10 ++-
 src/algorithms/time_evolution/trotter_gate.jl | 62 ++------------
 src/algorithms/time_evolution/trotter_mpo.jl  | 82 ++++++++++++++-----
 4 files changed, 73 insertions(+), 83 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 1ef38b566..8f710d5c1 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -145,7 +145,7 @@ end
 One iteration of simple update
 """
 function su_iter(
-        state::InfiniteState, gate::TrotterGates1stNeighbor,
+        state::InfiniteState, gate::TrotterNNGates,
         alg::SimpleUpdate, env::SUWeight
     )
     Nr, Nc, = size(state)
diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index 939bfb609..ed1d9947b 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -112,20 +112,22 @@ On the square lattice, the neighbor distances are
 function _check_hamiltonian_for_trotter(H::LocalOperator)
     dist = 0
     for (sites, op) in H.terms
-        @assert numin(op) <= 2 "Hamiltonians containing multi-site (> 2) terms are not supported."
+        @assert numin(op) <= 2 "Hamiltonians containing multi-site (> 2) terms are not currently supported."
         if numin(op) == 2
             dist = max(dist, sum(Tuple(sites[1] - sites[2]) .^ 2))
         end
     end
-    @assert dist > 0 "Hamiltonians with only one-site terms are not suited for current implementation of Trotter time evolution focusing on 2-site gates."
-    @assert dist <= 2 "Hamiltonians with 2-site terms on beyond 2nd-neighbor bonds are not supported."
+    @assert dist > 0 "Hamiltonians with only one-site terms are not suited for current implementation of Trotter time evolution."
+    @assert dist <= 2 "Hamiltonians with 2-site terms on beyond 2nd-neighbor bonds are not currently supported."
     return dist
 end
 
 function trotterize(H::LocalOperator, dt::Number)
     dist = _check_hamiltonian_for_trotter(H)
+    # For nearest neighbor Hamiltonian, convert to 2-site Trotter gates.
+    # For all other cases, convert to multi-site Trotter MPOs.
     gate = if dist == 1
-        TrotterGates1stNeighbor(H, dt)
+        TrotterNNGates(H, dt)
     elseif dist == 2
         TrotterMPOs2ndNeighbor(H, dt)
     end
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 1e9662db4..46cee3fa8 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -1,14 +1,5 @@
 """
-$(TYPEDEF)
-
-Abstract super type for the collection of 2-body Trotter evolution gates.
-"""
-abstract type TrotterGates end
-
-Base.getindex(gates::TrotterGates, args...) = Base.getindex(gates.gates, args...)
-
-"""
-Collection of 1st (nearest) neighbor 2-body Trotter gates.
+Collection of nearest neighbor 2-body Trotter gates.
 
 Before exponentiating, terms in the Hamiltonian are organized as
 ```
@@ -17,11 +8,12 @@ Before exponentiating, terms in the Hamiltonian are organized as
 where each `Xᵢⱼ` (or `Yᵢⱼ`) acts on a horizontal (or vertical) bond.
 The Trotter gates are `exp(-dt * Xᵢⱼ)`, `exp(-dt * Yᵢⱼ)`.
 """
-struct TrotterGates1stNeighbor{G} <: TrotterGates
+struct TrotterNNGates{G} <: TrotterNNGates
     gates::G
 end
+Base.getindex(gates::TrotterNNGates, args...) = Base.getindex(gates.gates, args...)
 
-function TrotterGates1stNeighbor(H::LocalOperator, dt::Number)
+function TrotterNNGates(H::LocalOperator, dt::Number)
     Nr, Nc = size(H)
     T = scalartype(H)
     gates = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
@@ -38,7 +30,7 @@ function TrotterGates1stNeighbor(H::LocalOperator, dt::Number)
         term = b_term + (s1term ⊗ unit2 + unit1 ⊗ s2term) / 4
         return exp(-dt * term)
     end
-    return TrotterGates1stNeighbor(gates)
+    return TrotterNNGates(gates)
 end
 
 """
@@ -119,47 +111,3 @@ function _get_bond_term(ham::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
     end
     return term
 end
-
-"""
-    _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
-
-Construct the term acting on the southeast 3-site cluster in `ham`.
-```
-    r-1     3
-            ↓
-    r   1-←-2
-        c   c+1
-```
-"""
-function _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
-    Nr, Nc = size(ham)
-    T = scalartype(ham)
-    @assert 1 <= row <= Nr && 1 <= col <= Nc
-    sites = [
-        CartesianIndex(row, col),
-        CartesianIndex(row, col + 1),
-        CartesianIndex(row - 1, col + 1),
-    ]
-    ss = map(sites) do site
-        _get_site_term(ham, site)
-    end
-    nb1x = _get_bond_term(ham, (sites[1], sites[2]))
-    nb1y = _get_bond_term(ham, (sites[2], sites[3]))
-    nb2 = _get_bond_term(ham, (sites[1], sites[3]))
-    # identity operator at each site
-    units = map(sites) do site
-        site_ = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
-        return id(T, physicalspace(ham)[site_])
-    end
-    # When iterating through ┘, └, ┌, ┐ clusters in the unit cell,
-    # each site / NN-bond / NNN-bond is counted 12 / 4 / 2 times, respectively.
-    term_site = (
-        ss[1] ⊗ units[2] ⊗ units[3] +
-            units[1] ⊗ ss[2] ⊗ units[3] +
-            units[1] ⊗ units[2] ⊗ ss[3]
-    ) / 12
-    @tensor term_nb1[i' j' k'; i j k] :=
-        (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4
-    @tensor term_nb2[i' j' k'; i j k] := (nb2[i' k'; i k] * units[2][j'; j]) / 2
-    return term_site + term_nb1 + term_nb2
-end
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
index 8a63aa922..0d8f4ca96 100644
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ b/src/algorithms/time_evolution/trotter_mpo.jl
@@ -6,7 +6,7 @@ Trotter evolution MPOs acting on 3 or more sites.
 """
 abstract type TrotterMPOs end
 
-Base.getindex(gate::TrotterMPOs, args...) = Base.getindex(gate.mpos, args...)
+Base.getindex(gate::TrotterMPOs, args...) = Base.getindex(gate.data, args...)
 
 """
     struct TrotterMPOs2ndNeighbor{T}
@@ -22,21 +22,22 @@ where `┘`, `┐`, `┌`, `└` refer to the following 3-site clusters
         |       |   |       |
     1---2       1   3       2---3
 ```
-`mpos[d][i, j]` is the `┘ᵢⱼ` MPO acting on the `[i, j]` southeast
+`data[d][i, j]` is the `┘ᵢⱼ` MPO acting on the `[i, j]` southeast
 cluster after the network is left-rotated by `90 x (d - 1)` degrees.
 """
 struct TrotterMPOs2ndNeighbor{T} <: TrotterMPOs
-    mpos::T
+    data::T
 end
 
 function TrotterMPOs2ndNeighbor(H::LocalOperator, dt::Number)
-    mpos = [
-        _get_gatempos_se(H, dt),
-        _get_gatempos_se(rotl90(H), dt),
-        _get_gatempos_se(rot180(H), dt),
-        _get_gatempos_se(rotr90(H), dt),
-    ]
-    return TrotterMPOs2ndNeighbor(mpos)
+    return TrotterMPOs2ndNeighbor(
+        [
+            _get_gatempos_se(H, dt),
+            _get_gatempos_se(rotl90(H), dt),
+            _get_gatempos_se(rot180(H), dt),
+            _get_gatempos_se(rotr90(H), dt),
+        ]
+    )
 end
 
 """
@@ -61,7 +62,52 @@ function gate_to_mpo3(
 end
 
 """
-Obtain the 3-site gate MPO on the southeast cluster at position `[row, col]`
+    _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
+
+Construct the term acting on the southeast 3-site cluster in `ham`.
+```
+    r-1     3
+            ↓
+    r   1-←-2
+        c   c+1
+```
+"""
+function _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
+    Nr, Nc = size(ham)
+    T = scalartype(ham)
+    @assert 1 <= row <= Nr && 1 <= col <= Nc
+    sites = [
+        CartesianIndex(row, col),
+        CartesianIndex(row, col + 1),
+        CartesianIndex(row - 1, col + 1),
+    ]
+    ss = map(sites) do site
+        _get_site_term(ham, site)
+    end
+    nb1x = _get_bond_term(ham, (sites[1], sites[2]))
+    nb1y = _get_bond_term(ham, (sites[2], sites[3]))
+    nb2 = _get_bond_term(ham, (sites[1], sites[3]))
+    # identity operator at each site
+    units = map(sites) do site
+        site_ = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
+        return id(T, physicalspace(ham)[site_])
+    end
+    # When iterating through ┘, └, ┌, ┐ clusters in the unit cell,
+    # each site / NN-bond / NNN-bond is counted 12 / 4 / 2 times, respectively.
+    term_site = (
+        ss[1] ⊗ units[2] ⊗ units[3] +
+            units[1] ⊗ ss[2] ⊗ units[3] +
+            units[1] ⊗ units[2] ⊗ ss[3]
+    ) / 12
+    @tensor term_nb1[i' j' k'; i j k] :=
+        (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4
+    @tensor term_nb2[i' j' k'; i j k] := (nb2[i' k'; i k] * units[2][j'; j]) / 2
+    return term_site + term_nb1 + term_nb2
+end
+
+
+"""
+Obtain 3-site gate MPOs on southeast cluster at all positions `[row, col]`
 ```
     r-1        g3
                 |
@@ -70,16 +116,10 @@ Obtain the 3-site gate MPO on the southeast cluster at position `[row, col]`
         c      c+1
 ```
 """
-function _get_gatempo_se(ham::LocalOperator, dt::Number, row::Int, col::Int)
-    term = _get_se3site_term(ham, row, col)
-    return gate_to_mpo3(exp(-dt * term))
-end
-
-"""
-Construct the 3-site gate MPOs on the southeast cluster 
-for 3-site simple update on square lattice.
-"""
 function _get_gatempos_se(ham::LocalOperator, dt::Number)
     Nr, Nc = size(ham.lattice)
-    return collect(_get_gatempo_se(ham, dt, r, c) for r in 1:Nr, c in 1:Nc)
+    return map(Iterators.product(1:Nr, 1:Nc)) do (row, col)
+        term = _get_se3site_term(ham, row, col)
+        return gate_to_mpo3(exp(-dt * term))
+    end
 end

From e20d63e037de064631ec1c18ac0dd9c33bb95448 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 4 Mar 2026 12:44:58 +0800
Subject: [PATCH 07/27] Restore force_3site

---
 src/algorithms/time_evolution/simpleupdate.jl |  4 +-
 src/algorithms/time_evolution/time_evolve.jl  |  6 +--
 test/timeevol/cluster_projectors.jl           | 50 +++++++++++++++++++
 3 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 8f710d5c1..5e3e7c7c6 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -12,6 +12,8 @@ $(TYPEDFIELDS)
     trunc::TruncationStrategy
     "When true (or false), the Trotter gate is `exp(-H dt)` (or `exp(-iH dt)`)"
     imaginary_time::Bool = true
+    "When true, force the usage of 3-site simple update"
+    force_3site::Bool = false
     "When true, assume bipartite unit cell structure"
     bipartite::Bool = false
     "(Only applicable to InfinitePEPO) 
@@ -55,7 +57,7 @@ function TimeEvolver(
     _timeevol_sanity_check(psi0, physicalspace(H), alg)
     dt′ = _get_dt(psi0, dt, alg.imaginary_time)
     # create Trotter gates
-    gate = trotterize(H, dt′)
+    gate = trotterize(H, dt′; force_mpo = alg.force_3site)
     if isa(gate, TrotterMPOs)
         @assert !alg.bipartite "Trotter MPOs are incompatible with bipartite lattice structure."
     end
diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index ed1d9947b..af1aa0d58 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -122,13 +122,13 @@ function _check_hamiltonian_for_trotter(H::LocalOperator)
     return dist
 end
 
-function trotterize(H::LocalOperator, dt::Number)
+function trotterize(H::LocalOperator, dt::Number; force_mpo::Bool = false)
     dist = _check_hamiltonian_for_trotter(H)
     # For nearest neighbor Hamiltonian, convert to 2-site Trotter gates.
     # For all other cases, convert to multi-site Trotter MPOs.
-    gate = if dist == 1
+    gate = if dist == 1 && !force_mpo
         TrotterNNGates(H, dt)
-    elseif dist == 2
+    elseif (dist == 1 && force_mpo) || dist == 2
         TrotterMPOs2ndNeighbor(H, dt)
     end
     return gate
diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index faff68774..0ad3ae2ce 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -100,3 +100,53 @@ end
         end
     end
 end
+
+@testset "Hubbard model with 2-site and 3-site SU" begin
+    Nr, Nc = 2, 2
+    ctmrg_tol = 1.0e-9
+    Random.seed!(1459)
+    # with U(1) spin rotation symmetry
+    Pspace = hubbard_space(Trivial, U1Irrep)
+    Vspace = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2, (1, 1 // 2) => 1, (1, -1 // 2) => 1)
+    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, Vspace'; unitcell = (Nr, Nc))
+    # make state bipartite
+    for r in 1:2
+        peps.A[_next(r, 2), 2] = copy(peps.A[r, 1])
+    end
+    wts = SUWeight(peps)
+    ham = real(
+        hubbard_model(
+            ComplexF64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t = 1.0, U = 8.0, mu = 0.0
+        ),
+    )
+    # usual 2-site simple update, and measure energy
+    dts = [1.0e-2, 1.0e-2]
+    tols = [1.0e-8, 1.0e-8]
+    for (n, (dt, tol)) in enumerate(zip(dts, tols))
+        trunc = truncerror(; atol = 1.0e-10) & truncrank(n == 1 ? 4 : 2)
+        alg = SimpleUpdate(; trunc, bipartite = true)
+        peps, wts, = time_evolve(peps, ham, dt, 10000, alg, wts; tol, check_interval = 1000)
+    end
+    normalize!.(peps.A, Inf)
+    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)
+    @info "2-site simple update energy = $e_site"
+    # continue with 3-site simple update; energy should not change much
+    dts = [1.0e-2]
+    tols = [1.0e-8]
+    trunc = truncerror(; atol = 1.0e-10) & truncrank(2)
+    alg = SimpleUpdate(; trunc, force_3site = true)
+    for (n, (dt, tol)) in enumerate(zip(dts, tols))
+        peps, wts, = time_evolve(peps, ham, dt, 5000, alg, wts; tol, check_interval = 1000)
+    end
+    normalize!.(peps.A, Inf)
+    env, = leading_boundary(env, peps; tol = ctmrg_tol, trunc = trunc_env)
+    e_site2 = cost_function(peps, env, ham) / (Nr * Nc)
+    @info "3-site simple update energy = $e_site2"
+    @test e_site ≈ e_site2 atol = 5.0e-4
+end

From ed578809a950ac3f893d04b8542f94cd1c5a709c Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 4 Mar 2026 13:02:46 +0800
Subject: [PATCH 08/27] More cleanups

---
 src/algorithms/time_evolution/trotter_gate.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 46cee3fa8..c8bb623fb 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -8,7 +8,7 @@ Before exponentiating, terms in the Hamiltonian are organized as
 where each `Xᵢⱼ` (or `Yᵢⱼ`) acts on a horizontal (or vertical) bond.
 The Trotter gates are `exp(-dt * Xᵢⱼ)`, `exp(-dt * Yᵢⱼ)`.
 """
-struct TrotterNNGates{G} <: TrotterNNGates
+struct TrotterNNGates{G}
     gates::G
 end
 Base.getindex(gates::TrotterNNGates, args...) = Base.getindex(gates.gates, args...)

From 6759928b7876a287ff5abdaf77ff05920c2b1452 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 4 Mar 2026 15:16:48 +0800
Subject: [PATCH 09/27] Fix SU gauge fixing

---
 src/algorithms/time_evolution/apply_gate.jl   | 30 +++++++++++++++++++
 src/algorithms/time_evolution/gaugefix_su.jl  | 21 ++-----------
 src/algorithms/time_evolution/simpleupdate.jl |  8 ++---
 src/algorithms/time_evolution/trotter_gate.jl |  2 ++
 4 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/src/algorithms/time_evolution/apply_gate.jl b/src/algorithms/time_evolution/apply_gate.jl
index 08981666a..cbfecc7f9 100644
--- a/src/algorithms/time_evolution/apply_gate.jl
+++ b/src/algorithms/time_evolution/apply_gate.jl
@@ -126,3 +126,33 @@ function _apply_gate(
     end
     return a, s, b, ϵ
 end
+
+"""
+$(SIGNATURES)
+
+Apply a trivial gate (indicated by `nothing`) on the reduced matrices `a`, `b`
+```
+    -1← a --- 3 --- b ← -4          -2         -3
+        ↓           ↓               ↓           ↓
+        -2         -3           -1← a --- 3 --- b ← -4
+```
+"""
+function _apply_gate(
+        a::AbstractTensorMap, b::AbstractTensorMap,
+        ::Nothing, trunc::TruncationStrategy
+    )
+    V = space(b, 1)
+    need_flip = isdual(V)
+    @tensor a2b2[-1 -2; -3 -4] := a[-1 -2 3] * b[3 -3 -4]
+    trunc = if trunc isa FixedSpaceTruncation
+        need_flip ? truncspace(flip(V)) : truncspace(V)
+    else
+        trunc
+    end
+    a, s, b, ϵ = svd_trunc!(a2b2; trunc, alg = LAPACK_QRIteration())
+    a, b = absorb_s(a, s, b)
+    if need_flip
+        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/gaugefix_su.jl b/src/algorithms/time_evolution/gaugefix_su.jl
index d9c7d4aaf..47e5dd6cc 100644
--- a/src/algorithms/time_evolution/gaugefix_su.jl
+++ b/src/algorithms/time_evolution/gaugefix_su.jl
@@ -17,31 +17,14 @@ $(TYPEDFIELDS)
     maxiter::Int = 100
 end
 
-"""
-A LocalOperator consisting of identity gates on all nearest neighbor bonds.
-"""
-function _trivial_gates(elt::Type{<:Number}, lattice::Matrix{S}) where {S <: ElementarySpace}
-    terms = []
-    for site1 in CartesianIndices(lattice)
-        r1, c1 = mod1.(Tuple(site1), size(lattice))
-        for d in (CartesianIndex(1, 0), CartesianIndex(0, 1))
-            site2 = site1 + d
-            r2, c2 = mod1.(Tuple(site2), size(lattice))
-            V1, V2 = lattice[r1, c1], lattice[r2, c2]
-            h = TensorKit.id(elt, V1 ⊗ V2)
-            push!(terms, (site1, site2) => h)
-        end
-    end
-    return LocalOperator(lattice, terms...)
-end
-
 """
     gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::SUGauge)
 
 Fix the gauge of `psi` using trivial simple update.
 """
 function gauge_fix(psi::InfiniteState, alg::SUGauge)
-    gates = _trivial_gates(scalartype(psi), physicalspace(psi))
+    Nr, Nc = size(psi)
+    gates = TrotterNNGates(fill(nothing, (2, Nr, Nc)))
     su_alg = SimpleUpdate(; trunc = FixedSpaceTruncation(), bipartite = _state_bipartite_check(psi))
     wts0 = SUWeight(psi)
     # use default constructor to avoid calculation of exp(-H * 0)
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 5e3e7c7c6..9e4f2a660 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -76,9 +76,9 @@ When `gate_ax = 1` (or `2`), the gate will be applied to
 the codomain (or domain) physicsl legs of `state`.
 """
 function _su_xbond!(
-        state::InfiniteState, gate::AbstractTensorMap{T, S, 2, 2}, env::SUWeight,
+        state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
         row::Int, col::Int, trunc::TruncationStrategy; gate_ax::Int = 1
-    ) where {T <: Number, S <: ElementarySpace}
+    )
     Nr, Nc, = size(state)
     cp1 = _next(col, Nc)
     # absorb environment weights
@@ -116,9 +116,9 @@ When `gate_ax = 1` (or `2`), the gate will be applied to
 the codomain (or domain) physicsl legs of `state`.
 """
 function _su_ybond!(
-        state::InfiniteState, gate::AbstractTensorMap{T, S, 2, 2}, env::SUWeight,
+        state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
         row::Int, col::Int, trunc::TruncationStrategy; gate_ax::Int = 1
-    ) where {T <: Number, S <: ElementarySpace}
+    )
     Nr, Nc, = size(state)
     rm1 = _prev(row, Nr)
     # absorb environment weights
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index c8bb623fb..4af9ab43c 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -13,6 +13,8 @@ struct TrotterNNGates{G}
 end
 Base.getindex(gates::TrotterNNGates, args...) = Base.getindex(gates.gates, args...)
 
+const NNGate{T, S} = AbstractTensorMap{T, S, 2, 2}
+
 function TrotterNNGates(H::LocalOperator, dt::Number)
     Nr, Nc = size(H)
     T = scalartype(H)

From deb0e2af264de83497f6ad6756c822fde0c60e0a Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 4 Mar 2026 15:41:47 +0800
Subject: [PATCH 10/27] Generalize gate_to_mpo to any number of sites

---
 src/algorithms/time_evolution/simpleupdate.jl |  6 ++---
 src/algorithms/time_evolution/trotter_mpo.jl  | 26 ++++++++++++-------
 test/timeevol/cluster_projectors.jl           |  6 ++---
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 9e4f2a660..b3b310186 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -12,8 +12,8 @@ $(TYPEDFIELDS)
     trunc::TruncationStrategy
     "When true (or false), the Trotter gate is `exp(-H dt)` (or `exp(-iH dt)`)"
     imaginary_time::Bool = true
-    "When true, force the usage of 3-site simple update"
-    force_3site::Bool = false
+    "When true, force the usage of MPO simple update for nearest neighbor Hamiltonians."
+    force_mpo::Bool = false
     "When true, assume bipartite unit cell structure"
     bipartite::Bool = false
     "(Only applicable to InfinitePEPO) 
@@ -57,7 +57,7 @@ function TimeEvolver(
     _timeevol_sanity_check(psi0, physicalspace(H), alg)
     dt′ = _get_dt(psi0, dt, alg.imaginary_time)
     # create Trotter gates
-    gate = trotterize(H, dt′; force_mpo = alg.force_3site)
+    gate = trotterize(H, dt′; force_mpo = alg.force_mpo)
     if isa(gate, TrotterMPOs)
         @assert !alg.bipartite "Trotter MPOs are incompatible with bipartite lattice structure."
     end
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
index 0d8f4ca96..bf5038eaf 100644
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ b/src/algorithms/time_evolution/trotter_mpo.jl
@@ -41,24 +41,30 @@ function TrotterMPOs2ndNeighbor(H::LocalOperator, dt::Number)
 end
 
 """
-Convert a 3-site gate to MPO form by SVD, 
+Convert an N-site gate to MPO form by SVD, 
 in which the axes are ordered as
 ```
+    site 1      mid sites      site N
     2               3               3
     ↓               ↓               ↓
-    g1 ←- 3    1 ←- g2 ←- 4    1 ←- g3
+    g1 ←- 3    1 ←- g ←- 4    1 ←- gN
     ↓               ↓               ↓
     1               2               2
 ```
 """
-function gate_to_mpo3(
-        gate::AbstractTensorMap{T, S, 3, 3}, trunc = trunctol(; atol = MPSKit.Defaults.tol)
-    ) where {T <: Number, S <: ElementarySpace}
+function gate_to_mpo(
+        gate::AbstractTensorMap{T, S, N, N}, trunc = trunctol(; atol = MPSKit.Defaults.tol)
+    ) where {T <: Number, S <: ElementarySpace, N}
     Os = MPSKit.decompose_localmpo(MPSKit.add_util_leg(gate), trunc)
-    g1 = removeunit(Os[1], 1)
-    g2 = Os[2]
-    g3 = removeunit(Os[3], 4)
-    return [g1, g2, g3]
+    return map(1:N) do i
+        if i == 1
+            return removeunit(Os[1], 1)
+        elseif i == N
+            return removeunit(Os[N], 4)
+        else
+            return Os[i]
+        end
+    end
 end
 
 """
@@ -120,6 +126,6 @@ function _get_gatempos_se(ham::LocalOperator, dt::Number)
     Nr, Nc = size(ham.lattice)
     return map(Iterators.product(1:Nr, 1:Nc)) do (row, col)
         term = _get_se3site_term(ham, row, col)
-        return gate_to_mpo3(exp(-dt * term))
+        return gate_to_mpo(exp(-dt * term))
     end
 end
diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index 0ad3ae2ce..5c49141b9 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -70,7 +70,7 @@ end
         flips = [isdual(space(M, 1)) for M in Ms1[2:end]]
         unit = id(Vphy)
         gate = reduce(⊗, fill(unit, 3))
-        gs = PEPSKit.gate_to_mpo3(gate)
+        gs = PEPSKit.gate_to_mpo(gate)
         @test mpo_to_gate3(gs) ≈ gate
         Ms2 = _flip_virtuals!(deepcopy(Ms1), flips)
         PEPSKit._apply_gatempo!(Ms2, gs)
@@ -87,7 +87,7 @@ end
         flips = [isdual(space(M, 1)) for M in Ms1[2:end]]
         unit = id(Vphy)
         gate = reduce(⊗, fill(unit, 3))
-        gs = PEPSKit.gate_to_mpo3(gate)
+        gs = PEPSKit.gate_to_mpo(gate)
         @test mpo_to_gate3(gs) ≈ gate
         for gate_ax in 1:2
             Ms2 = _flip_virtuals!(deepcopy(Ms1), flips)
@@ -140,7 +140,7 @@ end
     dts = [1.0e-2]
     tols = [1.0e-8]
     trunc = truncerror(; atol = 1.0e-10) & truncrank(2)
-    alg = SimpleUpdate(; trunc, force_3site = true)
+    alg = SimpleUpdate(; trunc, force_mpo = true)
     for (n, (dt, tol)) in enumerate(zip(dts, tols))
         peps, wts, = time_evolve(peps, ham, dt, 5000, alg, wts; tol, check_interval = 1000)
     end

From 47e23e9aff3909bbcd915720f2330d8948512ac9 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Thu, 5 Mar 2026 16:50:44 +0800
Subject: [PATCH 11/27] Multi-site simple update

---
 src/algorithms/time_evolution/apply_mpo.jl    |  11 -
 .../time_evolution/simpleupdate3site.jl       | 235 +++++++++++++-----
 src/algorithms/time_evolution/time_evolve.jl  |   9 +-
 src/algorithms/time_evolution/trotter_gate.jl |   2 +-
 src/algorithms/time_evolution/trotter_mpo.jl  | 148 +++++------
 test/timeevol/j1j2_finiteT.jl                 |   2 +-
 6 files changed, 256 insertions(+), 151 deletions(-)

diff --git a/src/algorithms/time_evolution/apply_mpo.jl b/src/algorithms/time_evolution/apply_mpo.jl
index 963540eb3..b60818530 100644
--- a/src/algorithms/time_evolution/apply_mpo.jl
+++ b/src/algorithms/time_evolution/apply_mpo.jl
@@ -283,17 +283,6 @@ e.g. Cluster in PEPS with `gate_ax = 1`:
         g1 -←-- g2 -←-- g3
         ↓       ↓       ↓
 ```
-
-In the cluster, the axes of each tensor use the MPS order
-```
-    PEPS:           PEPO:
-           3             3  4
-          ╱              | ╱
-    1 -- M -- 5     1 -- M -- 6
-       ╱ |             ╱ |
-      4  2            5  2
-    M[1 2 3 4; 5]  M[1 2 3 4 5; 6]
-```
 """
 function _apply_gatempo!(
         Ms::Vector{T1}, gs::Vector{T2}; gate_ax::Int = 1
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index 35a31206e..99f5b3a4f 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -13,62 +13,153 @@ function _unfuse_physicalspace(
     return O_unfused, F
 end
 
-const openaxs_se = [(NORTH, SOUTH, WEST), (EAST, SOUTH), (NORTH, EAST, WEST)]
-const invperms_se_peps = [((2,), (3, 5, 4, 1)), ((2,), (5, 3, 4, 1)), ((2,), (5, 3, 1, 4))]
-const perms_se_peps = map(invperms_se_peps) do (p1, p2)
-    p = invperm((p1..., p2...))
-    return (p[1:(end - 1)], (p[end],))
+"""
+Convert nearest neighbor vector `nn_vec` to direction labels.
+```
+                NORTH
+                (-1,0)
+                  ↑
+    WEST (0,-1)-←-∘-→-(0,+1) EAST
+                  ↓
+                (+1,0)
+                SOUTH
+```
+"""
+function _nn_vec_direction(nn_vec::CartesianIndex{2})
+    if nn_vec == CartesianIndex(-1, 0)
+        return NORTH
+    elseif nn_vec == CartesianIndex(0, 1)
+        return EAST
+    elseif nn_vec == CartesianIndex(1, 0)
+        return SOUTH
+    elseif nn_vec == CartesianIndex(0, -1)
+        return WEST
+    else
+        error("Input is not a nearest neighbor vector")
+    end
 end
-const invperms_se_pepo = [((2, 3), (4, 6, 5, 1)), ((2, 3), (6, 4, 5, 1)), ((2, 3), (6, 4, 1, 5))]
-const perms_se_pepo = map(invperms_se_pepo) do (p1, p2)
-    p = invperm((p1..., p2...))
-    return (p[1:(end - 1)], (p[end],))
+
+"""
+Find the permutation to permute `out_ax`, `in_ax` legs to
+the first and the last position of a tensor with `Nax` legs,
+then assign the last leg to domain, and the others to codomain.
+"""
+function _get_mpo_perm(out_ax::Int, in_ax::Int, Nax::Int)
+    perm = collect(1:Nax)
+    filter!(x -> x != out_ax && x != in_ax, perm)
+    pushfirst!(perm, out_ax)
+    push!(perm, in_ax)
+    return (Tuple(perm[1:(end - 1)]), (perm[end],))
 end
+
 """
-Obtain the 3-site cluster in the "southeast corner" of a square plaquette.
-``` 
-    r-1         M3
-                |
-                ↓
-    r   M1 -←- M2
-        c      c+1
+Obtain the cluster `Ms` along the (open and non-self-intersecting) path
+given by `sites` in `state`.
+
+Also returns:
+- `open_vaxs`: Open virtual axes (1 to 4) of each cluster tensor before permutation.
+- `bond_revs`: Position of the cluster bonds in the state, and whether the cluster path follows the "canonical" orientation of the bond (leftwards or downwards).
+- `invperms`: Permutations to restore the axes order of each cluster tensor.
+
+In the cluster, each tensor use the MPS axes order
+```
+    PEPS:           PEPO:
+           3             3  4
+          ╱              | ╱
+    o -- M -- i     o -- M -- i
+       ╱ |             ╱ |
+      4  2            5  2
+    M[o 2 3 4; i]  M[o 2 3 4 5; i]
 ```
+where `o` (`i`) connects to the previous (next) tensor.
 """
-function get_3site_se(state::InfiniteState, env::SUWeight, row::Int, col::Int)
+function _get_cluster(
+        state::InfiniteState, sites::Vector{CartesianIndex{2}},
+        env::Union{SUWeight, Nothing}
+    )
     Nr, Nc = size(state)
-    rm1, cp1 = _prev(row, Nr), _next(col, Nc)
-    coords_se = [(row, col), (row, cp1), (rm1, cp1)]
-    perms_se = isa(state, InfinitePEPS) ? perms_se_peps : perms_se_pepo
-    Ms = map(zip(coords_se, perms_se, openaxs_se)) do (coord, perm, openaxs)
-        M = absorb_weight(state.A[CartesianIndex(coord)], env, coord[1], coord[2], openaxs)
-        # permute to MPS axes order
+    # number of sites
+    Ns = length(sites)
+    # number of physical axes
+    Np = isa(state, InfinitePEPS) ? 1 : 2
+    # number of axes of each state tensor
+    Nax = 4 + Np
+    out_axs = map(2:Ns) do i
+        return _nn_vec_direction(sites[i - 1] - sites[i])
+    end
+    in_axs = map(1:(Ns - 1)) do i
+        return _nn_vec_direction(sites[i + 1] - sites[i])
+    end
+    all_vaxs = Tuple(1:4)
+    open_vaxs = map(1:Ns) do i
+        return if i == 1
+            filter(x -> x != in_axs[i], all_vaxs)
+        elseif i == Ns
+            filter(x -> x != out_axs[i - 1], all_vaxs)
+        else
+            filter(x -> x != out_axs[i - 1] && x != in_axs[i], all_vaxs)
+        end
+    end
+    bond_revs = map(zip(sites, Iterators.drop(sites, 1))) do (site1, site2)
+        diff = site1 - site2
+        if diff == CartesianIndex(0, -1)
+            r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)
+            return (1, r, c), false
+        elseif diff == CartesianIndex(0, 1)
+            r, c = mod1(site2[1], Nr), mod1(site2[2], Nc)
+            return (1, r, c), true
+        elseif diff == CartesianIndex(1, 0)
+            r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)
+            return (2, r, c), false
+        else # diff == CartesianIndex(-1, 0)
+            r, c = mod1(site2[1], Nr), mod1(site2[2], Nc)
+            return (2, r, c), true
+        end
+    end
+    perms = map(1:Ns) do i
+        out_ax, in_ax = if i == 1
+            # use direction opposite to `in` as `out`
+            mod1(2 + in_axs[i], 4), in_axs[i]
+        elseif i == Ns
+            # use direction opposite to `out` as `in`
+            out_axs[i - 1], mod1(2 + out_axs[i - 1], 4)
+        else
+            out_axs[i - 1], in_axs[i]
+        end
+        return _get_mpo_perm(out_ax + Np, in_ax + Np, Nax)
+    end
+    invperms = map(perms) do (p1, p2)
+        p = invperm((p1..., p2...))
+        return (p[1:Np], p[(Np + 1):end])
+    end
+    Ms = map(zip(sites, open_vaxs, perms)) do (site, vaxs, perm)
+        s = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
+        M = if env === nothing
+            state.A[s]
+        else
+            absorb_weight(state.A[s], env, s[1], s[2], vaxs)
+        end
         return permute(M, perm)
     end
-    return Ms
+    return Ms, open_vaxs, bond_revs, invperms
 end
 
-function _su3site_se!(
+function _su_cluster!(
         state::InfiniteState, gs::Vector{T}, env::SUWeight,
-        row::Int, col::Int, truncs::Vector{E};
+        sites::Vector{CartesianIndex{2}}, truncs::Vector{E};
         purified::Bool = true
     ) where {T <: AbstractTensorMap, E <: TruncationStrategy}
     Nr, Nc = size(state)
-    @assert 1 <= row <= Nr && 1 <= col <= Nc
-    rm1, cp1 = _prev(row, Nr), _next(col, Nc)
     # southwest 3-site cluster and arrow direction within it
-    Ms = get_3site_se(state, env, row, col)
+    Ms, open_vaxs, bond_revs, invperms = _get_cluster(state, sites, env)
     flips = [isdual(space(M, 1)) for M in Ms[2:end]]
     Vphys = [codomain(M, 2) for M in Ms]
     normalize!.(Ms, Inf)
     # flip virtual arrows in `Ms` to ←
     _flip_virtuals!(Ms, flips)
-    # sites in the cluster
-    coords = ((row, col), (row, cp1), (rm1, cp1))
-    # weights in the cluster
-    wt_idxs = ((1, row, col), (2, row, cp1))
     # apply gate MPOs and truncate
     gate_axs = purified ? (1:1) : (1:2)
-    ϵs = nothing
+    wts, ϵs = nothing, nothing
     for gate_ax in gate_axs
         _apply_gatempo!(Ms, gs; gate_ax)
         if isa(state, InfinitePEPO)
@@ -78,51 +169,65 @@ function _su3site_se!(
         if isa(state, InfinitePEPO)
             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 ? _fliptwist_s(wt) : wt, Inf)
-        end
     end
     # restore virtual arrows in `Ms`
     _flip_virtuals!(Ms, flips)
-    # update `state` from `Ms`
-    invperms_se = isa(state, InfinitePEPS) ? invperms_se_peps : invperms_se_pepo
-    for (M, coord, invperm, openaxs, Vphy) in zip(Ms, coords, invperms_se, openaxs_se, Vphys)
+    # update env weights
+    for (wt, (bond, rev), flip) in zip(wts, bond_revs, flips)
+        wt_new = flip ? _fliptwist_s(wt) : wt
+        wt_new = rev ? transpose(wt_new) : wt_new
+        @assert all(wt_new.data .>= 0)
+        env[CartesianIndex(bond)] = normalize(wt_new, Inf)
+    end
+    for (M, s, invperm, vaxs) in zip(Ms, sites, invperms, open_vaxs)
+        s′ = CartesianIndex(mod1(s[1], Nr), mod1(s[2], Nc))
         # restore original axes order
         M = permute(M, invperm)
         # remove weights on open axes of the cluster
-        M = absorb_weight(M, env, coord[1], coord[2], openaxs; inv = true)
-        state.A[CartesianIndex(coord)] = normalize(M, Inf)
+        M = absorb_weight(M, env, s′[1], s′[2], vaxs; inv = true)
+        # update state tensors
+        state.A[s′] = normalize(M, Inf)
     end
     return maximum(ϵs)
 end
 
-function su_iter(
-        state::InfiniteState, gatempos::TrotterMPOs2ndNeighbor,
-        alg::SimpleUpdate, env::SUWeight
+"""
+Get the `TruncationStrategy` for each bond in the cluster
+updated by the Trotter evolution MPO.
+"""
+function _get_cluster_trunc(
+        trunc::TruncationStrategy, sites::Vector{CartesianIndex{2}},
+        (Nrow, Ncol)::NTuple{2, Int}
     )
-    if state isa InfinitePEPO
-        @assert size(state, 3) == 1
+    return map(zip(sites, Iterators.drop(sites, 1))) do (site1, site2)
+        diff = site2 - site1
+        if diff == CartesianIndex(0, 1)
+            r, c = mod1(site1[1], Nrow), mod1(site1[2], Ncol)
+            return truncation_strategy(trunc, 1, r, c)
+        elseif diff == CartesianIndex(0, -1)
+            r, c = mod1(site2[1], Nrow), mod1(site2[2], Ncol)
+            return truncation_strategy(trunc, 1, r, c)
+        elseif diff == CartesianIndex(1, 0)
+            r, c = mod1(site2[1], Nrow), mod1(site2[2], Ncol)
+            return truncation_strategy(trunc, 2, r, c)
+        elseif diff == CartesianIndex(-1, 0)
+            r, c = mod1(site1[1], Nrow), mod1(site1[2], Ncol)
+            return truncation_strategy(trunc, 2, r, c)
+        else
+            error("The path `sites` contains a long-range bond.")
+        end
     end
-    Nr, Nc = size(state)[1:2]
-    (Nr >= 2 && Nc >= 2) || throw(
-        ArgumentError(
-            "iPEPS unit cell size for simple update should be no smaller than (2, 2)."
-        ),
+end
+
+function su_iter(
+        state::InfiniteState, gatempos::TrotterMPOs,
+        alg::SimpleUpdate, env::SUWeight
     )
     state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
-    trunc = alg.trunc
-    for i in 1:4
-        Nr, Nc = size(state2)[1:2]
-        for r in 1:Nr, c in 1:Nc
-            gs = gatempos[i][r, c]
-            truncs = [
-                truncation_strategy(trunc, 1, r, c)
-                truncation_strategy(trunc, 2, r, _next(c, Nc))
-            ]
-            ϵ = _su3site_se!(state2, gs, env2, r, c, truncs; alg.purified)
-        end
-        state2, env2 = rotl90(state2), rotl90(env2)
-        trunc = rotl90(trunc)
+    for (sites, gs) in gatempos.data
+        truncs = _get_cluster_trunc(alg.trunc, sites, size(state)[1:2])
+        ϵ′ = _su_cluster!(state2, gs, env2, sites, truncs; alg.purified)
+        ϵ = max(ϵ, ϵ′)
     end
     return state2, env2, ϵ
 end
diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index af1aa0d58..61b5dfa09 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -70,7 +70,10 @@ function _timeevol_sanity_check(
         ψ₀::InfiniteState, Pspaces::M, alg::A
     ) where {A <: TimeEvolution, M <: AbstractMatrix{<:ElementarySpace}}
     Nr, Nc, = size(ψ₀)
-    @assert (Nr >= 2 && Nc >= 2) "Unit cell size for simple update should be no smaller than (2, 2)."
+    @assert (Nr >= 2 && Nc >= 2) "Unit cell size for time evolution should be no smaller than (2, 2)."
+    if ψ₀ isa InfinitePEPO
+        @assert size(ψ₀, 3) == 1 "PEPO to be time evolved should have only one layer."
+    end
     @assert Pspaces == physicalspace(ψ₀) "Physical spaces of `ψ₀` do not match `Pspaces`."
     if hasfield(typeof(alg), :purified) && !alg.purified
         @assert ψ₀ isa InfinitePEPO "alg.purified = false is only applicable to PEPO."
@@ -127,9 +130,9 @@ function trotterize(H::LocalOperator, dt::Number; force_mpo::Bool = false)
     # For nearest neighbor Hamiltonian, convert to 2-site Trotter gates.
     # For all other cases, convert to multi-site Trotter MPOs.
     gate = if dist == 1 && !force_mpo
-        TrotterNNGates(H, dt)
+        trotterize_nn(H, dt)
     elseif (dist == 1 && force_mpo) || dist == 2
-        TrotterMPOs2ndNeighbor(H, dt)
+        trotterize_nnn(H, dt)
     end
     return gate
 end
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 4af9ab43c..8db4eddc5 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -15,7 +15,7 @@ Base.getindex(gates::TrotterNNGates, args...) = Base.getindex(gates.gates, args.
 
 const NNGate{T, S} = AbstractTensorMap{T, S, 2, 2}
 
-function TrotterNNGates(H::LocalOperator, dt::Number)
+function trotterize_nn(H::LocalOperator, dt::Number)
     Nr, Nc = size(H)
     T = scalartype(H)
     gates = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
index bf5038eaf..6bb7703e9 100644
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ b/src/algorithms/time_evolution/trotter_mpo.jl
@@ -1,45 +1,13 @@
 """
-$(TYPEDEF)
+    struct TrotterMPOs{T <: Vector}
 
-Abstract super type for the collection of
-Trotter evolution MPOs acting on 3 or more sites.
+Collection of Trotter evolution MPOs obtained from
+a Hamiltonian containing long-range or multi-site terms
 """
-abstract type TrotterMPOs end
-
-Base.getindex(gate::TrotterMPOs, args...) = Base.getindex(gate.data, args...)
-
-"""
-    struct TrotterMPOs2ndNeighbor{T}
-
-Collection of all Trotter evolution MPOs obtained from
-a Hamiltonian containing up to 2nd neighbor terms
-```
-    H = ∑ᵢⱼ(┘ᵢⱼ + ┐ᵢⱼ + ┌ᵢⱼ + └ᵢⱼ)
-```
-where `┘`, `┐`, `┌`, `└` refer to the following 3-site clusters 
-```
-        3   3---2   2---1   1
-        |       |   |       |
-    1---2       1   3       2---3
-```
-`data[d][i, j]` is the `┘ᵢⱼ` MPO acting on the `[i, j]` southeast
-cluster after the network is left-rotated by `90 x (d - 1)` degrees.
-"""
-struct TrotterMPOs2ndNeighbor{T} <: TrotterMPOs
+struct TrotterMPOs{T <: Vector}
     data::T
 end
 
-function TrotterMPOs2ndNeighbor(H::LocalOperator, dt::Number)
-    return TrotterMPOs2ndNeighbor(
-        [
-            _get_gatempos_se(H, dt),
-            _get_gatempos_se(rotl90(H), dt),
-            _get_gatempos_se(rot180(H), dt),
-            _get_gatempos_se(rotr90(H), dt),
-        ]
-    )
-end
-
 """
 Convert an N-site gate to MPO form by SVD, 
 in which the axes are ordered as
@@ -67,26 +35,84 @@ function gate_to_mpo(
     end
 end
 
-"""
-    _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
+# Specialized functions to trotterize Hamiltonians
+# with long-range/multi-site terms
+
+## Next-nearest neighbor H
 
-Construct the term acting on the southeast 3-site cluster in `ham`.
+"""
+Trotterize a Hamiltonian containing up to 2nd neighbor terms.
+```
+    H = ∑ᵢⱼ(Γᵢⱼ + ⅂ᵢⱼ + ⅃ᵢⱼ + Lᵢⱼ)
 ```
-    r-1     3
-            ↓
-    r   1-←-2
-        c   c+1
+where `Γ`, `⅂`, `⅃`, `L` refer to the following 3-site clusters 
 ```
+    NORTHWEST   NORTHEAST
+        2---1   3---2
+        |           |
+        3           1
+
+        1           3
+        |           |
+        2---3   1---2
+    SOUTHWEST   SOUTHEAST
+```
+acting on the the elemental square plaquette
+with southwest corner at `[i, j]`.
+"""
+function trotterize_nnn(H::LocalOperator, dt::Number)
+    Nr, Nc = size(H)
+    # iterate through corner `d` in outermost loop
+    terms = map(Iterators.product(1:Nr, 1:Nc, 1:4)) do (r, c, d)
+        return _get_nnn_mpo(H, dt, d, r, c)
+    end
+    return TrotterMPOs(vec(terms))
+end
+
 """
-function _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
+Get coordinates of sites in the 3-site triangular cluster
+used in Trotter evolution with next-nearest neighbor gates,
+with southwest corner at `[row, col]`.
+"""
+function _nnn_cluster_sites(dir::Int, row::Int, col::Int)
+    @assert 1 <= dir <= 4
+    return if dir == NORTHWEST
+        [
+            CartesianIndex(row - 1, col + 1),
+            CartesianIndex(row - 1, col),
+            CartesianIndex(row, col),
+        ]
+    elseif dir == NORTHEAST
+        [
+            CartesianIndex(row, col + 1),
+            CartesianIndex(row - 1, col + 1),
+            CartesianIndex(row - 1, col),
+        ]
+    elseif dir == SOUTHEAST
+        [
+            CartesianIndex(row, col),
+            CartesianIndex(row, col + 1),
+            CartesianIndex(row - 1, col + 1),
+        ]
+    else # dir == SOUTHWEST
+        [
+            CartesianIndex(row - 1, col),
+            CartesianIndex(row, col),
+            CartesianIndex(row, col + 1),
+        ]
+    end
+end
+
+"""
+Construct the evolution MPO acting on the 3-site triangular cluster
+in the square plaquette whose southwest corner is at `[row, col]`.
+`dir` takes values between 1 (`NORTHWEST`) and 4 (`SOUTHWEST`).
+"""
+function _get_nnn_mpo(ham::LocalOperator, dt::Number, dir::Int, row::Int, col::Int)
     Nr, Nc = size(ham)
     T = scalartype(ham)
     @assert 1 <= row <= Nr && 1 <= col <= Nc
-    sites = [
-        CartesianIndex(row, col),
-        CartesianIndex(row, col + 1),
-        CartesianIndex(row - 1, col + 1),
-    ]
+    sites = _nnn_cluster_sites(dir, row, col)
     ss = map(sites) do site
         _get_site_term(ham, site)
     end
@@ -98,7 +124,7 @@ function _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
         site_ = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
         return id(T, physicalspace(ham)[site_])
     end
-    # When iterating through ┘, └, ┌, ┐ clusters in the unit cell,
+    # When iterating through all triangular clusters in the unit cell,
     # each site / NN-bond / NNN-bond is counted 12 / 4 / 2 times, respectively.
     term_site = (
         ss[1] ⊗ units[2] ⊗ units[3] +
@@ -108,24 +134,6 @@ function _get_se3site_term(ham::LocalOperator, row::Int, col::Int)
     @tensor term_nb1[i' j' k'; i j k] :=
         (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4
     @tensor term_nb2[i' j' k'; i j k] := (nb2[i' k'; i k] * units[2][j'; j]) / 2
-    return term_site + term_nb1 + term_nb2
-end
-
-
-"""
-Obtain 3-site gate MPOs on southeast cluster at all positions `[row, col]`
-```
-    r-1        g3
-                |
-                ↓
-    r   g1 -←- g2
-        c      c+1
-```
-"""
-function _get_gatempos_se(ham::LocalOperator, dt::Number)
-    Nr, Nc = size(ham.lattice)
-    return map(Iterators.product(1:Nr, 1:Nc)) do (row, col)
-        term = _get_se3site_term(ham, row, col)
-        return gate_to_mpo(exp(-dt * term))
-    end
+    term = term_site + term_nb1 + term_nb2
+    return sites => gate_to_mpo(exp(-dt * term))
 end
diff --git a/test/timeevol/j1j2_finiteT.jl b/test/timeevol/j1j2_finiteT.jl
index ca74a6539..63576e6a1 100644
--- a/test/timeevol/j1j2_finiteT.jl
+++ b/test/timeevol/j1j2_finiteT.jl
@@ -26,9 +26,9 @@ wts0 = SUWeight(pepo0)
 # 7 = 1 (spin-0) + 2 x 3 (spin-1)
 trunc_pepo = truncrank(7) & truncerror(; atol = 1.0e-12)
 check_interval = 100
+dt, nstep = 1.0e-3, 600
 
 # PEPO approach
-dt, nstep = 1.0e-3, 600
 alg = SimpleUpdate(; trunc = trunc_pepo, purified = false)
 evolver = TimeEvolver(pepo0, ham, dt, nstep, alg, wts0)
 pepo, wts, info = time_evolve(evolver; check_interval)

From c292dcb1218d2d208359b72ad6fe752360a63a85 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Thu, 5 Mar 2026 17:11:38 +0800
Subject: [PATCH 12/27] Use sequential CTMRG to improve convergence in test

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

diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index 5c49141b9..2b0655991 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -109,7 +109,7 @@ end
     Pspace = hubbard_space(Trivial, U1Irrep)
     Vspace = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2, (1, 1 // 2) => 1, (1, -1 // 2) => 1)
     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_env0 = truncerror(; atol = 1.0e-10) & truncrank(8)
     trunc_env = truncerror(; atol = 1.0e-12) & truncrank(16)
     peps = InfinitePEPS(rand, Float64, Pspace, Vspace, Vspace'; unitcell = (Nr, Nc))
     # make state bipartite
@@ -123,8 +123,8 @@ end
         ),
     )
     # usual 2-site simple update, and measure energy
-    dts = [1.0e-2, 1.0e-2]
-    tols = [1.0e-8, 1.0e-8]
+    dts = [1.0e-2, 1.0e-2, 5.0e-3]
+    tols = [5.0e-6, 1.0e-7, 1.0e-8]
     for (n, (dt, tol)) in enumerate(zip(dts, tols))
         trunc = truncerror(; atol = 1.0e-10) & truncrank(n == 1 ? 4 : 2)
         alg = SimpleUpdate(; trunc, bipartite = true)
@@ -132,8 +132,8 @@ end
     end
     normalize!.(peps.A, Inf)
     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)
+    env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env0)
+    env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
     e_site = cost_function(peps, env, ham) / (Nr * Nc)
     @info "2-site simple update energy = $e_site"
     # continue with 3-site simple update; energy should not change much
@@ -145,7 +145,7 @@ end
         peps, wts, = time_evolve(peps, ham, dt, 5000, alg, wts; tol, check_interval = 1000)
     end
     normalize!.(peps.A, Inf)
-    env, = leading_boundary(env, peps; tol = ctmrg_tol, trunc = trunc_env)
+    env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
     e_site2 = cost_function(peps, env, ham) / (Nr * Nc)
     @info "3-site simple update energy = $e_site2"
     @test e_site ≈ e_site2 atol = 5.0e-4

From cdd0680cdbe5d12782daa0954c3d1ed5739375a5 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 6 Mar 2026 17:59:29 +0800
Subject: [PATCH 13/27] Make `trotterize` more universal

---
 src/PEPSKit.jl                                |   1 -
 src/algorithms/time_evolution/apply_gate.jl   |  23 ++
 src/algorithms/time_evolution/simpleupdate.jl |  99 +++++----
 .../time_evolution/simpleupdate3site.jl       |  13 --
 src/algorithms/time_evolution/time_evolve.jl  |  42 ----
 src/algorithms/time_evolution/trotter_gate.jl | 202 +++++++++++++++---
 src/algorithms/time_evolution/trotter_mpo.jl  | 139 ------------
 test/timeevol/cluster_projectors.jl           |   7 +-
 8 files changed, 261 insertions(+), 265 deletions(-)
 delete mode 100644 src/algorithms/time_evolution/trotter_mpo.jl

diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index a39e84e81..a0317a6b1 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -105,7 +105,6 @@ include("algorithms/truncation/fullenv_truncation.jl")
 include("algorithms/truncation/bond_truncation.jl")
 
 include("algorithms/time_evolution/trotter_gate.jl")
-include("algorithms/time_evolution/trotter_mpo.jl")
 include("algorithms/time_evolution/apply_gate.jl")
 include("algorithms/time_evolution/apply_mpo.jl")
 include("algorithms/time_evolution/time_evolve.jl")
diff --git a/src/algorithms/time_evolution/apply_gate.jl b/src/algorithms/time_evolution/apply_gate.jl
index cbfecc7f9..326a9333f 100644
--- a/src/algorithms/time_evolution/apply_gate.jl
+++ b/src/algorithms/time_evolution/apply_gate.jl
@@ -1,3 +1,26 @@
+"""
+Apply 1-site `gate` on the PEPS or PEPO tensor `a`.
+"""
+function _apply_site(
+        a::PEPSTensor, gate::AbstractTensorMap{T, S, 1, 1};
+        gate_ax::Int = 1
+    ) where {T, S}
+    @assert gate_ax == 1
+    return gate * a
+end
+
+function _apply_site(
+        a::PEPOTensor, gate::AbstractTensorMap{T, S, 1, 1};
+        gate_ax::Int = 1
+    ) where {T, S}
+    @assert gate_ax <= 2
+    if gate_ax == 1
+        return @plansor a′[p1 p2; n e s w] := gate[p1; p] * a[p p2; n e s w]
+    else
+        return @plansor a′[p1 p2; n e s w] := a[p1 p; n e s w] * gate[p; p2]
+    end
+end
+
 """
 $(SIGNATURES)
 
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index b3b310186..6181a8496 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -58,10 +58,8 @@ function TimeEvolver(
     dt′ = _get_dt(psi0, dt, alg.imaginary_time)
     # create Trotter gates
     gate = trotterize(H, dt′; force_mpo = alg.force_mpo)
-    if isa(gate, TrotterMPOs)
-        @assert !alg.bipartite "Trotter MPOs are incompatible with bipartite lattice structure."
-    end
     state = SUState(0, t0, psi0, env0)
+    # TODO: check gates for bipartite case
     return TimeEvolver(alg, dt, nstep, gate, state)
 end
 
@@ -77,9 +75,10 @@ the codomain (or domain) physicsl legs of `state`.
 """
 function _su_xbond!(
         state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
-        row::Int, col::Int, trunc::TruncationStrategy; gate_ax::Int = 1
+        row::Int, col::Int, trunc::TruncationStrategy; purified::Bool = true
     )
     Nr, Nc, = size(state)
+    gate_axs = purified ? (1:1) : (1:2)
     cp1 = _next(col, Nc)
     # absorb environment weights
     A, B = state.A[row, col], state.A[row, cp1]
@@ -88,9 +87,13 @@ function _su_xbond!(
     normalize!(A, Inf)
     normalize!(B, Inf)
     # apply gate
-    X, a, b, Y = _qr_bond(A, B; gate_ax)
-    a, s, b, ϵ = _apply_gate(a, b, gate, trunc)
-    A, B = _qr_bond_undo(X, a, b, Y)
+    ϵ, s = 0.0, nothing
+    for gate_ax in gate_axs
+        X, a, b, Y = _qr_bond(A, B; gate_ax)
+        a, s, b, ϵ′ = _apply_gate(a, b, gate, trunc)
+        ϵ = max(ϵ, ϵ′)
+        A, B = _qr_bond_undo(X, a, b, Y)
+    end
     # remove environment weights
     A = absorb_weight(A, env, row, col, (NORTH, SOUTH, WEST); inv = true)
     B = absorb_weight(B, env, row, cp1, (NORTH, SOUTH, EAST); inv = true)
@@ -117,9 +120,10 @@ the codomain (or domain) physicsl legs of `state`.
 """
 function _su_ybond!(
         state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
-        row::Int, col::Int, trunc::TruncationStrategy; gate_ax::Int = 1
+        row::Int, col::Int, trunc::TruncationStrategy; purified::Bool = true
     )
     Nr, Nc, = size(state)
+    gate_axs = purified ? (1:1) : (1:2)
     rm1 = _prev(row, Nr)
     # absorb environment weights
     A, B = state.A[row, col], state.A[rm1, col]
@@ -128,9 +132,13 @@ function _su_ybond!(
     normalize!(A, Inf)
     normalize!(B, Inf)
     # apply gate
-    X, a, b, Y = _qr_bond(rotr90(A), rotr90(B); gate_ax)
-    a, s, b, ϵ = _apply_gate(a, b, gate, trunc)
-    A, B = rotl90.(_qr_bond_undo(X, a, b, Y))
+    ϵ, s = 0.0, nothing
+    for gate_ax in gate_axs
+        X, a, b, Y = _qr_bond(rotr90(A), rotr90(B); gate_ax)
+        a, s, b, ϵ′ = _apply_gate(a, b, gate, trunc)
+        A, B = rotl90.(_qr_bond_undo(X, a, b, Y))
+        ϵ = max(ϵ, ϵ′)
+    end
     # remove environment weights
     A = absorb_weight(A, env, row, col, (EAST, SOUTH, WEST); inv = true)
     B = absorb_weight(B, env, rm1, col, (NORTH, EAST, WEST); inv = true)
@@ -147,38 +155,55 @@ end
 One iteration of simple update
 """
 function su_iter(
-        state::InfiniteState, gate::TrotterNNGates,
+        state::InfiniteState, gates::TrotterGates,
         alg::SimpleUpdate, env::SUWeight
     )
     Nr, Nc, = size(state)
     state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
-    gate_axs = alg.purified ? (1:1) : (1:2)
-    for r in 1:Nr, c in 1:Nc
-        (alg.bipartite && r > 1) && continue
-        # update x-bonds
-        trunc = truncation_strategy(alg.trunc, 1, r, c)
-        for gate_ax in gate_axs
-            ϵ′ = _su_xbond!(state2, gate[1, r, c], env2, r, c, trunc; gate_ax)
-            ϵ = max(ϵ, ϵ′)
-        end
-        if alg.bipartite
-            rp1, cp1 = _next(r, Nr), _next(c, Nc)
-            state2.A[rp1, cp1] = deepcopy(state2.A[r, c])
-            state2.A[rp1, c] = deepcopy(state2.A[r, cp1])
-            env2.data[1, rp1, cp1] = deepcopy(env2.data[1, r, c])
-        end
-        # update y-bonds
-        trunc = truncation_strategy(alg.trunc, 2, r, c)
-        for gate_ax in gate_axs
-            ϵ′ = _su_ybond!(state2, gate[2, r, c], env2, r, c, trunc; gate_ax)
+    for (sites, gs) in gates.data
+        if length(sites) == 1
+            # 1-site gate
+            site = sites[1]
+            r, c = mod1(site[1], Nr), mod1(site[2], Nc)
+            gate_axs = alg.purified ? (1:1) : (1:2)
+            for gate_ax in gate_axs
+                state2.A[r, c] = _apply_site(state2.A[r, c], gs; gate_ax)
+            end
+        elseif length(sites) == 2 && isa(gs, NNGate)
+            # 2-site gate not decomposed to MPO
+            site1, site2 = sites
+            r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)
+            (alg.bipartite && r > 1) && continue
+            if site1 - site2 == CartesianIndex(0, -1) # x-bonds (leftwards)
+                trunc = truncation_strategy(alg.trunc, 1, r, c)
+                ϵ′ = _su_xbond!(state2, gs, env2, r, c, trunc; purified = alg.purified)
+                ϵ = max(ϵ, ϵ′)
+                if alg.bipartite
+                    rp1, cp1 = _next(r, Nr), _next(c, Nc)
+                    state2.A[rp1, cp1] = deepcopy(state2.A[r, c])
+                    state2.A[rp1, c] = deepcopy(state2.A[r, cp1])
+                    env2.data[1, rp1, cp1] = deepcopy(env2.data[1, r, c])
+                end
+            elseif site1 - site2 == CartesianIndex(1, 0) # y-bonds (downwards)
+                trunc = truncation_strategy(alg.trunc, 2, r, c)
+                ϵ′ = _su_ybond!(state2, gs, env2, r, c, trunc; purified = alg.purified)
+                ϵ = max(ϵ, ϵ′)
+                if alg.bipartite
+                    rm1, cm1 = _prev(r, Nr), _prev(c, Nc)
+                    state2.A[rm1, cm1] = deepcopy(state2.A[r, c])
+                    state2.A[r, cm1] = deepcopy(state2.A[rm1, c])
+                    env2.data[2, rm1, cm1] = deepcopy(env2.data[2, r, c])
+                end
+            else
+                error("Non-standard direction of NN bonds for 2-site Trotter gates.")
+            end
+        else
+            # N-site MPO gate (N ≥ 2)
+            alg.bipartite && error("MPO gates are not compatible with bipartite states.")
+            truncs = _get_cluster_trunc(alg.trunc, sites, size(state)[1:2])
+            ϵ′ = _su_cluster!(state2, gs, env2, sites, truncs; purified = alg.purified)
             ϵ = max(ϵ, ϵ′)
         end
-        if alg.bipartite
-            rm1, cm1 = _prev(r, Nr), _prev(c, Nc)
-            state2.A[rm1, cm1] = deepcopy(state2.A[r, c])
-            state2.A[r, cm1] = deepcopy(state2.A[rm1, c])
-            env2.data[2, rm1, cm1] = deepcopy(env2.data[2, r, c])
-        end
     end
     return state2, env2, ϵ
 end
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index 99f5b3a4f..33c5c8e44 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -218,16 +218,3 @@ function _get_cluster_trunc(
         end
     end
 end
-
-function su_iter(
-        state::InfiniteState, gatempos::TrotterMPOs,
-        alg::SimpleUpdate, env::SUWeight
-    )
-    state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
-    for (sites, gs) in gatempos.data
-        truncs = _get_cluster_trunc(alg.trunc, sites, size(state)[1:2])
-        ϵ′ = _su_cluster!(state2, gs, env2, sites, truncs; alg.purified)
-        ϵ = max(ϵ, ϵ′)
-    end
-    return state2, env2, ϵ
-end
diff --git a/src/algorithms/time_evolution/time_evolve.jl b/src/algorithms/time_evolution/time_evolve.jl
index 61b5dfa09..ee5f236c1 100644
--- a/src/algorithms/time_evolution/time_evolve.jl
+++ b/src/algorithms/time_evolution/time_evolve.jl
@@ -94,45 +94,3 @@ function MPSKit.infinite_temperature_density_matrix(H::LocalOperator)
     end
     return InfinitePEPO(cat(A; dims = 3))
 end
-
-"""
-    _check_hamiltonian_for_trotter(H::LocalOperator)
-
-Assert that operator `H` contains only one-site and two-site terms,
-so that it is suited for Trotter time evolution algorithms.
-Returns the maximum squared distance covered by a two-site term in `H`.
-On the square lattice, the neighbor distances are
-```
-    1st nb.     2nd nb.     3rd nb.
-
-                    o
-                    |
-    o---o       o---o       o---o---o
-
-    dist² = 1   dist² = 2   dist² = 4
-```
-"""
-function _check_hamiltonian_for_trotter(H::LocalOperator)
-    dist = 0
-    for (sites, op) in H.terms
-        @assert numin(op) <= 2 "Hamiltonians containing multi-site (> 2) terms are not currently supported."
-        if numin(op) == 2
-            dist = max(dist, sum(Tuple(sites[1] - sites[2]) .^ 2))
-        end
-    end
-    @assert dist > 0 "Hamiltonians with only one-site terms are not suited for current implementation of Trotter time evolution."
-    @assert dist <= 2 "Hamiltonians with 2-site terms on beyond 2nd-neighbor bonds are not currently supported."
-    return dist
-end
-
-function trotterize(H::LocalOperator, dt::Number; force_mpo::Bool = false)
-    dist = _check_hamiltonian_for_trotter(H)
-    # For nearest neighbor Hamiltonian, convert to 2-site Trotter gates.
-    # For all other cases, convert to multi-site Trotter MPOs.
-    gate = if dist == 1 && !force_mpo
-        trotterize_nn(H, dt)
-    elseif (dist == 1 && force_mpo) || dist == 2
-        trotterize_nnn(H, dt)
-    end
-    return gate
-end
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 8db4eddc5..2552d172f 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -1,40 +1,15 @@
 """
-Collection of nearest neighbor 2-body Trotter gates.
+    struct TrotterGates{T <: Vector}
 
-Before exponentiating, terms in the Hamiltonian are organized as
-```
-    H = ∑ᵢⱼ (Xᵢⱼ + Yᵢⱼ)
-```
-where each `Xᵢⱼ` (or `Yᵢⱼ`) acts on a horizontal (or vertical) bond.
-The Trotter gates are `exp(-dt * Xᵢⱼ)`, `exp(-dt * Yᵢⱼ)`.
+Collection of Trotter evolution MPOs obtained from
+a Hamiltonian containing long-range or multi-site terms
 """
-struct TrotterNNGates{G}
-    gates::G
+struct TrotterGates{T <: Vector}
+    data::T
 end
-Base.getindex(gates::TrotterNNGates, args...) = Base.getindex(gates.gates, args...)
 
 const NNGate{T, S} = AbstractTensorMap{T, S, 2, 2}
 
-function trotterize_nn(H::LocalOperator, dt::Number)
-    Nr, Nc = size(H)
-    T = scalartype(H)
-    gates = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
-        # d = 1: horizontal bond; d = 2: vertical bond
-        site1 = CartesianIndex(r, c)
-        site2 = (d == 1) ? CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
-        s1term = _get_site_term(H, site1)
-        unit1 = TensorKit.id(T, space(s1term, 1))
-        s2term = _get_site_term(H, site2)
-        unit2 = TensorKit.id(T, space(s2term, 1))
-        b_term = _get_bond_term(H, (site1, site2))
-        # When iterating through horizontal and vertical bonds in the unit cell,
-        # each site / NN-bond is counted 4 / 1 times, respectively.
-        term = b_term + (s1term ⊗ unit2 + unit1 ⊗ s2term) / 4
-        return exp(-dt * term)
-    end
-    return TrotterNNGates(gates)
-end
-
 """
     is_equivalent_site(
         site1::CartesianIndex{2}, site2::CartesianIndex{2},
@@ -113,3 +88,170 @@ function _get_bond_term(ham::LocalOperator, bond::NTuple{2, CartesianIndex{2}})
     end
     return term
 end
+
+"""
+Get coordinates of sites in the 3-site triangular cluster
+used in Trotter evolution with next-nearest neighbor gates,
+with southwest corner at `[row, col]`.
+```
+    NORTHWEST   NORTHEAST
+        2---1   3---2
+        |           |
+        3           1
+
+        1           3
+        |           |
+        2---3   1---2
+    SOUTHWEST   SOUTHEAST
+```
+"""
+function _nnn_cluster_sites(dir::Int, row::Int, col::Int)
+    @assert 1 <= dir <= 4
+    return if dir == NORTHWEST
+        map(CartesianIndex, [(row - 1, col + 1), (row - 1, col), (row, col)])
+    elseif dir == NORTHEAST
+        map(CartesianIndex, [(row, col + 1), (row - 1, col + 1), (row - 1, col)])
+    elseif dir == SOUTHEAST
+        map(CartesianIndex, [(row, col), (row, col + 1), (row - 1, col + 1)])
+    else # dir == SOUTHWEST
+        map(CartesianIndex, [(row - 1, col), (row, col), (row, col + 1)])
+    end
+end
+
+"""
+Convert an N-site gate to MPO form by SVD, 
+in which the axes are ordered as
+```
+    site 1      mid sites      site N
+    2               3               3
+    ↓               ↓               ↓
+    g1 ←- 3    1 ←- g ←- 4    1 ←- gN
+    ↓               ↓               ↓
+    1               2               2
+```
+"""
+function gate_to_mpo(
+        gate::AbstractTensorMap{T, S, N, N}, trunc = trunctol(; atol = MPSKit.Defaults.tol)
+    ) where {T <: Number, S <: ElementarySpace, N}
+    Os = MPSKit.decompose_localmpo(MPSKit.add_util_leg(gate), trunc)
+    return map(1:N) do i
+        if i == 1
+            return removeunit(Os[1], 1)
+        elseif i == N
+            return removeunit(Os[N], 4)
+        else
+            return Os[i]
+        end
+    end
+end
+
+"""
+    _check_hamiltonian_for_trotter(H::LocalOperator)
+
+Assert that operator `H` contains only one-site and two-site terms.
+Returns the maximum squared distance covered by a two-site term in `H`.
+On the square lattice, the neighbor distances are
+```
+    1st nb.     2nd nb.     3rd nb.
+
+                    o
+                    |
+    o---o       o---o       o---o---o
+
+    dist² = 1   dist² = 2   dist² = 4
+```
+"""
+function _check_hamiltonian_for_trotter(H::LocalOperator)
+    dist = 0
+    for (sites, op) in H.terms
+        @assert numin(op) <= 2 "Hamiltonians containing multi-site (> 2) terms are not currently supported."
+        if numin(op) == 2
+            dist = max(dist, sum(Tuple(sites[1] - sites[2]) .^ 2))
+        end
+    end
+    @assert dist <= 2 "Hamiltonians with 2-site terms on beyond 2nd-neighbor bonds are not currently supported."
+    return dist
+end
+
+function _trotterize_1site!(gates::Vector, H::LocalOperator, dt::Number; atol::Real)
+    for site in CartesianIndices(size(H))
+        gate = _get_site_term(H, site)
+        (norm(gate) <= atol) && continue
+        push!(gates, [site] => exp(-dt * gate))
+    end
+    return gates
+end
+
+function _trotterize_nn2site!(gates::Vector, H::LocalOperator, dt::Number; atol::Real, force_mpo::Bool = false)
+    Nr, Nc = size(H)
+    T = scalartype(H)
+    for (d, c, r) in Iterators.product(1:2, 1:Nc, 1:Nr)
+        site1 = CartesianIndex(r, c)
+        site2 = (d == 1) ? CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
+        # group with 1-site terms
+        s1term = _get_site_term(H, site1)
+        unit1 = TensorKit.id(T, space(s1term, 1))
+        s2term = _get_site_term(H, site2)
+        unit2 = TensorKit.id(T, space(s2term, 1))
+        gate = _get_bond_term(H, (site1, site2))
+        gate = gate + (s1term ⊗ unit2 + unit1 ⊗ s2term) / 4
+        (norm(gate) <= atol) && continue
+        gate = exp(-dt * gate)
+        force_mpo && (gate = gate_to_mpo(gate))
+        push!(gates, [site1, site2] => gate)
+    end
+    return gates
+end
+
+function _trotterize_nnn2site!(gates::Vector, H::LocalOperator, dt::Number; atol::Real)
+    Nr, Nc = size(H)
+    T = scalartype(H)
+    for (c, r, d) in Iterators.product(1:Nc, 1:Nr, 1:4)
+        sites = _nnn_cluster_sites(d, r, c)
+        gate = _get_bond_term(H, (sites[1], sites[3]))
+        (norm(gate) <= atol) && continue
+        gate = exp(-(dt / 2) * gate) # account for double counting
+        # combine with identity at sites[2]
+        r2, c2 = mod1(sites[2][1], Nr), mod1(sites[2][2], Nc)
+        id_ = TensorKit.id(T, physicalspace(H)[r2, c2])
+        gate = permute(gate ⊗ id_, ((1, 3, 2), (4, 6, 5)))
+        push!(gates, sites => gate_to_mpo(gate))
+    end
+    return gates
+end
+
+"""
+Trotterize the evolution operator `exp(-H * dt)`.
+Currently, `H` can only contain the following terms:
+
+- 1-site terms
+- 2-site nearest neighbor (NN) terms
+- 2-site next-nearest neighbor (NNN) terms
+
+If `force_mpo = true`, 2-site nearest-neighbor gates are also decomposed to MPOs.
+"""
+function trotterize(H::LocalOperator, dt::Number; force_mpo::Bool = false)
+    Nr, Nc = size(H)
+    T = scalartype(H)
+    dist = _check_hamiltonian_for_trotter(H)
+    atol = eps(real(T))^(3 / 4)
+    gates = Vector{Pair{Any, Any}}()
+
+    # TODO: order of gates is fixed for more tight control.
+    # Consider directly iterating over H.terms in the future.
+
+    # 1-site gates are only constructed when H only has 1-site terms
+    dist == 0 && _trotterize_1site!(gates, H, dt; atol)
+    #= 
+    2-site NN gates on bonds [d, r, c], grouped with 1-site terms
+    - d = 1: horizontal bond ((r, c), (r, c+1))
+    - d = 2:   vertical bond ((r, c), (r-1, c))
+    =#
+    dist >= 1 && _trotterize_nn2site!(gates, H, dt; atol, force_mpo)
+    #= 
+    2-site NNN gates converted to 3-site MPOs on triangular clusters [d, r, c]
+    - d = 1 (NORTHWEST), ..., 4 (SOUTHWEST)
+    =#
+    dist >= 2 && _trotterize_nnn2site!(gates, H, dt; atol)
+    return TrotterGates(gates)
+end
diff --git a/src/algorithms/time_evolution/trotter_mpo.jl b/src/algorithms/time_evolution/trotter_mpo.jl
deleted file mode 100644
index 6bb7703e9..000000000
--- a/src/algorithms/time_evolution/trotter_mpo.jl
+++ /dev/null
@@ -1,139 +0,0 @@
-"""
-    struct TrotterMPOs{T <: Vector}
-
-Collection of Trotter evolution MPOs obtained from
-a Hamiltonian containing long-range or multi-site terms
-"""
-struct TrotterMPOs{T <: Vector}
-    data::T
-end
-
-"""
-Convert an N-site gate to MPO form by SVD, 
-in which the axes are ordered as
-```
-    site 1      mid sites      site N
-    2               3               3
-    ↓               ↓               ↓
-    g1 ←- 3    1 ←- g ←- 4    1 ←- gN
-    ↓               ↓               ↓
-    1               2               2
-```
-"""
-function gate_to_mpo(
-        gate::AbstractTensorMap{T, S, N, N}, trunc = trunctol(; atol = MPSKit.Defaults.tol)
-    ) where {T <: Number, S <: ElementarySpace, N}
-    Os = MPSKit.decompose_localmpo(MPSKit.add_util_leg(gate), trunc)
-    return map(1:N) do i
-        if i == 1
-            return removeunit(Os[1], 1)
-        elseif i == N
-            return removeunit(Os[N], 4)
-        else
-            return Os[i]
-        end
-    end
-end
-
-# Specialized functions to trotterize Hamiltonians
-# with long-range/multi-site terms
-
-## Next-nearest neighbor H
-
-"""
-Trotterize a Hamiltonian containing up to 2nd neighbor terms.
-```
-    H = ∑ᵢⱼ(Γᵢⱼ + ⅂ᵢⱼ + ⅃ᵢⱼ + Lᵢⱼ)
-```
-where `Γ`, `⅂`, `⅃`, `L` refer to the following 3-site clusters 
-```
-    NORTHWEST   NORTHEAST
-        2---1   3---2
-        |           |
-        3           1
-
-        1           3
-        |           |
-        2---3   1---2
-    SOUTHWEST   SOUTHEAST
-```
-acting on the the elemental square plaquette
-with southwest corner at `[i, j]`.
-"""
-function trotterize_nnn(H::LocalOperator, dt::Number)
-    Nr, Nc = size(H)
-    # iterate through corner `d` in outermost loop
-    terms = map(Iterators.product(1:Nr, 1:Nc, 1:4)) do (r, c, d)
-        return _get_nnn_mpo(H, dt, d, r, c)
-    end
-    return TrotterMPOs(vec(terms))
-end
-
-"""
-Get coordinates of sites in the 3-site triangular cluster
-used in Trotter evolution with next-nearest neighbor gates,
-with southwest corner at `[row, col]`.
-"""
-function _nnn_cluster_sites(dir::Int, row::Int, col::Int)
-    @assert 1 <= dir <= 4
-    return if dir == NORTHWEST
-        [
-            CartesianIndex(row - 1, col + 1),
-            CartesianIndex(row - 1, col),
-            CartesianIndex(row, col),
-        ]
-    elseif dir == NORTHEAST
-        [
-            CartesianIndex(row, col + 1),
-            CartesianIndex(row - 1, col + 1),
-            CartesianIndex(row - 1, col),
-        ]
-    elseif dir == SOUTHEAST
-        [
-            CartesianIndex(row, col),
-            CartesianIndex(row, col + 1),
-            CartesianIndex(row - 1, col + 1),
-        ]
-    else # dir == SOUTHWEST
-        [
-            CartesianIndex(row - 1, col),
-            CartesianIndex(row, col),
-            CartesianIndex(row, col + 1),
-        ]
-    end
-end
-
-"""
-Construct the evolution MPO acting on the 3-site triangular cluster
-in the square plaquette whose southwest corner is at `[row, col]`.
-`dir` takes values between 1 (`NORTHWEST`) and 4 (`SOUTHWEST`).
-"""
-function _get_nnn_mpo(ham::LocalOperator, dt::Number, dir::Int, row::Int, col::Int)
-    Nr, Nc = size(ham)
-    T = scalartype(ham)
-    @assert 1 <= row <= Nr && 1 <= col <= Nc
-    sites = _nnn_cluster_sites(dir, row, col)
-    ss = map(sites) do site
-        _get_site_term(ham, site)
-    end
-    nb1x = _get_bond_term(ham, (sites[1], sites[2]))
-    nb1y = _get_bond_term(ham, (sites[2], sites[3]))
-    nb2 = _get_bond_term(ham, (sites[1], sites[3]))
-    # identity operator at each site
-    units = map(sites) do site
-        site_ = CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
-        return id(T, physicalspace(ham)[site_])
-    end
-    # When iterating through all triangular clusters in the unit cell,
-    # each site / NN-bond / NNN-bond is counted 12 / 4 / 2 times, respectively.
-    term_site = (
-        ss[1] ⊗ units[2] ⊗ units[3] +
-            units[1] ⊗ ss[2] ⊗ units[3] +
-            units[1] ⊗ units[2] ⊗ ss[3]
-    ) / 12
-    @tensor term_nb1[i' j' k'; i j k] :=
-        (nb1x[i' j'; i j] * units[3][k' k] + units[1][i'; i] * nb1y[j' k'; j k]) / 4
-    @tensor term_nb2[i' j' k'; i j k] := (nb2[i' k'; i k] * units[2][j'; j]) / 2
-    term = term_site + term_nb1 + term_nb2
-    return sites => gate_to_mpo(exp(-dt * term))
-end
diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index 2b0655991..d48ec801c 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -135,8 +135,9 @@ end
     env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env0)
     env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
     e_site = cost_function(peps, env, ham) / (Nr * Nc)
-    @info "2-site simple update energy = $e_site"
-    # continue with 3-site simple update; energy should not change much
+    @info "2-site gate simple update energy = $e_site"
+    @test e_site ≈ -0.75 atol = 0.1
+    # continue with MPO simple update; energy should not change much
     dts = [1.0e-2]
     tols = [1.0e-8]
     trunc = truncerror(; atol = 1.0e-10) & truncrank(2)
@@ -147,6 +148,6 @@ end
     normalize!.(peps.A, Inf)
     env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
     e_site2 = cost_function(peps, env, ham) / (Nr * Nc)
-    @info "3-site simple update energy = $e_site2"
+    @info "2-site MPO simple update energy = $e_site2"
     @test e_site ≈ e_site2 atol = 5.0e-4
 end

From a74298a2f30be5badfec8aa0f2793c1a5f7b4bc2 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 6 Mar 2026 20:23:45 +0800
Subject: [PATCH 14/27] Fix trivial SU gauge fixing

---
 src/algorithms/time_evolution/gaugefix_su.jl  | 13 +++++++++++--
 src/algorithms/time_evolution/simpleupdate.jl |  2 +-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/algorithms/time_evolution/gaugefix_su.jl b/src/algorithms/time_evolution/gaugefix_su.jl
index 47e5dd6cc..20858ceb7 100644
--- a/src/algorithms/time_evolution/gaugefix_su.jl
+++ b/src/algorithms/time_evolution/gaugefix_su.jl
@@ -17,14 +17,23 @@ $(TYPEDFIELDS)
     maxiter::Int = 100
 end
 
+function _trivial_gates(Nrow::Int, Ncol::Int)
+    gates = map(Iterators.product(1:2, 1:Ncol, 1:Nrow)) do (d, c, r)
+        site1 = CartesianIndex(r, c)
+        site2 = (d == 1) ? CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
+        return [site1, site2] => nothing
+    end
+    return TrotterGates(vec(gates))
+end
+
 """
     gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::SUGauge)
 
 Fix the gauge of `psi` using trivial simple update.
 """
 function gauge_fix(psi::InfiniteState, alg::SUGauge)
-    Nr, Nc = size(psi)
-    gates = TrotterNNGates(fill(nothing, (2, Nr, Nc)))
+    Nr, Nc, = size(psi)
+    gates = _trivial_gates(Nr, Nc)
     su_alg = SimpleUpdate(; trunc = FixedSpaceTruncation(), bipartite = _state_bipartite_check(psi))
     wts0 = SUWeight(psi)
     # use default constructor to avoid calculation of exp(-H * 0)
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 6181a8496..be57a0765 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -169,7 +169,7 @@ function su_iter(
             for gate_ax in gate_axs
                 state2.A[r, c] = _apply_site(state2.A[r, c], gs; gate_ax)
             end
-        elseif length(sites) == 2 && isa(gs, NNGate)
+        elseif length(sites) == 2 && (isa(gs, NNGate) || gs === nothing)
             # 2-site gate not decomposed to MPO
             site1, site2 = sites
             r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)

From cc7e1656545e79885b546aeeedf25d87c2ec9071 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 6 Mar 2026 20:45:13 +0800
Subject: [PATCH 15/27] Change Hubbard SU test for MPO gate

---
 test/timeevol/cluster_projectors.jl | 46 +++++++++++------------------
 1 file changed, 18 insertions(+), 28 deletions(-)

diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index d48ec801c..45a66a06c 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -101,7 +101,7 @@ end
     end
 end
 
-@testset "Hubbard model with 2-site and 3-site SU" begin
+@testset "Hubbard model SU (MPO gate)" begin
     Nr, Nc = 2, 2
     ctmrg_tol = 1.0e-9
     Random.seed!(1459)
@@ -119,35 +119,25 @@ end
     wts = SUWeight(peps)
     ham = real(
         hubbard_model(
-            ComplexF64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t = 1.0, U = 8.0, mu = 0.0
+            ComplexF64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t = 1.0, U = 6.0, mu = 3.0
         ),
     )
     # usual 2-site simple update, and measure energy
-    dts = [1.0e-2, 1.0e-2, 5.0e-3]
-    tols = [5.0e-6, 1.0e-7, 1.0e-8]
-    for (n, (dt, tol)) in enumerate(zip(dts, tols))
-        trunc = truncerror(; atol = 1.0e-10) & truncrank(n == 1 ? 4 : 2)
-        alg = SimpleUpdate(; trunc, bipartite = true)
-        peps, wts, = time_evolve(peps, ham, dt, 10000, alg, wts; tol, check_interval = 1000)
-    end
-    normalize!.(peps.A, Inf)
-    env = CTMRGEnv(wts)
-    env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env0)
-    env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
-    e_site = cost_function(peps, env, ham) / (Nr * Nc)
-    @info "2-site gate simple update energy = $e_site"
-    @test e_site ≈ -0.75 atol = 0.1
-    # continue with MPO simple update; energy should not change much
-    dts = [1.0e-2]
-    tols = [1.0e-8]
-    trunc = truncerror(; atol = 1.0e-10) & truncrank(2)
-    alg = SimpleUpdate(; trunc, force_mpo = true)
-    for (n, (dt, tol)) in enumerate(zip(dts, tols))
-        peps, wts, = time_evolve(peps, ham, dt, 5000, alg, wts; tol, check_interval = 1000)
+    e_sites = map((true, false)) do force_mpo
+        dts = [1.0e-2, 1.0e-2, 5.0e-3]
+        tols = [1.0e-6, 1.0e-8, 1.0e-8]
+        for (n, (dt, tol)) in enumerate(zip(dts, tols))
+            trunc = truncerror(; atol = 1.0e-10) & truncrank(n == 1 ? 4 : 2)
+            alg = SimpleUpdate(; trunc, force_mpo)
+            peps, wts, = time_evolve(peps, ham, dt, 10000, alg, wts; tol, check_interval = 1000)
+        end
+        normalize!.(peps.A, Inf)
+        env = CTMRGEnv(wts)
+        env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env0)
+        env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
+        e_site = cost_function(peps, env, ham) / (Nr * Nc)
+        @info "Energy (force_mpo = $(force_mpo)): $e_site"
+        return e_site
     end
-    normalize!.(peps.A, Inf)
-    env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
-    e_site2 = cost_function(peps, env, ham) / (Nr * Nc)
-    @info "2-site MPO simple update energy = $e_site2"
-    @test e_site ≈ e_site2 atol = 5.0e-4
+    @test e_sites[1] ≈ e_sites[2] atol = 1.0e-4
 end

From 18d605ae29d3aa1fe0d1e5e11c6ad0fef3e9e341 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sat, 7 Mar 2026 09:48:53 +0800
Subject: [PATCH 16/27] Option for second order Trotter

---
 src/algorithms/time_evolution/simpleupdate.jl | 15 ++++++++-----
 src/algorithms/time_evolution/trotter_gate.jl | 22 +++++++++++++------
 test/timeevol/tf_ising_finiteT.jl             |  9 +++++---
 3 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index be57a0765..fd5be153e 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -48,16 +48,18 @@ Initialize a TimeEvolver with Hamiltonian `H` and simple update `alg`,
 starting from the initial state `psi0` and SUWeight environment `env0`.
 
 - The initial time is specified by `t0`.
+- Use `symmetrize_gates = true` for second-order Trotter decomposition.
 """
 function TimeEvolver(
         psi0::InfiniteState, H::LocalOperator, dt::Number, nstep::Int,
-        alg::SimpleUpdate, env0::SUWeight; t0::Number = 0.0
+        alg::SimpleUpdate, env0::SUWeight; t0::Number = 0.0,
+        symmetrize_gates::Bool = false
     )
     # sanity checks
     _timeevol_sanity_check(psi0, physicalspace(H), alg)
     dt′ = _get_dt(psi0, dt, alg.imaginary_time)
     # create Trotter gates
-    gate = trotterize(H, dt′; force_mpo = alg.force_mpo)
+    gate = trotterize(H, dt′; symmetrize_gates, force_mpo = alg.force_mpo)
     state = SUState(0, t0, psi0, env0)
     # TODO: check gates for bipartite case
     return TimeEvolver(alg, dt, nstep, gate, state)
@@ -302,8 +304,8 @@ end
 
 """
     time_evolve(
-        psi0::Union{InfinitePEPS, InfinitePEPO}, H::LocalOperator, 
-        dt::Number, nstep::Int, alg::SimpleUpdate, env0::SUWeight;
+        psi0::InfiniteState, H::LocalOperator, dt::Number, nstep::Int,
+        alg::SimpleUpdate, env0::SUWeight; symmetrize_gates::Bool = false,
         tol::Float64 = 0.0, t0::Number = 0.0, check_interval::Int = 500
     ) -> (psi, env, info)
 
@@ -311,6 +313,7 @@ Perform time evolution on the initial state `psi0` and initial environment `env0
 with Hamiltonian `H`, using SimpleUpdate algorithm `alg`, time step `dt` for 
 `nstep` number of steps. 
 
+- Use `symmetrize_gates = true` for second-order Trotter decomposition.
 - Setting `tol > 0` enables convergence check (for imaginary time evolution of InfinitePEPS only).
     For other usages it should not be changed.
 - Use `t0` to specify the initial time of the evolution.
@@ -320,9 +323,9 @@ with Hamiltonian `H`, using SimpleUpdate algorithm `alg`, time step `dt` for
 """
 function MPSKit.time_evolve(
         psi0::InfiniteState, H::LocalOperator, dt::Number, nstep::Int,
-        alg::SimpleUpdate, env0::SUWeight;
+        alg::SimpleUpdate, env0::SUWeight; symmetrize_gates::Bool = false,
         tol::Float64 = 0.0, t0::Number = 0.0, check_interval::Int = 500
     )
-    it = TimeEvolver(psi0, H, dt, nstep, alg, env0; t0)
+    it = TimeEvolver(psi0, H, dt, nstep, alg, env0; t0, symmetrize_gates)
     return time_evolve(it; tol, check_interval)
 end
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 2552d172f..4e774a350 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -228,30 +228,38 @@ Currently, `H` can only contain the following terms:
 - 2-site nearest neighbor (NN) terms
 - 2-site next-nearest neighbor (NNN) terms
 
-If `force_mpo = true`, 2-site nearest-neighbor gates are also decomposed to MPOs.
+## Keyword arguments
+
+- `symmetrize_gates::Bool`: if true, use second-order Trotter decomposition.
+- `force_mpo::Bool`: if true, 2-site nearest-neighbor gates are also decomposed to MPOs.
 """
-function trotterize(H::LocalOperator, dt::Number; force_mpo::Bool = false)
-    Nr, Nc = size(H)
-    T = scalartype(H)
+function trotterize(
+        H::LocalOperator, dt::Number;
+        symmetrize_gates::Bool = false, force_mpo::Bool = false
+    )
     dist = _check_hamiltonian_for_trotter(H)
+    T = scalartype(H)
     atol = eps(real(T))^(3 / 4)
+    dt′ = symmetrize_gates ? (dt / 2) : dt
     gates = Vector{Pair{Any, Any}}()
 
     # TODO: order of gates is fixed for more tight control.
     # Consider directly iterating over H.terms in the future.
 
     # 1-site gates are only constructed when H only has 1-site terms
-    dist == 0 && _trotterize_1site!(gates, H, dt; atol)
+    dist == 0 && _trotterize_1site!(gates, H, dt′; atol)
     #= 
     2-site NN gates on bonds [d, r, c], grouped with 1-site terms
     - d = 1: horizontal bond ((r, c), (r, c+1))
     - d = 2:   vertical bond ((r, c), (r-1, c))
     =#
-    dist >= 1 && _trotterize_nn2site!(gates, H, dt; atol, force_mpo)
+    dist >= 1 && _trotterize_nn2site!(gates, H, dt′; atol, force_mpo)
     #= 
     2-site NNN gates converted to 3-site MPOs on triangular clusters [d, r, c]
     - d = 1 (NORTHWEST), ..., 4 (SOUTHWEST)
     =#
-    dist >= 2 && _trotterize_nnn2site!(gates, H, dt; atol)
+    dist >= 2 && _trotterize_nnn2site!(gates, H, dt′; atol)
+
+    symmetrize_gates && push!(gates, reverse(gates)...)
     return TrotterGates(gates)
 end
diff --git a/test/timeevol/tf_ising_finiteT.jl b/test/timeevol/tf_ising_finiteT.jl
index 087fe7d46..529f7ba0f 100644
--- a/test/timeevol/tf_ising_finiteT.jl
+++ b/test/timeevol/tf_ising_finiteT.jl
@@ -45,9 +45,12 @@ dt, nstep = 1.0e-3, 400
 
 # when g = 2, β = 0.4 and 2β = 0.8 belong to two phases (without and with nonzero σᶻ)
 @testset "Finite-T SU (bipartite = $(bipartite))" for bipartite in (true, false)
+    # use second order Trotter decomposition
+    symmetrize_gates = true
+
     # PEPO approach: results at β, or T = 2.5
     alg = SimpleUpdate(; trunc = trunc_pepo, purified = false, bipartite)
-    pepo, wts, info = time_evolve(pepo0, ham, dt, nstep, alg, wts0)
+    pepo, wts, info = time_evolve(pepo0, ham, dt, nstep, alg, wts0; symmetrize_gates)
 
     ## BP gauge fixing
     bp_alg = BeliefPropagation(; maxiter = 100, tol = 1.0e-9, bipartite)
@@ -61,7 +64,7 @@ dt, nstep = 1.0e-3, 400
     @test isapprox(abs.(result_β), bm_β, rtol = 1.0e-2)
 
     # continue to get results at 2β, or T = 1.25
-    pepo, wts, info = time_evolve(pepo, ham, dt, nstep, alg, wts; t0 = β)
+    pepo, wts, info = time_evolve(pepo, ham, dt, nstep, alg, wts; t0 = β, symmetrize_gates)
     env = converge_env(InfinitePartitionFunction(pepo), 16)
     result_2β = measure_mag(pepo, env)
     @info "tr(σ(x,z)ρ) at T = $(1 / (2β)): $(result_2β)."
@@ -70,7 +73,7 @@ dt, nstep = 1.0e-3, 400
 
     # Purification approach: results at 2β, or T = 1.25
     alg = SimpleUpdate(; trunc = trunc_pepo, purified = true, bipartite)
-    pepo, wts, info = time_evolve(pepo0, ham, dt, 2 * nstep, alg, wts0)
+    pepo, wts, info = time_evolve(pepo0, ham, dt, 2 * nstep, alg, wts0; symmetrize_gates)
     env = converge_env(InfinitePEPS(pepo), 8)
     result_2β′ = measure_mag(pepo, env; purified = true)
     @info "⟨ρ|σ(x,z)|ρ⟩ at T = $(1 / (2β)): $(result_2β′)."

From db823bb3c16057a5f1ee787f6d1a5fcaf6542b68 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sat, 7 Mar 2026 10:10:29 +0800
Subject: [PATCH 17/27] Tweak SU log message and docstrings

---
 src/algorithms/time_evolution/simpleupdate.jl | 21 ++++++++-----------
 src/algorithms/time_evolution/trotter_gate.jl |  7 +++++--
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index fd5be153e..85e9c882a 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -44,8 +44,8 @@ end
         alg::SimpleUpdate, env0::SUWeight; t0::Number = 0.0
     )
 
-Initialize a TimeEvolver with Hamiltonian `H` and simple update `alg`, 
-starting from the initial state `psi0` and SUWeight environment `env0`.
+Initialize a `TimeEvolver` with Hamiltonian `H` and simple update `alg`, 
+starting from the initial state `psi0` and `SUWeight` environment `env0`.
 
 - The initial time is specified by `t0`.
 - Use `symmetrize_gates = true` for second-order Trotter decomposition.
@@ -227,7 +227,7 @@ end
         it::TimeEvolver{<:SimpleUpdate}, psi::InfiniteState, env::SUWeight
     ) -> (psi, env, info)
 
-Given the TimeEvolver iterator `it`, perform one step of time evolution
+Given the `TimeEvolver` iterator `it`, perform one step of time evolution
 on the input state `psi` and its environment `env`.
 """
 function MPSKit.timestep(
@@ -250,8 +250,8 @@ end
         tol::Float64 = 0.0, check_interval::Int = 500
     ) -> (psi, env, info)
 
-Perform time evolution to the end of TimeEvolver iterator `it`,
-or until convergence of SUWeight set by a positive `tol`.
+Perform time evolution to the end of `TimeEvolver` iterator `it`,
+or until convergence of `SUWeight` set by a positive `tol`.
 
 - Setting `tol > 0` enables convergence check (for imaginary time evolution of InfinitePEPS only).
     For other usages it should not be changed.
@@ -263,6 +263,7 @@ function MPSKit.time_evolve(
     )
     time_start = time()
     check_convergence = (tol > 0)
+    @info "--- Time evolution (simple update), dt = $(it.dt) ---"
     if check_convergence
         @assert (it.state.psi isa InfinitePEPS) && it.alg.imaginary_time "Only imaginary time evolution of InfinitePEPS allows convergence checking."
     end
@@ -276,11 +277,7 @@ function MPSKit.time_evolve(
         time1 = time()
         if showinfo
             @info "Space of x-weight at [1, 1] = $(space(env[1, 1, 1], 1))"
-            @info @sprintf(
-                "SU iter %-7d: dt = %s, |Δλ| = %.3e. Time = %.3f s/it",
-                # using `string` since `dt` can be complex
-                iter, string(it.dt), diff, time1 - time0
-            )
+            @info @sprintf("SU iter %-7d: |Δλ| = %.3e. Time = %.3f s/it", iter, diff, time1 - time0)
         end
         if check_convergence
             if (iter == it.nstep) && (diff >= tol)
@@ -292,7 +289,7 @@ function MPSKit.time_evolve(
         end
         if stop
             time_end = time()
-            @info @sprintf("Simple update finished. Total time elapsed: %.2f s", time_end - time_start)
+            @info @sprintf("Time evolution finished in %.2f s", time_end - time_start)
             return psi, env, info
         else
             env0 = env
@@ -310,7 +307,7 @@ end
     ) -> (psi, env, info)
 
 Perform time evolution on the initial state `psi0` and initial environment `env0`
-with Hamiltonian `H`, using SimpleUpdate algorithm `alg`, time step `dt` for 
+with Hamiltonian `H`, using `SimpleUpdate` algorithm `alg`, time step `dt` for 
 `nstep` number of steps. 
 
 - Use `symmetrize_gates = true` for second-order Trotter decomposition.
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 4e774a350..3e2c5d374 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -1,8 +1,11 @@
 """
     struct TrotterGates{T <: Vector}
 
-Collection of Trotter evolution MPOs obtained from
-a Hamiltonian containing long-range or multi-site terms
+Collection of Trotter evolution gates and MPOs obtained from
+a Hamiltonian containing long-range or multi-site terms.
+Each item in `data` is a pair `sites => gate`, where `sites` is a
+vector of `CartesianIndex`s storing the sites on which the
+Trotter `gate` acts.
 """
 struct TrotterGates{T <: Vector}
     data::T

From 94bb4a27309241a198d79a0283b674da14fc1105 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sat, 7 Mar 2026 10:34:17 +0800
Subject: [PATCH 18/27] Tweak Hubbard MPO-SU test

---
 src/operators/models.jl             | 15 +++++++++------
 test/timeevol/cluster_projectors.jl | 17 ++++++++---------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/operators/models.jl b/src/operators/models.jl
index 37eceaa72..3f302bbf2 100644
--- a/src/operators/models.jl
+++ b/src/operators/models.jl
@@ -79,15 +79,18 @@ function MPSKitModels.hubbard_model(
     # TODO: just add this
     @assert n == 0 "Currently no support for imposing a fixed particle number"
     N = MPSKitModels.e_number(T, particle_symmetry, spin_symmetry)
-    pspace = space(N, 1)
-    unit = TensorKit.id(pspace)
-    hopping =
+    spaces = fill(space(N, 1), (lattice.Nrows, lattice.Ncols))
+    hopping = (-t) * (
         MPSKitModels.e⁺e⁻(T, particle_symmetry, spin_symmetry) +
-        MPSKitModels.e⁻e⁺(T, particle_symmetry, spin_symmetry)
+            MPSKitModels.e⁻e⁺(T, particle_symmetry, spin_symmetry)
+    )
     interaction_term = MPSKitModels.nꜛnꜜ(T, particle_symmetry, spin_symmetry)
     site_term = U * interaction_term - mu * N
-    h = (-t) * hopping + (1 / 4) * (site_term ⊗ unit + unit ⊗ site_term)
-    return nearest_neighbour_hamiltonian(fill(pspace, size(lattice)), h)
+    return LocalOperator(
+        spaces,
+        (neighbor => hopping for neighbor in nearest_neighbours(lattice))...,
+        ((idx,) => site_term for idx in vertices(lattice))...,
+    )
 end
 
 function MPSKitModels.bose_hubbard_model(
diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index 45a66a06c..5a1964f98 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -112,16 +112,12 @@ end
     trunc_env0 = truncerror(; atol = 1.0e-10) & truncrank(8)
     trunc_env = truncerror(; atol = 1.0e-12) & truncrank(16)
     peps = InfinitePEPS(rand, Float64, Pspace, Vspace, Vspace'; unitcell = (Nr, Nc))
-    # make state bipartite
+    # make initial state bipartite
     for r in 1:2
         peps.A[_next(r, 2), 2] = copy(peps.A[r, 1])
     end
     wts = SUWeight(peps)
-    ham = real(
-        hubbard_model(
-            ComplexF64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t = 1.0, U = 6.0, mu = 3.0
-        ),
-    )
+    ham = hubbard_model(Float64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t = 1.0, U = 6.0, mu = 3.0)
     # usual 2-site simple update, and measure energy
     e_sites = map((true, false)) do force_mpo
         dts = [1.0e-2, 1.0e-2, 5.0e-3]
@@ -129,12 +125,15 @@ end
         for (n, (dt, tol)) in enumerate(zip(dts, tols))
             trunc = truncerror(; atol = 1.0e-10) & truncrank(n == 1 ? 4 : 2)
             alg = SimpleUpdate(; trunc, force_mpo)
-            peps, wts, = time_evolve(peps, ham, dt, 10000, alg, wts; tol, check_interval = 1000)
+            peps, wts, = time_evolve(
+                peps, ham, dt, 10000, alg, wts;
+                tol, symmetrize_gates = false, check_interval = 1000
+            )
         end
         normalize!.(peps.A, Inf)
         env = CTMRGEnv(wts)
-        env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env0)
-        env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc = trunc_env)
+        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)
         @info "Energy (force_mpo = $(force_mpo)): $e_site"
         return e_site

From b561eeb21d6ebb92dae4ca6b71c67e75f2ea9914 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sat, 7 Mar 2026 16:30:52 +0800
Subject: [PATCH 19/27] Minor cleanup

---
 src/algorithms/time_evolution/simpleupdate.jl | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 85e9c882a..d07911c9e 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -162,12 +162,13 @@ function su_iter(
     )
     Nr, Nc, = size(state)
     state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
+    purified = alg.purified
     for (sites, gs) in gates.data
         if length(sites) == 1
             # 1-site gate
             site = sites[1]
             r, c = mod1(site[1], Nr), mod1(site[2], Nc)
-            gate_axs = alg.purified ? (1:1) : (1:2)
+            gate_axs = purified ? (1:1) : (1:2)
             for gate_ax in gate_axs
                 state2.A[r, c] = _apply_site(state2.A[r, c], gs; gate_ax)
             end
@@ -178,7 +179,7 @@ function su_iter(
             (alg.bipartite && r > 1) && continue
             if site1 - site2 == CartesianIndex(0, -1) # x-bonds (leftwards)
                 trunc = truncation_strategy(alg.trunc, 1, r, c)
-                ϵ′ = _su_xbond!(state2, gs, env2, r, c, trunc; purified = alg.purified)
+                ϵ′ = _su_xbond!(state2, gs, env2, r, c, trunc; purified)
                 ϵ = max(ϵ, ϵ′)
                 if alg.bipartite
                     rp1, cp1 = _next(r, Nr), _next(c, Nc)
@@ -188,7 +189,7 @@ function su_iter(
                 end
             elseif site1 - site2 == CartesianIndex(1, 0) # y-bonds (downwards)
                 trunc = truncation_strategy(alg.trunc, 2, r, c)
-                ϵ′ = _su_ybond!(state2, gs, env2, r, c, trunc; purified = alg.purified)
+                ϵ′ = _su_ybond!(state2, gs, env2, r, c, trunc; purified)
                 ϵ = max(ϵ, ϵ′)
                 if alg.bipartite
                     rm1, cm1 = _prev(r, Nr), _prev(c, Nc)
@@ -203,7 +204,7 @@ function su_iter(
             # N-site MPO gate (N ≥ 2)
             alg.bipartite && error("MPO gates are not compatible with bipartite states.")
             truncs = _get_cluster_trunc(alg.trunc, sites, size(state)[1:2])
-            ϵ′ = _su_cluster!(state2, gs, env2, sites, truncs; purified = alg.purified)
+            ϵ′ = _su_cluster!(state2, gs, env2, sites, truncs; purified)
             ϵ = max(ϵ, ϵ′)
         end
     end

From dfb6a82984e3aa3c19b4319af9b3f6361c83df55 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sat, 7 Mar 2026 16:54:16 +0800
Subject: [PATCH 20/27] Fix logic in Hubbard MPO-SU test

---
 test/timeevol/cluster_projectors.jl | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/test/timeevol/cluster_projectors.jl b/test/timeevol/cluster_projectors.jl
index 5a1964f98..591e11952 100644
--- a/test/timeevol/cluster_projectors.jl
+++ b/test/timeevol/cluster_projectors.jl
@@ -109,31 +109,33 @@ end
     Pspace = hubbard_space(Trivial, U1Irrep)
     Vspace = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 2, (1, 1 // 2) => 1, (1, -1 // 2) => 1)
     Espace = Vect[FermionParity ⊠ U1Irrep]((0, 0) => 8, (1, 1 // 2) => 4, (1, -1 // 2) => 4)
-    trunc_env0 = truncerror(; atol = 1.0e-10) & truncrank(8)
-    trunc_env = truncerror(; atol = 1.0e-12) & truncrank(16)
-    peps = InfinitePEPS(rand, Float64, Pspace, Vspace, Vspace'; unitcell = (Nr, Nc))
+    truncs_env = collect(truncerror(; atol = 1.0e-12) & truncrank(χ) for χ in [8, 16])
+    peps0 = InfinitePEPS(rand, Float64, Pspace, Vspace, Vspace'; unitcell = (Nr, Nc))
     # make initial state bipartite
     for r in 1:2
-        peps.A[_next(r, 2), 2] = copy(peps.A[r, 1])
+        peps0.A[_next(r, 2), 2] = copy(peps0.A[r, 1])
     end
-    wts = SUWeight(peps)
+    wts0 = SUWeight(peps0)
     ham = hubbard_model(Float64, Trivial, U1Irrep, InfiniteSquare(Nr, Nc); t = 1.0, U = 6.0, mu = 3.0)
-    # usual 2-site simple update, and measure energy
+    # applying 2-site gates decomposed to MPO or not,
+    # resulting energy should be almost the same
     e_sites = map((true, false)) do force_mpo
-        dts = [1.0e-2, 1.0e-2, 5.0e-3]
-        tols = [1.0e-6, 1.0e-8, 1.0e-8]
+        dts = [1.0e-2, 1.0e-2]
+        tols = [1.0e-6, 1.0e-8]
+        peps, wts = deepcopy(peps0), deepcopy(wts0)
         for (n, (dt, tol)) in enumerate(zip(dts, tols))
             trunc = truncerror(; atol = 1.0e-10) & truncrank(n == 1 ? 4 : 2)
             alg = SimpleUpdate(; trunc, force_mpo)
             peps, wts, = time_evolve(
                 peps, ham, dt, 10000, alg, wts;
-                tol, symmetrize_gates = false, check_interval = 1000
+                tol, symmetrize_gates = true, check_interval = 1000
             )
         end
         normalize!.(peps.A, Inf)
         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)
+        for trunc in truncs_env
+            env, = leading_boundary(env, peps; alg = :sequential, tol = ctmrg_tol, trunc)
+        end
         e_site = cost_function(peps, env, ham) / (Nr * Nc)
         @info "Energy (force_mpo = $(force_mpo)): $e_site"
         return e_site

From a5f9613012b3de8dd415580cc313c6d7c5f1ae10 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Sun, 8 Mar 2026 11:02:13 +0800
Subject: [PATCH 21/27] Merge _su_xbond, _su_ybond

---
 src/algorithms/time_evolution/apply_gate.jl   |  21 +--
 src/algorithms/time_evolution/simpleupdate.jl | 156 +++++++-----------
 2 files changed, 68 insertions(+), 109 deletions(-)

diff --git a/src/algorithms/time_evolution/apply_gate.jl b/src/algorithms/time_evolution/apply_gate.jl
index 326a9333f..89b622a2e 100644
--- a/src/algorithms/time_evolution/apply_gate.jl
+++ b/src/algorithms/time_evolution/apply_gate.jl
@@ -1,24 +1,21 @@
 """
 Apply 1-site `gate` on the PEPS or PEPO tensor `a`.
 """
-function _apply_site(
-        a::PEPSTensor, gate::AbstractTensorMap{T, S, 1, 1};
-        gate_ax::Int = 1
+function _apply_sitegate(
+        a::PEPSTensor, gate::AbstractTensorMap{T, S, 1, 1}; purified::Bool = true
     ) where {T, S}
-    @assert gate_ax == 1
+    @assert purified
     return gate * a
 end
 
-function _apply_site(
-        a::PEPOTensor, gate::AbstractTensorMap{T, S, 1, 1};
-        gate_ax::Int = 1
+function _apply_sitegate(
+        a::PEPOTensor, gate::AbstractTensorMap{T, S, 1, 1}; purified::Bool = true
     ) where {T, S}
-    @assert gate_ax <= 2
-    if gate_ax == 1
-        return @plansor a′[p1 p2; n e s w] := gate[p1; p] * a[p p2; n e s w]
-    else
-        return @plansor a′[p1 p2; n e s w] := a[p1 p; n e s w] * gate[p; p2]
+    @plansor a′[p1 p2; n e s w] := gate[p1; p] * a[p p2; n e s w]
+    if !purified
+        @plansor a′[p1 p2; n e s w] := a′[p1 p; n e s w] * gate[p; p2]
     end
+    return a′
 end
 
 """
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index d07911c9e..29f7e6ada 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -40,8 +40,9 @@ end
 
 """
     TimeEvolver(
-        psi0::InfiniteState, H::LocalOperator, dt::Number, nstep::Int, 
-        alg::SimpleUpdate, env0::SUWeight; t0::Number = 0.0
+        psi0::Union{InfinitePEPS, InfinitePEPO}, H::LocalOperator, dt::Number,
+        nstep::Int, alg::SimpleUpdate, env0::SUWeight; t0::Number = 0.0,
+        symmetrize_gates::Bool = false
     )
 
 Initialize a `TimeEvolver` with Hamiltonian `H` and simple update `alg`, 
@@ -66,90 +67,54 @@ function TimeEvolver(
 end
 
 """
-Simple update of the x-bond between `[r,c]` and `[r,c+1]`.
-```
-        |           |
-    -- T[r,c] -- T[r,c+1] --
-        |           |
-```
-When `gate_ax = 1` (or `2`), the gate will be applied to 
-the codomain (or domain) physicsl legs of `state`.
+Optimized simple update of nearest neighbor bonds utilizing
+reduced bond tensors without decomposing the gate into a 2-site MPO.
+
+When `purified = true`, `gate` acts on the codomain physical legs of `state`.
+Otherwise, `gate` acts on both the codomain and the domain physical legs of `state`.
 """
-function _su_xbond!(
+function _su_nnbond!(
         state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
-        row::Int, col::Int, trunc::TruncationStrategy; purified::Bool = true
+        dir::Int, row::Int, col::Int, trunc::TruncationStrategy;
+        purified::Bool = true
     )
     Nr, Nc, = size(state)
-    gate_axs = purified ? (1:1) : (1:2)
-    cp1 = _next(col, Nc)
+    @assert dir == 1 || dir == 2
+    # position of bond tensors
+    siteA = CartesianIndex(row, col)
+    siteB = CartesianIndex((dir == 1) ? (row, _next(col, Nc)) : (_prev(row, Nr), col))
+    A, B = state.A[siteA], state.A[siteB]
     # absorb environment weights
-    A, B = state.A[row, col], state.A[row, cp1]
-    A = absorb_weight(A, env, row, col, (NORTH, SOUTH, WEST); inv = false)
-    B = absorb_weight(B, env, row, cp1, (NORTH, SOUTH, EAST); inv = false)
+    openaxsA = (dir == 1) ? (NORTH, SOUTH, WEST) : (EAST, SOUTH, WEST)
+    openaxsB = (dir == 1) ? (NORTH, SOUTH, EAST) : (NORTH, EAST, WEST)
+    A = absorb_weight(A, env, siteA[1], siteA[2], openaxsA; inv = false)
+    B = absorb_weight(B, env, siteB[1], siteB[2], openaxsB; inv = false)
     normalize!(A, Inf)
     normalize!(B, Inf)
     # apply gate
     ϵ, s = 0.0, nothing
+    if dir == 2
+        A, B = rotr90(A), rotr90(B)
+    end
+    gate_axs = purified ? (1:1) : (1:2)
     for gate_ax in gate_axs
         X, a, b, Y = _qr_bond(A, B; gate_ax)
         a, s, b, ϵ′ = _apply_gate(a, b, gate, trunc)
         ϵ = max(ϵ, ϵ′)
         A, B = _qr_bond_undo(X, a, b, Y)
     end
-    # remove environment weights
-    A = absorb_weight(A, env, row, col, (NORTH, SOUTH, WEST); inv = true)
-    B = absorb_weight(B, env, row, cp1, (NORTH, SOUTH, EAST); inv = true)
-    normalize!(A, Inf)
-    normalize!(B, Inf)
-    normalize!(s, Inf)
-    # update tensor dict and weight on current bond
-    state.A[row, col], state.A[row, cp1] = A, B
-    env.data[1, row, col] = s
-    return ϵ
-end
-
-"""
-Simple update of the y-bond between `[r,c]` and `[r-1,c]`.
-```
-        |
-    --T[r-1,c] --
-        |
-    -- T[r,c] ---
-        |
-```
-When `gate_ax = 1` (or `2`), the gate will be applied to 
-the codomain (or domain) physicsl legs of `state`.
-"""
-function _su_ybond!(
-        state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
-        row::Int, col::Int, trunc::TruncationStrategy; purified::Bool = true
-    )
-    Nr, Nc, = size(state)
-    gate_axs = purified ? (1:1) : (1:2)
-    rm1 = _prev(row, Nr)
-    # absorb environment weights
-    A, B = state.A[row, col], state.A[rm1, col]
-    A = absorb_weight(A, env, row, col, (EAST, SOUTH, WEST); inv = false)
-    B = absorb_weight(B, env, rm1, col, (NORTH, EAST, WEST); inv = false)
-    normalize!(A, Inf)
-    normalize!(B, Inf)
-    # apply gate
-    ϵ, s = 0.0, nothing
-    for gate_ax in gate_axs
-        X, a, b, Y = _qr_bond(rotr90(A), rotr90(B); gate_ax)
-        a, s, b, ϵ′ = _apply_gate(a, b, gate, trunc)
-        A, B = rotl90.(_qr_bond_undo(X, a, b, Y))
-        ϵ = max(ϵ, ϵ′)
+    if dir == 2
+        A, B = rotl90(A), rotl90(B)
     end
     # remove environment weights
-    A = absorb_weight(A, env, row, col, (EAST, SOUTH, WEST); inv = true)
-    B = absorb_weight(B, env, rm1, col, (NORTH, EAST, WEST); inv = true)
-    # update tensor dict and weight on current bond
+    A = absorb_weight(A, env, siteA[1], siteA[2], openaxsA; inv = true)
+    B = absorb_weight(B, env, siteB[1], siteB[2], openaxsB; inv = true)
     normalize!(A, Inf)
     normalize!(B, Inf)
     normalize!(s, Inf)
-    state.A[row, col], state.A[rm1, col] = A, B
-    env.data[2, row, col] = s
+    # update tensor dict and weight on current bond
+    state.A[siteA], state.A[siteB] = A, B
+    env.data[dir, row, col] = s
     return ϵ
 end
 
@@ -166,43 +131,40 @@ function su_iter(
     for (sites, gs) in gates.data
         if length(sites) == 1
             # 1-site gate
+            # TODO: special treatment for bipartite state
             site = sites[1]
             r, c = mod1(site[1], Nr), mod1(site[2], Nc)
-            gate_axs = purified ? (1:1) : (1:2)
-            for gate_ax in gate_axs
-                state2.A[r, c] = _apply_site(state2.A[r, c], gs; gate_ax)
-            end
+            state2.A[r, c] = _apply_sitegate(state2.A[r, c], gs; purified)
         elseif length(sites) == 2 && (isa(gs, NNGate) || gs === nothing)
             # 2-site gate not decomposed to MPO
             site1, site2 = sites
             r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)
             (alg.bipartite && r > 1) && continue
-            if site1 - site2 == CartesianIndex(0, -1) # x-bonds (leftwards)
-                trunc = truncation_strategy(alg.trunc, 1, r, c)
-                ϵ′ = _su_xbond!(state2, gs, env2, r, c, trunc; purified)
-                ϵ = max(ϵ, ϵ′)
-                if alg.bipartite
-                    rp1, cp1 = _next(r, Nr), _next(c, Nc)
-                    state2.A[rp1, cp1] = deepcopy(state2.A[r, c])
-                    state2.A[rp1, c] = deepcopy(state2.A[r, cp1])
-                    env2.data[1, rp1, cp1] = deepcopy(env2.data[1, r, c])
-                end
-            elseif site1 - site2 == CartesianIndex(1, 0) # y-bonds (downwards)
-                trunc = truncation_strategy(alg.trunc, 2, r, c)
-                ϵ′ = _su_ybond!(state2, gs, env2, r, c, trunc; purified)
-                ϵ = max(ϵ, ϵ′)
-                if alg.bipartite
-                    rm1, cm1 = _prev(r, Nr), _prev(c, Nc)
-                    state2.A[rm1, cm1] = deepcopy(state2.A[r, c])
-                    state2.A[r, cm1] = deepcopy(state2.A[rm1, c])
-                    env2.data[2, rm1, cm1] = deepcopy(env2.data[2, r, c])
-                end
+            d = if site1 - site2 == CartesianIndex(0, -1)
+                1 # x-bonds (leftwards)
+            elseif site1 - site2 == CartesianIndex(1, 0)
+                2 # y-bonds (downwards)
             else
                 error("Non-standard direction of NN bonds for 2-site Trotter gates.")
             end
+            trunc = truncation_strategy(alg.trunc, d, r, c)
+            ϵ′ = _su_nnbond!(state2, gs, env2, d, r, c, trunc; purified)
+            ϵ = max(ϵ, ϵ′)
+            (!alg.bipartite) && continue
+            if d == 1
+                rp1, cp1 = _next(r, Nr), _next(c, Nc)
+                state2.A[rp1, cp1] = deepcopy(state2.A[r, c])
+                state2.A[rp1, c] = deepcopy(state2.A[r, cp1])
+                env2.data[1, rp1, cp1] = deepcopy(env2.data[1, r, c])
+            else
+                rm1, cm1 = _prev(r, Nr), _prev(c, Nc)
+                state2.A[rm1, cm1] = deepcopy(state2.A[r, c])
+                state2.A[r, cm1] = deepcopy(state2.A[rm1, c])
+                env2.data[2, rm1, cm1] = deepcopy(env2.data[2, r, c])
+            end
         else
             # N-site MPO gate (N ≥ 2)
-            alg.bipartite && error("MPO gates are not compatible with bipartite states.")
+            alg.bipartite && error("Multi-site MPO gates are not compatible with bipartite states.")
             truncs = _get_cluster_trunc(alg.trunc, sites, size(state)[1:2])
             ϵ′ = _su_cluster!(state2, gs, env2, sites, truncs; purified)
             ϵ = max(ϵ, ϵ′)
@@ -302,17 +264,17 @@ end
 
 """
     time_evolve(
-        psi0::InfiniteState, H::LocalOperator, dt::Number, nstep::Int,
+        psi0::Union{InfinitePEPS, InfinitePEPO}, H::LocalOperator, dt::Number, nstep::Int,
         alg::SimpleUpdate, env0::SUWeight; symmetrize_gates::Bool = false,
         tol::Float64 = 0.0, t0::Number = 0.0, check_interval::Int = 500
     ) -> (psi, env, info)
 
-Perform time evolution on the initial state `psi0` and initial environment `env0`
-with Hamiltonian `H`, using `SimpleUpdate` algorithm `alg`, time step `dt` for 
-`nstep` number of steps. 
+Perform time evolution on the initial iPEPS or iPEPO `psi0` and
+initial environment `env0` with Hamiltonian `H`, using `SimpleUpdate`
+algorithm `alg`, time step `dt` for `nstep` number of steps. 
 
-- Use `symmetrize_gates = true` for second-order Trotter decomposition.
-- Setting `tol > 0` enables convergence check (for imaginary time evolution of InfinitePEPS only).
+- Set `symmetrize_gates = true` for second-order Trotter decomposition.
+- Set `tol > 0` to enable convergence check (for imaginary time evolution of iPEPS only).
     For other usages it should not be changed.
 - Use `t0` to specify the initial time of the evolution.
 - `check_interval` sets the interval to output information. Output during the evolution can be turned off by setting `check_interval <= 0`.

From 39717e1ebda956f477df5f814ef40db6e8b5c924 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Mon, 9 Mar 2026 20:36:25 +0800
Subject: [PATCH 22/27] Enable bipartite mode for 2-site MPO gates

---
 src/algorithms/time_evolution/simpleupdate.jl | 85 ++++++++++---------
 .../time_evolution/simpleupdate3site.jl       | 82 +++++++++++-------
 2 files changed, 94 insertions(+), 73 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 29f7e6ada..c453d8e05 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -73,48 +73,55 @@ reduced bond tensors without decomposing the gate into a 2-site MPO.
 When `purified = true`, `gate` acts on the codomain physical legs of `state`.
 Otherwise, `gate` acts on both the codomain and the domain physical legs of `state`.
 """
-function _su_nnbond!(
+function _su_iter!(
         state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
-        dir::Int, row::Int, col::Int, trunc::TruncationStrategy;
+        sites::Vector{CartesianIndex{2}}, truncs::Vector{E};
         purified::Bool = true
-    )
-    Nr, Nc, = size(state)
-    @assert dir == 1 || dir == 2
-    # position of bond tensors
-    siteA = CartesianIndex(row, col)
-    siteB = CartesianIndex((dir == 1) ? (row, _next(col, Nc)) : (_prev(row, Nr), col))
-    A, B = state.A[siteA], state.A[siteB]
-    # absorb environment weights
-    openaxsA = (dir == 1) ? (NORTH, SOUTH, WEST) : (EAST, SOUTH, WEST)
-    openaxsB = (dir == 1) ? (NORTH, SOUTH, EAST) : (NORTH, EAST, WEST)
-    A = absorb_weight(A, env, siteA[1], siteA[2], openaxsA; inv = false)
-    B = absorb_weight(B, env, siteB[1], siteB[2], openaxsB; inv = false)
-    normalize!(A, Inf)
-    normalize!(B, Inf)
+    ) where {E <: TruncationStrategy}
+    Nr, Nc = size(state)
+    @assert length(sites) == 2 && length(truncs) == 1
+    Ms, open_vaxs, = _get_cluster(state, sites, env; permute = false)
+    normalize!.(Ms, Inf)
+    # rotate
+    bond, rev = _nn_bondrev(sites..., (Nr, Nc))
+    A, B = if bond[1] == 1 # x-bond
+        rev ? map(rot180, Ms) : Ms
+    else # y-bond
+        rev ? map(rotl90, Ms) : map(rotr90, Ms)
+    end
     # apply gate
     ϵ, s = 0.0, nothing
-    if dir == 2
-        A, B = rotr90(A), rotr90(B)
-    end
     gate_axs = purified ? (1:1) : (1:2)
     for gate_ax in gate_axs
         X, a, b, Y = _qr_bond(A, B; gate_ax)
-        a, s, b, ϵ′ = _apply_gate(a, b, gate, trunc)
+        a, s, b, ϵ′ = _apply_gate(a, b, gate, truncs[1])
         ϵ = max(ϵ, ϵ′)
         A, B = _qr_bond_undo(X, a, b, Y)
     end
-    if dir == 2
-        A, B = rotl90(A), rotl90(B)
+    # rotate back
+    if bond[1] == 1 # x-bond
+        if rev
+            A, B = rot180(A), rot180(B)
+        end
+    else # y-bond
+        if rev
+            A, B = rotr90(A), rotr90(B)
+        else
+            A, B = rotl90(A), rotl90(B)
+        end
     end
     # remove environment weights
-    A = absorb_weight(A, env, siteA[1], siteA[2], openaxsA; inv = true)
-    B = absorb_weight(B, env, siteB[1], siteB[2], openaxsB; inv = true)
+    siteA, siteB = map(sites) do site
+        return CartesianIndex(mod1(site[1], Nr), mod1(site[2], Nc))
+    end
+    A = absorb_weight(A, env, siteA[1], siteA[2], open_vaxs[1]; inv = true)
+    B = absorb_weight(B, env, siteB[1], siteB[2], open_vaxs[2]; inv = true)
+    # update tensor dict and weight on current bond
     normalize!(A, Inf)
     normalize!(B, Inf)
     normalize!(s, Inf)
-    # update tensor dict and weight on current bond
     state.A[siteA], state.A[siteB] = A, B
-    env.data[dir, row, col] = s
+    env.data[bond...] = s
     return ϵ
 end
 
@@ -128,27 +135,21 @@ function su_iter(
     Nr, Nc, = size(state)
     state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
     purified = alg.purified
-    for (sites, gs) in gates.data
+    for (sites, gate) in gates.data
         if length(sites) == 1
             # 1-site gate
             # TODO: special treatment for bipartite state
             site = sites[1]
             r, c = mod1(site[1], Nr), mod1(site[2], Nc)
-            state2.A[r, c] = _apply_sitegate(state2.A[r, c], gs; purified)
-        elseif length(sites) == 2 && (isa(gs, NNGate) || gs === nothing)
-            # 2-site gate not decomposed to MPO
-            site1, site2 = sites
-            r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)
-            (alg.bipartite && r > 1) && continue
-            d = if site1 - site2 == CartesianIndex(0, -1)
-                1 # x-bonds (leftwards)
-            elseif site1 - site2 == CartesianIndex(1, 0)
-                2 # y-bonds (downwards)
-            else
-                error("Non-standard direction of NN bonds for 2-site Trotter gates.")
+            state2.A[r, c] = _apply_sitegate(state2.A[r, c], gate; purified)
+        elseif length(sites) == 2
+            (d, r, c), = _nn_bondrev(sites..., (Nr, Nc))
+            if alg.bipartite
+                length(sites) > 2 && error("Multi-site MPO gates are not compatible with bipartite states.")
+                r > 1 && continue
             end
-            trunc = truncation_strategy(alg.trunc, d, r, c)
-            ϵ′ = _su_nnbond!(state2, gs, env2, d, r, c, trunc; purified)
+            truncs = _get_cluster_trunc(alg.trunc, sites, size(state)[1:2])
+            ϵ′ = _su_iter!(state2, gate, env2, sites, truncs; purified)
             ϵ = max(ϵ, ϵ′)
             (!alg.bipartite) && continue
             if d == 1
@@ -166,7 +167,7 @@ function su_iter(
             # N-site MPO gate (N ≥ 2)
             alg.bipartite && error("Multi-site MPO gates are not compatible with bipartite states.")
             truncs = _get_cluster_trunc(alg.trunc, sites, size(state)[1:2])
-            ϵ′ = _su_cluster!(state2, gs, env2, sites, truncs; purified)
+            ϵ′ = _su_iter!(state2, gate, env2, sites, truncs; purified)
             ϵ = max(ϵ, ϵ′)
         end
     end
diff --git a/src/algorithms/time_evolution/simpleupdate3site.jl b/src/algorithms/time_evolution/simpleupdate3site.jl
index 33c5c8e44..c7792cca4 100644
--- a/src/algorithms/time_evolution/simpleupdate3site.jl
+++ b/src/algorithms/time_evolution/simpleupdate3site.jl
@@ -39,6 +39,30 @@ function _nn_vec_direction(nn_vec::CartesianIndex{2})
     end
 end
 
+"""
+Given `site1`, `site2` connected by a nearest neighbor bond,
+return the bond index and whether it is reversed from the
+standard orientation (`site1` on the west/south of `site2`).
+"""
+function _nn_bondrev(site1::CartesianIndex{2}, site2::CartesianIndex{2}, (Nrow, Ncol)::NTuple{2, Int})
+    diff = site1 - site2
+    if diff == CartesianIndex(0, -1)
+        r, c = mod1(site1[1], Nrow), mod1(site1[2], Ncol)
+        return (1, r, c), false
+    elseif diff == CartesianIndex(0, 1)
+        r, c = mod1(site2[1], Nrow), mod1(site2[2], Ncol)
+        return (1, r, c), true
+    elseif diff == CartesianIndex(1, 0)
+        r, c = mod1(site1[1], Nrow), mod1(site1[2], Ncol)
+        return (2, r, c), false
+    elseif diff == CartesianIndex(-1, 0)
+        r, c = mod1(site2[1], Nrow), mod1(site2[2], Ncol)
+        return (2, r, c), true
+    else
+        error("`site1` and `site2` are not nearest neighbors.")
+    end
+end
+
 """
 Find the permutation to permute `out_ax`, `in_ax` legs to
 the first and the last position of a tensor with `Nax` legs,
@@ -53,15 +77,12 @@ function _get_mpo_perm(out_ax::Int, in_ax::Int, Nax::Int)
 end
 
 """
-Obtain the cluster `Ms` along the (open and non-self-intersecting) path
-given by `sites` in `state`.
+Obtain the cluster `Ms` along the (open) path `sites` in `state`. 
 
-Also returns:
-- `open_vaxs`: Open virtual axes (1 to 4) of each cluster tensor before permutation.
-- `bond_revs`: Position of the cluster bonds in the state, and whether the cluster path follows the "canonical" orientation of the bond (leftwards or downwards).
-- `invperms`: Permutations to restore the axes order of each cluster tensor.
+When the `SUWeight` environment `env` is provided,
+it will be absorbed into tensors of `Ms`.
 
-In the cluster, each tensor use the MPS axes order
+When `permute = true`, permute tensors in `Ms` to MPS axis order
 ```
     PEPS:           PEPO:
            3             3  4
@@ -72,10 +93,20 @@ In the cluster, each tensor use the MPS axes order
     M[o 2 3 4; i]  M[o 2 3 4 5; i]
 ```
 where `o` (`i`) connects to the previous (next) tensor.
+Otherwise, axes order of each tensor in `Ms` are preserved.
+
+## Returns
+
+- `Ms`: Tensors in the cluster.
+- `open_vaxs`: Open virtual axes (1 to 4) of each cluster tensor before permutation.
+- `invperms`: Permutations to restore the axes order of each cluster tensor.
 """
+function _get_cluster(state, sites; permute::Bool = true)
+    return _get_cluster(state, sites, nothing; permute)
+end
 function _get_cluster(
         state::InfiniteState, sites::Vector{CartesianIndex{2}},
-        env::Union{SUWeight, Nothing}
+        env::Union{SUWeight, Nothing}; permute::Bool = true
     )
     Nr, Nc = size(state)
     # number of sites
@@ -100,22 +131,6 @@ function _get_cluster(
             filter(x -> x != out_axs[i - 1] && x != in_axs[i], all_vaxs)
         end
     end
-    bond_revs = map(zip(sites, Iterators.drop(sites, 1))) do (site1, site2)
-        diff = site1 - site2
-        if diff == CartesianIndex(0, -1)
-            r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)
-            return (1, r, c), false
-        elseif diff == CartesianIndex(0, 1)
-            r, c = mod1(site2[1], Nr), mod1(site2[2], Nc)
-            return (1, r, c), true
-        elseif diff == CartesianIndex(1, 0)
-            r, c = mod1(site1[1], Nr), mod1(site1[2], Nc)
-            return (2, r, c), false
-        else # diff == CartesianIndex(-1, 0)
-            r, c = mod1(site2[1], Nr), mod1(site2[2], Nc)
-            return (2, r, c), true
-        end
-    end
     perms = map(1:Ns) do i
         out_ax, in_ax = if i == 1
             # use direction opposite to `in` as `out`
@@ -139,19 +154,21 @@ function _get_cluster(
         else
             absorb_weight(state.A[s], env, s[1], s[2], vaxs)
         end
-        return permute(M, perm)
+        return permute ? TensorKit.permute(M, perm) : M
     end
-    return Ms, open_vaxs, bond_revs, invperms
+    return Ms, open_vaxs, invperms
 end
 
-function _su_cluster!(
-        state::InfiniteState, gs::Vector{T}, env::SUWeight,
+"""
+Simple update with an N-site MPO `gate` (N ≥ 2).
+"""
+function _su_iter!(
+        state::InfiniteState, gate::Vector{T}, env::SUWeight,
         sites::Vector{CartesianIndex{2}}, truncs::Vector{E};
         purified::Bool = true
     ) where {T <: AbstractTensorMap, E <: TruncationStrategy}
     Nr, Nc = size(state)
-    # southwest 3-site cluster and arrow direction within it
-    Ms, open_vaxs, bond_revs, invperms = _get_cluster(state, sites, env)
+    Ms, open_vaxs, invperms = _get_cluster(state, sites, env)
     flips = [isdual(space(M, 1)) for M in Ms[2:end]]
     Vphys = [codomain(M, 2) for M in Ms]
     normalize!.(Ms, Inf)
@@ -161,7 +178,7 @@ function _su_cluster!(
     gate_axs = purified ? (1:1) : (1:2)
     wts, ϵs = nothing, nothing
     for gate_ax in gate_axs
-        _apply_gatempo!(Ms, gs; gate_ax)
+        _apply_gatempo!(Ms, gate; gate_ax)
         if isa(state, InfinitePEPO)
             Ms = [first(_fuse_physicalspaces(M)) for M in Ms]
         end
@@ -173,6 +190,9 @@ function _su_cluster!(
     # restore virtual arrows in `Ms`
     _flip_virtuals!(Ms, flips)
     # update env weights
+    bond_revs = map(zip(sites, Iterators.drop(sites, 1))) do (site1, site2)
+        _nn_bondrev(site1, site2, (Nr, Nc))
+    end
     for (wt, (bond, rev), flip) in zip(wts, bond_revs, flips)
         wt_new = flip ? _fliptwist_s(wt) : wt
         wt_new = rev ? transpose(wt_new) : wt_new

From 9a5f142f564ea80dfec86267e56767f89a0faad9 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Tue, 10 Mar 2026 18:39:12 +0800
Subject: [PATCH 23/27] Try fixing LAPACKException in tf_ising_finiteT

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

diff --git a/test/timeevol/tf_ising_finiteT.jl b/test/timeevol/tf_ising_finiteT.jl
index 529f7ba0f..d3d668ad5 100644
--- a/test/timeevol/tf_ising_finiteT.jl
+++ b/test/timeevol/tf_ising_finiteT.jl
@@ -11,7 +11,7 @@ bm_2β = [0.5297, 0.8265]
 
 function converge_env(state, χ::Int)
     trunc1 = truncrank(4) & truncerror(; atol = 1.0e-12)
-    env0 = CTMRGEnv(ones, Float64, state, ℂ^1)
+    env0 = CTMRGEnv(rand, Float64, state, ℂ^2)
     env, = leading_boundary(env0, state; alg = :sequential, trunc = trunc1, tol = 1.0e-10)
     trunc2 = truncrank(χ) & truncerror(; atol = 1.0e-12)
     env, = leading_boundary(env, state; alg = :sequential, trunc = trunc2, tol = 1.0e-10)

From b23d5aa1130a92e513be21f36017a66c1dd5d9b6 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 11 Mar 2026 22:08:33 +0800
Subject: [PATCH 24/27] Add `lattice` field to TrotterGates

---
 src/algorithms/time_evolution/simpleupdate.jl |  4 +--
 src/algorithms/time_evolution/trotter_gate.jl | 28 +++++++++++++++----
 2 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index c453d8e05..8d105ee83 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -12,7 +12,7 @@ $(TYPEDFIELDS)
     trunc::TruncationStrategy
     "When true (or false), the Trotter gate is `exp(-H dt)` (or `exp(-iH dt)`)"
     imaginary_time::Bool = true
-    "When true, force the usage of MPO simple update for nearest neighbor Hamiltonians."
+    "When true, force decomposition of nearest neighbor gates to MPOs."
     force_mpo::Bool = false
     "When true, assume bipartite unit cell structure"
     bipartite::Bool = false
@@ -135,7 +135,7 @@ function su_iter(
     Nr, Nc, = size(state)
     state2, env2, ϵ = deepcopy(state), deepcopy(env), 0.0
     purified = alg.purified
-    for (sites, gate) in gates.data
+    for (sites, gate) in gates.terms
         if length(sites) == 1
             # 1-site gate
             # TODO: special treatment for bipartite state
diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 3e2c5d374..8644d0f5c 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -3,14 +3,27 @@
 
 Collection of Trotter evolution gates and MPOs obtained from
 a Hamiltonian containing long-range or multi-site terms.
-Each item in `data` is a pair `sites => gate`, where `sites` is a
-vector of `CartesianIndex`s storing the sites on which the
-Trotter `gate` acts.
+
+## Fields
+
+- `lattice::Matrix{S}`: The lattice on which the gates acts.
+- `terms::T`: The vector of `sites => gate` pairs, where `sites` is a
+vector of `CartesianIndex`s storing the sites on which the `gate` acts.
 """
-struct TrotterGates{T <: Vector}
-    data::T
+struct TrotterGates{T <: Vector, S}
+    lattice::Matrix{S}
+    terms::T
+    # TODO: check physical spaces of terms match `lattice`
 end
 
+"""
+    physicalspace(gates::TrotterGates)
+
+Return lattice of physical spaces on which the `TrotterGates` is defined.
+"""
+physicalspace(gates::TrotterGates) = gates.lattice
+Base.size(gates::TrotterGates) = size(physicalspace(gates))
+
 const NNGate{T, S} = AbstractTensorMap{T, S, 2, 2}
 
 """
@@ -260,9 +273,12 @@ function trotterize(
     #= 
     2-site NNN gates converted to 3-site MPOs on triangular clusters [d, r, c]
     - d = 1 (NORTHWEST), ..., 4 (SOUTHWEST)
+
+    TODO: when all plaquettes have NNN gates, group with
+    NN gates and 1-site gates to reduce number of gates
     =#
     dist >= 2 && _trotterize_nnn2site!(gates, H, dt′; atol)
 
     symmetrize_gates && push!(gates, reverse(gates)...)
-    return TrotterGates(gates)
+    return TrotterGates(physicalspace(H), gates)
 end

From c225a0fdeb3eb67e35d4d3715399f71b7062321d Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Wed, 11 Mar 2026 22:30:06 +0800
Subject: [PATCH 25/27] Fix SU gauge fixing again

---
 src/algorithms/time_evolution/apply_gate.jl   | 30 -------------------
 src/algorithms/time_evolution/gaugefix_su.jl  | 16 ++++++----
 src/algorithms/time_evolution/simpleupdate.jl |  2 +-
 3 files changed, 11 insertions(+), 37 deletions(-)

diff --git a/src/algorithms/time_evolution/apply_gate.jl b/src/algorithms/time_evolution/apply_gate.jl
index 89b622a2e..171f327c2 100644
--- a/src/algorithms/time_evolution/apply_gate.jl
+++ b/src/algorithms/time_evolution/apply_gate.jl
@@ -146,33 +146,3 @@ function _apply_gate(
     end
     return a, s, b, ϵ
 end
-
-"""
-$(SIGNATURES)
-
-Apply a trivial gate (indicated by `nothing`) on the reduced matrices `a`, `b`
-```
-    -1← a --- 3 --- b ← -4          -2         -3
-        ↓           ↓               ↓           ↓
-        -2         -3           -1← a --- 3 --- b ← -4
-```
-"""
-function _apply_gate(
-        a::AbstractTensorMap, b::AbstractTensorMap,
-        ::Nothing, trunc::TruncationStrategy
-    )
-    V = space(b, 1)
-    need_flip = isdual(V)
-    @tensor a2b2[-1 -2; -3 -4] := a[-1 -2 3] * b[3 -3 -4]
-    trunc = if trunc isa FixedSpaceTruncation
-        need_flip ? truncspace(flip(V)) : truncspace(V)
-    else
-        trunc
-    end
-    a, s, b, ϵ = svd_trunc!(a2b2; trunc, alg = LAPACK_QRIteration())
-    a, b = absorb_s(a, s, b)
-    if need_flip
-        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/gaugefix_su.jl b/src/algorithms/time_evolution/gaugefix_su.jl
index 20858ceb7..a7b9a9d01 100644
--- a/src/algorithms/time_evolution/gaugefix_su.jl
+++ b/src/algorithms/time_evolution/gaugefix_su.jl
@@ -17,13 +17,18 @@ $(TYPEDFIELDS)
     maxiter::Int = 100
 end
 
-function _trivial_gates(Nrow::Int, Ncol::Int)
-    gates = map(Iterators.product(1:2, 1:Ncol, 1:Nrow)) do (d, c, r)
+function _trivial_gates(elt::Type{<:Number}, lattice::Matrix{S}) where {S <: ElementarySpace}
+    Nr, Nc = size(lattice)
+    gates = map(Iterators.product(1:2, 1:Nc, 1:Nr)) do (d, c, r)
         site1 = CartesianIndex(r, c)
         site2 = (d == 1) ? CartesianIndex(r, c + 1) : CartesianIndex(r - 1, c)
-        return [site1, site2] => nothing
+        r1, c1 = mod1.(Tuple(site1), size(lattice))
+        r2, c2 = mod1.(Tuple(site2), size(lattice))
+        V1, V2 = lattice[r1, c1], lattice[r2, c2]
+        h = TensorKit.id(elt, V1 ⊗ V2)
+        return [site1, site2] => h
     end
-    return TrotterGates(vec(gates))
+    return TrotterGates(lattice, vec(gates))
 end
 
 """
@@ -32,8 +37,7 @@ end
 Fix the gauge of `psi` using trivial simple update.
 """
 function gauge_fix(psi::InfiniteState, alg::SUGauge)
-    Nr, Nc, = size(psi)
-    gates = _trivial_gates(Nr, Nc)
+    gates = _trivial_gates(scalartype(psi), physicalspace(psi))
     su_alg = SimpleUpdate(; trunc = FixedSpaceTruncation(), bipartite = _state_bipartite_check(psi))
     wts0 = SUWeight(psi)
     # use default constructor to avoid calculation of exp(-H * 0)
diff --git a/src/algorithms/time_evolution/simpleupdate.jl b/src/algorithms/time_evolution/simpleupdate.jl
index 8d105ee83..1a06fcbcc 100644
--- a/src/algorithms/time_evolution/simpleupdate.jl
+++ b/src/algorithms/time_evolution/simpleupdate.jl
@@ -74,7 +74,7 @@ When `purified = true`, `gate` acts on the codomain physical legs of `state`.
 Otherwise, `gate` acts on both the codomain and the domain physical legs of `state`.
 """
 function _su_iter!(
-        state::InfiniteState, gate::Union{NNGate, Nothing}, env::SUWeight,
+        state::InfiniteState, gate::NNGate, env::SUWeight,
         sites::Vector{CartesianIndex{2}}, truncs::Vector{E};
         purified::Bool = true
     ) where {E <: TruncationStrategy}

From 9aa956efeef975e354a6f24717dad49682efca91 Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 13 Mar 2026 09:20:53 +0800
Subject: [PATCH 26/27] Undo changes to hubbard_model

---
 src/operators/models.jl | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/operators/models.jl b/src/operators/models.jl
index 3f302bbf2..37eceaa72 100644
--- a/src/operators/models.jl
+++ b/src/operators/models.jl
@@ -79,18 +79,15 @@ function MPSKitModels.hubbard_model(
     # TODO: just add this
     @assert n == 0 "Currently no support for imposing a fixed particle number"
     N = MPSKitModels.e_number(T, particle_symmetry, spin_symmetry)
-    spaces = fill(space(N, 1), (lattice.Nrows, lattice.Ncols))
-    hopping = (-t) * (
+    pspace = space(N, 1)
+    unit = TensorKit.id(pspace)
+    hopping =
         MPSKitModels.e⁺e⁻(T, particle_symmetry, spin_symmetry) +
-            MPSKitModels.e⁻e⁺(T, particle_symmetry, spin_symmetry)
-    )
+        MPSKitModels.e⁻e⁺(T, particle_symmetry, spin_symmetry)
     interaction_term = MPSKitModels.nꜛnꜜ(T, particle_symmetry, spin_symmetry)
     site_term = U * interaction_term - mu * N
-    return LocalOperator(
-        spaces,
-        (neighbor => hopping for neighbor in nearest_neighbours(lattice))...,
-        ((idx,) => site_term for idx in vertices(lattice))...,
-    )
+    h = (-t) * hopping + (1 / 4) * (site_term ⊗ unit + unit ⊗ site_term)
+    return nearest_neighbour_hamiltonian(fill(pspace, size(lattice)), h)
 end
 
 function MPSKitModels.bose_hubbard_model(

From db9b51c5773a38dd31eb9c58e62017d69a6239cd Mon Sep 17 00:00:00 2001
From: Yue Zhengyuan <yuezy1997@icloud.com>
Date: Fri, 13 Mar 2026 09:37:54 +0800
Subject: [PATCH 27/27] Update `trotterize` docstrings

---
 src/algorithms/time_evolution/trotter_gate.jl | 31 ++++++++++++-------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/algorithms/time_evolution/trotter_gate.jl b/src/algorithms/time_evolution/trotter_gate.jl
index 8644d0f5c..bd1c66e73 100644
--- a/src/algorithms/time_evolution/trotter_gate.jl
+++ b/src/algorithms/time_evolution/trotter_gate.jl
@@ -189,6 +189,9 @@ function _check_hamiltonian_for_trotter(H::LocalOperator)
     return dist
 end
 
+"""
+Trotterize a trivial Hamiltonian `H` containing only 1-site terms.
+"""
 function _trotterize_1site!(gates::Vector, H::LocalOperator, dt::Number; atol::Real)
     for site in CartesianIndices(size(H))
         gate = _get_site_term(H, site)
@@ -198,6 +201,14 @@ function _trotterize_1site!(gates::Vector, H::LocalOperator, dt::Number; atol::R
     return gates
 end
 
+"""
+Trotterize nearest neighbor terms (grouped with 1-site terms)
+in the Hamiltonian `H`.
+
+Gate order: `(d, c, r)`
+- d = 1: horizontal bond ((r, c), (r, c+1))
+- d = 2:   vertical bond ((r, c), (r-1, c))
+"""
 function _trotterize_nn2site!(gates::Vector, H::LocalOperator, dt::Number; atol::Real, force_mpo::Bool = false)
     Nr, Nc = size(H)
     T = scalartype(H)
@@ -219,6 +230,12 @@ function _trotterize_nn2site!(gates::Vector, H::LocalOperator, dt::Number; atol:
     return gates
 end
 
+"""
+Trotterize a next-nearest neighbor terms in a Hamiltonian.
+
+Gate order: `(c, r, d)`
+- d = 1 (NORTHWEST), ..., 4 (SOUTHWEST) labels the triangular 3-site clusters.
+"""
 function _trotterize_nnn2site!(gates::Vector, H::LocalOperator, dt::Number; atol::Real)
     Nr, Nc = size(H)
     T = scalartype(H)
@@ -264,19 +281,9 @@ function trotterize(
 
     # 1-site gates are only constructed when H only has 1-site terms
     dist == 0 && _trotterize_1site!(gates, H, dt′; atol)
-    #= 
-    2-site NN gates on bonds [d, r, c], grouped with 1-site terms
-    - d = 1: horizontal bond ((r, c), (r, c+1))
-    - d = 2:   vertical bond ((r, c), (r-1, c))
-    =#
+    # 2-site NN gates grouped with 1-site terms
     dist >= 1 && _trotterize_nn2site!(gates, H, dt′; atol, force_mpo)
-    #= 
-    2-site NNN gates converted to 3-site MPOs on triangular clusters [d, r, c]
-    - d = 1 (NORTHWEST), ..., 4 (SOUTHWEST)
-
-    TODO: when all plaquettes have NNN gates, group with
-    NN gates and 1-site gates to reduce number of gates
-    =#
+    # 3-site NNN gate MPOs
     dist >= 2 && _trotterize_nnn2site!(gates, H, dt′; atol)
 
     symmetrize_gates && push!(gates, reverse(gates)...)