Add `diffrax` time integrators
astanziola opened this issue · comments
The current time integrator is only "optimal" for FourierSeries fields thanks to the k-space operator correction. For finite differences, it would probably be better to allow the user to experiment with different kind of integrators.
This is a good reason for using diffrax as the main ODE integrator library for time domain solvers.
Another good reason would be to expose to the user the advanced checkpointing methods enabled by diffrax.
An initial attempt to code the modified semi-implicit Eluer integrator used by the k-space equations is here: https://github.com/astanziola/diffrax/blob/main/diffrax/solver/semi_implicit_euler.py
It works pretty well, timings are roughly the same, and checkpointing works.
The API needs to be changed to make sure that the user can generally choose his favorite integrator.