Start using DifferentiationInterface.jl

[avik-pal]

Seeing Optimization.jl attempt to use DI, it seems like we have most of the features in place. We might be able to drop a lot of unnecessary code my making that shift!

[avik-pal]

@gdalle I am nearly done integrating DI into NonlinearSolve. The only feature I don't know how to use is specifying

1. A sparse jacobian as the direct prototype. SparseDiffTools has a JacPrototypeSparsityDetection(; jac_prototype)
2. If I already have the matrix coloring PrecomputedJacobianColorvec(; jac_prototype, f.colorvec, partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode)

[avik-pal]

Currently the code is like:

@inline function __sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
    if f.sparsity === nothing
        if f.jac_prototype === nothing
            is_extension_loaded(Val(:Symbolics)) && return SymbolicsSparsityDetection()
            return ApproximateJacobianSparsity()
        else
            jac_prototype = f.jac_prototype
        end
    elseif f.sparsity isa AbstractSparsityDetection
        f.jac_prototype === nothing && return f.sparsity
        jac_prototype = f.jac_prototype
    elseif f.sparsity isa AbstractMatrix
        jac_prototype = f.sparsity
    elseif f.jac_prototype isa AbstractMatrix
        jac_prototype = f.jac_prototype
    else
        error("`sparsity::typeof($(typeof(f.sparsity)))` & \
               `jac_prototype::typeof($(typeof(f.jac_prototype)))` is not supported. \
               Use `sparsity::AbstractMatrix` or `sparsity::AbstractSparsityDetection` or \
               set to `nothing`. `jac_prototype` can be set to `nothing` or an \
               `AbstractMatrix`.")
    end

    if SciMLBase.has_colorvec(f)
        return PrecomputedJacobianColorvec(; jac_prototype,
            f.colorvec,
            partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode)
    else
        return JacPrototypeSparsityDetection(; jac_prototype)
    end
end

ApproximateJacobianSparsity would map to DenseSparsityDetector, and I will clean up the above logic a bit, but we do allow users to pass in a prototype matrix and a colorvec.

[gdalle]

For passing a known sparsity pattern, you can use ADTypes.KnownJacobianSparsityDetector instead of a detector from Symbolics.jl or SparseConnectivityTracer.jl. Just put it inside AutoSparse and it should work fine.

For passing the color vector I don't have a good answer right now, because this logic is handled by SparseMatrixColorings.jl and the coloring result types there are a bit more complex than just a vector of colors (this is especially true for Hessian coloring, which didn't use to be a thing in SparseDiffTools.jl). Let me think about it a little.

Also, ApproximateJacobianSparsity doesn't exactly map to DI.DenseSparsityDetector but it might be a good drop-in replacement.

[avik-pal]

Thanks. I am fine with some manual combination of SparseMatrixColorings and DI if that is easy. This functionality is not exposed to the user so we have some leeway here.

[gdalle]

The main issue I see is that this code mixes sparsity detection and coloring, like SparseDiffTools does, whereas the AutoSparse interface separates those two. Any idea how to get around that?

[avik-pal]

the sparsity_detection_alg can be re-written to return an AutoSparse type. Essentially like this:

1. ad isa dense type.
   1. if sparsity is provided -> AutoSparse(ad;. ....) and some other cases with known sparsity patterns and such
   2. if no sparsity info provided use dense ad
2. if ad is already a sparse type
   1. Ideally there is no sparsity information
   2. If there is in the problem definition, try resolving it. (the details here are not super important, but I will document it at some point)

[avik-pal]

The only part that doesn't fit in here is that we can have the colorvec, so we just need to do the AD in that case

[gdalle]

Maybe I could add a KnownColorsColoringAlgorithm to SparseMatrixColorings, mirroring the KnownJacobianSparsityDetector?

[avik-pal]

I think that would do it

[ChrisRackauckas]

Approximate colorings can be used in order to give an accelerated Jacobian. Is there a way to do approximate colorings?

Also, using colorings on a dense Jacobian can be useful if the sparsity would is not enough and would slow down the LU. Can you sparse fill a dense matrix?

[gdalle]

At the moment there is no approximate coloring but I could add it to DI.

As for filling matrices with the results of sparse AD, SparseMatrixColorings should work with any mutable AbstractMatrix. However it is only optimized for SparseMatrixCSC right now. More structured matrix types could be added in package extensions.