From ab1c672cb187a18f6827dd7dfa87ce997885b29c Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 2 Oct 2024 15:03:52 +0200 Subject: [PATCH 01/21] Change x0 for IterSVD to fix element-wise convergence with U,S,V from IterSVDs --- src/utility/svd.jl | 6 +++++- test/ctmrg/fixed_iterscheme.jl | 13 ++++++++----- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/src/utility/svd.jl b/src/utility/svd.jl index 8f0abb73..eea8d04c 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -105,7 +105,11 @@ function TensorKit._compute_svddata!( Udata[c] = U[:, 1:howmany] Vdata[c] = V[1:howmany, :] else - x₀ = randn(eltype(b), size(b, 1)) + # x₀ = randn(eltype(b), size(b, 1)) # Leads to erroneous gauge fixing of U, S, V and thus failing element-wise conv. + # u, = TensorKit.MatrixAlgebra.svd!(deepcopy(b), TensorKit.SVD()) + # x₀ = sum(u[:, i] for i in 1:howmany) # Element-wise convergence works fine + # x₀ = dropdims(sum(b[:, 1:3]; dims=2); dims=2) # Summing too many columns also makes gauge fixing fail + x₀ = b[:, 1] # Leads so slower convergence of SVD than randn, but correct element-wise convergence S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, alg.alg) if info.converged < howmany # Fall back to dense SVD if not properly converged @warn "Iterative SVD did not converge for block $c, falling back to dense SVD" diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl index 1cfc9f83..4d9c0218 100644 --- a/test/ctmrg/fixed_iterscheme.jl +++ b/test/ctmrg/fixed_iterscheme.jl @@ -1,6 +1,6 @@ using Test using Random -using TensorKit +using TensorKit, KrylovKit using PEPSKit using PEPSKit: FixedSVD, @@ -46,13 +46,13 @@ unitcells = [(1, 1), (3, 4)] @test calc_elementwise_convergence(env_conv1, env_fixedsvd) ≈ 0 atol = 1e-6 end -# TODO: Why doesn't FixedSVD work with previous U, S and V from IterSVD? @testset "Element-wise consistency of TensorKit.SVD and IterSVD" begin ctm_alg_iter = CTMRG(; tol=1e-12, + maxiter=200, verbosity=2, ctmrgscheme=:simultaneous, - svd_alg=SVDAdjoint(; fwd_alg=IterSVD()), + svd_alg=SVDAdjoint(; fwd_alg=IterSVD(; alg=GKL(; tol=1e-14, krylovdim=χenv + 10))), ) ctm_alg_full = CTMRG(; tol=1e-12, @@ -65,7 +65,8 @@ end Random.seed!(91283219347) psi = InfinitePEPS(2, χbond) env_init = CTMRGEnv(psi, ComplexSpace(χenv)) - env_conv1 = leading_boundary(env_init, psi, ctm_alg_full) + env_conv1 = leading_boundary(env_init, psi, ctm_alg_iter) + # env_conv1 = leading_boundary(env_init, psi, ctm_alg_full); # do extra iteration to get SVD env_conv2_iter, info_iter = ctmrg_iter(psi, env_conv1, ctm_alg_iter) @@ -92,11 +93,13 @@ end # do iteration with FixedSVD env_fixedsvd_iter, = ctmrg_iter(psi, env_conv1, ctm_alg_fix_iter) env_fixedsvd_iter = fix_global_phases(env_conv1, env_fixedsvd_iter) - # @test calc_elementwise_convergence(env_conv1, env_fixedsvd_iter) ≈ 0 atol = 1e-6 # This should work, but doesn't! + @test calc_elementwise_convergence(env_conv1, env_fixedsvd_iter) ≈ 0 atol = 1e-6 # This doesn't work for x₀ = rand(size(b, 1))? + # @show calc_elementwise_convergence(env_conv1, env_fixedsvd_iter) env_fixedsvd_full, = ctmrg_iter(psi, env_conv1, ctm_alg_fix_full) env_fixedsvd_full = fix_global_phases(env_conv1, env_fixedsvd_full) @test calc_elementwise_convergence(env_conv1, env_fixedsvd_full) ≈ 0 atol = 1e-6 + # @show calc_elementwise_convergence(env_conv1, env_fixedsvd_full) # check matching decompositions atol = 1e-12 From 922c0f07b66eca6060e30cabf27fd8374b0216b5 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 2 Oct 2024 15:24:38 +0200 Subject: [PATCH 02/21] Remove setting restriction from PEPSOptimize --- src/algorithms/peps_opt.jl | 3 --- src/utility/svd.jl | 1 + 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/peps_opt.jl index 3b569e6b..b84467da 100644 --- a/src/algorithms/peps_opt.jl +++ b/src/algorithms/peps_opt.jl @@ -103,9 +103,6 @@ struct PEPSOptimize{G} if gradient_alg isa GradMode if S === :sequential && iterscheme(gradient_alg) === :fixed throw(ArgumentError(":sequential and :fixed are not compatible")) - elseif boundary_alg.projector_alg.svd_alg.fwd_alg isa IterSVD && - iterscheme(gradient_alg) === :fixed - throw(ArgumentError("IterSVD and :fixed are currently not compatible")) end end return new{G}(boundary_alg, optimizer, reuse_env, gradient_alg) diff --git a/src/utility/svd.jl b/src/utility/svd.jl index eea8d04c..94849daa 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -105,6 +105,7 @@ function TensorKit._compute_svddata!( Udata[c] = U[:, 1:howmany] Vdata[c] = V[1:howmany, :] else + # TODO: find better initial guess that leads to element-wise convergence # x₀ = randn(eltype(b), size(b, 1)) # Leads to erroneous gauge fixing of U, S, V and thus failing element-wise conv. # u, = TensorKit.MatrixAlgebra.svd!(deepcopy(b), TensorKit.SVD()) # x₀ = sum(u[:, i] for i in 1:howmany) # Element-wise convergence works fine From 826563ca6157dc7cdf18f70052ab07a930f2d681 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Wed, 2 Oct 2024 17:37:56 +0200 Subject: [PATCH 03/21] Implement parts of function handle CTMRG --- src/PEPSKit.jl | 1 + .../contractions/ctmrg_contractions.jl | 6 ++ src/algorithms/ctmrg/ctmrg.jl | 14 ++-- src/algorithms/ctmrg/sparse_environments.jl | 79 +++++++++++++++++++ test/ctmrg/fixed_iterscheme.jl | 2 +- 5 files changed, 96 insertions(+), 6 deletions(-) create mode 100644 src/algorithms/ctmrg/sparse_environments.jl diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 5a4fdb0a..72bb6da6 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -37,6 +37,7 @@ include("algorithms/contractions/localoperator.jl") include("algorithms/contractions/ctmrg_contractions.jl") include("algorithms/ctmrg/ctmrg.jl") +include("algorithms/ctmrg/sparse_environments.jl") include("algorithms/ctmrg/gaugefix.jl") include("algorithms/toolbox.jl") diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 69e4c9e4..a7ac85ca 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -165,6 +165,7 @@ end """ halfinfinite_environment(quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3}) + halfinfinite_environment(quadrant1::EnlargedCorner{A,C,E}, quadrant2::EnlargedCorner{A,C,E}) Contract two quadrants (enlarged corners) to form a half-infinite environment. @@ -182,6 +183,11 @@ function halfinfinite_environment( quadrant1[χ_in D_inabove D_inbelow; χ D1 D2] * quadrant2[χ D1 D2; χ_out D_outabove D_outbelow] end +function halfinfinite_environment( + quadrant1::EnlargedCorner{A,C,E}, quadrant2::EnlargedCorner{A,C,E} +) where {A,C,E} + return HalfInfiniteEnv(quadrant1, quadrant2)() +end # Renormalization contractions # ---------------------------- diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index bba37f25..3ebb0ac1 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -195,7 +195,7 @@ function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SequentialCTMRG) where {C,T} Q_nw[r, c] = enlarge_northwest_corner((r, c), envs, state) end - return copy(Q_sw), copy(Q_nw) + return stack((copy(Q_sw), copy(Q_nw)); dims=1) end function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SimultaneousCTMRG) where {C,T} Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) @@ -238,7 +238,7 @@ function ctmrg_projectors( for (r, c) in directions # SVD half-infinite environment r′ = _prev(r, size(envs.corners, 2)) - QQ = halfinfinite_environment(enlarged_envs[1][r, c], enlarged_envs[2][r′, c]) + QQ = halfinfinite_environment(enlarged_envs[1, r, c], enlarged_envs[2, r′, c]) trscheme = truncation_scheme(projector_alg, envs.edges[WEST, r′, c]) svd_alg = svd_algorithm(projector_alg, (WEST, r, c)) @@ -255,7 +255,7 @@ function ctmrg_projectors( # Compute projectors P_bottom[r, c], P_top[r, c] = build_projectors( - U, S, V, enlarged_envs[1][r, c], enlarged_envs[2][r′, c] + U, S, V, enlarged_envs[1, r, c], enlarged_envs[2, r′, c] ) end @@ -391,9 +391,13 @@ end # Auxiliary routines # ======================================================================================== # -# Build projectors from SVD and enlarged SW & NW corners +# Build projectors from SVD and dense enlarged corners function build_projectors( - U::AbstractTensorMap{E,3,1}, S, V::AbstractTensorMap{E,1,3}, Q, Q_next + U::AbstractTensorMap{E,3,1}, + S::AbstractTensorMap{E,1,1}, + V::AbstractTensorMap{E,1,3}, + Q::AbstractTensorMap{E,3,3}, + Q_next::AbstractTensorMap{E,3,3}, ) where {E<:ElementarySpace} isqS = sdiag_inv_sqrt(S) P_left = Q_next * V' * isqS diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl new file mode 100644 index 00000000..9f5b0c26 --- /dev/null +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -0,0 +1,79 @@ +""" + struct HalfInfiniteEnv{A,C,E} + +Half-infinite CTMRG environment tensor storage. + +``` + C_2 -- E_2 -- E_3 -- C_3 + | || || | + E_1 == ket_bra_1 == ket_bra_2 == E_4 + | || || | +``` +""" +struct HalfInfiniteEnv{A,A′,C,E} + ket_1::A + bra_1::A′ + ket_2::A + bra_2::A′ + C_1::C + C_2::C + E_1::E + E_2::E + E_3::E + E_4::E +end + +# Construct environment from two enlarged corners +function HalfInfiniteEnv(quadrant1::EnlargedCorner, quadrant2::EnlargedCorner) + return HalfInfiniteEnv( + quadrant1.ket_bra, + quadrant2.ket_bra, + quadrant1.C, + quadrant2.C, + quadrant1.E_1, + quadrant1.E_2, + quadrant2.E_1, + quadrant2.E_2, + ) +end + +# Contract half-infinite environment +function (env::HalfInfiniteEnv)() end + +""" + struct EnlargedCorner{A,C,E} + +Enlarged CTMRG corner tensor storage. + +``` + C -- E_2 -- + | || + E_1 == ket-bra == + | || +``` +""" +struct EnlargedCorner{A,A′,Ct,E} + ket::A + bra::A′ + C::Ct + E_1::E + E_2::E +end + +# Contract enlarged corner +function (Q::EnlargedCorner)() + return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +end + +# Compute left & right projectors from enlarged corner struct +function build_projectors( + U::AbstractTensorMap{E,3,1}, + S::AbstractTensorMap{E,1,1}, + V::AbstractTensorMap{E,1,3}, + Q::EnlargedCorner, + Q_next::EnlargedCorner, +) where {E<:ElementarySpace} + isqS = sdiag_inv_sqrt(S) + # TODO + return P_left, P_right +end \ No newline at end of file diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl index 4d9c0218..871e8711 100644 --- a/test/ctmrg/fixed_iterscheme.jl +++ b/test/ctmrg/fixed_iterscheme.jl @@ -13,7 +13,7 @@ using PEPSKit: # initialize parameters χbond = 2 χenv = 16 -svd_algs = [SVDAdjoint(; fwd_alg=TensorKit.SVD())] #, SVDAdjoint(; fwd_alg=IterSVD())] +svd_algs = [SVDAdjoint(; fwd_alg=TensorKit.SVD()), SVDAdjoint(; fwd_alg=IterSVD())] unitcells = [(1, 1), (3, 4)] # test for element-wise convergence after application of fixed step From 0afafe6597fb7c1d8db1783e1635d4008b3515c5 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 3 Oct 2024 11:12:12 +0200 Subject: [PATCH 04/21] Add left and right projector contractions --- .../contractions/ctmrg_contractions.jl | 47 +++++++++++++++++++ src/algorithms/ctmrg/sparse_environments.jl | 9 ++-- 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index a7ac85ca..636a3b85 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -163,6 +163,53 @@ end # Projector contractions # ---------------------- +""" + left_projector(E_1, C, E_2, V, isqS, ket::PEPSTensor, bra::PEPSTensor=ket) + +Contract the CTMRG left projector with the higher-dimensional subspace facing to the left. + +``` + C -- E_2 -- |~~| + | || |V'| -- isqS -- + E_1 == ket-bra == |~~| + | || +``` +""" +function left_projector(E_1, C, E_2, V, isqS, ket::PEPSTensor, bra::PEPSTensor=ket) + return @autoopt @tensor P_left[χ_in D_inabove D_inbelow; χ_out] := + E_1[χ_in D1 D2; χ1] * + C[χ1; χ2] * + E_2[χ2 D3 D4; χ3] * + ket[d; D3 D5 D_inabove D1] * + conj(bra[d; D4 D6 D_inbelow D2]) * + conj(V[χ4; χ3 D5 D6]) * + isqS[χ4; χ_out] +end + +""" + right_projector(E_1, C, E_2, U, isqS, ket::PEPSTensor, bra::PEPSTensor=ket) + +Contract the CTMRG right projector with the higher-dimensional subspace facing to the right. + +``` + |~~| -- E_2 -- C + -- isqS -- |U'| || | + |~~| == ket-bra == E_1 + || | +``` +""" +function right_projector(E_1, C, E_2, U, isqS, ket::PEPSTensor, bra::PEPSTensor=ket) + return @autoopt @tensor P_right[χ_in; χ_out D_outabove D_outbelow] := + isqS[χ_in; χ1] * + conj(U[χ1; χ2 D1 D2]) * + ket[d; D1 D5 D_outabove D1] * + conj(bra[d; D2 D6 D_outbelow D2]) * + E_2[χ2 D3 D4; χ3] * + C[χ3; χ4] * + E_1[χ4 D5 D6; χ_out] +end + + """ halfinfinite_environment(quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3}) halfinfinite_environment(quadrant1::EnlargedCorner{A,C,E}, quadrant2::EnlargedCorner{A,C,E}) diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index 9f5b0c26..ec086e87 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -60,7 +60,7 @@ struct EnlargedCorner{A,A′,Ct,E} E_2::E end -# Contract enlarged corner +# Contract enlarged corner (use NW corner as convention for connecting environment to PEPS tensor) function (Q::EnlargedCorner)() return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) end @@ -74,6 +74,9 @@ function build_projectors( Q_next::EnlargedCorner, ) where {E<:ElementarySpace} isqS = sdiag_inv_sqrt(S) - # TODO + P_left = left_projector(Q.E_1, Q.C, Q.E_2, V, isqS, Q.ket, Q.bra) + P_right = right_projector( + Q_next.E_1, Q_next.C, Q_next.E_2, U, isqS, Q_next.ket, Q_next.bra + ) return P_left, P_right -end \ No newline at end of file +end From 4fa2cd35a0db7ed28c77790badbd135c591800ab Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 3 Oct 2024 11:56:34 +0200 Subject: [PATCH 05/21] Add half-infinite environment contractions --- .../contractions/ctmrg_contractions.jl | 65 ++++++++++++++++--- src/algorithms/ctmrg/sparse_environments.jl | 44 ++++++++++--- 2 files changed, 90 insertions(+), 19 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 636a3b85..1e533c78 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -177,9 +177,9 @@ Contract the CTMRG left projector with the higher-dimensional subspace facing to """ function left_projector(E_1, C, E_2, V, isqS, ket::PEPSTensor, bra::PEPSTensor=ket) return @autoopt @tensor P_left[χ_in D_inabove D_inbelow; χ_out] := - E_1[χ_in D1 D2; χ1] * - C[χ1; χ2] * - E_2[χ2 D3 D4; χ3] * + E_1[χ_in D1 D2; χ1] * + C[χ1; χ2] * + E_2[χ2 D3 D4; χ3] * ket[d; D3 D5 D_inabove D1] * conj(bra[d; D4 D6 D_inbelow D2]) * conj(V[χ4; χ3 D5 D6]) * @@ -204,15 +204,14 @@ function right_projector(E_1, C, E_2, U, isqS, ket::PEPSTensor, bra::PEPSTensor= conj(U[χ1; χ2 D1 D2]) * ket[d; D1 D5 D_outabove D1] * conj(bra[d; D2 D6 D_outbelow D2]) * - E_2[χ2 D3 D4; χ3] * - C[χ3; χ4] * + E_2[χ2 D3 D4; χ3] * + C[χ3; χ4] * E_1[χ4 D5 D6; χ_out] end - """ halfinfinite_environment(quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3}) - halfinfinite_environment(quadrant1::EnlargedCorner{A,C,E}, quadrant2::EnlargedCorner{A,C,E}) + Contract two quadrants (enlarged corners) to form a half-infinite environment. @@ -222,6 +221,26 @@ Contract two quadrants (enlarged corners) to form a half-infinite environment. |~~~~~~~~~| == |~~~~~~~~~| | || || | ``` + +The environment can also be contracted directly from all its constituent tensors. + +``` + C_1 -- E_2 -- E_3 -- C_2 + | || || | + E_1 == ket_bra_1 == ket_bra_2 == E_4 + | || || | +``` + +Alternatively, contract environment with a vector `x` acting on it, e.g. as needed for iterative solvers. + +``` + C_1 -- E_2 -- E_3 -- C_2 + | || || | + E_1 == ket_bra_1 == ket_bra_2 == E_4 + | || || | + [~~~~~~x~~~~~~] + || | +``` """ function halfinfinite_environment( quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3} @@ -231,9 +250,35 @@ function halfinfinite_environment( quadrant2[χ D1 D2; χ_out D_outabove D_outbelow] end function halfinfinite_environment( - quadrant1::EnlargedCorner{A,C,E}, quadrant2::EnlargedCorner{A,C,E} -) where {A,C,E} - return HalfInfiniteEnv(quadrant1, quadrant2)() + E_1, C_1, E_2, E_3, C_2, E_4, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 +) where {P<:PEPSTensor} + return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := + E_1[χ_in D1 D2; χ1] * + C_1[χ1; χ2] * + E_2[χ2 D3 D4; χ3] * + ket_1[d1; D3 D9 D_inabove D1] * + conj(bra_1[d1; D4 D10 D_inbelow D2]) * + ket_2[d2; D5 D7 D_outabove D9] * + conj(bra_2[d2; D6 D8 D_outbelow D10]) * + E_3[χ3 D5 D6; χ4] * + C_2[χ4; χ5] * + E_4[χ5 D7 D8; χ_out] +end +function halfinfinite_environment( + E_1, C_1, E_2, E_3, C_2, E_4, x, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 +) where {P<:PEPSTensor} + return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := + E_1[χ_in D1 D2; χ1] * + C_1[χ1; χ2] * + E_2[χ2 D3 D4; χ3] * + ket_1[d1; D3 D9 D_inabove D1] * + conj(bra_1[d1; D4 D10 D_inbelow D2]) * + ket_2[d2; D5 D7 D11 D9] * + conj(bra_2[d2; D6 D8 D12 D10]) * + E_3[χ3 D5 D6; χ4] * + C_2[χ4; χ5] * + E_4[χ5 D7 D8; χ6] * + x[χ6 D11 D12; χ_out D_outabove D_outbelow] end # Renormalization contractions diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index ec086e87..0db09838 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -2,13 +2,6 @@ struct HalfInfiniteEnv{A,C,E} Half-infinite CTMRG environment tensor storage. - -``` - C_2 -- E_2 -- E_3 -- C_3 - | || || | - E_1 == ket_bra_1 == ket_bra_2 == E_4 - | || || | -``` """ struct HalfInfiniteEnv{A,A′,C,E} ket_1::A @@ -37,8 +30,41 @@ function HalfInfiniteEnv(quadrant1::EnlargedCorner, quadrant2::EnlargedCorner) ) end -# Contract half-infinite environment -function (env::HalfInfiniteEnv)() end +""" + (env::HalfInfiniteEnv)() + (env::HalfInfiniteEnv)(x) + +Contract half-infinite environment without or with a vector `x`. +""" +function (env::HalfInfiniteEnv)() + return halfinfinite_environment( + env.E_1, + env.C_1, + env.E_2, + env.E_3, + env.C_2, + env.E_4, + env.ket_1, + env.ket_2, + env.bra_1, + env.bra_2, + ) +end +function (env::HalfInfiniteEnv)(x) + return halfinfinite_environment( + env.E_1, + env.C_1, + env.E_2, + env.E_3, + env.C_2, + env.E_4, + x, + env.ket_1, + env.ket_2, + env.bra_1, + env.bra_2, + ) +end """ struct EnlargedCorner{A,C,E} From 58b92421adbf0995130ca8947716389707092be9 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 3 Oct 2024 12:06:34 +0200 Subject: [PATCH 06/21] Add EnlargedCorner directional contraction methods, fix docstrings --- src/algorithms/ctmrg/sparse_environments.jl | 25 ++++++++++----------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index 0db09838..3840a24e 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -1,5 +1,5 @@ """ - struct HalfInfiniteEnv{A,C,E} + struct HalfInfiniteEnv{A,A′,C,E} Half-infinite CTMRG environment tensor storage. """ @@ -67,16 +67,9 @@ function (env::HalfInfiniteEnv)(x) end """ - struct EnlargedCorner{A,C,E} + struct EnlargedCorner{A,A′,Ct,E} Enlarged CTMRG corner tensor storage. - -``` - C -- E_2 -- - | || - E_1 == ket-bra == - | || -``` """ struct EnlargedCorner{A,A′,Ct,E} ket::A @@ -86,10 +79,16 @@ struct EnlargedCorner{A,A′,Ct,E} E_2::E end -# Contract enlarged corner (use NW corner as convention for connecting environment to PEPS tensor) -function (Q::EnlargedCorner)() - return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) -end +""" + (Q::EnlargedCorner)(::Val{<:Int}) + +Contract enlarged corner where `Val(1)` dispatches the north-west, `Val(2)` the north-east +`Val(3)` the south-east and `Val(4)` the south-west contraction. +""" +(Q::EnlargedCorner)(::Val{1}) = enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +(Q::EnlargedCorner)(::Val{2}) = enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +(Q::EnlargedCorner)(::Val{3}) = enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +(Q::EnlargedCorner)(::Val{4}) = enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) # Compute left & right projectors from enlarged corner struct function build_projectors( From 1d91891a75a8fce7a2d8f2c07589ea1a4008e68b Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 3 Oct 2024 18:54:05 +0200 Subject: [PATCH 07/21] Further implement function handle details --- .../contractions/ctmrg_contractions.jl | 9 +- src/algorithms/ctmrg/ctmrg.jl | 4 + src/algorithms/ctmrg/sparse_environments.jl | 137 ++++++++++++------ src/utility/svd.jl | 66 ++++++--- test/ctmrg/svd_wrapper.jl | 31 ++++ 5 files changed, 177 insertions(+), 70 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 1e533c78..942f292a 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -211,7 +211,10 @@ end """ halfinfinite_environment(quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3}) - + function halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, + ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor} + function halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, x, + ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor} Contract two quadrants (enlarged corners) to form a half-infinite environment. @@ -250,7 +253,7 @@ function halfinfinite_environment( quadrant2[χ D1 D2; χ_out D_outabove D_outbelow] end function halfinfinite_environment( - E_1, C_1, E_2, E_3, C_2, E_4, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 + C_1, C_2, E_1, E_2, E_3, E_4, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 ) where {P<:PEPSTensor} return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := E_1[χ_in D1 D2; χ1] * @@ -265,7 +268,7 @@ function halfinfinite_environment( E_4[χ5 D7 D8; χ_out] end function halfinfinite_environment( - E_1, C_1, E_2, E_3, C_2, E_4, x, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 + C_1, C_2, E_1, E_2, E_3, E_4, x, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 ) where {P<:PEPSTensor} return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := E_1[χ_in D1 D2; χ1] * diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 3ebb0ac1..5520673e 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -202,6 +202,7 @@ function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SimultaneousCTMRG) where {C, Q = Zygote.Buffer(Array{Qtype,3}(undef, size(envs.corners))) drc_combinations = collect(Iterators.product(axes(envs.corners)...)) @fwdthreads for (dir, r, c) in drc_combinations + # TODO: generalize the corner computation to make it compatible with EnlargedCorners Q[dir, r, c] = if dir == NORTHWEST enlarge_northwest_corner((r, c), envs, state) elseif dir == NORTHEAST @@ -216,6 +217,9 @@ function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SimultaneousCTMRG) where {C, return copy(Q) end +function enlarge_corner((dir, r, c), envs, state, alg::CTMRG) # TODO: find a way to dispatch on CTMRG struct to switch on function handles +end + # ======================================================================================== # # Projector step # ======================================================================================== # diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index 3840a24e..783ef429 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -1,32 +1,73 @@ + +""" + struct EnlargedCorner{Ct,E,A,A′} + +Enlarged CTMRG corner tensor storage. +""" +struct EnlargedCorner{Ct,E,A,A′} + C::Ct + E_1::E + E_2::E + ket::A + bra::A′ +end + """ - struct HalfInfiniteEnv{A,A′,C,E} + (Q::EnlargedCorner)(::Val{<:Int}) + +Contract enlarged corner where `Val(1)` dispatches the north-west, `Val(2)` the north-east +`Val(3)` the south-east and `Val(4)` the south-west contraction. +""" +(Q::EnlargedCorner)(::Val{1}) = enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +(Q::EnlargedCorner)(::Val{2}) = enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +(Q::EnlargedCorner)(::Val{3}) = enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +(Q::EnlargedCorner)(::Val{4}) = enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + +# Compute left & right projectors from enlarged corner struct +function build_projectors( + U::AbstractTensorMap{E,3,1}, + S::AbstractTensorMap{E,1,1}, + V::AbstractTensorMap{E,1,3}, + Q::EnlargedCorner, + Q_next::EnlargedCorner, +) where {E<:ElementarySpace} + isqS = sdiag_inv_sqrt(S) + P_left = left_projector(Q.E_1, Q.C, Q.E_2, V, isqS, Q.ket, Q.bra) + P_right = right_projector( + Q_next.E_1, Q_next.C, Q_next.E_2, U, isqS, Q_next.ket, Q_next.bra + ) + return P_left, P_right +end + +""" + struct HalfInfiniteEnv{C,E,A,A′} Half-infinite CTMRG environment tensor storage. """ -struct HalfInfiniteEnv{A,A′,C,E} - ket_1::A - bra_1::A′ - ket_2::A - bra_2::A′ +struct HalfInfiniteEnv{C,E,A,A′} C_1::C C_2::C E_1::E E_2::E E_3::E E_4::E + ket_1::A + bra_1::A′ + ket_2::A + bra_2::A′ end # Construct environment from two enlarged corners function HalfInfiniteEnv(quadrant1::EnlargedCorner, quadrant2::EnlargedCorner) return HalfInfiniteEnv( - quadrant1.ket_bra, - quadrant2.ket_bra, quadrant1.C, quadrant2.C, quadrant1.E_1, quadrant1.E_2, quadrant2.E_1, quadrant2.E_2, + quadrant1.ket_bra, + quadrant2.ket_bra, ) end @@ -38,11 +79,11 @@ Contract half-infinite environment without or with a vector `x`. """ function (env::HalfInfiniteEnv)() return halfinfinite_environment( - env.E_1, env.C_1, + env.C_2, + env.E_1, env.E_2, env.E_3, - env.C_2, env.E_4, env.ket_1, env.ket_2, @@ -52,11 +93,11 @@ function (env::HalfInfiniteEnv)() end function (env::HalfInfiniteEnv)(x) return halfinfinite_environment( - env.E_1, env.C_1, + env.C_2, + env.E_1, env.E_2, env.E_3, - env.C_2, env.E_4, x, env.ket_1, @@ -66,42 +107,50 @@ function (env::HalfInfiniteEnv)(x) ) end -""" - struct EnlargedCorner{A,A′,Ct,E} +# TensorKit methods to make struct compatible with sparse SVD +TensorKit.InnerProductStyle(::HalfInfiniteEnv) = EuclideanProduct() +TensorKit.sectortype(::HalfInfiniteEnv) = Trivial +TensorKit.storagetype(env::HalfInfiniteEnv) = storagetype(env.ket_1) +TensorKit.spacetype(env::HalfInfiniteEnv) = spacetype(env.ket_1) -Enlarged CTMRG corner tensor storage. -""" -struct EnlargedCorner{A,A′,Ct,E} - ket::A - bra::A′ - C::Ct - E_1::E - E_2::E +function TensorKit.blocks(env::HalfInfiniteEnv) + return TensorKit.SingletonDict(Trivial() => env) +end +function TensorKit.blocksectors(::HalfInfiniteEnv) + return TensorKit.OneOrNoneIterator{Trivial}(true, Trivial()) end -""" - (Q::EnlargedCorner)(::Val{<:Int}) +function TensorKit.MatrixAlgebra.svd!(env::HalfInfiniteEnv, args...) + return TensorKit.MatrixAlgebra.svd!(env(), args...) +end -Contract enlarged corner where `Val(1)` dispatches the north-west, `Val(2)` the north-east -`Val(3)` the south-east and `Val(4)` the south-west contraction. -""" -(Q::EnlargedCorner)(::Val{1}) = enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) -(Q::EnlargedCorner)(::Val{2}) = enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) -(Q::EnlargedCorner)(::Val{3}) = enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) -(Q::EnlargedCorner)(::Val{4}) = enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +Base.eltype(env::HalfInfiniteEnv) = eltype(env.ket_1) +function Base.size(env::HalfInfiniteEnv) # Treat environment as matrix + χ_in = dim(space(env.E_1, 1)) + D_inabove = dim(space(env.ket_1, 2)) + D_inbelow = dim(space(env.bra_1, 2)) + χ_out = dim(space(env.E_4, 1)) + D_outabove = dim(space(env.ket_2, 2)) + D_outbelow = dim(space(env.bra_2, 2)) + return (χ_in * D_inabove * D_inbelow, χ_out * D_outabove * D_outbelow) +end +Base.size(env::HalfInfiniteEnv, i::Int) = size(env)[i] -# Compute left & right projectors from enlarged corner struct -function build_projectors( - U::AbstractTensorMap{E,3,1}, - S::AbstractTensorMap{E,1,1}, - V::AbstractTensorMap{E,1,3}, - Q::EnlargedCorner, - Q_next::EnlargedCorner, -) where {E<:ElementarySpace} - isqS = sdiag_inv_sqrt(S) - P_left = left_projector(Q.E_1, Q.C, Q.E_2, V, isqS, Q.ket, Q.bra) - P_right = right_projector( - Q_next.E_1, Q_next.C, Q_next.E_2, U, isqS, Q_next.ket, Q_next.bra +# TODO: implement VectorInterface +VectorInterface.scalartype(env::HalfInfiniteEnv) = scalartype(env.ket_1) + +# Wrapper around halfinfinite_environment contraction using EnlargedCorners (used in ctmrg_projectors) +function halfinfinite_environment(ec_1::EnlargedCorner, ec_2::EnlargedCorner) + return HalfInfiniteEnv( + ec_1.C, + ec_2.C, + ec_1.E_1, + ec_1.E_2, + ec_2.E_1, + ec_2.E_2, + ec_1.ket, + ec_2.ket, + ec_1.bra, + ec_2.bra, ) - return P_left, P_right end diff --git a/src/utility/svd.jl b/src/utility/svd.jl index 94849daa..54765e20 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -25,19 +25,20 @@ removes the divergences from the adjoint. end # Keep truncation algorithm separate to be able to specify CTMRG dependent information """ - PEPSKit.tsvd(t::AbstractTensorMap, alg; trunc=notrunc(), p=2) + PEPSKit.tsvd(t, alg; trunc=notrunc(), p=2) Wrapper around `TensorKit.tsvd` which dispatches on the `alg` argument. This is needed since a custom adjoint for `PEPSKit.tsvd` may be defined, depending on the algorithm. E.g., for `IterSVD` the adjoint for a truncated SVD from `KrylovKit.svdsolve` is used. """ -PEPSKit.tsvd(t::AbstractTensorMap, alg; kwargs...) = PEPSKit.tsvd!(copy(t), alg; kwargs...) -function PEPSKit.tsvd!( - t::AbstractTensorMap, alg::SVDAdjoint; trunc::TruncationScheme=notrunc(), p::Real=2 -) +PEPSKit.tsvd(t, alg; kwargs...) = PEPSKit.tsvd!(copy(t), alg; kwargs...) +function PEPSKit.tsvd!(t, alg::SVDAdjoint; trunc::TruncationScheme=notrunc(), p::Real=2) return TensorKit.tsvd!(t; alg=alg.fwd_alg, trunc, p) end +function TensorKit.tsvd!(f::HalfInfiniteEnv; trunc=NoTruncation(), p::Real=2, alg=IterSVD()) + return _tsvd!(f, alg, trunc, p) +end """ struct FixedSVD @@ -69,35 +70,45 @@ the iterative SVD didn't converge, the algorithm falls back to a dense SVD. @kwdef struct IterSVD alg::KrylovKit.GKL = KrylovKit.GKL(; tol=1e-14, krylovdim=25) fallback_threshold::Float64 = Inf + start_vector = random_start_vector +end + +# TODO: find better initial guess that leads to element-wise convergence and is compatible with function handles +function random_start_vector(f) + return randn(eltype(f), size(f, 1)) # Leads to erroneous gauge fixing of U, S, V and thus failing element-wise conv. + # u, = TensorKit.MatrixAlgebra.svd!(deepcopy(t), TensorKit.SVD()) + # return sum(u[:, i] for i in 1:howmany) # Element-wise convergence works fine + # return dropdims(sum(t[:, 1:3]; dims=2); dims=2) # Summing too many columns also makes gauge fixing fail + # return t[:, 1] # Leads so slower convergence of SVD than randn, but correct element-wise convergence end # Compute SVD data block-wise using KrylovKit algorithm function TensorKit._tsvd!( - t, alg::IterSVD, trunc::Union{NoTruncation,TruncationSpace}, p::Real=2 + f, alg::IterSVD, trunc::Union{NoTruncation,TruncationSpace}, p::Real=2 ) # early return - if isempty(blocksectors(t)) - truncerr = zero(real(scalartype(t))) - return _empty_svdtensors(t)..., truncerr + if isempty(blocksectors(f)) + truncerr = zero(real(scalartype(f))) + return _empty_svdtensors(f)..., truncerr end - Udata, Σdata, Vdata, dims = _compute_svddata!(t, alg, trunc) - U, S, V = _create_svdtensors(t, Udata, Σdata, Vdata, spacetype(t)(dims)) - truncerr = trunc isa NoTruncation ? abs(zero(scalartype(t))) : norm(U * S * V - t, p) + Udata, Σdata, Vdata, dims = _compute_svddata!(f, alg, trunc) + U, S, V = _create_svdtensors(f, Udata, Σdata, Vdata, spacetype(f)(dims)) + truncerr = trunc isa NoTruncation ? abs(zero(scalartype(f))) : norm(U * S * V - f, p) return U, S, V, truncerr end function TensorKit._compute_svddata!( - t::TensorMap, alg::IterSVD, trunc::Union{NoTruncation,TruncationSpace} + f, alg::IterSVD, trunc::Union{NoTruncation,TruncationSpace} ) - InnerProductStyle(t) === EuclideanProduct() || throw_invalid_innerproduct(:tsvd!) - I = sectortype(t) - A = storagetype(t) + InnerProductStyle(f) === EuclideanProduct() || throw_invalid_innerproduct(:tsvd!) + I = sectortype(f) + A = storagetype(f) Udata = SectorDict{I,A}() Vdata = SectorDict{I,A}() dims = SectorDict{I,Int}() local Sdata - for (c, b) in blocks(t) + for (c, b) in blocks(f) howmany = trunc isa NoTruncation ? minimum(size(b)) : blockdim(trunc.space, c) if howmany / minimum(size(b)) > alg.fallback_threshold # Use dense SVD for small blocks @@ -105,12 +116,7 @@ function TensorKit._compute_svddata!( Udata[c] = U[:, 1:howmany] Vdata[c] = V[1:howmany, :] else - # TODO: find better initial guess that leads to element-wise convergence - # x₀ = randn(eltype(b), size(b, 1)) # Leads to erroneous gauge fixing of U, S, V and thus failing element-wise conv. - # u, = TensorKit.MatrixAlgebra.svd!(deepcopy(b), TensorKit.SVD()) - # x₀ = sum(u[:, i] for i in 1:howmany) # Element-wise convergence works fine - # x₀ = dropdims(sum(b[:, 1:3]; dims=2); dims=2) # Summing too many columns also makes gauge fixing fail - x₀ = b[:, 1] # Leads so slower convergence of SVD than randn, but correct element-wise convergence + x₀ = alg.start_vector(b) S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, alg.alg) if info.converged < howmany # Fall back to dense SVD if not properly converged @warn "Iterative SVD did not converge for block $c, falling back to dense SVD" @@ -134,6 +140,7 @@ function TensorKit._compute_svddata!( return Udata, Sdata, Vdata, dims end +# Rrule with custom pullback to make KrylovKit rrule compatible with TensorMap symmetry blocks function ChainRulesCore.rrule( ::typeof(PEPSKit.tsvd!), t::AbstractTensorMap, @@ -194,6 +201,19 @@ function ChainRulesCore.rrule( return (U, S, V, ϵ), tsvd!_itersvd_pullback end +# Separate rule for SVD with function handle that uses KrylovKit.make_svdsolve_pullback +# but in turn cannot handle symmetric blocks +function ChainRulesCore.rrule( + ::typeof(PEPSKit.tsvd!), + f, + alg::SVDAdjoint{F,R,B}; + trunc::TruncationScheme=notrunc(), + p::Real=2, +) where {F<:Union{IterSVD,FixedSVD},R<:Union{GMRES,BiCGStab,Arnoldi},B} + return U, S, V, ϵ = PEPSKit.tsvd(t, alg; trunc, p) + # TODO: implement function handle adjoint wrapper with KrylovKit.make_svdsolve_pullback +end + """ struct NonTruncAdjoint diff --git a/test/ctmrg/svd_wrapper.jl b/test/ctmrg/svd_wrapper.jl index 47e1b9cf..927a1ed3 100644 --- a/test/ctmrg/svd_wrapper.jl +++ b/test/ctmrg/svd_wrapper.jl @@ -6,6 +6,7 @@ using KrylovKit using ChainRulesCore, Zygote using Accessors using PEPSKit +using PEPSKit: HalfInfiniteEnv # Gauge-invariant loss function function lossfun(A, alg, R=TensorMap(randn, space(A)), trunc=notrunc()) @@ -92,3 +93,33 @@ symm_R = TensorMap(randn, dtype, space(symm_r)) @test l_itersvd_fb ≈ l_fullsvd_tr @test g_fullsvd_tr[1] ≈ g_itersvd_fb[1] rtol = rtol end + +χbond = 2 +χenv = 6 +ctm_alg = CTMRG(; tol=1e-10, verbosity=2, svd_alg=SVDAdjoint()) +Random.seed!(91283219347) +H = heisenberg_XYZ(InfiniteSquare()) +psi = InfinitePEPS(2, χbond) +env = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg) +hienv = HalfInfiniteEnv( + env.corners[1], + env.corners[2], + env.edges[4], + env.edges[1], + env.edges[1], + env.edges[2], + psi[1], + psi[1], + psi[1], + psi[1], +) +hienv_dense = hienv() +env_R = TensorMap(randn, space(hienv_dense)) +PEPSKit.tsvd!(hienv, iter_alg) + +@testset "IterSVD with HalfInfiniteEnv function handle" begin + l_fullsvd, g_fullsvd = withgradient(A -> lossfun(A, full_alg, env_R), hienv_dense) + l_itersvd, g_itersvd = withgradient(A -> lossfun(A, iter_alg, env_R), hienv) + @test l_itersvd ≈ l_fullsvd + @test g_fullsvd[1] ≈ g_itersvd[1] rtol = rtol +end From 0cb1ccfaf6355bbf445b25321ab569d904f25408 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 4 Oct 2024 10:34:23 +0200 Subject: [PATCH 08/21] Add enlarge_corner function to enable sparse CTMRG dispatch, add renormalize_corner contraction --- .../contractions/ctmrg_contractions.jl | 26 +++++++- src/algorithms/ctmrg/ctmrg.jl | 61 ++++++++----------- src/algorithms/ctmrg/sparse_environments.jl | 46 ++++++++++++++ 3 files changed, 98 insertions(+), 35 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 942f292a..67317ed6 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -290,7 +290,8 @@ end # corners """ - renormalize_corner(quadrant, P_left, P_right) + renormalize_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) + renormalize_corner(E_1, C, E_2, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket) Apply projectors to each side of a quadrant. @@ -302,11 +303,34 @@ Apply projectors to each side of a quadrant. [P_right] | ``` + +Alternatively, provide the constituent tensors and perform the complete contraction. + +``` + C -- E_2 -- |~~~~~~| + | | |P_left| -- + E_1 == ket-bra == |~~~~~~| + | || + [~~P_right~~] + | +``` """ function renormalize_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) where {S} return @autoopt @tensor corner[χ_in; χ_out] := P_right[χ_in; χ1 D1 D2] * quadrant[χ1 D1 D2; χ2 D3 D4] * P_left[χ2 D3 D4; χ_out] end +function renormalize_corner( + E_1, C, E_2, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket +) + return @autoopt @tensor corner[χ_in; χ_out] := + P_right[χ_in; χ1 D1 D2] * + E_1[χ1 D3 D4; χ2] * + C[χ2; χ3] * + E_2[χ3 D5 D6; χ4] * + ket[d; D5 D7 D1 D3] * + conj(bra[d; D6 D8 D2 D4]) * + P_left[χ4 D7 D8; χ_out] +end """ renormalize_northwest_corner((row, col), enlarged_envs::CTMRGEnv, P_left, P_right) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 5520673e..af84e95f 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -175,6 +175,27 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N) # Expansion step # ======================================================================================== # +# TODO: find a way to dispatch on CTMRG struct to switch on function handles +""" + enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::CTMRG) + +Enlarge a CTMRG corner into the one of four directions `Val(1)` (NW), `Val(2)` (NE), +`Val(3)` (SE) or `Val(4)` (SW). This serves as a wrapper that can dispatch on the +directional index and different `CTMRG` algorithms. +""" +function enlarge_corner(::Val{1}, (r, c), envs, state, alg::CTMRG) + return enlarge_northwest_corner((r, c), envs, state) +end +function enlarge_corner(::Val{2}, (r, c), envs, state, alg::CTMRG) + return enlarge_northeast_corner((r, c), envs, state) +end +function enlarge_corner(::Val{3}, (r, c), envs, state, alg::CTMRG) + return enlarge_southeast_corner((r, c), envs, state) +end +function enlarge_corner(::Val{4}, (r, c), envs, state, alg::CTMRG) + return enlarge_southwest_corner((r, c), envs, state) +end + """ ctmrg_expand(state, envs, alg::CTMRG{M}) @@ -183,41 +204,13 @@ There are two modes of expansion: `M = :sequential` and `M = :simultaneous`. The first mode expands the environment in one direction at a time, for convenience towards the left. The second mode expands the environment in all four directions simultaneously. """ -function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SequentialCTMRG) where {C,T} - Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) - Q_sw = Zygote.Buffer(envs.corners, Qtype, axes(state)...) - Q_nw = Zygote.Buffer(envs.corners, Qtype, axes(state)...) - - directions = collect(Iterators.product(axes(state)...)) - # @fwdthreads for (r, c) in directions - for (r, c) in directions - Q_sw[r, c] = enlarge_southwest_corner((r, c), envs, state) - Q_nw[r, c] = enlarge_northwest_corner((r, c), envs, state) - end - - return stack((copy(Q_sw), copy(Q_nw)); dims=1) +function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG) + drc_combinations = collect(Iterators.product([4, 1], axes(state)...)) + return map(idx -> enlarge_corner(idx, envs, state, alg), drc_combinations) end -function ctmrg_expand(state, envs::CTMRGEnv{C,T}, ::SimultaneousCTMRG) where {C,T} - Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) - Q = Zygote.Buffer(Array{Qtype,3}(undef, size(envs.corners))) - drc_combinations = collect(Iterators.product(axes(envs.corners)...)) - @fwdthreads for (dir, r, c) in drc_combinations - # TODO: generalize the corner computation to make it compatible with EnlargedCorners - Q[dir, r, c] = if dir == NORTHWEST - enlarge_northwest_corner((r, c), envs, state) - elseif dir == NORTHEAST - enlarge_northeast_corner((r, c), envs, state) - elseif dir == SOUTHEAST - enlarge_southeast_corner((r, c), envs, state) - elseif dir == SOUTHWEST - enlarge_southwest_corner((r, c), envs, state) - end - end - - return copy(Q) -end - -function enlarge_corner((dir, r, c), envs, state, alg::CTMRG) # TODO: find a way to dispatch on CTMRG struct to switch on function handles +function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) + drc_combinations = collect(Iterators.product(1:4, axes(state)...)) + return map(idx -> enlarge_corner(idx, envs, state, alg), drc_combinations) end # ======================================================================================== # diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index 783ef429..f0dc2689 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -23,6 +23,48 @@ Contract enlarged corner where `Val(1)` dispatches the north-west, `Val(2)` the (Q::EnlargedCorner)(::Val{3}) = enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) (Q::EnlargedCorner)(::Val{4}) = enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +""" + enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::SparseCTMRG) + +Enlarge corner but return as a `EnlargedCorner` struct used in sparse CTMRG. +""" +function enlarge_corner(::Val{1}, (r, c), envs, state, alg::SparseCTMRG) + return EnlargedCorner( + envs.corners[NORTHWEST, _prev(r, end), _prev(c, end)], + envs.edges[WEST, r, _prev(c, end)], + envs.edges[NORTH, _prev(r, end), c], + state[r, c], + state[r, c], + ) +end +function enlarge_corner(::Val{2}, (r, c), envs, state, alg::SparseCTMRG) + return EnlargedCorner( + envs.corners[NORTHEAST, _prev(r, end), _next(c, end)], + envs.edges[NORTH, _prev(r, end), c], + envs.edges[EAST, r, _next(c, end)], + state[r, c], + state[r, c], + ) +end +function enlarge_corner(::Val{3}, (r, c), envs, state, alg::SparseCTMRG) + return EnlargedCorner( + envs.corners[SOUTHEAST, _next(r, end), _next(c, end)], + envs.edges[EAST, r, _next(c, end)], + envs.edges[SOUTH, _next(r, end), c], + state[r, c], + state[r, c], + ) +end +function enlarge_corner(::Val{4}, (r, c), envs, state, alg::SparseCTMRG) + return EnlargedCorner( + envs.corners[SOUTHWEST, _next(r, end), _prev(c, end)], + envs.edges[SOUTH, _next(r, end), c], + envs.edges[WEST, r, _prev(c, end)], + state[r, c], + state[r, c], + ) +end + # Compute left & right projectors from enlarged corner struct function build_projectors( U::AbstractTensorMap{E,3,1}, @@ -39,6 +81,10 @@ function build_projectors( return P_left, P_right end +function renormalize_corner(ec::EnlargedCorner, P_left, P_right) + return renormalize_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra) +end + """ struct HalfInfiniteEnv{C,E,A,A′} From 7bf94afaf5a4ab5f5d7458e4e90c27d111662f83 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 4 Oct 2024 11:41:38 +0200 Subject: [PATCH 09/21] Add sparse_env type parameter to CTMRG --- src/PEPSKit.jl | 1 + src/algorithms/ctmrg/ctmrg.jl | 70 ++++++++++++--------- src/algorithms/ctmrg/sparse_environments.jl | 1 + 3 files changed, 43 insertions(+), 29 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 72bb6da6..1f87ace5 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -72,6 +72,7 @@ module Defaults const fpgrad_maxiter = 30 const fpgrad_tol = 1e-6 const ctmrgscheme = :simultaneous + const sparse_env = false const reuse_env = true const trscheme = FixedSpaceTruncation() const iterscheme = :fixed diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index af84e95f..1844d836 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -46,7 +46,7 @@ end CTMRG(; tol=Defaults.ctmrg_tol, maxiter=Defaults.ctmrg_maxiter, miniter=Defaults.ctmrg_miniter, verbosity=0, svd_alg=SVDAdjoint(), trscheme=FixedSpaceTruncation(), - ctmrgscheme=Defaults.ctmrgscheme) + ctmrgscheme=Defaults.ctmrgscheme, sparse_env=Defaults.sparse_env) Algorithm struct that represents the CTMRG algorithm for contracting infinite PEPS. Each CTMRG run is converged up to `tol` where the singular value convergence of the @@ -63,13 +63,17 @@ CTMRG is implemented. It can either be `:sequential`, where the projectors are s computed on the western side, and then applied and rotated. Or with `:simultaneous` all projectors are computed and applied simultaneously on all sides, where in particular the corners get contracted with two projectors at the same time. + +The contraction and SVD of the CTMRG environments can be performed densely, or sparsely +if `sparse_env` is `true`. If sparsity is enabled, then the SVD will use only function handles +and the full enlarged corners and environments are never explicitly constructed. """ -struct CTMRG{S} +struct CTMRG{M,S} tol::Float64 maxiter::Int miniter::Int verbosity::Int - projector_alg::ProjectorAlg + projector_alg::P end function CTMRG(; tol=Defaults.ctmrg_tol, @@ -79,13 +83,15 @@ function CTMRG(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, ctmrgscheme::Symbol=Defaults.ctmrgscheme, + sparse_env::Bool=Defaults.sparse_env, ) - return CTMRG{ctmrgscheme}( + return CTMRG{ctmrgscheme,sparse_env}( tol, maxiter, miniter, verbosity, ProjectorAlg(; svd_alg, trscheme, verbosity) ) end -ctmrgscheme(::CTMRG{S}) where {S} = S +ctmrgscheme(::CTMRG{M,S}) where {M,S} = M +sparse_env(::CTMRG{M,S}) where {M,S} = S # aliases for the different CTMRG schemes const SequentialCTMRG = CTMRG{:sequential} @@ -175,27 +181,6 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N) # Expansion step # ======================================================================================== # -# TODO: find a way to dispatch on CTMRG struct to switch on function handles -""" - enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::CTMRG) - -Enlarge a CTMRG corner into the one of four directions `Val(1)` (NW), `Val(2)` (NE), -`Val(3)` (SE) or `Val(4)` (SW). This serves as a wrapper that can dispatch on the -directional index and different `CTMRG` algorithms. -""" -function enlarge_corner(::Val{1}, (r, c), envs, state, alg::CTMRG) - return enlarge_northwest_corner((r, c), envs, state) -end -function enlarge_corner(::Val{2}, (r, c), envs, state, alg::CTMRG) - return enlarge_northeast_corner((r, c), envs, state) -end -function enlarge_corner(::Val{3}, (r, c), envs, state, alg::CTMRG) - return enlarge_southeast_corner((r, c), envs, state) -end -function enlarge_corner(::Val{4}, (r, c), envs, state, alg::CTMRG) - return enlarge_southwest_corner((r, c), envs, state) -end - """ ctmrg_expand(state, envs, alg::CTMRG{M}) @@ -206,11 +191,15 @@ the left. The second mode expands the environment in all four directions simulta """ function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG) drc_combinations = collect(Iterators.product([4, 1], axes(state)...)) - return map(idx -> enlarge_corner(idx, envs, state, alg), drc_combinations) + return map(drc_combinations) do (dir, r, c) + enlarge_corner(Val(dir), (r, c), envs, state, alg) + end end function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) drc_combinations = collect(Iterators.product(1:4, axes(state)...)) - return map(idx -> enlarge_corner(idx, envs, state, alg), drc_combinations) + return map(drc_combinations) do (dir, r, c) + enlarge_corner(Val(dir), (r, c), envs, state, alg) + end end # ======================================================================================== # @@ -218,7 +207,10 @@ end # ======================================================================================== # """ - ctmrg_projectors(Q, env, alg::CTMRG{M}) + ctmrg_projectors(enlarged_envs, env, alg::CTMRG{M}) + +Compute the CTMRG projectors based from enlarged environments. +In the `:simultaneous` mode, the environment SVD is run in parallel. """ function ctmrg_projectors( enlarged_envs, envs::CTMRGEnv{C,E}, alg::SequentialCTMRG @@ -388,6 +380,26 @@ end # Auxiliary routines # ======================================================================================== # +""" + enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::CTMRG) + +Enlarge a CTMRG corner into the one of four directions `Val(1)` (NW), `Val(2)` (NE), +`Val(3)` (SE) or `Val(4)` (SW). This serves as a wrapper that can dispatch on the +directional index and different `CTMRG` algorithms. +""" +function enlarge_corner(::Val{1}, (r, c), envs, state, alg::CTMRG) + return enlarge_northwest_corner((r, c), envs, state) +end +function enlarge_corner(::Val{2}, (r, c), envs, state, alg::CTMRG) + return enlarge_northeast_corner((r, c), envs, state) +end +function enlarge_corner(::Val{3}, (r, c), envs, state, alg::CTMRG) + return enlarge_southeast_corner((r, c), envs, state) +end +function enlarge_corner(::Val{4}, (r, c), envs, state, alg::CTMRG) + return enlarge_southwest_corner((r, c), envs, state) +end + # Build projectors from SVD and dense enlarged corners function build_projectors( U::AbstractTensorMap{E,3,1}, diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index f0dc2689..abf4358e 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -1,3 +1,4 @@ +const SparseCTMRG = CTMRG{M,true} where {M} # TODO: Is this really a good way to dispatch on the sparse CTMRG methods? """ struct EnlargedCorner{Ct,E,A,A′} From c931def01846e492e90ecb656a5f0c23029b7ff6 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 4 Oct 2024 17:42:36 +0200 Subject: [PATCH 10/21] Implement normal and adjoint linear map (svdsolve does not yet work) --- .../contractions/ctmrg_contractions.jl | 48 ++++++++-- src/algorithms/ctmrg/ctmrg.jl | 6 +- src/algorithms/ctmrg/sparse_environments.jl | 92 ++++++++++++++----- src/utility/svd.jl | 7 +- test/ctmrg/svd_wrapper.jl | 29 ++++-- 5 files changed, 137 insertions(+), 45 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 67317ed6..1883b07c 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -248,14 +248,14 @@ Alternatively, contract environment with a vector `x` acting on it, e.g. as need function halfinfinite_environment( quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3} ) where {S} - return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := + return @autoopt @tensor env[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := quadrant1[χ_in D_inabove D_inbelow; χ D1 D2] * quadrant2[χ D1 D2; χ_out D_outabove D_outbelow] end function halfinfinite_environment( C_1, C_2, E_1, E_2, E_3, E_4, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 ) where {P<:PEPSTensor} - return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := + return @autoopt @tensor env[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := E_1[χ_in D1 D2; χ1] * C_1[χ1; χ2] * E_2[χ2 D3 D4; χ3] * @@ -268,9 +268,19 @@ function halfinfinite_environment( E_4[χ5 D7 D8; χ_out] end function halfinfinite_environment( - C_1, C_2, E_1, E_2, E_3, E_4, x, ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2 -) where {P<:PEPSTensor} - return @autoopt @tensor half[χ_in D_inabove D_inbelow; χ_out D_outabove D_outbelow] := + C_1, + C_2, + E_1, + E_2, + E_3, + E_4, + x::AbstractTensor{S,3}, + ket_1::P, + ket_2::P, + bra_1::P=ket_1, + bra_2::P=ket_2, +) where {S,P<:PEPSTensor} + return @autoopt @tensor env_x[χ_in D_inabove D_inbelow] := E_1[χ_in D1 D2; χ1] * C_1[χ1; χ2] * E_2[χ2 D3 D4; χ3] * @@ -281,7 +291,33 @@ function halfinfinite_environment( E_3[χ3 D5 D6; χ4] * C_2[χ4; χ5] * E_4[χ5 D7 D8; χ6] * - x[χ6 D11 D12; χ_out D_outabove D_outbelow] + x[χ6 D11 D12] +end +function halfinfinite_environment( + x::AbstractTensor{S,3}, + C_1, + C_2, + E_1, + E_2, + E_3, + E_4, + ket_1::P, + ket_2::P, + bra_1::P=ket_1, + bra_2::P=ket_2, +) where {S,P<:PEPSTensor} + return @autoopt @tensor env_x[χ_in D_inabove D_inbelow] := + x[χ1 D1 D2] * + conj(E_1[χ1 D3 D4; χ2]) * + conj(C_1[χ2; χ3]) * + conj(E_2[χ3 D5 D6; χ4]) * + conj(ket_1[d1; D5 D11 D1 D3]) * + bra_1[d1; D6 D12 D2 D4] * + conj(ket_2[d2; D7 D9 D_inabove D11]) * + bra_2[d2; D8 D10 D_inbelow D12] * + conj(E_3[χ4 D7 D8; χ5]) * + conj(C_2[χ5; χ6]) * + conj(E_4[χ6 D9 D10; χ_in]) end # Renormalization contractions diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 1844d836..f95c2317 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -73,7 +73,7 @@ struct CTMRG{M,S} maxiter::Int miniter::Int verbosity::Int - projector_alg::P + projector_alg::ProjectorAlg end function CTMRG(; tol=Defaults.ctmrg_tol, @@ -189,13 +189,13 @@ There are two modes of expansion: `M = :sequential` and `M = :simultaneous`. The first mode expands the environment in one direction at a time, for convenience towards the left. The second mode expands the environment in all four directions simultaneously. """ -function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG) +function ctmrg_expand(state, envs::CTMRGEnv, alg::SequentialCTMRG) drc_combinations = collect(Iterators.product([4, 1], axes(state)...)) return map(drc_combinations) do (dir, r, c) enlarge_corner(Val(dir), (r, c), envs, state, alg) end end -function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) +function ctmrg_expand(state, envs::CTMRGEnv, alg::SimultaneousCTMRG) drc_combinations = collect(Iterators.product(1:4, axes(state)...)) return map(drc_combinations) do (dir, r, c) enlarge_corner(Val(dir), (r, c), envs, state, alg) diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index abf4358e..b102b70a 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -1,5 +1,9 @@ const SparseCTMRG = CTMRG{M,true} where {M} # TODO: Is this really a good way to dispatch on the sparse CTMRG methods? +# -------------------------------------------------------- +# Sparse enlarged corner as building block for environment +# -------------------------------------------------------- + """ struct EnlargedCorner{Ct,E,A,A′} @@ -24,6 +28,10 @@ Contract enlarged corner where `Val(1)` dispatches the north-west, `Val(2)` the (Q::EnlargedCorner)(::Val{3}) = enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) (Q::EnlargedCorner)(::Val{4}) = enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) +# ------------------------- +# CTMRG compatibility layer +# ------------------------- + """ enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::SparseCTMRG) @@ -66,7 +74,6 @@ function enlarge_corner(::Val{4}, (r, c), envs, state, alg::SparseCTMRG) ) end -# Compute left & right projectors from enlarged corner struct function build_projectors( U::AbstractTensorMap{E,3,1}, S::AbstractTensorMap{E,1,1}, @@ -86,6 +93,10 @@ function renormalize_corner(ec::EnlargedCorner, P_left, P_right) return renormalize_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra) end +# ------------------ +# Sparse environment +# ------------------ + """ struct HalfInfiniteEnv{C,E,A,A′} @@ -120,11 +131,13 @@ end """ (env::HalfInfiniteEnv)() - (env::HalfInfiniteEnv)(x) + (env::HalfInfiniteEnv)(x, ::Val{false}) + (env::HalfInfiniteEnv)(x, ::Val{true}) -Contract half-infinite environment without or with a vector `x`. +Contract half-infinite environment. If a vector `x` is provided, the environment acts as a +linear map or adjoint linear map on `x` if `Val(true)` or `Val(false)` is passed, respectively. """ -function (env::HalfInfiniteEnv)() +function (env::HalfInfiniteEnv)() # Dense operator return halfinfinite_environment( env.C_1, env.C_2, @@ -138,7 +151,7 @@ function (env::HalfInfiniteEnv)() env.bra_2, ) end -function (env::HalfInfiniteEnv)(x) +function (env::HalfInfiniteEnv)(x, ::Val{false}) # Linear map: env() * x return halfinfinite_environment( env.C_1, env.C_2, @@ -153,22 +166,67 @@ function (env::HalfInfiniteEnv)(x) env.bra_2, ) end +function (env::HalfInfiniteEnv)(x, ::Val{true}) # Adjoint linear map: env()' * x + return halfinfinite_environment( + x, + env.C_1, + env.C_2, + env.E_1, + env.E_2, + env.E_3, + env.E_4, + env.ket_1, + env.ket_2, + env.bra_1, + env.bra_2, + ) +end + +# Wrapper around halfinfinite_environment contraction using EnlargedCorners (used in ctmrg_projectors) +function halfinfinite_environment(ec_1::EnlargedCorner, ec_2::EnlargedCorner) + return HalfInfiniteEnv( + ec_1.C, + ec_2.C, + ec_1.E_1, + ec_1.E_2, + ec_2.E_1, + ec_2.E_2, + ec_1.ket, + ec_2.ket, + ec_1.bra, + ec_2.bra, + ) +end + +# ------------------------------------------------------------- +# TensorKit methods to make environment compatible with IterSVD +# ------------------------------------------------------------- -# TensorKit methods to make struct compatible with sparse SVD TensorKit.InnerProductStyle(::HalfInfiniteEnv) = EuclideanProduct() TensorKit.sectortype(::HalfInfiniteEnv) = Trivial TensorKit.storagetype(env::HalfInfiniteEnv) = storagetype(env.ket_1) TensorKit.spacetype(env::HalfInfiniteEnv) = spacetype(env.ket_1) +function TensorKit.domain(env::HalfInfiniteEnv) + return domain(env.E_4) * domain(env.ket_2)[3] * domain(env.bra_2)[3]' +end +function TensorKit.codomain(env::HalfInfiniteEnv) + return codomain(env.E_1)[1] * domain(env.ket_1)[3]' * domain(env.bra_1)[3] +end +function TensorKit.space(env::HalfInfiniteEnv) + return codomain(env) ← domain(env) +end function TensorKit.blocks(env::HalfInfiniteEnv) return TensorKit.SingletonDict(Trivial() => env) end function TensorKit.blocksectors(::HalfInfiniteEnv) return TensorKit.OneOrNoneIterator{Trivial}(true, Trivial()) end - +function TensorKit.tsvd!(f::HalfInfiniteEnv; trunc=NoTruncation(), p::Real=2, alg=IterSVD()) + return _tsvd!(f, alg, trunc, p) +end function TensorKit.MatrixAlgebra.svd!(env::HalfInfiniteEnv, args...) - return TensorKit.MatrixAlgebra.svd!(env(), args...) + return TensorKit.MatrixAlgebra.svd!(env(), args...) # Construct environment densely as fallback end Base.eltype(env::HalfInfiniteEnv) = eltype(env.ket_1) @@ -182,22 +240,8 @@ function Base.size(env::HalfInfiniteEnv) # Treat environment as matrix return (χ_in * D_inabove * D_inbelow, χ_out * D_outabove * D_outbelow) end Base.size(env::HalfInfiniteEnv, i::Int) = size(env)[i] - -# TODO: implement VectorInterface VectorInterface.scalartype(env::HalfInfiniteEnv) = scalartype(env.ket_1) -# Wrapper around halfinfinite_environment contraction using EnlargedCorners (used in ctmrg_projectors) -function halfinfinite_environment(ec_1::EnlargedCorner, ec_2::EnlargedCorner) - return HalfInfiniteEnv( - ec_1.C, - ec_2.C, - ec_1.E_1, - ec_1.E_2, - ec_2.E_1, - ec_2.E_2, - ec_1.ket, - ec_2.ket, - ec_1.bra, - ec_2.bra, - ) +function random_start_vector(env::HalfInfiniteEnv) + return Tensor(randn, domain(env)) end diff --git a/src/utility/svd.jl b/src/utility/svd.jl index 54765e20..96c1af40 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -36,9 +36,6 @@ PEPSKit.tsvd(t, alg; kwargs...) = PEPSKit.tsvd!(copy(t), alg; kwargs...) function PEPSKit.tsvd!(t, alg::SVDAdjoint; trunc::TruncationScheme=notrunc(), p::Real=2) return TensorKit.tsvd!(t; alg=alg.fwd_alg, trunc, p) end -function TensorKit.tsvd!(f::HalfInfiniteEnv; trunc=NoTruncation(), p::Real=2, alg=IterSVD()) - return _tsvd!(f, alg, trunc, p) -end """ struct FixedSVD @@ -74,8 +71,8 @@ the iterative SVD didn't converge, the algorithm falls back to a dense SVD. end # TODO: find better initial guess that leads to element-wise convergence and is compatible with function handles -function random_start_vector(f) - return randn(eltype(f), size(f, 1)) # Leads to erroneous gauge fixing of U, S, V and thus failing element-wise conv. +function random_start_vector(t::Matrix) + return randn(eltype(t), size(t, 1)) # Leads to erroneous gauge fixing of U, S, V and thus failing element-wise conv. # u, = TensorKit.MatrixAlgebra.svd!(deepcopy(t), TensorKit.SVD()) # return sum(u[:, i] for i in 1:howmany) # Element-wise convergence works fine # return dropdims(sum(t[:, 1:3]; dims=2); dims=2) # Summing too many columns also makes gauge fixing fail diff --git a/test/ctmrg/svd_wrapper.jl b/test/ctmrg/svd_wrapper.jl index 927a1ed3..68a53aa8 100644 --- a/test/ctmrg/svd_wrapper.jl +++ b/test/ctmrg/svd_wrapper.jl @@ -100,7 +100,7 @@ ctm_alg = CTMRG(; tol=1e-10, verbosity=2, svd_alg=SVDAdjoint()) Random.seed!(91283219347) H = heisenberg_XYZ(InfiniteSquare()) psi = InfinitePEPS(2, χbond) -env = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg) +env = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg); hienv = HalfInfiniteEnv( env.corners[1], env.corners[2], @@ -114,12 +114,27 @@ hienv = HalfInfiniteEnv( psi[1], ) hienv_dense = hienv() -env_R = TensorMap(randn, space(hienv_dense)) -PEPSKit.tsvd!(hienv, iter_alg) +env_R = TensorMap(randn, space(hienv)) + +PEPSKit.tsvd!(hienv, iter_alg) # TODO: make the space mismatches work @testset "IterSVD with HalfInfiniteEnv function handle" begin - l_fullsvd, g_fullsvd = withgradient(A -> lossfun(A, full_alg, env_R), hienv_dense) - l_itersvd, g_itersvd = withgradient(A -> lossfun(A, iter_alg, env_R), hienv) - @test l_itersvd ≈ l_fullsvd - @test g_fullsvd[1] ≈ g_itersvd[1] rtol = rtol + # Equivalence of dense and sparse contractions + x₀ = PEPSKit.random_start_vector(hienv) + x′ = hienv(x₀, Val(false)) + x″ = hienv(x′, Val(true)) + x‴ = hienv(x″, Val(false)) + + a = hienv_dense * x₀ + b = hienv_dense' * a + c = hienv_dense * b + @test a ≈ x′ + @test b ≈ x″ + @test c ≈ x‴ + + # TODO: code up pullback + # l_fullsvd, g_fullsvd = withgradient(A -> lossfun(A, full_alg, env_R), hienv_dense) + # l_itersvd, g_itersvd = withgradient(A -> lossfun(A, iter_alg, env_R), hienv) + # @test l_itersvd ≈ l_fullsvd + # @test g_fullsvd[1] ≈ g_itersvd[1] rtol = rtol end From 9e033dc7c3c78c85b04b4af12caf68ff7a5a7cbe Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 4 Oct 2024 18:09:40 +0200 Subject: [PATCH 11/21] Add tsvd! rrule compatibility layer for HalfInfiniteEnv --- src/algorithms/ctmrg/sparse_environments.jl | 24 +++++++++++-- src/utility/svd.jl | 37 +++++++-------------- 2 files changed, 33 insertions(+), 28 deletions(-) diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index b102b70a..e79c5d84 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -198,9 +198,9 @@ function halfinfinite_environment(ec_1::EnlargedCorner, ec_2::EnlargedCorner) ) end -# ------------------------------------------------------------- -# TensorKit methods to make environment compatible with IterSVD -# ------------------------------------------------------------- +# ------------------------------------------------------------------------ +# Methods to make environment compatible with IterSVD and its reverse-rule +# ------------------------------------------------------------------------ TensorKit.InnerProductStyle(::HalfInfiniteEnv) = EuclideanProduct() TensorKit.sectortype(::HalfInfiniteEnv) = Trivial @@ -222,6 +222,9 @@ end function TensorKit.blocksectors(::HalfInfiniteEnv) return TensorKit.OneOrNoneIterator{Trivial}(true, Trivial()) end +function TensorKit.block(env::HalfInfiniteEnv, c::Sector) + return env +end function TensorKit.tsvd!(f::HalfInfiniteEnv; trunc=NoTruncation(), p::Real=2, alg=IterSVD()) return _tsvd!(f, alg, trunc, p) end @@ -245,3 +248,18 @@ VectorInterface.scalartype(env::HalfInfiniteEnv) = scalartype(env.ket_1) function random_start_vector(env::HalfInfiniteEnv) return Tensor(randn, domain(env)) end + +function Base.similar(env::HalfInfiniteEnv) + return HalfInfiniteEnv( + (similar(getfield(env, field)) for field in fieldnames(HalfInfiniteEnv))... + ) +end + +function Base.copyto!(dest::HalfInfiniteEnv, src::HalfInfiniteEnv) + for field in fieldnames(HalfInfiniteEnv) + for (bd, bs) in zip(blocks(getfield(dest, field)), blocks(getfield(src, field))) + copyto!(bd[2], bs[2]) + end + end + return dest +end diff --git a/src/utility/svd.jl b/src/utility/svd.jl index 96c1af40..d439b4f6 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -137,27 +137,27 @@ function TensorKit._compute_svddata!( return Udata, Sdata, Vdata, dims end -# Rrule with custom pullback to make KrylovKit rrule compatible with TensorMap symmetry blocks +# Rrule with custom pullback to make KrylovKit rrule compatible with TensorMaps & function handles function ChainRulesCore.rrule( ::typeof(PEPSKit.tsvd!), - t::AbstractTensorMap, + f, alg::SVDAdjoint{F,R,B}; trunc::TruncationScheme=notrunc(), p::Real=2, ) where {F<:Union{IterSVD,FixedSVD},R<:Union{GMRES,BiCGStab,Arnoldi},B} - U, S, V, ϵ = PEPSKit.tsvd(t, alg; trunc, p) + U, S, V, ϵ = PEPSKit.tsvd(f, alg; trunc, p) function tsvd!_itersvd_pullback((ΔU, ΔS, ΔV, Δϵ)) - Δt = similar(t) - for (c, b) in blocks(Δt) + Δf = similar(f) + for (c, b) in blocks(Δf) Uc, Sc, Vc = block(U, c), block(S, c), block(V, c) ΔUc, ΔSc, ΔVc = block(ΔU, c), block(ΔS, c), block(ΔV, c) Sdc = view(Sc, diagind(Sc)) ΔSdc = ΔSc isa AbstractZero ? ΔSc : view(ΔSc, diagind(ΔSc)) n_vals = length(Sdc) - lvecs = Vector{Vector{scalartype(t)}}(eachcol(Uc)) - rvecs = Vector{Vector{scalartype(t)}}(eachcol(Vc')) + lvecs = Vector{Vector{scalartype(f)}}(eachcol(Uc)) + rvecs = Vector{Vector{scalartype(f)}}(eachcol(Vc')) # Dummy objects only used for warnings minimal_info = KrylovKit.ConvergenceInfo(n_vals, nothing, nothing, -1, -1) # Only num. converged is used @@ -167,8 +167,8 @@ function ChainRulesCore.rrule( Δlvecs = fill(ZeroTangent(), n_vals) Δrvecs = fill(ZeroTangent(), n_vals) else - Δlvecs = Vector{Vector{scalartype(t)}}(eachcol(ΔUc)) - Δrvecs = Vector{Vector{scalartype(t)}}(eachcol(ΔVc')) + Δlvecs = Vector{Vector{scalartype(f)}}(eachcol(ΔUc)) + Δrvecs = Vector{Vector{scalartype(f)}}(eachcol(ΔVc')) end xs, ys = CRCExt.compute_svdsolve_pullback_data( @@ -179,17 +179,17 @@ function ChainRulesCore.rrule( lvecs, rvecs, minimal_info, - block(t, c), + block(f, c), :LR, minimal_alg, alg.rrule_alg, ) copyto!( b, - CRCExt.construct∂f_svd(HasReverseMode(), block(t, c), lvecs, rvecs, xs, ys), + CRCExt.construct∂f_svd(HasReverseMode(), block(f, c), lvecs, rvecs, xs, ys), ) end - return NoTangent(), Δt, NoTangent() + return NoTangent(), Δf, NoTangent() end function tsvd!_itersvd_pullback(::Tuple{ZeroTangent,ZeroTangent,ZeroTangent}) return NoTangent(), ZeroTangent(), NoTangent() @@ -198,19 +198,6 @@ function ChainRulesCore.rrule( return (U, S, V, ϵ), tsvd!_itersvd_pullback end -# Separate rule for SVD with function handle that uses KrylovKit.make_svdsolve_pullback -# but in turn cannot handle symmetric blocks -function ChainRulesCore.rrule( - ::typeof(PEPSKit.tsvd!), - f, - alg::SVDAdjoint{F,R,B}; - trunc::TruncationScheme=notrunc(), - p::Real=2, -) where {F<:Union{IterSVD,FixedSVD},R<:Union{GMRES,BiCGStab,Arnoldi},B} - return U, S, V, ϵ = PEPSKit.tsvd(t, alg; trunc, p) - # TODO: implement function handle adjoint wrapper with KrylovKit.make_svdsolve_pullback -end - """ struct NonTruncAdjoint From c3a88227fbe41f77575cb2697228324137dd1290 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 7 Oct 2024 11:42:36 +0200 Subject: [PATCH 12/21] Add sparse as kwarg of ProjectorAlg, use EnlargedCorners per default, fix some docstrings --- src/PEPSKit.jl | 2 +- .../contractions/ctmrg_contractions.jl | 14 +- src/algorithms/ctmrg/ctmrg.jl | 172 +++++++++++------- src/algorithms/ctmrg/sparse_environments.jl | 83 ++------- 4 files changed, 134 insertions(+), 137 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 1f87ace5..7dee4ebb 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -72,7 +72,7 @@ module Defaults const fpgrad_maxiter = 30 const fpgrad_tol = 1e-6 const ctmrgscheme = :simultaneous - const sparse_env = false + const sparse = false const reuse_env = true const trscheme = FixedSpaceTruncation() const iterscheme = :fixed diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 1883b07c..10a569e7 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -211,10 +211,12 @@ end """ halfinfinite_environment(quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3}) - function halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, - ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor} - function halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, x, - ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor} + halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, + ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor} + halfinfinite_environment(C_1, C_2, E_1, E_2, E_3, E_4, x, + ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor} + halfinfinite_environment(x, C_1, C_2, E_1, E_2, E_3, E_4, + ket_1::P, ket_2::P, bra_1::P=ket_1, bra_2::P=ket_2) where {P<:PEPSTensor} Contract two quadrants (enlarged corners) to form a half-infinite environment. @@ -234,7 +236,7 @@ The environment can also be contracted directly from all its constituent tensors | || || | ``` -Alternatively, contract environment with a vector `x` acting on it, e.g. as needed for iterative solvers. +Alternatively, contract environment with a vector `x` acting on it ``` C_1 -- E_2 -- E_3 -- C_2 @@ -244,6 +246,8 @@ Alternatively, contract environment with a vector `x` acting on it, e.g. as need [~~~~~~x~~~~~~] || | ``` + +or contract the adjoint environment with `x`, e.g. as needed for iterative solvers. """ function halfinfinite_environment( quadrant1::AbstractTensorMap{S,3,3}, quadrant2::AbstractTensorMap{S,3,3} diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index f95c2317..0cf3166c 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -8,18 +8,20 @@ have different spaces, this truncation style is different from `TruncationSpace` struct FixedSpaceTruncation <: TensorKit.TruncationScheme end """ - struct ProjectorAlg{S}(; svd_alg=TensorKit.SVD(), trscheme=TensorKit.notrunc(), - fixedspace=false, verbosity=0) - -Algorithm struct collecting all projector related parameters. The truncation scheme has to be -a `TensorKit.TruncationScheme`, and some SVD algorithms might have further restrictions on what -kind of truncation scheme can be used. If `fixedspace` is true, the truncation scheme is set to -`truncspace(V)` where `V` is the environment bond space, adjusted to the corresponding -environment direction/unit cell entry. + struct ProjectorAlg{S}(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, + sparse=Defaults.sparse, verbosity=0) + +Algorithm struct collecting all projector related parameters. + +The `svd_alg` sets the SVD algorithm for decomposing the CTM environment. The truncation scheme +has to be a `TensorKit.TruncationScheme`, and some SVD algorithms might have further restrictions +on what kind of truncation scheme can be used. If `sparse` is `true`, then the enlarged corners and +CTM environment are never computed densely and function handles are using for the SVD. """ @kwdef struct ProjectorAlg{S<:SVDAdjoint,T} svd_alg::S = Defaults.svd_alg trscheme::T = Defaults.trscheme + sparse::Bool = Defaults.sparse verbosity::Int = 0 end # TODO: add option for different projector styles (half-infinite, full-infinite, etc.) @@ -46,7 +48,7 @@ end CTMRG(; tol=Defaults.ctmrg_tol, maxiter=Defaults.ctmrg_maxiter, miniter=Defaults.ctmrg_miniter, verbosity=0, svd_alg=SVDAdjoint(), trscheme=FixedSpaceTruncation(), - ctmrgscheme=Defaults.ctmrgscheme, sparse_env=Defaults.sparse_env) + ctmrgscheme=Defaults.ctmrgscheme, sparse=Defaults.sparse) Algorithm struct that represents the CTMRG algorithm for contracting infinite PEPS. Each CTMRG run is converged up to `tol` where the singular value convergence of the @@ -65,10 +67,10 @@ are computed and applied simultaneously on all sides, where in particular the co contracted with two projectors at the same time. The contraction and SVD of the CTMRG environments can be performed densely, or sparsely -if `sparse_env` is `true`. If sparsity is enabled, then the SVD will use only function handles +if `sparse` is `true`. If sparsity is enabled, then the SVD will use only function handles and the full enlarged corners and environments are never explicitly constructed. """ -struct CTMRG{M,S} +struct CTMRG{S} tol::Float64 maxiter::Int miniter::Int @@ -83,15 +85,18 @@ function CTMRG(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, ctmrgscheme::Symbol=Defaults.ctmrgscheme, - sparse_env::Bool=Defaults.sparse_env, + sparse::Bool=Defaults.sparse, ) - return CTMRG{ctmrgscheme,sparse_env}( - tol, maxiter, miniter, verbosity, ProjectorAlg(; svd_alg, trscheme, verbosity) + return CTMRG{ctmrgscheme}( + tol, + maxiter, + miniter, + verbosity, + ProjectorAlg(; svd_alg, trscheme, sparse, verbosity), ) end -ctmrgscheme(::CTMRG{M,S}) where {M,S} = M -sparse_env(::CTMRG{M,S}) where {M,S} = S +ctmrgscheme(::CTMRG{S}) where {S} = S # aliases for the different CTMRG schemes const SequentialCTMRG = CTMRG{:sequential} @@ -189,16 +194,53 @@ There are two modes of expansion: `M = :sequential` and `M = :simultaneous`. The first mode expands the environment in one direction at a time, for convenience towards the left. The second mode expands the environment in all four directions simultaneously. """ -function ctmrg_expand(state, envs::CTMRGEnv, alg::SequentialCTMRG) +function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG) drc_combinations = collect(Iterators.product([4, 1], axes(state)...)) - return map(drc_combinations) do (dir, r, c) - enlarge_corner(Val(dir), (r, c), envs, state, alg) - end + return map(idx -> enlarge_corner(idx, envs, state), drc_combinations) end -function ctmrg_expand(state, envs::CTMRGEnv, alg::SimultaneousCTMRG) +function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) drc_combinations = collect(Iterators.product(1:4, axes(state)...)) - return map(drc_combinations) do (dir, r, c) - enlarge_corner(Val(dir), (r, c), envs, state, alg) + return map(idx -> enlarge_corner(idx, envs, state), drc_combinations) +end + +""" + enlarge_corner((dir, r, c), envs, state) + +Enlarge corner by constructing the corresponding `EnlargedCorner` struct. +""" +function enlarge_corner((dir, r, c), envs, state) + if dir == NORTHWEST + return EnlargedCorner( + envs.corners[NORTHWEST, _prev(r, end), _prev(c, end)], + envs.edges[WEST, r, _prev(c, end)], + envs.edges[NORTH, _prev(r, end), c], + state[r, c], + state[r, c], + ) + elseif dir == NORTHEAST + return EnlargedCorner( + envs.corners[NORTHEAST, _prev(r, end), _next(c, end)], + envs.edges[NORTH, _prev(r, end), c], + envs.edges[EAST, r, _next(c, end)], + state[r, c], + state[r, c], + ) + elseif dir == SOUTHEAST + return EnlargedCorner( + envs.corners[SOUTHEAST, _next(r, end), _next(c, end)], + envs.edges[EAST, r, _next(c, end)], + envs.edges[SOUTH, _next(r, end), c], + state[r, c], + state[r, c], + ) + elseif dir == SOUTHWEST + return EnlargedCorner( + envs.corners[SOUTHWEST, _next(r, end), _prev(c, end)], + envs.edges[SOUTH, _next(r, end), c], + envs.edges[WEST, r, _prev(c, end)], + state[r, c], + state[r, c], + ) end end @@ -222,6 +264,11 @@ function ctmrg_projectors( P_top = Zygote.Buffer(envs.edges, Prtype, axes(envs.corners, 2), axes(envs.corners, 3)) ϵ = zero(real(scalartype(envs))) + # If CTMRG is not sparse, construct all enlarged corners densely + sparse || enlarged_envs = map(Iterators.product(axes(enlarged_envs)...)) do (dir, r, c) + return enlarged_envs[dir, r, c](dir) + end + directions = collect(Iterators.product(axes(envs.corners, 2), axes(envs.corners, 3))) # @fwdthreads for (r, c) in directions for (r, c) in directions @@ -260,6 +307,11 @@ function ctmrg_projectors( # Corner type but with real numbers S = Zygote.Buffer(U.data, tensormaptype(spacetype(C), 1, 1, real(scalartype(E)))) + # If CTMRG is not sparse, construct all enlarged corners densely + sparse || enlarged_envs = map(Iterators.product(axes(enlarged_envs)...)) do (dir, r, c) + return enlarged_envs[dir, r, c](dir) + end + ϵ = zero(real(scalartype(envs))) drc_combinations = collect(Iterators.product(axes(envs.corners)...)) @fwdthreads for (dir, r, c) in drc_combinations @@ -308,6 +360,42 @@ function ctmrg_projectors( return (copy(P_left), copy(P_right)), (; err=ϵ, U=copy(U), S=copy(S), V=copy(V)) end +""" + build_projectors(U::AbstractTensorMap{E,3,1}, S::AbstractTensorMap{E,1,1}, V::AbstractTensorMap{E,1,3}, + Q::AbstractTensorMap{E,3,3}, Q_next::AbstractTensorMap{E,3,3}) where {E<:ElementarySpace} + + build_projectors(U::AbstractTensorMap{E,3,1}, S::AbstractTensorMap{E,1,1}, V::AbstractTensorMap{E,1,3}, + Q::EnlargedCorner, Q_next::EnlargedCorner) where {E<:ElementarySpace} + +Construct left and right projectors where the higher-dimensional is facing left and right, respectively. +""" +function build_projectors( + U::AbstractTensorMap{E,3,1}, + S::AbstractTensorMap{E,1,1}, + V::AbstractTensorMap{E,1,3}, + Q::AbstractTensorMap{E,3,3}, + Q_next::AbstractTensorMap{E,3,3}, +) where {E<:ElementarySpace} + isqS = sdiag_inv_sqrt(S) + P_left = Q_next * V' * isqS + P_right = isqS * U' * Q + return P_left, P_right +end +function build_projectors( + U::AbstractTensorMap{E,3,1}, + S::AbstractTensorMap{E,1,1}, + V::AbstractTensorMap{E,1,3}, + Q::EnlargedCorner, + Q_next::EnlargedCorner, +) where {E<:ElementarySpace} + isqS = sdiag_inv_sqrt(S) + P_left = left_projector(Q.E_1, Q.C, Q.E_2, V, isqS, Q.ket, Q.bra) + P_right = right_projector( + Q_next.E_1, Q_next.C, Q_next.E_2, U, isqS, Q_next.ket, Q_next.bra + ) + return P_left, P_right +end + # ======================================================================================== # # Renormalization step # ======================================================================================== # @@ -375,41 +463,3 @@ function ctmrg_renormalize(enlarged_envs, projectors, state, envs, ::Simultaneou return CTMRGEnv(copy(corners), copy(edges)) end - -# ======================================================================================== # -# Auxiliary routines -# ======================================================================================== # - -""" - enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::CTMRG) - -Enlarge a CTMRG corner into the one of four directions `Val(1)` (NW), `Val(2)` (NE), -`Val(3)` (SE) or `Val(4)` (SW). This serves as a wrapper that can dispatch on the -directional index and different `CTMRG` algorithms. -""" -function enlarge_corner(::Val{1}, (r, c), envs, state, alg::CTMRG) - return enlarge_northwest_corner((r, c), envs, state) -end -function enlarge_corner(::Val{2}, (r, c), envs, state, alg::CTMRG) - return enlarge_northeast_corner((r, c), envs, state) -end -function enlarge_corner(::Val{3}, (r, c), envs, state, alg::CTMRG) - return enlarge_southeast_corner((r, c), envs, state) -end -function enlarge_corner(::Val{4}, (r, c), envs, state, alg::CTMRG) - return enlarge_southwest_corner((r, c), envs, state) -end - -# Build projectors from SVD and dense enlarged corners -function build_projectors( - U::AbstractTensorMap{E,3,1}, - S::AbstractTensorMap{E,1,1}, - V::AbstractTensorMap{E,1,3}, - Q::AbstractTensorMap{E,3,3}, - Q_next::AbstractTensorMap{E,3,3}, -) where {E<:ElementarySpace} - isqS = sdiag_inv_sqrt(S) - P_left = Q_next * V' * isqS - P_right = isqS * U' * Q - return P_left, P_right -end diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index e79c5d84..f32842a3 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -1,5 +1,3 @@ -const SparseCTMRG = CTMRG{M,true} where {M} # TODO: Is this really a good way to dispatch on the sparse CTMRG methods? - # -------------------------------------------------------- # Sparse enlarged corner as building block for environment # -------------------------------------------------------- @@ -18,77 +16,22 @@ struct EnlargedCorner{Ct,E,A,A′} end """ - (Q::EnlargedCorner)(::Val{<:Int}) + (Q::EnlargedCorner)(dir::Int) -Contract enlarged corner where `Val(1)` dispatches the north-west, `Val(2)` the north-east -`Val(3)` the south-east and `Val(4)` the south-west contraction. +Contract enlarged corner where `dir` selects the correct contraction direction, +i.e. the way the environment and PEPS tensors connect. """ -(Q::EnlargedCorner)(::Val{1}) = enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) -(Q::EnlargedCorner)(::Val{2}) = enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) -(Q::EnlargedCorner)(::Val{3}) = enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) -(Q::EnlargedCorner)(::Val{4}) = enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - -# ------------------------- -# CTMRG compatibility layer -# ------------------------- - -""" - enlarge_corner(::Val{<:Int}, (r, c), envs, state, alg::SparseCTMRG) - -Enlarge corner but return as a `EnlargedCorner` struct used in sparse CTMRG. -""" -function enlarge_corner(::Val{1}, (r, c), envs, state, alg::SparseCTMRG) - return EnlargedCorner( - envs.corners[NORTHWEST, _prev(r, end), _prev(c, end)], - envs.edges[WEST, r, _prev(c, end)], - envs.edges[NORTH, _prev(r, end), c], - state[r, c], - state[r, c], - ) -end -function enlarge_corner(::Val{2}, (r, c), envs, state, alg::SparseCTMRG) - return EnlargedCorner( - envs.corners[NORTHEAST, _prev(r, end), _next(c, end)], - envs.edges[NORTH, _prev(r, end), c], - envs.edges[EAST, r, _next(c, end)], - state[r, c], - state[r, c], - ) -end -function enlarge_corner(::Val{3}, (r, c), envs, state, alg::SparseCTMRG) - return EnlargedCorner( - envs.corners[SOUTHEAST, _next(r, end), _next(c, end)], - envs.edges[EAST, r, _next(c, end)], - envs.edges[SOUTH, _next(r, end), c], - state[r, c], - state[r, c], - ) -end -function enlarge_corner(::Val{4}, (r, c), envs, state, alg::SparseCTMRG) - return EnlargedCorner( - envs.corners[SOUTHWEST, _next(r, end), _prev(c, end)], - envs.edges[SOUTH, _next(r, end), c], - envs.edges[WEST, r, _prev(c, end)], - state[r, c], - state[r, c], - ) -end - -function build_projectors( - U::AbstractTensorMap{E,3,1}, - S::AbstractTensorMap{E,1,1}, - V::AbstractTensorMap{E,1,3}, - Q::EnlargedCorner, - Q_next::EnlargedCorner, -) where {E<:ElementarySpace} - isqS = sdiag_inv_sqrt(S) - P_left = left_projector(Q.E_1, Q.C, Q.E_2, V, isqS, Q.ket, Q.bra) - P_right = right_projector( - Q_next.E_1, Q_next.C, Q_next.E_2, U, isqS, Q_next.ket, Q_next.bra - ) - return P_left, P_right +function (Q::EnlargedCorner)(dir::Int) + if dir == NORTHWEST + return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == NORTHWEST + return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == NORTHWEST + return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == NORTHWEST + return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + end end - function renormalize_corner(ec::EnlargedCorner, P_left, P_right) return renormalize_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra) end From 812481408cad3a88e33acdbbc03edc234215ad99 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 7 Oct 2024 16:44:52 +0200 Subject: [PATCH 13/21] Fix renormalize methods to make CTMRG run --- src/PEPSKit.jl | 2 +- .../contractions/ctmrg_contractions.jl | 146 ++++++++++++++---- src/algorithms/ctmrg/ctmrg.jl | 31 ++-- src/algorithms/ctmrg/sparse_environments.jl | 43 +++++- test/heisenberg.jl | 2 +- 5 files changed, 171 insertions(+), 53 deletions(-) diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 7dee4ebb..3973c6ab 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -36,8 +36,8 @@ include("environments/transferpepo_environments.jl") include("algorithms/contractions/localoperator.jl") include("algorithms/contractions/ctmrg_contractions.jl") -include("algorithms/ctmrg/ctmrg.jl") include("algorithms/ctmrg/sparse_environments.jl") +include("algorithms/ctmrg/ctmrg.jl") include("algorithms/ctmrg/gaugefix.jl") include("algorithms/toolbox.jl") diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index 10a569e7..f4d32d0b 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -331,7 +331,6 @@ end """ renormalize_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) - renormalize_corner(E_1, C, E_2, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket) Apply projectors to each side of a quadrant. @@ -343,78 +342,154 @@ Apply projectors to each side of a quadrant. [P_right] | ``` - -Alternatively, provide the constituent tensors and perform the complete contraction. - -``` - C -- E_2 -- |~~~~~~| - | | |P_left| -- - E_1 == ket-bra == |~~~~~~| - | || - [~~P_right~~] - | -``` """ function renormalize_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) where {S} return @autoopt @tensor corner[χ_in; χ_out] := P_right[χ_in; χ1 D1 D2] * quadrant[χ1 D1 D2; χ2 D3 D4] * P_left[χ2 D3 D4; χ_out] end -function renormalize_corner( - E_1, C, E_2, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket -) - return @autoopt @tensor corner[χ_in; χ_out] := - P_right[χ_in; χ1 D1 D2] * - E_1[χ1 D3 D4; χ2] * - C[χ2; χ3] * - E_2[χ3 D5 D6; χ4] * - ket[d; D5 D7 D1 D3] * - conj(bra[d; D6 D8 D2 D4]) * - P_left[χ4 D7 D8; χ_out] -end """ renormalize_northwest_corner((row, col), enlarged_envs::CTMRGEnv, P_left, P_right) + renormalize_northwest_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) where {S} + renormalize_northwest_corner(E_west, C_northwest, E_north, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket) Apply `renormalize_corner` to the enlarged northwest corner. +Alternatively, provide the constituent tensors and perform the complete contraction. + +``` + C_northwest -- E_north -- |~~~~~~| + | || |P_left| -- + E_west= == ket-bra == |~~~~~~| + | || + [~~~~~P_right~~~~~] + | +``` """ function renormalize_northwest_corner((row, col), enlarged_envs, P_left, P_right) - return renormalize_corner( + return renormalize_northwest_corner( enlarged_envs[NORTHWEST, row, col], P_left[NORTH, row, col], P_right[WEST, _next(row, end), col], ) end +function renormalize_northwest_corner( + quadrant::AbstractTensorMap{S,3,3}, P_left, P_right +) where {S} + return renormalize_corner(quadrant, P_left, P_right) +end +function renormalize_northwest_corner( + E_west, C_northwest, E_north, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket +) + return @autoopt @tensor corner[χ_in; χ_out] := + P_right[χ_in; χ1 D1 D2] * + E_west[χ1 D3 D4; χ2] * + C_northwest[χ2; χ3] * + E_north[χ3 D5 D6; χ4] * + ket[d; D5 D7 D1 D3] * + conj(bra[d; D6 D8 D2 D4]) * + P_left[χ4 D7 D8; χ_out] +end """ renormalize_northeast_corner((row, col), enlarged_envs::CTMRGEnv, P_left, P_right) + renormalize_northwest_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) where {S} + renormalize_northeast_corner(E_north, C_northeast, E_east, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket) Apply `renormalize_corner` to the enlarged northeast corner. +Alternatively, provide the constituent tensors and perform the complete contraction. + +``` + |~~~~~~~| -- E_north -- C_northeast + -- |P_right| || | + |~~~~~~~| == ket-bra == E_east + || | + [~~~~~~P_left~~~~~~] + | +``` """ function renormalize_northeast_corner((row, col), enlarged_envs, P_left, P_right) - return renormalize_corner( + return renormalize_northeast_corner( enlarged_envs[NORTHEAST, row, col], P_left[EAST, row, col], P_right[NORTH, row, _prev(col, end)], ) end +function renormalize_northeast_corner( + quadrant::AbstractTensorMap{S,3,3}, P_left, P_right +) where {S} + return renormalize_corner(quadrant, P_left, P_right) +end +function renormalize_northeast_corner( + E_north, C_northeast, E_east, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket +) + return @autoopt @tensor corner[χ_in; χ_out] := + P_right[χ_in; χ1 D1 D2] * + E_north[χ1 D3 D4; χ2] * + C_northeast[χ2; χ3] * + E_east[χ3 D5 D6; χ4] * + ket[d; D3 D5 D7 D1] * + conj(bra[d; D4 D6 D8 D2]) * + P_left[χ4 D7 D8; χ_out] +end """ renormalize_southeast_corner((row, col), enlarged_envs::CTMRGEnv, P_left, P_right) + renormalize_southeast_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) where {S} + renormalize_southeast_corner(E_east, C_southeast, E_south, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket) Apply `renormalize_corner` to the enlarged southeast corner. +Alternatively, provide the constituent tensors and perform the complete contraction. + +``` + | + [~~~~~P_right~~~~~] + || | + |~~~~~~| == ket-bra == E_east + -- |P_left| || | + |~~~~~~| -- E_south -- C_southeast +``` """ function renormalize_southeast_corner((row, col), enlarged_envs, P_left, P_right) - return renormalize_corner( + return renormalize_southeast_corner( enlarged_envs[SOUTHEAST, row, col], P_left[SOUTH, row, col], P_right[EAST, _prev(row, end), col], ) end +function renormalize_southeast_corner( + quadrant::AbstractTensorMap{S,3,3}, P_left, P_right +) where {S} + return renormalize_corner(quadrant, P_left, P_right) +end +function renormalize_southeast_corner( + E_east, C_southeast, E_south, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket +) + return @autoopt @tensor corner[χ_in; χ_out] := + P_right[χ_in; χ1 D1 D2] * + E_east[χ1 D3 D4; χ2] * + C_southeast[χ2; χ3] * + E_south[χ3 D5 D6; χ4] * + ket[d; D1 D3 D5 D7] * + conj(bra[d; D2 D4 D6 D8]) * + P_left[χ4 D7 D8; χ_out] +end """ renormalize_southwest_corner((row, col), enlarged_envs::CTMRGEnv, P_left, P_right) + renormalize_southwest_corner(quadrant::AbstractTensorMap{S,3,3}, P_left, P_right) where {S} + renormalize_southwest_corner(E_south, C_southwest, E_west, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket) Apply `renormalize_corner` to the enlarged southwest corner. +Alternatively, provide the constituent tensors and perform the complete contraction. + +``` + | + [~~~~~P_right~~~~~] + || | + E_west == ket-bra == |~~~~~~| + | || |P_left| -- + C_southwest -- E_south -- |~~~~~~| +``` """ function renormalize_southwest_corner((row, col), enlarged_envs, P_left, P_right) return renormalize_corner( @@ -423,6 +498,23 @@ function renormalize_southwest_corner((row, col), enlarged_envs, P_left, P_right P_right[SOUTH, row, _next(col, end)], ) end +function renormalize_southwest_corner( + quadrant::AbstractTensorMap{S,3,3}, P_left, P_right +) where {S} + return renormalize_southwest_corner(quadrant, P_left, P_right) +end +function renormalize_southwest_corner( + E_south, C_southwest, E_west, P_left, P_right, ket::PEPSTensor, bra::PEPSTensor=ket +) + return @autoopt @tensor corner[χ_in; χ_out] := + P_right[χ_in; χ1 D1 D2] * + E_south[χ1 D3 D4; χ2] * + C_southwest[χ2; χ3] * + E_west[χ3 D5 D6; χ4] * + ket[d; D7 D1 D3 D5] * + conj(bra[d; D8 D2 D4 D6]) * + P_left[χ4 D7 D8; χ_out] +end """ renormalize_bottom_corner((r, c), envs, projectors) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 0cf3166c..20d2584d 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -188,19 +188,27 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N) """ ctmrg_expand(state, envs, alg::CTMRG{M}) + ctmrg_expand(dirs, state, envs::CTMRGEnv, alg::CTMRG) Expand the environment by absorbing a new PEPS tensor. There are two modes of expansion: `M = :sequential` and `M = :simultaneous`. The first mode expands the environment in one direction at a time, for convenience towards the left. The second mode expands the environment in all four directions simultaneously. +Alternatively, one can provide directly the `dirs` in which the environment is grown. """ -function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG) - drc_combinations = collect(Iterators.product([4, 1], axes(state)...)) - return map(idx -> enlarge_corner(idx, envs, state), drc_combinations) +function ctmrg_expand(state, envs::CTMRGEnv, alg::SequentialCTMRG) + return ctmrg_expand([4, 1], state, envs, alg) end -function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) - drc_combinations = collect(Iterators.product(1:4, axes(state)...)) - return map(idx -> enlarge_corner(idx, envs, state), drc_combinations) +function ctmrg_expand(state, envs::CTMRGEnv, alg::SimultaneousCTMRG) + return ctmrg_expand(1:4, state, envs, alg) +end +function ctmrg_expand(dirs, state, envs::CTMRGEnv, alg::CTMRG) + drc_combinations = collect(Iterators.product(dirs, axes(state)...)) + if alg.projector_alg.sparse + return map(idx -> enlarge_corner(idx, envs, state), drc_combinations) + else # Construct quadrant densely if alg is not sparse + return map(idx -> enlarge_corner(idx, envs, state)(idx[1]), drc_combinations) + end end """ @@ -264,11 +272,6 @@ function ctmrg_projectors( P_top = Zygote.Buffer(envs.edges, Prtype, axes(envs.corners, 2), axes(envs.corners, 3)) ϵ = zero(real(scalartype(envs))) - # If CTMRG is not sparse, construct all enlarged corners densely - sparse || enlarged_envs = map(Iterators.product(axes(enlarged_envs)...)) do (dir, r, c) - return enlarged_envs[dir, r, c](dir) - end - directions = collect(Iterators.product(axes(envs.corners, 2), axes(envs.corners, 3))) # @fwdthreads for (r, c) in directions for (r, c) in directions @@ -307,11 +310,6 @@ function ctmrg_projectors( # Corner type but with real numbers S = Zygote.Buffer(U.data, tensormaptype(spacetype(C), 1, 1, real(scalartype(E)))) - # If CTMRG is not sparse, construct all enlarged corners densely - sparse || enlarged_envs = map(Iterators.product(axes(enlarged_envs)...)) do (dir, r, c) - return enlarged_envs[dir, r, c](dir) - end - ϵ = zero(real(scalartype(envs))) drc_combinations = collect(Iterators.product(axes(envs.corners)...)) @fwdthreads for (dir, r, c) in drc_combinations @@ -363,7 +361,6 @@ end """ build_projectors(U::AbstractTensorMap{E,3,1}, S::AbstractTensorMap{E,1,1}, V::AbstractTensorMap{E,1,3}, Q::AbstractTensorMap{E,3,3}, Q_next::AbstractTensorMap{E,3,3}) where {E<:ElementarySpace} - build_projectors(U::AbstractTensorMap{E,3,1}, S::AbstractTensorMap{E,1,1}, V::AbstractTensorMap{E,1,3}, Q::EnlargedCorner, Q_next::EnlargedCorner) where {E<:ElementarySpace} diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index f32842a3..5d076a6d 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -24,16 +24,45 @@ i.e. the way the environment and PEPS tensors connect. function (Q::EnlargedCorner)(dir::Int) if dir == NORTHWEST return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - elseif dir == NORTHWEST - return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - elseif dir == NORTHWEST - return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - elseif dir == NORTHWEST + elseif dir == NORTHEAST + return enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == SOUTHEAST + return enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == SOUTHWEST + return enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + end +end +function construct_corner(Q::EnlargedCorner, dir::Int) + if dir == NORTHWEST return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == NORTHEAST + return enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == SOUTHEAST + return enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) + elseif dir == SOUTHWEST + return enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) end end -function renormalize_corner(ec::EnlargedCorner, P_left, P_right) - return renormalize_corner(ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra) + +function renormalize_northwest_corner(ec::EnlargedCorner, P_left, P_right) + return renormalize_northwest_corner( + ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra + ) +end +function renormalize_northeast_corner(ec::EnlargedCorner, P_left, P_right) + return renormalize_northeast_corner( + ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra + ) +end +function renormalize_southeast_corner(ec::EnlargedCorner, P_left, P_right) + return renormalize_southeast_corner( + ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra + ) +end +function renormalize_southwest_corner(ec::EnlargedCorner, P_left, P_right) + return renormalize_southwest_corner( + ec.E_1, ec.C, ec.E_2, P_left, P_right, ec.ket, ec.bra + ) end # ------------------ diff --git a/test/heisenberg.jl b/test/heisenberg.jl index 4fa67362..8e80944e 100644 --- a/test/heisenberg.jl +++ b/test/heisenberg.jl @@ -27,7 +27,7 @@ opt_alg = PEPSOptimize(; Random.seed!(91283219347) H = heisenberg_XYZ(InfiniteSquare()) psi_init = InfinitePEPS(2, χbond) -env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg) +env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg); # find fixedpoint result = fixedpoint(psi_init, H, opt_alg, env_init) From 0f2ca223fa7fbb86497de915c655dcc98ead01a8 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 24 Oct 2024 10:56:54 +0200 Subject: [PATCH 14/21] Remove sparse args and fix AD --- src/algorithms/ctmrg/ctmrg.jl | 49 +++++++++++++++++++---------------- 1 file changed, 27 insertions(+), 22 deletions(-) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 20d2584d..3568c4bd 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -8,20 +8,17 @@ have different spaces, this truncation style is different from `TruncationSpace` struct FixedSpaceTruncation <: TensorKit.TruncationScheme end """ - struct ProjectorAlg{S}(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, - sparse=Defaults.sparse, verbosity=0) + struct ProjectorAlg{S}(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, verbosity=0) Algorithm struct collecting all projector related parameters. The `svd_alg` sets the SVD algorithm for decomposing the CTM environment. The truncation scheme has to be a `TensorKit.TruncationScheme`, and some SVD algorithms might have further restrictions -on what kind of truncation scheme can be used. If `sparse` is `true`, then the enlarged corners and -CTM environment are never computed densely and function handles are using for the SVD. +on what kind of truncation scheme can be used. """ @kwdef struct ProjectorAlg{S<:SVDAdjoint,T} svd_alg::S = Defaults.svd_alg trscheme::T = Defaults.trscheme - sparse::Bool = Defaults.sparse verbosity::Int = 0 end # TODO: add option for different projector styles (half-infinite, full-infinite, etc.) @@ -48,7 +45,7 @@ end CTMRG(; tol=Defaults.ctmrg_tol, maxiter=Defaults.ctmrg_maxiter, miniter=Defaults.ctmrg_miniter, verbosity=0, svd_alg=SVDAdjoint(), trscheme=FixedSpaceTruncation(), - ctmrgscheme=Defaults.ctmrgscheme, sparse=Defaults.sparse) + ctmrgscheme=Defaults.ctmrgscheme) Algorithm struct that represents the CTMRG algorithm for contracting infinite PEPS. Each CTMRG run is converged up to `tol` where the singular value convergence of the @@ -65,10 +62,6 @@ CTMRG is implemented. It can either be `:sequential`, where the projectors are s computed on the western side, and then applied and rotated. Or with `:simultaneous` all projectors are computed and applied simultaneously on all sides, where in particular the corners get contracted with two projectors at the same time. - -The contraction and SVD of the CTMRG environments can be performed densely, or sparsely -if `sparse` is `true`. If sparsity is enabled, then the SVD will use only function handles -and the full enlarged corners and environments are never explicitly constructed. """ struct CTMRG{S} tol::Float64 @@ -85,14 +78,13 @@ function CTMRG(; svd_alg=Defaults.svd_alg, trscheme=Defaults.trscheme, ctmrgscheme::Symbol=Defaults.ctmrgscheme, - sparse::Bool=Defaults.sparse, ) return CTMRG{ctmrgscheme}( tol, maxiter, miniter, verbosity, - ProjectorAlg(; svd_alg, trscheme, sparse, verbosity), + ProjectorAlg(; svd_alg, trscheme, verbosity), ) end @@ -188,7 +180,7 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N) """ ctmrg_expand(state, envs, alg::CTMRG{M}) - ctmrg_expand(dirs, state, envs::CTMRGEnv, alg::CTMRG) + ctmrg_expand(dirs, state, envs::CTMRGEnv) Expand the environment by absorbing a new PEPS tensor. There are two modes of expansion: `M = :sequential` and `M = :simultaneous`. @@ -196,19 +188,32 @@ The first mode expands the environment in one direction at a time, for convenien the left. The second mode expands the environment in all four directions simultaneously. Alternatively, one can provide directly the `dirs` in which the environment is grown. """ -function ctmrg_expand(state, envs::CTMRGEnv, alg::SequentialCTMRG) - return ctmrg_expand([4, 1], state, envs, alg) +function ctmrg_expand(state, envs::CTMRGEnv, ::SequentialCTMRG) + return ctmrg_expand([4, 1], state, envs) end -function ctmrg_expand(state, envs::CTMRGEnv, alg::SimultaneousCTMRG) - return ctmrg_expand(1:4, state, envs, alg) +function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) + return ctmrg_expand(1:4, state, envs) end -function ctmrg_expand(dirs, state, envs::CTMRGEnv, alg::CTMRG) +# function ctmrg_expand(dirs, state, envs::CTMRGEnv) # TODO: This doesn't AD due to length(::Nothing)... +# drc_combinations = collect(Iterators.product(dirs, axes(state)...)) +# return map(idx -> enlarge_corner(idx, envs, state)(idx[1]), drc_combinations) +# end +function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} + Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) + Q = Zygote.Buffer(Array{Qtype,3}(undef, size(envs.corners))) drc_combinations = collect(Iterators.product(dirs, axes(state)...)) - if alg.projector_alg.sparse - return map(idx -> enlarge_corner(idx, envs, state), drc_combinations) - else # Construct quadrant densely if alg is not sparse - return map(idx -> enlarge_corner(idx, envs, state)(idx[1]), drc_combinations) + @fwdthreads for (dir, r, c) in drc_combinations + Q[dir, r, c] = if dir == NORTHWEST + enlarge_northwest_corner((r, c), envs, state) + elseif dir == NORTHEAST + enlarge_northeast_corner((r, c), envs, state) + elseif dir == SOUTHEAST + enlarge_southeast_corner((r, c), envs, state) + elseif dir == SOUTHWEST + enlarge_southwest_corner((r, c), envs, state) + end end + return copy(Q) end """ From c2c37c3454fbbab236d792724d868f3a1b335ee2 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 24 Oct 2024 11:24:59 +0200 Subject: [PATCH 15/21] Use enlarge_corner in ctmrg_expand --- src/algorithms/ctmrg/ctmrg.jl | 16 +++++----------- src/algorithms/ctmrg/sparse_environments.jl | 11 ----------- 2 files changed, 5 insertions(+), 22 deletions(-) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 3568c4bd..69c82ac2 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -200,28 +200,22 @@ end # end function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) + # Qtype = EnlargedCorner{C,T,eltype(state),eltype(state)} Q = Zygote.Buffer(Array{Qtype,3}(undef, size(envs.corners))) drc_combinations = collect(Iterators.product(dirs, axes(state)...)) @fwdthreads for (dir, r, c) in drc_combinations - Q[dir, r, c] = if dir == NORTHWEST - enlarge_northwest_corner((r, c), envs, state) - elseif dir == NORTHEAST - enlarge_northeast_corner((r, c), envs, state) - elseif dir == SOUTHEAST - enlarge_southeast_corner((r, c), envs, state) - elseif dir == SOUTHWEST - enlarge_southwest_corner((r, c), envs, state) - end + ec = enlarge_corner((dir, r, c), state, envs) + Q[dir, r, c] = ec(dir) # Explicitly construct EnlargedCorner for now end return copy(Q) end """ - enlarge_corner((dir, r, c), envs, state) + enlarge_corner((dir, r, c), state, envs) Enlarge corner by constructing the corresponding `EnlargedCorner` struct. """ -function enlarge_corner((dir, r, c), envs, state) +function enlarge_corner((dir, r, c), state, envs) if dir == NORTHWEST return EnlargedCorner( envs.corners[NORTHWEST, _prev(r, end), _prev(c, end)], diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index 5d076a6d..ca99f2a3 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -32,17 +32,6 @@ function (Q::EnlargedCorner)(dir::Int) return enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) end end -function construct_corner(Q::EnlargedCorner, dir::Int) - if dir == NORTHWEST - return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - elseif dir == NORTHEAST - return enlarge_northeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - elseif dir == SOUTHEAST - return enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - elseif dir == SOUTHWEST - return enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) - end -end function renormalize_northwest_corner(ec::EnlargedCorner, P_left, P_right) return renormalize_northwest_corner( From ec9195a47663e1cf02e1cc1b5ed90ecb0f5a072c Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 24 Oct 2024 14:20:01 +0200 Subject: [PATCH 16/21] Format and remove HalfInfiniteEnv TensorKit methods --- src/algorithms/ctmrg/ctmrg.jl | 7 +-- src/algorithms/ctmrg/sparse_environments.jl | 60 ++------------------- 2 files changed, 6 insertions(+), 61 deletions(-) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 69c82ac2..59f4b969 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -80,11 +80,7 @@ function CTMRG(; ctmrgscheme::Symbol=Defaults.ctmrgscheme, ) return CTMRG{ctmrgscheme}( - tol, - maxiter, - miniter, - verbosity, - ProjectorAlg(; svd_alg, trscheme, verbosity), + tol, maxiter, miniter, verbosity, ProjectorAlg(; svd_alg, trscheme, verbosity) ) end @@ -200,7 +196,6 @@ end # end function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) - # Qtype = EnlargedCorner{C,T,eltype(state),eltype(state)} Q = Zygote.Buffer(Array{Qtype,3}(undef, size(envs.corners))) drc_combinations = collect(Iterators.product(dirs, axes(state)...)) @fwdthreads for (dir, r, c) in drc_combinations diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index ca99f2a3..c6b826b4 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -63,7 +63,7 @@ end Half-infinite CTMRG environment tensor storage. """ -struct HalfInfiniteEnv{C,E,A,A′} +struct HalfInfiniteEnv{C,E,A,A′} # TODO: subtype as AbstractTensorMap once TensorKit is updated C_1::C C_2::C E_1::E @@ -159,68 +159,18 @@ function halfinfinite_environment(ec_1::EnlargedCorner, ec_2::EnlargedCorner) ) end -# ------------------------------------------------------------------------ -# Methods to make environment compatible with IterSVD and its reverse-rule -# ------------------------------------------------------------------------ - -TensorKit.InnerProductStyle(::HalfInfiniteEnv) = EuclideanProduct() -TensorKit.sectortype(::HalfInfiniteEnv) = Trivial -TensorKit.storagetype(env::HalfInfiniteEnv) = storagetype(env.ket_1) -TensorKit.spacetype(env::HalfInfiniteEnv) = spacetype(env.ket_1) +# ----------------------------------------------------- +# AbstractTensorMap subtyping and IterSVD compatibility +# ----------------------------------------------------- function TensorKit.domain(env::HalfInfiniteEnv) return domain(env.E_4) * domain(env.ket_2)[3] * domain(env.bra_2)[3]' end + function TensorKit.codomain(env::HalfInfiniteEnv) return codomain(env.E_1)[1] * domain(env.ket_1)[3]' * domain(env.bra_1)[3] end -function TensorKit.space(env::HalfInfiniteEnv) - return codomain(env) ← domain(env) -end -function TensorKit.blocks(env::HalfInfiniteEnv) - return TensorKit.SingletonDict(Trivial() => env) -end -function TensorKit.blocksectors(::HalfInfiniteEnv) - return TensorKit.OneOrNoneIterator{Trivial}(true, Trivial()) -end -function TensorKit.block(env::HalfInfiniteEnv, c::Sector) - return env -end -function TensorKit.tsvd!(f::HalfInfiniteEnv; trunc=NoTruncation(), p::Real=2, alg=IterSVD()) - return _tsvd!(f, alg, trunc, p) -end -function TensorKit.MatrixAlgebra.svd!(env::HalfInfiniteEnv, args...) - return TensorKit.MatrixAlgebra.svd!(env(), args...) # Construct environment densely as fallback -end - -Base.eltype(env::HalfInfiniteEnv) = eltype(env.ket_1) -function Base.size(env::HalfInfiniteEnv) # Treat environment as matrix - χ_in = dim(space(env.E_1, 1)) - D_inabove = dim(space(env.ket_1, 2)) - D_inbelow = dim(space(env.bra_1, 2)) - χ_out = dim(space(env.E_4, 1)) - D_outabove = dim(space(env.ket_2, 2)) - D_outbelow = dim(space(env.bra_2, 2)) - return (χ_in * D_inabove * D_inbelow, χ_out * D_outabove * D_outbelow) -end -Base.size(env::HalfInfiniteEnv, i::Int) = size(env)[i] -VectorInterface.scalartype(env::HalfInfiniteEnv) = scalartype(env.ket_1) function random_start_vector(env::HalfInfiniteEnv) return Tensor(randn, domain(env)) end - -function Base.similar(env::HalfInfiniteEnv) - return HalfInfiniteEnv( - (similar(getfield(env, field)) for field in fieldnames(HalfInfiniteEnv))... - ) -end - -function Base.copyto!(dest::HalfInfiniteEnv, src::HalfInfiniteEnv) - for field in fieldnames(HalfInfiniteEnv) - for (bd, bs) in zip(blocks(getfield(dest, field)), blocks(getfield(src, field))) - copyto!(bd[2], bs[2]) - end - end - return dest -end From 77259d7edfab792cf763b89c04303ca16269b6d6 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 24 Oct 2024 17:21:24 +0200 Subject: [PATCH 17/21] Fix wrong daggers in relative gauge fix contractions --- .../contractions/ctmrg_contractions.jl | 16 ++++++++-------- src/utility/svd.jl | 7 +------ test/ctmrg/fixed_iterscheme.jl | 4 +--- 3 files changed, 10 insertions(+), 17 deletions(-) diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index f4d32d0b..e53bf2ac 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -872,7 +872,7 @@ end Multiply north left singular vectors with gauge signs from the right. """ function fix_gauge_north_left_vecs((row, col), U, signs) - return U[NORTH, row, col] * signs[NORTH, row, _next(col, end)] + return U[NORTH, row, col] * signs[NORTH, row, _next(col, end)]' end """ @@ -881,7 +881,7 @@ end Multiply east left singular vectors with gauge signs from the right. """ function fix_gauge_east_left_vecs((row, col), U, signs) - return U[EAST, row, col] * signs[EAST, _next(row, end), col] + return U[EAST, row, col] * signs[EAST, _next(row, end), col]' end """ @@ -890,7 +890,7 @@ end Multiply south left singular vectors with gauge signs from the right. """ function fix_gauge_south_left_vecs((row, col), U, signs) - return U[SOUTH, row, col] * signs[SOUTH, row, _prev(col, end)] + return U[SOUTH, row, col] * signs[SOUTH, row, _prev(col, end)]' end """ @@ -899,7 +899,7 @@ end Multiply west left singular vectors with gauge signs from the right. """ function fix_gauge_west_left_vecs((row, col), U, signs) - return U[WEST, row, col] * signs[WEST, _prev(row, end), col] + return U[WEST, row, col] * signs[WEST, _prev(row, end), col]' end # right singular vectors @@ -910,7 +910,7 @@ end Multiply north right singular vectors with gauge signs from the left. """ function fix_gauge_north_right_vecs((row, col), V, signs) - return signs[NORTH, row, _next(col, end)]' * V[NORTH, row, col] + return signs[NORTH, row, _next(col, end)] * V[NORTH, row, col] end """ @@ -919,7 +919,7 @@ end Multiply east right singular vectors with gauge signs from the left. """ function fix_gauge_east_right_vecs((row, col), V, signs) - return signs[EAST, _next(row, end), col]' * V[EAST, row, col] + return signs[EAST, _next(row, end), col] * V[EAST, row, col] end """ @@ -928,7 +928,7 @@ end Multiply south right singular vectors with gauge signs from the left. """ function fix_gauge_south_right_vecs((row, col), V, signs) - return signs[SOUTH, row, _prev(col, end)]' * V[SOUTH, row, col] + return signs[SOUTH, row, _prev(col, end)] * V[SOUTH, row, col] end """ @@ -937,5 +937,5 @@ end Multiply west right singular vectors with gauge signs from the left. """ function fix_gauge_west_right_vecs((row, col), V, signs) - return signs[WEST, _prev(row, end), col]' * V[WEST, row, col] + return signs[WEST, _prev(row, end), col] * V[WEST, row, col] end diff --git a/src/utility/svd.jl b/src/utility/svd.jl index d439b4f6..8b946f51 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -70,13 +70,8 @@ the iterative SVD didn't converge, the algorithm falls back to a dense SVD. start_vector = random_start_vector end -# TODO: find better initial guess that leads to element-wise convergence and is compatible with function handles function random_start_vector(t::Matrix) - return randn(eltype(t), size(t, 1)) # Leads to erroneous gauge fixing of U, S, V and thus failing element-wise conv. - # u, = TensorKit.MatrixAlgebra.svd!(deepcopy(t), TensorKit.SVD()) - # return sum(u[:, i] for i in 1:howmany) # Element-wise convergence works fine - # return dropdims(sum(t[:, 1:3]; dims=2); dims=2) # Summing too many columns also makes gauge fixing fail - # return t[:, 1] # Leads so slower convergence of SVD than randn, but correct element-wise convergence + return randn(eltype(t), size(t, 1)) end # Compute SVD data block-wise using KrylovKit algorithm diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl index 871e8711..49a1568a 100644 --- a/test/ctmrg/fixed_iterscheme.jl +++ b/test/ctmrg/fixed_iterscheme.jl @@ -1,5 +1,6 @@ using Test using Random +using LinearAlgebra using TensorKit, KrylovKit using PEPSKit using PEPSKit: @@ -66,7 +67,6 @@ end psi = InfinitePEPS(2, χbond) env_init = CTMRGEnv(psi, ComplexSpace(χenv)) env_conv1 = leading_boundary(env_init, psi, ctm_alg_iter) - # env_conv1 = leading_boundary(env_init, psi, ctm_alg_full); # do extra iteration to get SVD env_conv2_iter, info_iter = ctmrg_iter(psi, env_conv1, ctm_alg_iter) @@ -94,12 +94,10 @@ end env_fixedsvd_iter, = ctmrg_iter(psi, env_conv1, ctm_alg_fix_iter) env_fixedsvd_iter = fix_global_phases(env_conv1, env_fixedsvd_iter) @test calc_elementwise_convergence(env_conv1, env_fixedsvd_iter) ≈ 0 atol = 1e-6 # This doesn't work for x₀ = rand(size(b, 1))? - # @show calc_elementwise_convergence(env_conv1, env_fixedsvd_iter) env_fixedsvd_full, = ctmrg_iter(psi, env_conv1, ctm_alg_fix_full) env_fixedsvd_full = fix_global_phases(env_conv1, env_fixedsvd_full) @test calc_elementwise_convergence(env_conv1, env_fixedsvd_full) ≈ 0 atol = 1e-6 - # @show calc_elementwise_convergence(env_conv1, env_fixedsvd_full) # check matching decompositions atol = 1e-12 From bb4af77bfe9dea62de3b30dc3c078a17dfdd166a Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Thu, 24 Oct 2024 17:45:29 +0200 Subject: [PATCH 18/21] Fix :sequential mode in ctmrg_expand --- src/algorithms/ctmrg/ctmrg.jl | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 59f4b969..6e708857 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -196,11 +196,12 @@ end # end function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) - Q = Zygote.Buffer(Array{Qtype,3}(undef, size(envs.corners))) - drc_combinations = collect(Iterators.product(dirs, axes(state)...)) - @fwdthreads for (dir, r, c) in drc_combinations - ec = enlarge_corner((dir, r, c), state, envs) - Q[dir, r, c] = ec(dir) # Explicitly construct EnlargedCorner for now + Q = Zygote.Buffer(Array{Qtype,3}(undef, length(dirs), size(state)...)) + dirs_enum = [(i, dir) for (i, dir) in enumerate(dirs)] + drc_combinations = collect(Iterators.product(dirs_enum, axes(state)...)) + @fwdthreads for (d, r, c) in drc_combinations + ec = enlarge_corner((d[2], r, c), state, envs) + Q[d[1], r, c] = ec(d[2]) # Explicitly construct EnlargedCorner for now end return copy(Q) end From 899b10b61452f225425866ae1bde0974d31e5d4a Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Fri, 25 Oct 2024 10:07:01 +0200 Subject: [PATCH 19/21] Comment out HalfInfiniteEnv SVD test --- test/ctmrg/svd_wrapper.jl | 90 +++++++++++++++++++-------------------- 1 file changed, 45 insertions(+), 45 deletions(-) diff --git a/test/ctmrg/svd_wrapper.jl b/test/ctmrg/svd_wrapper.jl index 68a53aa8..01c4f39e 100644 --- a/test/ctmrg/svd_wrapper.jl +++ b/test/ctmrg/svd_wrapper.jl @@ -6,7 +6,7 @@ using KrylovKit using ChainRulesCore, Zygote using Accessors using PEPSKit -using PEPSKit: HalfInfiniteEnv +# using PEPSKit: HalfInfiniteEnv # Gauge-invariant loss function function lossfun(A, alg, R=TensorMap(randn, space(A)), trunc=notrunc()) @@ -94,47 +94,47 @@ symm_R = TensorMap(randn, dtype, space(symm_r)) @test g_fullsvd_tr[1] ≈ g_itersvd_fb[1] rtol = rtol end -χbond = 2 -χenv = 6 -ctm_alg = CTMRG(; tol=1e-10, verbosity=2, svd_alg=SVDAdjoint()) -Random.seed!(91283219347) -H = heisenberg_XYZ(InfiniteSquare()) -psi = InfinitePEPS(2, χbond) -env = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg); -hienv = HalfInfiniteEnv( - env.corners[1], - env.corners[2], - env.edges[4], - env.edges[1], - env.edges[1], - env.edges[2], - psi[1], - psi[1], - psi[1], - psi[1], -) -hienv_dense = hienv() -env_R = TensorMap(randn, space(hienv)) - -PEPSKit.tsvd!(hienv, iter_alg) # TODO: make the space mismatches work - -@testset "IterSVD with HalfInfiniteEnv function handle" begin - # Equivalence of dense and sparse contractions - x₀ = PEPSKit.random_start_vector(hienv) - x′ = hienv(x₀, Val(false)) - x″ = hienv(x′, Val(true)) - x‴ = hienv(x″, Val(false)) - - a = hienv_dense * x₀ - b = hienv_dense' * a - c = hienv_dense * b - @test a ≈ x′ - @test b ≈ x″ - @test c ≈ x‴ - - # TODO: code up pullback - # l_fullsvd, g_fullsvd = withgradient(A -> lossfun(A, full_alg, env_R), hienv_dense) - # l_itersvd, g_itersvd = withgradient(A -> lossfun(A, iter_alg, env_R), hienv) - # @test l_itersvd ≈ l_fullsvd - # @test g_fullsvd[1] ≈ g_itersvd[1] rtol = rtol -end +# TODO: Add when IterSVD is implemented for HalfInfiniteEnv +# χbond = 2 +# χenv = 6 +# ctm_alg = CTMRG(; tol=1e-10, verbosity=2, svd_alg=SVDAdjoint()) +# Random.seed!(91283219347) +# H = heisenberg_XYZ(InfiniteSquare()) +# psi = InfinitePEPS(2, χbond) +# env = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg); +# hienv = HalfInfiniteEnv( +# env.corners[1], +# env.corners[2], +# env.edges[4], +# env.edges[1], +# env.edges[1], +# env.edges[2], +# psi[1], +# psi[1], +# psi[1], +# psi[1], +# ) +# hienv_dense = hienv() +# env_R = TensorMap(randn, space(hienv)) + +# PEPSKit.tsvd!(hienv, iter_alg) + +# @testset "IterSVD with HalfInfiniteEnv function handle" begin +# # Equivalence of dense and sparse contractions +# x₀ = PEPSKit.random_start_vector(hienv) +# x′ = hienv(x₀, Val(false)) +# x″ = hienv(x′, Val(true)) +# x‴ = hienv(x″, Val(false)) + +# a = hienv_dense * x₀ +# b = hienv_dense' * a +# c = hienv_dense * b +# @test a ≈ x′ +# @test b ≈ x″ +# @test c ≈ x‴ + +# # l_fullsvd, g_fullsvd = withgradient(A -> lossfun(A, full_alg, env_R), hienv_dense) +# # l_itersvd, g_itersvd = withgradient(A -> lossfun(A, iter_alg, env_R), hienv) +# # @test l_itersvd ≈ l_fullsvd +# # @test g_fullsvd[1] ≈ g_itersvd[1] rtol = rtol +# end From e31dea04f2c03f374b925fab34ba083b60a615f1 Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 28 Oct 2024 11:05:37 +0100 Subject: [PATCH 20/21] Replace enlarge_corner with constructor --- src/algorithms/ctmrg/ctmrg.jl | 45 +-------------------- src/algorithms/ctmrg/sparse_environments.jl | 43 ++++++++++++++++++++ 2 files changed, 45 insertions(+), 43 deletions(-) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 6e708857..8167f37c 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -192,7 +192,7 @@ function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) end # function ctmrg_expand(dirs, state, envs::CTMRGEnv) # TODO: This doesn't AD due to length(::Nothing)... # drc_combinations = collect(Iterators.product(dirs, axes(state)...)) -# return map(idx -> enlarge_corner(idx, envs, state)(idx[1]), drc_combinations) +# return map(idx -> EnlargedCorner(state, envs, idx)(idx[1]), drc_combinations) # end function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) @@ -200,53 +200,12 @@ function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} dirs_enum = [(i, dir) for (i, dir) in enumerate(dirs)] drc_combinations = collect(Iterators.product(dirs_enum, axes(state)...)) @fwdthreads for (d, r, c) in drc_combinations - ec = enlarge_corner((d[2], r, c), state, envs) + ec = EnlargedCorner(state, envs, (d[2], r, c)) Q[d[1], r, c] = ec(d[2]) # Explicitly construct EnlargedCorner for now end return copy(Q) end -""" - enlarge_corner((dir, r, c), state, envs) - -Enlarge corner by constructing the corresponding `EnlargedCorner` struct. -""" -function enlarge_corner((dir, r, c), state, envs) - if dir == NORTHWEST - return EnlargedCorner( - envs.corners[NORTHWEST, _prev(r, end), _prev(c, end)], - envs.edges[WEST, r, _prev(c, end)], - envs.edges[NORTH, _prev(r, end), c], - state[r, c], - state[r, c], - ) - elseif dir == NORTHEAST - return EnlargedCorner( - envs.corners[NORTHEAST, _prev(r, end), _next(c, end)], - envs.edges[NORTH, _prev(r, end), c], - envs.edges[EAST, r, _next(c, end)], - state[r, c], - state[r, c], - ) - elseif dir == SOUTHEAST - return EnlargedCorner( - envs.corners[SOUTHEAST, _next(r, end), _next(c, end)], - envs.edges[EAST, r, _next(c, end)], - envs.edges[SOUTH, _next(r, end), c], - state[r, c], - state[r, c], - ) - elseif dir == SOUTHWEST - return EnlargedCorner( - envs.corners[SOUTHWEST, _next(r, end), _prev(c, end)], - envs.edges[SOUTH, _next(r, end), c], - envs.edges[WEST, r, _prev(c, end)], - state[r, c], - state[r, c], - ) - end -end - # ======================================================================================== # # Projector step # ======================================================================================== # diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index c6b826b4..44eca222 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -15,6 +15,49 @@ struct EnlargedCorner{Ct,E,A,A′} bra::A′ end +""" + EnlargedCorner(state, envs, coordinates) + +Construct an enlarged corner with the correct row and column indices based on the given +`coordinates` which are of the form `(dir, row, col)`. +""" +function EnlargedCorner(state, envs, coordinates) + dir, r, c = coordinates + if dir == NORTHWEST + return EnlargedCorner( + envs.corners[NORTHWEST, _prev(r, end), _prev(c, end)], + envs.edges[WEST, r, _prev(c, end)], + envs.edges[NORTH, _prev(r, end), c], + state[r, c], + state[r, c], + ) + elseif dir == NORTHEAST + return EnlargedCorner( + envs.corners[NORTHEAST, _prev(r, end), _next(c, end)], + envs.edges[NORTH, _prev(r, end), c], + envs.edges[EAST, r, _next(c, end)], + state[r, c], + state[r, c], + ) + elseif dir == SOUTHEAST + return EnlargedCorner( + envs.corners[SOUTHEAST, _next(r, end), _next(c, end)], + envs.edges[EAST, r, _next(c, end)], + envs.edges[SOUTH, _next(r, end), c], + state[r, c], + state[r, c], + ) + elseif dir == SOUTHWEST + return EnlargedCorner( + envs.corners[SOUTHWEST, _next(r, end), _prev(c, end)], + envs.edges[SOUTH, _next(r, end), c], + envs.edges[WEST, r, _prev(c, end)], + state[r, c], + state[r, c], + ) + end +end + """ (Q::EnlargedCorner)(dir::Int) From 1a8f7889c7e3d5d0a41d8cb7a1acdb23375addcb Mon Sep 17 00:00:00 2001 From: Paul Brehmer Date: Mon, 28 Oct 2024 19:59:15 +0100 Subject: [PATCH 21/21] Replace function call with TensorMap constructor --- src/algorithms/ctmrg/ctmrg.jl | 4 ++-- src/algorithms/ctmrg/sparse_environments.jl | 25 ++++++++++++--------- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 8167f37c..1dd07c09 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -192,7 +192,7 @@ function ctmrg_expand(state, envs::CTMRGEnv, ::SimultaneousCTMRG) end # function ctmrg_expand(dirs, state, envs::CTMRGEnv) # TODO: This doesn't AD due to length(::Nothing)... # drc_combinations = collect(Iterators.product(dirs, axes(state)...)) -# return map(idx -> EnlargedCorner(state, envs, idx)(idx[1]), drc_combinations) +# return map(idx -> TensorMap(EnlargedCorner(state, envs, idx), idx[1]), drc_combinations) # end function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} Qtype = tensormaptype(spacetype(C), 3, 3, storagetype(C)) @@ -201,7 +201,7 @@ function ctmrg_expand(dirs, state, envs::CTMRGEnv{C,T}) where {C,T} drc_combinations = collect(Iterators.product(dirs_enum, axes(state)...)) @fwdthreads for (d, r, c) in drc_combinations ec = EnlargedCorner(state, envs, (d[2], r, c)) - Q[d[1], r, c] = ec(d[2]) # Explicitly construct EnlargedCorner for now + Q[d[1], r, c] = TensorMap(ec, d[2]) # Explicitly construct EnlargedCorner for now end return copy(Q) end diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index 44eca222..9243067e 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -59,12 +59,12 @@ function EnlargedCorner(state, envs, coordinates) end """ - (Q::EnlargedCorner)(dir::Int) + TensorKit.TensorMap(Q::EnlargedCorner, dir::Int) -Contract enlarged corner where `dir` selects the correct contraction direction, -i.e. the way the environment and PEPS tensors connect. +Instantiate enlarged corner as `TensorMap` where `dir` selects the correct contraction +direction, i.e. the way the environment and PEPS tensors connect. """ -function (Q::EnlargedCorner)(dir::Int) +function TensorKit.TensorMap(Q::EnlargedCorner, dir::Int) if dir == NORTHWEST return enlarge_northwest_corner(Q.E_1, Q.C, Q.E_2, Q.ket, Q.bra) elseif dir == NORTHEAST @@ -134,14 +134,11 @@ function HalfInfiniteEnv(quadrant1::EnlargedCorner, quadrant2::EnlargedCorner) end """ - (env::HalfInfiniteEnv)() - (env::HalfInfiniteEnv)(x, ::Val{false}) - (env::HalfInfiniteEnv)(x, ::Val{true}) + TensorKit.TensorMap(env::HalfInfiniteEnv) -Contract half-infinite environment. If a vector `x` is provided, the environment acts as a -linear map or adjoint linear map on `x` if `Val(true)` or `Val(false)` is passed, respectively. +Instantiate half-infinite environment as `TensorMap` explicitly. """ -function (env::HalfInfiniteEnv)() # Dense operator +function TensorKit.TensorMap(env::HalfInfiniteEnv) # Dense operator return halfinfinite_environment( env.C_1, env.C_2, @@ -155,6 +152,14 @@ function (env::HalfInfiniteEnv)() # Dense operator env.bra_2, ) end + +""" + (env::HalfInfiniteEnv)(x, ::Val{false}) + (env::HalfInfiniteEnv)(x, ::Val{true}) + +Contract half-infinite environment with a vector `x`, such that the environment acts as a +linear map or adjoint linear map on `x` if `Val(true)` or `Val(false)` is passed, respectively. +""" function (env::HalfInfiniteEnv)(x, ::Val{false}) # Linear map: env() * x return halfinfinite_environment( env.C_1,