Vectorized decompression

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]

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

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,

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

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

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