Iterative VUMPS and VOMPS

Author: lkdvos

src/algorithms/approximate/vomps.jl

@@ -5,63 +5,105 @@ Base.@deprecate(approximate(ψ::MultilineMPS, toapprox::Tuple{<:MultilineMPO,<:M
                                   alg.verbosity, alg.alg_gauge, alg.alg_environments),
                             envs...; kwargs...))
 
-function approximate(ψ::MultilineMPS, toapprox::Tuple{<:MultilineMPO,<:MultilineMPS},
-                     alg::VOMPS, envs=environments(ψ, toapprox))
-    ϵ::Float64 = calc_galerkin(ψ, toapprox..., envs)
-    temp_ACs = similar.(ψ.AC)
-    scheduler = Defaults.scheduler[]
+function approximate(mps::MultilineMPS, toapprox::Tuple{<:MultilineMPO,<:MultilineMPS},
+                     alg::VOMPS, envs=environments(mps, toapprox))
     log = IterLog("VOMPS")
-    alg_environments = updatetol(alg.alg_environments, 0, ϵ)
-    recalculate!(envs, ψ, toapprox...; alg_environments.tol)
+    iter = 0
+    ϵ = calc_galerkin(mps, toapprox..., envs)
+    alg_environments = updatetol(alg.alg_environments, iter, ϵ)
+    recalculate!(envs, mps, toapprox...; alg_environments.tol)
 
-    LoggingExtras.withlevel(; alg.verbosity) do
+    state = VOMPSState(mps, toapprox, envs, iter, ϵ)
+    it = IterativeSolver(alg, state)
+
+    return LoggingExtras.withlevel(; alg.verbosity) do
         @infov 2 loginit!(log, ϵ)
-        for iter in 1:(alg.maxiter)
-            tmap!(eachcol(temp_ACs), 1:size(ψ, 2); scheduler) do col
-                return _vomps_localupdate(col, ψ, toapprox, envs)
+
+        for (mps, envs, ϵ) in it
+            if ϵ ≤ alg.tol
+                @infov 2 logfinish!(log, it.iter, ϵ)
+                return mps, envs, ϵ
+            end
+            if it.iter ≥ alg.maxiter
+                @warnv 1 logcancel!(log, it.iter, ϵ)
+                return mps, envs, ϵ
             end
+            @infov 3 logiter!(log, it.iter, ϵ)
+        end
 
-            alg_gauge = updatetol(alg.alg_gauge, iter, ϵ)
-            ψ = MultilineMPS(temp_ACs, ψ.C[:, end]; alg_gauge.tol, alg_gauge.maxiter)
+        # this should never be reached
+        return it.state.mps, it.state.envs, it.state.ϵ
+    end
+end
 
-            alg_environments = updatetol(alg.alg_environments, iter, ϵ)
-            recalculate!(envs, ψ, toapprox...; alg_environments.tol)
+# need to specialize a bunch of functions because different arguments are passed with tuples
+# TODO: can we avoid this?
+function Base.iterate(it::IterativeSolver{<:VOMPS},
+                      state::VOMPSState{<:Any,<:Tuple}=it.state)
+    ACs = localupdate_step!(it, state)
+    mps = gauge_step!(it, state, ACs)
+    envs = envs_step!(it, state, mps)
 
-            ψ, envs = alg.finalize(iter, ψ, toapprox, envs)::Tuple{typeof(ψ),typeof(envs)}
+    # finalizer step
+    mps, envs = it.finalize(state.iter, mps, state.operator, envs)::typeof((mps, envs))
 
-            ϵ = calc_galerkin(ψ, toapprox..., envs)
+    # error criterion
+    ϵ = calc_galerkin(mps, state.operator..., envs)
 
