Add EigenvalueProblem support with dense, Arpack, ArnoldiMethod, KrylovKit, and JacobiDavidson backends

Project.toml

@@ -33,6 +33,8 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"
+Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
+ArnoldiMethod = "ec485272-7323-5ecc-a04f-4719b315124d"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 cuSOLVER = "887afef0-6a32-4de5-add4-7827692ba8fc"
@@ -49,6 +51,7 @@ Ginkgo = "4c8bd3c9-ead9-4b5e-a625-08f1338ba0ec"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 HSL = "34c5aeac-e683-54a6-a0e9-6e0fdc586c50"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
+JacobiDavidson = "11c68b98-9c9b-11e8-267b-bbb95576cead"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LAPACK_jll = "51474c39-65e3-53ba-86ba-03b1b862ec14"
@@ -72,6 +75,8 @@ blis_jll = "6136c539-28a5-5bf0-87cc-b183200dce32"
 [extensions]
 LinearSolveAMDGPUExt = "AMDGPU"
 LinearSolveAlgebraicMultigridExt = "AlgebraicMultigrid"
+LinearSolveArpackExt = "Arpack"
+LinearSolveArnoldiMethodExt = "ArnoldiMethod"
 LinearSolveBLISExt = ["blis_jll", "LAPACK_jll"]
 LinearSolveBandedMatricesExt = "BandedMatrices"
 LinearSolveBlockDiagonalsExt = "BlockDiagonals"
@@ -89,6 +94,7 @@ LinearSolveGinkgoExt = ["Ginkgo", "SparseArrays"]
 LinearSolveHYPREExt = "HYPRE"
 LinearSolveHSLExt = ["HSL", "SparseArrays"]
 LinearSolveIterativeSolversExt = "IterativeSolvers"
+LinearSolveJacobiDavidsonExt = "JacobiDavidson"
 LinearSolveKernelAbstractionsExt = "KernelAbstractions"
 LinearSolveKrylovKitExt = "KrylovKit"
 LinearSolveMetalExt = "Metal"
@@ -115,6 +121,8 @@ AMDGPU = "1.2, 2"
 # the ext API is identical and the AMG testset passes on 0.5.1.
 AlgebraicMultigrid = "0.5, 1, 2"
 ArrayInterface = "7.19"
+ArnoldiMethod = "0.4"
+Arpack = "0.5"
 BandedMatrices = "1.8"
 BlockDiagonals = "0.2"
 cuSOLVER = "6"
@@ -139,6 +147,7 @@ HYPRE = "1.7"
 HSL = "0.5"
 InteractiveUtils = "1.10"
 IterativeSolvers = "0.9.4"
+JacobiDavidson = "0.1"
 KernelAbstractions = "0.9.30"
 Krylov = "0.10"
 KrylovKit = "0.10"
@@ -189,6 +198,8 @@ julia = "1.10"
 
 [extras]
 AlgebraicMultigrid = "2169fc97-5a83-5252-b627-83903c6c433c"
+Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
+ArnoldiMethod = "ec485272-7323-5ecc-a04f-4719b315124d"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 CliqueTrees = "60701a23-6482-424a-84db-faee86b9b1f8"
@@ -200,6 +211,7 @@ FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
+JacobiDavidson = "11c68b98-9c9b-11e8-267b-bbb95576cead"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 KrylovPreconditioners = "45d422c2-293f-44ce-8315-2cb988662dec"
 MultiFloats = "bdf0d083-296b-4888-a5b6-7498122e68a5"
@@ -218,4 +230,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["AlgebraicMultigrid", "BandedMatrices", "BlockDiagonals", "CliqueTrees", "ComponentArrays", "FastAlmostBandedMatrices", "FastLapackInterface", "FiniteDiff", "FixedSizeArrays", "ForwardDiff", "InteractiveUtils", "IterativeSolvers", "KrylovKit", "KrylovPreconditioners", "MultiFloats", "Pkg", "PureUMFPACK", "Random", "RecursiveFactorization", "STRUMPACK_jll", "SafeTestsets", "SciMLTesting", "SparseArrays", "Sparspak", "SpecializingFactorizations", "StaticArrays", "Test", "Zygote"]
+test = ["AlgebraicMultigrid", "Arpack", "ArnoldiMethod", "BandedMatrices", "BlockDiagonals", "CliqueTrees", "ComponentArrays", "FastAlmostBandedMatrices", "FastLapackInterface", "FiniteDiff", "FixedSizeArrays", "ForwardDiff", "InteractiveUtils", "IterativeSolvers", "JacobiDavidson", "KrylovKit", "KrylovPreconditioners", "MultiFloats", "Pkg", "PureUMFPACK", "Random", "RecursiveFactorization", "STRUMPACK_jll", "SafeTestsets", "SciMLTesting", "SparseArrays", "Sparspak", "SpecializingFactorizations", "StaticArrays", "Test", "Zygote"]

