solver and stepsize controller settings suggestions for stiff problem

[DylanEsguerra]

Hello

I have been using diffrax to solve pharmacokinetic models. The ones I am working with typically have a drug dose at an interval over time. For example every 2 weeks. I want to ensure my solver steps on these points and have been using PID controller with jump_ts for this. At the same time the problem is very stiff with reactions at different time scales. Does anyone have suggestions for a solver best suited for a stiff problem with a variable that changes due to a dosing schedule? Thank you for the help

[patrick-kidger]

For a stiff solver I'd probably recommend diffrax.Kvaerno5.

Using diffrax.PIDController is the right choice to pair with this, in particular by setting pcoeff and icoeff (see the documentation on choosing PID coefficients).

If you want to get outputs at particular times then you can do so with diffrax.SaveAt(ts=...). If you additionally want to be sure the solver steps to exactly these points (for highest possible accuracy in the numerical solve), then the modern way (just released!) is wrap the step size controller in a diffrax.ClipStepSizeController.

Thus put together, you'd end up with something like the following:

ts = [...]
saveat = diffrax.SaveAt(ts=ts)
solver = diffrax.Kvaerno5()
controller = diffrax.ClipStepSizeController(diffrax.PIDController(pcoeff=0.3, icoeff=0.4), step_ts=ts)
diffrax.diffeqsolve(solver=solver, stepsize_controller=controller, saveat=saveat, ...)

I hope that helps!

[DylanEsguerra]

Thank you for the advice, I was not aware of the importance of the icoeff and pcoeff parameters and was mostly focused on tuning tolerances.

Currently the Tsit() solver seems to work best for my problem and I cannot seem to get Kvaerno5() to start. This has been difficult since on some smaller similar models Kvaerno5() performed very well. Using the progress meter it never seems to get started and I am wondering if perhaps I am not giving it long of a chance to get warmed up?

I noticed on smaller problems it would take longer than Tsit() to get started but once it did, it was much faster.
Is this a trait of the Kvaerno5() solver to take awhile before showing 0% or 2% completion and then jump into solving quickly?

[patrick-kidger]

On the smaller problems -- it might just be compilation time? It's certainly the case that Kvaerno5 is a more complicated solver and takes a little longer to compile. (If so then there's also a separate issue with long compile times that we're aware of in #606, which we'll have fixed in #614.)

On larger problems -- as an implicit solver, the computational cost of Kvaerno5 scales as O(n^3) in the number of dimensions of y ∈ R^n. It might simply be that you're using too large a problem for an implicit solver to be viable.