Adapt simple update to DiagonalTensorMap (#124)

Author: Yue-Zhengyuan

examples/heisenberg_evol/heis_tools.jl

@@ -0,0 +1,58 @@
+using Test
+using Printf
+using Random
+import Statistics: mean
+using TensorKit
+using PEPSKit
+
+module MeasureHeis
+
+export measure_heis
+
+using TensorKit
+import MPSKitModels: S_x, S_y, S_z, S_exchange
+using PEPSKit
+
+"""
+Measure magnetization on each site
+"""
+function cal_mags(peps::InfinitePEPS, envs::CTMRGEnv)
+    N1, N2 = size(peps)
+    lattice = collect(space(t, 1) for t in peps.A)
+    # detect symmetry on physical axis
+    symm = sectortype(space(peps.A[1, 1]))
+    if symm == Trivial
+        Sas = real.([S_x(symm), im * S_y(symm), S_z(symm)])
+    elseif symm == U1Irrep
+        # only Sz preserves <Sz>
+        Sas = real.([S_z(symm)])
+    else
+        throw(ArgumentError("Unrecognized symmetry on physical axis"))
+    end
+    return [
+        collect(
+            expectation_value(
+                peps, LocalOperator(lattice, (CartesianIndex(r, c),) => Sa), envs
+            ) for (r, c) in Iterators.product(1:N1, 1:N2)
+        ) for Sa in Sas
+    ]
+end
+
+"""
+Measure physical quantities for Heisenberg model
+"""
+function measure_heis(peps::InfinitePEPS, H::LocalOperator, envs::CTMRGEnv)
+    results = Dict{String,Any}()
+    N1, N2 = size(peps)
+    results["e_site"] = costfun(peps, envs, H) / (N1 * N2)
+    results["mag"] = cal_mags(peps, envs)
+    results["mag_norm"] = collect(
+        norm([mags[r, c] for mags in results["mag"]]) for
+        (r, c) in Iterators.product(1:N1, 1:N2)
+    )
+    return results
+end
+
+end
+
+import .MeasureHeis: measure_heis

examples/heisenberg_evol/simpleupdate.jl

@@ -0,0 +1,52 @@
+include("heis_tools.jl")
+
+# benchmark data is from Phys. Rev. B 94, 035133 (2016)
+Dbond, χenv, symm = 4, 16, Trivial
+trscheme_peps = truncerr(1e-10) & truncdim(Dbond)
+trscheme_envs = truncerr(1e-10) & truncdim(χenv)
+N1, N2 = 2, 2
+# Heisenberg model Hamiltonian
+# (already only includes nearest neighbor terms)
+ham = heisenberg_XYZ(ComplexF64, symm, InfiniteSquare(N1, N2); Jx=1.0, Jy=1.0, Jz=1.0)
+# convert to real tensors
+ham = LocalOperator(ham.lattice, Tuple(ind => real(op) for (ind, op) in ham.terms)...)
+
+# random initialization of 2x2 iPEPS with weights and CTMRGEnv (using real numbers)
+if symm == Trivial
+    Pspace = ℂ^2
+    Vspace = ℂ^Dbond
+    Espace = ℂ^χenv
+elseif symm == U1Irrep
+    Pspace = ℂ[U1Irrep](1//2 => 1, -1//2 => 1)
+    Vspace = ℂ[U1Irrep](0 => Dbond ÷ 2, 1//2 => Dbond ÷ 4, -1//2 => Dbond ÷ 4)
+    Espace = ℂ[U1Irrep](0 => χenv ÷ 2, 1//2 => χenv ÷ 4, -1//2 => χenv ÷ 4)
+else
+    error("Not implemented")
+end
+Random.seed!(0)
+peps = InfiniteWeightPEPS(rand, Float64, Pspace, Vspace; unitcell=(N1, N2))
+# normalize vertex tensors
+for ind in CartesianIndices(peps.vertices)
+    peps.vertices[ind] /= norm(peps.vertices[ind], Inf)
+end
+
+# simple update
+dts = [1e-2, 1e-3, 4e-4]
+tols = [1e-6, 1e-8, 1e-8]
+maxiter = 10000
+for (dt, tol) in zip(dts, tols)
+    alg = SimpleUpdate(dt, tol, maxiter, trscheme_peps)
+    result = simpleupdate(peps, ham, alg; bipartite=true)
+    global peps = result[1]
+end
+# measure physical quantities with CTMRG
+peps_ = InfinitePEPS(peps)
+envs = CTMRGEnv(rand, Float64, peps_, Espace)
+ctm_alg = SequentialCTMRG(; tol=1e-10, verbosity=2, trscheme=trscheme_envs)
+envs = leading_boundary(envs, peps_, ctm_alg)
+meas = measure_heis(peps_, ham, envs)
+display(meas)
+@info @sprintf("Energy = %.8f\n", meas["e_site"])
+@info @sprintf("Staggered magnetization = %.8f\n", mean(meas["mag_norm"]))
+@test isapprox(meas["e_site"], -0.6675; atol=1e-3)
+@test isapprox(mean(meas["mag_norm"]), 0.3767; atol=1e-3)

src/PEPSKit.jl

@@ -49,7 +49,7 @@ include("algorithms/ctmrg/simultaneous.jl")
 include("algorithms/ctmrg/sequential.jl")
 include("algorithms/ctmrg/gaugefix.jl")
 
-include("algorithms/time_evolution/gatetools.jl")
+include("algorithms/time_evolution/evoltools.jl")
 include("algorithms/time_evolution/simpleupdate.jl")
 
 include("algorithms/toolbox.jl")

src/algorithms/ctmrg/sparse_environments.jl

@@ -229,7 +229,7 @@ function TensorKit.codomain(env::HalfInfiniteEnv)
 end
 
 function random_start_vector(env::HalfInfiniteEnv)
-    return Tensor(randn, domain(env))
+    return randn(domain(env))
 end
 
 # --------------------------------
@@ -377,5 +377,5 @@ function TensorKit.codomain(env::FullInfiniteEnv)
 end
 
 function random_start_vector(env::FullInfiniteEnv)
-    return Tensor(randn, domain(env))
+    return randn(domain(env))
 end

…c/algorithms/time_evolution/gatetools.jl …c/algorithms/time_evolution/evoltools.jlsrc/algorithms/time_evolution/gatetools.jl renamed to src/algorithms/time_evolution/evoltools.jl src/algorithms/time_evolution/gatetools.jl renamed to src/algorithms/time_evolution/evoltools.jl

@@ -49,3 +49,20 @@ function get_gateterm(gate::LocalOperator, bond::NTuple{2,CartesianIndex{2}})
         return gate.terms[label[1]].second
     end
 end
+
+const _axlabels = Set("trbl")
+
+"""
+Absorb environment weights into tensor `t` in position `[row, col]` in an iPEPS with `weights`. 
+A full weight is absorbed into axes of `t` on the boundary (specified by `open_axs`).
+But on internal bonds, square root of weights are absorbed. 
+"""
+function _absorb_weight(
+    t::PEPSTensor, row::Int, col::Int, open_axs::String, weights::SUWeight
+)
+    @assert all(c -> c in _axlabels, open_axs)
+    for (ax, ch) in zip(2:5, "trbl")
+        t = absorb_weight(t, row, col, ax, weights; sqrtwt=!(ch in open_axs))
+    end
+    return t
+end

src/algorithms/time_evolution/simpleupdate.jl

@@ -42,18 +42,11 @@ function _su_bondx!(
 ) where {T<:Number,S<:ElementarySpace}
     Nr, Nc = size(peps)
     @assert 1 <= row <= Nr && 1 <= col <= Nc
-    row2, col2 = row, _next(col, Nc)
-    T1, T2 = peps.vertices[row, col], peps.vertices[row2, col2]
+    cp1 = _next(col, Nc)
     # absorb environment weights
-    for ax in (2, 4, 5)
-        T1 = absorb_weight(T1, row, col, ax, peps.weights)
-    end
-    for ax in (2, 3, 4)
-        T2 = absorb_weight(T2, row2, col2, ax, peps.weights)
-    end
-    # absorb bond weight
-    T1 = absorb_weight(T1, row, col, 3, peps.weights; sqrtwt=true)
-    T2 = absorb_weight(T2, row2, col2, 5, peps.weights; sqrtwt=true)
+    T1, T2 = peps.vertices[row, col], peps.vertices[row, cp1]
+    T1 = _absorb_weight(T1, row, col, "tbl", peps.weights)
+    T2 = _absorb_weight(T2, row, cp1, "trb", peps.weights)
     #= QR and LQ decomposition
 
         2   1               1             2
@@ -72,17 +65,18 @@ function _su_bondx!(
     bL, Y = rightorth(T2, ((5, 1), (2, 3, 4)); alg=LQpos())
     #= apply gate
 
-            -2          -3
+            -2         -3
             ↑           ↑
             |----gate---|
             ↑           ↑
             1           2
             ↑           ↑
         -1← aR -← 3 -← bL ← -4
     =#
-    @tensor tmp[-1 -2; -3 -4] := gate[-2, -3, 1, 2] * aR[-1, 1, 3] * bL[3, 2, -4]
-    # SVD
-    aR, s, bL, ϵ = tsvd!(tmp; trunc=truncation_scheme(alg, space(T1, 3)))
+    @tensor tmp[-1 -2; -3 -4] := gate[-2 -3; 1 2] * aR[-1 1 3] * bL[3 2 -4]
+    aR, s, bL, ϵ = tsvd!(
+        tmp; trunc=truncation_scheme(alg, space(T1, 3)), alg=TensorKit.SVD()
+    )
     #=
             -2         -1              -1    -2
             |         ↗               ↗       |
@@ -97,11 +91,11 @@ function _su_bondx!(
         T1 = absorb_weight(T1, row, col, ax, peps.weights; invwt=true)
     end
     for ax in (2, 3, 4)
-        T2 = absorb_weight(T2, row2, col2, ax, peps.weights; invwt=true)
+        T2 = absorb_weight(T2, row, cp1, ax, peps.weights; invwt=true)
     end
     # update tensor dict and weight on current bond 
     # (max element of weight is normalized to 1)
-    peps.vertices[row, col], peps.vertices[row2, col2] = T1, T2
+    peps.vertices[row, col], peps.vertices[row, cp1] = T1, T2
     peps.weights[1, row, col] = s / norm(s, Inf)
     return ϵ
 end
@@ -181,17 +175,17 @@ function simpleupdate(
     check_int::Int=500,
 )
     time_start = time()
-    N1, N2 = size(peps)
+    Nr, Nc = size(peps)
     if bipartite
-        @assert N1 == N2 == 2
+        @assert Nr == Nc == 2
     end
     # exponentiating the 2-site Hamiltonian gate
     gate = get_gate(alg.dt, ham)
     wtdiff = 1.0
     wts0 = deepcopy(peps.weights)
     for count in 1:(alg.maxiter)
         time0 = time()
-        peps = su_iter(gate, peps, alg; bipartite=bipartite)
+        peps = su_iter(gate, peps, alg; bipartite)
         wtdiff = compare_weights(peps.weights, wts0)
         converge = (wtdiff < alg.tol)
         cancel = (count == alg.maxiter)
@@ -210,9 +204,7 @@ function simpleupdate(
             )
             cancel ? (@warn message) : (@info message)
         end
-        if converge || cancel
-            break
-        end
+        converge && break
     end
     return peps, wtdiff
 end

src/states/infiniteweightpeps.jl

@@ -1,11 +1,11 @@
 
 """
-    const PEPSWeight{S}
+    const PEPSWeight
 
 Default type for PEPS bond weights with 2 virtual indices, conventionally ordered as: ``wt : WS ← EN``. 
 `WS`, `EN` denote the west/south, east/north spaces for x/y-weights on the square lattice, respectively.
 """
-const PEPSWeight{T,S} = AbstractTensorMap{T,S,1,1} where {T<:Number,S<:ElementarySpace}
+const PEPSWeight{T,S} = AbstractTensorMap{T,S,1,1}
 
 """
     struct SUWeight{E<:PEPSWeight}
@@ -108,9 +108,10 @@ function InfiniteWeightPEPS(
 ) where {S<:ElementarySpace}
     vertices = InfinitePEPS(f, T, Pspace, Nspace, Espace; unitcell=unitcell).A
     Nr, Nc = unitcell
-    weights = collect(
-        id(d == 1 ? Espace : Nspace) for (d, r, c) in Iterators.product(1:2, 1:Nr, 1:Nc)
-    )
+    weights = map(Iterators.product(1:2, 1:Nr, 1:Nc)) do (d, r, c)
+        V = (d == 1 ? Espace : Nspace)
+        DiagonalTensorMap(ones(reduceddim(V)), V)
+    end
     return InfiniteWeightPEPS(vertices, SUWeight(weights))
 end
 

src/utility/util.jl

@@ -45,53 +45,63 @@ function _elementwise_mult(a₁::AbstractTensorMap, a₂::AbstractTensorMap)
     return dst
 end
 
-_safe_pow(a, pow, tol) = (pow < 0 && abs(a) < tol) ? zero(a) : a^pow
+_safe_pow(a::Real, pow::Real, tol::Real) = (pow < 0 && abs(a) < tol) ? zero(a) : a^pow
 
 """
-    sdiag_pow(S::AbstractTensorMap, pow::Real; tol::Real=eps(scalartype(S))^(3 / 4))
+    sdiag_pow(s, pow::Real; tol::Real=eps(scalartype(s))^(3 / 4))
 
-Compute `S^pow` for diagonal matrices `S`.
+Compute `s^pow` for a diagonal matrix `s`.
 """
-function sdiag_pow(S::AbstractTensorMap, pow::Real; tol::Real=eps(scalartype(S))^(3 / 4))
-    tol *= norm(S, Inf)  # Relative tol w.r.t. largest singular value (use norm(∘, Inf) to make differentiable)
-    Spow = similar(S)
-    for (k, b) in blocks(S)
+function sdiag_pow(s::DiagonalTensorMap, pow::Real; tol::Real=eps(scalartype(s))^(3 / 4))
+    # Relative tol w.r.t. largest singular value (use norm(∘, Inf) to make differentiable)
+    tol *= norm(s, Inf)
+    spow = DiagonalTensorMap(_safe_pow.(s.data, pow, tol), space(s, 1))
+    return spow
+end
+function sdiag_pow(
+    s::AbstractTensorMap{T,S,1,1}, pow::Real; tol::Real=eps(scalartype(s))^(3 / 4)
+) where {T,S}
+    # Relative tol w.r.t. largest singular value (use norm(∘, Inf) to make differentiable)
+    tol *= norm(s, Inf)
+    spow = similar(s)
+    for (k, b) in blocks(s)
         copyto!(
-            block(Spow, k), LinearAlgebra.diagm(_safe_pow.(LinearAlgebra.diag(b), pow, tol))
+            block(spow, k), LinearAlgebra.diagm(_safe_pow.(LinearAlgebra.diag(b), pow, tol))
         )
     end
-    return Spow
-end
-
-"""
-    absorb_s(u::AbstractTensorMap, s::AbstractTensorMap, vh::AbstractTensorMap)
-
-Given `tsvd` result `u`, `s` and `vh`, absorb singular values `s` into `u` and `vh` by:
-```
-    u -> u * sqrt(s), vh -> sqrt(s) * vh
-```
-"""
-function absorb_s(u::AbstractTensorMap, s::AbstractTensorMap, vh::AbstractTensorMap)
-    sqrt_s = sdiag_pow(s, 0.5)
-    return u * sqrt_s, sqrt_s * vh
+    return spow
 end
 
 function ChainRulesCore.rrule(
     ::typeof(sdiag_pow),
-    S::AbstractTensorMap,
+    s::AbstractTensorMap,
     pow::Real;
-    tol::Real=eps(scalartype(S))^(3 / 4),
+    tol::Real=eps(scalartype(s))^(3 / 4),
 )
-    tol *= norm(S, Inf)
-    spow = sdiag_pow(S, pow; tol)
-    spow_minus1_conj = scale!(sdiag_pow(S', pow - 1; tol), pow)
+    tol *= norm(s, Inf)
+    spow = sdiag_pow(s, pow; tol)
+    spow_minus1_conj = scale!(sdiag_pow(s', pow - 1; tol), pow)
     function sdiag_pow_pullback(c̄_)
         c̄ = unthunk(c̄_)
         return (ChainRulesCore.NoTangent(), _elementwise_mult(c̄, spow_minus1_conj))
     end
     return spow, sdiag_pow_pullback
 end
 
+"""
+    absorb_s(u::AbstractTensorMap, s::DiagonalTensorMap, vh::AbstractTensorMap)
+
+Given `tsvd` result `u`, `s` and `vh`, absorb singular values `s` into `u` and `vh` by:
+```
+    u -> u * sqrt(s), vh -> sqrt(s) * vh
+```
+"""
+function absorb_s(u::AbstractTensorMap, s::DiagonalTensorMap, vh::AbstractTensorMap)
+    @assert !isdual(space(s, 1))
+    sqrt_s = sdiag_pow(s, 0.5)
+    return u * sqrt_s, sqrt_s * vh
+end
+
 # Check whether diagonals contain degenerate values up to absolute or relative tolerance
 function is_degenerate_spectrum(
     S; atol::Real=0, rtol::Real=atol > 0 ? 0 : sqrt(eps(scalartype(S)))