Use Pivoted QR for Underdetermined Systems

Author: avik-pal

src/LinearSolve.jl

@@ -62,6 +62,9 @@ needs_concrete_A(alg::AbstractKrylovSubspaceMethod) = false
 needs_concrete_A(alg::AbstractSolveFunction) = false
 
 # Util
+is_underdetermined(x) = false
+is_underdetermined(A::AbstractMatrix) = size(A, 1) < size(A, 2)
+is_underdetermined(A::AbstractSciMLOperator) = size(A, 1) < size(A, 2)
 
 _isidentity_struct(A) = false
 _isidentity_struct(λ::Number) = isone(λ)
@@ -96,6 +99,7 @@ EnumX.@enumx DefaultAlgorithmChoice begin
     NormalCholeskyFactorization
     AppleAccelerateLUFactorization
     MKLLUFactorization
+    QRFactorizationPivoted
 end
 
 struct DefaultLinearSolver <: SciMLLinearSolveAlgorithm

src/default.jl

@@ -1,6 +1,6 @@
 needs_concrete_A(alg::DefaultLinearSolver) = true
 mutable struct DefaultLinearSolverInit{T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12,
-    T13, T14, T15, T16, T17, T18}
+    T13, T14, T15, T16, T17, T18, T19}
     LUFactorization::T1
     QRFactorization::T2
     DiagonalFactorization::T3
@@ -19,6 +19,7 @@ mutable struct DefaultLinearSolverInit{T1, T2, T3, T4, T5, T6, T7, T8, T9, T10,
     NormalCholeskyFactorization::T16
     AppleAccelerateLUFactorization::T17
     MKLLUFactorization::T18
+    QRFactorizationPivoted::T19
 end
 
 # Legacy fallback
@@ -168,8 +169,8 @@ function defaultalg(A, b, assump::OperatorAssumptions)
                        (A === nothing ? eltype(b) <: Union{Float32, Float64} :
                         eltype(A) <: Union{Float32, Float64})
                     DefaultAlgorithmChoice.RFLUFactorization
-                    #elseif A === nothing || A isa Matrix
-                    #    alg = FastLUFactorization()
+                #elseif A === nothing || A isa Matrix
+                #    alg = FastLUFactorization()
                 elseif usemkl && (A === nothing ? eltype(b) <: Union{Float32, Float64} :
                         eltype(A) <: Union{Float32, Float64})
                     DefaultAlgorithmChoice.MKLLUFactorization
@@ -199,9 +200,19 @@ function defaultalg(A, b, assump::OperatorAssumptions)
     elseif assump.condition === OperatorCondition.WellConditioned
         DefaultAlgorithmChoice.NormalCholeskyFactorization
     elseif assump.condition === OperatorCondition.IllConditioned
-        DefaultAlgorithmChoice.QRFactorization
+        if is_underdetermined(A)
+            # Underdetermined
+            DefaultAlgorithmChoice.QRFactorizationPivoted
+        else
+            DefaultAlgorithmChoice.QRFactorization
+        end
     elseif assump.condition === OperatorCondition.VeryIllConditioned
-        DefaultAlgorithmChoice.QRFactorization
+        if is_underdetermined(A)
+            # Underdetermined
+            DefaultAlgorithmChoice.QRFactorizationPivoted
+        else
+            DefaultAlgorithmChoice.QRFactorization
+        end
     elseif assump.condition === OperatorCondition.SuperIllConditioned
         DefaultAlgorithmChoice.SVDFactorization
     else
@@ -247,6 +258,12 @@ function algchoice_to_alg(alg::Symbol)
         NormalCholeskyFactorization()
     elseif alg === :AppleAccelerateLUFactorization
         AppleAccelerateLUFactorization()
+    elseif alg === :QRFactorizationPivoted
+        @static if VERSION ≥ v"1.7beta"
+            QRFactorization(ColumnNorm())
+        else
+            QRFactorization(Val(true))
+        end
     else
         error("Algorithm choice symbol $alg not allowed in the default")
     end
@@ -310,6 +327,7 @@ function defaultalg_symbol(::Type{T}) where {T}
     Symbol(split(string(SciMLBase.parameterless_type(T)), ".")[end])
 end
 defaultalg_symbol(::Type{<:GenericFactorization{typeof(ldlt!)}}) = :LDLtFactorization
+defaultalg_symbol(::Type{<:QRFactorization{ColumnNorm}}) = :QRFactorizationPivoted
 
 """
 if alg.alg === DefaultAlgorithmChoice.LUFactorization

src/factorization.jl

@@ -158,6 +158,16 @@ function QRFactorization(inplace = true)
     QRFactorization(pivot, 16, inplace)
 end
 
+@static if VERSION ≥ v"1.7beta"
+    function QRFactorization(pivot::LinearAlgebra.PivotingStrategy, inplace::Bool = true)
+        QRFactorization(pivot, 16, inplace)
+    end
+else
+    function QRFactorization(pivot::Val, inplace::Bool = true)
+        QRFactorization(pivot, 16, inplace)
+    end
+end
+
 function do_factorization(alg::QRFactorization, A, b, u)
     A = convert(AbstractMatrix, A)
     if alg.inplace && !(A isa SparseMatrixCSC) && !(A isa GPUArraysCore.AbstractGPUArray)

test/nonsquare.jl

@@ -56,3 +56,12 @@ solve(LinearProblem(A, b), (LinearSolve.NormalCholeskyFactorization())).u;
 solve(LinearProblem(A, b),
     assumptions = (OperatorAssumptions(false;
         condition = OperatorCondition.WellConditioned))).u;
+
+# Underdetermined
+m, n = 2, 3
+
+A = rand(m, n)
+b = rand(m)
+prob = LinearProblem(A, b)
+res = A \ b
+@test solve(prob).u ≈ res