Periodic indexing (#377)

* Add periodic getindex/setindex! utility

* add periodic getters and setters to object types

* simplify codebase to avoid modular indexing

* bypass non-scalar indexing calls

* Move indexing utils to a separate file

* Improve SUWeight to CTMRGEnv conversion

* Periodic indexing for SiteDependentTruncation

* Improve _is_bipartite

* Improve BPEnv indexing

* Drop unitcell argument in _nn_bondenv, _get_cluster_trunc

* Periodic `physicalspace` for LocalOperator, LocalCircuit

* Loose ends

---------

Co-authored-by: Yue Zhengyuan <yuezy1997@icloud.com>

Author: lkdvos

docs/src/examples/heisenberg_su/index.md

@@ -67,8 +67,8 @@ else
 end
 
 peps = InfinitePEPS(rand, Float64, physical_space, bond_space; unitcell = (Nr, Nc));
-peps.A[2, 2] = copy(peps.A[1, 1]) ## make initial random state bipartite
-peps.A[2, 1] = copy(peps.A[1, 2])
+peps[2, 2] = copy(peps[1, 1]) ## make initial random state bipartite
+peps[2, 1] = copy(peps[1, 2])
 wts = SUWeight(peps);
 ````
 

docs/src/examples/heisenberg_su/main.ipynb

@@ -97,8 +97,8 @@
     "end\n",
     "\n",
     "peps = InfinitePEPS(rand, Float64, physical_space, bond_space; unitcell = (Nr, Nc));\n",
-    "peps.A[2, 2] = copy(peps.A[1, 1]) ## make initial random state bipartite\n",
-    "peps.A[2, 1] = copy(peps.A[1, 2])\n",
+    "peps[2, 2] = copy(peps[1, 1]) ## make initial random state bipartite\n",
+    "peps[2, 1] = copy(peps[1, 2])\n",
     "wts = SUWeight(peps);"
    ],
    "metadata": {},

examples/heisenberg_su/main.jl

@@ -60,8 +60,8 @@ else
 end
 
 peps = InfinitePEPS(rand, Float64, physical_space, bond_space; unitcell = (Nr, Nc));
