Merge pull request #21 from vpuri3/vp/updates

LinearSolve interface

Author: ChrisRackauckas

.gitignore

@@ -3,3 +3,5 @@
 *.jl.mem
 /docs/build/
 Manifest.toml
+
+*.swp

Project.toml

@@ -5,15 +5,20 @@ version = "0.1.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
+KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [compat]
 ArrayInterface = "3"
+IterativeSolvers = "0.9.2"
 Krylov = "0.7"
 Reexport = "1"
 SciMLBase = "1.18.6"
@@ -22,7 +27,9 @@ UnPack = "1"
 julia = "1"
 
 [extras]
+DiffEqProblemLibrary = "a077e3f3-b75c-5d7f-a0c6-6bc4c8ec64a9"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "OrdinaryDiffEq", "DiffEqProblemLibrary"]

src/LinearSolve.jl

@@ -1,23 +1,48 @@
 module LinearSolve
 
-using ArrayInterface: lu_instance
+using ArrayInterface
+using RecursiveFactorization
 using Base: cache_dependencies, Bool
-using Krylov
 using LinearAlgebra
-using Reexport
+using SparseArrays
 using SciMLBase: AbstractDiffEqOperator, AbstractLinearAlgorithm
 using Setfield
 using UnPack
 
+# wrap
+import Krylov
+import KrylovKit
+import IterativeSolvers
+
+using Reexport
 @reexport using SciMLBase
 
-abstract type SciMLLinearSolveAlgorithm end
+abstract type SciMLLinearSolveAlgorithm <: SciMLBase.AbstractLinearAlgorithm end
+abstract type AbstractFactorization <: SciMLLinearSolveAlgorithm end
+abstract type AbstractKrylovSubspaceMethod <: SciMLLinearSolveAlgorithm end
 
 include("common.jl")
 include("factorization.jl")
-include("krylov.jl")
+include("wrappers.jl")
+include("default.jl")
+
+const IS_OPENBLAS = Ref(true)
+isopenblas() = IS_OPENBLAS[]
+
+function __init__()
+  @static if VERSION < v"1.7beta"
+    blas = BLAS.vendor()
+    IS_OPENBLAS[] = blas == :openblas64 || blas == :openblas
+  else
+    IS_OPENBLAS[] = occursin("openblas", BLAS.get_config().loaded_libs[1].libname)
+  end
+end
 
-export LUFactorization, SVDFactorization, QRFactorization
-export KrylovJL
+export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization
+export KrylovJL, KrylovJL_CG, KrylovJL_GMRES, KrylovJL_BICGSTAB,
+       KrylovJL_MINRES,
+       IterativeSolversJL, IterativeSolversJL_CG, IterativeSolversJL_GMRES,
+       IterativeSolversJL_BICGSTAB, IterativeSolversJL_MINRES
+export DefaultLinSolve
 
 end

src/common.jl

@@ -1,80 +1,90 @@
-struct LinearCache{TA,Tb,Tp,Talg,Tc,Tr,Tl}
+struct LinearCache{TA,Tb,Tu,Tp,Talg,Tc,Tl,Tr}
     A::TA
     b::Tb
+    u::Tu
     p::Tp
     alg::Talg
-    cacheval::Tc
-    isfresh::Bool
-    Pr::Tr
-    Pl::Tl
+    cacheval::Tc  # store alg cache here
+    isfresh::Bool # false => cacheval is set wrt A, true => update cacheval wrt A
+    Pl::Tl        # store final preconditioner here. not being used rn
+    Pr::Tr        # wrappers are using preconditioner in cache.alg for now
 end
 
-function set_A(cache, A)
+function set_A(cache::LinearCache, A)
     @set! cache.A = A
     @set! cache.isfresh = true
+    return cache
 end
 
-function set_b(cache, b)
+function set_b(cache::LinearCache, b)
     @set! cache.b = b
+    return cache
+end
+
+function set_u(cache::LinearCache, u)
+    @set! cache.u = u
+    return cache
 end
 
