Refactor of IDMRG with IterativeSolver

src/algorithms/groundstate/idmrg.jl

@@ -24,81 +24,6 @@ $(TYPEDFIELDS)
     alg_eigsolve::A = Defaults.alg_eigsolve()
 end
 
-function find_groundstate(ost::InfiniteMPS, H, alg::IDMRG, envs = environments(ost, H))
-    ϵ::Float64 = calc_galerkin(ost, H, ost, envs)
-    ψ = copy(ost)
-    log = IterLog("IDMRG")
-    local iter, E_current
-
-    LoggingExtras.withlevel(; alg.verbosity) do
-        @infov 2 begin
-            E_current = expectation_value(ψ, H, envs)
-            loginit!(log, ϵ, E_current)
-        end
-        for outer iter in 1:(alg.maxiter)
-            alg_eigsolve = updatetol(alg.alg_eigsolve, iter, ϵ)
-            C_current = ψ.C[0]
-
-            # left to right sweep
-            for pos in 1:length(ψ)
-                h = AC_hamiltonian(pos, ψ, H, ψ, envs)
-                _, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve)
-                if pos == length(ψ)
-                    # AC needed in next sweep
-                    ψ.AL[pos], ψ.C[pos] = left_orth(ψ.AC[pos])
-                else
-                    ψ.AL[pos], ψ.C[pos] = left_orth!(ψ.AC[pos])
-                end
-                transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
-            end
-
-            # right to left sweep
-            for pos in length(ψ):-1:1
-                h = AC_hamiltonian(pos, ψ, H, ψ, envs)
-                _, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve)
-
-                ψ.C[pos - 1], temp = right_orth!(_transpose_tail(ψ.AC[pos]; copy = (pos == 1)))
-                ψ.AR[pos] = _transpose_front(temp)
-
-                transfer_rightenv!(envs, ψ, H, ψ, pos - 1)
-            end
-
-            ϵ = norm(C_current - ψ.C[0])
-
-            if ϵ < alg.tol
-                @infov 2 begin
-                    E_next = expectation_value(ψ, H, envs)
-                    ΔE = E_next - E_current
-                    E_current = E_next
-                    logfinish!(log, iter, ϵ, ΔE)
-                end
-                break
-            end
-            if iter == alg.maxiter
-                @warnv 1 begin
-                    E_next = expectation_value(ψ, H, envs)
-                    ΔE = E_next - E_current
-                    E_current = E_next
-                    logcancel!(log, iter, ϵ, ΔE)
-                end
-            else
-                @infov 3 begin
-                    E_next = expectation_value(ψ, H, envs)
-                    ΔE = E_next - E_current
-                    E_current = E_next
-                    logiter!(log, iter, ϵ, ΔE)
-                end
-            end
-        end
-    end
-
-    alg_gauge = updatetol(alg.alg_gauge, iter, ϵ)
-    ψ′ = InfiniteMPS(ψ.AR[1:end]; alg_gauge.tol, alg_gauge.maxiter)
-    recalculate!(envs, ψ′, H, ψ′)
-
-    return ψ′, envs, ϵ
-end
-
 """
 $(TYPEDEF)
 
@@ -131,121 +56,204 @@ $(TYPEDFIELDS)
     trscheme::TruncationStrategy
 end
 
-function find_groundstate(ost::InfiniteMPS, H, alg::IDMRG2, envs = environments(ost, H))
-    length(ost) < 2 && throw(ArgumentError("unit cell should be >= 2"))
-    ϵ::Float64 = calc_galerkin(ost, H, ost, envs)
 
-    ψ = copy(ost)
-    log = IterLog("IDMRG2")
-    local iter
+# Internal state of the IDMRG algorithm
+struct IDMRGState{S, O, E, T}
+    mps::S
+    operator::O
+    envs::E
+    iter::Int
+    ϵ::Float64
+    energy::T
+end
+function IDMRGState{T}(mps::S, operator::O, envs::E, iter::Int, ϵ::Float64, energy) where {S, O, E, T}
+    return IDMRGState{S, O, E, T}(mps, operator, envs, iter, ϵ, T(energy))
+end
+
+function find_groundstate(mps, operator, alg::alg_type, envs = environments(mps, operator)) where {alg_type <: Union{<:IDMRG, <:IDMRG2}}
+    (length(mps) ≤ 1 && alg isa IDMRG2) && throw(ArgumentError("unit cell should be >= 2"))
+    log = alg isa IDMRG ? IterLog("IDMRG") : IterLog("IDMRG2")
+    mps = copy(mps)
+    iter = 0
+    ϵ = calc_galerkin(mps, operator, mps, envs)
+    E = zero(Base.promote_type(scalartype(mps), scalartype(operator)))
 
     LoggingExtras.withlevel(; alg.verbosity) do
-        @infov 2 loginit!(log, ϵ)
-        for outer iter in 1:(alg.maxiter)
-            alg_eigsolve = updatetol(alg.alg_eigsolve, iter, ϵ)
-            C_current = ψ.C[0]
-
-            # sweep from left to right
-            for pos in 1:(length(ψ) - 1)
-                ac2 = AC2(ψ, pos; kind = :ACAR)
-                h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs)
-                _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
-
-                al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
-                normalize!(c)
-
-                ψ.AL[pos] = al
-                ψ.C[pos] = complex(c)
-                ψ.AR[pos + 1] = _transpose_front(ar)
-                ψ.AC[pos + 1] = _transpose_front(c * ar)
-
-                transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
-                transfer_rightenv!(envs, ψ, H, ψ, pos)
+        @infov 2 begin
+            E = expectation_value(mps, operator, envs)
+            loginit!(log, ϵ, E)
+        end
+    end
+
+    state = IDMRGState(mps, operator, envs, iter, ϵ, E)
+    it = IterativeSolver(alg, state)
+
+    return LoggingExtras.withlevel(; alg.verbosity) do
+        for (mps, envs, ϵ, ΔE) in it
+            if ϵ ≤ alg.tol
+                @infov 2 logfinish!(log, it.iter, ϵ, ΔE)
+                break
+            end
+            if it.iter ≥ alg.maxiter
+                @warnv 1 logcancel!(log, it.iter, ϵ, ΔE)
+                break
             end
+            @infov 3 logiter!(log, it.iter, ϵ, ΔE)
+        end
 
-            # update the edge
-            ψ.AL[end] = ψ.AC[end] / ψ.C[end]
-            ψ.AC[1] = _mul_tail(ψ.AL[1], ψ.C[1])
-            ac2 = AC2(ψ, 0; kind = :ALAC)
-            h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs)
-            _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
+        alg_gauge = updatetol(alg.alg_gauge, it.state.iter, it.state.ϵ)
+        ψ′ = InfiniteMPS(it.state.mps.AR[1:end]; alg_gauge.tol, alg_gauge.maxiter)
+        envs = recalculate!(it.state.envs, ψ′, it.state.operator, ψ′)
+        return ψ′, envs, it.state.ϵ
+    end
+end
 
-            al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
-            normalize!(c)
+function Base.iterate(
+        it::IterativeSolver{alg_type}, state::IDMRGState{<:Any, <:Any, <:Any, T} = it.state
+    ) where {alg_type <: Union{<:IDMRG, <:IDMRG2}, T}
+    mps, envs, C_old, E_new = localupdate_step!(it, state)
+
+    # error criterion
+    C = mps.C[0]
+    space_C_old = _firstspace(C_old)
+    space_C = _firstspace(C)
+    if space_C != space_C_old
+        smallest = infimum(space_C_old, space_C)
+        e1 = isometry(space_C_old, smallest)
+        e2 = isometry(space_C, smallest)
+        ϵ = norm(e2' * C * e2 - e1' * C_old * e1)
+    else
+        ϵ = norm(C - C_old)
+    end
 
-            ψ.AL[end] = al
-            ψ.C[end] = complex(c)
-            ψ.AR[1] = _transpose_front(ar)
+    # New energy
+    ΔE = (E_new - state.energy) / 2
+    (alg_type <: IDMRG2 && length(mps) == 2) && (ΔE /= 2) # This extra factor gives the correct energy per unit cell. I have no clue why right now.
 
-            ψ.AC[end] = _mul_tail(al, c)
-            ψ.AC[1] = _transpose_front(c * ar)
-            ψ.AL[1] = ψ.AC[1] / ψ.C[1]
+    # update state
+    it.state = IDMRGState{T}(mps, state.operator, envs, state.iter + 1, ϵ, E_new)
 
-            C_current = complex(c)
+    return (mps, envs, ϵ, ΔE), it.state
+end
 
-            # update environments
-            transfer_leftenv!(envs, ψ, H, ψ, 1)
-            transfer_rightenv!(envs, ψ, H, ψ, 0)
+function MPSKit.localupdate_step!(
+        it::IterativeSolver{<:Union{IDMRG, IDMRG2}}, state
+    )
+    alg_eigsolve = updatetol(it.alg_eigsolve, state.iter, state.ϵ)
+    mps, envs, C_old, E = _localupdate_sweep_idmrg!(state.mps, state.operator, state.envs, alg_eigsolve, it.alg)
 
-            # sweep from right to left
-            for pos in (length(ψ) - 1):-1:1
-                ac2 = AC2(ψ, pos; kind = :ALAC)
-                h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs)
-                _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
+    return mps, envs, C_old, E
+end
 
-                al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
-                normalize!(c)
+function _localupdate_sweep_idmrg!(ψ, H, envs, alg_eigsolve, ::IDMRG)
+    local E
+    C_old = ψ.C[0]
+    # left to right sweep
+    for pos in 1:length(ψ)
+        h = AC_hamiltonian(pos, ψ, H, ψ, envs)
+        _, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve)
+        if pos == length(ψ)
+            # AC needed in next sweep
+            ψ.AL[pos], ψ.C[pos] = left_orth(ψ.AC[pos])
+        else
+            ψ.AL[pos], ψ.C[pos] = left_orth!(ψ.AC[pos])
+        end
+        transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
+    end
 
-                ψ.AL[pos] = al
-                ψ.AC[pos] = _mul_tail(al, c)
-                ψ.C[pos] = complex(c)
-                ψ.AR[pos + 1] = _transpose_front(ar)
-                ψ.AC[pos + 1] = _transpose_front(c * ar)
+    # right to left sweep
+    for pos in length(ψ):-1:1
+        h = AC_hamiltonian(pos, ψ, H, ψ, envs)
+        E, ψ.AC[pos] = fixedpoint(h, ψ.AC[pos], :SR, alg_eigsolve)
 
-                transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
-                transfer_rightenv!(envs, ψ, H, ψ, pos)
-            end
+        ψ.C[pos - 1], temp = right_orth!(_transpose_tail(ψ.AC[pos]; copy = (pos == 1)))
+        ψ.AR[pos] = _transpose_front(temp)
 
-            # update the edge
-            ψ.AC[end] = _mul_front(ψ.C[end - 1], ψ.AR[end])
-            ψ.AR[1] = _transpose_front(ψ.C[end] \ _transpose_tail(ψ.AC[1]))
-            ac2 = AC2(ψ, 0; kind = :ACAR)
-            h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs)
-            _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
-            al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
-            normalize!(c)
-
-            ψ.AL[end] = al
-            ψ.C[end] = complex(c)
-            ψ.AR[1] = _transpose_front(ar)
-
-            ψ.AR[end] = _transpose_front(ψ.C[end - 1] \ _transpose_tail(al * c))
-            ψ.AC[1] = _transpose_front(c * ar)
-
-            transfer_leftenv!(envs, ψ, H, ψ, 1)
-            transfer_rightenv!(envs, ψ, H, ψ, 0)
-
-            # update error
-            smallest = infimum(_firstspace(C_current), _firstspace(c))
-            e1 = isometry(_firstspace(C_current), smallest)
-            e2 = isometry(_firstspace(c), smallest)
-            ϵ = norm(e2' * c * e2 - e1' * C_current * e1)
-
-            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
+        transfer_rightenv!(envs, ψ, H, ψ, pos - 1)
+    end
+    return ψ, envs, C_old, E
+end
+
+
+function _localupdate_sweep_idmrg!(ψ, H, envs, alg_eigsolve, alg::IDMRG2)
+    # sweep from left to right
+    for pos in 1:(length(ψ) - 1)
+        ac2 = AC2(ψ, pos; kind = :ACAR)
+        h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs)
+        _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
+
+        al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
+        normalize!(c)
+
+        ψ.AL[pos] = al
+        ψ.C[pos] = complex(c)
+        ψ.AR[pos + 1] = _transpose_front(ar)
+        ψ.AC[pos + 1] = _transpose_front(c * ar)
+
+        transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
+        transfer_rightenv!(envs, ψ, H, ψ, pos)
     end
 
-    alg_gauge = updatetol(alg.alg_gauge, iter, ϵ)
-    ψ′ = InfiniteMPS(ψ.AR[1:end]; alg_gauge.tol, alg_gauge.maxiter)
-    recalculate!(envs, ψ′, H, ψ′)
+    # update the edge
+    ψ.AL[end] = ψ.AC[end] / ψ.C[end]
+    ψ.AC[1] = _mul_tail(ψ.AL[1], ψ.C[1])
+    ac2 = AC2(ψ, 0; kind = :ALAC)
+    h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs)
+    _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
+
+    al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
+    normalize!(c)
+
+    ψ.AL[end] = al
+    ψ.C[end] = complex(c)
+    ψ.AR[1] = _transpose_front(ar)
+
+    ψ.AC[end] = _mul_tail(al, c)
+    ψ.AC[1] = _transpose_front(c * ar)
+    ψ.AL[1] = ψ.AC[1] / ψ.C[1]
+
+    C_old = complex(c)
+
+    # update environments
+    transfer_leftenv!(envs, ψ, H, ψ, 1)
+    transfer_rightenv!(envs, ψ, H, ψ, 0)
+
+    # sweep from right to left
+    for pos in (length(ψ) - 1):-1:1
+        ac2 = AC2(ψ, pos; kind = :ALAC)
+        h_ac2 = AC2_hamiltonian(pos, ψ, H, ψ, envs)
+        _, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
+
+        al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
+        normalize!(c)
+
+        ψ.AL[pos] = al
+        ψ.AC[pos] = _mul_tail(al, c)
+        ψ.C[pos] = complex(c)
+        ψ.AR[pos + 1] = _transpose_front(ar)
+        ψ.AC[pos + 1] = _transpose_front(c * ar)
+
+        transfer_leftenv!(envs, ψ, H, ψ, pos + 1)
+        transfer_rightenv!(envs, ψ, H, ψ, pos)
+    end
 
-    return ψ′, envs, ϵ
+    # update the edge
+    ψ.AC[end] = _mul_front(ψ.C[end - 1], ψ.AR[end])
+    ψ.AR[1] = _transpose_front(ψ.C[end] \ _transpose_tail(ψ.AC[1]))
+    ac2 = AC2(ψ, 0; kind = :ACAR)
+    h_ac2 = AC2_hamiltonian(0, ψ, H, ψ, envs)
+    E, ac2′ = fixedpoint(h_ac2, ac2, :SR, alg_eigsolve)
+    al, c, ar = svd_trunc!(ac2′; trunc = alg.trscheme, alg = alg.alg_svd)
+    normalize!(c)
+
+    ψ.AL[end] = al
+    ψ.C[end] = complex(c)
+    ψ.AR[1] = _transpose_front(ar)
+
+    ψ.AR[end] = _transpose_front(ψ.C[end - 1] \ _transpose_tail(al * c))
+    ψ.AC[1] = _transpose_front(c * ar)
+
+    transfer_leftenv!(envs, ψ, H, ψ, 1)
+    transfer_rightenv!(envs, ψ, H, ψ, 0)
+    return ψ, envs, C_old, E
 end

src/utility/iterativesolvers.jl

@@ -7,7 +7,7 @@ mutable struct IterativeSolver{A, B}
 end
 
 function Base.getproperty(it::IterativeSolver{A, B}, name::Symbol) where {A, B}
-    name === :alg || name === :state && return getfield(it, name)
+    (name === :alg || name === :state) && return getfield(it, name)
 
     alg = getfield(it, :alg)
     name in propertynames(alg) && return getproperty(alg, name)