-peps.A[2, 2] = copy(peps.A[1, 1]) ## make initial random state bipartite
-peps.A[2, 1] = copy(peps.A[1, 2])
+peps[2, 2] = copy(peps[1, 1]) ## make initial random state bipartite
+peps[2, 1] = copy(peps[1, 2])
 wts = SUWeight(peps);
 
 md"""

src/PEPSKit.jl

@@ -41,6 +41,7 @@ using DocStringExtensions
 include("Defaults.jl")  # Include first to allow for docstring interpolation with Defaults values
 
 include("utility/util.jl")
+include("utility/indexing.jl")
 include("utility/diffable_threads.jl")
 include("utility/eigh.jl")
 include("utility/svd.jl")
@@ -133,6 +134,7 @@ using .Defaults: set_scheduler!
 export set_scheduler!
 export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint
 export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
+export corner, edge, setcorner!, setedge!
 export FixedSpaceTruncation, SiteDependentTruncation
 export HalfInfiniteProjector, FullInfiniteProjector
 export C4vCTMRG, C4vEighProjector, C4vQRProjector

src/algorithms/bp/beliefpropagation.jl

@@ -104,10 +104,10 @@ function update_message(I::CartesianIndex{3}, network::InfiniteSquareNetwork, en
     (1 <= dir <= 4) || throw(ArgumentError("Invalid direction $dir"))
 
     A = network[row, col]
-    dir == SOUTH || (M_north = env[NORTH, _prev(row, end), col])
-    dir == WEST || (M_east = env[EAST, row, _next(col, end)])
-    dir == NORTH || (M_south = env[SOUTH, _next(row, end), col])
-    dir == EAST || (M_west = env[WEST, row, _prev(col, end)])
+    dir == SOUTH || (M_north = env[NORTH, row - 1, col])
+    dir == WEST || (M_east = env[EAST, row, col + 1])
+    dir == NORTH || (M_south = env[SOUTH, row + 1, col])
+    dir == EAST || (M_west = env[WEST, row, col - 1])
 
     return if dir == NORTH
         contract_north_message(A, M_west, M_north, M_east)

src/algorithms/bp/gaugefix.jl

@@ -25,7 +25,7 @@ function gauge_fix(psi::InfinitePEPS, alg::BPGauge, env::BPEnv)
         # copy 1st column to 2nd column to eliminate differences
         # caused by order of applying gauge transformations
         for r in 1:2
-            psi′[_next(r, 2), 2] = copy(psi′[r, 1])
+            psi′[r + 1, 2] = copy(psi′[r, 1])
         end
     end
     return psi′, XXinv
@@ -52,7 +52,7 @@ end
 function _sqrt_bp_messages(I::CartesianIndex{3}, env::BPEnv)
     dir, row, col = Tuple(I)
     @assert dir == NORTH || dir == EAST
-    M12 = env[dir, dir == NORTH ? _prev(row, end) : row, dir == EAST ? _next(col, end) : col]
+    M12 = env[dir, dir == NORTH ? row - 1 : row, dir == EAST ? col + 1 : col]
     sqrtM12, isqrtM12 = sqrt_invsqrt(twist(M12, 1))
     M21 = env[dir + 2, row, col]
     sqrtM21, isqrtM21 = sqrt_invsqrt(M21)
@@ -90,10 +90,10 @@ function _bp_gauge_fix!(I::CartesianIndex{3}, psi::InfinitePEPS, env::BPEnv)
     end
     if dir == NORTH
         psi[row, col] = absorb_north_message(psi[row, col], X)
-        psi[_prev(row, end), col] = absorb_south_message(psi[_prev(row, end), col], invX)
+        psi[row - 1, col] = absorb_south_message(psi[row - 1, col], invX)
     elseif dir == EAST
         psi[row, col] = absorb_east_message(psi[row, col], X)
-        psi[row, _next(col, end)] = absorb_west_message(psi[row, _next(col, end)], invX)
+        psi[row, col + 1] = absorb_west_message(psi[row, col + 1], invX)
     end
     return psi, X, invX
 end
@@ -122,9 +122,9 @@ to a belief propagation environment.
 function BPEnv(wts::SUWeight)
     messages = map(Iterators.product(1:4, axes(wts, 2), axes(wts, 3))) do (d, r, c)
         wt = if d == NORTH
-            twist(wts[2, _next(r, end), c], 1)
+            twist(wts[2, r + 1, c], 1)
         elseif d == EAST
-            twist(wts[1, r, _prev(c, end)], 1)
+            twist(wts[1, r, c - 1], 1)
         elseif d == SOUTH
             copy(wts[2, r, c])
         else # WEST

src/algorithms/contractions/absorb_weight.jl

@@ -70,18 +70,16 @@ square root (or inverse square root if `inv = true`).
 function weight_to_absorb(
         weights::SUWeight, row::Int, col::Int, ax::Int; inv::Bool = false
     )
-    _, Nr, Nc = size(weights)
-    r, c = mod1(row, Nr), mod1(col, Nc)
     pow = inv ? -1 / 2 : 1 / 2
     wt = sdiag_pow(
         if ax == NORTH
-            weights[2, r, c]
+            weights[2, row, col]
         elseif ax == EAST
-            weights[1, r, c]
+            weights[1, row, col]
         elseif ax == SOUTH
-            weights[2, _next(r, Nr), c]
+            weights[2, row + 1, col]
         else # WEST
-            weights[1, r, _prev(c, Nc)]
+            weights[1, row, col - 1]
         end,
         pow,
     )

src/algorithms/contractions/bondenv/benv_ctm.jl

@@ -24,20 +24,14 @@ Axis order: `[DX1 DY1; DX0 DY0]`, as in
 ```
 """
 function bondenv_ctm(row::Int, col::Int, X, Y, env::CTMRGEnv)
-    Nr, Nc = size(env.corners)[[2, 3]]
-    cm1 = _prev(col, Nc)
-    cp1 = _next(col, Nc)
-    cp2 = _next(cp1, Nc)
-    rm1 = _prev(row, Nr)
-    rp1 = _next(row, Nr)
-    c1 = env.corners[1, rm1, cm1]
-    c2 = env.corners[2, rm1, cp2]
-    c3 = env.corners[3, rp1, cp2]
-    c4 = env.corners[4, rp1, cm1]
-    e1X, e1Y = env.edges[1, rm1, col], env.edges[1, rm1, cp1]
-    e2 = env.edges[2, row, cp2]
-    e3X, e3Y = env.edges[3, rp1, col], env.edges[3, rp1, cp1]
-    e4 = env.edges[4, row, cm1]
+    c1 = corner(env, 1, row - 1, col - 1)
+    c2 = corner(env, 2, row - 1, col + 2)
+    c3 = corner(env, 3, row + 1, col + 2)
+    c4 = corner(env, 4, row + 1, col - 1)
+    e1X, e1Y = edge(env, 1, row - 1, col), edge(env, 1, row - 1, col + 1)
+    e2 = edge(env, 2, row, col + 2)
+    e3X, e3Y = edge(env, 3, row + 1, col), edge(env, 3, row + 1, col + 1)
+    e4 = edge(env, 4, row, col - 1)
     #= index labels
 
     C1--χ4--E1X---------χ6---------E1Y--χ8---C2     r-1

