From 8fe91832f4096f438c77ccf8d4950f59807253de Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 10 Jul 2024 11:35:04 +0200 Subject: [PATCH] Add FixedSVD, refactor build_projectors, improve SVD passing --- src/algorithms/ctmrg.jl | 28 +++-- src/algorithms/ctmrg_all_sides.jl | 167 +++++++++++++----------------- src/utility/svd.jl | 18 ++++ src/utility/util.jl | 14 ++- 4 files changed, 116 insertions(+), 111 deletions(-) diff --git a/src/algorithms/ctmrg.jl b/src/algorithms/ctmrg.jl index 331c7415..8306f3c8 100644 --- a/src/algorithms/ctmrg.jl +++ b/src/algorithms/ctmrg.jl @@ -52,7 +52,7 @@ function CTMRG(; verbosity=0, svd_alg=SVDAdjoint(), trscheme=FixedSpaceTruncation(), - ctmrgscheme=:AllSides + ctmrgscheme=:AllSides, ) return CTMRG{ctmrgscheme}( tol, maxiter, miniter, verbosity, ProjectorAlg(; svd_alg, trscheme, verbosity) @@ -133,7 +133,6 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG) return envfix end - """ ctmrg_iter(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T} @@ -165,7 +164,7 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T} corners::typeof(env.corners) = copy(env.corners) edges::typeof(env.edges) = copy(env.edges) ϵ = 0.0 - Pleft, Pright = Zygote.Buffer.(projector_type(T, size(state))) # Use Zygote.Buffer instead of @diffset to avoid ZeroTangent errors in _setindex + P_top, P_bottom = Zygote.Buffer.(projector_type(T, size(state))) # Use Zygote.Buffer instead of @diffset to avoid ZeroTangent errors in _setindex for col in 1:size(state, 2) cprev = _prev(col, size(state, 2)) @@ -192,11 +191,10 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T} # SVD half-infinite environment trscheme = if alg.trscheme isa FixedSpaceTruncation - truncspace(space(env.edges[WEST, row, cnext], 1)) + truncspace(space(env.edges[WEST, row, col], 1)) else alg.trscheme end - @tensor QQ[-1 -2 -3; -4 -5 -6] := Q_sw[-1 -2 -3; 1 2 3] * Q_nw[1 2 3; -4 -5 -6] @autoopt @tensor QQ[χ_EB D_EBabove D_EBbelow; χ_ET D_ETabove D_ETbelow] := Q_sw[χ_EB D_EBabove D_EBbelow; χ D1 D2] * Q_nw[χ D1 D2; χ_ET D_ETabove D_ETbelow] @@ -213,9 +211,9 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T} end # Compute projectors - Pl, Pr = build_projectors(U, S, V, Q_sw, Q_nw) - Pleft[row, col] = Pl - Pright[row, col] = Pr + Pt, Pb = build_projectors(U, S, V, Q_sw, Q_nw) + P_top[row, col] = Pt + P_bottom[row, col] = Pb end # Use projectors to grow the corners & edges @@ -223,8 +221,8 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T} rprev = _prev(row, size(state, 1)) C_sw, C_nw, T_w = grow_env_left( state[row, col], - Pleft[rprev, col], - Pright[row, col], + P_left[rprev, col], + P_right[row, col], env.corners[SOUTHWEST, row, cprev], env.corners[NORTHWEST, row, cprev], env.edges[SOUTH, row, col], @@ -237,7 +235,7 @@ function left_move(state, env::CTMRGEnv{C,T}, alg::ProjectorAlg) where {C,T} end end - return CTMRGEnv(corners, edges), (; Pleft=copy(Pleft), Pright=copy(Pright), ϵ) + return CTMRGEnv(corners, edges), (; P_left=copy(P_top), P_right=copy(P_bottom), ϵ) end # Compute enlarged corners @@ -276,12 +274,12 @@ end # 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 + U::AbstractTensorMap{E,3,1}, S, V::AbstractTensorMap{E,1,3}, Q, Q_next ) where {E<:ElementarySpace} isqS = sdiag_inv_sqrt(S) - P_bottom = Q_NW * V' * isqS - P_top = isqS * U' * Q_SW - return P_bottom, P_top + P_left = Q_next * V' * isqS + P_right = isqS * U' * Q + return P_left, P_right end # Apply projectors to entire left half-environment to grow SW & NW corners, and W edge diff --git a/src/algorithms/ctmrg_all_sides.jl b/src/algorithms/ctmrg_all_sides.jl index f18d2044..636a8322 100644 --- a/src/algorithms/ctmrg_all_sides.jl +++ b/src/algorithms/ctmrg_all_sides.jl @@ -1,15 +1,20 @@ + +struct IndexSelector{N} <: TensorKit.TruncationScheme + index::NTuple{N,Int} +end + # One CTMRG iteration with both-sided application of projectors function ctmrg_iter(state, env::CTMRGEnv, alg::CTMRG{:AllSides}) # Compute enlarged corners Q = enlarge_corners_edges(state, env) # Compute projectors if none are supplied - Pleft, Pright, info = build_projectors(Q, env, alg) + P_left, P_right, info = build_projectors(Q, env, alg) # Apply projectors and normalize - corners, edges = renormalize_corners_edges(state, env, Q, Pleft, Pright) + corners, edges = renormalize_corners_edges(state, env, Q, P_left, P_right) - return CTMRGEnv(corners, edges), (; Pleft, Pright, info...) + return CTMRGEnv(corners, edges), (; P_left, P_right, info...) end # Compute enlarged corners and edges for all directions and unit cell entries @@ -52,90 +57,66 @@ function enlarge_corners_edges(state, env::CTMRGEnv) end # Build projectors from SVD and enlarged corners -function build_projectors(Q, env::CTMRGEnv, alg::ProjectorAlg) # TODO: Add projector type annotations - Pleft, Pright = Zygote.Buffer.(projector_type(env.edges)) +function build_projectors(Q, env::CTMRGEnv, alg::ProjectorAlg{A,T}) where {A,T} # TODO: Add projector type annotations + P_left, P_right = Zygote.Buffer.(projector_type(env.edges)) U, V = Zygote.Buffer.(projector_type(env.edges)) S = Zygote.Buffer(env.corners) - ϵ1, ϵ2, ϵ3, ϵ4 = 0, 0, 0, 0 - for c in 1:size(env.corners, 3), r in 1:size(env.corners, 2) - Pl1, Pr1, info1 = build_projectors( - Q[1, r, c], - Q[2, r, _next(c, end)], - alg; - trunc=truncspace(space(env.edges[1, r, c], 1)), - envindex=(1, r, c), - ) - Pl2, Pr2, info2 = build_projectors( - Q[2, r, c], - Q[3, _next(r, end), c], - alg; - trunc=truncspace(space(env.edges[2, r, c], 1)), - envindex=(2, r, c), - ) - Pl3, Pr3, info3 = build_projectors( - Q[3, r, c], - Q[4, r, _prev(c, end)], - alg; - trunc=truncspace(space(env.edges[3, r, c], 1)), - envindex=(3, r, c), - ) - Pl4, Pr4, info4 = build_projectors( - Q[4, r, c], - Q[1, _prev(r, end), c], - alg; - trunc=truncspace(space(env.edges[4, r, c], 1)), - envindex=(4, r, c), - ) - - Pleft[NORTH, r, c] = Pl1 - Pright[NORTH, r, c] = Pr1 - U[NORTH, r, c] = info1.U - S[NORTH, r, c] = info1.S - V[NORTH, r, c] = info1.V - - Pleft[EAST, r, c] = Pl2 - Pright[EAST, r, c] = Pr2 - U[EAST, r, c] = info2.U - S[EAST, r, c] = info2.S - V[EAST, r, c] = info2.V - - Pleft[SOUTH, r, c] = Pl3 - Pright[SOUTH, r, c] = Pr3 - U[SOUTH, r, c] = info3.U - S[SOUTH, r, c] = info3.S - V[SOUTH, r, c] = info3.V + ϵ = 0.0 + rsize, csize = size(env.corners)[2:3] + for dir in 1:4, r in 1:rsize, c in 1:csize + # Row-column index of next enlarged corner + next_rc in if dir == 1 + (r, _next(c, csize)) + elseif dir == 2 + (_next(r, rsize), c) + elseif dir == 3 + (r, _prev(c, csize)) + elseif dir == 4 + (_prev(r, rsize), c) + end - Pleft[WEST, r, c] = Pl4 - Pright[WEST, r, c] = Pr4 - U[WEST, r, c] = info4.U - S[WEST, r, c] = info4.S - V[WEST, r, c] = info4.V - end - return copy(Pleft), copy(Pright), (; ϵ=max(ϵ1, ϵ2, ϵ3, ϵ4), U=copy(U), S=copy(S), V=copy(V)) -end -function build_projectors(Qleft, Qright, alg::ProjectorAlg; kwargs...) - # SVD half-infinite environment - U, S, V, ϵ = PEPSKit.tsvd!(Qleft * Qright, alg.svd_alg; kwargs...) - ϵ /= norm(S) + # SVD half-infinite environment + trscheme = if T <: FixedSpaceTruncation + truncspace(space(env.edges[dir, r, c], 1)) + else + alg.trscheme + end + svd_alg = if A <: SVDAdjoint{<:FixedSVD} + idx = (dir, r, c) + fwd_alg = alg.svd_alg.fwd_alg + fix_svd = FixedSVD(fwd_alg.U[idx...], fwd_alg.S[idx...], fwd_alg.V[idx...]) + return SVDAdjoint(; fwd_alg=fix_svd, rrule_alg=nothing, broadening=nothing) + else + alg.svd_alg + end + @autoopt @tensor QQ[χ_EB D_EBabove D_EBbelow; χ_ET D_ETabove D_ETbelow] := + Q[dir, r, c][χ_EB D_EBabove D_EBbelow; χ D1 D2] * + Q[_next(dir, 4), next_rc...][χ D1 D2; χ_ET D_ETabove D_ETbelow] + U, S, V, ϵ_local = PEPSKit.tsvd!(QQ, svd_alg; trunc=trscheme) + ϵ = max(ϵ, ϵ_local / norm(S)) - # Compute SVD truncation error and check for degenerate singular values - ignore_derivatives() do - if alg.verbosity > 1 && is_degenerate_spectrum(S) - svals = TensorKit.SectorDict(c => diag(b) for (c, b) in blocks(S)) - @warn("degenerate singular values detected: ", svals) + # Compute SVD truncation error and check for degenerate singular values + ignore_derivatives() do + if alg.verbosity > 0 && is_degenerate_spectrum(S) + svals = TensorKit.SectorDict(c => diag(b) for (c, b) in blocks(S)) + @warn("degenerate singular values detected: ", svals) + end end - end - # Compute projectors - isqS = sdiag_inv_sqrt(S) - @tensor Pl[-1 -2 -3; -4] := Qright[-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]) * Qleft[2 3 4; -2 -3 -4] + # Compute projectors + Pl, Pr = build_projectors(U, S, V, Q_sw, Q_nw) + P_left[dir, r, c] = Pl + P_right[dir, r, c] = Pr + U[dir, r, c] = info.U + S[dir, r, c] = info.S + V[dir, r, c] = info.V + end - return Pl, Pr, (; ϵ, U, S, V) + return copy(P_left), copy(P_right), (; ϵ, U=copy(U), S=copy(S), V=copy(V)) end # Apply projectors to renormalize corners and edges -function renormalize_corners_edges(state, env::CTMRGEnv, Q, Pleft, Pright) +function renormalize_corners_edges(state, env::CTMRGEnv, Q, P_left, P_right) corners::typeof(env.corners) = copy(env.corners) edges::typeof(env.edges) = copy(env.edges) for c in 1:size(state, 2), r in 1:size(state, 1) @@ -144,46 +125,46 @@ function renormalize_corners_edges(state, env::CTMRGEnv, Q, Pleft, Pright) cprev = _prev(c, size(state, 2)) cnext = _next(c, size(state, 2)) @diffset @tensor corners[NORTHWEST, r, c][-1; -2] := - Pright[WEST, rnext, c][-1; 1 2 3] * + P_right[WEST, rnext, c][-1; 1 2 3] * Q[NORTHWEST, r, c][1 2 3; 4 5 6] * - Pleft[NORTH, r, c][4 5 6; -2] + P_left[NORTH, r, c][4 5 6; -2] @diffset @tensor corners[NORTHEAST, r, c][-1; -2] := - Pright[NORTH, r, cprev][-1; 1 2 3] * + P_right[NORTH, r, cprev][-1; 1 2 3] * Q[NORTHEAST, r, c][1 2 3; 4 5 6] * - Pleft[EAST, r, c][4 5 6; -2] + P_left[EAST, r, c][4 5 6; -2] @diffset @tensor corners[SOUTHEAST, r, c][-1; -2] := - Pright[EAST, rprev, c][-1; 1 2 3] * + P_right[EAST, rprev, c][-1; 1 2 3] * Q[SOUTHEAST, r, c][1 2 3; 4 5 6] * - Pleft[SOUTH, r, c][4 5 6; -2] + P_left[SOUTH, r, c][4 5 6; -2] @diffset @tensor corners[SOUTHWEST, r, c][-1; -2] := - Pright[SOUTH, r, cnext][-1; 1 2 3] * + P_right[SOUTH, r, cnext][-1; 1 2 3] * Q[SOUTHWEST, r, c][1 2 3; 4 5 6] * - Pleft[WEST, r, c][4 5 6; -2] + P_left[WEST, r, c][4 5 6; -2] @diffset @tensor edges[NORTH, r, c][-1 -2 -3; -4] := env.edges[NORTH, rprev, c][1 2 3; 4] * state[r, c][9; 2 5 -2 7] * conj(state[r, c][9; 3 6 -3 8]) * - Pleft[NORTH, r, c][4 5 6; -4] * - Pright[NORTH, r, cprev][-1; 1 7 8] + P_left[NORTH, r, c][4 5 6; -4] * + P_right[NORTH, r, cprev][-1; 1 7 8] @diffset @tensor edges[EAST, r, c][-1 -2 -3; -4] := env.edges[EAST, r, _next(c, end)][1 2 3; 4] * state[r, c][9; 7 2 5 -2] * conj(state[r, c][9; 8 3 6 -3]) * - Pleft[EAST, r, c][4 5 6; -4] * - Pright[EAST, rprev, c][-1; 1 7 8] + P_left[EAST, r, c][4 5 6; -4] * + P_right[EAST, rprev, c][-1; 1 7 8] @diffset @tensor edges[SOUTH, r, c][-1 -2 -3; -4] := env.edges[SOUTH, _next(r, end), c][1 2 3; 4] * state[r, c][9; -2 7 2 5] * conj(state[r, c][9; -3 8 3 6]) * - Pleft[SOUTH, r, c][4 5 6; -4] * - Pright[SOUTH, r, cnext][-1; 1 7 8] + P_left[SOUTH, r, c][4 5 6; -4] * + P_right[SOUTH, r, cnext][-1; 1 7 8] @diffset @tensor edges[WEST, r, c][-1 -2 -3; -4] := env.edges[WEST, r, _prev(c, end)][1 2 3; 4] * state[r, c][9; 5 -2 7 2] * conj(state[r, c][9; 6 -3 8 3]) * - Pleft[WEST, r, c][4 5 6; -4] * - Pright[WEST, rnext, c][-1; 1 7 8] + P_left[WEST, r, c][4 5 6; -4] * + P_right[WEST, rnext, c][-1; 1 7 8] end @diffset corners[:, :, :] ./= norm.(corners[:, :, :]) diff --git a/src/utility/svd.jl b/src/utility/svd.jl index 3c0afdaa..2e9618a2 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -168,6 +168,24 @@ function ChainRulesCore.rrule( return (U, S, V, ϵ), tsvd!_itersvd_pullback end +""" + struct FixedSVD + +SVD struct containing a pre-computed decomposition or even multiple ones. +The call to `tsvd` just returns the pre-computed U, S and V. In the reverse +pass, the SVD adjoint is computed with these exact U, S, and V. +""" +struct FixedSVD{Ut,St,Vt} + U::Ut + S::St + V::Vt +end + +# Return pre-computed SVD +function TensorKit._tsvd!(_, alg::FixedSVD, ::NoTruncation, ::Real=2) + return alg.U, alg.S, alg.V, 0 +end + """ struct NonTruncAdjoint(; lorentz_broadening = 0.0) diff --git a/src/utility/util.jl b/src/utility/util.jl index 0797a412..e66b059b 100644 --- a/src/utility/util.jl +++ b/src/utility/util.jl @@ -51,16 +51,24 @@ end """ projector_type(T::DataType, size) + projector_type(edges::Array{<:AbstractTensorMap}) Create two arrays of specified `size` that contain undefined tensors representing left and right acting projectors, respectively. The projector types are inferred from the TensorMap type `T` which avoids having to recompute transpose tensors. +Alternatively, supply an array of edge tensors from which left and right projectors +are intialized explicitly with zeros. """ function projector_type(T::DataType, size) - Pleft = Array{T,length(size)}(undef, size) + P_left = Array{T,length(size)}(undef, size) Prtype = tensormaptype(spacetype(T), numin(T), numout(T), storagetype(T)) - Pright = Array{Prtype,length(size)}(undef, size) - return Pleft, Pright + P_right = Array{Prtype,length(size)}(undef, size) + return P_left, P_right +end +function projector_type(edges::Array{<:AbstractTensorMap}) + P_left = map(e -> TensorMap(zeros, scalartype(e), space(e)), edges) + P_right = map(e -> TensorMap(zeros, scalartype(e), domain(e), codomain(e)), edges) + return P_left, P_right end # There are no rrules for rotl90 and rotr90 in ChainRules.jl