diff --git a/src/algorithms/derivatives/hamiltonian_derivatives.jl b/src/algorithms/derivatives/hamiltonian_derivatives.jl index 164f2d58c..51e680d04 100644 --- a/src/algorithms/derivatives/hamiltonian_derivatives.jl +++ b/src/algorithms/derivatives/hamiltonian_derivatives.jl @@ -47,8 +47,8 @@ function AC_hamiltonian( backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) @assert below === above "JordanMPO assumptions break" - GL = leftenv(envs, site, below) - GR = rightenv(envs, site, below) + GL = leftenv(envs, site, below; backend, allocator) + GR = rightenv(envs, site, below; backend, allocator) W = operator[site] H_AC = JordanMPO_AC_Hamiltonian(GL, W, GR; backend, allocator) return prepare ? prepare_operator!!(H_AC) : H_AC @@ -209,8 +209,8 @@ function AC2_hamiltonian( backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) @assert below === above "JordanMPO assumptions break" - GL = leftenv(envs, site, below) - GR = rightenv(envs, site + 1, below) + GL = leftenv(envs, site, below; backend, allocator) + GR = rightenv(envs, site + 1, below; backend, allocator) W1, W2 = operator[site], operator[site + 1] H_AC2 = JordanMPO_AC2_Hamiltonian(GL, W1, W2, GR; backend, allocator) return prepare ? prepare_operator!!(H_AC2) : H_AC2 diff --git a/src/algorithms/derivatives/mpo_derivatives.jl b/src/algorithms/derivatives/mpo_derivatives.jl index 88b837d91..67e47dbee 100644 --- a/src/algorithms/derivatives/mpo_derivatives.jl +++ b/src/algorithms/derivatives/mpo_derivatives.jl @@ -220,13 +220,85 @@ function prepare_operator!!(H::MPO_C_Hamiltonian{<:MPSTensor, <:MPSTensor}) rightenv = H.rightenv isa TensorMap ? H.rightenv : TensorMap(H.rightenv) return prepared_operator_type(typeof(H))(leftenv, rightenv, H.backend, H.allocator) end -function prepare_operator!!(H::MPO_AC_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPSTensor}) - backend, allocator = H.backend, H.allocator + +# Scratch space of `prepare_operator!!` +# ------------------------------------ +# The environments of a prepared derivative are the dense, fused form of an environment-operator +# contraction. Neither that contraction nor its dense copy is kept - only the repartitioned +# result is - so both are taken from the allocator and handed straight back. +# +# `GL_O` and `O_GR` are allocated here rather than by `:=`, which always allocates the result of +# a `@plansor` block on the heap. + +# A `TensorMap` is already dense, so `repartition` is the only tensor that has to be kept. +@inline function _fuse_env(t::TensorMap, N₁::Int, N₂::Int, backend, allocator) + return repartition(fuse_legs(t, N₁, N₂), 2, 2; copy = true, backend, allocator) +end +@inline function _fuse_env(t::AbstractBlockTensorMap, N₁::Int, N₂::Int, backend, allocator) + TT = TensorKit.tensormaptype(spacetype(t), numout(t), numin(t), storagetype(t)) + S = spacetype(t) + Nout, Nin = numout(t), numin(t) + V = ProductSpace{S, Nout}(BlockTensorKit.oplus.(codomain(t).spaces)) ← + ProductSpace{S, Nin}(BlockTensorKit.oplus.(domain(t).spaces)) + + tdense = TensorOperations.tensoralloc(TT, V, Val(true), allocator) + BlockTensorKit.issparse(t) && zerovector!(tdense) + BlockTensorKit._copy_subblocks!(tdense, t) + + env = repartition(fuse_legs(tdense, N₁, N₂), 2, 2; copy = true, backend, allocator) + TensorOperations.tensorfree!(tdense, allocator) + + return env +end + +function _prepare_GL_O(GL, O, backend, allocator) + cp = allocator_checkpoint!(allocator) + + TC = TensorOperations.promote_contract(scalartype(GL), scalartype(O)) + GL_O = TensorOperations.tensoralloc_contract( + TC, + GL, ((1, 3), (2,)), false, + O, ((1,), (2, 3, 4)), false, + ((1, 3, 5), (2, 4)), Val(true), allocator + ) + @plansor backend = backend allocator = allocator begin + GL_O[-1 -2 -3; -4 -5] = GL[-1 1; -4] * O[1 -2; -5 -3] + end + leftenv = _fuse_env(GL_O, 1, 2, backend, allocator) + + TensorOperations.tensorfree!(GL_O, allocator) + allocator_reset!(allocator, cp) + + return leftenv +end + +function _prepare_O_GR(O, GR, backend, allocator) + cp = allocator_checkpoint!(allocator) + + TC = TensorOperations.promote_contract(scalartype(O), scalartype(GR)) + O_GR = TensorOperations.tensoralloc_contract( + TC, + O, ((1, 2, 3), (4,)), false, + GR, ((2,), (1, 3)), false, + ((4, 3), (5, 2, 1)), Val(true), allocator + ) @plansor backend = backend allocator = allocator begin - GL_O[-1 -2 -3; -4 -5] := H.leftenv[-1 1; -4] * H.operators[1][1 -2; -5 -3] + O_GR[-1 -2; -4 -5 -3] = O[-3 -5; -2 1] * GR[-1 1; -4] end - leftenv = GL_O isa TensorMap ? GL_O : TensorMap(GL_O) - leftenv = repartition(fuse_legs(leftenv, 1, 2), 2, 2) + rightenv = _fuse_env(O_GR, 2, 1, backend, allocator) + + TensorOperations.tensorfree!(O_GR, allocator) + allocator_reset!(allocator, cp) + + return rightenv +end + +function prepare_operator!!(H::MPO_AC_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPSTensor}) + backend, allocator = H.backend, H.allocator + + GL, O = H.leftenv, H.operators[1] + leftenv = _prepare_GL_O(GL, O, backend, allocator) + rightenv = H.rightenv isa TensorMap ? H.rightenv : TensorMap(H.rightenv) return prepared_operator_type(typeof(H))(leftenv, rightenv, backend, allocator) @@ -236,15 +308,11 @@ function prepare_operator!!( H::MPO_AC2_Hamiltonian{<:MPSTensor, <:MPOTensor, <:MPOTensor, <:MPSTensor} ) backend, allocator = H.backend, H.allocator - @plansor backend = backend allocator = allocator begin - GL_O[-1 -2 -3; -4 -5] := H.leftenv[-1 1; -4] * H.operators[1][1 -2; -5 -3] - O_GR[-1 -2; -4 -5 -3] := H.operators[2][-3 -5; -2 1] * H.rightenv[-1 1; -4] - end - leftenv = GL_O isa TensorMap ? GL_O : TensorMap(GL_O) - leftenv = repartition(fuse_legs(leftenv, 1, 2), 2, 2) - rightenv = O_GR isa TensorMap ? O_GR : TensorMap(O_GR) - rightenv = repartition(fuse_legs(rightenv, 2, 1), 2, 2) + GL, O₁, O₂, GR = H.leftenv, H.operators[1], H.operators[2], H.rightenv + leftenv = _prepare_GL_O(GL, O₁, backend, allocator) + rightenv = _prepare_O_GR(O₂, GR, backend, allocator) + return prepared_operator_type(typeof(H))(leftenv, rightenv, backend, allocator) end diff --git a/src/algorithms/excitation/exci_transfer_system.jl b/src/algorithms/excitation/exci_transfer_system.jl index c2b4b217d..b0c2aa551 100644 --- a/src/algorithms/excitation/exci_transfer_system.jl +++ b/src/algorithms/excitation/exci_transfer_system.jl @@ -1,6 +1,7 @@ function left_excitation_transfer_system( GBL, H::InfiniteMPOHamiltonian, exci; - mom = exci.momentum, solver = Defaults.linearsolver + mom = exci.momentum, solver = Defaults.linearsolver, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) len = length(H) found = zerovector(GBL) @@ -17,7 +18,7 @@ function left_excitation_transfer_system( # this would require to check the finite state machine, and discard non-connected # terms. H_partial = map(h -> getindex(h, 1:i, 1, 1, 1:i), parent(H)) - T = TransferMatrix(exci.right_gs.AR, H_partial, exci.left_gs.AL) + T = TransferMatrix(exci.right_gs.AR, H_partial, exci.left_gs.AL; backend, allocator) start = scale!(last(found[1:i] * T), cis(-mom * len)) if istrivial(exci) && isidentitylevel(H, i) regularize!(start, ρ_right, ρ_left) @@ -27,13 +28,14 @@ function left_excitation_transfer_system( if !isemptylevel(H, i) if isidentitylevel(H, i) - T = TransferMatrix(exci.right_gs.AR, exci.left_gs.AL) + T = TransferMatrix(exci.right_gs.AR, exci.left_gs.AL; backend, allocator) if istrivial(exci) T = regularize(T, ρ_left, ρ_right) end else T = TransferMatrix( - exci.right_gs.AR, map(h -> h[i, 1, 1, i], parent(H)), exci.left_gs.AL + exci.right_gs.AR, map(h -> h[i, 1, 1, i], parent(H)), exci.left_gs.AL; + backend, allocator ) end @@ -49,8 +51,8 @@ end function right_excitation_transfer_system( GBR, H::InfiniteMPOHamiltonian, exci; - mom = exci.momentum, - solver = Defaults.linearsolver + mom = exci.momentum, solver = Defaults.linearsolver, + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) len = length(H) found = zerovector(GBR) @@ -67,7 +69,7 @@ function right_excitation_transfer_system( # this would require to check the finite state machine, and discard non-connected # terms. H_partial = map(h -> h[i:end, 1, 1, i:end], parent(H)) - T = TransferMatrix(exci.left_gs.AL, H_partial, exci.right_gs.AR) + T = TransferMatrix(exci.left_gs.AL, H_partial, exci.right_gs.AR; backend, allocator) start = scale!(first(T * found[i:odim]), cis(mom * len)) if istrivial(exci) && isidentitylevel(H, i) regularize!(start, ρ_left, ρ_right) @@ -77,13 +79,14 @@ function right_excitation_transfer_system( if !isemptylevel(H, i) if isidentitylevel(H, i) - tm = TransferMatrix(exci.left_gs.AL, exci.right_gs.AR) + tm = TransferMatrix(exci.left_gs.AL, exci.right_gs.AR; backend, allocator) if istrivial(exci) tm = regularize(tm, ρ_left, ρ_right) end else tm = TransferMatrix( - exci.left_gs.AL, map(h -> h[i, 1, 1, i], parent(H)), exci.right_gs.AR + exci.left_gs.AL, map(h -> h[i, 1, 1, i], parent(H)), exci.right_gs.AR; + backend, allocator ) end diff --git a/src/algorithms/excitation/quasiparticleexcitation.jl b/src/algorithms/excitation/quasiparticleexcitation.jl index 3ba679af0..c54ceae7f 100644 --- a/src/algorithms/excitation/quasiparticleexcitation.jl +++ b/src/algorithms/excitation/quasiparticleexcitation.jl @@ -28,19 +28,25 @@ Used as the `algorithm` argument of [`excitations`](@ref). * [Haegeman et al. Phys. Rev. Let. 111 (2013)](@cite haegeman2013) """ -struct QuasiparticleAnsatz{A, E} <: Algorithm +struct QuasiparticleAnsatz{A, E, B} <: Algorithm "algorithm used for the eigenvalue solvers" alg::A "algorithm used for the quasiparticle environments" alg_environments::E + + "backend for tensor contractions and index manipulations" + backend::B end +QuasiparticleAnsatz(alg, alg_environments) = + QuasiparticleAnsatz(alg, alg_environments, Defaults.backend()) function QuasiparticleAnsatz(; alg_environments = Defaults.alg_environments(; dynamic_tols = false), + backend = Defaults.backend(), kwargs... ) alg = Defaults.alg_eigsolve(; dynamic_tols = false, kwargs...) - return QuasiparticleAnsatz(alg, alg_environments) + return QuasiparticleAnsatz(alg, alg_environments, backend) end ################################################################################ @@ -51,7 +57,7 @@ function excitations(H, alg::QuasiparticleAnsatz, ϕ₀::InfiniteQP, lenvs, renv E = effective_excitation_renormalization_energy(H, ϕ₀, lenvs, renvs) H_eff = EffectiveExcitationHamiltonian(H, lenvs, renvs, E) Es, ϕs, convhist = eigsolve(ϕ₀, num, :SR, alg.alg) do ϕ - return H_eff(ϕ, alg.alg_environments) + return H_eff(ϕ, alg.alg_environments; alg.backend) end convhist.converged < num && @warn "excitation failed to converge: normres = $(convhist.normres)" @@ -157,7 +163,7 @@ function excitations( E = effective_excitation_renormalization_energy(H, ϕ₀, lenvs, renvs) H_eff = EffectiveExcitationHamiltonian(H, lenvs, renvs, E) Es, ϕs, convhist = eigsolve(ϕ₀, num, :SR, alg.alg) do ϕ - return H_eff(ϕ, alg.alg_environments) + return H_eff(ϕ, alg.alg_environments; alg.backend) end convhist.converged < num && @@ -300,24 +306,37 @@ end # to allow Multiline checks Base.length(H::EffectiveExcitationHamiltonian) = length(H.operator) -function (H::EffectiveExcitationHamiltonian)(ϕ::QP, alg_environments = DefaultAlgorithm()) - qp_envs = environments(ϕ, H.operator, ϕ, alg_environments; lenvs = H.lenvs, renvs = H.renvs) - return effective_excitation_hamiltonian(H.operator, ϕ, qp_envs, H.energy) +function (H::EffectiveExcitationHamiltonian)( + ϕ::QP, alg_environments = DefaultAlgorithm(); + backend::AbstractBackend = DefaultBackend() + ) + qp_envs = environments( + ϕ, H.operator, ϕ, alg_environments; lenvs = H.lenvs, renvs = H.renvs, backend + ) + return effective_excitation_hamiltonian(H.operator, ϕ, qp_envs, H.energy; backend) end function (H::Multiline{<:EffectiveExcitationHamiltonian})( - ϕ::MultilineQP, alg_environments = DefaultAlgorithm() + ϕ::MultilineQP, alg_environments = DefaultAlgorithm(); kwargs... ) - return Multiline(map((x, y) -> x(y, alg_environments), parent(H), parent(ϕ))) + return Multiline(map((x, y) -> x(y, alg_environments; kwargs...), parent(H), parent(ϕ))) end function effective_excitation_hamiltonian(H, ϕ, envs = environments(ϕ, H)) E₀ = effective_excitation_renormalization_energy(H, ϕ, envs.leftenvs, envs.rightenvs) return effective_excitation_hamiltonian(H, ϕ, envs, E₀) end -function effective_excitation_hamiltonian(H, ϕ, qp_envs, E) +function effective_excitation_hamiltonian( + H, ϕ, qp_envs, E, scheduler = Defaults.scheduler[]; + backend::AbstractBackend = DefaultBackend() + ) ϕ′ = similar(ϕ) - tforeach(1:length(ϕ); scheduler = Defaults.scheduler[]) do loc - ϕ′[loc] = _effective_excitation_local_apply(loc, ϕ, H, E[loc], qp_envs) + # This is the site that fans out, so it is the one that reads the scheduler and derives + # the allocator from it + allocator = default_allocator(ϕ.left_gs, scheduler) + tforeach(1:length(ϕ); scheduler) do loc + ϕ′[loc] = _effective_excitation_local_apply( + loc, ϕ, H, E[loc], qp_envs; backend, allocator + ) return nothing end return ϕ′ @@ -338,10 +357,13 @@ function effective_excitation_hamiltonian(H::MultilineMPO, ϕ::MultilineQP, envs ) end -function _effective_excitation_local_apply(site, ϕ, H::MPOHamiltonian, E::Number, envs) +function _effective_excitation_local_apply( + site, ϕ, H::MPOHamiltonian, E::Number, envs; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) B = ϕ[site] - GL = leftenv(envs.leftenvs, site, ϕ.left_gs) - GR = rightenv(envs.rightenvs, site, ϕ.right_gs) + GL = leftenv(envs.leftenvs, site, ϕ.left_gs; backend, allocator) + GR = rightenv(envs.rightenvs, site, ϕ.right_gs; backend, allocator) # renormalize first -> allocates destination B′ = scale(B, -E) @@ -366,13 +388,16 @@ function _effective_excitation_local_apply(site, ϕ, H::MPOHamiltonian, E::Numbe return B′ end -function _effective_excitation_local_apply(site, ϕ, H::MPO, E::Number, envs) +function _effective_excitation_local_apply( + site, ϕ, H::MPO, E::Number, envs; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) left_gs = ϕ.left_gs right_gs = ϕ.right_gs B = ϕ[site] - GL = leftenv(envs.leftenvs, site, ϕ.left_gs) - GR = rightenv(envs.rightenvs, site, ϕ.right_gs) + GL = leftenv(envs.leftenvs, site, ϕ.left_gs; backend, allocator) + GR = rightenv(envs.rightenvs, site, ϕ.right_gs; backend, allocator) @plansor T[-1 -2; -3 -4] := GL[-1 5; 4] * B[4 2; -3 1] * H[site][5 -2; 2 3] * GR[1 3; -4] diff --git a/src/algorithms/groundstate/idmrg.jl b/src/algorithms/groundstate/idmrg.jl index 1b9c3c494..519933630 100644 --- a/src/algorithms/groundstate/idmrg.jl +++ b/src/algorithms/groundstate/idmrg.jl @@ -133,7 +133,7 @@ function _find_groundstate_idmrg(mps, operator, alg::alg_type, envs) where {alg_ alg_gauge = adapt_solver(alg.alg_gauge; iter = it.state.iter, g_global = it.state.ϵ) ψ′ = InfiniteMPS(it.state.mps.AR; alg_gauge.tol, alg_gauge.maxiter) - envs = recalculate!(it.state.envs, ψ′, it.state.operator, ψ′) + envs = recalculate!(it.state.envs, ψ′, it.state.operator, ψ′; alg.backend) return ψ′, envs, it.state.ϵ end end @@ -210,7 +210,7 @@ function _localupdate_sweep_idmrg!( ψ.AL[pos], ψ.C[pos] = left_orth!(ψ.AC[pos]; positive = true) end end - @timeit timeroutput "transfer_env" transfer_leftenv!(envs, ψ, H, ψ, pos + 1) + @timeit timeroutput "transfer_env" transfer_leftenv!(envs, ψ, H, ψ, pos + 1; backend, allocator) end # right to left sweep @@ -223,7 +223,7 @@ function _localupdate_sweep_idmrg!( ψ.C[pos - 1], temp = right_orth!(_transpose_tail(ψ.AC[pos]; copy = (pos == 1)); positive = true) ψ.AR[pos] = _transpose_front(temp) end - @timeit timeroutput "transfer_env" transfer_rightenv!(envs, ψ, H, ψ, pos - 1) + @timeit timeroutput "transfer_env" transfer_rightenv!(envs, ψ, H, ψ, pos - 1; backend, allocator) end return ψ, envs, C_old, E end @@ -252,8 +252,8 @@ function _localupdate_sweep_idmrg2!( ψ.AC[pos + 1] = _transpose_front(c * ar) end @timeit timeroutput "transfer_env" begin - transfer_leftenv!(envs, ψ, H, ψ, pos + 1) - transfer_rightenv!(envs, ψ, H, ψ, pos) + transfer_leftenv!(envs, ψ, H, ψ, pos + 1; backend, allocator) + transfer_rightenv!(envs, ψ, H, ψ, pos; backend, allocator) end end @@ -282,8 +282,8 @@ function _localupdate_sweep_idmrg2!( # update environments @timeit timeroutput "transfer_env" begin - transfer_leftenv!(envs, ψ, H, ψ, 1) - transfer_rightenv!(envs, ψ, H, ψ, 0) + transfer_leftenv!(envs, ψ, H, ψ, 1; backend, allocator) + transfer_rightenv!(envs, ψ, H, ψ, 0; backend, allocator) end # sweep from right to left @@ -304,8 +304,8 @@ function _localupdate_sweep_idmrg2!( ψ.AC[pos + 1] = _transpose_front(c * ar) end @timeit timeroutput "transfer_env" begin - transfer_leftenv!(envs, ψ, H, ψ, pos + 1) - transfer_rightenv!(envs, ψ, H, ψ, pos) + transfer_leftenv!(envs, ψ, H, ψ, pos + 1; backend, allocator) + transfer_rightenv!(envs, ψ, H, ψ, pos; backend, allocator) end end @@ -330,8 +330,8 @@ function _localupdate_sweep_idmrg2!( end @timeit timeroutput "transfer_env" begin - transfer_leftenv!(envs, ψ, H, ψ, 1) - transfer_rightenv!(envs, ψ, H, ψ, 0) + transfer_leftenv!(envs, ψ, H, ψ, 1; backend, allocator) + transfer_rightenv!(envs, ψ, H, ψ, 0; backend, allocator) end return ψ, envs, C_old, E end diff --git a/src/algorithms/groundstate/vumps.jl b/src/algorithms/groundstate/vumps.jl index 7be8c096a..aa51f44b5 100644 --- a/src/algorithms/groundstate/vumps.jl +++ b/src/algorithms/groundstate/vumps.jl @@ -70,7 +70,10 @@ function dominant_eigsolve( mps = copy(mps) ϵ = calc_galerkin(mps, operator, mps, envs; alg.backend) alg_environments = adapt_solver(alg.alg_environments; iter, g_global = ϵ) - recalculate!(envs, mps, operator, mps, alg_environments; timeroutput) + recalculate!( + envs, mps, operator, mps, alg_environments; + timeroutput, alg.backend + ) state = VUMPSState(mps, operator, envs, iter, ϵ, which, timeroutput) it = IterativeSolver(alg, state) @@ -172,9 +175,11 @@ end function gauge_step!(it::IterativeSolver{<:VUMPS}, state, ACs::AbstractVector) alg_gauge = adapt_solver(it.alg_gauge; iter = state.iter, g_global = state.ϵ) + # the gauge sweep is serial, so safe to use non-threadsafe allocator + allocator = default_allocator(state.mps, SerialScheduler()) mps = gaugefix!( state.mps, ACs, state.mps.C[end]; - order = :R, timeroutput = state.timeroutput, alg_gauge..., + order = :R, timeroutput = state.timeroutput, it.backend, allocator, alg_gauge..., ) mul!.(mps.AC, mps.AL, mps.C) return mps @@ -186,5 +191,8 @@ end function envs_step!(it::IterativeSolver{<:VUMPS}, state, mps) alg_environments = adapt_solver(it.alg_environments; iter = state.iter, g_global = state.ϵ) - return recalculate!(state.envs, mps, state.operator, mps, alg_environments; state.timeroutput) + return recalculate!( + state.envs, mps, state.operator, mps, alg_environments; + state.timeroutput, it.backend + ) end diff --git a/src/environments/finite_envs.jl b/src/environments/finite_envs.jl index 6605ccc74..46afc6828 100644 --- a/src/environments/finite_envs.jl +++ b/src/environments/finite_envs.jl @@ -107,7 +107,10 @@ it directly. See also [`leftenv`](@ref) and [`environments`](@ref). """ -function rightenv(ca::FiniteEnvironments, ind, state) +function rightenv( + ca::FiniteEnvironments, ind, state; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) a = findfirst(i -> !(state.AR[i] === ca.rdependencies[i]), length(state):-1:(ind + 1)) a = isnothing(a) ? nothing : length(state) - a + 1 @@ -115,8 +118,9 @@ function rightenv(ca::FiniteEnvironments, ind, state) #we need to recalculate for j in a:-1:(ind + 1) above = isnothing(ca.above) ? state.AR[j] : ca.above.AR[j] - ca.GRs[j] = TransferMatrix(above, ca.operator[j], state.AR[j]) * - ca.GRs[j + 1] + ca.GRs[j] = TransferMatrix( + above, ca.operator[j], state.AR[j]; backend, allocator + ) * ca.GRs[j + 1] ca.rdependencies[j] = state.AR[j] end end @@ -134,15 +138,19 @@ it directly. See also [`rightenv`](@ref) and [`environments`](@ref). """ -function leftenv(ca::FiniteEnvironments, ind, state) +function leftenv( + ca::FiniteEnvironments, ind, state; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) a = findfirst(i -> !(state.AL[i] === ca.ldependencies[i]), 1:(ind - 1)) if !isnothing(a) #we need to recalculate for j in a:(ind - 1) above = isnothing(ca.above) ? state.AL[j] : ca.above.AL[j] - ca.GLs[j + 1] = ca.GLs[j] * - TransferMatrix(above, ca.operator[j], state.AL[j]) + ca.GLs[j + 1] = ca.GLs[j] * TransferMatrix( + above, ca.operator[j], state.AL[j]; backend, allocator + ) ca.ldependencies[j] = state.AL[j] end end diff --git a/src/environments/infinite_envs.jl b/src/environments/infinite_envs.jl index 2a57a10ea..5a0852c76 100644 --- a/src/environments/infinite_envs.jl +++ b/src/environments/infinite_envs.jl @@ -15,8 +15,11 @@ end Base.length(envs::InfiniteEnvironments) = length(envs.GLs) -leftenv(envs::InfiniteEnvironments, site::Int, state) = envs.GLs[site] -rightenv(envs::InfiniteEnvironments, site::Int, state) = envs.GRs[site] +# `backend`/`allocator` are accepted and ignored here: these environments are already +# materialised, so there is nothing to contract. They exist so that callers which do have an +# allocator (the derivatives) can pass it uniformly, without dispatching on environment type. +leftenv(envs::InfiniteEnvironments, site::Int, state; kwargs...) = envs.GLs[site] +rightenv(envs::InfiniteEnvironments, site::Int, state; kwargs...) = envs.GRs[site] function environments( below::InfiniteMPS, operator::Union{InfiniteMPO, InfiniteMPOHamiltonian}, above; @@ -58,10 +61,14 @@ end function recalculate!( envs::InfiniteEnvironments, below, operator::Union{InfiniteMPO, InfiniteMPOHamiltonian}, above = below; - timeroutput = NoTimerOutput(), kwargs... + timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), scheduler = Defaults.scheduler[], + kwargs... ) + # `backend` and `scheduler` are named explicitly so that they are not swept into + # `environment_alg`'s keyword arguments, which describe the linear solver. alg = environment_alg(below, operator, above; kwargs...) - return recalculate!(envs, below, operator, above, alg; timeroutput) + return recalculate!(envs, below, operator, above, alg; timeroutput, backend, scheduler) end function recalculate!( envs::InfiniteEnvironments, below::InfiniteMPS, @@ -85,6 +92,7 @@ function recalculate!( operator::Union{InfiniteMPO, InfiniteMPOHamiltonian}, above::InfiniteMPS, alg; timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), scheduler = Defaults.scheduler[], ) if !issamespace(envs, below, operator, above) # TODO: in-place initialization? @@ -94,17 +102,20 @@ function recalculate!( end tree_point = timer_treepoint(timeroutput) - @sync begin - @spawn begin - sub_timeroutput = subtimer(timeroutput) - @timeit sub_timeroutput "left_envs" compute_leftenvs!(envs, below, operator, above, alg) - merge_subtimer!(timeroutput, sub_timeroutput; tree_point) - end - @spawn begin - sub_timeroutput = subtimer(timeroutput) - @timeit sub_timeroutput "right_envs" compute_rightenvs!(envs, below, operator, above, alg) - merge_subtimer!(timeroutput, sub_timeroutput; tree_point) + # concurrency depends on scheduler, but each half takes its own dedicated scratch allocator + tforeach(1:2; scheduler) do half + allocator = default_allocator(below, SerialScheduler()) + sub_timeroutput = subtimer(timeroutput) + if isone(half) + @timeit sub_timeroutput "left_envs" compute_leftenvs!( + envs, below, operator, above, alg; backend, allocator + ) + else + @timeit sub_timeroutput "right_envs" compute_rightenvs!( + envs, below, operator, above, alg; backend, allocator + ) end + merge_subtimer!(timeroutput, sub_timeroutput; tree_point) end normalize!(envs, below, operator, above) @@ -125,15 +136,17 @@ end function compute_leftenvs!( envs::InfiniteEnvironments, below::InfiniteMPS, operator::InfiniteMPO, above::InfiniteMPS, - alg + alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) # compute eigenvector - T = TransferMatrix(above.AL, operator, below.AL) + T = TransferMatrix(above.AL, operator, below.AL; backend, allocator) λ, envs.GLs[1] = fixedpoint(flip(T), envs.GLs[1], :LM, alg) # push through unitcell for i in 2:length(operator) - envs.GLs[i] = envs.GLs[i - 1] * - TransferMatrix(above.AL[i - 1], operator[i - 1], below.AL[i - 1]) + envs.GLs[i] = envs.GLs[i - 1] * TransferMatrix( + above.AL[i - 1], operator[i - 1], below.AL[i - 1]; backend, allocator + ) end return λ, envs end @@ -141,15 +154,16 @@ end function compute_rightenvs!( envs::InfiniteEnvironments, below::InfiniteMPS, operator::InfiniteMPO, above::InfiniteMPS, - alg + alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) # compute eigenvector - T = TransferMatrix(above.AR, operator, below.AR) + T = TransferMatrix(above.AR, operator, below.AR; backend, allocator) λ, envs.GRs[end] = fixedpoint(T, envs.GRs[end], :LM, alg) # push through unitcell for i in reverse(1:(length(operator) - 1)) envs.GRs[i] = TransferMatrix( - above.AR[i + 1], operator[i + 1], below.AR[i + 1] + above.AR[i + 1], operator[i + 1], below.AR[i + 1]; backend, allocator ) * envs.GRs[i + 1] end return λ, envs @@ -207,7 +221,8 @@ end function compute_leftenvs!( envs::InfiniteEnvironments, below::InfiniteMPS, operator::InfiniteMPOHamiltonian, above::InfiniteMPS, - alg + alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) L = check_length(below, above, operator) GLs = envs.GLs @@ -222,38 +237,38 @@ function compute_leftenvs!( # TODO: check if this is necessary # leftutil = similar(above.AL[1], space(GL[1], 2)[1]) # fill_data!(leftutil, one) - # @plansor GL[1][1][-1 -2; -3] = ρ_left[-1; -3] * leftutil[-2] + # @plansor backend = backend allocator = allocator GL[1][1][-1 -2; -3] = ρ_left[-1; -3] * leftutil[-2] - (L > 1) && left_cyclethrough!(1, GLs, below, operator, above) + (L > 1) && left_cyclethrough!(1, GLs, below, operator, above; backend, allocator) for i in 2:vsize prev = copy(GLs[1][i]) zerovector!(GLs[1][i]) - left_cyclethrough!(i, GLs, below, operator, above) + left_cyclethrough!(i, GLs, below, operator, above; backend, allocator) if isidentitylevel(operator, i) # identity matrices; do the hacky renormalization - T = regularize(TransferMatrix(above.AL, below.AL), ρ_left, ρ_right) + T = regularize(TransferMatrix(above.AL, below.AL; backend, allocator), ρ_left, ρ_right) GLs[1][i], convhist = linsolve(flip(T), GLs[1][i], prev, alg, 1, -1) convhist.converged == 0 && @warn "GL$i failed to converge: normres = $(convhist.normres)" - (L > 1) && left_cyclethrough!(i, GLs, below, operator, above) + (L > 1) && left_cyclethrough!(i, GLs, below, operator, above; backend, allocator) # go through the unitcell, again subtracting fixpoints for site in 1:L - @plansor GLs[site][i][-1 -2; -3] -= GLs[site][i][1 -2; 2] * + @plansor backend = backend allocator = allocator GLs[site][i][-1 -2; -3] -= GLs[site][i][1 -2; 2] * r_LL(above, site - 1)[2; 1] * l_LL(above, site)[-1; -3] end else if !isemptylevel(operator, i) diag = map(h -> h[i, 1, 1, i], operator[:]) - T = TransferMatrix(above.AL, diag, below.AL) + T = TransferMatrix(above.AL, diag, below.AL; backend, allocator) GLs[1][i], convhist = linsolve(flip(T), GLs[1][i], prev, alg, 1, -1) convhist.converged == 0 && @warn "GL$i failed to converge: normres = $(convhist.normres)" end - (L > 1) && left_cyclethrough!(i, GLs, below, operator, above) + (L > 1) && left_cyclethrough!(i, GLs, below, operator, above; backend, allocator) end end @@ -262,13 +277,15 @@ end function left_cyclethrough!( index::Int, GL, - below::InfiniteMPS, H::InfiniteMPOHamiltonian, above::InfiniteMPS = below + below::InfiniteMPS, H::InfiniteMPOHamiltonian, above::InfiniteMPS = below; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) # TODO: efficient transfer matrix slicing for large unitcells leftinds = 1:index for site in eachindex(GL) GL[site + 1][index] = GL[site][leftinds] * TransferMatrix( - above.AL[site], H[site][leftinds, 1, 1, index], below.AL[site] + above.AL[site], H[site][leftinds, 1, 1, index], below.AL[site]; + backend, allocator ) end return GL @@ -277,7 +294,8 @@ end function compute_rightenvs!( envs::InfiniteEnvironments, below::InfiniteMPS, operator::InfiniteMPOHamiltonian, above::InfiniteMPS, - alg + alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) L = check_length(above, operator, below) GRs = envs.GRs @@ -292,39 +310,39 @@ function compute_rightenvs!( # TODO: check if this is necessary # rightutil = similar(state.AL[1], space(GR[end], 2)[end]) # fill_data!(rightutil, one) - # @plansor GR[end][end][-1 -2; -3] = r_RR(state)[-1; -3] * rightutil[-2] + # @plansor backend = backend allocator = allocator GR[end][end][-1 -2; -3] = r_RR(state)[-1; -3] * rightutil[-2] - (L > 1) && right_cyclethrough!(vsize, GRs, below, operator, above) # populate other sites + (L > 1) && right_cyclethrough!(vsize, GRs, below, operator, above; backend, allocator) # populate other sites for i in (vsize - 1):-1:1 prev = copy(GRs[end][i]) zerovector!(GRs[end][i]) - right_cyclethrough!(i, GRs, below, operator, above) + right_cyclethrough!(i, GRs, below, operator, above; backend, allocator) if isidentitylevel(operator, i) # identity matrices; do the hacky renormalization # subtract fixpoints - T = regularize(TransferMatrix(above.AR, below.AR), ρ_left, ρ_right) + T = regularize(TransferMatrix(above.AR, below.AR; backend, allocator), ρ_left, ρ_right) GRs[end][i], convhist = linsolve(T, GRs[end][i], prev, alg, 1, -1) convhist.converged == 0 && @warn "GR$i failed to converge: normres = $(convhist.normres)" - L > 1 && right_cyclethrough!(i, GRs, below, operator, above) + L > 1 && right_cyclethrough!(i, GRs, below, operator, above; backend, allocator) # go through the unitcell, again subtracting fixpoints for site in 1:L - @plansor GRs[site][i][-1 -2; -3] -= GRs[site][i][1 -2; 2] * + @plansor backend = backend allocator = allocator GRs[site][i][-1 -2; -3] -= GRs[site][i][1 -2; 2] * l_RR(above, site + 1)[2; 1] * r_RR(above, site)[-1; -3] end else if !isemptylevel(operator, i) diag = map(b -> b[i, 1, 1, i], operator[:]) - T = TransferMatrix(above.AR, diag, below.AR) + T = TransferMatrix(above.AR, diag, below.AR; backend, allocator) GRs[end][i], convhist = linsolve(T, GRs[end][i], prev, alg, 1, -1) convhist.converged == 0 && @warn "GR$i failed to converge: normres = $(convhist.normres)" end - (L > 1) && right_cyclethrough!(i, GRs, below, operator, above) + (L > 1) && right_cyclethrough!(i, GRs, below, operator, above; backend, allocator) end end @@ -333,13 +351,15 @@ end function right_cyclethrough!( index::Int, GR, - below::InfiniteMPS, operator::InfiniteMPOHamiltonian, above::InfiniteMPS = below + below::InfiniteMPS, operator::InfiniteMPOHamiltonian, above::InfiniteMPS = below; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) # TODO: efficient transfer matrix slicing for large unitcells for site in reverse(eachindex(GR)) rightinds = index:length(GR[site]) GR[site - 1][index] = TransferMatrix( - above.AR[site], operator[site][index, 1, 1, rightinds], below.AR[site] + above.AR[site], operator[site][index, 1, 1, rightinds], below.AR[site]; + backend, allocator ) * GR[site][rightinds] end return GR @@ -356,14 +376,24 @@ end # Transfer operations # ------------------- -function transfer_leftenv!(envs::InfiniteEnvironments, below, operator, above, site::Int) - T = TransferMatrix(above.AL[site - 1], operator[site - 1], below.AL[site - 1]) +function transfer_leftenv!( + envs::InfiniteEnvironments, below, operator, above, site::Int; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + T = TransferMatrix( + above.AL[site - 1], operator[site - 1], below.AL[site - 1]; backend, allocator + ) envs.GLs[site] = envs.GLs[site - 1] * T return envs end -function transfer_rightenv!(envs::InfiniteEnvironments, below, operator, above, site::Int) - T = TransferMatrix(above.AR[site + 1], operator[site + 1], below.AR[site + 1]) +function transfer_rightenv!( + envs::InfiniteEnvironments, below, operator, above, site::Int; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + T = TransferMatrix( + above.AR[site + 1], operator[site + 1], below.AR[site + 1]; backend, allocator + ) envs.GRs[site] = T * envs.GRs[site + 1] return envs end diff --git a/src/environments/multiline_envs.jl b/src/environments/multiline_envs.jl index a8e1d73f4..b39636905 100644 --- a/src/environments/multiline_envs.jl +++ b/src/environments/multiline_envs.jl @@ -71,30 +71,40 @@ function TensorKit.normalize!(envs::MultilineEnvironments, below, (operator, abo return envs end -function leftenv(envs::MultilineEnvironments, col::Int, state) - return leftenv.(parent(envs), col, parent(state)) +function leftenv(envs::MultilineEnvironments, col::Int, state; kwargs...) + return leftenv.(parent(envs), col, parent(state); kwargs...) end -function rightenv(envs::MultilineEnvironments, col::Int, state) - return rightenv.(parent(envs), col, parent(state)) +function rightenv(envs::MultilineEnvironments, col::Int, state; kwargs...) + return rightenv.(parent(envs), col, parent(state); kwargs...) end -function transfer_leftenv!(envs::MultilineEnvironments, below, operator, above, site::Int) +function transfer_leftenv!( + envs::MultilineEnvironments, below, operator, above, site::Int; kwargs... + ) for row in 1:size(above, 1) - transfer_leftenv!(envs[row], below[row + 1], operator[row], above[row], site) + transfer_leftenv!( + envs[row], below[row + 1], operator[row], above[row], site; kwargs... + ) end return envs end -function transfer_leftenv!(envs::MultilineEnvironments, below, (O, above)::Tuple, site::Int) - return transfer_leftenv!(envs, below, O, above, site) +function transfer_leftenv!( + envs::MultilineEnvironments, below, (O, above)::Tuple, site::Int; kwargs... + ) + return transfer_leftenv!(envs, below, O, above, site; kwargs...) end -function transfer_rightenv!(envs::MultilineEnvironments, below, operator, above, site::Int) +function transfer_rightenv!( + envs::MultilineEnvironments, below, operator, above, site::Int; kwargs... + ) for row in 1:size(above, 1) - transfer_rightenv!(envs[row], below[row + 1], operator[row], above[row], site) + transfer_rightenv!( + envs[row], below[row + 1], operator[row], above[row], site; kwargs... + ) end return envs end function transfer_rightenv!( - envs::MultilineEnvironments, below, (O, above)::Tuple, site::Int + envs::MultilineEnvironments, below, (O, above)::Tuple, site::Int; kwargs... ) - return transfer_rightenv!(envs, below, O, above, site) + return transfer_rightenv!(envs, below, O, above, site; kwargs...) end diff --git a/src/environments/multiple_envs.jl b/src/environments/multiple_envs.jl index 9935d52d5..2b12d9cd6 100644 --- a/src/environments/multiple_envs.jl +++ b/src/environments/multiple_envs.jl @@ -55,22 +55,24 @@ function Base.getproperty(envs::MultipleEnvironments, prop::Symbol) end end +# `backend`/`allocator` are forwarded to each summand, the same way `leftenv`/`rightenv` +# forward them for `MultilineEnvironments`: these only dispatch, they contract nothing. function transfer_rightenv!( envs::MultipleEnvironments{<:InfiniteEnvironments}, - below, operator, above, pos::Int + below, operator, above, pos::Int; kwargs... ) for (subH, subenv) in zip(operator, envs.envs) - transfer_rightenv!(subenv, below, subH, above, pos) + transfer_rightenv!(subenv, below, subH, above, pos; kwargs...) end return envs end function transfer_leftenv!( envs::MultipleEnvironments{<:InfiniteEnvironments}, - below, operator, above, pos::Int + below, operator, above, pos::Int; kwargs... ) for (subH, subenv) in zip(operator, envs.envs) - transfer_leftenv!(subenv, below, subH, above, pos) + transfer_leftenv!(subenv, below, subH, above, pos; kwargs...) end return envs end diff --git a/src/environments/qp_envs.jl b/src/environments/qp_envs.jl index 9429da3a6..a874d3c9a 100644 --- a/src/environments/qp_envs.jl +++ b/src/environments/qp_envs.jl @@ -18,125 +18,161 @@ end Base.length(envs::InfiniteQPEnvironments) = length(envs.leftenvs) -function leftenv(envs::InfiniteQPEnvironments, site::Int, state) - return leftenv(envs.leftenvs, site, state) +function leftenv(envs::InfiniteQPEnvironments, site::Int, state; kwargs...) + return leftenv(envs.leftenvs, site, state; kwargs...) end -function rightenv(envs::InfiniteQPEnvironments, site::Int, state) - return rightenv(envs.rightenvs, site, state) +function rightenv(envs::InfiniteQPEnvironments, site::Int, state; kwargs...) + return rightenv(envs.rightenvs, site, state; kwargs...) end function environments( exci::Union{InfiniteQP, MultilineQP}, operator::Union{InfiniteMPO, InfiniteMPOHamiltonian, MultilineMPO}, above = exci; lenvs = environments(exci.left_gs, operator, exci.left_gs), renvs = istopological(exci) ? environments(exci.right_gs, operator, exci.right_gs) : lenvs, + backend::AbstractBackend = DefaultBackend(), scheduler = Defaults.scheduler[], kwargs... ) + # `backend` and `scheduler` are named explicitly so that they are not swept into + # `environment_alg`'s keyword arguments, which describe the linear solver. Same + # convention as `recalculate!`. alg = environment_alg(exci, operator, above; kwargs...) - return environments(exci, operator, above, alg; lenvs, renvs) + return environments(exci, operator, above, alg; lenvs, renvs, backend, scheduler) end function environments( - qp::MultilineQP, operator::MultilineMPO, above, alg; lenvs, renvs = lenvs + qp::MultilineQP, operator::MultilineMPO, above, alg; lenvs, renvs = lenvs, + backend::AbstractBackend = DefaultBackend(), scheduler = Defaults.scheduler[] ) (rows = size(qp, 1)) == size(operator, 1) || throw(ArgumentError("Incompatible sizes")) envs = map(1:rows) do row return environments( - qp[row], operator[row], qp[row], alg; lenvs = lenvs[row], renvs = renvs[row] + qp[row], operator[row], qp[row], alg; lenvs = lenvs[row], renvs = renvs[row], + backend, scheduler ) end return Multiline(PeriodicVector(envs)) end function environments( - exci::InfiniteQP, H::InfiniteMPOHamiltonian, above, alg; lenvs, renvs = lenvs + exci::InfiniteQP, H::InfiniteMPOHamiltonian, above, alg; lenvs, renvs = lenvs, + backend::AbstractBackend = DefaultBackend(), scheduler = Defaults.scheduler[] ) - ids = findall(Base.Fix1(isidentitylevel, H), 2:(size(H[1], 1) - 1)) .+ 1 solver = resolve_environment_solver(alg, exci, H, exci) - AL = exci.left_gs.AL - AR = exci.right_gs.AR - lBs = PeriodicVector([allocate_GBL(exci, H, exci, i) for i in 1:length(exci)]) rBs = PeriodicVector([allocate_GBR(exci, H, exci, i) for i in 1:length(exci)]) - - zerovector!(lBs[1]) - for pos in 1:length(exci) - lBs[pos + 1] = lBs[pos] * TransferMatrix(AR[pos], H[pos], AL[pos]) / - cis(exci.momentum) - lBs[pos + 1] += leftenv(lenvs, pos, exci.left_gs) * - TransferMatrix(exci[pos], H[pos], AL[pos]) / cis(exci.momentum) - - if istrivial(exci) && !isempty(ids) # regularization of trivial excitations - ρ_left = l_RL(exci.left_gs, pos + 1) - ρ_right = r_RL(exci.left_gs, pos) - for i in ids - regularize!(lBs[pos + 1][i], ρ_right, ρ_left) - end + envs = InfiniteQPEnvironments(lBs, rBs, lenvs, renvs) + + # concurrency depends on scheduler, but each half takes its own dedicated scratch allocator + tforeach(1:2; scheduler) do half + allocator = default_allocator(exci.left_gs, SerialScheduler()) + if isone(half) + compute_leftenvs!(envs, exci, H, above, solver; backend, allocator) + else + compute_rightenvs!(envs, exci, H, above, solver; backend, allocator) end + return nothing end - zerovector!(rBs[end]) - for pos in length(exci):-1:1 - rBs[pos - 1] = TransferMatrix(AL[pos], H[pos], AR[pos]) * - rBs[pos] * cis(exci.momentum) - rBs[pos - 1] += TransferMatrix(exci[pos], H[pos], AR[pos]) * - rightenv(renvs, pos, exci.right_gs) * cis(exci.momentum) - - if istrivial(exci) && !isempty(ids) - ρ_left = l_LR(exci.left_gs, pos) - ρ_right = r_LR(exci.left_gs, pos - 1) - for i in ids - regularize!(rBs[pos - 1][i], ρ_left, ρ_right) - end - end + return envs +end + +# regularization of trivial excitations +function regularize_GBL!(GBL, exci::InfiniteQP, ids, pos::Int) + (istrivial(exci) && !isempty(ids)) || return GBL + ρ_left = l_RL(exci.left_gs, pos + 1) + ρ_right = r_RL(exci.left_gs, pos) + for i in ids + regularize!(GBL[i], ρ_right, ρ_left) + end + return GBL +end +function regularize_GBR!(GBR, exci::InfiniteQP, ids, pos::Int) + (istrivial(exci) && !isempty(ids)) || return GBR + ρ_left = l_LR(exci.left_gs, pos) + ρ_right = r_LR(exci.left_gs, pos - 1) + for i in ids + regularize!(GBR[i], ρ_left, ρ_right) end + return GBR +end - @sync begin - Threads.@spawn $lBs[1] = left_excitation_transfer_system( - $lBs[1], $H, $exci; solver = $solver - ) - Threads.@spawn $rBs[end] = right_excitation_transfer_system( - $rBs[end], $H, $exci; solver = $solver - ) +function compute_leftenvs!( + envs::InfiniteQPEnvironments, exci::InfiniteQP, H::InfiniteMPOHamiltonian, above, alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + lBs, lenvs = envs.leftBenvs, envs.leftenvs + ids = findall(Base.Fix1(isidentitylevel, H), 2:(size(H[1], 1) - 1)) .+ 1 + AL, AR = exci.left_gs.AL, exci.right_gs.AR + + # push through unitcell to obtain the inhomogeneity + zerovector!(lBs[1]) + for pos in 1:length(exci) + lBs[pos + 1] = lBs[pos] * TransferMatrix(AR[pos], H[pos], AL[pos]; backend, allocator) / + cis(exci.momentum) + lBs[pos + 1] += leftenv(lenvs, pos, exci.left_gs; backend, allocator) * + TransferMatrix(exci[pos], H[pos], AL[pos]; backend, allocator) / cis(exci.momentum) + regularize_GBL!(lBs[pos + 1], exci, ids, pos) end + # solve the fixed point equation + lBs[1] = left_excitation_transfer_system( + lBs[1], H, exci; solver = alg, backend, allocator + ) + + # push the solution through the unitcell lB_cur = lBs[1] for i in 1:(length(exci) - 1) - lB_cur = lB_cur * TransferMatrix(AR[i], H[i], AL[i]) / cis(exci.momentum) - - if istrivial(exci) && !isempty(ids) - ρ_left = l_RL(exci.left_gs, i + 1) - ρ_right = r_RL(exci.left_gs, i) - for k in ids - regularize!(lB_cur[k], ρ_right, ρ_left) - end - end - + lB_cur = lB_cur * TransferMatrix(AR[i], H[i], AL[i]; backend, allocator) / cis(exci.momentum) + regularize_GBL!(lB_cur, exci, ids, i) lBs[i + 1] += lB_cur end + return envs +end + +function compute_rightenvs!( + envs::InfiniteQPEnvironments, exci::InfiniteQP, H::InfiniteMPOHamiltonian, above, alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + rBs, renvs = envs.rightBenvs, envs.rightenvs + ids = findall(Base.Fix1(isidentitylevel, H), 2:(size(H[1], 1) - 1)) .+ 1 + AL, AR = exci.left_gs.AL, exci.right_gs.AR + + # push through unitcell to obtain the inhomogeneity + zerovector!(rBs[end]) + for pos in length(exci):-1:1 + rBs[pos - 1] = TransferMatrix(AL[pos], H[pos], AR[pos]; backend, allocator) * + rBs[pos] * cis(exci.momentum) + rBs[pos - 1] += TransferMatrix(exci[pos], H[pos], AR[pos]; backend, allocator) * + rightenv(renvs, pos, exci.right_gs; backend, allocator) * cis(exci.momentum) + regularize_GBR!(rBs[pos - 1], exci, ids, pos) + end + + # solve the fixed point equation + rBs[end] = right_excitation_transfer_system( + rBs[end], H, exci; solver = alg, backend, allocator + ) + + # push the solution through the unitcell rB_cur = rBs[end] for i in length(exci):-1:2 - rB_cur = TransferMatrix(AL[i], H[i], AR[i]) * rB_cur * cis(exci.momentum) - - if istrivial(exci) && !isempty(ids) - ρ_left = l_LR(exci.left_gs, i) - ρ_right = r_LR(exci.left_gs, i - 1) - for k in ids - regularize!(rB_cur[k], ρ_left, ρ_right) - end - end - + rB_cur = TransferMatrix(AL[i], H[i], AR[i]; backend, allocator) * rB_cur * cis(exci.momentum) + regularize_GBR!(rB_cur, exci, ids, i) rBs[i - 1] += rB_cur end - return InfiniteQPEnvironments(lBs, rBs, lenvs, renvs) + return envs end function environments( exci::FiniteQP, H::FiniteMPOHamiltonian, above = exci, alg = nothing; lenvs = environments(exci.left_gs, H, exci.left_gs), - renvs = istopological(exci) ? environments(exci.right_gs, H, exci.right_gs) : lenvs + renvs = istopological(exci) ? environments(exci.right_gs, H, exci.right_gs) : lenvs, + backend::AbstractBackend = DefaultBackend() ) + # the sweeps below are serial, so a single allocator serves the whole chain + allocator = default_allocator(exci.left_gs, SerialScheduler()) + AL = exci.left_gs.AL AR = exci.right_gs.AR @@ -147,113 +183,151 @@ function environments( zerovector!(lBs[1]) for pos in 1:(length(exci) - 1) - lBs[pos + 1] = lBs[pos] * TransferMatrix(AR[pos], H[pos], AL[pos]) + lBs[pos + 1] = lBs[pos] * TransferMatrix(AR[pos], H[pos], AL[pos]; backend, allocator) lBs[pos + 1] += leftenv(lenvs, pos, exci.left_gs) * - TransferMatrix(exci[pos], H[pos], AL[pos]) + TransferMatrix(exci[pos], H[pos], AL[pos]; backend, allocator) end zerovector!(rBs[end]) for pos in length(exci):-1:2 - rBs[pos - 1] = TransferMatrix(AL[pos], H[pos], AR[pos]) * rBs[pos] - rBs[pos - 1] += TransferMatrix(exci[pos], H[pos], AR[pos]) * + rBs[pos - 1] = TransferMatrix(AL[pos], H[pos], AR[pos]; backend, allocator) * rBs[pos] + rBs[pos - 1] += TransferMatrix(exci[pos], H[pos], AR[pos]; backend, allocator) * rightenv(renvs, pos, exci.right_gs) end return InfiniteQPEnvironments(lBs, rBs, lenvs, renvs) end -function environments(exci::InfiniteQP, O::InfiniteMPO, above, alg; lenvs, renvs) +function environments( + exci::InfiniteQP, O::InfiniteMPO, above, alg; lenvs, renvs, + backend::AbstractBackend = DefaultBackend(), scheduler = Defaults.scheduler[] + ) istopological(exci) && @warn "there is a phase ambiguity in topologically nontrivial statmech excitations" solver = resolve_environment_solver(alg, exci, O, exci) - left_gs = exci.left_gs - right_gs = exci.right_gs - GBL = PeriodicVector([allocate_GBL(exci, O, exci, i) for i in 1:length(exci)]) GBR = PeriodicVector([allocate_GBR(exci, O, exci, i) for i in 1:length(exci)]) + envs = InfiniteQPEnvironments(GBL, GBR, lenvs, renvs) + + # concurrency depends on scheduler, but each half takes its own dedicated scratch allocator + tforeach(1:2; scheduler) do half + allocator = default_allocator(exci.left_gs, SerialScheduler()) + if isone(half) + compute_leftenvs!(envs, exci, O, above, solver; backend, allocator) + else + compute_rightenvs!(envs, exci, O, above, solver; backend, allocator) + end + return nothing + end - left_regularization = map(1:length(exci)) do site + return envs +end + +function compute_leftenvs!( + envs::InfiniteQPEnvironments, exci::InfiniteQP, O::InfiniteMPO, above, alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + GBL, lenvs = envs.leftBenvs, envs.leftenvs + left_gs, right_gs = exci.left_gs, exci.right_gs + + regularization = map(1:length(exci)) do site GL = leftenv(lenvs, site, left_gs) GR = rightenv(lenvs, site, left_gs) return inv(contract_mpo_expval(left_gs.AC[site], GL, O[site], GR)) end - right_regularization = map(1:length(exci)) do site - GL = leftenv(renvs, site, right_gs) - GR = rightenv(renvs, site, right_gs) - return inv(contract_mpo_expval(right_gs.AC[site], GL, O[site], GR)) - end + # push through unitcell to obtain the inhomogeneity # GBL[i] lives on the left MPO bond of site i. Applying site i therefore # produces an object in GBL[i + 1]. This matters when the MPO bond spaces # vary within the unit cell, as for a finite-ring shift MPO. gbl = zerovector!(GBL[1]) for col in 1:length(exci) - gbl = gbl * TransferMatrix(right_gs.AR[col], O[col], left_gs.AL[col]) + gbl = gbl * TransferMatrix(right_gs.AR[col], O[col], left_gs.AL[col]; backend, allocator) gbl += leftenv(lenvs, col, left_gs) * - TransferMatrix(exci[col], O[col], left_gs.AL[col]) - gbl *= left_regularization[col] * cis(-exci.momentum) + TransferMatrix(exci[col], O[col], left_gs.AL[col]; backend, allocator) + gbl *= regularization[col] * cis(-exci.momentum) GBL[col + 1] = gbl end + # solve the fixed point equation + T_RL = TransferMatrix(right_gs.AR, O, left_gs.AL; backend, allocator) + if istrivial(exci) + @plansor rvec[-1 -2; -3] := rightenv(lenvs, 0, left_gs)[-1 -2; 1] * + conj(left_gs.C[0][-3; 1]) + @plansor lvec[-1 -2; -3] := leftenv(lenvs, 1, left_gs)[-1 -2; 1] * + left_gs.C[0][1; -3] + T_RL = regularize(T_RL, lvec, rvec) + end + + GBL[1], convhist = linsolve( + flip(T_RL), gbl, gbl, alg, 1, + -cis(-length(exci) * exci.momentum) * prod(regularization) + ) + convhist.converged == 0 && + @warn "GBL failed to converge: normres = $(convhist.normres)" + + # push the solution through the unitcell + left_cur = GBL[1] + for col in 1:(length(exci) - 1) + left_cur = regularization[col] * left_cur * + TransferMatrix(right_gs.AR[col], O[col], left_gs.AL[col]; backend, allocator) * + cis(-exci.momentum) + GBL[col + 1] += left_cur + end + + return envs +end + +function compute_rightenvs!( + envs::InfiniteQPEnvironments, exci::InfiniteQP, O::InfiniteMPO, above, alg; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + GBR, renvs = envs.rightBenvs, envs.rightenvs + left_gs, right_gs = exci.left_gs, exci.right_gs + + regularization = map(1:length(exci)) do site + GL = leftenv(renvs, site, right_gs) + GR = rightenv(renvs, site, right_gs) + return inv(contract_mpo_expval(right_gs.AC[site], GL, O[site], GR)) + end + + # push through unitcell to obtain the inhomogeneity # GBR[i] lives on the right MPO bond of site i. Applying site i from the # right therefore produces an object in GBR[i - 1]. gbr = zerovector!(GBR[end]) for col in reverse(1:length(exci)) - gbr = TransferMatrix(left_gs.AL[col], O[col], right_gs.AR[col]) * gbr - gbr += TransferMatrix(exci[col], O[col], right_gs.AR[col]) * + gbr = TransferMatrix(left_gs.AL[col], O[col], right_gs.AR[col]; backend, allocator) * gbr + gbr += TransferMatrix(exci[col], O[col], right_gs.AR[col]; backend, allocator) * rightenv(renvs, col, right_gs) - gbr *= right_regularization[col] * cis(exci.momentum) + gbr *= regularization[col] * cis(exci.momentum) GBR[col - 1] = gbr end - T_RL = TransferMatrix(right_gs.AR, O, left_gs.AL) - T_LR = TransferMatrix(left_gs.AL, O, right_gs.AR) - + # solve the fixed point equation + T_LR = TransferMatrix(left_gs.AL, O, right_gs.AR; backend, allocator) if istrivial(exci) - @plansor rvec[-1 -2; -3] := rightenv(lenvs, 0, left_gs)[-1 -2; 1] * - conj(left_gs.C[0][-3; 1]) - @plansor lvec[-1 -2; -3] := leftenv(lenvs, 1, left_gs)[-1 -2; 1] * - left_gs.C[0][1; -3] - - T_RL = regularize(T_RL, lvec, rvec) - @plansor rvec[-1 -2; -3] := rightenv(renvs, 0, right_gs)[1 -2; -3] * right_gs.C[0][-1; 1] @plansor lvec[-1 -2; -3] := conj(right_gs.C[0][-3; 1]) * leftenv(renvs, 1, right_gs)[-1 -2; 1] - T_LR = regularize(T_LR, lvec, rvec) end - GBL[1], convhist = linsolve( - flip(T_RL), gbl, gbl, solver, 1, - -cis(-length(exci) * exci.momentum) * prod(left_regularization) - ) - - convhist.converged == 0 && - @warn "GBL failed to converge: normres = $(convhist.normres)" - GBR[end], convhist = linsolve( - T_LR, gbr, gbr, GMRES(), 1, - -cis(length(exci) * exci.momentum) * prod(right_regularization) + T_LR, gbr, gbr, alg, 1, + -cis(length(exci) * exci.momentum) * prod(regularization) ) convhist.converged == 0 && @warn "GBR failed to converge: normres = $(convhist.normres)" - left_cur = GBL[1] + # push the solution through the unitcell right_cur = GBR[end] - for col in 1:(length(exci) - 1) - left_cur = left_regularization[col] * left_cur * - TransferMatrix(right_gs.AR[col], O[col], left_gs.AL[col]) * - cis(-exci.momentum) - GBL[col + 1] += left_cur - - col = length(exci) - col + 1 - right_cur = TransferMatrix(left_gs.AL[col], O[col], right_gs.AR[col]) * right_cur * - cis(exci.momentum) * right_regularization[col] + for col in reverse(2:length(exci)) + right_cur = TransferMatrix(left_gs.AL[col], O[col], right_gs.AR[col]; backend, allocator) * + right_cur * cis(exci.momentum) * regularization[col] GBR[col - 1] += right_cur end - return InfiniteQPEnvironments(GBL, GBR, lenvs, renvs) + return envs end diff --git a/src/states/ortho.jl b/src/states/ortho.jl index 68b062d35..515f899f2 100644 --- a/src/states/ortho.jl +++ b/src/states/ortho.jl @@ -112,8 +112,13 @@ gaugefix! function gaugefix!( ψ::InfiniteMPS, A, C₀ = ψ.C[end]; - order = :LR, timeroutput = NoTimerOutput(), kwargs... + order = :LR, timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator(), + kwargs... ) + # `backend` and `allocator` are named explicitly so that they are not swept into the + # algorithm constructor's keyword arguments, which describe the gauge itself. Same + # convention as `recalculate!`. alg = if order === :LR || order === :RL MixedCanonical(; order, kwargs...) elseif order === :L @@ -124,20 +129,21 @@ function gaugefix!( throw(ArgumentError("Invalid order: $order")) end - return gaugefix!(ψ, A, C₀, alg; timeroutput) + return gaugefix!(ψ, A, C₀, alg; timeroutput, backend, allocator) end # expert mode: actual implementation function gaugefix!( ψ::InfiniteMPS, A, C₀, alg::MixedCanonical; - timeroutput = NoTimerOutput() + timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) if alg.order === :LR - gaugefix!(ψ, A, C₀, alg.alg_leftcanonical; timeroutput) - gaugefix!(ψ, ψ.AL, ψ.C[end], alg.alg_rightcanonical; timeroutput) + gaugefix!(ψ, A, C₀, alg.alg_leftcanonical; timeroutput, backend, allocator) + gaugefix!(ψ, ψ.AL, ψ.C[end], alg.alg_rightcanonical; timeroutput, backend, allocator) elseif alg.order === :RL - gaugefix!(ψ, A, C₀, alg.alg_rightcanonical; timeroutput) - gaugefix!(ψ, ψ.AR, ψ.C[end], alg.alg_leftcanonical; timeroutput) + gaugefix!(ψ, A, C₀, alg.alg_rightcanonical; timeroutput, backend, allocator) + gaugefix!(ψ, ψ.AR, ψ.C[end], alg.alg_leftcanonical; timeroutput, backend, allocator) else throw(ArgumentError("Invalid order: $(alg.order)")) end @@ -145,16 +151,18 @@ function gaugefix!( end function gaugefix!( ψ::InfiniteMPS, A, C₀, alg::LeftCanonical; - timeroutput = NoTimerOutput() + timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) - uniform_leftorth!((ψ.AL, ψ.C), A, C₀, alg; timeroutput) + uniform_leftorth!((ψ.AL, ψ.C), A, C₀, alg; timeroutput, backend, allocator) return ψ end function gaugefix!( ψ::InfiniteMPS, A, C₀, alg::RightCanonical; - timeroutput = NoTimerOutput() + timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) - uniform_rightorth!((ψ.AR, ψ.C), A, C₀, alg; timeroutput) + uniform_rightorth!((ψ.AR, ψ.C), A, C₀, alg; timeroutput, backend, allocator) return ψ end @@ -214,7 +222,8 @@ end function uniform_leftorth!( (AL, C), A, C₀, alg::LeftCanonical; - timeroutput = NoTimerOutput() + timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) C[end] = normalize!(C₀) return LoggingExtras.withlevel(; alg.verbosity) do @@ -222,7 +231,7 @@ function uniform_leftorth!( log = IterLog("LC") A_tail = _transpose_tail.(A) # pre-transpose A CA_tail = similar.(A_tail) # pre-allocate workspace - state = (; AL, C, A, A_tail, CA_tail, iter = 0, ϵ = Inf, timeroutput) + state = (; AL, C, A, A_tail, CA_tail, iter = 0, ϵ = Inf, timeroutput, backend, allocator) it = IterativeSolver(alg, state) loginit!(log, it.ϵ) @@ -250,27 +259,30 @@ function Base.iterate(it::IterativeSolver{LeftCanonical}, state = it.state) iter = state.iter + 1 it.state = (; state.AL, state.C, state.A, state.A_tail, state.CA_tail, iter, ϵ, timeroutput, + state.backend, state.allocator, ) return (it.state.AL, it.state.C), it.state end function gauge_eigsolve_step!(it::IterativeSolver{LeftCanonical}, state) - (; AL, C, A, iter, ϵ) = state + (; AL, C, A, iter, ϵ, backend, allocator) = state if iter ≥ it.eig_miniter alg_eigsolve = adapt_solver(it.alg_eigsolve; iter = 1, g_global = ϵ^2) - _, vec = fixedpoint(flip(TransferMatrix(A, AL)), C[end], :LM, alg_eigsolve) + _, vec = fixedpoint( + flip(TransferMatrix(A, AL; backend, allocator)), C[end], :LM, alg_eigsolve + ) _, C[end] = left_orth!(vec; alg = it.alg_orth) end return C[end] end function gauge_orth_step!(it::IterativeSolver{LeftCanonical}, state) - (; AL, C, A_tail, CA_tail) = state + (; AL, C, A_tail, CA_tail, backend, allocator) = state for i in 1:length(AL) # repartition!(A_tail[i], AL[i]) mul!(CA_tail[i], C[i - 1], A_tail[i]) - repartition!(AL[i], CA_tail[i]) + repartition!(AL[i], CA_tail[i], One(), Zero(), backend, allocator) AL[i], C[i] = left_orth!(AL[i]; alg = it.alg_orth) end normalize!(C[end]) @@ -279,14 +291,15 @@ end function uniform_rightorth!( (AR, C), A, C₀, alg::RightCanonical; - timeroutput = NoTimerOutput() + timeroutput = NoTimerOutput(), + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) C[end] = normalize!(C₀) return LoggingExtras.withlevel(; alg.verbosity) do # initialize algorithm and temporary variables log = IterLog("RC") AC_tail = _similar_tail.(A) # pre-allocate workspace - state = (; AR, C, A, AC_tail, iter = 0, ϵ = Inf, timeroutput) + state = (; AR, C, A, AC_tail, iter = 0, ϵ = Inf, timeroutput, backend, allocator) it = IterativeSolver(alg, state) loginit!(log, it.ϵ) @@ -312,28 +325,30 @@ function Base.iterate(it::IterativeSolver{RightCanonical}, state = it.state) ϵ = oftype(state.ϵ, norm(C₀ - C₁)) iter = state.iter + 1 - it.state = (; state.AR, state.C, state.A, state.AC_tail, iter, ϵ, timeroutput) + it.state = (; state.AR, state.C, state.A, state.AC_tail, iter, ϵ, timeroutput, state.backend, state.allocator) return (it.state.AR, it.state.C), it.state end function gauge_eigsolve_step!(it::IterativeSolver{RightCanonical}, state) - (; AR, C, A, iter, ϵ) = state + (; AR, C, A, iter, ϵ, backend, allocator) = state if iter ≥ it.eig_miniter alg_eigsolve = adapt_solver(it.alg_eigsolve; iter = 1, g_global = ϵ^2) - _, vec = fixedpoint(TransferMatrix(A, AR), C[end], :LM, alg_eigsolve) + _, vec = fixedpoint( + TransferMatrix(A, AR; backend, allocator), C[end], :LM, alg_eigsolve + ) C[end], _ = right_orth!(vec; alg = it.alg_orth) end return C[end] end function gauge_orth_step!(it::IterativeSolver{RightCanonical}, state) - (; A, AR, C, AC_tail) = state + (; A, AR, C, AC_tail, backend, allocator) = state for i in length(AR):-1:1 AC = mul!(AR[i], A[i], C[i]) # use AR as temporary storage for A * C - tmp = repartition!(AC_tail[i], AC) + tmp = repartition!(AC_tail[i], AC, One(), Zero(), backend, allocator) C[i - 1], tmp = right_orth!(tmp; alg = it.alg_orth) - repartition!(AR[i], tmp) # TODO: avoid doing this every iteration + repartition!(AR[i], tmp, One(), Zero(), backend, allocator) # TODO: avoid doing this every iteration end normalize!(C[end]) return C[end] diff --git a/src/transfermatrix/transfer.jl b/src/transfermatrix/transfer.jl index 9cf61e363..776491256 100644 --- a/src/transfermatrix/transfer.jl +++ b/src/transfermatrix/transfer.jl @@ -18,14 +18,15 @@ apply a transfer matrix to the left. @generated function transfer_left( v::AbstractTensorMap{<:Any, S, 1, N₁}, A::GenericMPSTensor{S, N₂}, - Abar::GenericMPSTensor{S, N₂} + Abar::GenericMPSTensor{S, N₂}; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) where {S, N₁, N₂} t_out = tensorexpr(:v, -1, -(2:(N₁ + 1))) t_top = tensorexpr(:A, 2:(N₂ + 1), -(N₁ + 1)) t_bot = tensorexpr(:Abar, (1, (3:(N₂ + 1))...), -1) t_in = tensorexpr(:v, 1, (-(2:N₁)..., 2)) return macroexpand( - @__MODULE__, :(return @plansor $t_out := $t_in * $t_top * conj($t_bot)) + @__MODULE__, :(return @plansor backend = backend allocator = allocator $t_out := $t_in * $t_top * conj($t_bot)) ) end @@ -43,120 +44,125 @@ apply a transfer matrix to the right. @generated function transfer_right( v::AbstractTensorMap{<:Any, S, 1, N₁}, A::GenericMPSTensor{S, N₂}, - Abar::GenericMPSTensor{S, N₂} + Abar::GenericMPSTensor{S, N₂}; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) where {S, N₁, N₂} t_out = tensorexpr(:v, -1, -(2:(N₁ + 1))) t_top = tensorexpr(:A, (-1, reverse(3:(N₂ + 1))...), 1) t_bot = tensorexpr(:Abar, (-(N₁ + 1), reverse(3:(N₂ + 1))...), 2) t_in = tensorexpr(:v, 1, (-(2:N₁)..., 2)) return macroexpand( - @__MODULE__, :(return @plansor $t_out := $t_top * conj($t_bot) * $t_in) + @__MODULE__, :(return @plansor backend = backend allocator = allocator $t_out := $t_top * conj($t_bot) * $t_in) ) end # transfer, but the upper A is an excited tensor -function transfer_left(v::MPSBondTensor, A::MPOTensor, Ab::MPSTensor) - return @plansor t[-1; -2 -3] := v[1; 2] * A[2 3; -2 -3] * conj(Ab[1 3; -1]) +function transfer_left(v::MPSBondTensor, A::MPOTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1; -2 -3] := v[1; 2] * A[2 3; -2 -3] * conj(Ab[1 3; -1]) end -function transfer_right(v::MPSBondTensor, A::MPOTensor, Ab::MPSTensor) - return @plansor t[-1; -2 -3] := A[-1 3; -2 1] * v[1; 2] * conj(Ab[-3 3; 2]) +function transfer_right(v::MPSBondTensor, A::MPOTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1; -2 -3] := A[-1 3; -2 1] * v[1; 2] * conj(Ab[-3 3; 2]) end # transfer, but the upper A is an excited tensor and there is an mpo leg being passed through -function transfer_left(v::MPSTensor, A::MPOTensor, Ab::MPSTensor) - return @plansor t[-1 -2; -3 -4] := v[1 3; 4] * A[4 5; -3 -4] * τ[3 2; 5 -2] * conj(Ab[1 2; -1]) +function transfer_left(v::MPSTensor, A::MPOTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1 -2; -3 -4] := v[1 3; 4] * A[4 5; -3 -4] * τ[3 2; 5 -2] * conj(Ab[1 2; -1]) end -function transfer_right(v::MPSTensor, A::MPOTensor, Ab::MPSTensor) - return @plansor t[-1 -2; -3 -4] := A[-1 4; -3 5] * τ[-2 3; 4 2] * conj(Ab[-4 3; 1]) * v[5 2; 1] +function transfer_right(v::MPSTensor, A::MPOTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1 -2; -3 -4] := A[-1 4; -3 5] * τ[-2 3; 4 2] * conj(Ab[-4 3; 1]) * v[5 2; 1] end # the transfer operation of a density matrix with a utility leg in its codomain is ill defined - how should one braid the utility leg? # hence the checks - to make sure that this operation is uniquely defined -function transfer_left(v::MPSTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}) where {S} +function transfer_left(v::MPSTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) where {S} check_unambiguous_braiding(space(v, 2)) - return @plansor v[-1 -2; -3] := v[1 2; 4] * A[4 5; -3] * τ[2 3; 5 -2] * conj(Ab[1 3; -1]) + return @plansor backend = backend allocator = allocator v[-1 -2; -3] := v[1 2; 4] * A[4 5; -3] * τ[2 3; 5 -2] * conj(Ab[1 3; -1]) end -function transfer_right(v::MPSTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}) where {S} +function transfer_right(v::MPSTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) where {S} check_unambiguous_braiding(space(v, 2)) - return @plansor v[-1 -2; -3] := A[-1 2; 1] * τ[-2 4; 2 3] * conj(Ab[-3 4; 5]) * v[1 3; 5] + return @plansor backend = backend allocator = allocator v[-1 -2; -3] := A[-1 2; 1] * τ[-2 4; 2 3] * conj(Ab[-3 4; 5]) * v[1 3; 5] end # similar for Matrix Product Density Operators function transfer_left( - v::MPSTensor{S}, A::GenericMPSTensor{S, 3}, Ab::GenericMPSTensor{S, 3} + v::MPSTensor{S}, A::GenericMPSTensor{S, 3}, Ab::GenericMPSTensor{S, 3}; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) where {S} check_unambiguous_braiding(space(v, 2)) - return @plansor v[-1 -2; -3] ≔ v[1 3; 8] * A[8 7 6; -3] * τ[3 2; 7 5] * τ[5 4; 6 -2] * + return @plansor backend = backend allocator = allocator v[-1 -2; -3] ≔ v[1 3; 8] * A[8 7 6; -3] * τ[3 2; 7 5] * τ[5 4; 6 -2] * conj(Ab[1 2 4; -1]) end function transfer_right( - v::MPSTensor{S}, A::GenericMPSTensor{S, 3}, Ab::GenericMPSTensor{S, 3} + v::MPSTensor{S}, A::GenericMPSTensor{S, 3}, Ab::GenericMPSTensor{S, 3}; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) where {S} check_unambiguous_braiding(space(v, 2)) - return @plansor v[-1 -2; -3] ≔ A[-1 4 2; 1] * τ[-2 6; 4 5] * τ[5 7; 2 3] * + return @plansor backend = backend allocator = allocator v[-1 -2; -3] ≔ A[-1 4 2; 1] * τ[-2 6; 4 5] * τ[5 7; 2 3] * conj(Ab[-3 6 7; 8]) * v[1 3; 8] end # the transfer operation with a utility leg in both the domain and codomain is also ill defined - only due to the codomain utility space -function transfer_left(v::MPOTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}) where {S} +function transfer_left(v::MPOTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) where {S} check_unambiguous_braiding(space(v, 2)) - return @plansor t[-1 -2; -3 -4] := v[1 2; -3 4] * A[4 5; -4] * τ[2 3; 5 -2] * conj(Ab[1 3; -1]) + return @plansor backend = backend allocator = allocator t[-1 -2; -3 -4] := v[1 2; -3 4] * A[4 5; -4] * τ[2 3; 5 -2] * conj(Ab[1 3; -1]) end -function transfer_right(v::MPOTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}) where {S} +function transfer_right(v::MPOTensor{S}, A::MPSTensor{S}, Ab::MPSTensor{S}; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) where {S} check_unambiguous_braiding(space(v, 2)) - return @plansor t[-1 -2; -3 -4] := A[-1 2; 1] * τ[-2 4; 2 3] * conj(Ab[-4 4; 5]) * v[1 3; -3 5] + return @plansor backend = backend allocator = allocator t[-1 -2; -3 -4] := A[-1 2; 1] * τ[-2 4; 2 3] * conj(Ab[-4 4; 5]) * v[1 3; -3 5] end #transfer for 2 mpo tensors -function transfer_left(v::MPSBondTensor, A::MPOTensor, B::MPOTensor) - return @plansor t[-1; -2] := v[1; 2] * A[2 3; 4 -2] * conj(B[1 3; 4 -1]) +function transfer_left(v::MPSBondTensor, A::MPOTensor, B::MPOTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1; -2] := v[1; 2] * A[2 3; 4 -2] * conj(B[1 3; 4 -1]) end -function transfer_right(v::MPSBondTensor, A::MPOTensor, B::MPOTensor) - return @plansor t[-1; -2] := A[-1 3; 4 1] * conj(B[-2 3; 4 2]) * v[1; 2] +function transfer_right(v::MPSBondTensor, A::MPOTensor, B::MPOTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1; -2] := A[-1 3; 4 1] * conj(B[-2 3; 4 2]) * v[1; 2] end # ---------------------------------------------------- # | transfers for (vector, operator, tensor, tensor) | # ---------------------------------------------------- -transfer_left(v, ::Nothing, A, B) = transfer_left(v, A, B); -transfer_right(v, ::Nothing, A, B) = transfer_right(v, A, B); +transfer_left(v, ::Nothing, A, B; kwargs...) = transfer_left(v, A, B; kwargs...); +transfer_right(v, ::Nothing, A, B; kwargs...) = transfer_right(v, A, B; kwargs...); #mpo transfer -function transfer_left(x::MPSTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor) - return @plansor y[-1 -2; -3] := x[1 2; 4] * A[4 5; -3] * O[2 3; 5 -2] * conj(Ab[1 3; -1]) +function transfer_left(x::MPSTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator y[-1 -2; -3] := x[1 2; 4] * A[4 5; -3] * O[2 3; 5 -2] * conj(Ab[1 3; -1]) end -function transfer_right(v::MPSTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor) - return @plansor v[-1 -2; -3] := A[-1 2; 1] * O[-2 4; 2 3] * conj(Ab[-3 4; 5]) * v[1 3; 5] +function transfer_right(v::MPSTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator v[-1 -2; -3] := A[-1 2; 1] * O[-2 4; 2 3] * conj(Ab[-3 4; 5]) * v[1 3; 5] end #mpo transfer, but with A an excitation-tensor -function transfer_left(v::MPSTensor, O::MPOTensor, A::MPOTensor, Ab::MPSTensor) - return @plansor t[-1 -2; -3 -4] := v[4 2; 1] * A[1 3; -3 -4] * O[2 5; 3 -2] * conj(Ab[4 5; -1]) +function transfer_left(v::MPSTensor, O::MPOTensor, A::MPOTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1 -2; -3 -4] := v[4 2; 1] * A[1 3; -3 -4] * O[2 5; 3 -2] * conj(Ab[4 5; -1]) end -function transfer_right(v::MPSTensor, O::MPOTensor, A::MPOTensor, Ab::MPSTensor) - return @plansor t[-1 -2; -3 -4] := A[-1 4; -3 5] * O[-2 2; 4 3] * conj(Ab[-4 2; 1]) * v[5 3; 1] +function transfer_right(v::MPSTensor, O::MPOTensor, A::MPOTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator t[-1 -2; -3 -4] := A[-1 4; -3 5] * O[-2 2; 4 3] * conj(Ab[-4 2; 1]) * v[5 3; 1] end #mpo transfer, with an excitation leg -function transfer_left(v::MPOTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor) - return @plansor v[-1 -2; -3 -4] := v[4 2; -3 1] * A[1 3; -4] * O[2 5; 3 -2] * conj(Ab[4 5; -1]) +function transfer_left(v::MPOTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator v[-1 -2; -3 -4] := v[4 2; -3 1] * A[1 3; -4] * O[2 5; 3 -2] * conj(Ab[4 5; -1]) end -function transfer_right(v::MPOTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor) - return @plansor v[-1 -2; -3 -4] := A[-1 4; 5] * O[-2 2; 4 3] * +function transfer_right(v::MPOTensor, O::MPOTensor, A::MPSTensor, Ab::MPSTensor; backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator()) + return @plansor backend = backend allocator = allocator v[-1 -2; -3 -4] := A[-1 4; 5] * O[-2 2; 4 3] * conj(Ab[-4 2; 1]) * v[5 3; -3 1] end # density matrix transfer # braid handedness dependent on MPO-to-MPS conversion function transfer_left( - x::MPSTensor, O::MPOTensor, A::GenericMPSTensor{<:Any, 3}, Ab::GenericMPSTensor{<:Any, 3} + x::MPSTensor, O::MPOTensor, A::GenericMPSTensor{<:Any, 3}, Ab::GenericMPSTensor{<:Any, 3}; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) - return @plansor y[-1 -2; -3] ≔ x[1 2; 6] * A[6 7 8; -3] * O[2 3; 7 5] * τ[5 4; 8 -2] * + return @plansor backend = backend allocator = allocator y[-1 -2; -3] ≔ x[1 2; 6] * A[6 7 8; -3] * O[2 3; 7 5] * τ[5 4; 8 -2] * conj(Ab[1 3 4; -1]) end function transfer_right( - x::MPSTensor, O::MPOTensor, A::GenericMPSTensor{<:Any, 3}, Ab::GenericMPSTensor{<:Any, 3} + x::MPSTensor, O::MPOTensor, A::GenericMPSTensor{<:Any, 3}, Ab::GenericMPSTensor{<:Any, 3}; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() ) - return @plansor y[-1 -2; -3] ≔ A[-1 4 2; 1] * O[-2 6; 4 5] * τ[5 7; 2 3] * + return @plansor backend = backend allocator = allocator y[-1 -2; -3] ≔ A[-1 4 2; 1] * O[-2 6; 4 5] * τ[5 7; 2 3] * conj(Ab[-3 6 7; 8]) * x[1 3; 8] end diff --git a/src/transfermatrix/transfermatrix.jl b/src/transfermatrix/transfermatrix.jl index 31a33ab06..401e97f31 100644 --- a/src/transfermatrix/transfermatrix.jl +++ b/src/transfermatrix/transfermatrix.jl @@ -1,12 +1,14 @@ abstract type AbstractTransferMatrix end; # single site transfer -struct SingleTransferMatrix{A <: AbstractTensorMap, B, C <: AbstractTensorMap} <: +struct SingleTransferMatrix{A <: AbstractTensorMap, B, C <: AbstractTensorMap, Bk, Al} <: AbstractTransferMatrix above::A middle::B below::C isflipped::Bool + backend::Bk + allocator::Al end #the product of transfer matrices is its own type @@ -34,7 +36,9 @@ end #flip em function TensorKit.flip(tm::SingleTransferMatrix) - return SingleTransferMatrix(tm.above, tm.middle, tm.below, !tm.isflipped) + return SingleTransferMatrix( + tm.above, tm.middle, tm.below, !tm.isflipped, tm.backend, tm.allocator + ) end; TensorKit.flip(tm::ProductTransferMatrix) = ProductTransferMatrix(flip.(reverse(tm.tms))); TensorKit.flip(tm::RegTransferMatrix) = RegTransferMatrix(flip(tm.tm), tm.rvec, tm.lvec); @@ -47,21 +51,29 @@ Base.:*(vec, tm::AbstractTransferMatrix) = flip(tm)(vec); (d::ProductTransferMatrix)(vec) = foldr((a, b) -> a(b), d.tms; init = vec); function (d::SingleTransferMatrix)(vec) return if d.isflipped - transfer_left(vec, d.middle, d.above, d.below) + transfer_left(vec, d.middle, d.above, d.below; d.backend, d.allocator) else - transfer_right(vec, d.middle, d.above, d.below) + transfer_right(vec, d.middle, d.above, d.below; d.backend, d.allocator) end end; (d::RegTransferMatrix)(vec) = regularize!(d.tm * vec, d.lvec, d.rvec); # constructors -TransferMatrix(a) = TransferMatrix(a, nothing, a); -TransferMatrix(a, b) = TransferMatrix(a, nothing, b); -function TransferMatrix(a::AbstractTensorMap, b, c::AbstractTensorMap, isflipped = false) - return SingleTransferMatrix(a, b, c, isflipped) +TransferMatrix(a; kwargs...) = TransferMatrix(a, nothing, a; kwargs...); +TransferMatrix(a, b; kwargs...) = TransferMatrix(a, nothing, b; kwargs...); +function TransferMatrix( + a::AbstractTensorMap, b, c::AbstractTensorMap, isflipped = false; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + return SingleTransferMatrix(a, b, c, isflipped, backend, allocator) end -function TransferMatrix(a::AbstractVector, b, c::AbstractVector, isflipped = false) - tot = ProductTransferMatrix(convert(Vector, TransferMatrix.(a, b, c))) +function TransferMatrix( + a::AbstractVector, b, c::AbstractVector, isflipped = false; + backend::AbstractBackend = DefaultBackend(), allocator = DefaultAllocator() + ) + tot = ProductTransferMatrix( + convert(Vector, TransferMatrix.(a, b, c; backend, allocator)) + ) return isflipped ? flip(tot) : tot end