ext/LinearSolveArnoldiMethodExt.jl

@@ -0,0 +1,66 @@
+module LinearSolveArnoldiMethodExt
+
+using LinearAlgebra
+using LinearSolve
+import ArnoldiMethod: partialschur, partialeigen, LM, LR, SR, LI, SI
+using SciMLBase: SciMLBase, ReturnCode
+
+function SciMLBase.solve(
+        prob::LinearSolve.EigenvalueProblem,
+        alg::LinearSolve.ArnoldiMethodJL,
+        args...; kwargs...
+    )
+    prob.B === nothing ||
+        error("ArnoldiMethod backend currently supports standard eigenvalue problems only.")
+    nev = LinearSolve.default_num_eigenpairs(prob)
+    which = prob.shift === nothing ? _arnoldi_target(prob.eigentarget) : LM()
+    A = prob.shift === nothing ? prob.A : _shift_invert_operator(prob.A, prob.shift)
+    kw = (; nev, which, prob.kwargs..., alg.kwargs..., kwargs...)
+    decomp, history = partialschur(A; kw...)
+    values, vectors = partialeigen(decomp)
+    if prob.shift !== nothing
+        values = prob.shift .+ inv.(values)
+    end
+    values, vectors = LinearSolve._select_eigenpairs(
+        values, vectors, nev, prob.eigentarget, prob.shift
+    )
+    retcode = history.converged ? ReturnCode.Success : ReturnCode.ConvergenceFailure
+    return LinearSolve.build_eigenvalue_solution(
+        prob, alg, values, vectors; retcode, stats = history
+    )
+end
+
+# ArnoldiMethod exposes its own `Target` types (preferred over ARPACK-style
+# symbols per its own documentation) for all but smallest-magnitude, which it
+# does not support at all.
+function _arnoldi_target(w::LinearSolve.EigenvalueTarget.T)
+    T = LinearSolve.EigenvalueTarget
+    return w == T.LargestMagnitude ? LM() :
+        w == T.LargestRealPart ? LR() :
+        w == T.SmallestRealPart ? SR() :
+        w == T.LargestImaginaryPart ? LI() :
+        w == T.SmallestImaginaryPart ? SI() :
+        throw(ArgumentError("ArnoldiMethod does not support `eigentarget = EigenvalueTarget.SmallestMagnitude`; use a different backend (e.g. `KrylovKitEigen()` or `ArpackJL()`) or supply `shift` for shift-and-invert."))
+end
+
+function _shift_invert_operator(A, shift)
+    F = factorize(A - shift * I)
+    T = promote_type(eltype(A), typeof(shift))
+    return ShiftInvertMap{typeof(F), T}(F, size(A, 1))
+end
+
+struct ShiftInvertMap{F, T}
+    F::F
+    n::Int
+end
+
+Base.size(A::ShiftInvertMap) = (A.n, A.n)
+Base.size(A::ShiftInvertMap, dim::Integer) = dim <= 2 ? A.n : 1
+Base.eltype(::Type{<:ShiftInvertMap{F, T}}) where {F, T} = T
+
+function LinearAlgebra.mul!(y, A::ShiftInvertMap, x)
+    copyto!(y, A.F \ x)
+    return y
+end
+
+end

ext/LinearSolveArpackExt.jl

