add oop jac forwarddiff_color_jacobian

Project.toml

@@ -30,6 +30,7 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparsityDetection = "684fba80-ace3-11e9-3d08-3bc7ed6f96df"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [targets]
-test = ["Test", "IterativeSolvers", "Random", "SafeTestsets", "Zygote", "SparsityDetection"]
+test = ["Test", "IterativeSolvers", "Random", "SafeTestsets", "Zygote", "SparsityDetection", "StaticArrays"]

README.md

@@ -19,7 +19,7 @@ Suppose we had the function
 
 ```julia
 fcalls = 0
-function f(dx,x)
+function f(dx,x) # in-place
   global fcalls += 1
   for i in 2:length(x)-1
     dx[i] = x[i-1] - 2x[i] + x[i+1]
@@ -28,6 +28,17 @@ function f(dx,x)
   dx[end] = x[end-1] - 2x[end]
   nothing
 end
+
+function g(x) # out-of-place
+  global fcalls += 1
+  dx = zero(x)
+  for i in 2:length(x)-1
+    dx[i] = x[i-1] - 2x[i] + x[i+1]
+  end
+  dx[1] = -2x[1] + x[2]
+  dx[end] = x[end-1] - 2x[end]
+  dx
+end
 ```
 
 For this function, we know that the sparsity pattern of the Jacobian is a
@@ -62,7 +73,10 @@ In addition, a faster forward-mode autodiff call can be utilized as well:
 
 ```julia
 forwarddiff_color_jacobian!(jac, f, x, colorvec = colors)
+jacout = forwarddiff_color_jacobian(g, x, colorvec = colors, sparsity = similar(jac), jac_prototype = similar(jac))
 ```
+One can specify the type of the output jacobian by giving an additional `jac_prototype` to 
+the out-of place `forwarddiff_color_jacobian` function, otherwise it will become a dense matrix.
 
 If one only need to compute products, one can use the operators. For example,
 

src/SparseDiffTools.jl

@@ -25,6 +25,7 @@ export  contract_color,
         matrix2graph,
         matrix_colors,
         forwarddiff_color_jacobian!,
+        forwarddiff_color_jacobian,
         ForwardColorJacCache,
         auto_jacvec,auto_jacvec!,
         num_jacvec,num_jacvec!,

src/differentiation/compute_jacobian_ad.jl

@@ -41,48 +41,89 @@ function ForwardColorJacCache(f,x,_chunksize = nothing;
         _dx = dx
     end
 
-
     ForwardColorJacCache(t,fx,_dx,p,colorvec,sparsity,getsize(chunksize))
 end
 
 generate_chunked_partials(x,colorvec,N::Integer) = generate_chunked_partials(x,colorvec,Val(N))
 function generate_chunked_partials(x,colorvec,::Val{chunksize}) where chunksize
+    maxcolor = maximum(colorvec)
+    num_of_chunks = Int(ceil(maxcolor / chunksize))
+    padding_size = (chunksize - (maxcolor % chunksize)) % chunksize
+    partials = colorvec .== (1:maxcolor)'
+    padding_matrix = BitMatrix(undef, length(x), padding_size)
+    partials = hcat(partials, padding_matrix)
 
-    num_of_chunks = Int(ceil(maximum(colorvec) / chunksize))
+    chunked_partials = map(i -> Tuple.(eachrow(partials[:,(i-1)*chunksize+1:i*chunksize])),1:num_of_chunks)
+    chunked_partials
 
-    padding_size = (chunksize - (maximum(colorvec) % chunksize)) % chunksize
+end
 
-    partials = BitMatrix(undef, length(x), maximum(colorvec))
-    partial = BitMatrix(undef, length(x), chunksize)
-    chunked_partials = Array{Array{Tuple{Vararg{Bool,chunksize}},1},1}(
-                                                          undef, num_of_chunks)
+function forwarddiff_color_jacobian(f,
+                x::AbstractArray{<:Number};
+                dx = nothing,
+                colorvec = 1:length(x),
+                sparsity = nothing,
+                jac_prototype = nothing)
+    forwarddiff_color_jacobian(f,x,ForwardColorJacCache(f,x,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype)
+end
 
-    for color_i in 1:maximum(colorvec)
-        for j in 1:length(x)
-            partials[j,color_i] = colorvec[j]==color_i
-        end
-    end
+function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)
+    t = jac_cache.t
+    fx = jac_cache.fx
+    dx = jac_cache.dx
+    p = jac_cache.p
+    colorvec = jac_cache.colorvec
+    sparsity = jac_cache.sparsity
+    chunksize = jac_cache.chunksize
+    color_i = 1
+    maxcolor = maximum(colorvec)
 
-    padding_matrix = BitMatrix(undef, length(x), padding_size)
-    partials = hcat(partials, padding_matrix)
+    vecx = vec(x)
+
+    ncols=length(x)
+    J = jac_prototype isa Nothing ? false .* x .* x' : zero(jac_prototype)
 
-    for i in 1:num_of_chunks
-        partial[:,1] .= partials[:,(i-1)*chunksize+1]
-        for j in 2:chunksize
-            partial[:,j] .= partials[:,(i-1)*chunksize+j]
-        end
-        chunked_partials[i] = Tuple.(eachrow(partial))
+    if !(sparsity isa Nothing)
+        rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
+        rows_index = [rows_index[i] for i in 1:length(rows_index)]
+        cols_index = [cols_index[i] for i in 1:length(cols_index)]
     end
 
-    chunked_partials
-
+    for i in eachindex(p)
+        partial_i = p[i]
+        t = reshape(Dual{typeof(f)}.(vecx, partial_i),size(t))
+        fx = f(t)
+        if !(sparsity isa Nothing)
+            for j in 1:chunksize
+                dx = vec(partials.(fx, j))
+                pick_inds = [i for i in 1:length(rows_index) if colorvec[cols_index[i]] == color_i]
+                rows_index_c = rows_index[pick_inds]
+                cols_index_c = cols_index[pick_inds]
+                len_rows = length(pick_inds)
+                unused_rows = setdiff(1:ncols,rows_index_c)
+                perm_rows = sortperm(vcat(rows_index_c,unused_rows))
+                cols_index_c = vcat(cols_index_c,zeros(Int,ncols-len_rows))[perm_rows]
+                Ji = [j==cols_index_c[i] ? dx[i] : false for i in 1:ncols, j in 1:ncols]
+                J = J + Ji
+                color_i += 1
+                (color_i > maxcolor) && return J
+            end
+        else
+            for j in 1:chunksize
+                col_index = (i-1)*chunksize + j
+                (col_index > maxcolor) && return J
+                J = J + mapreduce(i -> i==col_index ? partials.(vec(fx), j) : zeros(ncols), hcat, 1:ncols)
+            end
+        end
+    end
+    J
 end
 
 function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                 f,
                 x::AbstractArray{<:Number};
                 dx = nothing,
-                colorvec = eachindex(x),
+                colorvec = 1:length(x),
                 sparsity = ArrayInterface.has_sparsestruct(J) ? J : nothing)
     forwarddiff_color_jacobian!(J,f,x,ForwardColorJacCache(f,x,dx=dx,colorvec=colorvec,sparsity=sparsity))
 end
