Simple update for PEPO

README.md

@@ -33,7 +33,7 @@ pkg> add PEPSKit
 - Contraction of infinite PEPS using the corner transfer matrix renormalization group (CTMRG) and boundary MPS methods
 - Native support for symmetric tensors through [TensorKit](https://github.com/Jutho/TensorKit.jl), including fermionic tensors
 - PEPS optimization using automatic differentiation (AD) provided through [Zygote](https://fluxml.ai/Zygote.jl/stable/)
-- Imaginary time evolution algorithms
+- Imaginary time evolution algorithms (including ground state and finite-temperature simulations)
 - Support for PEPS with generic unit cells
 - Support for classical 2D partition functions and projected entangled-pair operators (PEPOs)
 - Extensible system for custom states, operators and algorithms

src/PEPSKit.jl

@@ -15,7 +15,8 @@ using LoggingExtras
 
 using MPSKit
 using MPSKit: MPOTensor, GenericMPSTensor, MPSBondTensor, TransferMatrix
-import MPSKit: leading_boundary, loginit!, logiter!, logfinish!, logcancel!, physicalspace
+import MPSKit: tensorexpr, leading_boundary, loginit!, logiter!, logfinish!, logcancel!, physicalspace
+import MPSKit: infinite_temperature_density_matrix
 
 using MPSKitModels
 using FiniteDifferences
@@ -105,6 +106,7 @@ export InfinitePartitionFunction
 export InfinitePEPS, InfiniteTransferPEPS
 export SUWeight
 export InfinitePEPO, InfiniteTransferPEPO
+export infinite_temperature_density_matrix
 export initialize_mps, initializePEPS
 export ReflectDepth, ReflectWidth, Rotate, RotateReflect
 export symmetrize!, symmetrize_retract_and_finalize!

src/algorithms/contractions/bondenv/benv_tools.jl

@@ -1,2 +1,3 @@
 const BondEnv{T, S} = AbstractTensorMap{T, S, 2, 2} where {T <: Number, S <: ElementarySpace}
 const PEPSOrth{T, S} = AbstractTensor{T, S, 4} where {T <: Number, S <: ElementarySpace}
+const PEPOOrth{T, S} = AbstractTensor{T, S, 5} where {T <: Number, S <: ElementarySpace}

src/algorithms/contractions/localoperator.jl

@@ -1,6 +1,5 @@
 # Contraction of local operators on arbitrary lattice locations
 # -------------------------------------------------------------
-import MPSKit: tensorexpr
 
 # currently need this because MPSKit restricts tensor names to symbols
 _totuple(t) = t isa Tuple ? t : tuple(t)

src/algorithms/time_evolution/evoltools.jl

@@ -1,3 +1,16 @@
+const InfiniteState = Union{InfinitePEPS, InfinitePEPO}
+
+function MPSKit.infinite_temperature_density_matrix(H::LocalOperator)
+    T = promote_type((scalartype(t.second) for t in H.terms)...)
+    A = map(physicalspace(H)) do Vp
+        ψ = permute(TensorKit.id(T, Vp), (1, 2))
+        Vv = oneunit(Vp) # trivial (1D) virtual space
+        virt = ones(T, domain(ψ) ← Vv ⊗ Vv ⊗ Vv' ⊗ Vv')
+        return ψ * virt
+    end
+    return InfinitePEPO(cat(A; dims = 3))
+end
+
 """
     get_expham(H::LocalOperator, dt::Number)
 
@@ -71,75 +84,121 @@ end
 """
 $(SIGNATURES)
 
-Use QR decomposition on two tensors connected by a bond
-to get the reduced tensors
+Use QR decomposition on two tensors `A`, `B` connected by a bond to get the reduced tensors.
+When `A`, `B` are PEPSTensors,
 ```
-        2                   1
-        |                   |
-    5 - A - 3   ====>   4 - X ← 2   1 ← a - 3
-        | ↘                 |            ↘
-        4   1               3             2
-
-        2                               1
-        |                               |
-    5 - B - 3   ====>   1 - b → 3   4 → Y - 2
-        | ↘                  ↘          |
-        4   1                 2         3
+        2                   1                                   1
+        |                   |                                   |
+    5 -A/B- 3   ====>   4 - X ← 2   1 ← a - 3   1 - b → 3   4 → Y - 2
+        | ↘                 |            ↘           ↘          |
+        4   1               3             2           2         3
+```
+When `A`, `B` are PEPOTensors, 
+- If `gate_ax = 1`
+```
+    2   3                1  2                                1  2
+      ↘ |                 ↘ |                                 ↘ |
+    6 -A/B- 4   ====>   5 - X ← 3   1 ← a - 3   1 - b → 3   5 → Y - 3
+        | ↘                 |            ↘           ↘          |
+        5   1               4             2           2         4
+```
+- If `gate_ax = 2`
+```
+    2   3                   2         2           2             2
+      ↘ |                   |          ↘           ↘            |
+    6 -A/B- 4   ====>   5 - X ← 3   1 ← a - 3   1 - b → 3   5 → Y - 3
+        | ↘                 | ↘                                 | ↘
+        5   1               4  1                                4  1
 ```
 """
-function _qr_bond(A::PEPSTensor, B::PEPSTensor)
-    X, a = leftorth(A, ((2, 4, 5), (1, 3)))
-    Y, b = leftorth(B, ((2, 3, 4), (1, 5)))
+function _qr_bond(A::PT, B::PT; gate_ax::Int = 1) where {PT <: Union{PEPSTensor, PEPOTensor}}
+    @assert 1 <= gate_ax <= numout(A)
+    permA, permB, permX, permY = if A isa PEPSTensor
+        ((2, 4, 5), (1, 3)), ((2, 3, 4), (1, 5)), (1, 4, 2, 3), Tuple(1:4)
+    else
+        if gate_ax == 1
+            ((2, 3, 5, 6), (1, 4)), ((2, 3, 4, 5), (1, 6)), (1, 2, 5, 3, 4), Tuple(1:5)
+        else
+            ((1, 3, 5, 6), (2, 4)), ((1, 3, 4, 5), (2, 6)), (1, 2, 5, 3, 4), Tuple(1:5)
+        end
+    end
+    X, a = leftorth(A, permA)
+    Y, b = leftorth(B, permB)
     @assert !isdual(space(a, 1))
     @assert !isdual(space(b, 1))
-    X = permute(X, (1, 4, 2, 3))
-    Y = permute(Y, (1, 2, 3, 4))
+    X, Y = permute(X, permX), permute(Y, permY)
     b = permute(b, ((3, 2), (1,)))
     return X, a, b, Y
 end
 
 """
 $(SIGNATURES)
 
-Reconstruct the tensors connected by a bond from their QR results
-obtained from `_qr_bond`
+Reconstruct the tensors connected by a bond from their `_qr_bond` results.
+For PEPSTensors,
 ```
         -2                             -2
         |                               |
     -5- X - 1 - a - -3     -5 - b - 1 - Y - -3
         |        ↘               ↘      |
         -4        -1              -1   -4
 ```
+For PEPOTensors
+```
+    -2  -3                          -2  -3
+      ↘ |                             ↘ |
+    -6- X - 1 - a - -4     -6 - b - 1 - Y - -4
+        |        ↘               ↘      |
+        -5        -1              -1   -5
+
+        -3   -2              -2        -3
+        |      ↘               ↘        |
+    -6- X - 1 - a - -4     -6 - b - 1 - Y - -4
+        | ↘                             | ↘
+        -5 -1                          -5  -1
+```
 """
 function _qr_bond_undo(X::PEPSOrth, a::AbstractTensorMap, b::AbstractTensorMap, Y::PEPSOrth)
     @tensor A[-1; -2 -3 -4 -5] := X[-2 1 -4 -5] * a[1 -1 -3]
     @tensor B[-1; -2 -3 -4 -5] := b[-5 -1 1] * Y[-2 -3 -4 1]
     return A, B
 end
+function _qr_bond_undo(X::PEPOOrth, a::AbstractTensorMap, b::AbstractTensorMap, Y::PEPOOrth)
+    if !isdual(space(a, 2))
+        @tensor A[-1 -2; -3 -4 -5 -6] := X[-2 -3 1 -5 -6] * a[1 -1 -4]
+        @tensor B[-1 -2; -3 -4 -5 -6] := b[-6 -1 1] * Y[-2 -3 -4 -5 1]
+    else
+        @tensor A[-1 -2; -3 -4 -5 -6] := X[-1 -3 1 -5 -6] * a[1 -2 -4]
+        @tensor B[-1 -2; -3 -4 -5 -6] := b[-6 -2 1] * Y[-1 -3 -4 -5 1]
+    end
+    return A, B
+end
 
 """
 $(SIGNATURES)
 
 Apply 2-site `gate` on the reduced matrices `a`, `b`
 ```
-    -1← a --- 3 --- b ← -4
-        ↓           ↓
-        1           2
-        ↓           ↓
-        |----gate---|
-        ↓           ↓
-        -2         -3
+    -1← a --- 3 --- b ← -4          -2         -3
+        ↓           ↓               ↓           ↓
+        1           2               |----gate---|
+        ↓           ↓       or      ↓           ↓
+        |----gate---|               1           2
+        ↓           ↓               ↓           ↓
+        -2         -3           -1← a --- 3 --- b ← -4
 ```
 """
 function _apply_gate(
-        a::AbstractTensorMap{T, S},
-        b::AbstractTensorMap{T, S},
-        gate::AbstractTensorMap{T, S, 2, 2},
-        trscheme::TruncationScheme,
+        a::AbstractTensorMap{T, S}, b::AbstractTensorMap{T, S},
+        gate::AbstractTensorMap{T, S, 2, 2}, trscheme::TruncationScheme
     ) where {T <: Number, S <: ElementarySpace}
     V = space(b, 1)
     need_flip = isdual(V)
-    @tensor a2b2[-1 -2; -3 -4] := gate[-2 -3; 1 2] * a[-1 1 3] * b[3 2 -4]
+    if isdual(space(a, 2))
+        @tensor a2b2[-1 -2; -3 -4] := gate[1 2; -2 -3] * a[-1 1 3] * b[3 2 -4]
+    else
+        @tensor a2b2[-1 -2; -3 -4] := gate[-2 -3; 1 2] * a[-1 1 3] * b[3 2 -4]
+    end
     trunc = (trscheme isa FixedSpaceTruncation) ? truncspace(V) : trscheme
     a, s, b, ϵ = tsvd!(a2b2; trunc, alg = TensorKit.SVD())
     a, b = absorb_s(a, s, b)