patrick-kidger
commented
6 months ago "general PDE solving" is a huge field.
As a quick rule of thumb, all ODEs are basically alike, and the same set of techniques work well on nearly all of them. This similarity is what allows packages likes Diffrax to exist.
Unfortunately, many PDEs are each special snowflakes that demand their own custom solver. For this reason general PDE solving is outside the scope of Diffrax!
As for the Fokker-Planck equation, it will probably depend on how many dimensions you're working in. For low dimensions then I imagine you could semidiscretise to an ODE in exactly as is done in the example you link, and have that work fine. In high dimensions I imagine other more specialised techniques will become necessary.
Beyond that I'm not really equipped to say, I'm afraid: I've not made a serious study of numerical methods for PDEs, so I suspect other references will be able to give you better answers than I can.
lockwo
Are the methods for simulating the heat PDE as presented in https://docs.kidger.site/diffrax/examples/nonlinear_heat_pde/ where you discretize \pdv y(x, t) into d/dt y_i(t) for N points a good/valid approach for general PDE solving? For context, specifically of interest to me is the Fokker-Planck PDE (https://en.wikipedia.org/wiki/Fokker–Planck_equation).