patrick-kidger
3D heat diffusion
Please excuse the somewhat open-ended question!
Could you comment on the suitability of Diffrax for modelling the transient heat diffusion equation in 3D? I have followed the 1D example in the docs, where you obtain a system of coupled ODEs by discretising the spatial domain. I understand this could be extended to 3D space, indeed Crank-Nicolson is often applied to 3D space, but I am aware it would require me to write my own solver. I wonder whether this is not the sort of thing you originally wrote Diffrax for. Other libraries (FEniCSx, FiPy) exist for this sort of thing, but they are outside the JAX ecosystem. Diffrax still appeals for performance reasons and because I like the documentation.
(just linking other relevant PDE issues: https://github.com/patrick-kidger/diffrax/issues/394, https://github.com/patrick-kidger/diffrax/issues/177)
I think that should totally be possible.
The scope of Diffrax is essentially anything with a finite-difference discretisation down a temporal axis. That means ODEs and SDEs, of course. But it also means a lot of the methods for solving parabolic PDEs etc are in-scope as an application well, in specifically a bring-your-own-spatial-numerics kind of way. (Whilst letting Diffrax handle the temporal part -- adaptive stepping, interpolations, backprop through variable-length loops, etc.) The 1D example in the documentation is a great example of the kind of thing that is in-scope.
I hope that helps!
Thank you for the reply. In the end, I have decided to use FiPy for ease, but maybe I will return to this one day. It would have been fun to write a 3D Crank-Nicolson solver for Diffrax to timestep but I don't think I have the time.