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.

[ChrisRackauckas]

I think the main case of it is using a chosen color vector with a dense matrix. The reason for that is again if sparse LU is slow on the matrix (sparse matrices really need to be "sufficiently large" for sparse factorization to make sense) but coloring still makes sense to cut down the Jacobian cost.

[avik-pal]

Is it possible to query DI to get a Jacobian structure (basically the result of DI.jacobian but not computing the actual jacobian)? It is needed to ensure that we can create a concrete cache on init.

Note that similar doesn't work because jacobian might give a SMatrix.

[amontoison]

Currently, we have four variants of coloring in SparseMatrixColorings.jl.
For some of them, we can easily use a user-provided coloring (vector of integers) and perform the decompression without loss of performance.

• Row coloring (optimal)
• Column coloring (optimal)
• Star coloring (possible but slow without the structure StarSet)
• Acyclic coloring (not possible; the structure TreeSet is needed)

StarSet and TreeSet are structures generated during the coloring.

Star and acyclic coloring are generally used on Hessians, while row and column coloring do not exploit symmetry and are mostly used for Jacobians.
I think just supporting colorvec as input for an expected row or column coloring is the best alternative.
You can use it for both Jacobians and Hessians if you don't want to perform a new coloring.
You will keep all the features from SparseDiffTools.jl.

[gdalle]

Is it possible to query DI to get a Jacobian structure (basically the result of DI.jacobian but not computing the actual jacobian)? It is needed to ensure that we can create a concrete cache on init.

Not at the moment no. Are we talking about sparse AD or general AD?
If you're creating a cache that means you're gonna run jacobian many times, so I think a good strategy is just to actually run it once to get a cache?

Note that similar doesn't work because jacobian might give a SMatrix.

But then what kind of cache do you want to get in this case? similar will give you an MMatrix, which seems like the only reasonable choice?

[avik-pal]

Not at the moment no. Are we talking about sparse AD or general AD?
If you're creating a cache that means you're gonna run jacobian many times, so I think a good strategy is just to actually run it once to get a cache?

Sorry I meant cache for NonlinearSolve. I am currently running it once to get the structure as you mention.

My structure I mean, if the returned J is mutable then DI.jacobian! should work on it and if it is immutable then I shouldn't be able to do DI.jacobian!

[gdalle]

Then I would recommend you just copy the output of DI.jacobian? But keep in mind that I make no guarantee about the exact return type of the Jacobian matrix, because the backends don't either. So even if you give SVectors as input, with one backend you may get SMatrix, with another MMatrix, and with a third one a plain Matrix. The only thing that is guaranteed is the correctness of the values inside.

[avik-pal]

Right. This is what I will do (am doing on master) and seems like a fine solution since users caring too much about performance (1 jac call) should be using the caching interface anyways

[gdalle]

@avik-pal does this do what you want in terms of fixing the vector of colors? And it is explained clearly enough?
gdalle/SparseMatrixColorings.jl#127

[avik-pal]

I think so. Do you need to make the conversion to sparse array? If you have a BandedMatrix it seems like a bad idea to make it Sparse because then the jacobian function with presumably return a sparse array instead of a banded matrix and linear solve will hit a regression

[gdalle]

At the moment yes, there will be a conversion to SparseMatrixCSC because custom matrix types (like banded) are not yet supported by SparseMatrixColorings.jl. However, with ConstantColoringAlgorithm, this conversion to SparseMatrixCSC will only happen once (at creation).

What matters to me right now is getting the necessary functionality up and running, then we can see what the benchmarks tell us about performance regressions. But a fair warning that may not be obvious to everyone: when switching to the DI ecosystem, there will be performance regressions in some cases, and performance improvements in others. A good example is the Brusselator AD benchmark, where sparsity detection is 100x faster, coloring 10x faster, and differentiation 2x slower (although this has probably improved with the latest DI and SMC releases, I'll have to run it again).
To adapt to your use cases and milk the last drops of performance, first we need to get DI into NonlinearSolve, and then we'll identify the bottlenecks together. Most functionalities in DI have never been benchmarked on realistic use cases, so I'll definitely need your help for that once the switch passes correctness tests.

[avik-pal]

Let me make a branch using DI for sparse jacobians, but I will hold off on merging that for now.

[gdalle]

FYI the new version of SparseMatrixColorings that includes ConstantColoringAlgorithm is out

[gdalle]

I re-ran the Brusselator sparse AD benchmark with the latest versions of everything, and DI is now on par with SparseDiffTools: https://docs.sciml.ai/SciMLBenchmarksOutput/dev/AutomaticDifferentiationSparse/BrusselatorSparseAD/