Merge pull request #61 from huanglangwen/optimize

Optimize for BandedBlockBanded Matrix

Author: ChrisRackauckas

Project.toml

@@ -17,7 +17,8 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f"
 
 [compat]
-ArrayInterface = "1.1"
+ArrayInterface = ">= 1.1"
+DiffEqDiffTools = ">= 1.3"
 julia = "1"
 
 [extras]

src/differentiation/compute_jacobian_ad.jl

@@ -102,24 +102,24 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     color_i = 1
     fill!(J, zero(eltype(J)))
 
+    if DiffEqDiffTools._use_findstructralnz(sparsity)
+        rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
+    else
+        rows_index = nothing
+        cols_index = nothing
+    end
+
+    ncols=size(J,2)
+
     for i in eachindex(p)
         partial_i = p[i]
         t .= Dual{typeof(f)}.(x, partial_i)
         f(fx,t)
-        if ArrayInterface.has_sparsestruct(sparsity)
-            rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
+        if !(sparsity isa Nothing)
             for j in 1:chunksize
                 dx .= partials.(fx, j)
                 if ArrayInterface.fast_scalar_indexing(dx)
-                    for k in 1:length(cols_index)
-                        if colorvec[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
+                    DiffEqDiffTools._colorediteration!(J,sparsity,rows_index,cols_index,dx,colorvec,color_i,ncols)
                 else
                     #=
                     J.nzval[rows_index] .+= (colorvec[cols_index] .== color_i) .* dx[rows_index]
@@ -134,12 +134,12 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                     end
                 end
                 color_i += 1
-                (color_i > maximum(colorvec)) && continue
+                (color_i > maximum(colorvec)) && return
             end
         else
             for j in 1:chunksize
                 col_index = (i-1)*chunksize + j
-                (col_index > maximum(colorvec)) && continue
+                (col_index > maximum(colorvec)) && return
                 J[:, col_index] .= partials.(fx, j)
             end
         end

test/test_ad.jl

@@ -2,6 +2,7 @@ using SparseDiffTools
 using ForwardDiff: Dual, jacobian
 using SparseArrays, Test
 using LinearAlgebra
+using BlockBandedMatrices
 
 fcalls = 0
 function f(dx,x)
@@ -61,3 +62,22 @@ forwarddiff_color_jacobian!(_denseJ1, f, x, jac_cache)
 _Jt = similar(Tridiagonal(J))
 forwarddiff_color_jacobian!(_Jt, f, x, colorvec = repeat(1:3,10), sparsity = _Jt)
 @test _Jt ≈ J
+
+#https://github.com/JuliaDiffEq/DiffEqDiffTools.jl/issues/67#issuecomment-516871956
+function f(out, x)
+	x = reshape(x, 100, 100)
+	out = reshape(out, 100, 100)
+	for i in 1:100
+		for j in 1:100
+			out[i, j] = x[i, j] + x[max(i -1, 1), j] + x[min(i+1, size(x, 1)), j] +  x[i, max(j-1, 1)]  + x[i, min(j+1, size(x, 2))]
+		end
+	end
+	return vec(out)
+end
+x = rand(10000)
+J = BandedBlockBandedMatrix(Ones(10000, 10000), (fill(100, 100), fill(100, 100)), (1, 1), (1, 1))
+Jsparse = sparse(J)
+colors = matrix_colors(J)
+forwarddiff_color_jacobian!(J, f, x, colorvec=colors)
+forwarddiff_color_jacobian!(Jsparse, f, x, colorvec=colors)
+@test J ≈ Jsparse