@@ -100,6 +141,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     sparsity = jac_cache.sparsity
     chunksize = jac_cache.chunksize
     color_i = 1
+    maxcolor = maximum(colorvec)
     fill!(J, zero(eltype(J)))
 
     if DiffEqDiffTools._use_findstructralnz(sparsity)
@@ -124,6 +166,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
             for j in 1:chunksize
                 dx .= partials.(fx, j)
                 if ArrayInterface.fast_scalar_indexing(dx)
+                    #dx is implicitly used in vecdx
                     DiffEqDiffTools._colorediteration!(J,sparsity,rows_index,cols_index,vecdx,colorvec,color_i,ncols)
                 else
                     #=
@@ -139,12 +182,12 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                     end
                 end
                 color_i += 1
-                (color_i > maximum(colorvec)) && return
+                (color_i > maxcolor) && return
             end
         else
             for j in 1:chunksize
                 col_index = (i-1)*chunksize + j
-                (col_index > maximum(colorvec)) && return
+                (col_index > maxcolor) && return
                 J[:, col_index] .= partials.(vecfx, j)
             end
         end

test/test_ad.jl

@@ -3,6 +3,7 @@ using ForwardDiff: Dual, jacobian
 using SparseArrays, Test
 using LinearAlgebra
 using BlockBandedMatrices
+using StaticArrays
 
 fcalls = 0
 function f(dx,x)
@@ -15,6 +16,23 @@ function f(dx,x)
     nothing
 end
 
+function oopf(x)
+    global fcalls += 1
+    dx = zero(x)
+    for i in 2:length(x)-1
+        dx[i] = x[i-1] - 2x[i] + x[i+1]
+    end
+    dx[1] = -2x[1] + x[2]
+    dx[end] = x[end-1] - 2x[end]
+    dx
+end
+
+function staticf(x,N=length(x))
+    global fcalls += 1
+    SVector{N}([i == 1 ? -2x[1]+x[2] : (i == N ? x[N-1]-2x[N] : x[i-1]-2x[i]+x[i+1]) for i in 1:N])
+end
+
+
 function second_derivative_stencil(N)
     A = zeros(N,N)
     for i in 1:N, j in 1:N
@@ -39,6 +57,29 @@ forwarddiff_color_jacobian!(_J1, f, x, colorvec = repeat(1:3,10))
 @test fcalls == 1
 
 fcalls = 0
+_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = _J)
+@test _J1 ≈ J
+@test fcalls == 1
+
+fcalls = 0
+_J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity = _J)
+@test _J1 ≈ J
+@test fcalls == 1
+
+fcalls = 0
+_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = SMatrix{30,30}(_J))
+@test _J1 ≈ J
+@test fcalls == 1
+
+_J1 = forwarddiff_color_jacobian(staticf, SVector{30}(x), jac_prototype = SMatrix{30,30}(_J))
+@test _J1 ≈ J
+_J1 = forwarddiff_color_jacobian(oopf, x, jac_prototype = _J)
+@test _J1 ≈ J
+_J1 = forwarddiff_color_jacobian(oopf, x)
+@test _J1 ≈ J
+
+fcalls = 0
+_J1 = similar(_J)
 jac_cache = ForwardColorJacCache(f,x,colorvec = repeat(1:3,10), sparsity = _J1)
 forwarddiff_color_jacobian!(_J1, f, x, jac_cache)
 @test _J1 ≈ J