|
| 1 | +module LinearSolveSparseArraysExt |
| 2 | + |
| 3 | +using LinearSolve |
| 4 | +import LinearSolve: SciMLBase, LinearAlgebra, PrecompileTools, init_cacheval |
| 5 | +using LinearSolve: DefaultLinearSolver, DefaultAlgorithmChoice |
| 6 | +using SparseArrays |
| 7 | +using SparseArrays: AbstractSparseMatrixCSC, nonzeros, rowvals, getcolptr |
| 8 | + |
| 9 | +# Specialize QR for the non-square case |
| 10 | +# Missing ldiv! definitions: https://github.com/JuliaSparse/SparseArrays.jl/issues/242 |
| 11 | +function LinearSolve._ldiv!(x::Vector, |
| 12 | + A::Union{SparseArrays.QR, LinearAlgebra.QRCompactWY, |
| 13 | + SparseArrays.SPQR.QRSparse, |
| 14 | + SparseArrays.CHOLMOD.Factor}, b::Vector) |
| 15 | +x .= A \ b |
| 16 | +end |
| 17 | + |
| 18 | +function LinearSolve._ldiv!(x::AbstractVector, |
| 19 | + A::Union{SparseArrays.QR, LinearAlgebra.QRCompactWY, |
| 20 | + SparseArrays.SPQR.QRSparse, |
| 21 | + SparseArrays.CHOLMOD.Factor}, b::AbstractVector) |
| 22 | +x .= A \ b |
| 23 | +end |
| 24 | + |
| 25 | +# Ambiguity removal |
| 26 | +function LinearSolve._ldiv!(::SVector, |
| 27 | + A::Union{SparseArrays.CHOLMOD.Factor, LinearAlgebra.QR, |
| 28 | + LinearAlgebra.QRCompactWY, SparseArrays.SPQR.QRSparse}, |
| 29 | + b::AbstractVector) |
| 30 | +(A \ b) |
| 31 | +end |
| 32 | +function LinearSolve._ldiv!(::SVector, |
| 33 | + A::Union{SparseArrays.CHOLMOD.Factor, LinearAlgebra.QR, |
| 34 | + LinearAlgebra.QRCompactWY, SparseArrays.SPQR.QRSparse}, |
| 35 | + b::SVector) |
| 36 | +(A \ b) |
| 37 | +end |
| 38 | + |
| 39 | +function LinearSolve.pattern_changed(fact, A::SparseArrays.SparseMatrixCSC) |
| 40 | + !(SparseArrays.decrement(SparseArrays.getcolptr(A)) == |
| 41 | + fact.colptr && SparseArrays.decrement(SparseArrays.getrowval(A)) == |
| 42 | + fact.rowval) |
| 43 | +end |
| 44 | + |
| 45 | +const PREALLOCATED_UMFPACK = SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC(0, 0, [1], |
| 46 | + Int[], Float64[])) |
| 47 | + |
| 48 | +function init_cacheval(alg::UMFPACKFactorization, |
| 49 | + A, b, u, Pl, Pr, |
| 50 | + maxiters::Int, abstol, reltol, |
| 51 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 52 | + nothing |
| 53 | +end |
| 54 | + |
| 55 | +function init_cacheval(alg::UMFPACKFactorization, A::SparseMatrixCSC{Float64, Int}, b, u, |
| 56 | + Pl, Pr, |
| 57 | + maxiters::Int, abstol, reltol, |
| 58 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 59 | + PREALLOCATED_UMFPACK |
| 60 | +end |
| 61 | + |
| 62 | +function init_cacheval(alg::UMFPACKFactorization, A::AbstractSparseArray, b, u, Pl, Pr, |
| 63 | + maxiters::Int, abstol, |
| 64 | + reltol, |
| 65 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 66 | + A = convert(AbstractMatrix, A) |
| 67 | + zerobased = SparseArrays.getcolptr(A)[1] == 0 |
| 68 | + return SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC(size(A)..., getcolptr(A), |
| 69 | + rowvals(A), nonzeros(A))) |
| 70 | +end |
| 71 | + |
| 72 | +function SciMLBase.solve!(cache::LinearCache, alg::UMFPACKFactorization; kwargs...) |
| 73 | + A = cache.A |
| 74 | + A = convert(AbstractMatrix, A) |
| 75 | + if cache.isfresh |
| 76 | + cacheval = @get_cacheval(cache, :UMFPACKFactorization) |
| 77 | + if alg.reuse_symbolic |
| 78 | + # Caches the symbolic factorization: https://github.com/JuliaLang/julia/pull/33738 |
| 79 | + if alg.check_pattern && pattern_changed(cacheval, A) |
| 80 | + fact = lu( |
| 81 | + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), |
| 82 | + nonzeros(A)), |
| 83 | + check = false) |
| 84 | + else |
| 85 | + fact = lu!(cacheval, |
| 86 | + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), |
| 87 | + nonzeros(A)), check = false) |
| 88 | + end |
| 89 | + else |
| 90 | + fact = lu(SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), nonzeros(A)), |
| 91 | + check = false) |
| 92 | + end |
| 93 | + cache.cacheval = fact |
| 94 | + cache.isfresh = false |
| 95 | + end |
| 96 | + |
| 97 | + F = @get_cacheval(cache, :UMFPACKFactorization) |
| 98 | + if F.status == SparseArrays.UMFPACK.UMFPACK_OK |
| 99 | + y = ldiv!(cache.u, F, cache.b) |
| 100 | + SciMLBase.build_linear_solution(alg, y, nothing, cache) |
| 101 | + else |
| 102 | + SciMLBase.build_linear_solution( |
| 103 | + alg, cache.u, nothing, cache; retcode = ReturnCode.Infeasible) |
| 104 | + end |
| 105 | +end |
| 106 | + |
| 107 | +const PREALLOCATED_CHOLMOD = cholesky(SparseMatrixCSC(0, 0, [1], Int[], Float64[])) |
| 108 | + |
| 109 | +function init_cacheval(alg::CHOLMODFactorization, |
| 110 | + A, b, u, Pl, Pr, |
| 111 | + maxiters::Int, abstol, reltol, |
| 112 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 113 | + nothing |
| 114 | +end |
| 115 | + |
| 116 | +function init_cacheval(alg::CHOLMODFactorization, |
| 117 | + A::Union{SparseMatrixCSC{T, Int}, Symmetric{T, SparseMatrixCSC{T, Int}}}, b, u, |
| 118 | + Pl, Pr, |
| 119 | + maxiters::Int, abstol, reltol, |
| 120 | + verbose::Bool, assumptions::OperatorAssumptions) where {T <: |
| 121 | + Union{Float32, Float64}} |
| 122 | + PREALLOCATED_CHOLMOD |
| 123 | +end |
| 124 | + |
| 125 | +function SciMLBase.solve!(cache::LinearCache, alg::CHOLMODFactorization; kwargs...) |
| 126 | + A = cache.A |
| 127 | + A = convert(AbstractMatrix, A) |
| 128 | + |
| 129 | + if cache.isfresh |
| 130 | + cacheval = @get_cacheval(cache, :CHOLMODFactorization) |
| 131 | + fact = cholesky(A; check = false) |
| 132 | + if !LinearAlgebra.issuccess(fact) |
| 133 | + ldlt!(fact, A; check = false) |
| 134 | + end |
| 135 | + cache.cacheval = fact |
| 136 | + cache.isfresh = false |
| 137 | + end |
| 138 | + |
| 139 | + cache.u .= @get_cacheval(cache, :CHOLMODFactorization) \ cache.b |
| 140 | + SciMLBase.build_linear_solution(alg, cache.u, nothing, cache) |
| 141 | +end |
| 142 | + |
| 143 | +function LinearSolve.defaultalg( |
| 144 | + A::Symmetric{<:Number, <:SparseMatrixCSC}, b, ::OperatorAssumptions{Bool}) |
| 145 | + DefaultLinearSolver(DefaultAlgorithmChoice.CHOLMODFactorization) |
| 146 | +end |
| 147 | + |
| 148 | +function LinearSolve.defaultalg(A::AbstractSparseMatrixCSC{Tv, Ti}, b, |
| 149 | + assump::OperatorAssumptions{Bool}) where {Tv, Ti} |
| 150 | + if assump.issq |
| 151 | + DefaultLinearSolver(DefaultAlgorithmChoice.SparspakFactorization) |
| 152 | + else |
| 153 | + error("Generic number sparse factorization for non-square is not currently handled") |
| 154 | + end |
| 155 | +end |
| 156 | + |
| 157 | +function LinearSolve.defaultalg(A::AbstractSparseMatrixCSC{<:Union{Float64, ComplexF64}, Ti}, b, |
| 158 | + assump::OperatorAssumptions{Bool}) where {Ti} |
| 159 | + if assump.issq |
| 160 | + if length(b) <= 10_000 && length(nonzeros(A)) / length(A) < 2e-4 |
| 161 | + DefaultLinearSolver(DefaultAlgorithmChoice.KLUFactorization) |
| 162 | + else |
| 163 | + DefaultLinearSolver(DefaultAlgorithmChoice.UMFPACKFactorization) |
| 164 | + end |
| 165 | + else |
| 166 | + DefaultLinearSolver(DefaultAlgorithmChoice.QRFactorization) |
| 167 | + end |
| 168 | +end |
| 169 | + |
| 170 | +PrecompileTools.@compile_workload begin |
| 171 | + A = sprand(4, 4, 0.3) + I |
| 172 | + b = rand(4) |
| 173 | + prob = LinearProblem(A, b) |
| 174 | + sol = solve(prob, KLUFactorization()) |
| 175 | + sol = solve(prob, UMFPACKFactorization()) |
| 176 | +end |
| 177 | + |
| 178 | +end |
0 commit comments