Stan differential equation functionality development priorities

Tagging @charlesm93, @wds15, @yizhang, and @bbbales2 to start for this question, since I know you guys are working on various things related to the topic.

I was wondering about the general state of priorities right now for ODEs in Stan and what present areas of development focus were. Over the past few months, among the in-progress tasks I’ve seen include variadic arguments into the ODE function, DDEs, and non-linear root solvers. I’ve seen that the ability to use DifferentialEquations.jl is one of the things that is making Julia-native PPL’s more appealing for some whose work is more focused on differential equation system inferences, so general improvement in this area would probably be of interest to a lot of PPL users.

Generally, it looks like these angles are being worked on?

Edit: Some of these bullet points are not, so this is now more intended to be taken as a list of ripe areas of contribution for improving Stan’s differential equation functionality.

  1. Variadic arguments in ODE function
  2. DDE/SDE/other diff eq type solvers
  3. Additional ODE solvers (don’t need too many, but having a couple more options other than rke45 and bdf might be useful for versatility). An auto-switching algorithm that automatically switches between stiff and non-stiff would be the most useful addition in this area.
  4. PDE solvers
  5. Adjoint method for propagating derivatives to ODEs
  6. Deterministic event handling, as @wds15 mentions below, which would help with forcing function manipulations that happen a lot in ecology and pharmacology.
  7. Dynamic event handling, also mentioned by @wds15 below. Non-linear root-finding would be a very useful subset of this, as that would help larger ODE models be initiated at a non-analytic steady state prior to a perturbation response.

Of most use to me are introductions of the non-linear root solving, which I think @charlesm93 is working on, and the variadic argument update, which I think is being headed by @bbbales2, as they directly correspond to the next project I’m going to be working on for my PhD thesis, so I figured it was high time I took a look under the Stan hood and learn about contributing to one of those areas. However, I was also just generally interested in the state of what was in the works on the Stan differential equation front. Thanks!

3 Likes

Variadic arguments in ODE function

I’m working on this now. I want to have it done for the next release (roughly 2 months out?). There’s a design doc here and discussion here.

Adjoint ODEs might be able to hit the next release too. I’ve historically been unnecessarily pessimistic about the difficulty of doing that. Anyway if these are gonna go in the next release I probably need to finish the math stuff this month.

DDE/SDE/other diff eq type solvers

No developments in core Stan that I know of. @Charles_Driver is the SDE guy I know off the top of my head so he might have opinions here.

Non-linear Newton-Krylov root solvers for non-linear ODEs
Additional ODE solvers

I would like more ODE solvers. DifferentialEquations.jl seems to be the gold standard. If there are any steps we can do to identify and close the gaps, I’m in for it.

PDE solvers

Of most use to me are introductions of the non-linear root solving

If you want to do the variadic version of the algebra solver, that would be great. Or anything really.

If you want to set up the adjoint stuff, that’d also be helpful, but you’ll have to deal with the work in progress variadic ODEs branch.

@jtimonen has done some experiments with fixed-step size ODE solvers for correcting approximate solves. People using Stan have found it useful to just do their own solvers: Costs/benefits of solving large ODE systems as systems of difference equations/in discrete time , so we were looking at ways they could diagnose these approximation.

1 Like

Also when I say I’m working on this now… That means at this moment I am working on it, but if you want to work on it as well, you’re more than welcome, though it’s at the boring make-tons-and-tons-of-tests stage of development :/.

@bbbales2 I’m happy to get started with helping with the C++ testing and contributing to some of the more boring, rote tasks as part of the learning process. I can individually contact you about that to not clutter up the forum.

As for getting additional ODE solvers, I have more of a higher-up understanding of it since differential equations, dynamical systems, and numerical analysis are not close to my specific research field so I won’t be able to commit taking the lead on that. But for a conceptual start, I think adding the Vern7 or DP8 algorithm would be an additional non-stiff solver to use for those that want something more accurate than rk45 but still non-stiff, a BS3/ode23 algorithm for people who need speed, and another stiff solver to compare to BDF, such as a Rosenbrock/W method. An auto-switching method that can go between stiff and non-stiff would offer lots of versatility.
Chris Rackauckas lists the extensive amount of solvers that DifferentialEquations.jl offers pretty neatly here as something that Stan can then drawn from for its ODE development efforts: https://docs.sciml.ai/dev/solvers/ode_solve/#ode_solve-1.

