Using functors for sparse Jacobian evaluations

General Usage
,
1

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] = 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?

2

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[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).

3

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] = 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))