@@ -0,0 +1,45 @@
+module LinearSolveArpackExt
+
+using LinearSolve
+using Arpack
+using SciMLBase: SciMLBase, ReturnCode
+
+function SciMLBase.solve(
+        prob::LinearSolve.EigenvalueProblem,
+        alg::LinearSolve.ArpackJL,
+        args...; kwargs...
+    )
+    nev = LinearSolve.default_num_eigenpairs(prob)
+    base = (; nev, which = _arpack_which(prob.eigentarget))
+    if prob.shift !== nothing
+        base = (; base..., sigma = prob.shift)
+    end
+    kw = (; base..., prob.kwargs..., alg.kwargs..., kwargs...)
+    # `Arpack.eigs` takes the generalized-problem matrix `B` positionally, not
+    # as a keyword argument.
+    values, vectors, nconv, niter, nmult, resid = if prob.B === nothing
+        Arpack.eigs(prob.A; kw...)
+    else
+        Arpack.eigs(prob.A, prob.B; kw...)
+    end
+    retcode = nconv >= length(values) ? ReturnCode.Success : ReturnCode.ConvergenceFailure
+    stats = (; nconv, niter, nmult)
+    return LinearSolve.build_eigenvalue_solution(
+        prob, alg, values, vectors; retcode, resid, stats
+    )
+end
+
+# Arpack.eigs requires a raw ARPACK-style Symbol for `which`; this mapping is
+# purely a private adapter to that third-party API, not a general LinearSolve
+# concept.
+function _arpack_which(w::LinearSolve.EigenvalueTarget.T)
+    T = LinearSolve.EigenvalueTarget
+    return w == T.LargestMagnitude ? :LM :
+        w == T.SmallestMagnitude ? :SM :
+        w == T.LargestRealPart ? :LR :
+        w == T.SmallestRealPart ? :SR :
+        w == T.LargestImaginaryPart ? :LI :
+        :SI
+end
+
+end

ext/LinearSolveJacobiDavidsonExt.jl