src/algorithms/contractions/bp_contractions.jl

@@ -105,14 +105,14 @@ function contract_local_operator1x1(
         env::BPEnv,
     )
     row, col = Tuple(ind)
-    M_north = env.messages[NORTH, _prev(row, end), mod1(col, end)]
-    M_east = env.messages[EAST, mod1(row, end), _next(col, end)]
-    M_south = env.messages[SOUTH, _next(row, end), mod1(col, end)]
-    M_west = env.messages[WEST, mod1(row, end), _prev(col, end)]
+    M_north = env[NORTH, row - 1, col]
+    M_east = env[EAST, row, col + 1]
+    M_south = env[SOUTH, row + 1, col]
+    M_west = env[WEST, row, col - 1]
 
     return @autoopt @tensor begin
-        ket[mod1(row, end), mod1(col, end)][dt; DNt DEt DSt DWt] *
-            conj(bra[mod1(row, end), mod1(col, end)][db; DNb DEb DSb DWb]) *
+        ket[row, col][dt; DNt DEt DSt DWt] *
+            conj(bra[row, col][db; DNb DEb DSb DWb]) *
             O[db; dt] *
             M_north[DNt; DNb] *
             M_east[DEt; DEb] *
@@ -125,14 +125,14 @@ function contract_local_norm1x1(
         ind::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv
     )
     row, col = Tuple(ind)
-    M_north = env.messages[NORTH, _prev(row, end), mod1(col, end)]
-    M_east = env.messages[EAST, mod1(row, end), _next(col, end)]
-    M_south = env.messages[SOUTH, _next(row, end), mod1(col, end)]
-    M_west = env.messages[WEST, mod1(row, end), _prev(col, end)]
+    M_north = env[NORTH, row - 1, col]
+    M_east = env[EAST, row, col + 1]
+    M_south = env[SOUTH, row + 1, col]
+    M_west = env[WEST, row, col - 1]
 
     return @autoopt @tensor begin
-        ket[mod1(row, end), mod1(col, end)][d; DNt DEt DSt DWt] *
-            conj(bra[mod1(row, end), mod1(col, end)][d; DNb DEb DSb DWb]) *
+        ket[row, col][d; DNt DEt DSt DWt] *
+            conj(bra[row, col][d; DNb DEb DSb DWb]) *
             M_north[DNt; DNb] *
             M_east[DEt; DEb] *
             M_south[DSb; DSt] *
@@ -148,17 +148,17 @@ function contract_local_operator2x1(
         env::BPEnv,
     )
     row, col = Tuple(coord)
-    M_north = env.messages[NORTH, _prev(row, end), mod1(col, end)]
-    M_northeast = env.messages[EAST, mod1(row, end), _next(col, end)]
-    M_southeast = env.messages[EAST, _next(row, end), _next(col, end)]
-    M_south = env.messages[SOUTH, mod1(row + 2, end), mod1(col, end)]
-    M_southwest = env.messages[WEST, _next(row, end), _prev(col, end)]
-    M_northwest = env.messages[WEST, mod1(row, end), _prev(col, end)]
+    M_north = env[NORTH, row - 1, col]
+    M_northeast = env[EAST, row, col + 1]
+    M_southeast = env[EAST, row + 1, col + 1]
+    M_south = env[SOUTH, row + 2, col]
+    M_southwest = env[WEST, row + 1, col - 1]
+    M_northwest = env[WEST, row, col - 1]
 
-    return @autoopt @tensor ket[mod1(row, end), mod1(col, end)][dNt; DNt DNEt DMt DNWt] *
-        ket[_next(row, end), mod1(col, end)][dSt; DMt DSEt DSt DSWt] *
-        conj(bra[mod1(row, end), mod1(col, end)][dNb; DNb DNEb DMb DNWb]) *
-        conj(bra[_next(row, end), mod1(col, end)][dSb; DMb DSEb DSb DSWb]) *
+    return @autoopt @tensor ket[row, col][dNt; DNt DNEt DMt DNWt] *
+        ket[row + 1, col][dSt; DMt DSEt DSt DSWt] *
+        conj(bra[row, col][dNb; DNb DNEb DMb DNWb]) *
+        conj(bra[row + 1, col][dSb; DMb DSEb DSb DSWb]) *
         M_north[DNt; DNb] *
         M_northeast[DNEt; DNEb] *
         M_southeast[DSEt; DSEb] *
