Merge pull request #146 from sjdaines/forwardiff_color_jacobian_fast_sparse_path

Add fast path for forwarddiff_color_jacobian! with sparse J

Author: ChrisRackauckas

Project.toml

@@ -22,7 +22,7 @@ Adapt = "1, 2.0, 3.0"
 ArrayInterface = "2.8, 3.0"
 Compat = "2.2, 3"
 DataStructures = "0.17, 0.18"
-FiniteDiff = "2"
+FiniteDiff = "2.8.1"
 ForwardDiff = "0.10"
 LightGraphs = "1.3"
 Requires = "0.5, 1.0"

src/differentiation/compute_jacobian_ad.jl

@@ -264,6 +264,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     chunksize = jac_cache.chunksize
     color_i = 1
     maxcolor = maximum(colorvec)
+    
     fill!(J, zero(eltype(J)))
 
     if FiniteDiff._use_findstructralnz(sparsity)
@@ -273,6 +274,9 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
         cols_index = 1:size(J,2)
     end
 
+    # fast path if J and sparsity are both SparseMatrixCSC and have the same sparsity pattern
+    sparseCSC_common_sparsity = FiniteDiff._use_sparseCSC_common_sparsity(J, sparsity)
+
     vecx = vec(x)
     vect = vec(t)
     vecfx= vec(fx)
@@ -287,9 +291,14 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
         if !(sparsity isa Nothing)
             for j in 1:chunksize
                 dx .= partials.(fx, j)
+
                 if ArrayInterface.fast_scalar_indexing(dx)
                     #dx is implicitly used in vecdx
-                    FiniteDiff._colorediteration!(J,sparsity,rows_index,cols_index,vecdx,colorvec,color_i,ncols)
+                    if sparseCSC_common_sparsity
+                        FiniteDiff._colorediteration!(J,vecdx,colorvec,color_i,ncols)
+                    else
+                        FiniteDiff._colorediteration!(J,sparsity,rows_index,cols_index,vecdx,colorvec,color_i,ncols)
+                    end
                 else
                     #=
                     J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* dx[rows_index]
@@ -316,3 +325,6 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     end
     return J
 end
+
+
+