@@ -0,0 +1,68 @@
+module LinearSolveJacobiDavidsonExt
+
+using LinearSolve
+using LinearAlgebra
+using JacobiDavidson
+using SciMLBase: SciMLBase, ReturnCode
+
+function SciMLBase.solve(
+        prob::LinearSolve.EigenvalueProblem,
+        alg::LinearSolve.JacobiDavidsonJL,
+        args...; kwargs...
+    )
+    # JacobiDavidson.jl's `jdqz` (generalized solver) is broken upstream: it
+    # references an undefined `verbose` variable that is absent from its
+    # signature. Until that is fixed, only the standard solver `jdqr` is wired
+    # up here; point users at the backends that do support generalized problems.
+    prob.B === nothing ||
+        error("The JacobiDavidson backend currently supports standard eigenvalue problems only. Use `ArpackJL()` or `KrylovKitEigen()` for generalized problems.")
+
+    n = size(prob.A, 2)
+    nev = LinearSolve.default_num_eigenpairs(prob)
+    target = _jd_target(prob)
+    # Search-subspace bounds, capped at the problem size. Users may override
+    # `subspace_dimensions` (and any other jdqr keyword) via the algorithm.
+    hi = min(max(2 * nev + 10, 20), n)
+    lo = min(max(nev + 2, 8), hi)
+    defaults = (; pairs = nev, target = target, subspace_dimensions = lo:hi)
+    kw = (; defaults..., prob.kwargs..., alg.kwargs..., kwargs...)
+
+    out = JacobiDavidson.jdqr(prob.A; kw...)
+    values, vectors = _jd_standard_pairs(prob.A, out[1])
+
+    values, vectors = LinearSolve._select_eigenpairs(
+        values, vectors, nev, prob.eigentarget, prob.shift
+    )
+    retcode = length(values) >= nev ? ReturnCode.Success : ReturnCode.ConvergenceFailure
+    return LinearSolve.build_eigenvalue_solution(
+        prob, alg, values, vectors; retcode, stats = out[end]
+    )
+end
+
+# Map the problem's spectral selector onto a JacobiDavidson `Target`. A supplied
+# `shift` is the natural interior target (`Near`), which is Jacobi-Davidson's
+# strength; otherwise `eigentarget` selects an extremal target.
+function _jd_target(prob)
+    prob.shift !== nothing && return JacobiDavidson.Near(ComplexF64(prob.shift))
+    T = LinearSolve.EigenvalueTarget
+    w = prob.eigentarget
+    return w == T.LargestMagnitude ? JacobiDavidson.LargestMagnitude(0.0 + 0.0im) :
+        w == T.SmallestMagnitude ? JacobiDavidson.SmallestMagnitude(0.0 + 0.0im) :
+        w == T.LargestRealPart ? JacobiDavidson.LargestRealPart(0.0 + 0.0im) :
+        w == T.SmallestRealPart ? JacobiDavidson.SmallestRealPart(0.0 + 0.0im) :
+        w == T.LargestImaginaryPart ? JacobiDavidson.LargestImaginaryPart(0.0 + 0.0im) :
+        JacobiDavidson.SmallestImaginaryPart(0.0 + 0.0im)
+end
+
+# jdqr yields a partial Schur decomposition `A*Q = Q*R`. Eigenpairs are recovered
+# from the small projected `R = Q'AQ`: if `R*y = λ*y` then `A*(Q*y) = λ*(Q*y)`.
+function _jd_standard_pairs(A, pschur)
+    T = complex(float(eltype(A)))
+    k = length(pschur.values)
+    k == 0 && return (T[], Matrix{T}(undef, size(A, 1), 0))
+    Q = pschur.Q[:, 1:k]
+    F = eigen(Q' * (A * Q))
+    return (F.values, Q * F.vectors)
+end
+
+end

ext/LinearSolveKrylovKitExt.jl

@@ -61,4 +61,61 @@ end
 
 LinearSolve.update_tolerances_internal!(cache, alg::KrylovKitJL, atol, rtol) = nothing
 
+function SciMLBase.solve(
+        prob::LinearSolve.EigenvalueProblem,
+        alg::LinearSolve.KrylovKitEigen,
+        args...; kwargs...
+    )
+    nev = LinearSolve.default_num_eigenpairs(prob)
+    which = _krylovkit_which(prob.eigentarget)
+    kw = (; prob.kwargs..., alg.kwargs..., kwargs...)
+    values, vectors, info = if prob.shift !== nothing
+        _shift_invert_eigsolve(prob, nev, kw)
+    elseif prob.B === nothing
+        KrylovKit.eigsolve(prob.A, nev, which; kw...)
+    else
+        KrylovKit.geneigsolve((prob.A, prob.B), nev, which; kw...)
+    end
+    if prob.shift !== nothing
+        values = prob.shift .+ inv.(values)
+    end
+    vecmat = reduce(hcat, vectors)
+    values, vecmat = LinearSolve._select_eigenpairs(
+        values, vecmat, nev, prob.eigentarget, prob.shift
+    )
+    retcode = info.converged >= length(values) ? ReturnCode.Success : ReturnCode.ConvergenceFailure
+    return LinearSolve.build_eigenvalue_solution(
+        prob, alg, values, vecmat; retcode, resid = info.normres, stats = info
+    )
+end
+
+# KrylovKit.eigsolve requires a raw ARPACK-style Symbol for `which`; this
+# mapping is purely a private adapter to that third-party API.
+function _krylovkit_which(w::LinearSolve.EigenvalueTarget.T)
+    T = LinearSolve.EigenvalueTarget
+    return w == T.LargestMagnitude ? :LM :
+        w == T.SmallestMagnitude ? :SM :
+        w == T.LargestRealPart ? :LR :
+        w == T.SmallestRealPart ? :SR :
+        w == T.LargestImaginaryPart ? :LI :
+        :SI
+end
+
+function _shift_invert_eigsolve(prob, nev, kw)
+    A, B, shift = prob.A, prob.B, prob.shift
+    T = isnothing(B) ? promote_type(eltype(A), typeof(shift)) :
+        promote_type(eltype(A), eltype(B), typeof(shift))
+    if isnothing(B)
+        F = factorize(A - shift * I)
+        op = x -> F \ x
+    elseif B isa LinearAlgebra.UniformScaling
+        F = factorize(A - shift * B)
+        op = x -> F \ (B.λ * x)
+    else
+        F = factorize(A - shift * B)
+        op = x -> F \ (B * x)
+    end
+    return KrylovKit.eigsolve(op, size(A, 2), nev, :LM, T; kw...)
+end
+
 end

src/LinearSolve.jl

@@ -410,6 +410,7 @@ function defaultalg_symbol end
 include("verbosity.jl")
 include("blas_logging.jl")
 include("generic_lufact.jl")
+include("eigenvalue.jl")
 include("common.jl")
 include("extension_algs.jl")
 include("factorization.jl")
@@ -573,4 +574,8 @@ export LinearSolveAdjoint
 
 export LinearVerbosity
 
+export AbstractEigenvalueAlgorithm,
+    DenseEigen, ArpackJL, ArnoldiMethod, ArnoldiMethodJL,
+    KrylovKitEigen, JacobiDavidsonJL
+
 end