We try to do as much as possible on the forums/Github. Even video calls kinda suck. Information gets lost, people get bored and forget, etc.

Much better for the project that the information just gets blasted out everywhere 24/7. It doesn’t really matter if you say something wrong. It can get corrected.

I’ll get what I have of the variadic stuff up later today. What is your Github handle so I can @ you?

Yes I would really like something like this.

If you know what all these algorithms are, that’s more than I know.

Yeah, 100%, I meant that more as in continuing with the variadic testing discussion would be more suited for Github, apologies for my lack of clarity. My Github handle is @wallyxie.

Unfortunate part for me is that the implementation of those algorithms would be separate enough from my funding and thesis topics that I wouldn’t be able to take time away from finishing up my thesis to really dig into the implementation, but I’m happy to get the ideas started at this point. Could be a worthwhile project for someone working at the intersection of numerical analysis and HMC, perhaps someone working in Mark Girolami’s group.

Hi!

The other major item is the adjoint method for propagating derivatives through solutions to ODEs; see this discussion.

Non-linear Newton-Krylov root solvers for non-linear ODEs

I did this in the summer and added it to the math library. People didn’t bug me enough for me to add this to the Stan language (especially given we already have another solver), but I’ll expose it to the language if this helps you.

1 Like

We have someone working on this? I didn’t know that.

Oups – I misread. We have a Newton-Krylov solver, but it’s not specifically applied to non-linear ODEs.

@charlesm93 I don’t think Krylov solver is implemented there. It’s Newton with direct solver.

Oh, that would be great.

Edit: I see the later correction. Yeah, I meant specifically to initiate from roots for non-linear ODEs before driving a response.

And I’ll go ahead and add that as a bullet on my initial post.

Never mind, looks like I was wrong. Would be a good thing to add though. Also thought someone earlier was already working on additional ODE solvers, but I must have mis-read and jumbled my memory. Anyhow, those would be some good areas for contributions to expand Stan ODE functionality. Will edit my post to reflect that those bullets are more areas of need rather than active areas of development.

Can you provide (or point to) a bit of background on the solvers you like to add to Stan?

First of all, there is also the Adams-moulton non-stiff solver in Stan. That one was not documented for a long time, but it’s there (and now documented).

We also have a differential algebraic solver in Stan Math which was never exported to the language as the respective PR in Stan faded away unfortunately. Hopefully @yizhang or someone else can add that thing to the stanc3 compiler.

Other than that I do think that an adjoint sensitivity method would be super important for Stan. That’s probably the most important thing to get into Stan from my perspective. I very much hope that we can enable that easily for the Adams-Moulton and the BDF solvers soonish.

Another interesting spin which we intend to try out as part of the variadic ODE stuff is to allow for different tolerance targets for the function itself vs the gradients. Going with lower precision for the gradients should be doable and will certainly speedup things - but that needs testing.

I am really not sure if we need more solvers. These will confuse users a lot unless there are clear scenarios where one outperforms better than another one, which we should ideally be able to articulate.

We would gain a lot of speed if we could symbolically generate the Jacobian of the ODE RHS. In my testing that gives you 3x speed increases on smallish examples (3 states, 5 parameters). Though, I have no hope that we get that any time soon.

1 Like

Hi Sebastian,

I actually somehow entirely missed that Adams-Moulton has been available since 2018! Foolish oversight on my part, and checking the documentation, lo and behold integrate_ode_adams it is there. It seems like people still like to say that Stan has only non-stiff solver available, so I’m going to push back against that where I see it online.

With respect to this, I am now in agreement that we do not need more non-stiff ordinary differential equation solvers after learning that integrate_ode_adams exists. I do think that perhaps the addition of an auto-switching algorithm that swaps between non-stiff and stiff as necessary could be useful for users uncertain about the stiffness of their ODEs or those that want to just plug and play with their respective system with less initial thought to the numerical analysis. I have not found as much formal literature about auto-switching algorithms, but it seems like they are one of the more popular ODE algorithms chose for use from DifferentialEquations.jl because of their ease of use, and a list of them is presented here: https://docs.sciml.ai/dev/solvers/ode_solve/#Pre-Built-Stiffness-Detecting-and-Auto-Switching-Algorithms-1

