diff --git a/Project.toml b/Project.toml index cb1675223..11fb333be 100644 --- a/Project.toml +++ b/Project.toml @@ -43,6 +43,7 @@ Random = "1" Statistics = "1" TensorKit = "0.14.6" TensorOperations = "5" +TestExtras = "0.3" VectorInterface = "0.4, 0.5" Zygote = "0.6, 0.7" julia = "1.10" @@ -53,6 +54,7 @@ FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [targets] -test = ["Test", "SafeTestsets", "ChainRulesTestUtils", "FiniteDifferences", "QuadGK"] +test = ["Test", "TestExtras", "SafeTestsets", "ChainRulesTestUtils", "FiniteDifferences", "QuadGK"] diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 6c2b83a4c..a8476abb3 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -15,7 +15,7 @@ import MPSKit: leading_boundary, loginit!, logiter!, logfinish!, logcancel!, phy using MPSKitModels using FiniteDifferences -using OhMyThreads: tmap +using OhMyThreads: tmap, tmap! using DocStringExtensions include("Defaults.jl") # Include first to allow for docstring interpolation with Defaults values diff --git a/src/algorithms/contractions/ctmrg_contractions.jl b/src/algorithms/contractions/ctmrg_contractions.jl index c8b6ff3b7..40b7f8dc2 100644 --- a/src/algorithms/contractions/ctmrg_contractions.jl +++ b/src/algorithms/contractions/ctmrg_contractions.jl @@ -1310,12 +1310,19 @@ end function renormalize_west_edge( E_west::CTMRG_PEPS_EdgeTensor, P_bottom, P_top, A::PEPSSandwich ) - return @autoopt @tensor edge[χ_S D_Eab D_Ebe; χ_N] := - E_west[χ1 D1 D2; χ2] * - ket(A)[d; D3 D_Eab D5 D1] * - conj(bra(A)[d; D4 D_Ebe D6 D2]) * - P_bottom[χ2 D3 D4; χ_N] * - P_top[χ_S; χ1 D5 D6] + # starting with P_top to save one permute in the end + return @tensor begin + # already putting χE in front here to make next permute cheaper + PE[χS χNW DSb DWb; DSt DWt] := P_top[χS; χSW DSt DSb] * E_west[χSW DWt DWb; χNW] + + PEket[χS χNW DNt DEt; DSb DWb d] := + PE[χS χNW DSb DWb; DSt DWt] * ket(A)[d; DNt DEt DSt DWt] + + corner[χS DEt DEb; χNW DNt DNb] := + PEket[χS χNW DNt DEt; DSb DWb d] * conj(bra(A)[d; DNb DEb DSb DWb]) + + edge[χS DEt DEb; χN] := corner[χS DEt DEb; χNW DNt DNb] * P_bottom[χNW DNt DNb; χN] + end end function renormalize_west_edge(E_west::CTMRG_PF_EdgeTensor, P_bottom, P_top, A::PFTensor) return @autoopt @tensor edge[χ_S D_E; χ_N] := diff --git a/src/algorithms/correlators.jl b/src/algorithms/correlators.jl index 9b5ce5a0e..76d32a37e 100644 --- a/src/algorithms/correlators.jl +++ b/src/algorithms/correlators.jl @@ -120,7 +120,7 @@ function end_correlator_numerator( C_southeast = env.corners[SOUTHEAST, _next(r, end), _next(c, end)] sandwich = (above[mod1(r, end), mod1(c, end)], below[mod1(r, end), mod1(c, end)]) - return @autoopt @tensor contractcheck = true V[χSW DWt dstring DWb; χNW] * + return @autoopt @tensor V[χSW DWt dstring DWb; χNW] * E_south[χSSE DSt DSb; χSW] * E_east[χNEE DEt DEb; χSEE] * E_north[χNW DNt DNb; χNNE] * diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index 0c688f3be..f781f5807 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -110,14 +110,12 @@ end function leading_boundary( env₀::CTMRGEnv, network::InfiniteSquareNetwork, alg::CTMRGAlgorithm ) - CS = map(x -> tsvd(x)[2], env₀.corners) - TS = map(x -> tsvd(x)[2], env₀.edges) - - η = one(real(scalartype(network))) - env = deepcopy(env₀) log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG")) - return LoggingExtras.withlevel(; alg.verbosity) do + env = deepcopy(env₀) + CS = map(x -> tsvd(x)[2], env₀.corners) + TS = map(x -> tsvd(x)[2], env₀.edges) + η = one(real(scalartype(network))) ctmrg_loginit!(log, η, network, env₀) local info for iter in 1:(alg.maxiter) diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl index 6ec8e7581..adeb94d51 100644 --- a/src/algorithms/ctmrg/sequential.jl +++ b/src/algorithms/ctmrg/sequential.jl @@ -24,12 +24,12 @@ For a full description, see [`leading_boundary`](@ref). The supported keywords a * `svd_alg::Union{<:SVDAdjoint,NamedTuple}` * `projector_alg::Symbol=:$(Defaults.projector_alg)` """ -struct SequentialCTMRG <: CTMRGAlgorithm +struct SequentialCTMRG{P<:ProjectorAlgorithm} <: CTMRGAlgorithm tol::Float64 maxiter::Int miniter::Int verbosity::Int - projector_alg::ProjectorAlgorithm + projector_alg::P end function SequentialCTMRG(; kwargs...) return CTMRGAlgorithm(; alg=:sequential, kwargs...) @@ -81,14 +81,18 @@ for a specific `coordinate` (where `dir=WEST` is already implied in the `:sequen """ function sequential_projectors(col::Int, network, env::CTMRGEnv, alg::ProjectorAlgorithm) coordinates = eachcoordinate(env)[:, col] - proj_and_info = dtmap(coordinates) do (r, c) + T_dst = Base.promote_op( + sequential_projectors, NTuple{3,Int}, typeof(network), typeof(env), typeof(alg) + ) + proj_and_info = similar(coordinates, T_dst) + proj_and_info′::typeof(proj_and_info) = dtmap!!(proj_and_info, coordinates) do (r, c) trscheme = truncation_scheme(alg, env.edges[WEST, _prev(r, size(env, 2)), c]) proj, info = sequential_projectors( (WEST, r, c), network, env, @set(alg.trscheme = trscheme) ) return proj, info end - return _split_proj_and_info(proj_and_info) + return _split_proj_and_info(proj_and_info′) end function sequential_projectors( coordinate::NTuple{3,Int}, network, env::CTMRGEnv, alg::HalfInfiniteProjector diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl index 0dd47ef42..01d6139ee 100644 --- a/src/algorithms/ctmrg/simultaneous.jl +++ b/src/algorithms/ctmrg/simultaneous.jl @@ -23,12 +23,12 @@ For a full description, see [`leading_boundary`](@ref). The supported keywords a * `svd_alg::Union{<:SVDAdjoint,NamedTuple}` * `projector_alg::Symbol=:$(Defaults.projector_alg)` """ -struct SimultaneousCTMRG <: CTMRGAlgorithm +struct SimultaneousCTMRG{P<:ProjectorAlgorithm} <: CTMRGAlgorithm tol::Float64 maxiter::Int miniter::Int verbosity::Int - projector_alg::ProjectorAlgorithm + projector_alg::P end function SimultaneousCTMRG(; kwargs...) return CTMRGAlgorithm(; alg=:simultaneous, kwargs...) @@ -37,9 +37,15 @@ end CTMRG_SYMBOLS[:simultaneous] = SimultaneousCTMRG function ctmrg_iteration(network, env::CTMRGEnv, alg::SimultaneousCTMRG) - enlarged_corners = dtmap(eachcoordinate(network, 1:4)) do idx - return TensorMap(EnlargedCorner(network, env, idx)) - end # expand environment + coordinates = eachcoordinate(network, 1:4) + T_corners = Base.promote_op( + TensorMap ∘ EnlargedCorner, typeof(network), typeof(env), eltype(coordinates) + ) + enlarged_corners′ = similar(coordinates, T_corners) + enlarged_corners::typeof(enlarged_corners′) = + dtmap!!(enlarged_corners′, eachcoordinate(network, 1:4)) do idx + return TensorMap(EnlargedCorner(network, env, idx)) + end # expand environment projectors, info = simultaneous_projectors(enlarged_corners, env, alg.projector_alg) # compute projectors on all coordinates env′ = renormalize_simultaneously(enlarged_corners, projectors, network, env) # renormalize enlarged corners return env′, info @@ -72,25 +78,36 @@ enlarged corners or on a specific `coordinate`. function simultaneous_projectors( enlarged_corners::Array{E,3}, env::CTMRGEnv, alg::ProjectorAlgorithm ) where {E} - proj_and_info = dtmap(eachcoordinate(env, 1:4)) do coordinate - coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...) - trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...]) - return simultaneous_projectors( - coordinate, enlarged_corners, @set(alg.trscheme = trscheme) - ) - end + coordinates = eachcoordinate(env, 1:4) + T_dst = Base.promote_op( + simultaneous_projectors, + NTuple{3,Int}, + typeof(enlarged_corners), + typeof(env), + typeof(alg), + ) + proj_and_info′ = similar(coordinates, T_dst) + proj_and_info::typeof(proj_and_info′) = + dtmap!!(proj_and_info′, coordinates) do coordinate + return simultaneous_projectors(coordinate, enlarged_corners, env, alg) + end return _split_proj_and_info(proj_and_info) end function simultaneous_projectors( - coordinate, enlarged_corners::Array{E,3}, alg::HalfInfiniteProjector + coordinate, enlarged_corners::Array{E,3}, env, alg::HalfInfiniteProjector ) where {E} - coordinate′ = _next_coordinate(coordinate, size(enlarged_corners)[2:3]...) + coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...) + trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...]) + alg′ = @set alg.trscheme = trscheme ec = (enlarged_corners[coordinate...], enlarged_corners[coordinate′...]) - return compute_projector(ec, coordinate, alg) + return compute_projector(ec, coordinate, alg′) end function simultaneous_projectors( - coordinate, enlarged_corners::Array{E,3}, alg::FullInfiniteProjector + coordinate, enlarged_corners::Array{E,3}, env, alg::FullInfiniteProjector ) where {E} + coordinate′ = _next_coordinate(coordinate, size(env)[2:3]...) + trscheme = truncation_scheme(alg, env.edges[coordinate[1], coordinate′[2:3]...]) + alg′ = @set alg.trscheme = trscheme rowsize, colsize = size(enlarged_corners)[2:3] coordinate2 = _next_coordinate(coordinate, rowsize, colsize) coordinate3 = _next_coordinate(coordinate2, rowsize, colsize) @@ -101,7 +118,7 @@ function simultaneous_projectors( enlarged_corners[coordinate2...], enlarged_corners[coordinate3...], ) - return compute_projector(ec, coordinate, alg) + return compute_projector(ec, coordinate, alg′) end """ @@ -112,22 +129,33 @@ Renormalize all enlarged corners and edges simultaneously. function renormalize_simultaneously(enlarged_corners, projectors, network, env) P_left, P_right = projectors coordinates = eachcoordinate(env, 1:4) - corners_edges = dtmap(coordinates) do (dir, r, c) - if dir == NORTH - corner = renormalize_northwest_corner((r, c), enlarged_corners, P_left, P_right) - edge = renormalize_north_edge((r, c), env, P_left, P_right, network) - elseif dir == EAST - corner = renormalize_northeast_corner((r, c), enlarged_corners, P_left, P_right) - edge = renormalize_east_edge((r, c), env, P_left, P_right, network) - elseif dir == SOUTH - corner = renormalize_southeast_corner((r, c), enlarged_corners, P_left, P_right) - edge = renormalize_south_edge((r, c), env, P_left, P_right, network) - elseif dir == WEST - corner = renormalize_southwest_corner((r, c), enlarged_corners, P_left, P_right) - edge = renormalize_west_edge((r, c), env, P_left, P_right, network) + T_CE = Tuple{cornertype(env),edgetype(env)} + corners_edges′ = similar(coordinates, T_CE) + corners_edges::typeof(corners_edges′) = + dtmap!!(corners_edges′, coordinates) do (dir, r, c) + if dir == NORTH + corner = renormalize_northwest_corner( + (r, c), enlarged_corners, P_left, P_right + ) + edge = renormalize_north_edge((r, c), env, P_left, P_right, network) + elseif dir == EAST + corner = renormalize_northeast_corner( + (r, c), enlarged_corners, P_left, P_right + ) + edge = renormalize_east_edge((r, c), env, P_left, P_right, network) + elseif dir == SOUTH + corner = renormalize_southeast_corner( + (r, c), enlarged_corners, P_left, P_right + ) + edge = renormalize_south_edge((r, c), env, P_left, P_right, network) + elseif dir == WEST + corner = renormalize_southwest_corner( + (r, c), enlarged_corners, P_left, P_right + ) + edge = renormalize_west_edge((r, c), env, P_left, P_right, network) + end + return corner / norm(corner), edge / norm(edge) end - return corner / norm(corner), edge / norm(edge) - end return CTMRGEnv(map(first, corners_edges), map(last, corners_edges)) end diff --git a/src/algorithms/ctmrg/sparse_environments.jl b/src/algorithms/ctmrg/sparse_environments.jl index f53747d7c..1085be036 100644 --- a/src/algorithms/ctmrg/sparse_environments.jl +++ b/src/algorithms/ctmrg/sparse_environments.jl @@ -56,6 +56,8 @@ function EnlargedCorner(network::InfiniteSquareNetwork, env, coordinates) network[r, c], dir, ) + else + throw(ArgumentError(lazy"Invalid direction $dir")) end end @@ -73,6 +75,8 @@ function TensorKit.TensorMap(Q::EnlargedCorner) return enlarge_southeast_corner(Q.E_1, Q.C, Q.E_2, Q.A) elseif Q.dir == SOUTHWEST return enlarge_southwest_corner(Q.E_1, Q.C, Q.E_2, Q.A) + else + throw(ArgumentError(lazy"Invalid direction $dir")) end end diff --git a/src/utility/diffable_threads.jl b/src/utility/diffable_threads.jl index b01375482..ff9c96af6 100644 --- a/src/utility/diffable_threads.jl +++ b/src/utility/diffable_threads.jl @@ -8,6 +8,8 @@ All calls of `dtmap` inside of PEPSKit use the threading scheduler stored inside """ dtmap(args...; scheduler=Defaults.scheduler[]) = tmap(args...; scheduler) +dtmap!!(args...; scheduler=Defaults.scheduler[]) = tmap!(args...; scheduler) + # Follows the `map` rrule from ChainRules.jl but specified for the case of one AbstractArray that is being mapped # https://github.com/JuliaDiff/ChainRules.jl/blob/e245d50a1ae56ce46fc8c1f0fe9b925964f1146e/src/rulesets/Base/base.jl#L243 function ChainRulesCore.rrule( @@ -33,6 +35,22 @@ function ChainRulesCore.rrule( return y, dtmap_pullback end +function ChainRulesCore.rrule( + config::RuleConfig{>:HasReverseMode}, + ::typeof(dtmap!!), + f, + C′::AbstractArray, + A::AbstractArray; + kwargs..., +) + C, dtmap_pullback = rrule(config, dtmap, f, A; kwargs...) + function dtmap!!_pullback(dy) + dtmap, df, dA = dtmap_pullback(dy) + return dtmap, df, NoTangent, dA + end + return C, dtmap!!_pullback +end + """ @fwdthreads(ex) diff --git a/src/utility/svd.jl b/src/utility/svd.jl index 25033c9fb..569dd3c5f 100644 --- a/src/utility/svd.jl +++ b/src/utility/svd.jl @@ -290,7 +290,11 @@ function TensorKit._compute_svddata!( V = V[1:howmany, :] else x₀ = alg.start_vector(b) - S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, alg.alg) + svd_alg = alg.alg + if howmany > alg.alg.krylovdim + svd_alg = @set svd_alg.krylovdim = round(Int, howmany * 1.2) + end + S, lvecs, rvecs, info = KrylovKit.svdsolve(b, x₀, howmany, :LR, svd_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" U, S, V = TensorKit.MatrixAlgebra.svd!(b, TensorKit.SDD()) diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl index f7a4d43d5..afa3c67ea 100644 --- a/test/ctmrg/fixed_iterscheme.jl +++ b/test/ctmrg/fixed_iterscheme.jl @@ -1,4 +1,5 @@ using Test +using TestExtras: @constinferred using Accessors using Random using LinearAlgebra @@ -37,7 +38,7 @@ atol = 1e-5 env_conv1, = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg) # do extra iteration to get SVD - env_conv2, info = ctmrg_iteration(n, env_conv1, ctm_alg) + env_conv2, info = @constinferred ctmrg_iteration(n, env_conv1, ctm_alg) env_fix, signs = gauge_fix(env_conv1, env_conv2) @test calc_elementwise_convergence(env_conv1, env_fix) ≈ 0 atol = atol @@ -47,7 +48,7 @@ atol = 1e-5 ctm_alg_fix = @set ctm_alg_fix.projector_alg.trscheme = notrunc() # do iteration with FixedSVD - env_fixedsvd, = ctmrg_iteration(n, env_conv1, ctm_alg_fix) + env_fixedsvd, = @constinferred ctmrg_iteration(n, env_conv1, ctm_alg_fix) env_fixedsvd = fix_global_phases(env_conv1, env_fixedsvd) @test calc_elementwise_convergence(env_conv1, env_fixedsvd) ≈ 0 atol = atol end