Merge pull request #393 from avik-pal/ap/banded

Proper handling for BandedMatrices

Author: ChrisRackauckas

.github/workflows/Downstream.yml

@@ -19,6 +19,7 @@ jobs:
           - {user: SciML, repo: OrdinaryDiffEq.jl, group: InterfaceII}
           - {user: SciML, repo: ModelingToolkit.jl, group: All}
           - {user: SciML, repo: SciMLSensitivity.jl, group: Core1}
+          - {user: SciML, repo: BoundaryValueDiffEq.jl, group: All}
 
     steps:
       - uses: actions/checkout@v4

Project.toml

@@ -1,7 +1,7 @@
 name = "LinearSolve"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 authors = ["SciML"]
-version = "2.10.0"
+version = "2.11.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -31,17 +31,20 @@ SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [weakdeps]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
-Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Pardiso = "46dd5b70-b6fb-5a00-ae2d-e8fea33afaf2"
+RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 
 [extensions]
+LinearSolveBandedMatricesExt = "BandedMatrices"
 LinearSolveBlockDiagonalsExt = "BlockDiagonals"
 LinearSolveCUDAExt = "CUDA"
 LinearSolveEnzymeExt = "Enzyme"
@@ -51,9 +54,11 @@ LinearSolveKernelAbstractionsExt = "KernelAbstractions"
 LinearSolveKrylovKitExt = "KrylovKit"
 LinearSolveMetalExt = "Metal"
 LinearSolvePardisoExt = "Pardiso"
+LinearSolveRecursiveArrayToolsExt = "RecursiveArrayTools"
 
 [compat]
 ArrayInterface = "7.4.11"
+BandedMatrices = "1"
 BlockDiagonals = "0.1"
 ConcreteStructs = "0.2"
 DocStringExtensions = "0.8, 0.9"
@@ -69,6 +74,7 @@ Krylov = "0.9"
 KrylovKit = "0.5, 0.6"
 PrecompileTools = "1"
 Preferences = "1"
+RecursiveArrayTools = "2"
 RecursiveFactorization = "0.2.8"
 Reexport = "1"
 Requires = "1"
@@ -80,6 +86,7 @@ UnPack = "1"
 julia = "1.6"
 
 [extras]
+BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -95,6 +102,7 @@ Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 MultiFloats = "bdf0d083-296b-4888-a5b6-7498122e68a5"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 

ext/LinearSolveBandedMatricesExt.jl

@@ -0,0 +1,58 @@
+module LinearSolveBandedMatricesExt
+
+using BandedMatrices, LinearAlgebra, LinearSolve
+import LinearSolve: defaultalg,
+    do_factorization, init_cacheval, DefaultLinearSolver, DefaultAlgorithmChoice
+
+# Defaults for BandedMatrices
+function defaultalg(A::BandedMatrix, b, ::OperatorAssumptions)
+    return DefaultLinearSolver(DefaultAlgorithmChoice.DirectLdiv!)
+end
+
+function defaultalg(A::Symmetric{<:Number, <:BandedMatrix}, b, ::OperatorAssumptions)
+    return DefaultLinearSolver(DefaultAlgorithmChoice.CholeskyFactorization)
+end
+
+# BandedMatrices `qr` doesn't allow other args without causing an ambiguity
+do_factorization(alg::QRFactorization, A::BandedMatrix, b, u) = alg.inplace ? qr!(A) : qr(A)
+
+function do_factorization(alg::LUFactorization, A::BandedMatrix, b, u)
+    _pivot = alg.pivot isa NoPivot ? Val(false) : Val(true)
+    return lu!(A, _pivot; check = false)
+end
+
+# For BandedMatrix
+for alg in (:SVDFactorization, :MKLLUFactorization, :DiagonalFactorization,
+    :SparspakFactorization, :KLUFactorization, :UMFPACKFactorization,
+    :GenericLUFactorization, :RFLUFactorization, :BunchKaufmanFactorization,
+    :CHOLMODFactorization, :NormalCholeskyFactorization, :LDLtFactorization,
+    :AppleAccelerateLUFactorization, :CholeskyFactorization)
+    @eval begin
+        function init_cacheval(::$(alg), ::BandedMatrix, b, u, Pl, Pr, maxiters::Int,
+            abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
+            return nothing
+        end
+    end
+end
+
+function init_cacheval(::LUFactorization, A::BandedMatrix, b, u, Pl, Pr, maxiters::Int,
+    abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
+    return lu(similar(A, 0, 0))
+end
+
+# For Symmetric BandedMatrix
+for alg in (:SVDFactorization, :MKLLUFactorization, :DiagonalFactorization,
+    :SparspakFactorization, :KLUFactorization, :UMFPACKFactorization,
+    :GenericLUFactorization, :RFLUFactorization, :BunchKaufmanFactorization,
+    :CHOLMODFactorization, :NormalCholeskyFactorization,
+    :AppleAccelerateLUFactorization, :QRFactorization, :LUFactorization)
+    @eval begin
+        function init_cacheval(::$(alg), ::Symmetric{<:Number, <:BandedMatrix}, b, u, Pl,
+            Pr, maxiters::Int, abstol, reltol, verbose::Bool,
+            assumptions::OperatorAssumptions)
+            return nothing
+        end
+    end
+end
+
+end

ext/LinearSolveRecursiveArrayToolsExt.jl

@@ -0,0 +1,12 @@
+module LinearSolveRecursiveArrayToolsExt
+
+using LinearSolve, RecursiveArrayTools
+import LinearSolve: init_cacheval
+
+# Krylov.jl tries to init with `ArrayPartition(undef, ...)`. Avoid hitting that!
+function init_cacheval(alg::LinearSolve.KrylovJL, A, b::ArrayPartition, u, Pl, Pr,
+    maxiters::Int, abstol, reltol, verbose::Bool, ::LinearSolve.OperatorAssumptions)
+    return nothing
+end
+
+end

src/common.jl

@@ -14,7 +14,7 @@ However, in practice this computation is very expensive and thus not possible fo
 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 
+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

src/default.jl

@@ -93,7 +93,7 @@ end
 end
 
 function defaultalg(A::GPUArraysCore.AbstractGPUArray, b, assump::OperatorAssumptions)
-    if assump.condition === OperatorConodition.IllConditioned || !assump.issq
+    if assump.condition === OperatorCondition.IllConditioned || !assump.issq
         DefaultLinearSolver(DefaultAlgorithmChoice.QRFactorization)
     else
         @static if VERSION >= v"1.8-"
@@ -163,7 +163,7 @@ function defaultalg(A, b, assump::OperatorAssumptions)
                     DefaultAlgorithmChoice.GenericLUFactorization
                 elseif VERSION >= v"1.8" && appleaccelerate_isavailable()
                     DefaultAlgorithmChoice.AppleAccelerateLUFactorization
-                elseif (length(b) <= 100 || (isopenblas() && length(b) <= 500) || 
+                elseif (length(b) <= 100 || (isopenblas() && length(b) <= 500) ||
                        (usemkl && length(b) <= 200)) &&
                        (A === nothing ? eltype(b) <: Union{Float32, Float64} :
                         eltype(A) <: Union{Float32, Float64})