Refactor bond direction helpers to use lattice direction constants

Author: Yue-Zhengyuan

src/algorithms/time_evolution/get_cluster.jl

@@ -51,46 +51,66 @@ 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`).
+return `(dir, r, c)` where `dir` ∈ {NORTH, EAST, SOUTH, WEST}
+is the direction from `site1` to `site2`.
+`(r, c)` is the position of the source endpoint:
+the west site for x-bonds (`dir == EAST` or `dir == WEST`),
+the south site for y-bonds (`dir == NORTH` or `dir == SOUTH`).
 """
-function _nn_bondrev(site1::CartesianIndex{2}, site2::CartesianIndex{2})
+function _nn_bonddir(site1::CartesianIndex{2}, site2::CartesianIndex{2})
     diff = site1 - site2
-    if diff == CartesianIndex(0, -1)
+    if diff == CartesianIndex(0, -1)     # site1 west of site2 → site2 to the EAST
         r, c = site1[1], site1[2]
-        return (1, r, c), false
-    elseif diff == CartesianIndex(0, 1)
-        r, c = site2[1], site2[2]
-        return (1, r, c), true
-    elseif diff == CartesianIndex(1, 0)
+        return (EAST, r, c)
+    elseif diff == CartesianIndex(1, 0)  # site1 south of site2 → site2 to the NORTH
         r, c = site1[1], site1[2]
-        return (2, r, c), false
-    elseif diff == CartesianIndex(-1, 0)
+        return (NORTH, r, c)
+    elseif diff == CartesianIndex(-1, 0) # site1 north of site2 → site2 to the SOUTH
+        r, c = site2[1], site2[2]
+        return (SOUTH, r, c)
+    elseif diff == CartesianIndex(0, 1)  # site1 east of site2 → site2 to the WEST
         r, c = site2[1], site2[2]
-        return (2, r, c), true
+        return (WEST, r, c)
     else
         error("`site1` and `site2` are not nearest neighbors.")
     end
 end
 
-function _bond_rotation(x, bonddir::Int, rev::Bool; inv::Bool = false)
-    return if bonddir == 1 # x-bond
-        rev ? rot180(x) : x
-    elseif bonddir == 2 # y-bond
-        if rev
-            inv ? rotr90(x) : rotl90(x)
-        else
-            inv ? rotl90(x) : rotr90(x)
-        end
-    else
-        error("`bonddir` must be 1 (for x-bonds) or 2 (for y-bonds).")
+"""
+Apply the tensor rotation that maps a bond from orientation `dir_from`
+to orientation `dir_to`, where the directions are one of
+`NORTH`, `EAST`, `SOUTH`, `WEST`.
+"""
+function _rotate_by_dir(x, dir_from::Int, dir_to::Int)
+    delta = mod(dir_to - dir_from, 4)
+    return if delta == 0
+        x
+    elseif delta == 1
+        rotr90(x)
+    elseif delta == 2
+        rot180(x)
+    else # delta == 3
+        rotl90(x)
     end
 end
-function _bond_rotation(x::CartesianIndex{2}, bonddir::Int, rev::Bool, unitcell::NTuple{2, Int})
-    return if bonddir == 1
-        rev ? siterot180(x, unitcell) : x
-    else
-        rev ? siterotl90(x, unitcell) : siterotr90(x, unitcell)
+
+"""
+Map a site coordinate from the lattice frame where the bond has orientation
+`dir_from` to the frame where the bond has orientation `dir_to`.
+"""
+function _siterot_by_dir(
+        site::CartesianIndex{2}, dir_from::Int, dir_to::Int,
+        unitcell::NTuple{2, Int}
+    )
+    delta = mod(dir_to - dir_from, 4)
+    return if delta == 0
+        site
+    elseif delta == 1
+        siterotr90(site, unitcell)
+    elseif delta == 2
+        siterot180(site, unitcell)
+    else # delta == 3
+        siterotl90(site, unitcell)
     end
 end
 

src/algorithms/time_evolution/ntupdate.jl

@@ -98,7 +98,8 @@ function ntu_iter(
             state2[site] = _apply_sitegate(state2[site], gate)
             info′ = (; fid = 1.0)
         elseif length(sites) == 2
-            (d, r, c), = _nn_bondrev(sites...)
+            (dir, r, c) = _nn_bonddir(sites...)
+            d = dir in (EAST, WEST) ? 1 : 2
             alg.bipartite && iseven(r) && continue
             state2, wts, info′ = _ntu_iter(state2, gate, wts, sites, alg)
             (!alg.bipartite) && continue

src/algorithms/time_evolution/ntupdate3site.jl

@@ -64,15 +64,14 @@ function _bond_truncate(
         (stype1, stype2)::NTuple{2, Symbol},
         alg::NeighbourUpdate; gate::Union{NNGate, Nothing} = nothing
     )
-    # rotate bond to standard x direction `A ← B`
+    # get bond direction and rotate state to standard x-direction `A ← B`
+    (dir, _, _) = _nn_bonddir(site1, site2)
     ucell = size(state)[1:2]
-    bond, rev = _nn_bondrev(site1, site2)
-    dir = first(bond)
-    state2 = _bond_rotation(state, dir, rev; inv = false)
-    wts2 = _bond_rotation(wts, dir, rev; inv = false)
+    state2 = _rotate_by_dir(state, dir, EAST)
+    wts2 = _rotate_by_dir(wts, dir, EAST)
 
-    # rotated bond tensors
-    siteA = _bond_rotation(site1, dir, rev, ucell)
+    # site1 position in the rotated frame
+    siteA = _siterot_by_dir(site1, dir, EAST, ucell)
     row, col = siteA[1], siteA[2]
     A, B = state2[row, col], state2[row, col + 1]
 
@@ -123,8 +122,8 @@ function _bond_truncate(
     state2[row, col + 1] = normalize!(B, Inf)
     wts2[1, row, col] = normalize!(s, Inf)
 
-    # rotate back tensors and bond weight
-    state2 = _bond_rotation(state2, dir, rev; inv = true)
-    wts2 = _bond_rotation(wts2, dir, rev; inv = true)
+    # rotate back state and weights
+    state2 = _rotate_by_dir(state2, EAST, dir)
+    wts2 = _rotate_by_dir(wts2, EAST, dir)
     return state2, wts2, info
 end

src/algorithms/time_evolution/simpleupdate.jl

@@ -86,9 +86,8 @@ function _su_iter!(
     Ms, open_vaxs, = _get_cluster(state, sites)
     _absorb_weight!(Ms, sites, open_vaxs, env)
     # rotate
-    bond, rev = _nn_bondrev(sites...)
-    dir = first(bond)
-    A, B = _bond_rotation.(Ms, dir, rev; inv = false)
+    (dir, r, c) = _nn_bonddir(sites...)
+    A, B = _rotate_by_dir.(Ms, dir, EAST)
     # apply gate
     ϵ = 0.0
     local s
@@ -102,17 +101,18 @@ function _su_iter!(
         alg.purified && break # only apply gate to 1st physical leg
     end
     # rotate back
-    A = _bond_rotation(A, dir, rev; inv = true)
-    B = _bond_rotation(B, dir, rev; inv = true)
-    rev && (s = transpose(s))
+    A = _rotate_by_dir(A, EAST, dir)
+    B = _rotate_by_dir(B, EAST, dir)
+    (dir in (WEST, SOUTH)) && (s = transpose(s))
     # remove environment weights
     siteA, siteB = sites
     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
     state[siteA] = normalize!(A, Inf)
     state[siteB] = normalize!(B, Inf)
-    env[bond...] = normalize!(s, Inf)
+    d = dir in (EAST, WEST) ? 1 : 2
+    env[d, r, c] = normalize!(s, Inf)
     return ϵ
 end
 
@@ -131,7 +131,8 @@ function su_iter(
             site = only(sites)
             state2[site] = _apply_sitegate(state2[site], gate; alg.purified)
         elseif length(sites) == 2
-            (d, r, c), = _nn_bondrev(sites...)
+            (dir, r, c) = _nn_bonddir(sites...)
+            d = dir in (EAST, WEST) ? 1 : 2
             alg.bipartite && iseven(r) && continue
             ϵ′ = _su_iter!(state2, gate, env2, sites, alg)
             ϵ = max(ϵ, ϵ′)

src/algorithms/time_evolution/simpleupdate3site.jl

@@ -47,13 +47,14 @@ function _su_iter!(
     # restore virtual arrows in `Ms`
     _flip_virtuals!(Ms, flips)
     # update env weights
-    bond_revs = map(sites, Iterators.drop(sites, 1)) do site1, site2
-        _nn_bondrev(site1, site2)
+    bond_dirs = map(sites, Iterators.drop(sites, 1)) do site1, site2
+        _nn_bonddir(site1, site2)
     end
-    for (wt, (bond, rev), flip) in zip(wts, bond_revs, flips)
+    for (wt, (dir, r, c), flip) in zip(wts, bond_dirs, flips)
         wt_new = flip ? _fliptwist_s(wt) : wt
-        wt_new = rev ? transpose(wt_new) : wt_new
-        env[CartesianIndex(bond)] = normalize!(wt_new, Inf)
+        wt_new = (dir in (WEST, SOUTH)) ? transpose(wt_new) : wt_new
+        d = dir in (EAST, WEST) ? 1 : 2
+        env[d, r, c] = normalize!(wt_new, Inf)
     end
     # update state tensors
     for (M, s, invperm, vaxs) in zip(Ms, sites, invperms, open_vaxs)
@@ -72,9 +73,10 @@ function _get_cluster_trunc(
         trunc::TruncationStrategy, sites::Vector{CartesianIndex{2}}
     )
     return map(sites, Iterators.drop(sites, 1)) do site1, site2
-        (d, r, c), rev = _nn_bondrev(site1, site2)
+        (dir, r, c) = _nn_bonddir(site1, site2)
+        d = dir in (EAST, WEST) ? 1 : 2
         t = truncation_strategy(trunc, d, r, c)
-        if rev && isa(t, TruncationSpace)
+        if (dir in (WEST, SOUTH)) && isa(t, TruncationSpace)
             t = truncspace(flip(t.space)')
         end
         return t