@@ -176,16 +176,16 @@ function contract_local_operator1x2(
         env::BPEnv,
     )
     row, col = Tuple(coord)
-    M_west = env.messages[WEST, mod1(row, end), _prev(col, end)]
-    M_northwest = env.messages[NORTH, _prev(row, end), mod1(col, end)]
-    M_northeast = env.messages[NORTH, _prev(row, end), _next(col, end)]
-    M_east = env.messages[EAST, mod1(row, end), mod1(col + 2, end)]
-    M_southeast = env.messages[SOUTH, _next(row, end), _next(col, end)]
-    M_southwest = env.messages[SOUTH, _next(row, end), mod1(col, end)]
-    A_west = ket[mod1(row, end), mod1(col, end)]
-    Ā_west = bra[mod1(row, end), mod1(col, end)]
-    A_east = ket[mod1(row, end), _next(col, end)]
-    Ā_east = bra[mod1(row, end), _next(col, end)]
+    M_west = env[WEST, row, col - 1]
+    M_northwest = env[NORTH, row - 1, col]
+    M_northeast = env[NORTH, row - 1, col + 1]
+    M_east = env[EAST, row, col + 2]
+    M_southeast = env[SOUTH, row + 1, col + 1]
+    M_southwest = env[SOUTH, row + 1, col]
+    A_west = ket[row, col]
+    Ā_west = bra[row, col]
+    A_east = ket[row, col + 1]
+    Ā_east = bra[row, col + 1]
 
     return @autoopt @tensor begin
         A_west[dWt; DNWt DMt DSWt DWt] *
@@ -206,17 +206,17 @@ function contract_local_norm2x1(
         coord::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, env::BPEnv
     )
     row, col = Tuple(coord)
-    M_north = env.messages[NORTH, _prev(row, end), mod1(col, end)]
-    M_northeast = env.messages[EAST, mod1(row, end), _next(col, end)]
-    M_southeast = env.messages[EAST, _next(row, end), _next(col, end)]
-    M_south = env.messages[SOUTH, mod1(row + 2, end), mod1(col, end)]
-    M_southwest = env.messages[WEST, _next(row, end), _prev(col, end)]
-    M_northwest = env.messages[WEST, mod1(row, end), _prev(col, end)]
+    M_north = env[NORTH, row - 1, col]
+    M_northeast = env[EAST, row, col + 1]
+    M_southeast = env[EAST, row + 1, col + 1]
+    M_south = env[SOUTH, row + 2, col]
+    M_southwest = env[WEST, row + 1, col - 1]
+    M_northwest = env[WEST, row, col - 1]
 
-    return @autoopt @tensor ket[mod1(row, end), mod1(col, end)][dN; DNt DNEt DMt DNWt] *
-        ket[_next(row, end), mod1(col, end)][dS; DMt DSEt DSt DSWt] *
-        conj(bra[mod1(row, end), mod1(col, end)][dN; DNb DNEb DMb DNWb]) *
-        conj(bra[_next(row, end), mod1(col, end)][dS; DMb DSEb DSb DSWb]) *
+    return @autoopt @tensor ket[row, col][dN; DNt DNEt DMt DNWt] *
+        ket[row + 1, col][dS; DMt DSEt DSt DSWt] *
+        conj(bra[row, col][dN; DNb DNEb DMb DNWb]) *
+        conj(bra[row + 1, col][dS; DMb DSEb DSb DSWb]) *
         M_north[DNt; DNb] *
         M_northeast[DNEt; DNEb] *
         M_southeast[DSEt; DSEb] *
@@ -230,17 +230,17 @@ function contract_local_norm1x2(
     )
     row, col = Tuple(coord)
 
-    M_west = env.messages[WEST, mod1(row, end), _prev(col, end)]
-    M_northwest = env.messages[NORTH, _prev(row, end), mod1(col, end)]
-    M_northeast = env.messages[NORTH, _prev(row, end), _next(col, end)]
-    M_east = env.messages[EAST, mod1(row, end), mod1(col + 2, end)]
-    M_southeast = env.messages[SOUTH, _next(row, end), _next(col, end)]
-    M_southwest = env.messages[SOUTH, _next(row, end), mod1(col, end)]
+    M_west = env[WEST, row, col - 1]
+    M_northwest = env[NORTH, row - 1, col]
+    M_northeast = env[NORTH, row - 1, col + 1]
+    M_east = env[EAST, row, col + 2]
+    M_southeast = env[SOUTH, row + 1, col + 1]
+    M_southwest = env[SOUTH, row + 1, col]
 
