# Using functors for sparse Jacobian evaluations

**URL:** <https://discourse.julialang.org/t/using-functors-for-sparse-jacobian-evaluations/64111>\
**Category:** General Usage\
**Tags:** forwarddiff, functors\
**Created:** [July 5, 2021, 10:16pm UTC](https://discourse.julialang.org/t/using-functors-for-sparse-jacobian-evaluations/64111 "2021-07-05T22:16:38Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![jdlara-berkeley](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jdlara-berkeley/32/4096_2.png) [@jdlara-berkeley](https://discourse.julialang.org/u/jdlara-berkeley)\
**Post date:** [July 5, 2021, 10:16pm UTC](https://discourse.julialang.org/t/using-functors-for-sparse-jacobian-evaluations/64111/1 "2021-07-05T22:16:38Z")

</div>

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.

```nohighlight
using ForwardDiff
using Random
using SparseArrays

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

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

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

function system(du, u, common, t) where T
    du[1] = component1(u, common)
    du[2] = component2(u, common)
    du[3] = 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?

---

<div class="post-metadata">

**Author:** ![stillyslalom](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stillyslalom/32/45687_2.png) [@stillyslalom](https://discourse.julialang.org/u/stillyslalom)\
**Post date:** [July 6, 2021, 12:03am UTC](https://discourse.julialang.org/t/using-functors-for-sparse-jacobian-evaluations/64111/2 "2021-07-06T00:03:26Z")

</div>

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]`).

```julia
function component3(u, common) where T
    return sum(common .* [-0.1, 0.2, 0.01]) + 0.02 * u[3]
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)`.

---

<div class="post-metadata">

**Author:** ![jdlara-berkeley](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jdlara-berkeley/32/4096_2.png) [@jdlara-berkeley](https://discourse.julialang.org/u/jdlara-berkeley)\
**Post date:** [July 6, 2021, 12:13am UTC](https://discourse.julialang.org/t/using-functors-for-sparse-jacobian-evaluations/64111/3 "2021-07-06T00:13:38Z")

</div>

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

```nohighlight
using ForwardDiff
using Random
using SparseArrays

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

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

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

function system(du, u, common, t)
    du[1] = component1(u, common)
    du[2] = component2(u, common)
    du[3] = 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))

```