Use mixed mode AD for Jac evaluation

[SouthEndMusic]

Fixes #2176 (?)

I haven't really looked into scaling yet, but I found an avenue for performance improvement: evaluating the Jacobian with mixed mode AD where the reverse mode is AutoMooncake(). Some notes:

• I haven't been able to make CArrays compatible with Mooncake.jl, so for now the state is a Vector{Float64} and wrapped in a CArray where needed
• Mooncake doesn't like @threaded code but it should be fine for our application, an issue came out of my discussions with a developer: Add support for reverse mode AD of multi-threaded map! operation chalk-lab/Mooncake.jl#791. This is probably not going to be solved before merging this PR. There is a workaround with @mooncake_overlay.
• How many rhs calls are needed for either the forward or reverse mode depends on the sparse matrix coloring, I added a loop in get_diff_eval which tries different colorings to obtain the cheapest way to evaluate the Jacobian
• I found a bug in the DifferentiationInterface code for sparse matrix coloring for mixed mode AD which makes certain small models fail, I've been told this is easy to fix so it will probably be fixed shortly: Sparse Jacobian preparation with mixed mode AD and GreedyColoringAlgorithm(; postprocessing=true) fails JuliaDiff/DifferentiationInterface.jl#857
• I made the mixed mode Jacobian evaluation opt-in as an experimental feature because it adds a lot of precompile time.

[gdalle]

Just popping by to say that the bug you found is indeed easy to fix, but that bidirectional coloring is still fairly new and not widely used. In particular, the number of colors it returns for each side is fairly unpredictable, as you have noticed. I'll be tinkering with ways to weight forward and reverse mode differently, in order to favor forward colors when possible

[SouthEndMusic]

@gdalle In this PR I loop over the orderings for the mixed mode coloring and check the Jacobian evaluation time:

Ribasim/core/src/model.jl

Lines 80 to 104 in 33f082c

	if solver.sparse && experimental.optimize_jacobian
	backend_jac = nothing
	jac_prep = nothing
	t_min_jac_eval = Inf
	for order in
	(NaturalOrder, LargestFirst, SmallestLast, IncidenceDegree, DynamicLargestFirst)
	backend = get_jac_ad_backend(solver; order, mixed_mode = true)
	jac_prep_option = jac_prep_from_backend(backend)
	J = Float64.(sparsity_pattern(jac_prep_option))
	args = (J, u_raw, p, t, jac_prep_option, backend)
	# First evaluate only for precompilation purposes
	jac_from_jac_prep(args...)
	t_jac_eval = @elapsed for _ in 1:10
	jac_from_jac_prep(args...)
	end
	if t_jac_eval < t_min_jac_eval
	t_min_jac_eval = t_jac_eval
	backend_jac = backend
	jac_prep = jac_prep_option
	end
	end
	else
	backend_jac = get_jac_ad_backend(solver)
	jac_prep = jac_prep_from_backend(backend_jac)
	end

As an (opt-in) feature you could do something similar internally where you time a forward and reverse pass to use as weights in the coloring algorithm

[gdalle]

Hi @SouthEndMusic,
@amontoison and I have been working on ways to target specifically the number of row colors or the number of column colors to be minimized during postprocessing. This should allow a bit more control over how many reverse vs. forward autodiff passes are necessary for your mixed-mode Jacobian, but we're still unsure whether our heuristics work. Would you mind trying out the code from the following PR on your matrices, to see if it makes a difference?

Details

using StableRNGs, SparseArrays
using SparseMatrixColorings  # branch https://github.com/gdalle/SparseMatrixColorings.jl/tree/am/postprocessing_v2

J = sprand(StableRNG(0), Bool, 1000, 2000, 0.001)  # replace with your Jacobian sparsity pattern

problem = ColoringProblem(; structure = :nonsymmetric, partition = :bidirectional)

order = RandomOrder(StableRNG(0), 0)

algo_r = GreedyColoringAlgorithm(
    order; postprocessing = true, postprocessing_minimizes = :row_colors
)
algo_c = GreedyColoringAlgorithm(
    order; postprocessing = true, postprocessing_minimizes = :column_colors
)
algo_a = GreedyColoringAlgorithm(
    order; postprocessing = true, postprocessing_minimizes = :all_colors
)

result_r = coloring(J, problem, algo_r);
result_c = coloring(J, problem, algo_c);
result_a = coloring(J, problem, algo_a);

count_colors(result) = length(row_groups(result)), length(column_groups(result))

count_colors(result_r)
count_colors(result_a)
count_colors(result_c)