feat: use DifferentiationInterface for sparse AD

[avik-pal]

The only part where we don't use DI is structured Jacobians. That still uses SparseDiffTools

TODO:

• construct a case with structured jacobian that is small enough for testing. something block diagonal is the easiest to generate
• merge feat: use DI for structured Jacobians #470
• start removing sparsedifftools from the tutorials/docs
  • docs
  • tests
• needs Coloring and decompression for structured matrices gdalle/SparseMatrixColorings.jl#132

[gdalle]

Congrats on the adaptation and thank you for trusting me with this!
I've never used LinearSolve before so take my comments with a grain of salt, but I find it needlessly complicated to mix the SparseDiffTools legacy syntax with the new sparse API of ADTypes. Do you think it is a problem?

[avik-pal]

but I find it needlessly complicated to mix the SparseDiffTools legacy syntax with the new sparse API of ADTypes.

This preserves support for structured matrices. Once the support for that lands in SparseMatrixColorings we can just delete that check and everything will just work.

[avik-pal]

To clarify the rationale behind marking AutoSparse as deprecated:

Previously we had sparse backends (simple wrapper over dense backend) and we had sparsity specification inside the problem (sparsity, jac_prototype, colorvec). But now the AutoSparse backend holds the sparsity detection + coloring information. So this creates chances for ambiguity, say user says sparsity = Symbolics..Detector() and AutoSparse(..., Tracer...Detector()) what should take precedence? We could choose to make the autodiff via the solver have higher precedence. But the other way is also perfectly valid. This mostly causes confusion without a clear benefit (at least to me)

Now what does the new-API look like? Users always provide a dense_ad to the solver. If the function has sparsity information in the form of a detector/colorvec/prototype (basically any hint that the jacobian is sparse) we construct a AutoSparse backend using that information. From DI's perspective it always sees the final autodiff that we constructed and not the one user passed in.

[gdalle]

Right, so essentially you reconstruct an AutoSparse from

• the backend provided to the solver
• the sparsity detector / coloring provided to the function

That seems reasonable for now, but it might lead to a breaking change when we swing back to full AutoSparse support?

[gdalle]

This preserves support for structured matrices. Once the support for that lands in SparseMatrixColorings we can just delete that check and everything will just work.

Just to clarify, with structured matrices, there are two things we can optimize:

• the coloring, for which an optimal solution is often known
• the decompression, because storage is e.g. bandwise instead of columnwise

At the moment, if I understand correctly:

• the coloring is in ArrayInterface.jl, reused by SparseDiffTools.jl?
• the decompression is in these extensions of FiniteDiff.jl, reused by SparseDiffTools.jl

I'll try to come up with a prototype putting everything in SparseMatrixColorings.jl

[gdalle]

On this branch you have optimized coloring and decompression for Diagonal, Bidiagonal, Tridiagonal and BandedMatrix: gdalle/SparseMatrixColorings.jl#132. Wanna try it out?

[avik-pal]

On this branch you have optimized coloring and decompression for Diagonal, Bidiagonal, Tridiagonal and BandedMatrix: gdalle/SparseMatrixColorings.jl#132. Wanna try it out?

I will stack a PR on top of this and try it out

[avik-pal]

That seems reasonable for now, but it might lead to a breaking change when we swing back to full AutoSparse support?

What do you mean by full AutoSparse support?

I am generally against having 2 APIs to accomplish the exact same thing unless each API provides additional disjoint benefits. If you are referring to the coloring_algorithm, we can always rename colorvec to coloring and construct AutoSparse based on the type of coloring.

@ChrisRackauckas what do you think about this new API proposal for v4?

[gdalle]

What do you mean by full AutoSparse support?

I meant that for users who pass an AutoSparse backend to NonLinearSolve, it might be surprising to see that the choices of coloring_algorithm / sparsity_detector inside the backend are discarded, in favor of the choices colorvec / sparsity given in the function itself?

If you want to deprecate this way of handling sparsity, I would almost prefer throwing an error if someone gives you an AutoSparse.

[avik-pal]

If you want to deprecate this way of handling sparsity, I would almost prefer throwing an error if someone gives you an AutoSparse.

Yes I agree. The function currently throws a depwarn, that will become an error in v4

[gdalle]

I added support for BlockBandedMatrix and BandedBlockBandedMatrix to the same PR. Now we have the exact same colorings as ArrayInterface, although decompression could use a little more work (but it's not terrible).

[gdalle]

What else would you need to drop SparseDiffTools completely? In terms of functionality it should be all there, performance may still lag a bit behind but honestly I'm not even sure.

[avik-pal]

see the other PR. tests pass locally for me, so we can go ahead and remove sparsedifftools. I will make some additional pushes to remove sparse diff tools from the tests and docs

[avik-pal]

Added a big table to clarify the selection mechanism https://github.com/SciML/NonlinearSolve.jl/blob/ap/sparse_di/docs/src/basics/sparsity_detection.md

[oscardssmith]

Is there anything we should be doing to pass down the sparsity info that we get down to LinearSolve? It seems like knowing that should be able to make the solve faster.

[avik-pal]

We are giving LinearSolve the exact matrix it needs to solve

[avik-pal]

We assume for SparseMatrixCSC jacobians the sparsity struture cannot arbitrarily change over iterations

[gdalle]

Yeah that goes without saying for me but if you prepare DI's Jacobian with a given sparsity pattern, you'll get incorrect outputs if the sparsity changes between points

[gdalle]

Fair warning: you're probably going to see an epic slowdown when computing sparse Jacobians with FiniteDiff. The reason is that, in DI, sparse Jacobians rely on pushforwards, and I don't have a native pushforward operator in FiniteDiff so I make do with a derivative closure.

[oscardssmith]

Can a native pushforward be added to FiniteDiff? it seems like it should be pretty straightforward...

[gdalle]

Knock yourselves out! I'll happily use it.

[gdalle]

In the meantime, maybe I'll try to use the sparse Jacobian inside FiniteDiff as a workaround

[oscardssmith]

JuliaDiff/FiniteDiff.jl#191

[gdalle]

@avik-pal with the very newest versions of DifferentiationInterface (0.6.4) and SparseMatrixColorings (0.4.4), your structured tests from #470 should also pass.
If you just want the thing to work but you're okay with slightly suboptimal colorings until gdalle/SparseMatrixColorings.jl#139 is merged, then you can remove the dependency on a specific branch and use the latest from the registry. I added tests for structured matrices in gdalle/SparseMatrixColorings.jl#137 and they already succeed with the current version of the package, without the optimized implementation of gdalle/SparseMatrixColorings.jl#139.