-            if ϵ <= alg.tol
-                @infov 2 logfinish!(log, iter, ϵ)
-                break
-            end
-            if iter == alg.maxiter
-                @warnv 1 logcancel!(log, iter, ϵ)
-            else
-                @infov 3 logiter!(log, iter, ϵ)
-            end
-        end
+    # update state
+    it.state = VOMPSState(mps, state.operator, envs, state.iter + 1, ϵ)
+
+    return (mps, envs, ϵ), it.state
+end
+
+# TODO: ac_proj and c_proj should be rewritten to also be simply ∂AC/∂C functions
+# once these have better support for different above/below mps
+function localupdate_step!(::IterativeSolver{<:VOMPS}, state::VOMPSState{<:Any,<:Tuple},
+                           ::SerialScheduler)
+    alg_orth = QRpos()
+    eachsite = 1:length(state.mps)
+    ACs = similar(state.mps.AC)
+    dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs
+
+    foreach(eachsite) do site
+        AC = circshift([ac_proj(row, loc, state.mps, state.toapprox, state.envs)
+                        for row in 1:size(state.mps, 1)], 1)
+        C = circshift([c_proj(row, loc, state.mps, state.toapprox, state.envs)
+                       for row in 1:size(state.mps, 1)], 1)
+        dst_ACs[site] = regauge!(AC, C; alg=alg_orth)
+        return nothing
     end
 
-    return ψ, envs, ϵ
+    return ACs
 end
+function localupdate_step!(::IterativeSolver{<:VOMPS}, state::VOMPSState{<:Any,<:Tuple},
+                           scheduler)
+    alg_orth = QRpos()
+    eachsite = 1:length(state.mps)
 
-function _vomps_localupdate(loc, ψ, Oϕ, envs, factalg=QRpos())
-    local tmp_AC, tmp_C
-    if Defaults.scheduler[] isa SerialScheduler
-        tmp_AC = circshift([ac_proj(row, loc, ψ, Oϕ, envs) for row in 1:size(ψ, 1)], 1)
-        tmp_C = circshift([c_proj(row, loc, ψ, Oϕ, envs) for row in 1:size(ψ, 1)], 1)
-    else
+    ACs = similar(state.mps.AC)
+    dst_ACs = state.mps isa Multiline ? eachcol(ACs) : ACs
+
+    tforeach(eachsite; scheduler) do site
+        local AC, C
         @sync begin
             Threads.@spawn begin
-                tmp_AC = circshift([ac_proj(row, loc, ψ, Oϕ, envs)
-                                    for row in 1:size(ψ, 1)], 1)
+                AC = circshift([ac_proj(row, site, state.mps, state.operator, state.envs)
+                                for row in 1:size(state.mps, 1)], 1)
             end
             Threads.@spawn begin
-                tmp_C = circshift([c_proj(row, loc, ψ, Oϕ, envs) for row in 1:size(ψ, 1)],
-                                  1)
+                C = circshift([c_proj(row, site, state.mps, state.operator, state.envs)
+                               for row in 1:size(state.mps, 1)], 1)
             end
         end
+        dst_ACs[site] = regauge!(AC, C; alg=alg_orth)
+        return nothing
     end
-    return regauge!.(tmp_AC, tmp_C; alg=factalg)
+
+    return ACs
+end
+
+function envs_step!(it::IterativeSolver{<:VOMPS}, state::VOMPSState{<:Any,<:Tuple}, mps)
+    alg_environments = updatetol(it.alg_environments, state.iter, state.ϵ)
+    return recalculate!(state.envs, mps, state.operator...; alg_environments.tol)
 end

src/algorithms/groundstate/vumps.jl

@@ -35,77 +35,124 @@ $(TYPEDFIELDS)
     finalize::F = Defaults._finalize
 end
 
-function find_groundstate(ψ::InfiniteMPS, H, alg::VUMPS, envs=environments(ψ, H))
+# Internal state of the VUMPS algorithm
+struct VUMPSState{S,O,E}
+    mps::S
+    operator::O
+    envs::E
+    iter::Int
+    ϵ::Float64
+    which::Symbol
+end
+
+function find_groundstate(mps::InfiniteMPS, operator, alg::VUMPS,
+                          envs=environments(mps, operator))
+    return dominant_eigsolve(operator, mps, alg, envs; which=:SR)
+end
+
+function dominant_eigsolve(operator, mps, alg::VUMPS, envs=environments(mps, operator);
+                           which)
     log = IterLog("VUMPS")
