Merge pull request #124 from sjdaines/fast_sparse_path

Add fast path for _colorediteration! with sparse J

Author: ChrisRackauckas

src/iteration_utils.jl

@@ -16,10 +16,27 @@ end
     end
 end
 
+# fast version for the case where J and sparsity have the same sparsity pattern
+@inline function _colorediteration!(Jsparsity::SparseMatrixCSC,vfx,colorvec,color_i,ncols)
+    @inbounds for col_index in 1:ncols
+        if colorvec[col_index] == color_i
+            @inbounds for spidx in nzrange(Jsparsity, col_index)
+                row_index = Jsparsity.rowval[spidx]
+                Jsparsity.nzval[spidx]=vfx[row_index]
+            end
+        end
+    end
+end
+
 #override default setting of using findstructralnz
 _use_findstructralnz(sparsity) = ArrayInterface.has_sparsestruct(sparsity)
 _use_findstructralnz(::SparseMatrixCSC) = false
 
+# test if J, sparsity are both SparseMatrixCSC and have the same sparsity pattern of stored values
+_use_sparseCSC_common_sparsity(J, sparsity) = false
+_use_sparseCSC_common_sparsity(J::SparseMatrixCSC, sparsity::SparseMatrixCSC) =
+    ((J.colptr == sparsity.colptr) && (J.rowval == sparsity.rowval))
+
 function __init__()
     @require BlockBandedMatrices="ffab5731-97b5-5995-9138-79e8c1846df0" begin
         @require BlockArrays="8e7c35d0-a365-5155-bbbb-fb81a777f24e" begin

src/jacobians.jl

@@ -348,6 +348,9 @@ function finite_difference_jacobian!(
         fill!(J,false)
     end
 
+    # fast path if J and sparsity are both SparseMatrixCSC and have the same sparsity pattern
+    sparseCSC_common_sparsity = _use_sparseCSC_common_sparsity(J, sparsity)
+    
     if fdtype == Val(:forward)
         vfx1 = _vec(fx1)
 
@@ -378,7 +381,11 @@ function finite_difference_jacobian!(
                 # J is a sparse matrix, so decompress on the fly
                 @. vfx1 = (vfx1 - vfx) / epsilon
                 if ArrayInterface.fast_scalar_indexing(x1)
-                    _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n)
+                    if sparseCSC_common_sparsity
+                        _colorediteration!(J,vfx1,colorvec,color_i,n)
+                    else
+                        _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n)
+                    end
                 else
                     #=
                     J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* vfx1[rows_index]
@@ -417,8 +424,12 @@ function finite_difference_jacobian!(
                 f(fx1, x1)
                 f(fx, x)
                 @. vfx1 = (vfx1 - vfx) / 2epsilon
-                if ArrayInterface.fast_scalar_indexing(x1)
-                    _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n)
+                if ArrayInterface.fast_scalar_indexing(x1)                    
+                    if sparseCSC_common_sparsity
+                        _colorediteration!(J,vfx1,colorvec,color_i,n)
+                    else
+                        _colorediteration!(J,sparsity,rows_index,cols_index,vfx1,colorvec,color_i,n)
+                    end
                 else
                     if J isa SparseMatrixCSC
                         @. void_setindex!((J.nzval,), getindex((J.nzval,), rows_index) + (getindex((_color,), cols_index) == color_i) * getindex((vfx1,), rows_index), rows_index)
@@ -443,8 +454,12 @@ function finite_difference_jacobian!(
                 @. x1 = x1 + im * epsilon * (_color == color_i)
                 f(fx,x1)
                 @. vfx = imag(vfx) / epsilon
-                if ArrayInterface.fast_scalar_indexing(x1)
-                    _colorediteration!(J,sparsity,rows_index,cols_index,vfx,colorvec,color_i,n)
+                if ArrayInterface.fast_scalar_indexing(x1)                    
+                    if sparseCSC_common_sparsity
+                        _colorediteration!(J,vfx,colorvec,color_i,n)
+                    else
+                        _colorediteration!(J,sparsity,rows_index,cols_index,vfx,colorvec,color_i,n)
+                    end
                 else
                    if J isa SparseMatrixCSC
                         @. void_setindex!((J.nzval,), getindex((J.nzval,), rows_index) + (getindex((_color,), cols_index) == color_i) * getindex((vfx,),rows_index), rows_index)