diff --git a/Project.toml b/Project.toml index 8028adf6..d5e3a930 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec" VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] MPSKit = "0.10" diff --git a/README.md b/README.md index 32f36fab..c6c11f91 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ PEPSKit.jl logo -# PepsKit.jl +# PEPSKit.jl [![docs][docs-dev-img]][docs-dev-url] ![CI][ci-url] [![codecov][codecov-img]][codecov-url] diff --git a/examples/heisenberg.jl b/examples/heisenberg.jl index 7e8581d7..ad31abbb 100644 --- a/examples/heisenberg.jl +++ b/examples/heisenberg.jl @@ -1,140 +1,47 @@ -using Revise, PEPSKit, TensorKit, Zygote, MPSKit -using MPSKitModels, LinearAlgebra, OptimKit -using PEPSKit: - NORTH, SOUTH, WEST, EAST, NORTHWEST, NORTHEAST, SOUTHEAST, SOUTHWEST, @diffset -using JLD2, ChainRulesCore +using LinearAlgebra +using TensorKit, MPSKitModels, OptimKit +using PEPSKit -function two_site_rho(r::Int, c::Int, ψ::InfinitePEPS, env::PEPSKit.CTMRGEnv) - cp = mod1(c + 1, size(ψ, 2)) - @tensor ρ[-11, -20; -12, -18] := - env.corners[NORTHWEST, r, c][1, 3] * - env.edges[WEST, r, c][2, 7, 9, 1] * - env.corners[SOUTHWEST, r, c][4, 2] * - env.edges[NORTH, r, c][3, 5, 8, 13] * - env.edges[SOUTH, r, c][14, 6, 10, 4] * - ψ[r, c][-12, 5, 15, 6, 7] * - conj(ψ[r, c][-11, 8, 19, 10, 9]) * - env.edges[NORTH, r, cp][13, 16, 22, 23] * - env.edges[SOUTH, r, cp][28, 17, 21, 14] * - ψ[r, cp][-18, 16, 25, 17, 15] * - conj(ψ[r, cp][-20, 22, 26, 21, 19]) * - env.corners[NORTHEAST, r, cp][23, 24] * - env.edges[EAST, r, cp][24, 25, 26, 27] * - env.corners[SOUTHEAST, r, cp][27, 28] - return ρ -end - -function iCtmGsEh( - ψ::InfinitePEPS, env::PEPSKit.CTMRGEnv, H::AbstractTensorMap{S,2,2} -) where {S} - #Es = Matrix{eltype(H)}(undef,size(ψ,1),size(ψ,2)) - E = 0.0 - for r in 1:size(ψ, 1), c in 1:size(ψ, 2) - ρ = two_site_rho(r, c, ψ, env) - @tensor nn = ρ[1 2; 1 2] - @tensor Eh = H[1 2; 3 4] * ρ[1 2; 3 4] - Eh = Eh / nn - E = E + Eh - #@diffset Es[r,c] = Eh; - end - return real(E) -end - -function H_expectation_value( - ψ::InfinitePEPS, env::PEPSKit.CTMRGEnv, H::AbstractTensorMap{S,2,2} -) where {S} - Eh = iCtmGsEh(ψ, env, H) - ψ1 = rotl90(ψ) - env1 = PEPSKit.rotate_north(env, EAST) - Ev = iCtmGsEh(ψ1, env1, H) - E = real(Eh + Ev) - return E -end - -function SqLatHeisenberg() +# Square lattice Heisenberg Hamiltonian +# Sublattice-rotate to get (1, 1, 1) → (-1, 1, -1), transformed to GS with single-site unit cell +function square_lattice_heisenberg(; Jx=-1.0, Jy=1.0, Jz=-1.0) Sx, Sy, Sz, _ = spinmatrices(1//2) - Dphys = ComplexSpace(2) - σx = TensorMap(Sx, Dphys, Dphys) - σy = TensorMap(Sy, Dphys, Dphys) - σz = TensorMap(Sz, Dphys, Dphys) + Vphys = ℂ^2 + σx = TensorMap(Sx, Vphys, Vphys) + σy = TensorMap(Sy, Vphys, Vphys) + σz = TensorMap(Sz, Vphys, Vphys) @tensor H[-1 -3; -2 -4] := - -σx[-1, -2] * σx[-3, -4] + σy[-1, -2] * σy[-3, -4] + -σz[-1, -2] * σz[-3, -4] + Jx * σx[-1, -2] * σx[-3, -4] + + Jy * σy[-1, -2] * σy[-3, -4] + + Jz * σz[-1, -2] * σz[-3, -4] return H end -H = SqLatHeisenberg() - -function cfun(x) - (ψ, env) = x - - function fun(peps) - env = leading_boundary(peps, alg_ctm, env) - x = H_expectation_value(peps, env, H) - return x - end - env = leading_boundary(ψ, alg_ctm, env) - E = H_expectation_value(ψ, env, H) - ∂E = fun'(ψ) - - @assert !isnan(norm(∂E)) - return E, ∂E -end - -# my_retract is not an in place function which should not change x -function my_retract(x, dx, α::Number) - (ϕ, env0) = x - ψ = deepcopy(ϕ) - env = deepcopy(env0) - ψ.A .+= dx.A .* α - #env = leading_boundary(ψ, alg_ctm,env) - return (ψ, env), dx -end - -my_inner(x, dx1, dx2) = real(dot(dx1, dx2)) - -function my_add!(Y, X, a) - Y.A .+= X.A .* a - return Y -end - -function my_scale!(η, β) - rmul!(η.A, β) - return η -end - -function init_psi(d::Int, D::Int, Lx::Int, Ly::Int) +# Initialize InfinitePEPS with random & complex entries by default +function init_peps(d, D, Lx, Ly; finit=randn, dtype=ComplexF64) Pspaces = fill(ℂ^d, Lx, Ly) Nspaces = fill(ℂ^D, Lx, Ly) Espaces = fill(ℂ^D, Lx, Ly) - - Sspaces = adjoint.(circshift(Nspaces, (1, 0))) - Wspaces = adjoint.(circshift(Espaces, (0, -1))) - - A = map(Pspaces, Nspaces, Espaces, Sspaces, Wspaces) do P, N, E, S, W - return TensorMap(rand, ComplexF64, P ← N * E * S * W) - end - - return InfinitePEPS(A) -end - -alg_ctm = CTMRG(; verbose=1, tol=1e-4, trscheme=truncdim(10), miniter=4, maxiter=200) - -function main(; d=2, D=2, Lx=1, Ly=1) - ψ = init_psi(d, D, Lx, Ly) - env = leading_boundary(ψ, alg_ctm) - optimize( - cfun, - (ψ, env), - ConjugateGradient(; verbosity=2); - inner=my_inner, - retract=my_retract, - (scale!)=my_scale!, - (add!)=my_add!, - ) - return ψ -end - -main() + return InfinitePEPS(Pspaces, Nspaces, Espaces; finit, dtype) +end + +# Parameters +H = square_lattice_heisenberg() +χbond = 2 +χenv = 20 +ctmalg = CTMRG(; trscheme=truncdim(χenv), tol=1e-10, miniter=4, maxiter=100, verbose=2) +optalg = PEPSOptimize{NaiveAD}(; + optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2), + fpgrad_tol=1e-6, + fpgrad_maxiter=100, + verbose=2, +) + +# Ground state search +ψinit = init_peps(2, χbond, 1, 1) +envinit = leading_boundary(ψinit, ctmalg, CTMRGEnv(ψinit; χenv)) +result = groundsearch(H, ctmalg, optalg, ψinit, envinit) +@show result.E₀ diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index b2dcebbe..d70a1ea8 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -1,17 +1,13 @@ module PEPSKit +using LinearAlgebra, Base.Threads, Base.Iterators, Printf using Accessors using VectorInterface -using TensorKit, - KrylovKit, MPSKit, OptimKit, Base.Threads, Base.Iterators, Parameters, Printf -using ChainRulesCore - -using LinearAlgebra: LinearAlgebra - -export CTMRG, CTMRG2 -export leading_boundary +using TensorKit, KrylovKit, MPSKit, OptimKit +using ChainRulesCore, Zygote include("utility/util.jl") +include("utility/rotations.jl") include("states/abstractpeps.jl") include("states/infinitepeps.jl") @@ -28,19 +24,24 @@ include("environments/ctmrgenv.jl") include("environments/boundarympsenv.jl") include("algorithms/ctmrg.jl") -include("algorithms/expval.jl") +include("algorithms/peps_opt.jl") include("utility/symmetrization.jl") include("algorithms/pepo_opt.jl") -include("utility/rotations.jl") - -#default settings +# Default settings module Defaults - const maxiter = 100 - const tol = 1e-12 + const ctmrg_maxiter = 100 + const ctmrg_miniter = 4 + const ctmrg_tol = 1e-12 + const grad_maxiter = 100 + const grad_tol = 4 end +export CTMRG, CTMRGEnv +export leading_boundary, expectation_value +export PEPSOptimize, NaiveAD, ManualIter, LinSolve +export groundsearch export InfinitePEPS, InfiniteTransferPEPS export InfinitePEPO, InfiniteTransferPEPO export initializeMPS, initializePEPS diff --git a/src/algorithms/ctmrg.jl b/src/algorithms/ctmrg.jl index 932ed866..4ee06ac3 100644 --- a/src/algorithms/ctmrg.jl +++ b/src/algorithms/ctmrg.jl @@ -1,321 +1,265 @@ -@with_kw struct CTMRG #<: Algorithm +@kwdef struct CTMRG #<: Algorithm trscheme::TruncationScheme = TensorKit.notrunc() - tol::Float64 = Defaults.tol - maxiter::Integer = Defaults.maxiter - miniter::Integer = 4 - verbose::Integer = 0 + tol::Float64 = Defaults.ctmrg_tol + maxiter::Int = Defaults.ctmrg_maxiter + miniter::Int = Defaults.ctmrg_miniter + verbose::Int = 0 fixedspace::Bool = false end -@with_kw struct CTMRG2 #<: Algorithm - trscheme::TruncationScheme = TensorKit.notrunc() - tol::Float64 = Defaults.tol - maxiter::Integer = Defaults.maxiter - miniter::Integer = 4 - verbose::Integer = 0 -end - -function MPSKit.leading_boundary(peps::InfinitePEPS, alg::CTMRG, envs=CTMRGEnv(peps)) - return MPSKit.leading_boundary(peps, peps, alg, envs) -end; - -function MPSKit.leading_boundary( - peps_above::InfinitePEPS, - peps_below::InfinitePEPS, - alg::CTMRG, - envs=CTMRGEnv(peps_above, peps_below), -) - err = Inf - iter = 1 - - #for convergence criterium we use the on site contracted boundary - #this convergences, though the value depends on the bond dimension χ - old_norm = 1.0 - new_norm = old_norm - ϵ₁ = 1.0 - while (err > alg.tol && iter <= alg.maxiter) || iter <= alg.miniter - ϵ = 0.0 - for i in 1:4 - envs, ϵ₀ = left_move(peps_above, peps_below, alg, envs) - ϵ = max(ϵ, ϵ₀) - envs = rotate_north(envs, EAST) - peps_above = envs.peps_above - peps_below = envs.peps_below +# Compute CTMRG environment for a given state +function MPSKit.leading_boundary(state, alg::CTMRG, envinit=CTMRGEnv(state)) + normold = 1.0 + CSold = tsvd(envinit.corners[NORTHWEST]; alg=TensorKit.SVD())[2] + TSold = tsvd(envinit.edges[NORTH]; alg=TensorKit.SVD())[2] + ϵold = 1.0 + + env = deepcopy(envinit) + for i in 1:(alg.maxiter) + env, iterinfo = ctmrg_iter(state, env, alg) # Grow and renormalize in all 4 directions + + # Compute convergence criteria and take max (TODO: How should we handle logging all of this?) + Δϵ = abs((ϵold - iterinfo.ϵ) / ϵold) + normnew = contract_ctmrg(state, env) + Δnorm = abs(normold - normnew) + CSnew = tsvd(env.corners[NORTHWEST]; alg=TensorKit.SVD())[2] + ΔCS = norm(CSnew - CSold) + TSnew = tsvd(env.edges[NORTH]; alg=TensorKit.SVD())[2] + ΔTS = norm(TSnew - TSold) + (max(Δnorm, ΔCS, ΔTS) < alg.tol && i > alg.miniter) && break # Converge if maximal Δ falls below tolerance + + # Print verbose info + ignore_derivatives() do + alg.verbose > 1 && @printf( + "CTMRG iter: %3d norm: %.2e Δnorm: %.2e ΔCS: %.2e ΔTS: %.2e ϵ: %.2e Δϵ: %.2e\n", + i, + abs(normnew), + Δnorm, + ΔCS, + ΔTS, + iterinfo.ϵ, + Δϵ + ) + alg.verbose > 0 && + i == alg.maxiter && + @warn( + "CTMRG reached maximal number of iterations at (Δnorm=$Δnorm, ΔCS=$ΔCS, ΔTS=$ΔTS)" + ) end - new_norm = contract_ctrmg(envs) - - err = abs(old_norm - new_norm) - dϵ = abs((ϵ₁ - ϵ) / ϵ₁) - @ignore_derivatives alg.verbose > 1 && @printf( - "CTMRG: \titeration: %4d\t\terror: %.2e\t\tnorm: %.10e\t\tϵ: %.2e\t\tdϵ: %.2e\n", - iter, - err, - abs(new_norm), - ϵ, - dϵ - ) - - old_norm = new_norm - ϵ₁ = ϵ - iter += 1 + # Update convergence criteria + normold = normnew + CSold = CSnew + TSold = TSnew + ϵold = iterinfo.ϵ end - #@ignore_derivatives @show iter, new_norm, err - @ignore_derivatives iter > alg.maxiter && - alg.verbose > 0 && - @warn "maxiter $(alg.maxiter) reached: error was $(err)" - - return envs + return env end -function left_move( - peps_above::InfinitePEPS{PType}, - peps_below::InfinitePEPS{PType}, - alg::CTMRG, - envs::CTMRGEnv, -) where {PType} - corners::typeof(envs.corners) = copy(envs.corners) - edges::typeof(envs.edges) = copy(envs.edges) - - above_projector_type = tensormaptype(spacetype(PType), 1, 3, storagetype(PType)) - below_projector_type = tensormaptype(spacetype(PType), 3, 1, storagetype(PType)) +# One CTMRG iteration x′ = f(A, x) +function ctmrg_iter(state, env::CTMRGEnv, alg::CTMRG) ϵ = 0.0 - n0 = 1.0 - n1 = 1.0 - for col in 1:size(peps_above, 2) - cop = mod1(col + 1, size(peps_above, 2)) - com = mod1(col - 1, size(peps_above, 2)) + Pleft = Vector{Matrix{typeof(env.edges[1])}}(undef, 4) + Pright = Vector{Matrix{typeof(transpose(env.edges[1]))}}(undef, 4) + + for i in 1:4 + env, info = left_move(state, env, alg) + state = rotate_north(state, EAST) + env = rotate_north(env, EAST) + ϵ = max(ϵ, info.ϵ) + @diffset Pleft[i] = info.Pleft + @diffset Pright[i] = info.Pright + end - above_projs = Vector{above_projector_type}(undef, size(peps_above, 1)) - below_projs = Vector{below_projector_type}(undef, size(peps_above, 1)) + iterinfo = (; ϵ, Pleft=copy(Pleft), Pright=copy(Pright)) + return env, iterinfo +end - # find all projectors - for row in 1:size(peps_above, 1) - rop = mod1(row + 1, size(peps_above, 1)) - peps_above_nw = peps_above[row, col] - peps_above_sw = rotate_north(peps_above[rop, col], WEST) - peps_below_nw = peps_below[row, col] - peps_below_sw = rotate_north(peps_below[rop, col], WEST) +# Grow environment, compute projectors and renormalize +function left_move(state, env::CTMRGEnv, alg::CTMRG) + corners::typeof(env.corners) = copy(env.corners) + edges::typeof(env.edges) = copy(env.edges) + ϵ = 0.0 - Q1 = northwest_corner( - envs.edges[SOUTH, mod1(row + 1, end), col], - envs.corners[SOUTHWEST, mod1(row + 1, end), col], - envs.edges[WEST, mod1(row + 1, end), col], - peps_above_sw, - peps_below_sw, + Pleft = similar(state.A, typeof(env.edges[1])) + Pright = similar(state.A, typeof(transpose(env.edges[1]))) + for col in 1:size(state, 2) + cnext = _next(col, size(state, 2)) + + # Compute projectors + for row in 1:size(state, 1) + rnext = _next(row, size(state, 1)) + state_nw = state[row, col] + state_sw = rotate_north(state[rnext, col], WEST) + + # Enlarged corners + Q_sw = northwest_corner( + env.edges[SOUTH, _next(row, end), col], + env.corners[SOUTHWEST, _next(row, end), col], + env.edges[WEST, _next(row, end), col], + state_sw, ) - Q2 = northwest_corner( - envs.edges[WEST, row, col], - envs.corners[NORTHWEST, row, col], - envs.edges[NORTH, row, col], - peps_above_nw, - peps_below_nw, + Q_nw = northwest_corner( + env.edges[WEST, row, col], + env.corners[NORTHWEST, row, col], + env.edges[NORTH, row, col], + state_nw, ) + # SVD half-infinite environment trscheme = if alg.fixedspace == true - truncspace(space(envs.edges[WEST, row, cop], 1)) + truncspace(space(env.edges[WEST, row, cnext], 1)) else alg.trscheme end - #@ignore_derivatives @show norm(Q1*Q2) - - (U, S, V) = tsvd(Q1 * Q2; trunc=trscheme, alg=SVD()) - - @ignore_derivatives n0 = norm(Q1 * Q2)^2 - @ignore_derivatives n1 = norm(U * S * V)^2 - @ignore_derivatives ϵ = max(ϵ, (n0 - n1) / n0) - - isqS = sdiag_inv_sqrt(S) - #Q = isqS*U'*Q1; - #P = Q2*V'*isqS; - @tensor Q[-1; -2 -3 -4] := isqS[-1; 1] * conj(U[2 3 4; 1]) * Q1[2 3 4; -2 -3 -4] - @tensor P[-1 -2 -3; -4] := Q2[-1 -2 -3; 1 2 3] * conj(V[4; 1 2 3]) * isqS[4; -4] + (U, S, V) = tsvd(Q_sw * Q_nw; trunc=trscheme, alg=TensorKit.SVD()) # TODO: Add field in CTMRG to choose SVD function + + # Compute SVD truncation error and check for degenerate singular values + ignore_derivatives() do + if unique(x -> round(x; digits=14), diag(S.data)) != diag(S.data) && + alg.verbose > 0 + println("degenerate singular values detected", diag(S.data)) + end + n0 = norm(Q_sw * Q_nw)^2 + n1 = norm(U * S * V)^2 + ϵ = max(ϵ, (n0 - n1) / n0) + end - @diffset above_projs[row] = Q - @diffset below_projs[row] = P + # Compute projectors + Pl, Pr = build_projectors(U, S, V, Q_sw, Q_nw) + @diffset Pleft[row, col] = Pl + @diffset Pright[row, col] = Pr end - #use the projectors to grow the corners/edges - for row in 1:size(peps_above, 1) - Q = above_projs[row] - P = below_projs[mod1(row - 1, end)] - rop = mod1(row + 1, size(peps_above, 1)) - rom = mod1(row - 1, size(peps_above, 1)) - - @diffset @tensor corners[NORTHWEST, rop, cop][-1; -2] := - envs.corners[NORTHWEST, rop, col][1, 2] * - envs.edges[NORTH, rop, col][2, 3, 4, -2] * - Q[-1; 1 3 4] - @diffset @tensor corners[SOUTHWEST, rom, cop][-1; -2] := - envs.corners[SOUTHWEST, rom, col][1, 4] * - envs.edges[SOUTH, rom, col][-1, 2, 3, 1] * - P[4 2 3; -2] - @diffset @tensor edges[WEST, row, cop][-1 -2 -3; -4] := - envs.edges[WEST, row, col][1 2 3; 4] * - peps_above[row, col][9; 5 -2 7 2] * - conj(peps_below[row, col][9; 6 -3 8 3]) * - P[4 5 6; -4] * - Q[-1; 1 7 8] + # Use projectors to grow the corners & edges + for row in 1:size(state, 1) + rprev = _prev(row, size(state, 1)) + rnext = _next(row, size(state, 1)) + C_sw, C_nw, T_w = grow_env_left( + state[row, col], + Pleft[_prev(row, end), col], + Pright[row, col], + env.corners[SOUTHWEST, rprev, col], + env.corners[NORTHWEST, rnext, col], + env.edges[SOUTH, rprev, col], + env.edges[WEST, row, col], + env.edges[NORTH, rnext, col], + ) + @diffset corners[SOUTHWEST, rprev, cnext] = C_sw + @diffset corners[NORTHWEST, rnext, cnext] = C_nw + @diffset edges[WEST, row, cnext] = T_w end - @diffset corners[NORTHWEST, :, cop] ./= norm.(corners[NORTHWEST, :, cop]) - @diffset edges[WEST, :, cop] ./= norm.(edges[WEST, :, cop]) - @diffset corners[SOUTHWEST, :, cop] ./= norm.(corners[SOUTHWEST, :, cop]) + @diffset corners[SOUTHWEST, :, cnext] ./= norm.(corners[SOUTHWEST, :, cnext]) + @diffset corners[NORTHWEST, :, cnext] ./= norm.(corners[NORTHWEST, :, cnext]) + @diffset edges[WEST, :, cnext] ./= norm.(edges[WEST, :, cnext]) end - return CTMRGEnv(peps_above, peps_below, corners, edges), ϵ + return CTMRGEnv(corners, edges), (; ϵ, Pleft, Pright) end -function northwest_corner(E4, C1, E1, peps_above, peps_below=peps_above) +# Compute enlarged NW corner +function northwest_corner(E4, C1, E1, peps) @tensor corner[-1 -2 -3; -4 -5 -6] := E4[-1 1 2; 3] * C1[3; 4] * E1[4 5 6; -4] * - peps_above[7; 5 -5 -2 1] * - conj(peps_below[7; 6 -6 -3 2]) + peps[7; 5 -5 -2 1] * + conj(peps[7; 6 -6 -3 2]) end -function northeast_corner(E1, C2, E2, peps_above, peps_below=peps_above) + +function northeast_corner(E1, C2, E2, peps) @tensor corner[-1 -2 -3; -4 -5 -6] := E1[-1 1 2; 3] * C2[3; 4] * E2[4 5 6; -4] * - peps_above[7; 1 5 -5 -2] * - conj(peps_below[7; 2 6 -6 -3]) + peps[7; 1 5 -5 -2] * + conj(peps[7; 2 6 -6 -3]) end -function southeast_corner(E2, C3, E3, peps_above, peps_below=peps_above) + +function southeast_corner(E2, C3, E3, peps) @tensor corner[-1 -2 -3; -4 -5 -6] := E2[-1 1 2; 3] * C3[3; 4] * E3[4 5 6; -4] * - peps_above[7; -2 1 5 -5] * - conj(peps_below[7; -3 2 6 -6]) + peps[7; -2 1 5 -5] * + conj(peps[7; -3 2 6 -6]) end -#= - -function MPSKit.leading_boundary(peps::InfinitePEPS,alg::CTMRG2,envs = CTMRGEnv(peps)) - err = Inf - iter = 1 - - old_norm = 1.0 - - while (err > alg.tol && iter <= alg.maxiter) || iter<4 - - for dir in 1:4 - envs = left_move(peps,alg,envs); - - envs = rotate_north(envs,EAST); - peps = envs.peps; - end - new_norm = abs(contract_ctrmg(peps,envs,1,1)) - #@show new_norm - err = abs(old_norm-new_norm) - @ignore_derivatives mod(alg.verbose,alg.miniter)==0 && mod(iter,alg.verbose+1)==0 && @info "$(iter) $(err) $(new_norm)" - - old_norm = new_norm - iter += 1 - end - @ignore_derivatives iter > alg.maxiter && @warn "maxiter $(alg.maxiter) reached: error was $(err)" - - envs +# Build projectors from SVD and enlarged SW & NW corners +function build_projectors( + U::AbstractTensorMap{E,3,1}, S, V::AbstractTensorMap{E,1,3}, Q_sw, Q_nw +) where {E<:ElementarySpace} + isqS = sdiag_inv_sqrt(S) + @tensor Pl[-1 -2 -3; -4] := Q_nw[-1 -2 -3; 1 2 3] * conj(V[4; 1 2 3]) * isqS[4; -4] + @tensor Pr[-1; -2 -3 -4] := isqS[-1; 1] * conj(U[2 3 4; 1]) * Q_sw[2 3 4; -2 -3 -4] + return Pl, Pr end -# the actual left_move is dependent on the type of ctmrg, so this seems natural -function left_move(peps::InfinitePEPS{PType},alg::CTMRG2,envs::CTMRGEnv) where PType - corners::typeof(envs.corners) = copy(envs.corners); - edges::typeof(envs.edges) = copy(envs.edges); - - above_projector_type = tensormaptype(spacetype(PType),1,3,storagetype(PType)); - below_projector_type = tensormaptype(spacetype(PType),3,1,storagetype(PType)); - - for col in 1:size(peps,2) - cop = mod1(col+1, size(peps,2)) - above_projs = Vector{above_projector_type}(undef,size(peps,1)); - below_projs = Vector{below_projector_type}(undef,size(peps,1)); - - # find all projectors - for row in 1:size(peps,1) - rop = mod1(row+1, size(peps,1)) - peps_nw = peps[row,col]; - peps_sw = rotate_north(peps[rop,col],WEST); - - Q1 = northwest_corner(envs.edges[WEST,row,col],envs.corners[NORTHWEST,row,col], envs.edges[NORTH,row,col],peps_nw); - Q2 = northeast_corner(envs.edges[NORTH,row,cop],envs.corners[NORTHEAST,row,cop],envs.edges[EAST,row,cop],peps[row,cop]) - Q3 = southeast_corner(envs.edges[EAST,rop,cop],envs.corners[SOUTHEAST,rop,cop],envs.edges[SOUTH,rop,cop],peps[rop,cop]) - Q4 = northwest_corner(envs.edges[SOUTH,rop,col],envs.corners[SOUTHWEST,rop,col],envs.edges[WEST,rop,col],peps_sw); - Qnorth = Q1*Q2 - Qsouth = Q3*Q4 - (U,S,V) = tsvd(Qsouth*Qnorth, trunc = alg.trscheme); - #@ignore_derivatives @show ϵ = real(norm(Qsouth*Qnorth)^2-norm(U*S*V)^2) - #@ignore_derivatives @info ϵ - isqS = sdiag_inv_sqrt(S); - Q = isqS*U'*Qsouth; - P = Qnorth*V'*isqS; - - @diffset above_projs[row] = Q; - @diffset below_projs[row] = P; - end - - #use the projectors to grow the corners/edges - for row in 1:size(peps,1) - Q = above_projs[row]; - P = below_projs[mod1(row-1,end)]; - - @diffset @tensor corners[NORTHWEST,row+1,col+1][-1;-2] := envs.corners[NORTHWEST,row+1,col][1,2] * envs.edges[NORTH,row+1,col][2,3,4,-2]*Q[-1;1 3 4] - @diffset @tensor corners[SOUTHWEST,row-1,col+1][-1;-2] := envs.corners[SOUTHWEST,row-1,col][1,4] * envs.edges[SOUTH,row-1,col][-1,2,3,1]*P[4 2 3;-2] - @diffset @tensor edges[WEST,row,col+1][-1 -2 -3;-4] := envs.edges[WEST,row,col][1 2 3;4]* - peps[row,col][9;5 -2 7 2]* - conj(peps[row,col][9;6 -3 8 3])* - P[4 5 6;-4]* - Q[-1;1 7 8] - end - - @diffset corners[NORTHWEST,:,col+1]./=norm.(corners[NORTHWEST,:,col+1]); - @diffset corners[SOUTHWEST,:,col+1]./=norm.(corners[SOUTHWEST,:,col+1]); - @diffset edges[WEST,:,col+1]./=norm.(edges[WEST,:,col+1]); - end - - CTMRGEnv(peps,corners,edges); +# Apply projectors to entire left half-environment to grow SW & NW corners, and W edge +function grow_env_left(peps, Pl, Pr, C_sw, C_nw, T_s, T_w, T_n) + # @diffset @tensor corners[NORTHWEST, rop, cop][-1; -2] := + # envs.corners[NORTHWEST, rop, col][1, 2] * + # envs.edges[NORTH, rop, col][2, 3, 4, -2] * + # Q[-1; 1 3 4] + # @diffset @tensor corners[SOUTHWEST, rom, cop][-1; -2] := + # envs.corners[SOUTHWEST, rom, col][1, 4] * + # envs.edges[SOUTH, rom, col][-1, 2, 3, 1] * + # P[4 2 3; -2] + # @diffset @tensor edges[WEST, row, cop][-1 -2 -3; -4] := + # envs.edges[WEST, row, col][1 2 3; 4] * + # peps_above[row, col][9; 5 -2 7 2] * + # conj(peps_below[row, col][9; 6 -3 8 3]) * + # P[4 5 6; -4] * + # Q[-1; 1 7 8] + @tensor C_sw′[-1; -2] := C_sw[1; 4] * T_s[-1 2 3; 1] * Pl[4 2 3; -2] + @tensor C_nw′[-1; -2] := C_nw[1; 2] * T_n[2 3 4; -2] * Pr[-1; 1 3 4] + @tensor T_w′[-1 -2 -3; -4] := + T_w[1 2 3; 4] * + peps[9; 5 -2 7 2] * + conj(peps[9; 6 -3 8 3]) * + Pl[4 5 6; -4] * + Pr[-1; 1 7 8] + return C_sw′, C_nw′, T_w′ end -=# -function contract_ctrmg( - envs::CTMRGEnv, peps_above=envs.peps_above, peps_below=envs.peps_below -) +# Compute norm of the entire CMTRG enviroment +function contract_ctmrg(peps, env::CTMRGEnv) total = 1.0 + 0im - for r in 1:size(peps_above, 1), c in 1:size(peps_above, 2) - total *= @tensor envs.edges[WEST, r, c][1 2 3; 4] * - envs.corners[NORTHWEST, r, c][4; 5] * - envs.edges[NORTH, r, c][5 6 7; 8] * - envs.corners[NORTHEAST, r, c][8; 9] * - envs.edges[EAST, r, c][9 10 11; 12] * - envs.corners[SOUTHEAST, r, c][12; 13] * - envs.edges[SOUTH, r, c][13 14 15; 16] * - envs.corners[SOUTHWEST, r, c][16; 1] * - peps_above[r, c][17; 6 10 14 2] * - conj(peps_below[r, c][17; 7 11 15 3]) + for r in 1:size(peps, 1), c in 1:size(peps, 2) + total *= @tensor env.edges[WEST, r, c][1 2 3; 4] * + env.corners[NORTHWEST, r, c][4; 5] * + env.edges[NORTH, r, c][5 6 7; 8] * + env.corners[NORTHEAST, r, c][8; 9] * + env.edges[EAST, r, c][9 10 11; 12] * + env.corners[SOUTHEAST, r, c][12; 13] * + env.edges[SOUTH, r, c][13 14 15; 16] * + env.corners[SOUTHWEST, r, c][16; 1] * + peps[r, c][17; 6 10 14 2] * + conj(peps[r, c][17; 7 11 15 3]) + total *= tr( - envs.corners[NORTHWEST, r, c] * - envs.corners[NORTHEAST, r, mod1(c - 1, end)] * - envs.corners[SOUTHEAST, mod1(r - 1, end), mod1(c - 1, end)] * - envs.corners[SOUTHWEST, mod1(r - 1, end), c], + env.corners[NORTHWEST, r, c] * + env.corners[NORTHEAST, r, mod1(c - 1, end)] * + env.corners[SOUTHEAST, mod1(r - 1, end), mod1(c - 1, end)] * + env.corners[SOUTHWEST, mod1(r - 1, end), c], ) - total /= @tensor envs.edges[WEST, r, c][1 10 11; 4] * - envs.corners[NORTHWEST, r, c][4; 5] * - envs.corners[NORTHEAST, r, mod1(c - 1, end)][5; 6] * - envs.edges[EAST, r, mod1(c - 1, end)][6 10 11; 7] * - envs.corners[SOUTHEAST, r, mod1(c - 1, end)][7; 8] * - envs.corners[SOUTHWEST, r, c][8; 1] - - total /= @tensor envs.corners[NORTHWEST, r, c][1; 2] * - envs.edges[NORTH, r, c][2 10 11; 3] * - envs.corners[NORTHEAST, r, c][3; 4] * - envs.corners[SOUTHEAST, mod1(r - 1, end), c][4; 5] * - envs.edges[SOUTH, mod1(r - 1, end), c][5 10 11; 6] * - envs.corners[SOUTHWEST, mod1(r - 1, end), c][6; 1] + total /= @tensor env.edges[WEST, r, c][1 10 11; 4] * + env.corners[NORTHWEST, r, c][4; 5] * + env.corners[NORTHEAST, r, mod1(c - 1, end)][5; 6] * + env.edges[EAST, r, mod1(c - 1, end)][6 10 11; 7] * + env.corners[SOUTHEAST, r, mod1(c - 1, end)][7; 8] * + env.corners[SOUTHWEST, r, c][8; 1] + + total /= @tensor env.corners[NORTHWEST, r, c][1; 2] * + env.edges[NORTH, r, c][2 10 11; 3] * + env.corners[NORTHEAST, r, c][3; 4] * + env.corners[SOUTHEAST, mod1(r - 1, end), c][4; 5] * + env.edges[SOUTH, mod1(r - 1, end), c][5 10 11; 6] * + env.corners[SOUTHWEST, mod1(r - 1, end), c][6; 1] end return total diff --git a/src/algorithms/expval.jl b/src/algorithms/expval.jl deleted file mode 100644 index 99a7b1a5..00000000 --- a/src/algorithms/expval.jl +++ /dev/null @@ -1,34 +0,0 @@ -function MPSKit.expectation_value(envs::CTMRGEnv, ham::AbstractTensorMap{S,1,1}) where {S} - @assert envs.peps_above == envs.peps_below - peps = envs.peps_above - result = Matrix{eltype(ham)}(undef, size(peps, 1), size(peps, 2)) - - for r in 1:size(peps, 1), c in 1:size(peps, 2) - e = @tensor envs.edges[WEST, r, c][1 2 3; 4] * - envs.corners[NORTHWEST, r, c][4; 5] * - envs.edges[NORTH, r, c][5 6 7; 8] * - envs.corners[NORTHEAST, r, c][8; 9] * - envs.edges[EAST, r, c][9 10 11; 12] * - envs.corners[SOUTHEAST, r, c][12; 13] * - envs.edges[SOUTH, r, c][13 14 15; 16] * - envs.corners[SOUTHWEST, r, c][16; 1] * - peps[r, c][17; 6 10 14 2] * - conj(peps[r, c][18; 7 11 15 3]) * - ham[18; 17] - - n = @tensor envs.edges[WEST, r, c][1 2 3; 4] * - envs.corners[NORTHWEST, r, c][4; 5] * - envs.edges[NORTH, r, c][5 6 7; 8] * - envs.corners[NORTHEAST, r, c][8; 9] * - envs.edges[EAST, r, c][9 10 11; 12] * - envs.corners[SOUTHEAST, r, c][12; 13] * - envs.edges[SOUTH, r, c][13 14 15; 16] * - envs.corners[SOUTHWEST, r, c][16; 1] * - peps[r, c][17; 6 10 14 2] * - conj(peps[r, c][17; 7 11 15 3]) - - @diffset result[r, c] = e / n - end - - return result -end diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl new file mode 100644 index 00000000..96889c73 --- /dev/null +++ b/src/algorithms/peps_opt.jl @@ -0,0 +1,259 @@ +abstract type GradMode end + +struct NaiveAD <: GradMode end +struct GeomSum <: GradMode end +struct ManualIter <: GradMode end +struct LinSolve <: GradMode end + +# Algorithm struct containing parameters for PEPS optimization +# TODO: have an interface for general cost functions? (could merge energyfun and reuse_env) +@kwdef struct PEPSOptimize{G<:GradMode} + optimizer::OptimKit.OptimizationAlgorithm = LBFGS( + 4; maxiter=100, gradtol=1e-4, verbosity=2 + ) + energyfun::Function = next_neighbor_energy # Energy function returning real scalar + reuse_env::Bool = true # Reuse environment of previous optimization as initial guess for next + fpgrad_tol::Float64 = Defaults.grad_tol # Convergence tolerance for gradient FP iteration + fpgrad_maxiter::Int = Defaults.grad_maxiter # Maximal number of FP iterations + verbose::Int = 0 +end + +# Find ground-state PEPS, environment and energy +function groundsearch( + H, ctmalg::CTMRG, optalg::PEPSOptimize, ψinit::InfinitePEPS, envinit::CTMRGEnv +) + (peps₀, env₀), E₀, ∂E, info = optimize( + x -> ctmrg_gradient(x, H, ctmalg, optalg), + (ψinit, envinit), + optalg.optimizer; + inner=my_inner, + retract=my_retract, + (add!)=my_add!, + (scale!)=my_scale!, + ) + return (; peps₀, env₀, E₀, ∂E, info) +end + +# Function returning energy and CTMRG gradient for each optimization step +function ctmrg_gradient(x, H, ctmalg::CTMRG, optalg::PEPSOptimize) + peps, env = x + cfun = optalg.reuse_env ? costfun! : costfun + E = cfun(peps, env, H, ctmalg, optalg) + ∂E∂A = gradient(cfun, peps, env, H, ctmalg, optalg)[1] + @assert !isnan(norm(∂E∂A)) + return E, ∂E∂A +end + +# Energy cost function with proper backwards rule depending only on final CTMRG fixed-point +# Mutates environment to reuse previous environments in optimization +function costfun!(peps, env, H, ctmalg::CTMRG, optalg::PEPSOptimize) + env′ = leading_boundary(peps, ctmalg, env) + @diffset env = env′ + return optalg.energyfun(peps, env′, H) +end + +# Non-mutating version, recomputing environment from random initial guess in every optimization step +function costfun(peps, env, H, ctmalg::CTMRG, optalg::PEPSOptimize) + env′ = deepcopy(env) # Create copy to make non-mutating + return costfun!(peps, env′, H, ctmalg, optalg) +end + +# Energy gradient backwards rule (does not apply to NaiveAD gradient mode) +function ChainRulesCore.rrule( + ::typeof(costfun!), peps, env, H, ctmalg::CTMRG, optalg::PEPSOptimize{G} +) where {G<:Union{GeomSum,ManualIter,LinSolve}} + env = leading_boundary(peps, env, ctmalg) + E, Egrad = withgradient(optalg.energyfun, peps, env, H) + ∂F∂x = CTMRGEnv(Egrad[2]...) + _, xvjp = pullback(ctmrg_gauged_iter, peps, env, ctmalg) + ∂f∂A(v) = xvjp(v)[1] + ∂f∂x(v) = CTMRGEnv(xvjp(v)[2]...) # Function wrapper to compute v*∂f∂A vJP as CTMRGEnv + + function costfun!_pullback(_) + y₀ = CTMRGEnv(peps, dim(space(env.edges[1])[1])) + dx, = fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, optalg) + return NoTangent(), Egrad[1] + dx, NoTangent(), NoTangent(), NoTangent() + end + + return E, costfun!_pullback +end + +# Contraction of CTMRGEnv and PEPS tensors with open physical bonds +function one_site_rho(peps::InfinitePEPS, env::CTMRGEnv{C,T}) where {C,T} + ρunitcell = similar(peps.A, tensormaptype(spacetype(C), 1, 1, storagetype(C))) + + for r in size(env.corners, 2), c in size(env.corners, 3) + @tensor ρ[-1; -2] := + env.corners[NORTHWEST, r, c][1; 2] * + env.edges[NORTH, r, c][2 3 4; 5] * + env.corners[NORTHEAST, r, c][5; 6] * + env.edges[EAST, r, c][6 7 8; 9] * + env.corners[SOUTHEAST, r, c][9; 10] * + env.edges[SOUTH, r, c][10 11 12; 13] * + env.corners[SOUTHWEST, r, c][13; 14] * + env.edges[WEST, r, c][14 15 16; 1] * + peps[r, c][-1; 3 7 11 15] * + conj(peps[r, c][-2; 4 8 12 16]) + @diffset ρunitcell[r, c] = ρ + end + + return ρunitcell +end + +# Horizontally extended contraction of CTMRGEnv and PEPS tensors with open physical bonds +function two_site_rho(peps::InfinitePEPS, env::CTMRGEnv{C,T}) where {C,T} + ρunitcell = similar(peps.A, tensormaptype(spacetype(C), 2, 2, storagetype(C))) + + for r in size(env.corners, 2), c in size(env.corners, 3) + cnext = _next(c, size(peps, 2)) + @tensor ρ[-11 -20; -12 -18] := + env.corners[NORTHWEST, r, c][1; 3] * + env.edges[NORTH, r, c][3 5 8; 13] * + env.edges[NORTH, r, cnext][13 16 22; 23] * + env.corners[NORTHEAST, r, cnext][23; 24] * + env.edges[EAST, r, cnext][24 25 26; 27] * + env.corners[SOUTHEAST, r, cnext][27; 28] * + env.edges[SOUTH, r, cnext][28 17 21; 14] * + env.edges[SOUTH, r, c][14 6 10; 4] * + env.corners[SOUTHWEST, r, c][4; 2] * + env.edges[WEST, r, c][2 7 9; 1] * + peps[r, c][-12; 5 15 6 7] * + conj(peps[r, c][-11; 8 19 10 9]) * + peps[r, cnext][-18; 16 25 17 15] * + conj(peps[r, cnext][-20; 22 26 21 19]) + @diffset ρunitcell[r, c] = ρ + end + + return ρunitcell +end + +# 1-site operator expectation values on unit cell +function MPSKit.expectation_value( + peps::InfinitePEPS, env::CTMRGEnv, op::AbstractTensorMap{S,1,1} +) where {S<:ElementarySpace} + result = similar(peps.A, eltype(op)) + ρ = one_site_rho(peps, env) + + for r in 1:size(peps, 1), c in 1:size(peps, 2) + o = @tensor ρ[r, c][1; 2] * op[1; 2] + n = @tensor ρ[r, c][1; 1] + @diffset result[r, c] = o / n + end + + return result +end + +# 2-site operator expectation values on unit cell +function MPSKit.expectation_value( + peps::InfinitePEPS, env::CTMRGEnv, op::AbstractTensorMap{S,2,2} +) where {S<:ElementarySpace} + result = similar(peps.A, eltype(op)) + ρ = two_site_rho(peps, env) + + for r in 1:size(peps, 1), c in 1:size(peps, 2) + o = @tensor ρ[r, c][1 2; 3 4] * op[1 2; 3 4] + n = @tensor ρ[r, c][1 2; 1 2] + @diffset result[r, c] = o / n + end + + return result +end + +# ⟨H⟩ from vertical and horizontal next-neighbor contributions +function next_neighbor_energy( + peps::InfinitePEPS, env::CTMRGEnv, H::AbstractTensorMap{S,2,2} +) where {S<:ElementarySpace} + Eh = sum(expectation_value(peps, env, H)) + Ev = sum(expectation_value(rotl90(peps), rotl90(env), H)) + return real(Eh + Ev) +end + +# Compute energy and energy gradient, by explicitly evaluating geometric series +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, _, alg::PEPSOptimize{GeomSum}) + g = ∂F∂x + dx = ∂f∂A(g) # n=0 term: ∂F∂x ∂f∂A + err = 0.0 + for i in 1:(alg.maxiter) + updateenv!(g, ∂f∂x(g)) + Σₙ = ∂f∂A(g) + dx += Σₙ + err = norm(Σₙ) # TODO: normalize this error? + alg.verbose && @show err + err < alg.tol && break + + if i == alg.maxiter + @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Σₙ‖ = $(norm(Σₙ))" + end + end + return dx, err +end + +# Manual iteration to solve gradient linear problem +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::PEPSOptimize{ManualIter}) + y = deepcopy(y₀) # Do not mutate y₀ + err = 0.0 + for i in 1:(alg.maxiter) + y′ = ∂F∂x + ∂f∂x(y) + norma = norm(y.corners[NORTHWEST]) + err = norm(y′.corners[NORTHWEST] - y.corners[NORTHWEST]) / norma # Normalize error to get comparable convergence tolerance + alg.verbose && @show err + updateenv!(y, y′) + err < alg.tol && break + + if i == alg.maxiter + @warn "gradient fixed-point iteration reached maximal number of iterations at ‖Cᵢ₊₁ - Cᵢ‖ = $err" + end + end + return ∂f∂A(y), err +end + +# Use proper iterative solver to solve gradient problem +function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::PEPSOptimize) + # spaces = [space.(getfield(∂F∂x, f)) for f in fieldnames(CTMRGEnv)] + # sizes = [map(x -> size(x.data), getfield(∂F∂x, f)) for f in fieldnames(CTMRGEnv)] + # op = LinearMap(vecdim(∂F∂x)) do v + # env = unvec(v, spaces..., sizes...) + # x = env - ∂f∂x(env) + # vec(x) + # end + # envvec = vec(y₀) + # info = gmres!(envvec, op, vec(∂F∂x); reltol=alg.grad_convtol, maxiter=alg.grad_maxiter) + # y = unvec(envvec, spaces..., sizes...) + + # ∂f∂A(y), info +end + +# Update PEPS unit cell in non-mutating way +function my_retract(x, dx, α) + peps = deepcopy(x[1]) + peps.A .+= dx.A .* α + env = deepcopy(x[2]) + return (peps, env), dx +end + +# Take real valued part of standard dot product +my_inner(_, η₁, η₂) = real(dot(η₁, η₂)) + +# Add unit cell elements element-wise +function my_add!(Y, X, a) + Y.A .+= X.A .* a + return Y +end + +# Scale all unit cell elements +function my_scale!(η, β) + rmul!(η.A, β) + return η +end + +# my_retract is not an in place function which should not change x +# function my_retract(x, dx, α::Number) +# (ϕ, env0) = x +# ψ = deepcopy(ϕ) +# env = deepcopy(env0) +# ψ.A .+= dx.A .* α +# #env = leading_boundary(ψ, alg_ctm,env) +# return (ψ, env), dx +# end + +# my_inner(x, dx1, dx2) = real(dot(dx1, dx2)) diff --git a/src/environments/ctmrgenv.jl b/src/environments/ctmrgenv.jl index a00627e4..dcb64a1a 100644 --- a/src/environments/ctmrgenv.jl +++ b/src/environments/ctmrgenv.jl @@ -1,52 +1,44 @@ -struct CTMRGEnv{P,C,T} - peps_above::InfinitePEPS{P} - peps_below::InfinitePEPS{P} +struct CTMRGEnv{C,T} corners::Array{C,3} edges::Array{T,3} end # initialize ctmrg environments with some random tensors -CTMRGEnv(peps::InfinitePEPS) = CTMRGEnv(peps, peps); - -function CTMRGEnv(peps_above::InfinitePEPS{P}, peps_below::InfinitePEPS{P}) where {P} - ou = oneunit(space(peps_above, 1, 1)) # the bogus space - +function CTMRGEnv(peps::InfinitePEPS{P}; χenv=1) where {P} + envspace = field(space(peps, 1, 1))^χenv # Environment space C_type = tensormaptype(spacetype(P), 1, 1, storagetype(P)) - T_type = tensormaptype(spacetype(P), 3, 1, storagetype(P)) # debatable how we should do the legs? + T_type = tensormaptype(spacetype(P), 3, 1, storagetype(P)) - #first index is de - corners = Array{C_type}(undef, 4, size(peps_above)...) - edges = Array{T_type}(undef, 4, size(peps_above)...) + # First index is direction + corners = Array{C_type}(undef, 4, size(peps)...) + edges = Array{T_type}(undef, 4, size(peps)...) - for dir in 1:4, i in 1:size(peps_above, 1), j in 1:size(peps_above, 2) - @diffset corners[dir, i, j] = TensorMap(randn, scalartype(P), ou, ou) + for dir in 1:4, i in 1:size(peps, 1), j in 1:size(peps, 2) + @diffset corners[dir, i, j] = TensorMap(randn, scalartype(P), envspace, envspace) @diffset edges[dir, i, j] = TensorMap( randn, scalartype(P), - ou * space(peps_above[i, j], dir + 1)' * space(peps_below[i, j], dir + 1), - ou, + envspace * space(peps[i, j], dir + 1)' * space(peps[i, j], dir + 1), + envspace, ) end @diffset corners[:, :, :] ./= norm.(corners[:, :, :]) @diffset edges[:, :, :] ./= norm.(edges[:, :, :]) - return CTMRGEnv(peps_above, peps_below, corners, edges) + return CTMRGEnv(corners, edges) end -function Base.rotl90(envs::CTMRGEnv{P,C,T}) where {P,C,T} - n_peps_above = rotl90(envs.peps_above) - n_peps_below = rotl90(envs.peps_below) - n_corners = Array{C,3}(undef, size(envs.corners)...) - n_edges = Array{T,3}(undef, size(envs.edges)...) +function Base.rotl90(env::CTMRGEnv{C,T}) where {C,T} + corners′ = similar(env.corners) + edges′ = similar(env.edges) for dir in 1:4 - dirm = mod1(dir - 1, 4) - @diffset n_corners[dirm, :, :] .= rotl90(envs.corners[dir, :, :]) - @diffset n_edges[dirm, :, :] .= rotl90(envs.edges[dir, :, :]) + @diffset corners′[_prev(dir, 4), :, :] .= rotl90(env.corners[dir, :, :]) + @diffset edges′[_prev(dir, 4), :, :] .= rotl90(env.edges[dir, :, :]) end - return CTMRGEnv(n_peps_above, n_peps_below, n_corners, n_edges) + return CTMRGEnv(corners′, edges′) end -Base.eltype(envs::CTMRGEnv) = eltype(envs.corners[1]) +Base.eltype(env::CTMRGEnv) = eltype(env.corners[1]) diff --git a/src/states/abstractpeps.jl b/src/states/abstractpeps.jl index 490f4695..acacaf4f 100644 --- a/src/states/abstractpeps.jl +++ b/src/states/abstractpeps.jl @@ -1,9 +1,3 @@ -# abstractpeps.jl - -########################### -## Abstract tensor types ## -########################### - """ const PEPSTensor{S} @@ -23,6 +17,7 @@ function PEPSTensor( ) where {T,S<:ElementarySpace} return TensorMap(f, T, Pspace ← Nspace ⊗ Espace ⊗ Sspace ⊗ Wspace) end + function PEPSTensor( f, ::Type{T}, @@ -39,14 +34,10 @@ end const PEPOTensor{S} Default type for PEPO tensors with a single incoming and outgoing physical index, and 4 - virtual indices, conventionally ordered as: O : P ⊗ P' ← N ⊗ E ⊗ S ⊗ W. +virtual indices, conventionally ordered as: O : P ⊗ P' ← N ⊗ E ⊗ S ⊗ W. """ const PEPOTensor{S} = AbstractTensorMap{S,2,4} where {S<:ElementarySpace} -########################## -## Abstract state types ## -########################## - """ abstract type AbstractPEPS end @@ -62,4 +53,4 @@ Abstract supertype for a 2D projected entangled pairs operator. abstract type AbstractPEPO end Base.rotl90(t::PEPSTensor) = permute(t, ((1,), (3, 4, 5, 2))) -Base.rotl90(t::PEPOTensor) = permute(t, ((1, 2), (4, 5, 6, 3))); +Base.rotl90(t::PEPOTensor) = permute(t, ((1, 2), (4, 5, 6, 3))) \ No newline at end of file diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl index dc717656..88a6dd19 100644 --- a/src/states/infinitepeps.jl +++ b/src/states/infinitepeps.jl @@ -1,14 +1,10 @@ -# not everything is a PeriodicArray anymore -_next(i, total) = mod1(i + 1, total) -_prev(i, total) = mod1(i - 1, total) - """ struct InfinitePEPS{T<:PEPSTensor} Represents an infinite projected entangled-pair state on a 2D square lattice. """ struct InfinitePEPS{T<:PEPSTensor} <: AbstractPEPS - A::Array{T,2} # TODO: switch back to PeriodicArray? + A::Array{T,2} function InfinitePEPS(A::Array{T,2}) where {T<:PEPSTensor} for (d, w) in Tuple.(CartesianIndices(A)) @@ -40,7 +36,9 @@ Allow users to pass in arrays of spaces. function InfinitePEPS( Pspaces::AbstractArray{S,2}, Nspaces::AbstractArray{S,2}, - Espaces::AbstractArray{S,2}=Nspaces, + Espaces::AbstractArray{S,2}=Nspaces; + finit=rand, + dtype=ComplexF64, ) where {S<:ElementarySpace} size(Pspaces) == size(Nspaces) == size(Espaces) || throw(ArgumentError("Input spaces should have equal sizes.")) @@ -49,7 +47,7 @@ function InfinitePEPS( Wspaces = adjoint.(circshift(Espaces, (0, -1))) A = map(Pspaces, Nspaces, Espaces, Sspaces, Wspaces) do P, N, E, S, W - return PEPSTensor(randn, ComplexF64, P, N, E, S, W) + return PEPSTensor(finit, dtype, P, N, E, S, W) end return InfinitePEPS(A) @@ -87,7 +85,7 @@ VectorInterface.scalartype(T::InfinitePEPS) = scalartype(T.A) ## Copy Base.copy(T::InfinitePEPS) = InfinitePEPS(copy(T.A)) -Base.similar(T::InfinitePEPS) = InfinitePEPS(similar(T.A)) +Base.similar(T::InfinitePEPS, args...) = InfinitePEPS(similar(T.A, args...)) Base.repeat(T::InfinitePEPS, counts...) = InfinitePEPS(repeat(T.A, counts...)) Base.getindex(T::InfinitePEPS, args...) = Base.getindex(T.A, args...) @@ -95,4 +93,4 @@ Base.setindex!(T::InfinitePEPS, args...) = (Base.setindex!(T.A, args...); T) Base.axes(T::InfinitePEPS, args...) = axes(T.A, args...) TensorKit.space(t::InfinitePEPS, i, j) = space(t[i, j], 1) -Base.rotl90(t::InfinitePEPS) = InfinitePEPS(rotl90(rotl90.(t.A))); +Base.rotl90(t::InfinitePEPS) = InfinitePEPS(rotl90(rotl90.(t.A))) diff --git a/src/utility/rotations.jl b/src/utility/rotations.jl index 58d96c3d..983a8292 100644 --- a/src/utility/rotations.jl +++ b/src/utility/rotations.jl @@ -1,13 +1,12 @@ -const NORTH = 1; -const EAST = 2; -const SOUTH = 3; -const WEST = 4; +const NORTH = 1 +const EAST = 2 +const SOUTH = 3 +const WEST = 4 -const NORTHWEST = 1; -const NORTHEAST = 2; -const SOUTHEAST = 3; -const SOUTHWEST = 4; +const NORTHWEST = 1 +const NORTHEAST = 2 +const SOUTHEAST = 3 +const SOUTHWEST = 4 +# Rotate tensor to any direction by successive application of Base.rotl90 rotate_north(t, dir) = mod1(dir, 4) == NORTH ? t : rotate_north(rotl90(t), dir - 1) - -#Base.rotl90(t::Tuple) = rotl90.(t); diff --git a/src/utility/util.jl b/src/utility/util.jl index 5ebf8d4e..25c1a70c 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -1,34 +1,42 @@ +# Get next and previous directional CTM enviroment index, respecting periodicity +_next(i, total) = mod1(i + 1, total) +_prev(i, total) = mod1(i - 1, total) + +# Element-wise multiplication of TensorMaps respecting block structure +function _elementwise_mult(a::AbstractTensorMap, b::AbstractTensorMap) + dst = similar(a) + for (k, block) in blocks(dst) + copyto!(block, blocks(a)[k] .* blocks(b)[k]) + end + return dst +end + +# Compute √S⁻¹ for diagonal TensorMaps function sdiag_inv_sqrt(S::AbstractTensorMap) - toret = similar(S) + invsq = similar(S) if sectortype(S) == Trivial - copyto!(toret.data, LinearAlgebra.diagm(LinearAlgebra.diag(S.data) .^ (-1 / 2))) + copyto!(invsq.data, LinearAlgebra.diagm(LinearAlgebra.diag(S.data) .^ (-1 / 2))) else for (k, b) in blocks(S) copyto!( - blocks(toret)[k], LinearAlgebra.diagm(LinearAlgebra.diag(b) .^ (-1 / 2)) + blocks(invsq)[k], LinearAlgebra.diagm(LinearAlgebra.diag(b) .^ (-1 / 2)) ) end end - return toret + return invsq end + function ChainRulesCore.rrule(::typeof(sdiag_inv_sqrt), S::AbstractTensorMap) - toret = sdiag_inv_sqrt(S) - return toret, - c̄ -> (ChainRulesCore.NoTangent(), -1 / 2 * _elementwise_mult(c̄, toret'^3)) -end -function _elementwise_mult(a::AbstractTensorMap, b::AbstractTensorMap) - dst = similar(a) - for (k, block) in blocks(dst) - copyto!(block, blocks(a)[k] .* blocks(b)[k]) + invsq = sdiag_inv_sqrt(S) + function sdiag_inv_sqrt_pullback(c̄) + return (ChainRulesCore.NoTangent(), -1 / 2 * _elementwise_mult(c̄, invsq'^3)) end - return dst + return invsq, sdiag_inv_sqrt_pullback end -#rotl90 appeared to lose PeriodicArray'ness, which in turn caused zygote problems -#Base.rotl90(a::Array) = Array(rotl90(a)); -#Base.rotr90(a::Array) = Array(rotr90(a)); +# rotl90 is set to non_differentiable in ChainRules function ChainRulesCore.rrule(::typeof(rotl90), a::AbstractMatrix) function pb(x) if !iszero(x) @@ -45,8 +53,7 @@ function ChainRulesCore.rrule(::typeof(rotl90), a::AbstractMatrix) return rotl90(a), pb end -structure(t) = codomain(t) ← domain(t); - +# Differentiable setindex! alternative function _setindex(a::AbstractArray, v, args...) b::typeof(a) = copy(a) b[args...] = v @@ -56,7 +63,7 @@ end function ChainRulesCore.rrule(::typeof(_setindex), a::AbstractArray, tv, args...) t = _setindex(a, tv, args...) - function toret(v) + function _setindex_pullback(v) if iszero(v) backwards_tv = ZeroTangent() backwards_a = ZeroTangent() @@ -80,9 +87,11 @@ function ChainRulesCore.rrule(::typeof(_setindex), a::AbstractArray, tv, args... NoTangent(), backwards_a, backwards_tv, fill(ZeroTangent(), length(args))... ) end - return t, toret + return t, _setindex_pullback end +# Allows in-place operations during AD that copies when differentiating +# Especially needed to set tensors in unit cell of environments macro diffset(ex) return esc(parse_ex(ex)) end