diff --git a/Project.toml b/Project.toml index b9c7add4..bebbbe47 100644 --- a/Project.toml +++ b/Project.toml @@ -13,8 +13,10 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Requires = "ae029012-a4dd-5104-9daa-d747884805df" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f" +ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" [compat] +ArrayInterface = "1.1" julia = "1" [extras] diff --git a/README.md b/README.md index 19c344c3..b4c1a15d 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,9 @@ gmres!(res,J,v) ### Matrix Coloring +This library extends the common `ArrayInterface.matrix_colors` function to allow +for coloring sparse matrices using graphical techniques. + Matrix coloring allows you to reduce the number of times finite differencing requires an `f` call to `maximum(colors)+1`, or reduces automatic differentiation to using `maximum(colors)` partials. Since normally these values are `length(x)`, @@ -101,7 +104,8 @@ this can be significant savings. The API for computing the color vector is: ```julia -matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color(); partition_by_rows::Bool = false) +matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color(); + partition_by_rows::Bool = false) ``` The first argument is the abstract matrix which represents the sparsity pattern diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index 7e3a8b42..fc1e8a0e 100644 --- a/src/SparseDiffTools.jl +++ b/src/SparseDiffTools.jl @@ -9,11 +9,12 @@ using Requires using VertexSafeGraphs using LinearAlgebra -using SparseArrays +using SparseArrays, ArrayInterface using BlockBandedMatrices: blocksize, nblocks using ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD +using ArrayInterface: matrix_colors export contract_color, greedy_d1, diff --git a/src/coloring/high_level.jl b/src/coloring/high_level.jl index 08978e03..83a16c44 100644 --- a/src/coloring/high_level.jl +++ b/src/coloring/high_level.jl @@ -1,9 +1,9 @@ -abstract type ColoringAlgorithm end -struct GreedyD1Color <: ColoringAlgorithm end -struct BacktrackingColor <: ColoringAlgorithm end -struct ContractionColor <: ColoringAlgorithm end -struct GreedyStar1Color <: ColoringAlgorithm end -struct GreedyStar2Color <: ColoringAlgorithm end +abstract type SparseDiffToolsColoringAlgorithm <: ArrayInterface.ColoringAlgorithm end +struct GreedyD1Color <: SparseDiffToolsColoringAlgorithm end +struct BacktrackingColor <: SparseDiffToolsColoringAlgorithm end +struct ContractionColor <: SparseDiffToolsColoringAlgorithm end +struct GreedyStar1Color <: SparseDiffToolsColoringAlgorithm end +struct GreedyStar2Color <: SparseDiffToolsColoringAlgorithm end """ matrix_colors(A,alg::ColoringAlgorithm = GreedyD1Color()) @@ -13,70 +13,8 @@ struct GreedyStar2Color <: ColoringAlgorithm end The coloring defaults to a greedy distance-1 coloring. """ -function matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color(); partition_by_rows::Bool = false) +function ArrayInterface.matrix_colors(A::AbstractMatrix,alg::SparseDiffToolsColoringAlgorithm = GreedyD1Color(); partition_by_rows::Bool = false) _A = A isa SparseMatrixCSC ? A : sparse(A) # Avoid the copy A_graph = matrix2graph(_A, partition_by_rows) color_graph(A_graph,alg) end - -""" - matrix_colors(A::Union{Array,UpperTriangular,LowerTriangular}) - - The color vector for dense matrix and triangular matrix is simply - `[1,2,3,...,size(A,2)]` -""" -function matrix_colors(A::Union{Array,UpperTriangular,LowerTriangular}) - eachindex(1:size(A,2)) # Vector size matches number of rows -end - -function _cycle(repetend,len) - repeat(repetend,div(len,length(repetend))+1)[1:len] -end - -function matrix_colors(A::Diagonal) - fill(1,size(A,2)) -end - -function matrix_colors(A::Bidiagonal) - _cycle(1:2,size(A,2)) -end - -function matrix_colors(A::Union{Tridiagonal,SymTridiagonal}) - _cycle(1:3,size(A,2)) -end - -function matrix_colors(A::BandedMatrix) - u,l=bandwidths(A) - width=u+l+1 - _cycle(1:width,size(A,2)) -end - -function matrix_colors(A::BlockBandedMatrix) - l,u=blockbandwidths(A) - blockwidth=l+u+1 - nblock=nblocks(A,2) - cols=[blocksize(A,(1,i))[2] for i in 1:nblock] - blockcolors=_cycle(1:blockwidth,nblock) - #the reserved number of colors of a block is the maximum length of columns of blocks with the same block color - ncolors=[maximum(cols[i:blockwidth:nblock]) for i in 1:blockwidth] - endinds=cumsum(ncolors) - startinds=[endinds[i]-ncolors[i]+1 for i in 1:blockwidth] - colors=[(startinds[blockcolors[i]]:endinds[blockcolors[i]])[1:cols[i]] for i in 1:nblock] - vcat(colors...) -end - -function matrix_colors(A::BandedBlockBandedMatrix) - l,u=blockbandwidths(A) - lambda,mu=subblockbandwidths(A) - blockwidth=l+u+1 - subblockwidth=lambda+mu+1 - nblock=nblocks(A,2) - cols=[blocksize(A,(1,i))[2] for i in 1:nblock] - blockcolors=_cycle(1:blockwidth,nblock) - #the reserved number of colors of a block is the min of subblockwidth and the largest length of columns of blocks with the same block color - ncolors=[min(subblockwidth,maximum(cols[i:blockwidth:nblock])) for i in 1:min(blockwidth,nblock)] - endinds=cumsum(ncolors) - startinds=[endinds[i]-ncolors[i]+1 for i in 1:min(blockwidth,nblock)] - colors=[_cycle(startinds[blockcolors[i]]:endinds[blockcolors[i]],cols[i]) for i in 1:nblock] - vcat(colors...) -end diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index e00d9e53..468fe7d7 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -21,7 +21,7 @@ getsize(N::Integer) = N function ForwardColorJacCache(f,x,_chunksize = nothing; dx = nothing, color=1:length(x), - sparsity::Union{SparseMatrixCSC,Nothing}=nothing) + sparsity::Union{AbstractArray,Nothing}=nothing) if _chunksize === nothing chunksize = default_chunk_size(maximum(color)) @@ -81,7 +81,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, x::AbstractArray{<:Number}; dx = nothing, color = eachindex(x), - sparsity = J isa SparseMatrixCSC ? J : nothing) + sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing) forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,color=color,sparsity=sparsity)) end @@ -98,18 +98,37 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, sparsity = jac_cache.sparsity color_i = 1 chunksize = length(first(first(jac_cache.p))) + fill!(J, zero(eltype(J))) - for i in 1:length(p) + for i in eachindex(p) partial_i = p[i] t .= Dual{typeof(f)}.(x, partial_i) f(fx,t) - if sparsity isa SparseMatrixCSC - rows_index, cols_index, val = findnz(sparsity) + if ArrayInterface.has_sparsestruct(sparsity) + rows_index, cols_index = ArrayInterface.findstructralnz(sparsity) for j in 1:chunksize dx .= partials.(fx, j) - for k in 1:length(cols_index) - if color[cols_index[k]] == color_i - J[rows_index[k], cols_index[k]] = dx[rows_index[k]] + if ArrayInterface.fast_scalar_indexing(dx) + for k in 1:length(cols_index) + if color[cols_index[k]] == color_i + if J isa SparseMatrixCSC + J.nzval[k] = dx[rows_index[k]] + else + J[rows_index[k],cols_index[k]] = dx[rows_index[k]] + end + end + end + else + #= + J.nzval[rows_index] .+= (color[cols_index] .== color_i) .* dx[rows_index] + or + J[rows_index, cols_index] .+= (color[cols_index] .== color_i) .* dx[rows_index] + += means requires a zero'd out start + =# + if J isa SparseMatrixCSC + @. setindex!((J.nzval,),getindex((J.nzval,),rows_index) + (getindex((color,),cols_index) == color_i) * getindex((dx,),rows_index),rows_index) + else + @. setindex!((J,),getindex((J,),rows_index, cols_index) + (getindex((color,),cols_index) == color_i) * getindex((dx,),rows_index),rows_index, cols_index) end end color_i += 1 diff --git a/test/test_ad.jl b/test/test_ad.jl index 3055ce4f..29207476 100644 --- a/test/test_ad.jl +++ b/test/test_ad.jl @@ -1,6 +1,7 @@ using SparseDiffTools using ForwardDiff: Dual, jacobian using SparseArrays, Test +using LinearAlgebra fcalls = 0 function f(dx,x) @@ -56,3 +57,7 @@ jac_cache = ForwardColorJacCache(f,x,color = repeat(1:3,10), sparsity = _J1) forwarddiff_color_jacobian!(_denseJ1, f, x, jac_cache) @test _denseJ1 ≈ J @test fcalls == 1 + +_Jt = similar(Tridiagonal(J)) +forwarddiff_color_jacobian!(_Jt, f, x, color = repeat(1:3,10), sparsity = _Jt) +@test _Jt ≈ J