Add Neighborhood Tensor Update (NTU)

src/PEPSKit.jl

@@ -3,7 +3,7 @@ module PEPSKit
 using LinearAlgebra, Statistics, Base.Threads, Base.Iterators, Printf
 using Random
 using Compat
-using Accessors: @set, @reset
+using Accessors: @set, @reset, @insert
 using VectorInterface
 import VectorInterface as VI
 
@@ -89,6 +89,7 @@ include("algorithms/contractions/bp_contractions.jl")
 include("algorithms/contractions/bondenv/benv_tools.jl")
 include("algorithms/contractions/bondenv/gaugefix.jl")
 include("algorithms/contractions/bondenv/als_solve.jl")
+include("algorithms/contractions/bondenv/benv_ntu.jl")
 include("algorithms/contractions/bondenv/benv_ctm.jl")
 include("algorithms/contractions/correlator/peps.jl")
 include("algorithms/contractions/correlator/pepo_1layer.jl")
@@ -103,14 +104,17 @@ include("algorithms/ctmrg/c4v.jl")
 
 include("algorithms/truncation/truncationschemes.jl")
 include("algorithms/truncation/fullenv_truncation.jl")
+include("algorithms/truncation/bond_tensor.jl")
 include("algorithms/truncation/bond_truncation.jl")
 
-include("algorithms/time_evolution/trotter_gate.jl")
 include("algorithms/time_evolution/apply_gate.jl")
 include("algorithms/time_evolution/apply_mpo.jl")
+include("algorithms/time_evolution/trotter_gate.jl")
 include("algorithms/time_evolution/time_evolve.jl")
 include("algorithms/time_evolution/simpleupdate.jl")
 include("algorithms/time_evolution/simpleupdate3site.jl")
+include("algorithms/time_evolution/ntupdate.jl")
+include("algorithms/time_evolution/ntupdate3site.jl")
 include("algorithms/time_evolution/gaugefix_su.jl")
 
 include("algorithms/bp/beliefpropagation.jl")
@@ -143,14 +147,16 @@ export fixedpoint
 
 export absorb_weight
 export ALSTruncation, FullEnvTruncation
-export SimpleUpdate
+export NNEnv, NNpEnv, NNNEnv
+export SimpleUpdate, NeighbourUpdate
 export TimeEvolver, timestep, time_evolve
 
 export InfiniteSquareNetwork
 export InfinitePartitionFunction
 export InfinitePEPS, InfiniteTransferPEPS
 export SUWeight
 export InfinitePEPO, InfiniteTransferPEPO
+export infinite_temperature_density_matrix
 
 export BPEnv, BeliefPropagation
 export BPGauge, SUGauge

src/algorithms/bp/beliefpropagation.jl

