Add method to mirror ForwardDiff's syntax for writing Jacobian of an oop function to an allocated matrix

src/differentiation/compute_jacobian_ad.jl

@@ -78,7 +78,31 @@ end
     forwarddiff_color_jacobian(f,x,ForwardColorJacCache(f,x,chunksize,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype)
 end
 
+@inline function forwarddiff_color_jacobian(J::AbstractArray{<:Number}, f,
+                x::AbstractArray{<:Number};
+                colorvec = 1:length(x),
+                sparsity = nothing,
+                jac_prototype = nothing,
+                chunksize = nothing,
+                dx = similar(x, size(J, 1))) #if dx is nothing, we will estimate dx at the cost of a function call
+    if sparsity === nothing && jac_prototype === nothing || !ArrayInterface.ismutable(x)
+        cfg = chunksize === nothing ? ForwardDiff.JacobianConfig(f, x) : ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize)))
+        return ForwardDiff.jacobian(f, x, cfg)
+    end
+    forwarddiff_color_jacobian(J,f,x,ForwardColorJacCache(f,x,chunksize,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype)
+end
+
 function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)
+    dx = jac_cache.dx
+    vecx = vec(x)
+    sparsity = jac_cache.sparsity
+
+    J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false .* vec(dx) .* vecx' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype)
+
+    forwarddiff_color_jacobian(J, f, x, jac_cache, jac_prototype)
+end
+
+function forwarddiff_color_jacobian(J::AbstractArray{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)
     t = jac_cache.t
     dx = jac_cache.dx
     p = jac_cache.p
@@ -90,7 +114,6 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw
 
     vecx = vec(x)
 
-    J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false .* vec(dx) .* vecx' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype)
     nrows,ncols = size(J)
 
     if !(sparsity isa Nothing)
@@ -134,6 +157,63 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw
     J
 end
 
+function forwarddiff_color_jacobian(J::SparseMatrixCSC{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)
+    t = jac_cache.t
+    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)
+
+    vecx = vec(x)
+
+    nrows,ncols = size(J)
+
+    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
+
+    for i in eachindex(p)
+        partial_i = p[i]
+        t = reshape(Dual{typeof(ForwardDiff.Tag(f,eltype(vecx)))}.(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]
+                Ji = sparse(rows_index_c, cols_index_c, dx[rows_index_c],nrows,ncols)
+                # J = J + Ji
+                if j == 1 && i == 1
+                    J .= Ji # overwrite pre-allocated matrix
+                else
+                    J .+= Ji
+                end
+                color_i += 1
+                (color_i > maxcolor) && return J
+            end
+        else
+            for j in 1:chunksize
+                col_index = (i-1)*chunksize + j
+                (col_index > ncols) && return J
+                Ji = mapreduce(i -> i==col_index ? partials.(vec(fx), j) : adapt(parameterless_type(J),zeros(eltype(J),nrows)), hcat, 1:ncols)
+                # J = J + (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
+                if j == 1 && i == 1
+                    J .= (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) # overwrite pre-allocated matrix
+                else
+                    J .+= (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
+                end
+            end
+        end
+    end
+    J
+end
+
 function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
                 f,
                 x::AbstractArray{<:Number};

test/test_ad.jl

@@ -115,6 +115,14 @@ _J1 = forwarddiff_color_jacobian(oopf, x, colorvec = repeat(1:3,10), sparsity =
 @test _J1 ≈ J
 @test fcalls == 1
 
+#oop with in-place Jacobian
+fcalls = 0
+_oop_jacout = sparse(1.01 .* J) # want to be nonzero to check that the pre-allocated matrix is overwritten properly
+forwarddiff_color_jacobian(_oop_jacout, oopf, x; colorvec = repeat(1:3,10), sparsity = _J, jac_prototype = _J)
+@test _oop_jacout ≈ J
+@test typeof(_oop_jacout) == typeof(_J)
+@test fcalls == 1
+
 @info "4th passed"
 
 fcalls = 0