Merge pull request #423 from avik-pal/ap/almost_banded_matrices

Add support for almost banded matrix

Project.toml

@@ -1,7 +1,7 @@
 name = "LinearSolve"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 authors = ["SciML"]
-version = "2.19.0"
+version = "2.20.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -34,6 +34,7 @@ BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FastAlmostBandedMatrices = "9d29842c-ecb8-4973-b1e9-a27b1157504e"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -54,6 +55,7 @@ LinearSolveKrylovKitExt = "KrylovKit"
 LinearSolveMetalExt = "Metal"
 LinearSolvePardisoExt = "Pardiso"
 LinearSolveRecursiveArrayToolsExt = "RecursiveArrayTools"
+LinearSolveFastAlmostBandedMatricesExt = ["FastAlmostBandedMatrices"]
 
 [compat]
 ArrayInterface = "7.4.11"
@@ -92,6 +94,7 @@ julia = "1.9"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FastAlmostBandedMatrices = "9d29842c-ecb8-4973-b1e9-a27b1157504e"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
@@ -110,4 +113,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "FiniteDiff", "BandedMatrices"]
+test = ["Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "FiniteDiff", "BandedMatrices", "FastAlmostBandedMatrices"]

ext/LinearSolveFastAlmostBandedMatricesExt.jl

@@ -0,0 +1,33 @@
+module LinearSolveFastAlmostBandedMatricesExt
+
+using FastAlmostBandedMatrices, LinearAlgebra, LinearSolve
+import LinearSolve: defaultalg,
+    do_factorization, init_cacheval, DefaultLinearSolver, DefaultAlgorithmChoice
+
+function defaultalg(A::AlmostBandedMatrix, b, oa::OperatorAssumptions)
+    if oa.issq
+        return DefaultLinearSolver(DefaultAlgorithmChoice.DirectLdiv!)
+    else
+        return DefaultLinearSolver(DefaultAlgorithmChoice.QRFactorization)
+    end
+end
+
+# For BandedMatrix
+for alg in (:SVDFactorization, :MKLLUFactorization, :DiagonalFactorization,
+    :SparspakFactorization, :KLUFactorization, :UMFPACKFactorization,
+    :GenericLUFactorization, :RFLUFactorization, :BunchKaufmanFactorization,
+    :CHOLMODFactorization, :NormalCholeskyFactorization, :LDLtFactorization,
+    :AppleAccelerateLUFactorization, :CholeskyFactorization, :LUFactorization)
+    @eval begin
+        function init_cacheval(::$(alg), ::AlmostBandedMatrix, b, u, Pl, Pr, maxiters::Int,
+                abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
+            return nothing
+        end
+    end
+end
+
+function do_factorization(alg::QRFactorization, A::AlmostBandedMatrix, b, u)
+    return alg.inplace ? qr!(A) : qr(A)
+end
+
+end

src/factorization.jl

@@ -12,6 +12,8 @@ _ldiv!(x, A, b) = ldiv!(x, A, b)
 
 function _ldiv!(x::Vector, A::Factorization, b::Vector)
     # workaround https://github.com/JuliaLang/julia/issues/43507
+    # Fallback if working with non-square matrices
+    length(x) != length(b) && return ldiv!(x, A, b)
     copyto!(x, b)
     ldiv!(A, x)
 end

test/banded.jl

@@ -1,4 +1,4 @@
-using BandedMatrices, LinearAlgebra, LinearSolve, Test
+using FastAlmostBandedMatrices, BandedMatrices, LinearAlgebra, LinearSolve, Test
 
 # Square Case
 n = 8
@@ -16,6 +16,12 @@ sol1 = solve(LinearProblem(A1, b1; u0 = x1))
 sol2 = solve(LinearProblem(A2, b2; u0 = x2))
 @test sol2.u ≈ A2 \ b2
 
+A = AlmostBandedMatrix(BandedMatrix(fill(2.0, n, n), (1, 1)), fill(3.0, 2, n))
+A[band(0)] .+= 1:n
+
+sol1ab = solve(LinearProblem(A, b; u0 = x1))
+@test sol1ab.u ≈ Matrix(A) \ b
+
 # Square Symmetric
 A1s = Symmetric(A1)
 A2s = Symmetric(A2)
@@ -31,8 +37,18 @@ b = rand(8)
 
 @test_throws ErrorException solve(LinearProblem(A, b)).u
 
+A = AlmostBandedMatrix(BandedMatrix(fill(2.0, n - 2, n), (1, 1)), fill(3.0, 2, n))
+A[band(0)] .+= 1:(n - 2)
+
+@test_throws ErrorException solve(LinearProblem(A, b)).u
+
 # Overdetermined
 A = BandedMatrix(ones(10, 8), (2, 0))
 b = rand(10)
 
 @test_nowarn solve(LinearProblem(A, b))
+
+A = AlmostBandedMatrix(BandedMatrix(fill(2.0, n + 2, n), (1, 1)), fill(3.0, 2, n))
+A[band(0)] .+= 1:n
+
+@test_nowarn solve(LinearProblem(A, b))