add sparsity to forwarddiff_color_jacobian! and ForwardColorJacCache

Author: huanglangwen

src/differentiation/compute_jacobian_ad.jl

@@ -1,9 +1,10 @@
-struct ForwardColorJacCache{T,T2,T3,T4,T5}
+struct ForwardColorJacCache{T,T2,T3,T4,T5,T6}
     t::T
     fx::T2
     dx::T3
     p::T4
     color::T5
+    sparsity::T6
 end
 
 function default_chunk_size(maxcolor)
@@ -19,7 +20,8 @@ getsize(N::Integer) = N
 
 function ForwardColorJacCache(f,x,_chunksize = nothing;
                               dx = nothing,
-                              color=1:length(x))
+                              color=1:length(x),
+                              sparsity::Union{SparseMatrixCSC,Nothing}=nothing)
 
     if _chunksize === nothing
         chunksize = default_chunk_size(maximum(color))
@@ -38,7 +40,7 @@ function ForwardColorJacCache(f,x,_chunksize = nothing;
     end
 
     p = generate_chunked_partials(x,color,chunksize)
-    ForwardColorJacCache(t,fx,_dx,p,color)
+    ForwardColorJacCache(t,fx,_dx,p,color,sparsity)
 end
 
 generate_chunked_partials(x,color,N::Integer) = generate_chunked_partials(x,color,Val(N))
@@ -78,8 +80,9 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                 f,
                 x::AbstractArray{<:Number};
                 dx = nothing,
-                color = eachindex(x))
-    forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,color=color))
+                color = eachindex(x),
+                sparsity = nothing)
+    forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,color=color,sparsity=sparsity))
 end
 
 function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
@@ -92,15 +95,16 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     dx = jac_cache.dx
     p = jac_cache.p
     color = jac_cache.color
+    sparsity = jac_cache.sparsity
     color_i = 1
     chunksize = length(first(first(jac_cache.p)))
 
     for i in 1:length(p)
         partial_i = p[i]
         t .= Dual{typeof(f)}.(x, partial_i)
         f(fx,t)
-        if J isa SparseMatrixCSC
-            rows_index, cols_index, val = findnz(J)
+        if sparsity isa SparseMatrixCSC
+            rows_index, cols_index, val = findnz(sparsity)
             for j in 1:chunksize
                 dx .= partials.(fx, j)
                 for k in 1:length(cols_index)

test/test_ad.jl

@@ -31,12 +31,12 @@ _J = sparse(J)
 
 fcalls = 0
 _J1 = similar(_J)
-forwarddiff_color_jacobian!(_J1, f, x, color = repeat(1:3,10))
+forwarddiff_color_jacobian!(_J1, f, x, color = repeat(1:3,10), sparsity=_J1)
 @test _J1 ≈ J
 @test fcalls == 1
 
 fcalls = 0
-jac_cache = ForwardColorJacCache(f,x,color = repeat(1:3,10))
+jac_cache = ForwardColorJacCache(f,x,color = repeat(1:3,10), sparsity=_J1)
 forwarddiff_color_jacobian!(_J1, f, x, jac_cache)
 @test _J1 ≈ J
 @test fcalls == 1

test/test_integration.jl

@@ -37,7 +37,14 @@ J = DiffEqDiffTools.finite_difference_jacobian(f, rand(30))
 
 #Jacobian computed with coloring vectors
 fcalls = 0
-_J = 200 .* true_jac
-DiffEqDiffTools.finite_difference_jacobian!(_J, f, rand(30), color = colors)
+_J = similar(true_jac)
+DiffEqDiffTools.finite_difference_jacobian!(_J, f, rand(30), color = colors, sparsity=sparse(_J))
 @test fcalls == 4
 @test _J ≈ J
+
+fcalls = 0
+_J = similar(true_jac)
+_denseJ = collect(_J)
+DiffEqDiffTools.finite_difference_jacobian!(_denseJ, f, rand(30), color = colors, sparsity=sparse(_J))
+@test fcalls == 4
+@test _denseJ ≈ J