Support for structured sparsity in sparse hessian

[baggepinnen]

The following example works fine for sparse sparsity, but sparsity::Diagonal does not appear to be supported.

function fscalar(x)
    return dot(x, x)
end

x = randn(5)
sparsity = hessian_sparsity(fscalar, x)
colors = matrix_colors(tril(sparsity))

sparsity = Diagonal(Vector(diag(sparsity)))
hescache3 = ForwardColorHesCache(fscalar, x, colors, sparsity)

res = zeros(5, 5)
SparseDiffTools.numauto_color_hessian!(res, fscalar, x, hescache3)

julia> SparseDiffTools.numauto_color_hessian!(res, fscalar, x, hescache3)
ERROR: MethodError: no method matching findnz(::Diagonal{Bool, Vector{Bool}})
Closest candidates are:
  findnz(::SparseArrays.AbstractSparseMatrixCSC{Tv, Ti}) where {Tv, Ti} at ~/.julia/juliaup/julia-1.8.0-rc4+0.x64/share/julia/stdlib/v1.8/SparseArrays/src/sparsematrix.jl:1596
  findnz(::SparseVector{Tv, Ti}) where {Tv, Ti} at ~/.julia/juliaup/julia-1.8.0-rc4+0.x64/share/julia/stdlib/v1.8/SparseArrays/src/sparsevector.jl:784
Stacktrace:
 [1] numauto_color_hessian!(H::Matrix{Float64}, f::Function, x::Vector{Float64}, hes_cache::ForwardColorHesCache{Diagonal{Bool, Vector{Bool}}, Vector{Int64}, Int64, Matrix{Float64}, SparseDiffTools.var"#43#46"{typeof(fscalar)}, ForwardDiff.GradientConfig{ForwardDiff.Tag{typeof(fscalar), Float64}, Float64, 5, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(fscalar), Float64}, Float64, 5}}}, Vector{Float64}}; safe::Bool)
   @ SparseDiffTools ~/.julia/dev/SparseDiffTools/src/differentiation/compute_hessian_ad.jl:62
 [2] numauto_color_hessian!(H::Matrix{Float64}, f::Function, x::Vector{Float64}, hes_cache::ForwardColorHesCache{Diagonal{Bool, Vector{Bool}}, Vector{Int64}, Int64, Matrix{Float64}, SparseDiffTools.var"#43#46"{typeof(fscalar)}, ForwardDiff.GradientConfig{ForwardDiff.Tag{typeof(fscalar), Float64}, Float64, 5, Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(fscalar), Float64}, Float64, 5}}}, Vector{Float64}})
   @ SparseDiffTools ~/.julia/dev/SparseDiffTools/src/differentiation/compute_hessian_ad.jl:48
 [3] top-level scope
   @ REPL[109]:1