Support for modified Semi-implicit euler solver

[astanziola]

Hello! I'm trying to include diffrax to my wave simulation library as @patrick-kidger kindly suggested a while ago (and by the way, thanks a lot for this and lineax!) as I really would like to use the advanced checkpointing features :smile:

However, my library also features a pseudo-spectral solver that requires a slight modification of the semi-implicit euler solver, where the discrete update equations also feature an extra $\alpha$ factor, that is:

$$ \begin{cases} u_{n+1} = \alpha(\alpha u_n + \Delta t f(vn)) \ v{n+1} = \alpha(\alpha vn + \Delta t g(u{n+1})) \end{cases} $$

For the interested reader, this kind of integrator appears in equation 2.27 of this manual, and it uses the $\alpha$ term to correct the time-stepping error when absorption is added.

I don't think there's a way to do this with the current SemiImplicitEuler solver, so I wrote the modified version here: https://github.com/astanziola/diffrax/blob/semi-implicit-corrected/diffrax/solver/semi_implicit_euler_corrected.py

I'm very happy to make a PR, but before writing proper tests etc. I wanted to check first if this is really something that can't be done with the current solvers, and most importantly if it is of interest to add such a solver to the library.

Thanks!

[patrick-kidger]

You're correct, this kind of thing is best done by using the AbstractSolver API, for defining entirely new solvers.

This seems like quite a niche solver, if I understand correctly. So I'm not sure this is necessarily worth the maintenance overhead of including it in Diffrax itself?

Nonetheless I appreciate your willingness to contribute!

[astanziola]

Sure, that makes totally sense. I was not sure how much is planned to extend the solvers library. Thanks again!