Merge pull request #289 from SciML/conditioning_assumption

Expand OperatorAssumptions to allow for choosing conditioning

Author: ChrisRackauckas

Project.toml

@@ -6,6 +6,7 @@ version = "1.40.0"
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 FastLapackInterface = "29a986be-02c6-4525-aec4-84b980013641"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"

docs/pages.jl

@@ -6,6 +6,7 @@ pages = ["index.md",
     "Basics" => Any["basics/LinearProblem.md",
                     "basics/common_solver_opts.md",
                     "basics/CachingAPI.md",
+                    "basics/OperatorAssumptions.md",
                     "basics/Preconditioners.md",
                     "basics/FAQ.md"],
     "Solvers" => Any["solvers/solvers.md"],

docs/src/basics/FAQ.md

@@ -9,6 +9,18 @@ Note that if `\` is good enough for you, great! We still tend to use `\` in the
 However, if you're building a package, you may want to consider using LinearSolve.jl for the improved
 efficiency and ability to choose solvers.
 
+## I'm seeing some dynamic dispatches in the default algorithm choice, how do I reduce that?
+
+Make sure you set the `OperatorAssumptions` to get the full performance, especially the `issquare` choice
+as otherwise that will need to be determined at runtime.
+
+## I found a faster algorithm that can be used than what LinearSolve.jl chose?
+
+What assumptions are made as part of your method? If your method only works on well-conditioned operators, then
+make sure you set the `WellConditioned` assumption in the `assumptions`. See the 
+[OperatorAssumptions page for more details](@ref assumptions). If using the right assumptions does not improve
+the performance to the expected state, please open an issue and we will improve the default algorithm.
+
 ## Python's NumPy/SciPy just calls fast Fortran/C code, why would LinearSolve.jl be any better?
 
 This is addressed in the [JuliaCon 2022 video](https://www.youtube.com/watch?v=JWI34_w-yYw&t=182s). This happens in

docs/src/basics/OperatorAssumptions.md

@@ -0,0 +1,15 @@
+# [Linear Solve Operator Assumptions](@id assumptions)
+
+```@docs
+OperatorAssumptions
+OperatorCondition
+```
+
+## Condition Number Specifications
+
+```@docs
+IllConditioned
+VeryIllConditioned
+SuperIllConditioned
+WellConditioned
+```

docs/src/basics/common_solver_opts.md

@@ -14,7 +14,8 @@ The following are the options these algorithms take, along with their defaults.
     algorithms can write and change `b` upon usage. Care must be taken as the
     original input will be modified. Default is `false`.
   - `verbose`: Whether to print extra information. Defaults to `false`.
-
+  - `assumptions`: Sets the assumptions of the operator in order to effect the default
+    choice algorithm. See the [Operator Assumptions page for more details](@ref assumptions).
 ## Iterative Solver Controls
 
 Error controls are not used by all algorithms. Specifically, direct solves always

src/LinearSolve.jl

@@ -20,6 +20,7 @@ using KLU
 using Sparspak
 using FastLapackInterface
 using DocStringExtensions
+using EnumX
 import GPUArraysCore
 import Preferences
 

src/common.jl

@@ -1,9 +1,67 @@
-struct OperatorAssumptions{issq} end
-function OperatorAssumptions(issquare = nothing)
+"""
+`OperatorCondition`
+
+Specifies the assumption of matrix conditioning for the default linear solver choices. Condition number
+is defined as the ratio of eigenvalues. The numerical stability of many linear solver algorithms
+can be dependent on the condition number of the matrix. The condition number can be computed as:
+
+```julia
+using LinearAlgebra
+cond(rand(100,100))
+```
+
+However, in practice this computation is very expensive and thus not possible for most practical cases.
+Therefore, OperatorCondition lets one share to LinearSolve the expected conditioning. The higher the
+expected condition number, the safer the algorithm needs to be and thus there is a trade-off between
+numerical performance and stability. By default the method assumes the operator may be ill-conditioned
+for the standard linear solvers to converge (such as LU-factorization), though more extreme 
+ill-conditioning or well-conditioning could be the case and specified through this assumption.
+"""
+EnumX.@enumx OperatorCondition begin
+    """
+    `OperatorCondition.IllConditioned`
+
+    The default assumption of LinearSolve. Assumes that the operator can have minor ill-conditioning
+    and thus needs to use safe algorithms.
+    """
+    IllConditioned
+    """
+    `OperatorCondition.VeryIllConditioned`
+
+    Assumes that the operator can have fairly major ill-conditioning and thus the standard linear algebra
+    algorithms cannot be used.
+    """
+    VeryIllConditioned
+    """
+    `OperatorCondition.SuperIllConditioned`
+
+    Assumes that the operator can have fairly extreme ill-conditioning and thus the most stable algorithm
+    is used.
+    """
+    SuperIllConditioned
+    """
+    `OperatorCondition.WellConditioned`
+
+    Assumes that the operator can have fairly contained conditioning and thus the fastest algorithm is
+    used.
+    """
+    WellConditioned
+end
+
+"""
+    OperatorAssumptions(issquare = nothing; condition::OperatorCondition.T = IllConditioned)
+
+Sets the operator `A` assumptions used as part of the default algorithm
+"""
+struct OperatorAssumptions{issq,condition} end
+function OperatorAssumptions(issquare = nothing; condition::OperatorCondition.T = OperatorCondition.IllConditioned)
     issq = something(_unwrap_val(issquare), Nothing)
-    OperatorAssumptions{issq}()
+    condition = _unwrap_val(condition)
+    OperatorAssumptions{issq,condition}()
 end
-__issquare(::OperatorAssumptions{issq}) where {issq} = issq
+__issquare(::OperatorAssumptions{issq,condition}) where {issq,condition} = issq
+__conditioning(::OperatorAssumptions{issq,condition}) where {issq,condition} = condition
+
 
 struct LinearCache{TA, Tb, Tu, Tp, Talg, Tc, Tl, Tr, Ttol, issq}
     A::TA

src/default.jl

@@ -66,22 +66,30 @@ end
     end
 end
 
-function defaultalg(A::GPUArraysCore.AbstractGPUArray, b, ::OperatorAssumptions{true})
+function defaultalg(A::GPUArraysCore.AbstractGPUArray, b, assump::OperatorAssumptions{true})
     if VERSION >= v"1.8-"
         LUFactorization()
     else
         QRFactorization()
     end
 end
 
-function defaultalg(A, b::GPUArraysCore.AbstractGPUArray, ::OperatorAssumptions{true})
+function defaultalg(A::GPUArraysCore.AbstractGPUArray, b, assump::OperatorAssumptions{true,OperatorCondition.IllConditioned})
+    QRFactorization()
+end
+
+function defaultalg(A, b::GPUArraysCore.AbstractGPUArray, assump::OperatorAssumptions{true})
     if VERSION >= v"1.8-"
         LUFactorization()
     else
         QRFactorization()
     end
 end
 
+function defaultalg(A, b::GPUArraysCore.AbstractGPUArray, assump::OperatorAssumptions{true,OperatorCondition.IllConditioned})
+    QRFactorization()
+end
+
 function defaultalg(A::SciMLBase.AbstractSciMLOperator, b,
                     assumptions::OperatorAssumptions{true})
     if has_ldiv!(A)
@@ -121,6 +129,11 @@ function defaultalg(A::GPUArraysCore.AbstractGPUArray, b::GPUArraysCore.Abstract
     end
 end
 
+function defaultalg(A::GPUArraysCore.AbstractGPUArray, b::GPUArraysCore.AbstractGPUArray,
+                    ::OperatorAssumptions{true,OperatorCondition.IllConditioned})
+    QRFactorization()
+end
+
 function defaultalg(A::GPUArraysCore.AbstractGPUArray, b, ::OperatorAssumptions{false})
     QRFactorization()
 end
@@ -136,13 +149,13 @@ function defaultalg(A::GPUArraysCore.AbstractGPUArray, b::GPUArraysCore.Abstract
 end
 
 # Allows A === nothing as a stand-in for dense matrix
-function defaultalg(A, b, ::OperatorAssumptions{true})
+function defaultalg(A, b, assump::OperatorAssumptions{true})
     # Special case on Arrays: avoid BLAS for RecursiveFactorization.jl when
     # it makes sense according to the benchmarks, which is dependent on
     # whether MKL or OpenBLAS is being used
     if (A === nothing && !(b isa GPUArraysCore.AbstractGPUArray)) || A isa Matrix
         if (A === nothing || eltype(A) <: Union{Float32, Float64, ComplexF32, ComplexF64}) &&
-           ArrayInterface.can_setindex(b)
+           ArrayInterface.can_setindex(b) && __conditioning(assump) === OperatorCondition.IllConditioned
             if length(b) <= 10
                 alg = GenericLUFactorization()
             elseif (length(b) <= 100 || (isopenblas() && length(b) <= 500)) &&
@@ -154,6 +167,10 @@ function defaultalg(A, b, ::OperatorAssumptions{true})
             else
                 alg = LUFactorization()
             end
+        elseif __conditioning(assump) === OperatorCondition.VeryIllConditioned
+            alg = QRFactorization()
+        elseif __conditioning(assump) === OperatorCondition.SuperIllConditioned
+            alg = SVDFactorization(false, LinearAlgebra.QRIteration())
         else
             alg = LUFactorization()
         end
@@ -170,10 +187,18 @@ function defaultalg(A, b, ::OperatorAssumptions{true})
     alg
 end
 
-function defaultalg(A, b, ::OperatorAssumptions{false})
+function defaultalg(A, b, ::OperatorAssumptions{false,OperatorCondition.IllConditioned})
     QRFactorization()
 end
 
+function defaultalg(A, b, ::OperatorAssumptions{false,OperatorCondition.VeryIllConditioned})
+    QRFactorization()
+end
+
+function defaultalg(A, b, ::OperatorAssumptions{false,OperatorCondition.SuperIllConditioned})
+    SVDFactorization(false, LinearAlgebra.QRIteration())
+end
+
 ## Catch high level interface
 
 function SciMLBase.solve(cache::LinearCache, alg::Nothing,

test/default_algs.jl

@@ -6,7 +6,7 @@ using LinearSolve, LinearAlgebra, SparseArrays, Test
       DiagonalFactorization
 
 @test LinearSolve.defaultalg(nothing, zeros(5),
-                             LinearSolve.OperatorAssumptions{false}()) isa QRFactorization
+                             LinearSolve.OperatorAssumptions(false)) isa QRFactorization
 
 @test LinearSolve.defaultalg(sprand(1000, 1000, 0.01), zeros(1000)) isa KLUFactorization
 @test LinearSolve.defaultalg(sprand(11000, 11000, 0.001), zeros(11000)) isa