feat: standalone linear_solve primitives

docs/src/changelog.md

@@ -29,6 +29,9 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
 - `changebonds(::InfiniteMPOHamiltonian, ::SvdCut)` compresses infinite Hamiltonians with the
   almost-Schmidt decomposition of [Parker et al.](https://arxiv.org/abs/1909.06341), and
   `left_canonicalize!`/`right_canonicalize!` bring an `InfiniteMPOHamiltonian` into canonical form. ([#539](https://github.com/QuantumKitHub/MPSKit.jl/pull/539))
+- `linear_solve`/`linear_solve!` for solving `(a₀ + a₁·A)·x = b` with finite MPS `x` and `b`, using the
+  DMRG-style sweeps `DMRGSolve` and `DMRGSolve2`; see [`linear_solve`](@ref) for details.
+- `DynamicalDMRG` gained a `trunc` field for a bond-adaptive `propagator`; see [`DynamicalDMRG`](@ref).
 
 ### Changed
 
@@ -46,6 +49,8 @@ When releasing a new version, move the "Unreleased" changes to a new version sec
   its allocator across steps, rather than copying the state at every step. A `finalize` callback
   that keeps the state of an earlier step therefore sees it updated by later steps, as with
   `time_evolve!`. ([#537](https://github.com/QuantumKitHub/MPSKit.jl/pull/537))
+- `propagator`/`DynamicalDMRG` are now implemented on top of `linear_solve`, which changes the default
+  convergence criterion, `tol` and local solver; `Jeckelmann` now throws for real `z`.
 
 ### Deprecated
 

docs/src/man/algorithms.md

@@ -444,16 +444,38 @@ state, by a new state.
 approximate
 ```
 
+## `linear_solve`
+
+Solving a linear system ``(a₀ + a₁ A) x = b`` for an MPS ``x``, with ``A`` an MPO and ``b`` an MPS.
+The shift ``a₀ + a₁ A`` follows the convention of `KrylovKit.linsolve` and is applied implicitly, so
+no shifted operator is ever formed.
+The sweeps are DMRG-like: [`DMRGSolve`](@ref) keeps the bond dimension of the initial guess fixed,
+while [`DMRGSolve2`](@ref) updates two sites at a time and is therefore rank-adaptive.
+Each is parameterized by how the local subproblem is posed — [`Galerkin`](@ref) solves the effective
+system directly, [`LeastSquares`](@ref) minimizes ``\|(a₀ + a₁ A) x - b\|²`` through the normal
+equations.
+
+The resolvent is the canonical application: ``\frac{1}{z - H}|ψ₀⟩`` is `linear_solve(ψ₀, H, ψ₀; a₀ = z, a₁ = -1)`.
+
+```@docs; canonical=false
+linear_solve
+linear_solve!
+DMRGSolve
+DMRGSolve2
+Galerkin
+LeastSquares
+```
+
 ## Varia
 
 What follows is a medley of lesser known (or used) algorithms and don't entirely fit under
 one of the above categories.
 
 ### Dynamical DMRG
 
-Dynamical DMRG has been described in other papers and is a way to find the propagator. The
-basic idea is that to calculate ``G(z) = ⟨ V | (H-z)^{-1} | V ⟩ `` , one can variationally
-find ``(H-z) |W ⟩ = | V ⟩ `` and then the propagator simply equals ``G(z) = ⟨ V | W ⟩``.
+Dynamical DMRG has been described in other papers and is a way to find the propagator.
+The basic idea is that to calculate ``G(z) = ⟨ V | (H-z)^{-1} | V ⟩ `` , one can variationally find ``(H-z) |W ⟩ = | V ⟩ `` and then the propagator simply equals ``G(z) = ⟨ V | W ⟩``.
+This is a thin wrapper around [`linear_solve`](@ref), see there for the sweep, the convergence criterion and the local solver options.
 
 ```@docs; canonical=false
 propagator

src/MPSKit.jl

@@ -43,6 +43,8 @@ export NoiseSchedule, FunctionalSchedule, ExponentialDecay, Warmup, DMRG3S
 export Zipup
 export propagator
 export DynamicalDMRG, NaiveInvert, Jeckelmann
+export linear_solve, linear_solve!
+export DMRGSolve, DMRGSolve2, Galerkin, LeastSquares
 export exact_diagonalization, fidelity_susceptibility
 
 # toolbox:
@@ -209,6 +211,9 @@ include("algorithms/ED.jl")
 
 include("algorithms/unionalg.jl")
 
+include("algorithms/linear_solve/linear_solve.jl")
+include("algorithms/linear_solve/flinear_solve.jl")
+
 include("utility/show.jl")
 
 function __init__()

src/algorithms/approximate/fvomps.jl

@@ -38,20 +38,18 @@ end
 
 function Base.iterate(it::IterativeSolver{<:Union{DMRG, DMRG2}}, state::ApproximateState)
     iter = state.iter + 1
-    state = ApproximateState(
-        state.mps, state.operator, state.envs, state.iter, 0.0, state.allocator
-    )
+    state = setproperties(state; ϵ = 0.0)
     state = sweep!(it, state, Val(:right), iter)
     state = sweep!(it, state, Val(:left), iter)
     ψ, envs = it.finalize(
         iter, state.mps, state.operator, state.envs
     )::Tuple{typeof(state.mps), typeof(state.envs)}
-    it.state = ApproximateState(ψ, state.operator, envs, iter, state.ϵ, state.allocator)
+    it.state = setproperties(state; mps = ψ, envs, iter)
     return (ψ, envs, state.ϵ), it.state
 end
 
 function sweep!(it::IterativeSolver{<:DMRG2}, state::ApproximateState, direction, iter)
-    fwd, bwd = _sweep_ranges(it.alg, state.mps)
+    fwd, bwd = sweep_ranges(it.alg, state.mps)
     sites = direction === Val(:right) ? fwd : bwd
     ψ, Oϕ, envs, ϵ = state.mps, state.operator, state.envs, state.ϵ
     for pos in sites
@@ -64,11 +62,11 @@ function sweep!(it::IterativeSolver{<:DMRG2}, state::ApproximateState, direction
         ψ.AC[pos] = (al, complex(c))
         ψ.AC[pos + 1] = (complex(c), _transpose_front(ar))
     end
-    return ApproximateState(ψ, Oϕ, envs, state.iter, ϵ, state.allocator)
+    return setproperties(state; mps = ψ, ϵ)
 end
 
 function sweep!(it::IterativeSolver{<:DMRG}, state::ApproximateState, direction, iter)
-    fwd, bwd = _sweep_ranges(it.alg, state.mps)
+    fwd, bwd = sweep_ranges(it.alg, state.mps)
     sites = direction === Val(:right) ? fwd : bwd
     ψ, Oϕ, envs, ϵ = state.mps, state.operator, state.envs, state.ϵ
     for pos in sites
@@ -78,5 +76,5 @@ function sweep!(it::IterativeSolver{<:DMRG}, state::ApproximateState, direction,
 
         ψ.AC[pos] = AC′
     end
-    return ApproximateState(ψ, Oϕ, envs, state.iter, ϵ, state.allocator)
+    return setproperties(state; mps = ψ, ϵ)
 end

src/algorithms/approximate/idmrg.jl

@@ -46,10 +46,7 @@ function Base.iterate(it::IterativeSolver{<:Union{IDMRG, IDMRG2}}, state::IDMRGS
     state = sweep!(it, state, Val(:right), iter)
     state = sweep!(it, state, Val(:left), iter)
     ϵ = bond_change(C_current, state.mps.C[:, 0])
-    it.state = IDMRGState(
-        state.mps, state.operator, state.envs, iter, ϵ, state.truncation_errors, state.energy,
-        state.timeroutput, state.allocator,
-    )
+    it.state = setproperties(state; iter, ϵ)
     return (it.state.mps, it.state.envs, it.state.ϵ), it.state
 end
 

src/algorithms/derivatives/mpo_derivatives.jl

@@ -138,7 +138,7 @@ function (h::MPO_AC2_Hamiltonian{<:MPSBondTensor, Nothing, Nothing, <:MPSBondTen
     )
     backend, allocator = h.backend, h.allocator
     @plansor backend = backend allocator = allocator begin
-        y[-1 -2; -3 -4] ≔ h.leftenv[-1; 1] * x[1 -2; 2 -4] * h.rightenv[2 -3]
+        y[-1 -2; -3 -4] ≔ h.leftenv[-1; 1] * x[1 -2; 2 -4] * h.rightenv[2; -3]
     end
     return y isa AbstractBlockTensorMap ? only(y) : y
 end

src/algorithms/groundstate/dmrg.jl

@@ -243,20 +243,11 @@ function local_update!(
     return ψ, ϵ_local, ϵ_trunc, decay_rate
 end
 
-# Per-algorithm sweep geometry: single-site DMRG updates all `length(ψ)` sites (endpoints once,
-# interior twice), whereas two-site DMRG2 updates the `length(ψ) - 1` bonds. `_num_updates`
-# gives the number of per-update bookkeeping slots and `_sweep_ranges` the forward/backward
-# index ranges; everything else in the sweep is shared.
-_num_updates(::DMRG, ψ) = length(ψ)
-_num_updates(::DMRG2, ψ) = length(ψ) - 1
+sweep_sites(::DMRG) = 1
+sweep_sites(::DMRG2) = 2
 
-_sweep_ranges(::DMRG, ψ) = (1:(length(ψ) - 1), length(ψ):-1:2)
-_sweep_ranges(::DMRG2, ψ) = (1:(length(ψ) - 1), (length(ψ) - 2):-1:1)
-
-# per-bond truncation errors of the last cut from the per-update slots: single-site slot `pos`
-# holds the backward cut of bond `pos - 1`, which is the last one made there
-_bond_truncation_errors(::DMRG, ϵ_truncs) = ϵ_truncs[2:end]
-_bond_truncation_errors(::DMRG2, ϵ_truncs) = copy(ϵ_truncs)
+# single-site slot `pos` holds the backward cut of bond `pos - 1`, which is the last one made there
+bond_truncation_errors(alg, ϵ_truncs) = sweep_sites(alg) == 1 ? ϵ_truncs[2:end] : copy(ϵ_truncs)
 
 inner_alg_gauge(alg::Union{DMRG, DMRG2}) = alg_gauge(alg.alg_gauge)
 
@@ -310,7 +301,7 @@ end
 
 function DMRGState(ψ, H, alg::Union{DMRG, DMRG2}, envs, allocator, timeroutput)
     Tr = real(scalartype(ψ))
-    n = _num_updates(alg, ψ)
+    n = num_updates(alg, ψ)
     local_errors = ones(Tr, n)
     return DMRGState(
         ψ, H, envs, 0, maximum(local_errors), local_errors, zeros(Tr, n),
@@ -319,7 +310,7 @@ function DMRGState(ψ, H, alg::Union{DMRG, DMRG2}, envs, allocator, timeroutput)
 end
 
 function sweep!(it::IterativeSolver{<:Union{DMRG, DMRG2}}, state::DMRGState, direction, iter)
-    fwd, bwd = _sweep_ranges(it.alg, state.mps)
+    fwd, bwd = sweep_ranges(it.alg, state.mps)
     sites = direction === Val(:right) ? fwd : bwd
     ψ, ϵ = state.mps, state.ϵ
     for pos in sites
@@ -332,10 +323,7 @@ function sweep!(it::IterativeSolver{<:Union{DMRG, DMRG2}}, state::DMRGState, dir
         # The next local solver uses the errors accumulated so far in this sweep.
         ϵ = maximum(state.local_errors)
     end
-    return DMRGState(
-        ψ, state.operator, state.envs, state.iter, ϵ, state.local_errors,
-        state.truncation_errors, state.decay_rates, state.timeroutput, state.allocator,
-    )
+    return setproperties(state; mps = ψ, ϵ)
 end
 
 function Base.iterate(it::IterativeSolver{<:Union{DMRG, DMRG2}}, state::DMRGState)
@@ -348,10 +336,7 @@ function Base.iterate(it::IterativeSolver{<:Union{DMRG, DMRG2}}, state::DMRGStat
     ψ, envs = @timeit timeroutput "finalize" it.finalize(
         iter, state.mps, state.operator, state.envs
     )::Tuple{typeof(state.mps), typeof(state.envs)}
-    it.state = DMRGState(
-        ψ, state.operator, envs, iter, state.ϵ, state.local_errors,
-        state.truncation_errors, state.decay_rates, timeroutput, state.allocator,
-    )
+    it.state = setproperties(state; mps = ψ, envs, iter)
     return (ψ, envs, state.ϵ), it.state
 end
 
@@ -384,7 +369,7 @@ function find_groundstate_sweep!(
     state = it.state
     info = AlgorithmInfo(;
         converged = sweep_converged(alg, state), galerkin = state.ϵ,
-        truncation_errors = _bond_truncation_errors(alg, state.truncation_errors),
+        truncation_errors = bond_truncation_errors(alg, state.truncation_errors),
         numiter = state.iter,
     )
     return state.mps, state.envs, info

src/algorithms/linear_solve/flinear_solve.jl

@@ -0,0 +1,166 @@
+# Finite-MPS sweeps for `linear_solve`. Convergence is measured, as in `find_groundstate!`, by the
+# local residual of the original system `(a₀ + a₁·A)·x = b`, evaluated before each local solve.
+
+sweep_sites(::DMRGSolve) = 1
+sweep_sites(::DMRGSolve2) = 2
+
+# environments beyond `⟨x|A|x⟩`: `LeastSquares` also needs `A²` and the mixed sandwich `⟨x|A|b⟩`
+linear_solve_environments(::Galerkin, x, A, b, openvs) = (; openvs, rhsenvs = environments(x, b))
+function linear_solve_environments(::LeastSquares, x, A, b, openvs)
+    Asq, sqenvs = squaredenvs(x, A, openvs)
+    return (; openvs, rhsenvs = environments(x, b), Asq, sqenvs, abenvs = environments(x, A, b))
+end
+
+# `⟨∂x|A|b⟩` through the generic MPO derivative, since the JordanMPO one assumes `below === above`
+function mixed_projection(alg, pos, x, A, b, envs; backend, allocator)
+    W = A isa WindowMPOHamiltonian ? A.finite_ham : A
+    GL = leftenv(envs, pos, x; backend, allocator)
+    if sweep_sites(alg) == 1
+        GR = rightenv(envs, pos, x; backend, allocator)
+        return MPO_AC_Hamiltonian(GL, W[pos], GR, backend, allocator) * b.AC[pos]
+    else
+        GR = rightenv(envs, pos + 1, x; backend, allocator)
+        return MPO_AC2_Hamiltonian(GL, W[pos], W[pos + 1], GR, backend, allocator) * AC2(b, pos)
+    end
+end
+
+function local_hamiltonian(alg, pos, x, A, envs; kwargs...)
+    return sweep_sites(alg) == 1 ? AC_hamiltonian(pos, x, A, x, envs; kwargs...) :
+        AC2_hamiltonian(pos, x, A, x, envs; kwargs...)
+end
+
+# the local system `(α₀ + α₁·M)·v = rhs` actually handed to the Krylov solver
+local_linear_system(::Galerkin, alg, pos, x, A, b, ctx, A_eff, b_eff, a₀, a₁; kwargs...) =
+    (A_eff, b_eff, a₀, a₁)
+function local_linear_system(::LeastSquares, alg, pos, x, A, b, ctx, A_eff, b_eff, a₀, a₁; kwargs...)
+    Asq_eff = local_hamiltonian(alg, pos, x, ctx.Asq, ctx.sqenvs; kwargs...)
+    N_eff = LinearCombination((A_eff, Asq_eff), (2 * real(conj(a₀) * a₁), abs2(a₁)))
+    Ab_eff = mixed_projection(alg, pos, x, A, b, ctx.abenvs; kwargs...)
+    return N_eff, conj(a₀) * b_eff + conj(a₁) * Ab_eff, abs2(a₀), one(a₁)
+end
+
+# returns the updated local tensor and the original-system residual before the update
+function local_linear_solve(alg, pos, x, state; kind = :ACAR, adapt_kwargs...)
+    A, b, ctx, (a₀, a₁) = state.operator, state.rhs, state.context, state.shift
+    backend, allocator = alg.backend, state.allocator
+    onesite = sweep_sites(alg) == 1
+    A_eff = local_hamiltonian(alg, pos, x, A, ctx.openvs; backend, allocator)
+    b_eff = onesite ? AC_projection(pos, x, b, ctx.rhsenvs; backend, allocator) :
+        AC2_projection(pos, x, b, ctx.rhsenvs; backend, allocator)
+    v = onesite ? x.AC[pos] : AC2(x, pos; kind)
+    # `‖(a₀ + a₁·A)·x − b‖`, reusing the buffer of `A·x`
+    res = norm(add!!(add!!(A_eff * v, v, a₀, a₁), b_eff, -1))
+    M, rhs, α₀, α₁ = local_linear_system(
+        alg.formulation, alg, pos, x, A, b, ctx, A_eff, b_eff, a₀, a₁; backend, allocator
+    )
+    v′, info = KrylovKit.linsolve(M, rhs, v, adapt_solver(alg.solver; adapt_kwargs...), α₀, α₁)
+    # only warn when the local solve also misses the accuracy the sweep needs
+    info.converged == 0 && info.normres > alg.tol * state.normb &&
+        @warn "linear_solve: local solve at $pos did not converge (normres = $(info.normres))"
+    return v′, res
+end
+
+function local_update!(pos, direction, x, alg::DMRGSolve, state, g_global, iter)
+    AC′, res = local_linear_solve(alg, pos, x, state; iter, g_global)
+    x.AC[pos] = AC′
+    return x, res, 0.0
+end
+function local_update!(pos, direction, x, alg::DMRGSolve2, state, g_global, iter)
+    kind = direction === Val(:right) ? :ACAR : :ALAC
+    eps_trunc = state.truncation_errors[pos]
+    AC2′, res = local_linear_solve(alg, pos, x, state; kind, iter, g_global, eps_trunc)
+    x, ϵ_trunc = gauge2!(x, pos, direction, AC2′, alg.alg_gauge; normalize = false)
+    return x, res, ϵ_trunc
+end
+
+# Internal state of the finite linear-solve sweeps, where `ϵ` is the largest local residual
+# relative to `‖b‖`
+struct LinearSolveState{S, O, B, E, C, T, A}
+    mps::S
+    operator::O
+    rhs::B
+    envs::E
+    context::C
+    shift::T
+    normb::Float64
+    iter::Int
+    ϵ::Float64
+    residuals::Vector{Float64}
+    truncation_errors::Vector{Float64}
+    allocator::A
+end
+
+function LinearSolveState(x, A, b, alg::Union{DMRGSolve, DMRGSolve2}, envs, a₀, a₁)
+    n = num_updates(alg, x)
+    context = linear_solve_environments(alg.formulation, x, A, b, envs)
+    allocator = default_allocator(x, SerialScheduler())
+    # guard the relative residual against a zero right-hand side
+    normb = norm(b)
+    iszero(normb) && (normb = one(normb))
+    return LinearSolveState(
+        x, A, b, envs, context, (a₀, a₁), normb, 0, 1.0, ones(n), zeros(n), allocator
+    )
+end
+
+function sweep!(
+        it::IterativeSolver{<:Union{DMRGSolve, DMRGSolve2}}, state::LinearSolveState,
+        direction, iter
+    )
+    fwd, bwd = sweep_ranges(it.alg, state.mps)
+    x = state.mps
+    for pos in (direction === Val(:right) ? fwd : bwd)
+        # the next local solver uses the residuals accumulated so far, as in `DMRG`
+        g_global = maximum(state.residuals) * state.normb
+        x, res, state.truncation_errors[pos] = local_update!(
+            pos, direction, x, it.alg, state, g_global, iter
+        )
+        state.residuals[pos] = res / state.normb
+    end
+    return setproperties(state; mps = x, ϵ = maximum(state.residuals))
+end
+
+function Base.iterate(
+        it::IterativeSolver{<:Union{DMRGSolve, DMRGSolve2}}, state::LinearSolveState
+    )
+    iter = state.iter + 1
+    state = sweep!(it, state, Val(:right), iter)
+    state = sweep!(it, state, Val(:left), iter)
+    x, envs = it.finalize(
+        iter, state.mps, state.operator, state.envs
+    )::Tuple{typeof(state.mps), typeof(state.envs)}
+    it.state = setproperties(state; mps = x, envs, iter)
+    return (x, envs, state.ϵ), it.state
+end
+
+# the residual cannot drop below the level set by the discarded weight
+sweep_converged(alg, state::LinearSolveState) =
+    state.ϵ <= max(alg.tol, maximum(state.truncation_errors) / state.normb)
+
+function linear_solve!(
+        x::AbstractFiniteMPS, A, b, alg::Union{DMRGSolve, DMRGSolve2},
+        envs = environments(x, A, x); a₀ = 0, a₁ = 1
+    )
+    log = IterLog(alg)
+    it = IterativeSolver(alg, LinearSolveState(x, A, b, alg, envs, a₀, a₁))
+
+    with_verbosity(; alg.verbosity) do
+        @log_initialization loginit!(log, it.ϵ)
+        for (_, _, ϵ) in Iterators.take(it, alg.maxiter)
+            if sweep_converged(alg, it.state)
+                @log_convergence logfinish!(log, it.iter, ϵ)
+                break
+            elseif it.iter == alg.maxiter
+                @log_nonconvergence logcancel!(log, it.iter, ϵ)
+            else
+                @log_iteration logiter!(log, it.iter, ϵ)
+            end
+        end
+    end
+
+    state = it.state
+    info = AlgorithmInfo(;
+        converged = sweep_converged(alg, state), galerkin = state.ϵ, numiter = state.iter,
+        truncation_errors = alg isa DMRGSolve2 ? copy(state.truncation_errors) : nothing,
+    )
+    return state.mps, state.envs, info
+end