Advanced implementation of the numerical solver

[jiaxi98]

For higher order PDE, the condition number of the discretization matrix can be very big for small grid size. Use precondition method to improve the accuracy.

[jiaxi98]

jax-cfd package would be a good template to learn from:
https://github.com/google/jax-cfd/blob/main/jax_cfd/spectral/time_stepping.py