tumaer/jaxfluids

Differentiable Fluid Dynamics Package

JAX-Fluids: A Differentiable Fluid Dynamics Package

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 JAX-Fluids, feel free to read our paper. And also check out the documentation of JAX-Fluids.

Authors:

• Deniz A. Bezgin
• Aaron B. Buhendwa
• Nikolaus A. Adams

Correspondence via mail.

JAX-FLUIDS 2.0 IS COMING IN 2024!

JAX-Fluids 2.0 is coming in 2024! Here, is a short overview on what JAX-Fluids 2.0 will bring to the table:

• Support for parallel simulations
• Diffuse-interface model for two-phase simulations
• Many new & robust high-order numerics
• New boundary conditions
• Performance improvements for forward simulations and backward passes
• Completely updated & more intuitive I/O
• And many more...

Stay tuned and feel free to reach out!

Physical models and numerical methods

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

Example simulations

Space shuttle at Mach 2 - Immersed solid boundary method via level-set

Shock-bubble interaction with diffuse-interface method - approx. 800M cells on TPUv3-64

Shock-bubble interaction with level-set method - approx. 2B cells on TPUv3-256

Shock-induced collapse of air bubbles in water (click link for video)

https://www.youtube.com/watch?v=mt8HjZhm60U

Pip Installation

Before installing JAX-Fluids, please ensure that you have an updated and upgraded pip version.

pip install --upgrade pip

CPU-only support

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.

GPU and CPU support

If you want to install JAX-Fluids with CPU AND GPU support, you must first install JAX with GPU support. There are two ways to do this:

1. installing CUDA & CUDNN via pip,
2. installing CUDA & CUDNN yourself.

See JAX installation for details.

We recommend using CUDA & CUDNN using pip wheels:

pip install --upgrade pip

# CUDA 12 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

After having installed CUDA, run the following to install JAX-Fluids

git clone https://github.com/tumaer/JAXFLUIDS.git
cd JAXFLUIDS
pip install -e .

For more information on JAX on GPU please refer to the github of JAX

Quickstart

This github contains five jupyter-notebooks which will get you started quickly. They demonstrate how to run simple simulations like a 1D sod shock tube or a 2D supersonic cylinder flow. Furthermore, they show how you can easily switch the numerical and/or case setup in order to, e.g., increase the order of the spatial reconstruction stencil or decrease the resolution of the simulation.

Upcoming features

• 5-Equation diffuse interface model for multiphase flows
• CPU/GPU/TPU parallelization based on homogenous domain decomposition
• Lagrangian particles

Documentation

Check out the documentation of JAX-Fluids.

Acknowledgements

We gratefully acknowledge access to TPU compute resources granted by Google's TRC program.

Citation

https://doi.org/10.1016/j.cpc.2022.108527

@article{BEZGIN2022108527,
   title = {JAX-Fluids: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows},
   journal = {Computer Physics Communications},
   pages = {108527},
   year = {2022},
   issn = {0010-4655},
   doi = {https://doi.org/10.1016/j.cpc.2022.108527},
   url = {https://www.sciencedirect.com/science/article/pii/S0010465522002466},
   author = {Deniz A. Bezgin and Aaron B. Buhendwa and Nikolaus A. Adams},
   keywords = {Computational fluid dynamics, Machine learning, Differential programming, Navier-Stokes equations, Level-set, Turbulence, Two-phase flows}
} 

License

This project is licensed under the GNU General Public License v3 - see the LICENSE file or for details https://www.gnu.org/licenses/.

jaxfluids's Issues

Is it possible for JAXFLUIDS to import a CAD model and generate AMR grids during the run

Hi. I saw the test case: Space shuttle at Mach 2 - Immersed solid boundary method via level-set. I wonder how to import a CAD model to the JAXFLUIDS solver. Thanks in advance.

Encountering AttributeError with jaxfluids and JAX 0.4.14

Problem

I'm using the jaxfluids library, which is a computational fluid dynamics library that requires JAX. I have installed JAX and jaxlib with versions jax=0.4.14 and jaxlib=0.4.14, respectively. However, when running the provided examples, I encountered the following AttributeError:

AttributeError: module 'jax.numpy' has no attribute 'DeviceArray'

I noticed that this issue has been discussed in various contexts, including GPT-3's responses. Is there a recommended solution for resolving this issue with the mentioned versions of JAX?

Solution

After researching and experimenting, I found a workaround to resolve this issue. I used PyCharm's "Replace in file" feature to replace all instances of jnp.DeviceArray with jnp.Array. This helped me successfully run the examples without encountering the AttributeError.

Here's how I performed the replacement:

Open the file where jnp.DeviceArray is used.
Use the "Replace in file" feature in PyCharm.
Replace all occurrences of jnp.DeviceArray with jnp.Array.
Please note that this is a specific workaround I tried, and it may not be the only solution. However, it worked for me with the mentioned versions of JAX. Others who are facing similar issues might find this workaround helpful.

Result

Success

weakly compressible low mach

Hi,

I am just a bit curious what the plans of jaxfluids are :-) Is there any activity ongoing (or planned) for low Mach turbulent flow as well?

Best Regards
Fab

Fluid-Fluid-Solid simulations?

Hello,

Is it currently possible to do triple-phase simulations involving two fluid phases and a solid phase?

import error

Dear developers
I get an error when I use from jaxfluids.simulation_manager import SimulationManager。The reson of this error is that from jax.config import config is used in /JAXFLUIDS/src/jaxfluids/simulation_manager.py.In fact, jax suggest using this way.
import jax jax.config.update(your_config_options_here)
Thank you!