-function set_p(cache, p)
+function set_p(cache::LinearCache, p)
     @set! cache.p = p
-    # @set! cache.isfresh = true
+#   @set! cache.isfresh = true
+    return cache
 end
 
-function set_cacheval(cache::LinearCache, alg)
+function set_cacheval(cache::LinearCache, alg_cache)
     if cache.isfresh
-        @set! cache.cacheval = alg
+        @set! cache.cacheval = alg_cache
         @set! cache.isfresh = false
     end
     return cache
 end
 
-function SciMLBase.init(
-    prob::LinearProblem,
-    alg,
-    args...;
-    alias_A = false,
-    alias_b = false,
-    kwargs...,
-)
-    @unpack A, b, p = prob
-    if alg isa LUFactorization
-        fact = lu_instance(A)
-        Tfact = typeof(fact)
-    else
-        fact = nothing
-        Tfact = Any
-    end
-    Pr = nothing
-    Pl = nothing
+init_cacheval(alg::SciMLLinearSolveAlgorithm, A, b, u) = nothing
+
+function SciMLBase.init(prob::LinearProblem, alg, args...;
+                        alias_A = false, alias_b = false,
+                        kwargs...,
+                       )
+    @unpack A, b, u0, p = prob
 
-    A = alias_A ? A : copy(A)
-    b = alias_b ? b : copy(b)
+    u0 = (u0 === nothing) ? zero(b) : u0
+
+    cacheval = init_cacheval(alg, A, b, u0)
+    isfresh = cacheval === nothing
+    Tc = isfresh ? Any : typeof(cacheval)
+
+    Pl = LinearAlgebra.I
+    Pr = LinearAlgebra.I
+
+    A = alias_A ? A : deepcopy(A)
+    b = alias_b ? b : deepcopy(b)
 
     cache = LinearCache{
         typeof(A),
         typeof(b),
+        typeof(u0),
         typeof(p),
         typeof(alg),
-        Tfact,
-        typeof(Pr),
+        Tc,
         typeof(Pl),
+        typeof(Pr),
     }(
         A,
         b,
+        u0,
         p,
         alg,
-        fact,
-        true,
-        Pr,
+        cacheval,
+        isfresh,
         Pl,
+        Pr,
     )
     return cache
 end
 
-SciMLBase.solve(prob::LinearProblem, alg, args...; kwargs...) =
-    solve(init(prob, alg, args...; kwargs...))
+SciMLBase.solve(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
+                args...; kwargs...) = solve(init(prob, alg, args...; kwargs...))
 
-SciMLBase.solve(cache) = solve(cache, cache.alg)
+SciMLBase.solve(cache::LinearCache, args...; kwargs...) =
+    solve(cache, cache.alg, args...; kwargs...)

src/default.jl

@@ -0,0 +1,35 @@
+## Default algorithm
+
+function SciMLBase.solve(cache::LinearCache, alg::Nothing,
+                         args...; kwargs...)
+    @unpack A = cache
+    if A isa Matrix
+        if ArrayInterface.can_setindex(x) && (size(A,1) <= 100 ||
+                                              (isopenblas() && size(A,1) <= 500)
+                                             )
+            alg = GenericFactorization(;fact_alg=:(RecursiveFactorization.lu!))
+            SciMLBase.solve(cache, alg, args...; kwargs...)
+        else
+            alg = LUFactorization()
+            SciMLBase.solve(cache, alg, args...; kwargs...)
+        end
+    elseif A isa Tridiagonal
+        alg = GenericFactorization(;fact_alg=lu!)
+        SciMLBase.solve(cache, alg, args...; kwargs...)
+    elseif A isa SymTridiagonal
+        alg = GenericFactorization(;fact_alg=ldlt!)
+        SciMLBase.solve(cache, alg, args...; kwargs...)
+    elseif A isa SparseMatrixCSC
+        alg = LUFactorization()
+        SciMLBase.solve(cache, alg, args...; kwargs...)
+    elseif ArrayInterface.isstructured(A)
+        alg = GenericFactorization()
+        SciMLBase.solve(cache, alg, args...; kwargs...)
+    elseif !(A isa AbstractDiffEqOperator)
+        alg = QRFactorization()
+        SciMLBase.solve(cache, alg, args...; kwargs...)
+    else
+        alg = IterativeSolversJL_GMRES()
+        SciMLBase.solve(cache, alg, args...; kwargs...)
+    end
+end

