This is a small (compact single-script), didactic implementation of the Lattice Boltzmann Method (LBM) in JAX, which should scale very efficiently.
Tip
You may directly jump into the code here. The example showcased below is generated at the end of the file.
The general mathematical formulation follows Ataei and Salehipour, 2024 and the obstacle handling and test case is based on https://github.com/pmocz/latticeboltzmann-python.