Strict in-place sparse Jacobians now cause errors when they don't anticipate the mass matrix

[gdalle]

Describe the bug 🐞

In #2567, it seems that the sparsity pattern is modified using the mass_matrix outside of the DI pipeline (see #2567 (comment)). This is an issue because:

1. during preparation, the sparsity_detector from AutoSparse won't be aware of this change
2. during preparation, the coloring_algorithm from AutoSparse will compute a coloring for the unmodified sparsity pattern, and prepare decompression accordingly
3. during execution, the function SparseMatrixColorings.decompress! will either error or return incorrect results because it expected the unmodified sparsity pattern

This problem was detected by Deltares/Ribasim#2201

Expected behavior

Any modification of the sparsity pattern should be done in a way that allows the sparsity_detector to see it.

Minimal Reproducible Example 👇

Not available yet.

Error & Stacktrace ⚠️

The typical stack trace could look something like:

ERROR: DimensionMismatch: all inputs to eachindex must have the same indices, got Base.OneTo(28) and Base.OneTo(26)
Stacktrace:
  [1] throw_eachindex_mismatch_indices(::IndexLinear, ::Base.OneTo{Int64}, ::Vararg{Base.OneTo{Int64}})
    @ Base ./abstractarray.jl:325
  [2] eachindex
    @ ./abstractarray.jl:393 [inlined]
  [3] eachindex
    @ ./abstractarray.jl:382 [inlined]
  [4] decompress!
    @ ~/.julia/packages/SparseMatrixColorings/BED6Y/src/decompression.jl:361 [inlined]
  [5] _sparse_jacobian_aux!(::Tuple{typeof(Ribasim.water_balance!), Vector{Float64}}, 
  [6] jacobian!
    @ ~/.julia/packages/DifferentiationInterface/7eD1K/ext/DifferentiationInterfaceSparseMatrixColoringsExt/jacobian.jl:237 [inlined]

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

Not available yet

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

Not available yet

• Output of versioninfo()

Not available yet

Additional context

This should only happen for people who have upgraded to OrdinaryDiffEqDifferentiation v1.5.0, released a few hours ago

Ping @jClugstor @oscardssmith @ChrisRackauckas

[gdalle]

Actually I'm not sure I fully understand what goes on.

The modifications to the sparsity pattern are done here, provided that sparsity isa SparseMatrixCSC:

OrdinaryDiffEq.jl/lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

Lines 109 to 129 in da8d130

	if sparsity isa SparseMatrixCSC
	if prob.f.mass_matrix isa UniformScaling
	idxs = diagind(sparsity)
	@. @view(sparsity[idxs]) = 1
	if !isnothing(jac_prototype)
	@. @view(jac_prototype[idxs]) = 1
	end
	else
	idxs = findall(!iszero, prob.f.mass_matrix)
	for idx in idxs
	sparsity[idx] = prob.f.mass_matrix[idx]
	end
	if !isnothing(jac_prototype)
	for idx in idxs
	jac_prototype[idx] = f.mass_matrix[idx]
	end
	end
	end
	end

And the creation of the sparsity detector happens here:

OrdinaryDiffEq.jl/lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

Lines 132 to 138 in da8d130

	sparsity = sparsity isa MatrixOperator ? sparsity.A : sparsity
	color_alg = DiffEqBase.has_colorvec(prob.f) ?
	SparseMatrixColorings.ConstantColoringAlgorithm(
	sparsity, prob.f.colorvec) : SparseMatrixColorings.GreedyColoringAlgorithm()
	sparsity_detector = ADTypes.KnownJacobianSparsityDetector(sparsity)

So it seems that whenever the modifications happen, they are passed along to the detector ADTypes.KnownJacobianSparsityDetector, since sparsity isa SparseMatrixCSC implies !(sparsity isa MatrixOperator).

Still, the error seen in Deltares/Ribasim#2201 had something to do with the diagonal of the sparsity pattern being tampered with after detection, so my gut tells me the problem is here somehow. Perhaps the error is in the way the Jacobian prototype is built:

OrdinaryDiffEq.jl/lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl

Lines 764 to 779 in da8d130

	J = if f.jac_prototype === nothing
	if alg_autodiff(alg) isa AutoSparse
	if isnothing(f.sparsity)
	!isnothing(jac_config) ?
	convert.(
	eltype(u), SparseMatrixColorings.sparsity_pattern(jac_config[1])) :
	spzeros(eltype(u), length(u), length(u))
	elseif eltype(f.sparsity) == Bool
	convert.(eltype(u), f.sparsity)
	else
	f.sparsity
	end
	else
	ArrayInterface.zeromatrix(u)
	end
	else

[ChrisRackauckas]

Well the diagonal has to be added. How do you make that work in DI?

[ChrisRackauckas]

To handle the DI jank we could do something like SparseMatrixColorings.sparsity_pattern(abs.(jac_config[1]) .+ abs.(f.mass_matrix)). It's a bit weird but it would get the job done.

[gdalle]

I think there is no issue if the sparsity that gets computed with the added mass_matrix also gets passed to the detector. Thus, the question is what happens when sparsity isa MatrixOperator, because then the detector gets sparsity.A instead. And even when !(sparsity isa SparseMatrixCSC), because those changes to the pattern still happen in other places.

[ChrisRackauckas]

We can ad the mass matrix part wherever it's needed. The function for the nonlinear solve won't exist yet though since that construction has the sparsity pattern as part of it. But we can just jank it to be abs.(MM) .+ abs.(f(u,p,t)) for the first version as a safe fix, and figure out a nicer handling later.

[gdalle]

The fundamental problem is the following: when the user provides an in-place sparse Jacobian callable jac!(J, x), it is tailored to their f! and the sparsity pattern of its Jacobian J. Later on, inside OrdinaryDiffEq, you essentially define something like f!_bis = f! + identity, and then you try jac!(J_plus_identity, x), which doesn't work whenever J and J_plus_identity have different sparsity patterns.

[ChrisRackauckas]

We could just define that the prototype sparsity needs to include the structural nonzeros of the mass matrix as well. We were able to handle that implicitly behind the scenes before though, I'm not sure what tooling is not robust to it? Does DI not differentiate between the sparsity pattern and the Jacobian to be written to?

[gdalle]

Does DI not differentiate between the sparsity pattern and the Jacobian to be written to?

This part of the pipeline actually happens in SparseMatrixColorings, which handles the decompression of a compressed Jacobian B (computed color by color) into a sparse or structured matrix A. There are two methods for this decompression:

decompress!(A::SparseMatrixCSC, B, coloring_result)
decompress!(A::AbstractMatrix,  B, coloring_result)

For the first one, which is the one we're hitting, the assumption is that A will have the same pattern as the sparsity pattern S initially provided to the coloring algorithm.

In the case of Deltares/Ribasim#2201,

• S is the true sparsity pattern of f, without the added diagonal
• A is the jacobian prototype where OrdinaryDiffEq has added a diagonal

This assumption of identical sparsity patterns is not only intuitive, it is also essential to enable fast decompression because we can reason directly on the vector of nonzeros(A). Otherwise the process would be much slower.

If Deltares/Ribasim#2201 was using the automatic sparse Jacobians from OrdinaryDiffEq, this bug probably wouldn't occur because they are prepared using an augmented sparsity pattern S + I. But of course some users will still want to provide their own Jacobians, and those will be tailored to a decompression into A = similar(S, Float64), not A + I.

[gdalle]

Essentially, the question boils down to: should DI be able to compute

jacobian!(J, prep, backend, f, x)

when J isa SparseMatrixCSC BUT its sparsity pattern is different than the detected pattern S used during the coloring inside prepare_jacobian?

• If sparsity_pattern(J) is a subset of sparsity_pattern(S), this should definitely error
• If sparsity_pattern(J) is a superset of sparsity_pattern(S), it is a subjective design choice. Right now it errors, but we could make it work, at the cost of lower performance. And there is also a question of what to do with the remaining entries (fill them with zeros?)

[oscardssmith]

Isn't it just a bug that we aren't preparing with jac+MM?