Jacobian reuse strategy

[huanglangwen]

Many implicit ODE solvers suffer from too many Jacobian evaluations when solving mid-sized stiff ODE problems (~1000?). It is possible to reduce the time cost of Jacobian evaluations by reusing Jacobians from previous steps and utilizing ODE solvers that are less sensitive to the accuracy of Jacobians like those with Newton steps and Rosenbrock-W methods.

As we already have the solvers, we need to implement Jacobian reuse strategies for them. Currently, solvers like TRBDF2 and KenCarp4 are really good in terms of Jacobian reuse, but there is still space for improvements as CVODE_BDF from Sundials suite performs better in some cases. For Rosenbrock-W methods, new reuse strategies are needed since they don't have information like Newton iteration steps to indicate the accuracy of the Jacobian.

Here is what I found from CVODE 1.0.0, it re-evaluates the Jacobian when at least one of the following conditions is satisfied:

1. nst==0 , number of internal steps == 0 .
2. nst - nstlj > CVD_MSBJ , number of steps after last Jac eval exceed a threshold.
3. convfail == FAIL_BAD_J , the value is passed if
   • The previous Newton corrector iteration did not converge and the linear solver's setup routine indicated that its Jacobian-related data is not current.
   • During the previous Newton corrector iteration, the linear solver's solve routine failed in a recoverable manner and the linear solver's setup routine indicated that its Jacobian-related data is not current.
4. dgamma < CVD_DGMAX , the change of gamma between Jac eval exceed a threshold.
5. convfail == FAIL_OTHER
   during the current internal step try, the previous Newton iteration failed to converge even though the linear solver was using current Jacobian-related data.

@ChrisRackauckas @YingboMa

[devmotion]

In my opinion, an important step before/while adding more branches and reuse strategies would be to clean and simplify the current implementation in derivative_utils.jl. There's a lot of redundancy and inconsistency between in-place and out-of-place versions. I assume, further development will benefit from this a lot.

[ChrisRackauckas]

Indeed that needs a cleanup. In particular, I think we should benchmark ForwardDiff.jacobian vs SparseDiffTools' Jacobian. I am not convinced we ever actually need to call into ForwardDiff.jacobian directly, and calling SparseDiffTools all of the time could actually be beneficial (and it's safer for a lot of things like GPUs). If it benchmarks fine on the no-color no-sparsity case, we can just get rid of most of the branches, and instead just branch between DiffEqDiffTools and SparseDiffTools.

(Note that this SparseDiffTools stuff will soon be using https://github.com/YingboMa/ForwardDiff2.jl which will definitely be faster for our use cases)

[huanglangwen]

But oop methods simply don't have the place to cache the previous Jacobian.

[ChrisRackauckas]

That needs to change. The OOP methods need to have Jacobian re-use as well, both for consistency and because it would speed them up. That might be a good place to start since that too would get rid of a lot of branching.

[ChrisRackauckas]

We also need an OOP version of SparseDiffTools' jacobian method. @ranjanan did you ever finish that up?

[huanglangwen]

I'm curious why a whole bunch of oop methods (constant cache + perform_step) exists because one can always transform it into iip?

[ChrisRackauckas]

No you cannot always transform to IIP. If a model is small enough, the fastest method by far is OOP with static arrays. If you use TrackedArray neural networks or Zygote, then you need to use OOP. The TPU stuff we had working for a bit needed OOP, so I am sure compilation to MLIR would require OOP. The most stable form for WASM and direct ptx compilation is making use of the fully stack allocated OOP forms as well. So there's a ton of uses for them, but yes any code with standard arrays should use the in place form.

[ChrisRackauckas]

In that sense, the most interesting reason for W reuse in OOP Ros W is for applications to neural differential equations. That might be something to keep in mind for write-ups.

[huanglangwen]

If OOP methods requires strict OOP all the way, how to write updating W in an OOP fashion? For example, if I add a jacobian term in XXXConstantCache and update it just like IIP, will it work on TPU/GPU/wasm?

[ChrisRackauckas]

Updating a field of a mutable struct is fine. Updating the values of an existing array is not.

[devmotion]

It's the same as updating the state u: oop methods just replace integrator.u with a new u whereas iip methods update the entries of integrator.u.

[huanglangwen]

Okay, so I need to change all the files under caches/ and perform_step/ ?