JAX-Fluids is a fully-differentiable CFD solver for 3D, compressible two-phase flows. We developed this package with the intention to push and facilitate research at the intersection of ML and CFD. It is easy to use - running a simulation only requires a couple lines of code. Written entirely in JAX, the solver runs on CPU/GPU/TPU and enables automatic differentiation for end-to-end optimization of numerical models.
To learn more about implementation details and details on numerical methods provided by Jaxfluids, feel free to read our paper.
Authors:
Correspondence via mail.
To get started with Jaxfluids, we highly recommend to have a dive into the JuPyter notebooks. All of which can be opened on Google Colab:
- SOD Shocktube:
- Bow Shock:
- Taylor-Green Vortex:
- How to Set Up a Case:
- How to Change the Numerical Setup:
- Automatic Differentiation in Jaxfluids:
- Neural Networks in Jaxfluids:
- Cylinder Flow:
- Laminar Channel Flow:
JAX-Fluids solves the Navier-Stokes-equations using the finite-volume-method on a Cartesian grid. The current version provides the following features:
- Explicit time stepping (Euler, RK2, RK3)
- High-order adaptive spatial reconstruction (WENO-3/5/7, WENO-CU6, WENO-3NN, TENO)
- Riemann solvers (Lax-Friedrichs, Rusanov, HLL, HLLC, Roe)
- Implicit turbulence sub-grid scale model ALDM
- Two-phase simulations via level-set method
- Immersed solid boundaries via level-set method
- Forcings for temperature, mass flow rate and kinetic energy spectrum
- Boundary conditions: Symmetry, Periodic, Wall, Dirichlet, Neumann
- CPU/GPU/TPU capability
Before installing JAX-Fluids, please ensure that you have an updated and upgraded pip version.
To install the CPU-only version of JAX-Fluids, you can run
git clone https://github.com/tumaer/JAXFLUIDS.git
cd JAXFLUIDS
pip install .Note: if you want to install JAX-Fluids in editable mode, e.g., for code development on your local machine, run
pip install --editable .Note: if you want to use jaxlib on a Mac with M1 chip, check the discussion here.
If you want to install JAX-Fluids with CPU and GPU support, you must first install CUDA - we have tested JAX-Fluids with CUDA 11.1 or newer. After installing CUDA, run the following
git clone https://github.com/tumaer/JAXFLUIDS.git
cd JAXFLUIDS
pip install .[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlFor more information on JAX on GPU please refer to the github of JAX
- 5-Equation diffuse interface model for multiphase flows
- CPU/GPU/TPU parallelization based on homogenous domain decomposition
This project is licensed under the GNU General Public License v3 - see the LICENSE file or for details https://www.gnu.org/licenses/.