Proper handling of static arrays

Author: avik-pal

Project.toml

@@ -63,8 +63,8 @@ Aqua = "0.8"
 ArrayInterface = "7.4.11"
 BandedMatrices = "1"
 BlockDiagonals = "0.1"
-ConcreteStructs = "0.2"
 CUDA = "5"
+ConcreteStructs = "0.2"
 DocStringExtensions = "0.9"
 EnumX = "1"
 Enzyme = "0.11"
@@ -77,15 +77,15 @@ GPUArraysCore = "0.1"
 HYPRE = "1.4.0"
 InteractiveUtils = "1.6"
 IterativeSolvers = "0.9.3"
-Libdl = "1.6"
-LinearAlgebra = "1.9"
 JET = "0.8"
 KLU = "0.3.0, 0.4"
 KernelAbstractions = "0.9"
 Krylov = "0.9"
 KrylovKit = "0.6"
-Metal = "0.5"
+Libdl = "1.6"
+LinearAlgebra = "1.9"
 MPI = "0.20"
+Metal = "0.5"
 MultiFloats = "1"
 Pardiso = "0.5"
 Pkg = "1"
@@ -102,8 +102,8 @@ SciMLOperators = "0.3"
 Setfield = "1"
 SparseArrays = "1.9"
 Sparspak = "0.3.6"
-StaticArraysCore = "1"
 StaticArrays = "1"
+StaticArraysCore = "1"
 Test = "1"
 UnPack = "1"
 julia = "1.9"

src/common.jl

@@ -119,6 +119,15 @@ default_alias_b(::Any, ::Any, ::Any) = false
 default_alias_A(::AbstractKrylovSubspaceMethod, ::Any, ::Any) = true
 default_alias_b(::AbstractKrylovSubspaceMethod, ::Any, ::Any) = true
 
+function __init_u0_from_Ab(A, b)
+    u0 = similar(b, size(A, 2))
+    fill!(u0, false)
+    return u0
+end
+function __init_u0_from_Ab(A::SMatrix{S1, S2}, b) where {S1, S2}
+    return zeros(SVector{S2, eltype(b)})
+end
+
 function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
     args...;
     alias_A = default_alias_A(alg, prob.A, prob.b),
@@ -133,7 +142,7 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
     kwargs...)
     @unpack A, b, u0, p = prob
 
-    A = if alias_A
+    A = if alias_A || A isa SMatrix
         A
     elseif A isa Array || A isa SparseMatrixCSC
         copy(A)
@@ -143,34 +152,29 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
 
     b = if b isa SparseArrays.AbstractSparseArray && !(A isa Diagonal)
         Array(b) # the solution to a linear solve will always be dense!
-    elseif alias_b
+    elseif alias_b || b isa SVector
         b
     elseif b isa Array || b isa SparseMatrixCSC
         copy(b)
     else
         deepcopy(b)
     end
 
-    u0 = if u0 !== nothing
-        u0
-    else
-        u0 = similar(b, size(A, 2))
-        fill!(u0, false)
-    end
+    u0_ = u0 !== nothing ? u0 : __init_u0_from_Ab(A, b)
 
     # Guard against type mismatch for user-specified reltol/abstol
     reltol = real(eltype(prob.b))(reltol)
     abstol = real(eltype(prob.b))(abstol)
 
-    cacheval = init_cacheval(alg, A, b, u0, Pl, Pr, maxiters, abstol, reltol, verbose,
+    cacheval = init_cacheval(alg, A, b, u0_, Pl, Pr, maxiters, abstol, reltol, verbose,
         assumptions)
     isfresh = true
     Tc = typeof(cacheval)
 
     cache = LinearCache{
         typeof(A),
         typeof(b),
-        typeof(u0),
+        typeof(u0_),
         typeof(p),
         typeof(alg),
         Tc,
@@ -180,7 +184,7 @@ function SciMLBase.init(prob::LinearProblem, alg::SciMLLinearSolveAlgorithm,
         typeof(assumptions.issq),
     }(A,
         b,
-        u0,
+        u0_,
         p,
         alg,
         cacheval,

src/default.jl

@@ -36,6 +36,14 @@ function defaultalg(A, b, assump::OperatorAssumptions{Nothing})
     defaultalg(A, b, OperatorAssumptions(issq, assump.condition))
 end
 
+function defaultalg(A::SMatrix{S1, S2}, b, assump::OperatorAssumptions{Bool}) where {S1, S2}
+    if S1 == S2
+        return LUFactorization()
+    else
+        return SVDFactorization()  # QR(...) \ b is not defined currently
+    end
+end
+
 function defaultalg(A::Tridiagonal, b, assump::OperatorAssumptions{Bool})
     if assump.issq
         DefaultLinearSolver(DefaultAlgorithmChoice.LUFactorization)
@@ -175,10 +183,6 @@ function defaultalg(A, b, assump::OperatorAssumptions{Bool})
                 DefaultAlgorithmChoice.LUFactorization
             end
 
-            # For static arrays GMRES allocates a lot. Use factorization
-        elseif A isa StaticArray
-            DefaultAlgorithmChoice.LUFactorization
-
             # This catches the cases where a factorization overload could exist
             # For example, BlockBandedMatrix
         elseif A !== nothing && ArrayInterface.isstructured(A)
@@ -190,9 +194,6 @@ function defaultalg(A, b, assump::OperatorAssumptions{Bool})
         end
     elseif assump.condition === OperatorCondition.WellConditioned
         DefaultAlgorithmChoice.NormalCholeskyFactorization
-    elseif A isa StaticArray
-        # Static Array doesn't have QR() \ b defined
-        DefaultAlgorithmChoice.SVDFactorization
     elseif assump.condition === OperatorCondition.IllConditioned
         if is_underdetermined(A)
             # Underdetermined
@@ -269,8 +270,7 @@ function SciMLBase.init(prob::LinearProblem, alg::Nothing,
     args...;
     assumptions = OperatorAssumptions(issquare(prob.A)),
     kwargs...)
-    alg = defaultalg(prob.A, prob.b, assumptions)
-    SciMLBase.init(prob, alg, args...; assumptions, kwargs...)
+    SciMLBase.init(prob, defaultalg(prob.A, prob.b, assumptions), args...; assumptions, kwargs...)
 end
 
 function SciMLBase.solve!(cache::LinearCache, alg::Nothing,

test/static_arrays.jl

@@ -1,4 +1,4 @@
-using LinearSolve, StaticArrays, LinearAlgebra
+using LinearSolve, StaticArrays, LinearAlgebra, Test
 
 A = SMatrix{5, 5}(Hermitian(rand(5, 5) + I))
 b = SVector{5}(rand(5))