diff --git a/Project.toml b/Project.toml index c6a144b32..a23ab45ef 100644 --- a/Project.toml +++ b/Project.toml @@ -26,6 +26,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SciMLLogging = "a6db7da4-7206-11f0-1eab-35f2a5dbe1d1" SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseColumnPivotedQR = "a57abbd0-fea5-4d57-96be-5e525945e8e4" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" @@ -62,7 +63,6 @@ PartitionedSolvers = "11b65f7f-80ac-401b-9ef2-3db765482d62" PureUMFPACK = "b7e1f0a2-3c4d-4e5f-9a0b-1c2d3e4f5a6b" RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4" STRUMPACK_jll = "86fbd0b9-476f-557c-b766-62c724b42d8c" -SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseMatricesCSR = "a0a7dd2c-ebf4-11e9-1f05-cf50bc540ca1" Sparspak = "e56a9233-b9d6-4f03-8d0f-1825330902ac" SpecializingFactorizations = "fa08b7a1-13d3-4faf-875d-5cbc1520e3f3" diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index b7925bb23..d21e86fcc 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -4,6 +4,7 @@ using LinearSolve using LinearSolve: SciMLLinearSolveAlgorithm, __init, LinearVerbosity, DefaultLinearSolver, DefaultAlgorithmChoice, defaultalg, reinit! using LinearAlgebra +using SparseArrays using ForwardDiff using ForwardDiff: Dual, Partials using SciMLBase @@ -636,10 +637,9 @@ function update_partials_list!(partial_matrix::AbstractVector{T}, list_cache) wh return list_cache end -function update_partials_list!(partial_matrix, list_cache) +function update_partials_list!(partial_matrix::AbstractMatrix{T}, list_cache) where {T} p = length(first(partial_matrix)) m, n = size(partial_matrix) - for k in 1:p for i in 1:m for j in 1:n @@ -655,7 +655,7 @@ function partials_to_list(partial_matrix::AbstractVector{T}) where {T} return [[partial[i] for partial in partial_matrix] for i in p] end -function partials_to_list(partial_matrix) +function partials_to_list(partial_matrix::AbstractMatrix{T}) where {T} p = length(first(partial_matrix)) m, n = size(partial_matrix) res_list = fill(zeros(typeof(partial_matrix[1, 1][1]), m, n), p) @@ -671,4 +671,31 @@ function partials_to_list(partial_matrix) return res_list end +# Specializations for sparse matrices + +function partials_to_list(partial_matrix::SparseMatrixCSC) + nz = nonzeros(partial_matrix) + m, n = size(partial_matrix) + T = eltype(partial_matrix) + p = ForwardDiff.npartials(T) + V = ForwardDiff.valtype(T) # use type for concrete array below in empty-nz case (e.g. all-zero Jacobian at init) + return [SparseMatrixCSC(m, n, copy(partial_matrix.colptr), copy(partial_matrix.rowval), + V[nz[i][k] for i in eachindex(nz)]) for k in 1:p] +end + +function update_partials_list!(partial_matrix::SparseMatrixCSC, list_cache) + nz = nonzeros(partial_matrix) + if length(nz) != length(nonzeros(first(list_cache))) # TODO: more precise? + list_cache .= partials_to_list(partial_matrix) # sparsity pattern changed + else + for k in eachindex(list_cache) + nz_k = nonzeros(list_cache[k]) + @inbounds for i in eachindex(nz, nz_k) + nz_k[i] = nz[i][k] + end + end + end + return list_cache +end + end