# Using functors for sparse Jacobian evaluations

I am trying to use ForwardDiff to calculate the Jacobian of a function so that I can reuse the resulting Jacobian function and use it as an input for several packages. For instance, pass it to NLsolve and later use it in DiffEq as well. I know that the Jacobian will be sparse, so I do one evaluation to get the sparsity and use the in-place evaluations later.

The MWE below works, and @code_warntype is making the type inference correctly in every step. However, I still get allocations when evaluating the Jacobian function on a set of values.

``````using ForwardDiff
using Random
using SparseArrays

function component1(u, common) where T
common = 1.0*u
return -0.1 * u
end

function component2(u, common) where T
common = -0.8 * u
return -0.4 * u
end

function component3(u, common) where T
return sum(common .* [-0.1, 0.2, 0.01]) + 0.02 * u
end

function system(du, u, common, t) where T
du = component1(u, common)
du = component2(u, common)
du = component3(u, common)
nothing
end

struct modelf{T, U}
common::Vector{T}
common_dual::Vector{U}
du::Vector{T}
du_dual::Vector{U}
function modelf(du0::Vector{T}) where T <: Number
# Forced the CHUNK to 3 for the MWE.
U = ForwardDiff.Dual{typeof(ForwardDiff.Tag(system, T)), T, 3}
new{T, U}(zeros(T, 3), zeros(U, 3), similar(du0), zeros(U, length(du0)))
end
end

get_common(m::modelf, ::Type{Float64}) = m.common
get_common(m::modelf, ::Type{T}) where T <: ForwardDiff.Dual = reinterpret(T, m.common_dual)
get_du(m::modelf, ::Type{Float64}) = m.du
get_du(m::modelf, ::Type{T}) where T <: ForwardDiff.Dual = reinterpret(T, m.du_dual)

function (f::modelf)(u::AbstractArray{T}, t) where T
system(get_du(f, T), u, get_common(f, T), t)
return get_du(f, T)
end

struct modelJ{T}
Jf::T
Jv::SparseMatrixCSC{Float64, Int64}
x::Vector{Float64}
function modelJ(x)
model = modelf(zeros(3))
fu = (u) -> model(u, 0)
jconfig = ForwardDiff.JacobianConfig(fu , x, ForwardDiff.Chunk{3}())
Jf = (Jv, x) -> ForwardDiff.jacobian!(Jv, fu, x, jconfig)
n = length(x)
sJ = sparse(Jf(zeros(n, n), x))
new{typeof(Jf)}(Jf, sJ, x)
end
end

function (J::modelJ{T})(x) where T <: Function
J.x .= x
return J.Jf(J.Jv, x)
end

test = modelJ(rand(3))

@time test(rand(3))

0.000014 seconds (9 allocations: 656 bytes)
3×3 SparseMatrixCSC{Float64, Int64} with 5 stored entries:
-0.1    ⋅     ⋅
⋅   -0.4    ⋅
-0.1  -0.16  0.02
``````

Any ideas on where the allocations might be coming from?

I don’t know ForwardDiff’s internals well enough to guess where the rest of the allocations are coming from, but you’re getting at least two from `component3` (one from the constant array, one from the result of `common .* [-0.1, 0.2, 0.01]`).

``````function component3(u, common) where T
return sum(common .* [-0.1, 0.2, 0.01]) + 0.02 * u
end
``````

You can trim those allocations by using `mapreduce(*, +, common, (-0.1, 0.2, 0.01))`, or with `LinearAlgebra.dot`, you can do `common ⋅ (-0.1, 0.2, 0.01)`.

Thanks, you are right. The allocations are coming from the component 3 call. This version only produces 2 allocations.

``````using ForwardDiff
using Random
using SparseArrays

function component1(u, common)
common = 1.0*u
return -0.1 * u
end

function component2(u, common)
common = -0.8 * u
return -0.4 * u
end

function component3(u, common)
return mapreduce(*, +, common, (-0.1, 0.2, 0.01)) + 0.02 * u
end

function system(du, u, common, t)
du = component1(u, common)
du = component2(u, common)
du = component3(u, common)
nothing
end

struct modelf{T, U}
common::Vector{T}
common_dual::Vector{U}
du::Vector{T}
du_dual::Vector{U}
function modelf(du0::Vector{T}) where T <: Number
# Forced the CHUNK to 3 for the MWE.
U = ForwardDiff.Dual{typeof(ForwardDiff.Tag(system, T)), T, 3}
new{T, U}(zeros(T, 3), zeros(U, 3), similar(du0), zeros(U, length(du0)))
end
end

get_common(m::modelf, ::Type{Float64}) = m.common
get_common(m::modelf, ::Type{T}) where T <: ForwardDiff.Dual = reinterpret(T, m.common_dual)
get_du(m::modelf, ::Type{Float64}) = m.du
get_du(m::modelf, ::Type{T}) where T <: ForwardDiff.Dual = reinterpret(T, m.du_dual)

function (f::modelf)(u::AbstractArray{T}, t) where T
system(get_du(f, T), u, get_common(f, T), t)
return get_du(f, T)
end

struct modelJ{T}
Jf::T
Jv::SparseMatrixCSC{Float64, Int64}
x::Vector{Float64}
function modelJ(x)
model = modelf(zeros(3))
fu = (u) -> model(u, 0)
jconfig = ForwardDiff.JacobianConfig(fu, x, ForwardDiff.Chunk{3}())
Jf = (Jv, x) -> ForwardDiff.jacobian!(Jv, fu, x, jconfig)
n = length(x)
sJ = sparse(Jf(zeros(n, n), x))
new{typeof(Jf)}(Jf, sJ, x)
end
end

function (J::modelJ{T})(x) where T <: Function
J.x .= x
return J.Jf(J.Jv, x)
end

test = modelJ(rand(3))

@time test(rand(3))
``````