Creating Support for Delay Differential Equations

[PraneethMerugu]

Hey all!

I am Praneeth Merugu, an undergraduate at Georgia State, and I have been using diffrax for quite awhile. I have noticed that adding support for delay differential equations is relatively low hanging fruit given the awesome tools Patrick has made. I recently had the curiosity and researched how current literature handles them, as well as delaydiffeq.jl. From my research, it seems as though we have 3 issues:

1. Function History - Because DDEs reference previous function values and their derivatives, my implementation preserves the dense output of the solver at each step. This means that the solver that is being wrapped will preserve its order during the delay
2. Discontinuities in the delay - Since discontinuities are a major source of error within DDE solvers, much of the difficulty lies here. To ensure that we preserve the order of our solver, we have to make sure our solver doesnt overshoot a discontinuity, but rather solves exactly at that value. By setting up a optimization problem within a step size controller, we were able to acheive discontinuity tracking for variable delay diff eqs.
3. Memory Management - Since Jax hates reallocating memory (tracer errors :( ), I employed solve functions that either manage a circular buffer of dense histories, or splits the solve into a jax.lax.scan where it is solved in chunks.

I would love to work on this further and eventually make a contribution of some sorts. All feedback is welcome.

https://github.com/PraneethMerugu/diffraxdde

Thanks!
Praneeth Merugu

[lockwo]

You might be interested in #169

[patrick-kidger]

Yup, I'd love to find the time to finish up DDEs some time 😅