src/factorization.jl

@@ -1,5 +1,15 @@
+function SciMLBase.solve(cache::LinearCache, alg::AbstractFactorization)
+    if cache.isfresh
+        fact = init_cacheval(alg, cache.A, cache.b, cache.u)
+        cache = set_cacheval(cache, fact)
+    end
+
+    ldiv!(cache.u,cache.cacheval, cache.b)
+end
+
+## LUFactorization
 
-struct LUFactorization{P} <: AbstractLinearAlgorithm
+struct LUFactorization{P} <: AbstractFactorization
     pivot::P
 end
 
@@ -12,14 +22,16 @@ function LUFactorization()
     LUFactorization(pivot)
 end
 
-function SciMLBase.solve(cache::LinearCache, alg::LUFactorization)
-    cache.A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
-        error("LU is not defined for $(typeof(prob.A))")
-    cache = set_cacheval(cache, lu!(cache.A, alg.pivot))
-    ldiv!(cache.cacheval, cache.b)
+function init_cacheval(alg::LUFactorization, A, b, u)
+    A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
+        error("LU is not defined for $(typeof(A))")
+    fact = lu!(A, alg.pivot)
+    return fact
 end
 
-struct QRFactorization{P} <: AbstractLinearAlgorithm
+## QRFactorization
+
+struct QRFactorization{P} <: AbstractFactorization
     pivot::P
     blocksize::Int
 end
@@ -33,26 +45,44 @@ function QRFactorization()
     QRFactorization(pivot, 16)
 end
 
-function SciMLBase.solve(cache::LinearCache, alg::QRFactorization)
-    cache.A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
-        error("QR is not defined for $(typeof(prob.A))")
-    cache = set_cacheval(
-        cache,
-        qr!(cache.A.A, alg.pivot; blocksize = alg.blocksize),
-    )
-    ldiv!(cache.cacheval, cache.b)
+function init_cacheval(alg::QRFactorization, A, b, u)
+    A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
+        error("QR is not defined for $(typeof(A))")
+
+    fact = qr!(A.A, alg.pivot; blocksize = alg.blocksize)
+    return fact
 end
 
-struct SVDFactorization{A} <: AbstractLinearAlgorithm
+## SVDFactorization
+
+struct SVDFactorization{A} <: AbstractFactorization
     full::Bool
     alg::A
 end
 
 SVDFactorization() = SVDFactorization(false, LinearAlgebra.DivideAndConquer())
 
-function SciMLBase.solve(cache::LinearCache, alg::SVDFactorization)
-    cache.A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
-        error("SVD is not defined for $(typeof(prob.A))")
-    cache = set_cacheval(cache, svd!(cache.A; full = alg.full, alg = alg.alg))
-    ldiv!(cache.cacheval, cache.b)
+function init_cacheval(alg::SVDFactorization, A, b, u)
+    A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
+        error("SVD is not defined for $(typeof(A))")
+
+    fact = svd!(A; full = alg.full, alg = alg.alg)
+    return fact
+end
+
+## GenericFactorization
+
+struct GenericFactorization{F} <: AbstractFactorization
+    fact_alg::F
+end
+
+GenericFactorization(;fact_alg = LinearAlgebra.factorize) =
+    GenericFactorization(fact_alg)
+
+function init_cacheval(alg::GenericFactorization, A, b, u)
+    A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
+        error("GenericFactorization is not defined for $(typeof(A))")
+
+    fact = alg.fact_alg(A)
+    return fact
 end