Fix default solve algorithm handling

SciML · Nov 24, 2021 · c4f8ed5 · c4f8ed5
1 parent cdfe484
commit c4f8ed5
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 8 deletions.
diff --git a/src/common.jl b/src/common.jl
@@ -40,9 +40,11 @@ function set_cacheval(cache::LinearCache, alg_cache)
     return cache
 end
 
-init_cacheval(alg::SciMLLinearSolveAlgorithm, A, b, u) = nothing
+init_cacheval(alg::Union{SciMLLinearSolveAlgorithm,Nothing}, A, b, u) = nothing
 
-function SciMLBase.init(prob::LinearProblem, alg, args...;
+SciMLBase.init(prob::LinearProblem, args...; kwargs...) = SciMLBase.init(prob,nothing,args...;kwargs...)
+
+function SciMLBase.init(prob::LinearProblem, alg::Union{SciMLLinearSolveAlgorithm,Nothing}, args...;
                         alias_A = false, alias_b = false,
                         kwargs...,
                        )
@@ -83,7 +85,9 @@ function SciMLBase.init(prob::LinearProblem, alg, args...;
     return cache
 end
 
-SciMLBase.solve(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
+SciMLBase.solve(prob::LinearProblem, args...; kwargs...) = solve(init(prob, nothing, args...; kwargs...))
+
+SciMLBase.solve(prob::LinearProblem, alg::Union{SciMLLinearSolveAlgorithm,Nothing},
                 args...; kwargs...) = solve(init(prob, alg, args...; kwargs...))
 
 SciMLBase.solve(cache::LinearCache, args...; kwargs...) =

diff --git a/src/default.jl b/src/default.jl
@@ -3,11 +3,15 @@
 function SciMLBase.solve(cache::LinearCache, alg::Nothing,
                          args...; kwargs...)
     @unpack A = cache
+    if A isa DiffEqArrayOperator
+        A = A.A
+    end
+
     if A isa Matrix
-        if ArrayInterface.can_setindex(x) && (size(A,1) <= 100 ||
+        if ArrayInterface.can_setindex(cache.b) && (size(A,1) <= 100 ||
                                               (isopenblas() && size(A,1) <= 500)
                                              )
-            alg = GenericFactorization(;fact_alg=:(RecursiveFactorization.lu!))
+            alg = GenericFactorization(;fact_alg=RecursiveFactorization.lu!)
             SciMLBase.solve(cache, alg, args...; kwargs...)
         else
             alg = LUFactorization()

diff --git a/src/factorization.jl b/src/factorization.jl
@@ -25,6 +25,10 @@ end
 function init_cacheval(alg::LUFactorization, A, b, u)
     A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("LU is not defined for $(typeof(A))")
+
+    if A isa AbstractDiffEqOperator
+        A = A.A
+    end
     fact = lu!(A, alg.pivot)
     return fact
 end
@@ -49,7 +53,10 @@ 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)
+    if A isa AbstractDiffEqOperator
+        A = A.A
+    end
+    fact = qr!(A, alg.pivot; blocksize = alg.blocksize)
     return fact
 end
 
@@ -66,6 +73,10 @@ function init_cacheval(alg::SVDFactorization, A, b, u)
     A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("SVD is not defined for $(typeof(A))")
 
+    if A isa AbstractDiffEqOperator
+        A = A.A
+    end
+
     fact = svd!(A; full = alg.full, alg = alg.alg)
     return fact
 end
@@ -83,6 +94,9 @@ function init_cacheval(alg::GenericFactorization, A, b, u)
     A isa Union{AbstractMatrix,AbstractDiffEqOperator} ||
         error("GenericFactorization is not defined for $(typeof(A))")
 
+    if A isa AbstractDiffEqOperator
+        A = A.A
+    end
     fact = alg.fact_alg(A)
     return fact
 end
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -32,12 +32,19 @@ function test_interface(alg, prob1, prob2)
     return
 end
 
+@testset "Default Linear Solver" begin
+    test_interface(nothing, prob1, prob2)
+
+    A1 = prob1.A; b1 = prob1.b; x1 = prob1.u0
+    y = solve(prob1)
+    @test A1 *  y  ≈ b1
+end
+
 @testset "Concrete Factorizations" begin
     for alg in (
                 LUFactorization(),
                 QRFactorization(),
-                SVDFactorization(),
-                #nothing
+                SVDFactorization()
                )
         @testset "$alg" begin
             test_interface(alg, prob1, prob2)