For not confusing users, perhaps it would be also in a future version of Stan for each solver to not have its own integrate_ode function, but having each solver being an argument in a more general integrate_ode function with an auto-switching algorithm being the default? If there is a specific reason for needing a separate function for each solver, I apologize for not having been familiar with it. Even if the general consensus is that adding solvers is not a major priority, I feel that if someone did want to add more, that needing to have a new function for each one would be unwieldy for development and use alike.

With respect to other kinds of differential equations, it would be good to have a roadmap for implementing/reviving DAE, DDE, and SDE solvers.

I was honestly less familiar with adjoint sensitivity earlier, but having read more about sensitivity computation over the past few hours, I definitely agree with its prioritization.

The Adams-Moulton was forgotten to get documented. It’s now in the doc since 2.23 for the first time, but is around for much longer (this not how it is supposed to be).

An auto-switching solver between stiff / non-stiff would be nice, yeah. I did follow the Sundials discussion on that and they deemed it unnecessary as people should know their problem (this what they concluded). Since problems are either stiff or non-stiff, I tend to agree that it’s fair to say “just try each method and if in doubt just use a stiff solver”. Maybe some problems would benefit though from switching.

I find that setting the right tolerances is really key to good performance - this is why I am keen on getting extra tolerance control over the state vs the gradients.

And yeah - for larger systems - we absolutely need the adjoint method, but I think @bbbales2 is on that (like on so many other things).

EDIT: I forgot… what would be really useful for things I do are event handlers. So, I would like to have

  • deterministic events: “stop integrator at time td and add a dose d to one compartment, then continue integrating”; different events are useful here like turning on/off forcing functions
  • dynamic events: determine the time-point when the solution crosses some value (root-finding problem)

the deterministic events can be coded in Stan directly, but that takes a lot of Stan code to do it…

2 Likes

Oh I was curious about why this was, thanks for checking.

Yes, completely agree with both of those points you mentioned. Edited them into the first post list. Additions on those fronts would expand the appeal and utility of Stan to some other dynamical modeling heavy fields. Deterministic event capabilities would make a lot of more complicated ecological models easier to translate into Stan, and dynamic event handling would additionally be of great use for computational neuroscientists and biophysics. The non-linear root-finding I mentioned earlier is better re-framed under the broader dynamic events umbrella.

In case someone is keen adding new ODE integrators, I worked on a C interface for FATODE. The interface was designed to be consistent with that of Stan Math. Some of the integrators there are superior for certain stiff problems. there’s some overlapping to what sundials’ ARKode offers, which unfortunately doesn’t do sensivity.

2 Likes

One of the reasons that sundials’ CVODE doesn’t do stiffness detection is that we’re way beyond the era of LSODA where a dominant eigenvalue of Jacobian was deemed sufficient to quantify stiffness. On one hand many nonlinear system defy stiffness detection by exhibiting cyclic transition region where stiffness emerges and disappears, on the other hand the detection technique usually involves 1) calculatingJacobian and 2) estimating eigenvalue which in turn requires many matrix-vector multiplication. For large dynamic system this becomes prohibitively expensive. My experience is that auto-switching is a pain to developers and can be a foot gun to users.

3 Likes

Thank you, @yizhang, for this information. I learned something new today and did not know those limitations existed. If you are able to take a cursory look, what is your opinion on the implementation of DifferentialEquations.jl’s auto-switching ODE solvers (linked to their description page above)?

Edit: Chris Rackauckas let me know he has a paper about his auto-switching methods in the works after I contacted him. They switch from non-stiff to stiff easily enough, but still are figuring out how to make it more easily go back from stiff to non-stiff as necessary. Will post a link to that paper when it comes out.

I don’t have time right now, unfortunately, but would be keen on helping to add solvers if time permits in the months ahead. In the mean time, if someone else runs into this thread later and wants to take the lead on this (maybe another graduate student with a summer project), do go ahead.