-    ϵ::Float64 = calc_galerkin(ψ, H, ψ, envs)
-    ACs = similar.(ψ.AC)
-    alg_environments = updatetol(alg.alg_environments, 0, ϵ)
-    recalculate!(envs, ψ, H, ψ; alg_environments.tol)
+    iter = 0
+    ϵ = calc_galerkin(mps, operator, mps, envs)
+    alg_environments = updatetol(alg.alg_environments, iter, ϵ)
+    recalculate!(envs, mps, operator, mps; alg_environments.tol)
 
-    state = (; ψ, H, envs, ACs, iter=0, ϵ)
+    state = VUMPSState(mps, operator, envs, iter, ϵ, which)
     it = IterativeSolver(alg, state)
 
     return LoggingExtras.withlevel(; alg.verbosity) do
-        @infov 2 loginit!(log, ϵ, sum(expectation_value(ψ, H, envs)))
+        @infov 2 loginit!(log, ϵ, sum(expectation_value(mps, operator, envs)))
 
-        for (ψ, envs, ϵ) in it
+        for (mps, envs, ϵ) in it
             if ϵ ≤ alg.tol
-                @infov 2 logfinish!(log, it.iter, ϵ, expectation_value(ψ, H, envs))
-                return ψ, envs, ϵ
+                @infov 2 logfinish!(log, it.iter, ϵ, expectation_value(mps, operator, envs))
+                return mps, envs, ϵ
             end
             if it.iter ≥ alg.maxiter
-                @warnv 1 logcancel!(log, it.iter, ϵ, expectation_value(ψ, H, envs))
-                return ψ, envs, ϵ
+                @warnv 1 logcancel!(log, it.iter, ϵ, expectation_value(mps, operator, envs))
+                return mps, envs, ϵ
             end
-            @infov 3 logiter!(log, it.iter, ϵ, expectation_value(ψ, H, envs))
+            @infov 3 logiter!(log, it.iter, ϵ, expectation_value(mps, operator, envs))
         end
 
-        return it.state.ψ, it.state.envs, it.state.ϵ
+        # this should never be reached
+        return it.state.mps, it.state.envs, it.state.ϵ
     end
 end
 
 function Base.iterate(it::IterativeSolver{<:VUMPS}, state=it.state)
-    # eigsolver step
-    alg_eigsolve = updatetol(it.alg_eigsolve, state.iter, state.ϵ)
-    scheduler = Defaults.scheduler[]
-    ACs = tmap!(state.ACs, 1:length(state.ψ); scheduler) do site
-        return _vumps_localupdate(site, state.ψ, state.H, state.envs, alg_eigsolve)
-    end
-
-    # gauge step
-    alg_gauge = updatetol(it.alg_gauge, state.iter, state.ϵ)
-    ψ = InfiniteMPS(ACs, state.ψ.C[end]; alg_gauge.tol, alg_gauge.maxiter)
-
-    # environment step
-    alg_environments = updatetol(it.alg_environments, state.iter, state.ϵ)
-    envs = recalculate!(state.envs, ψ, state.H, ψ; alg_environments.tol)
+    ACs = localupdate_step!(it, state)
+    mps = gauge_step!(it, state, ACs)
+    envs = envs_step!(it, state, mps)
 
     # finalizer step
-    ψ′, envs′ = it.finalize(state.iter, ψ, state.H, envs)::Tuple{typeof(ψ),typeof(envs)}
+    mps, envs = it.finalize(state.iter, mps, state.operator, envs)::typeof((mps, envs))
 
     # error criterion
-    ϵ = calc_galerkin(ψ′, state.H, ψ′, envs′)
+    ϵ = calc_galerkin(mps, state.operator, mps, envs)
 
     # update state
-    it.state = (; ψ=ψ′, H=state.H, envs=envs′, ACs, iter=state.iter + 1, ϵ)
+    it.state = VUMPSState(mps, state.operator, envs, state.iter + 1, ϵ, state.which)
 
