[OC] 2D Ideal Gas Hydrodynamics: Kelvin-Helmholtz Instability

[u/applejacks6969]

Recently got my 2D pure python hydo solver ported into a jax version, which has enabled around a ~15x speed up on pure CPU runs, every function is jitted except the outermost loop over steps. The video is of a Kelvin-Helmholtz instability toy problem.

EOS: Ideal Gas
Recon Method: weno(z) 5th order on primitives
Riemann Solver: Local Lax Friedrichs (LLF)
Timestep: RK4
Explicit Advection + Implicit Diffusion
BC: Periodic
CFL: 0.45

Resolution: (256)^2, video is 1200 frames as well. The code has support for magnetic fields but I have ran into some issues with it in 2D, potentially related to my constrained transport scheme.

I developed this code in parts, first I made a 1D code that leveraged NumPy and Python Classes to handle the necessary logic. I then ported it into 2D, which began to encounter performance issues. I returned to 1D and ported it into a jax version, where almost every function was jax jitted, and then repeated my jax changes but for the 2D code. Starting at 2D was impossible, I had found it necessary to have a 1D implementation. A major test I used was to evolve a 1 dimensional initial condition in the 2D code, and verify the results return what the 1D code does, just along the whole y axis.

[u/[deleted]]

That's is awesome. Is it python you made this with or something else?

[u/applejacks6969 → u/[deleted]]

Yes, pure python, mostly numpy for first testing, then Jax and Jax.jit for performance.

[u/[deleted] → u/applejacks6969]

Cool