LBM simulation built in python using JAX and numpy.
This code was developed with Python 3.10.X, tested on Windows and WSL.
Install required libraries:
pip install -r requirements.txt
If you're running on Linux and wish to use the GPU, install JAX with all the cuda helpers:
pip install -U "jax[cuda12]"
Switch between rendering a plot vs recording renders to disk:
RECORD = Falsefor rendering the simulationRECORD = Truefor saving simulation to disk
Switch between the NumPy and JAX implementations:
USE_LIBRARY = NPfor NumPyUSE_LIBRARY = JAXfor JAX
Switch between using gpu and cpu:
USE_DEVICE = cpufor CPUUSE_DEVICE = cudafor GPU via cuda
Note that JAX currently only offers GPU support on Linux, so it is not possible to use the GPU if you are running on windows. NumPy on the otherhand, does not support GPU support on any platform.
On GPU:
On CPU: