|
| 1 | +module LinearSolveSparseArraysExt |
| 2 | + |
| 3 | +using LinearSolve, LinearAlgebra |
| 4 | +using SparseArrays |
| 5 | +using SparseArrays: AbstractSparseMatrixCSC, nonzeros, rowvals, getcolptr |
| 6 | + |
| 7 | +# Can't `using KLU` because cannot have a dependency in there without |
| 8 | +# requiring the user does `using KLU` |
| 9 | +# But there's no reason to require it because SparseArrays will already |
| 10 | +# load SuiteSparse and thus all of the underlying KLU code |
| 11 | +include("../src/KLU/klu.jl") |
| 12 | + |
| 13 | +LinearSolve.issparsematrixcsc(A::AbstractSparseMatrixCSC) = true |
| 14 | +LinearSolve.issparsematrix(A::AbstractSparseArray) = true |
| 15 | +function LinearSolve.make_SparseMatrixCSC(A::AbstractSparseArray) |
| 16 | + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), nonzeros(A)) |
| 17 | +end |
| 18 | +function LinearSolve.makeempty_SparaseMatrixCSC(A::AbstractSparseArray) |
| 19 | + SparseMatrixCSC(0, 0, [1], Int[], eltype(A)[]) |
| 20 | +end |
| 21 | + |
| 22 | +function LinearSolve.init_cacheval(alg::RFLUFactorization, |
| 23 | + A::Union{AbstractSparseArray, LinearSolve.SciMLOperators.AbstractSciMLOperator}, b, u, Pl, Pr, |
| 24 | + maxiters::Int, |
| 25 | + abstol, reltol, verbose::Bool, assumptions::OperatorAssumptions) |
| 26 | + nothing, nothing |
| 27 | +end |
| 28 | + |
| 29 | +function LinearSolve.init_cacheval( |
| 30 | + alg::QRFactorization, A::Symmetric{<:Number, <:SparseMatrixCSC}, b, u, Pl, Pr, |
| 31 | + maxiters::Int, abstol, reltol, verbose::Bool, |
| 32 | + assumptions::OperatorAssumptions) |
| 33 | + return nothing |
| 34 | +end |
| 35 | + |
| 36 | +function LinearSolve.handle_sparsematrixcsc_lu(A::AbstractSparseMatrixCSC) |
| 37 | + lu(SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), nonzeros(A)), |
| 38 | + check = false) |
| 39 | +end |
| 40 | + |
| 41 | +function LinearSolve.defaultalg( |
| 42 | + A::Symmetric{<:Number, <:SparseMatrixCSC}, b, ::OperatorAssumptions{Bool}) |
| 43 | + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.CHOLMODFactorization) |
| 44 | +end |
| 45 | + |
| 46 | +function LinearSolve.defaultalg(A::AbstractSparseMatrixCSC{Tv, Ti}, b, |
| 47 | + assump::OperatorAssumptions{Bool}) where {Tv, Ti} |
| 48 | + if assump.issq |
| 49 | + DefaultLinearSolver(DefaultAlgorithmChoice.SparspakFactorization) |
| 50 | + else |
| 51 | + error("Generic number sparse factorization for non-square is not currently handled") |
| 52 | + end |
| 53 | +end |
| 54 | + |
| 55 | +function LinearSolve.init_cacheval(alg::GenericFactorization, |
| 56 | + A::Union{Hermitian{T, <:SparseMatrixCSC}, |
| 57 | + Symmetric{T, <:SparseMatrixCSC}}, b, u, Pl, Pr, |
| 58 | + maxiters::Int, abstol, reltol, verbose::Bool, |
| 59 | + assumptions::OperatorAssumptions) where {T} |
| 60 | + newA = copy(convert(AbstractMatrix, A)) |
| 61 | + LinearSolve.do_factorization(alg, newA, b, u) |
| 62 | +end |
| 63 | + |
| 64 | +const PREALLOCATED_UMFPACK = SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC(0, 0, [1], |
| 65 | + Int[], Float64[])) |
| 66 | + |
| 67 | +function LinearSolve.init_cacheval( |
| 68 | + alg::UMFPACKFactorization, A::SparseMatrixCSC{Float64, Int}, b, u, |
| 69 | + Pl, Pr, |
| 70 | + maxiters::Int, abstol, reltol, |
| 71 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 72 | + PREALLOCATED_UMFPACK |
| 73 | +end |
| 74 | + |
| 75 | +function LinearSolve.init_cacheval( |
| 76 | + alg::UMFPACKFactorization, A::AbstractSparseArray, b, u, Pl, Pr, |
| 77 | + maxiters::Int, abstol, |
| 78 | + reltol, |
| 79 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 80 | + A = convert(AbstractMatrix, A) |
| 81 | + zerobased = SparseArrays.getcolptr(A)[1] == 0 |
| 82 | + return SparseArrays.UMFPACK.UmfpackLU(SparseMatrixCSC(size(A)..., getcolptr(A), |
| 83 | + rowvals(A), nonzeros(A))) |
| 84 | +end |
| 85 | + |
| 86 | +function SciMLBase.solve!( |
| 87 | + cache::LinearSolve.LinearCache, alg::UMFPACKFactorization; kwargs...) |
| 88 | + A = cache.A |
| 89 | + A = convert(AbstractMatrix, A) |
| 90 | + if cache.isfresh |
| 91 | + cacheval = LinearSolve.@get_cacheval(cache, :UMFPACKFactorization) |
| 92 | + if alg.reuse_symbolic |
| 93 | + # Caches the symbolic factorization: https://github.com/JuliaLang/julia/pull/33738 |
| 94 | + if alg.check_pattern && pattern_changed(cacheval, A) |
| 95 | + fact = lu( |
| 96 | + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), |
| 97 | + nonzeros(A)), |
| 98 | + check = false) |
| 99 | + else |
| 100 | + fact = lu!(cacheval, |
| 101 | + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), |
| 102 | + nonzeros(A)), check = false) |
| 103 | + end |
| 104 | + else |
| 105 | + fact = lu(SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), nonzeros(A)), |
| 106 | + check = false) |
| 107 | + end |
| 108 | + cache.cacheval = fact |
| 109 | + cache.isfresh = false |
| 110 | + end |
| 111 | + |
| 112 | + F = LinearSolve.@get_cacheval(cache, :UMFPACKFactorization) |
| 113 | + if F.status == SparseArrays.UMFPACK.UMFPACK_OK |
| 114 | + y = ldiv!(cache.u, F, cache.b) |
| 115 | + SciMLBase.build_linear_solution(alg, y, nothing, cache) |
| 116 | + else |
| 117 | + SciMLBase.build_linear_solution( |
| 118 | + alg, cache.u, nothing, cache; retcode = ReturnCode.Infeasible) |
| 119 | + end |
| 120 | +end |
| 121 | + |
| 122 | +const PREALLOCATED_KLU = KLU.KLUFactorization(SparseMatrixCSC(0, 0, [1], Int[], |
| 123 | + Float64[])) |
| 124 | + |
| 125 | +function LinearSolve.init_cacheval( |
| 126 | + alg::KLUFactorization, A::SparseMatrixCSC{Float64, Int}, b, u, Pl, |
| 127 | + Pr, |
| 128 | + maxiters::Int, abstol, reltol, |
| 129 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 130 | + PREALLOCATED_KLU |
| 131 | +end |
| 132 | + |
| 133 | +function LinearSolve.init_cacheval( |
| 134 | + alg::KLUFactorization, A::AbstractSparseArray, b, u, Pl, Pr, |
| 135 | + maxiters::Int, abstol, |
| 136 | + reltol, |
| 137 | + verbose::Bool, assumptions::OperatorAssumptions) |
| 138 | + A = convert(AbstractMatrix, A) |
| 139 | + return KLU.KLUFactorization(SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), |
| 140 | + nonzeros(A))) |
| 141 | +end |
| 142 | + |
| 143 | +# TODO: guard this against errors |
| 144 | +function SciMLBase.solve!(cache::LinearSolve.LinearCache, alg::KLUFactorization; kwargs...) |
| 145 | + A = cache.A |
| 146 | + A = convert(AbstractMatrix, A) |
| 147 | + if cache.isfresh |
| 148 | + cacheval = LinearSolve.@get_cacheval(cache, :KLUFactorization) |
| 149 | + if alg.reuse_symbolic |
| 150 | + if alg.check_pattern && pattern_changed(cacheval, A) |
| 151 | + fact = KLU.klu( |
| 152 | + SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), |
| 153 | + nonzeros(A)), |
| 154 | + check = false) |
| 155 | + else |
| 156 | + fact = KLU.klu!(cacheval, nonzeros(A), check = false) |
| 157 | + end |
| 158 | + else |
| 159 | + # New fact each time since the sparsity pattern can change |
| 160 | + # and thus it needs to reallocate |
| 161 | + fact = KLU.klu(SparseMatrixCSC(size(A)..., getcolptr(A), rowvals(A), |
| 162 | + nonzeros(A))) |
| 163 | + end |
| 164 | + cache.cacheval = fact |
| 165 | + cache.isfresh = false |
| 166 | + end |
| 167 | + F = LinearSolve.@get_cacheval(cache, :KLUFactorization) |
| 168 | + if F.common.status == KLU.KLU_OK |
| 169 | + y = ldiv!(cache.u, F, cache.b) |
| 170 | + SciMLBase.build_linear_solution(alg, y, nothing, cache) |
| 171 | + else |
| 172 | + SciMLBase.build_linear_solution( |
| 173 | + alg, cache.u, nothing, cache; retcode = ReturnCode.Infeasible) |
| 174 | + end |
| 175 | +end |
| 176 | + |
| 177 | +const PREALLOCATED_CHOLMOD = cholesky(SparseMatrixCSC(0, 0, [1], Int[], Float64[])) |
| 178 | + |
| 179 | +function LinearSolve.init_cacheval(alg::CHOLMODFactorization, |
| 180 | + A::Union{SparseMatrixCSC{T, Int}, Symmetric{T, SparseMatrixCSC{T, Int}}}, b, u, |
| 181 | + Pl, Pr, |
| 182 | + maxiters::Int, abstol, reltol, |
| 183 | + verbose::Bool, assumptions::OperatorAssumptions) where {T <: |
| 184 | + Union{Float32, Float64}} |
| 185 | + PREALLOCATED_CHOLMOD |
| 186 | +end |
| 187 | + |
| 188 | +function LinearSolve.init_cacheval(alg::NormalCholeskyFactorization, |
| 189 | + A::Union{AbstractSparseArray, LinearSolve.GPUArraysCore.AnyGPUArray, |
| 190 | + Symmetric{<:Number, <:AbstractSparseArray}}, b, u, Pl, Pr, |
| 191 | + maxiters::Int, abstol, reltol, verbose::Bool, |
| 192 | + assumptions::OperatorAssumptions) |
| 193 | + LinearSolve.ArrayInterface.cholesky_instance(convert(AbstractMatrix, A)) |
| 194 | +end |
| 195 | + |
| 196 | +# Specialize QR for the non-square case |
| 197 | +# Missing ldiv! definitions: https://github.com/JuliaSparse/SparseArrays.jl/issues/242 |
| 198 | +function LinearSolve._ldiv!(x::Vector, |
| 199 | + A::Union{SparseArrays.QR, LinearAlgebra.QRCompactWY, |
| 200 | + SparseArrays.SPQR.QRSparse, |
| 201 | + SparseArrays.CHOLMOD.Factor}, b::Vector) |
| 202 | + x .= A \ b |
| 203 | +end |
| 204 | + |
| 205 | +function LinearSolve._ldiv!(x::AbstractVector, |
| 206 | + A::Union{SparseArrays.QR, LinearAlgebra.QRCompactWY, |
| 207 | + SparseArrays.SPQR.QRSparse, |
| 208 | + SparseArrays.CHOLMOD.Factor}, b::AbstractVector) |
| 209 | + x .= A \ b |
| 210 | +end |
| 211 | + |
| 212 | +# Ambiguity removal |
| 213 | +function LinearSolve._ldiv!(::LinearSolve.SVector, |
| 214 | + A::Union{SparseArrays.CHOLMOD.Factor, LinearAlgebra.QR, |
| 215 | + LinearAlgebra.QRCompactWY, SparseArrays.SPQR.QRSparse}, |
| 216 | + b::AbstractVector) |
| 217 | + (A \ b) |
| 218 | +end |
| 219 | +function LinearSolve._ldiv!(::LinearSolve.SVector, |
| 220 | + A::Union{SparseArrays.CHOLMOD.Factor, LinearAlgebra.QR, |
| 221 | + LinearAlgebra.QRCompactWY, SparseArrays.SPQR.QRSparse}, |
| 222 | + b::LinearSolve.SVector) |
| 223 | + (A \ b) |
| 224 | +end |
| 225 | + |
| 226 | +function pattern_changed(fact, A::SparseArrays.SparseMatrixCSC) |
| 227 | + !(SparseArrays.decrement(SparseArrays.getcolptr(A)) == |
| 228 | + fact.colptr && SparseArrays.decrement(SparseArrays.getrowval(A)) == |
| 229 | + fact.rowval) |
| 230 | +end |
| 231 | + |
| 232 | +function LinearSolve.defaultalg( |
| 233 | + A::AbstractSparseMatrixCSC{<:Union{Float64, ComplexF64}, Ti}, b, |
| 234 | + assump::OperatorAssumptions{Bool}) where {Ti} |
| 235 | + if assump.issq |
| 236 | + if length(b) <= 10_000 && length(nonzeros(A)) / length(A) < 2e-4 |
| 237 | + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KLUFactorization) |
| 238 | + else |
| 239 | + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.UMFPACKFactorization) |
| 240 | + end |
| 241 | + else |
| 242 | + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.QRFactorization) |
| 243 | + end |
| 244 | +end |
| 245 | + |
| 246 | +LinearSolve.PrecompileTools.@compile_workload begin |
| 247 | + A = sprand(4, 4, 0.3) + I |
| 248 | + b = rand(4) |
| 249 | + prob = LinearProblem(A, b) |
| 250 | + sol = solve(prob, KLUFactorization()) |
| 251 | + sol = solve(prob, UMFPACKFactorization()) |
| 252 | +end |
| 253 | + |
| 254 | +end |
0 commit comments