@@ -59,7 +59,7 @@ function leading_boundary(env₀::BPEnv, network::InfiniteSquareNetwork, alg::Be
 end
 function leading_boundary(env₀::BPEnv, state::InfiniteState, alg::BeliefPropagation)
     if alg.bipartite
-        @assert _state_bipartite_check(state)
+        _is_bipartite(state) || error("Input state is not bipartite")
     end
     return leading_boundary(env₀, InfiniteSquareNetwork(state), alg)
 end

src/algorithms/bp/gaugefix.jl

@@ -7,16 +7,6 @@ Algorithm for gauging PEPS with belief propagation fixed point messages.
     # TODO: add options
 end
 
-function _bpenv_bipartite_check(env::BPEnv)
-    for (r, c) in Iterators.product(1:2, 1:2)
-        r′, c′ = _next(r, 2), _next(c, 2)
-        if !all(env[:, r, c] .== env[:, r′, c′])
-            return false
-        end
-    end
-    return true
-end
-
 """
     gauge_fix(psi::Union{InfinitePEPS, InfinitePEPO}, alg::BPGauge, env::BPEnv)
 
@@ -25,7 +15,7 @@ an [`InfinitePEPO`](@ref) interpreted as purified state with two physical legs)
 using fixed point environment `env` of belief propagation.
 """
 function gauge_fix(psi::InfinitePEPS, alg::BPGauge, env::BPEnv)
-    bipartite = _state_bipartite_check(psi) && _bpenv_bipartite_check(env)
+    bipartite = _is_bipartite(psi) && _is_bipartite(env)
     psi′ = copy(psi)
     XXinv = map(eachcoordinate(psi, 1:2)) do I
         _, X, Xinv = _bp_gauge_fix!(CartesianIndex(I), psi′, env)

src/algorithms/contractions/bondenv/als_solve.jl

@@ -2,139 +2,148 @@
 In the following, the names `Ra`, `Sa` etc comes from 
 the fast full update article Physical Review B 92, 035142 (2015)
 =#
-
 """
-$(SIGNATURES)
+Contract the virtual legs between
+```
+    -- DX --a-- D --b-- DY --
+            ↓       ↓
+            da      db
+```
+"""
+function _combine_ket(a::MPSTensor, b::MPSTensor)
+    return @tensor ket[DX DY; da db] := a[DX da; D] * b[D db; DY]
+end
 
-Construct the tensor
-```
-    ┌-----------------------------------┐
-    |   ┌----┐                          |
-    └---|    |- DX0     Db0 - b -- DY0 -┘
-        |    |                ↓
-        |benv|                db
-        |    |                ↓
-    ┌---|    |- DX1     Db1 - b† - DY1 -┐
-    |   └----┘                          |
-    └-----------------------------------┘
-```
-"""
-function _tensor_Ra(benv::BondEnv, b::MPSTensor)
-    return @autoopt @tensor Ra[DX1 Db1; DX0 Db0] := (
-        benv[DX1 DY1; DX0 DY0] * b[Db0 db; DY0] * conj(b[Db1 db; DY1])
-    )
+function _combine_ket_for_svd(a::MPSTensor, b::MPSTensor)
+    return @tensor ket[DX da; db DY] := a[DX da; D] * b[D db; DY]
 end
 
 """
-$(SIGNATURES)
-
-Construct the tensor
+Construct the norm with bra bond tensors removed
 ```
-    ┌-----------------------------------┐
-    |   ┌----┐                          |
-    └---|    |- DX0 -- (a2 b2) -- DY0 --┘
-        |    |         ↓     ↓
-        |benv|         da    db
-        |    |               ↓
-    ┌---|    |- DX1   Db1 -- b† - DY1 --┐
-    |   └----┘                          |
-    └-----------------------------------┘
+    ┌benv-------┐
+    ├---a---b---┤
+    |   ↓   ↓   |
+    ├--       --┤
+    └-----------┘
 ```
 """
-function _tensor_Sa(
-        benv::BondEnv, b::MPSTensor, a2b2::AbstractTensorMap{T, S, 2, 2}
-    ) where {T <: Number, S <: ElementarySpace}
-    return @autoopt @tensor Sa[DX1 da; Db1] := (
-        benv[DX1 DY1; DX0 DY0] * conj(b[Db1 db; DY1]) * a2b2[DX0 DY0; da db]
-    )
+function _benv_ket(benv::BondEnv, ket::AbstractTensorMap{T, S, 2, 2}) where {T, S}
+    return benv * twistdual(ket, 1:2)
 end
 
 """
-$(SIGNATURES)
+    _als_tensor_R(benv::BondEnv, xs::Vector{<:MPSTensor}, i::Int)
 
-Construct the tensor
-```
-    ┌-----------------------------------┐
-    |   ┌----┐                          |
-    └---|    |- DX0 - a -- Da0     DY0 -┘
-        |    |        ↓
-        |benv|        da
-        |    |        ↓
-    ┌---|    |- DX1 - a† - Da1     DY1 -┐
-    |   └----┘                          |
-    └-----------------------------------┘
-```
-"""
-function _tensor_Rb(benv::BondEnv, a::MPSTensor)
-    return @autoopt @tensor Rb[Da1 DY1; Da0 DY0] := (
-        benv[DX1 DY1; DX0 DY0] * a[DX0 da; Da0] * conj(a[DX1 da; Da1])
-    )
+Construct the bond environment around the `i`th bond tensor
+in two-site ALS optimization.
+```
+    i = 1           i = 2
+    ┌benv-------┐   ┌benv-------┐
+    ├--   --b---┤   ├---a--   --┤
+    |       ↓   |   |   ↓       |
+    ├--   --b̄---┤   ├---ā--   --┤
+    └-----------┘   └-----------┘
+```
+"""
+_als_tensor_R(benv, xs, i::Int) = _als_tensor_R(benv, xs, Val(i))
+function _als_tensor_R(benv::BondEnv, xs::Vector{<:MPSTensor}, ::Val{1})
+    return @tensor Ra[DX1 D1; DX0 D0] :=
+        benv[DX1 DY1; DX0 DY0] * xs[2][D0 db; DY0] * conj(xs[2][D1 db; DY1])
+end
+function _als_tensor_R(benv::BondEnv, xs::Vector{<:MPSTensor}, ::Val{2})
+    return @tensor Rb[D1 DY1; D0 DY0] :=
+        benv[DX1 DY1; DX0 DY0] * xs[1][DX0 da; D0] * conj(xs[1][DX1 da; D1])
 end
 
 """
-$(SIGNATURES)
-
-Construct the tensor
+Calculate the 2-site norm
 ```
-    ┌-----------------------------------┐
-    |   ┌----┐                          |
-    └---|    |- DX0 -- (a2 b2) -- DY0 --┘
-        |    |         ↓     ↓
-        |benv|         da    db
-        |    |         ↓
-    ┌---|    |- DX1 -- a† - Da1   DY1 --┐
-    |   └----┘                          |
-    └-----------------------------------┘
+    ┌benv-------┐
+    ├---a---b---┤
+    |   ↓   ↓   |
+    ├---ā---b̄---┤
+    └-----------┘
 ```
+using pre-calcuated partial contraction results.
 """
-function _tensor_Sb(
-        benv::BondEnv, a::MPSTensor, a2b2::AbstractTensorMap{T, S, 2, 2}
-    ) where {T <: Number, S <: ElementarySpace}
-    return @autoopt @tensor Sb[Da1 db; DY1] := (
-        benv[DX1 DY1; DX0 DY0] * conj(a[DX1 da; Da1]) * a2b2[DX0 DY0; da db]
-    )
+function _als_norm(
+        ket::AbstractTensorMap{T, S, 2, 2}, benv_ket::AbstractTensorMap{T, S, 2, 2}
+    ) where {T, S}
+    return @tensor benv_ket[DX1 DY1; da db] * conj(ket[DX1 DY1; da db])
+end
+function _als_norm(a::MPSTensor, Ra::BondEnv)
+    return @tensor Ra[DX1 D1; DX0 D0] * a[DX0 da; D0] * conj(a[DX1 da; D1])
 end
 
 """
-$(SIGNATURES)
+    _als_tensor_S(
+        benv_ket::AbstractTensorMap{T, S, 2, 2},
+        xs::Vector{<:MPSTensor}, i::Int
+    ) where {T <: Number, S <: ElementarySpace}
 
-Calculate the inner product <a1,b1|a2,b2>
+Construct the overlap but with one of the bra bond tensor removed.
 ```
-    ┌--------------------------------┐
-    |   ┌----┐                       |
-    └---|    |- DX0 - (a2 b2) - DY0 -┘
-        |    |        ↓    ↓
-        |benv|        da   db
-        |    |        ↓    ↓
-    ┌---|    |- DX1 - (a1 b1)†- DY1 -┐
-    |   └----┘                       |
-    └--------------------------------┘
+    i = 1           i = 2
+    ┌benv-------┐   ┌benv-------┐
+    ├---a₂==b₂--┤   ├---a₂==b₂--┤
+    |   ↓   ↓   |   |   ↓   ↓   |
+    ├--   --b̄---┤   ├---ā--   --┤
+    └-----------┘   └-----------┘
 ```
+The ket part is provided by the partial contraction `benv_ket`.
 """
-function inner_prod(
-        benv::BondEnv, a1b1::AbstractTensorMap{T, S, 2, 2}, a2b2::AbstractTensorMap{T, S, 2, 2}
+_als_tensor_S(benv_ket, xs, i::Int) = _als_tensor_S(benv_ket, xs, Val(i))
+function _als_tensor_S(
+        benv_ket::AbstractTensorMap{T, S, 2, 2},
+        xs::Vector{<:MPSTensor}, ::Val{1}
     ) where {T <: Number, S <: ElementarySpace}
-    return @autoopt @tensor benv[DX1 DY1; DX0 DY0] *
-        conj(a1b1[DX1 DY1; da db]) * a2b2[DX0 DY0; da db]
+    return @tensor Sa[DX1 da; D1] :=
+        benv_ket[DX1 DY1; da db] * conj(xs[2][D1 db; DY1])
+end
+function _als_tensor_S(
+        benv_ket::AbstractTensorMap{T, S, 2, 2},
+        xs::Vector{<:MPSTensor}, ::Val{2}
+    ) where {T <: Number, S <: ElementarySpace}
+    return @tensor contractcheck = true Sb[D1 db; DY1] :=
+        benv_ket[DX1 DY1; da db] * conj(xs[1][DX1 da; D1])
 end
 
 """
-$(SIGNATURES)
+Calculate the inner product (overlap)
+```
+    ┌benv-------┐
+    ├---a₂--b₂--┤
+    |   ↓   ↓   |
+    ├---ā---b̄---┤
+    └-----------┘
+```
+using pre-calculated partial contraction results.
+"""
+function _als_overlap(a::MPSTensor, Sa::MPSTensor)
+    # applies to b, Sb as well
+    # @tensor Sb[D1 db; DY1] * conj(b[D1 db; DY1])
+    return @tensor Sa[DX1 da; D1] * conj(a[DX1 da; D1])
+end
 
-Contract the axis between `a` and `b` tensors
+"""
+Calculate the 2-site ALS inner product ⟨a₁,b₁|a₂,b₂⟩
 ```
-    -- DX - a - D - b - DY --
-            ↓       ↓
-            da      db
+    ┌benv-------┐
+    ├---a₂--b₂--┤
+    |   ↓   ↓   |
+    ├---ā₁--b̄₁--┤
+    └-----------┘
 ```
+where `|bra⟩ = |a₁,b₁⟩` and `|ket⟩ = |a₂,b₂⟩`,
+with virtual leg between a, b contracted.
 """
-function _combine_ab(
-        a::MPSTensor, b::AbstractTensorMap{T, S, 1, 2}
+function inner_prod(
+        benv::BondEnv, bra::AbstractTensorMap{T, S, 2, 2},
+        ket::AbstractTensorMap{T, S, 2, 2}
     ) where {T <: Number, S <: ElementarySpace}
-    return @tensor ab[DX DY; da db] := a[DX da; D] * b[D; db DY]
-end
-function _combine_ab(a::MPSTensor, b::MPSTensor)
-    return @tensor ab[DX DY; da db] := a[DX da; D] * b[D db; DY]
+    return @autoopt @tensor benv[DX1 DY1; DX0 DY0] *
+        conj(bra[DX1 DY1; da db]) * ket[DX0 DY0; da db]
 end
 
 """
@@ -161,10 +170,30 @@ function cost_function_als(benv, ψ1, ψ2)
     return cost, fid
 end
 
+# applies to Rb, Sb, b as well
+# b22 is the pre-calculated untruncated norm
+function cost_function_als(Ra::BondEnv, Sa::MPSTensor, a::MPSTensor, b22::Real)
+    b11 = real(_als_norm(a, Ra))
+    b12 = _als_overlap(a, Sa)
+    cost = b11 + b22 - 2 * real(b12)
+    fid = abs2(b12) / abs(b11 * b22)
+    return cost, fid
+end
+
 """
 $(SIGNATURES)
 
 Solve the equations `Rx x = Sx` with initial guess `x0`.
+
+In ALS over `a`, `b`, if we fix `b`, the cost function can
+be expressed in the `Ra`, `Sa` tensors as
+```
+    f(a†,a) = a† Ra a - a† Sa - Sa† a + const
+```
+Therefore `f` is minimized when
+```
+    ∂f/∂ā = Ra a - Sa = 0
+```
 """
 function _solve_als(
         Rx::AbstractTensorMap{T, S, N, N},

src/algorithms/contractions/bondenv/benv_ctm.jl

@@ -9,8 +9,8 @@ Construct the environment (norm) tensor
         -1   0          1    2
 ```
 with `X` at position `[row, col]`.
-`X, Y` are unitary tensors produced when finding the reduced site tensors 
-with `_qr_bond`; and `XX = X' X` and `YY = Y' Y` (stacked together).
+`X, Y` are unitary tensors produced when finding the reduced site tensors,
+and `XX = X' X` and `YY = Y' Y` (stacked together).
 
 Axis order: `[DX1 DY1; DX0 DY0]`, as in
 ```
@@ -24,8 +24,8 @@ Axis order: `[DX1 DY1; DX0 DY0]`, as in
 ```
 """
 function bondenv_ctm(
-        row::Int, col::Int, X::T, Y::T, env::CTMRGEnv
-    ) where {T}
+        row::Int, col::Int, X::TX, Y::TY, env::CTMRGEnv
+    ) where {TX, TY}
     Nr, Nc = size(env.corners)[[2, 3]]
     cm1 = _prev(col, Nc)
     cp1 = _next(col, Nc)