-    return (ψ′, envs′, ϵ), it.state
+    return (mps, envs, ϵ), it.state
 end
 
-function _vumps_localupdate(loc, ψ, H, envs, eigalg, factalg=QRpos())
-    local AC′, C′
-    if Defaults.scheduler[] isa SerialScheduler
-        _, AC′ = fixedpoint(∂∂AC(loc, ψ, H, envs), ψ.AC[loc], :SR, eigalg)
-        _, C′ = fixedpoint(∂∂C(loc, ψ, H, envs), ψ.C[loc], :SR, eigalg)
-    else
-        @sync begin
-            Threads.@spawn begin
-                _, AC′ = fixedpoint(∂∂AC(loc, ψ, H, envs), ψ.AC[loc], :SR, eigalg)
-            end
-            Threads.@spawn begin
-                _, C′ = fixedpoint(∂∂C(loc, ψ, H, envs), ψ.C[loc], :SR, eigalg)
-            end
+function localupdate_step!(it::IterativeSolver{<:VUMPS}, state,
+                           scheduler=Defaults.scheduler[])
+    alg_eigsolve = updatetol(it.alg_eigsolve, state.iter, state.ϵ)
+    alg_orth = QRpos()
+
+    mps = state.mps
+    eachsite = 1:length(mps)
+    src_Cs = mps isa Multiline ? eachcol(mps.C) : mps.C
+    src_ACs = mps isa Multiline ? eachcol(mps.AC) : mps.AC
+    ACs = similar(mps.AC)
+    dst_ACs = mps isa Multiline ? eachcol(ACs) : ACs
+
+    tforeach(eachsite, src_ACs, src_Cs; scheduler) do site, AC₀, C₀
+        dst_ACs[site] = _localupdate_vumps_step!(site, mps, state.operator, state.envs,
+                                                 AC₀, C₀; parallel=false, alg_orth,
+                                                 state.which, alg_eigsolve)
+        return nothing
+    end
+
+    return ACs
+end
+
+function _localupdate_vumps_step!(site, mps, operator, envs, AC₀, C₀;
+                                  parallel::Bool=false, alg_orth=QRpos(),
+                                  alg_eigsolve=Defaults.eigsolver, which)
+    if !parallel
+        _, AC = fixedpoint(∂∂AC(site, mps, operator, envs), AC₀, which, alg_eigsolve)
+        _, C = fixedpoint(∂∂C(site, mps, operator, envs), C₀, which, alg_eigsolve)
+        return regauge!(AC, C; alg=alg_orth)
+    end
+
+    local AC, C
+    @sync begin
+        @spawn begin
+            _, AC = fixedpoint(∂∂AC(site, mps, operator, envs),
+                               AC₀, which, alg_eigsolve)
+        end
+        @spawn begin
+            _, C = fixedpoint(∂∂C(site, mps, operator, envs),
+                              C₀, which, alg_eigsolve)
         end
     end
-    return regauge!(AC′, C′; alg=factalg)
+    return regauge!(AC, C; alg=alg_orth)
+end
+
+function gauge_step!(it::IterativeSolver{<:VUMPS}, state, ACs::AbstractVector)
+    alg_gauge = updatetol(it.alg_gauge, state.iter, state.ϵ)
+    return InfiniteMPS(ACs, state.mps.C[end]; alg_gauge.tol, alg_gauge.maxiter)
+end
+function gauge_step!(it::IterativeSolver{<:VUMPS}, state, ACs::AbstractMatrix)
+    alg_gauge = updatetol(it.alg_gauge, state.iter, state.ϵ)
+    return MultilineMPS(ACs, @view(state.mps.C[:, end]); alg_gauge.tol, alg_gauge.maxiter)
+end
+
+function envs_step!(it::IterativeSolver{<:VUMPS}, state, mps)
+    alg_environments = updatetol(it.alg_environments, state.iter, state.ϵ)
+    return recalculate!(state.envs, mps, state.operator, mps; alg_environments.tol)
 end