-    A_west = ket[mod1(row, end), mod1(col, end)]
-    Ā_west = bra[mod1(row, end), mod1(col, end)]
-    A_east = ket[mod1(row, end), _next(col, end)]
-    Ā_east = bra[mod1(row, end), _next(col, end)]
+    A_west = ket[row, col]
+    Ā_west = bra[row, col]
+    A_east = ket[row, col + 1]
+    Ā_east = bra[row, col + 1]
 
     return @autoopt @tensor begin
         A_west[dW; DNWt DMt DSWt DWt] *

src/algorithms/contractions/correlator/pepo_1layer.jl

@@ -5,12 +5,12 @@ function start_correlator(
     (size(ρ, 3) == 1) ||
         throw(ArgumentError("The input PEPO ρ must have only one layer."))
     r, c = Tuple(i)
-    E_north = env.edges[NORTH, _prev(r, end), mod1(c, end)]
-    E_south = env.edges[SOUTH, _next(r, end), mod1(c, end)]
-    E_west = env.edges[WEST, mod1(r, end), _prev(c, end)]
-    C_northwest = env.corners[NORTHWEST, _prev(r, end), _prev(c, end)]
-    C_southwest = env.corners[SOUTHWEST, _next(r, end), _prev(c, end)]
-    t = twistdual(ρ[mod1(r, end), mod1(c, end)], 1:2)
+    E_north = edge(env, NORTH, r - 1, c)
+    E_south = edge(env, SOUTH, r + 1, c)
+    E_west = edge(env, WEST, r, c - 1)
+    C_northwest = corner(env, NORTHWEST, r - 1, c - 1)
+    C_southwest = corner(env, SOUTHWEST, r + 1, c - 1)
+    t = twistdual(ρ[r, c], 1:2)
     # TODO: part of these contractions is duplicated between the two output tensors,
     # so could be optimized further
     @autoopt @tensor Vn[χSE De; χNE] :=
@@ -32,12 +32,12 @@ function end_correlator_numerator(
     (size(ρ, 3) == 1) ||
         throw(ArgumentError("The input PEPO ρ must have only one layer."))
     r, c = Tuple(j)
-    E_north = env.edges[NORTH, _prev(r, end), mod1(c, end)]
-    E_east = env.edges[EAST, mod1(r, end), _next(c, end)]
-    E_south = env.edges[SOUTH, _next(r, end), mod1(c, end)]
-    C_northeast = env.corners[NORTHEAST, _prev(r, end), _next(c, end)]
-    C_southeast = env.corners[SOUTHEAST, _next(r, end), _next(c, end)]
-    t = twistdual(ρ[mod1(r, end), mod1(c, end)], 1:2)
+    E_north = edge(env, NORTH, r - 1, c)
+    E_east = edge(env, EAST, r, c + 1)
+    E_south = edge(env, SOUTH, r + 1, c)
+    C_northeast = corner(env, NORTHEAST, r - 1, c + 1)
+    C_southeast = corner(env, SOUTHEAST, r + 1, c + 1)
+    t = twistdual(ρ[r, c], 1:2)
     return @autoopt @tensor V[χSW DW dstring; χNW] *
         E_south[χSSE DS; χSW] * E_east[χNEE DE; χSEE] * E_north[χNW DN; χNNE] *
         C_northeast[χNNE; χNEE] * C_southeast[χSEE; χSSE] *
@@ -48,9 +48,9 @@ function end_correlator_denominator(
         j::CartesianIndex{2}, V::CTMRGEdgeTensor{T, S, 2}, env::CTMRGEnv
     ) where {T, S}
     r, c = Tuple(j)
-    C_northeast = env.corners[NORTHEAST, _prev(r, end), _next(c, end)]
-    E_east = env.edges[EAST, mod1(r, end), _next(c, end)]
-    C_southeast = env.corners[SOUTHEAST, _next(r, end), _next(c, end)]
+    C_northeast = corner(env, NORTHEAST, r - 1, c + 1)
+    E_east = edge(env, EAST, r, c + 1)
+    C_southeast = corner(env, SOUTHEAST, r + 1, c + 1)
     return @autoopt @tensor V[χS DE; χN] * C_northeast[χN; χNE] *
         E_east[χNE DE; χSE] * C_southeast[χSE; χS]
 end