Merge pull request #244 from SciML/diagonal

Specialize diagonal matrices

Author: ChrisRackauckas

docs/src/solvers/solvers.md

@@ -84,6 +84,7 @@ customized per-package, details given below describe a subset of important array
 LinearSolve.jl contains some linear solvers built in.
 
 - `SimpleLUFactorization`: a simple LU-factorization implementation without BLAS. Fast for small matrices.
+- `DiagonalFactorization`: a special implementation only for solving `Diagonal` matrices fast.
 
 ### FastLapackInterface.jl
 

src/LinearSolve.jl

@@ -84,7 +84,7 @@ end
 export LUFactorization, SVDFactorization, QRFactorization, GenericFactorization,
        GenericLUFactorization, SimpleLUFactorization, RFLUFactorization,
        UMFPACKFactorization, KLUFactorization, FastLUFactorization, FastQRFactorization,
-       SparspakFactorization
+       SparspakFactorization, DiagonalFactorization
 
 export LinearSolveFunction
 

src/default.jl

@@ -33,6 +33,15 @@ end
 function defaultalg(A::SymTridiagonal, b, ::OperatorAssumptions{true})
     GenericFactorization(; fact_alg = ldlt!)
 end
+function defaultalg(A::Diagonal, b, ::OperatorAssumptions{true})
+    DiagonalFactorization()
+end
+function defaultalg(A::Diagonal, b, ::OperatorAssumptions{false})
+    DiagonalFactorization()
+end
+function defaultalg(A::Diagonal, b, ::OperatorAssumptions{Nothing})
+    DiagonalFactorization()
+end
 
 @static if INCLUDE_SPARSE
     function defaultalg(A::SparseMatrixCSC, b, ::OperatorAssumptions{true})

src/factorization.jl

@@ -374,6 +374,28 @@ function SciMLBase.solve(cache::LinearCache, alg::RFLUFactorization{P, T};
     SciMLBase.build_linear_solution(alg, y, nothing, cache)
 end
 
+## DiagonalFactorization
+
+struct DiagonalFactorization <: AbstractFactorization end
+
+function init_cacheval(alg::DiagonalFactorization, A, b, u, Pl, Pr, maxiters::Int,
+                       abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions)
+    nothing
+end
+
+function SciMLBase.solve(cache::LinearCache, alg::DiagonalFactorization;
+                         kwargs...) where {P, T}
+    A = cache.A
+    if cache.u isa Vector && cache.b isa Vector
+        @simd ivdep for i in eachindex(cache.u)
+            cache.u[i] = A.diag[i] \ cache.b[i]
+        end
+    else
+        cache.u .= A.diag .\ cache.b
+    end
+    SciMLBase.build_linear_solution(alg, cache.u, nothing, cache)
+end
+
 ## FastLAPACKFactorizations
 
 struct WorkspaceAndFactors{W, F}