+ 
-# XLB: A Hardware-Accelerated Differentiable Lattice Boltzmann Simulation Framework based on JAX for Physics-based Machine Learning
+# XLB: Distributed Multi-GPU Lattice Boltzmann Simulation Framework for Differentiable Scientific Machine Learning
-XLB (Accelerated LB) is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) solver that leverages hardware acceleration. It's built on top of the [JAX](https://github.com/google/jax) library and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning.
+XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It's built on top of the [JAX](https://github.com/google/jax) library and is specifically designed to solve fluid dynamics problems in a computationally efficient and differentiable manner. Its unique combination of features positions it as an exceptionally suitable tool for applications in physics-based machine learning.
+
+## Accompanying Paper
+
+Please refer to the [accompanying paper](https://arxiv.org/abs/2311.16080) for benchmarks, validation, and more details about the library.
+
+## Citing XLB
+
+If you use XLB in your research, please cite the following paper:
+
+```
+@article{ataei2023xlb,
+ title={{XLB}: Distributed Multi-GPU Lattice Boltzmann Simulation Framework for Differentiable Scientific Machine Learning},
+ author={Ataei, Mohammadmehdi and Salehipour, Hesam},
+ journal={arXiv preprint arXiv:2311.16080},
+ year={2023},
+}
+```
## Key Features
-- **Integration with JAX Ecosystem:** The solver can be easily integrated with JAX's robust ecosystem of machine learning libraries such as [Flax](https://github.com/google/flax), [Haiku](https://github.com/deepmind/dm-haiku), [Optax](https://github.com/deepmind/optax), and many more.
-- **Scalability:** XLB is capable of scaling on distributed multi-GPU systems, enabling the execution of large-scale simulations on hundreds of GPUs with billions of voxels.
+- **Integration with JAX Ecosystem:** The library can be easily integrated with JAX's robust ecosystem of machine learning libraries such as [Flax](https://github.com/google/flax), [Haiku](https://github.com/deepmind/dm-haiku), [Optax](https://github.com/deepmind/optax), and many more.
+- **Differentiable LBM Kernels:** XLB provides differentiable LBM kernels that can be used in differentiable physics and deep learning applications.
+- **Scalability:** XLB is capable of scaling on distributed multi-GPU systems, enabling the execution of large-scale simulations on hundreds of GPUs with billions of cells.
- **Support for Various LBM Boundary Conditions and Kernels:** XLB supports several LBM boundary conditions and collision kernels.
-- **User-Friendly Interface:** Written entirely in Python, XLB emphasizes a highly accessible interface that allows users to extend the solver with ease and quickly set up and run new simulations.
-- **Leverages JAX Array and Shardmap:** The solver incorporates the new JAX array unified array type and JAX shardmap, providing users with a numpy-like interface. This allows users to focus solely on the semantics, leaving performance optimizations to the compiler.
+- **User-Friendly Interface:** Written entirely in Python, XLB emphasizes a highly accessible interface that allows users to extend the library with ease and quickly set up and run new simulations.
+- **Leverages JAX Array and Shardmap:** The library incorporates the new JAX array unified array type and JAX shardmap, providing users with a numpy-like interface. This allows users to focus solely on the semantics, leaving performance optimizations to the compiler.
- **Platform Versatility:** The same XLB code can be executed on a variety of platforms including multi-core CPUs, single or multi-GPU systems, TPUs, and it also supports distributed runs on multi-GPU systems or TPU Pod slices.
+- **Visualization:** XLB provides a variety of visualization options including in-situ rendering using [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze).
-## Documentation
-The documentation can be found [here](https://autodesk.github.io/XLB/) (in preparation)
## Showcase
-The following examples showcase the capabilities of XLB:
-
+
- Lid-driven Cavity flow at Re=100,000 (~25 million voxels) + On GPU in-situ rendering using PhantomGaze library (no I/O). Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with ~10 million cells.
+
- 
+ 
- Airflow in to, out of, and within a building (~400 million voxels) + DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million cells
+
+
+
+ Airflow in to, out of, and within a building (~400 million cells) +
-
+ 
- DrivAer model in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million voxels +The stages of a fluid density field from an initial state to the emergence of the "XLB" pattern through deep learning optimization at timestep 200 (see paper for details)
+
- 
+
- Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with approx. 100 million voxels + Lid-driven Cavity flow at Re=100,000 (~25 million cells)
## Capabilities ### LBM + - BGK collision model (Standard LBM collision model) - KBC collision model (unconditionally stable for flows with high Reynolds number) +### Machine Learning + +- Easy integration with JAX's ecosystem of machine learning libraries +- Differentiable LBM kernels +- Differentiable boundary conditions + ### Lattice Models + - D2Q9 - D3Q19 - D3Q27 (Must be used for KBC simulation runs) +### Compute Capabilities +- Distributed Multi-GPU support +- Mixed-Precision support (store vs compute) +- Out-of-core support (coming soon) + ### Output + - Binary and ASCII VTK output (based on PyVista library) +- In-situ rendering using [PhantomGaze](https://github.com/loliverhennigh/PhantomGaze) library +- [Orbax](https://github.com/google/orbax)-based distributed asynchronous checkpointing - Image Output - 3D mesh voxelizer using trimesh ### Boundary conditions + - **Equilibrium BC:** In this boundary condition, the fluid populations are assumed to be in at equilibrium. Can be used to set prescribed velocity or pressure. - **Full-Way Bounceback BC:** In this boundary condition, the velocity of the fluid populations is reflected back to the fluid side of the boundary, resulting in zero fluid velocity at the boundary. @@ -80,11 +122,7 @@ The following examples showcase the capabilities of XLB: - **Regularized BC:** This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary. This BC is more stable than Zouhe BC, but computationally more expensive. - **Extrapolation Outflow BC:** A type of outflow boundary condition that uses extrapolation to avoid strong wave reflections. -### Compute Capabilities -- Distributed Multi-GPU support -- JAX shard-map and JAX Array support -- Mixed-Precision support (store vs compute) -- [Orbax](https://github.com/google/orbax)-based distributed asynchronous checkpointing +- **Interpolated Bounceback BC:** Interpolated bounce-back boundary condition due to Bouzidi for a lattice Boltzmann method simulation. ## Installation Guide @@ -116,5 +154,3 @@ cd XLB export PYTHONPATH=. python3 examples/cavity2d.py ``` -## Citing XLB -Accompanying paper will be available soon. \ No newline at end of file diff --git a/assets/XLB_diff.png b/assets/XLB_diff.png new file mode 100644 index 0000000..8b53d18 Binary files /dev/null and b/assets/XLB_diff.png differ diff --git a/assets/airfoil.gif b/assets/airfoil.gif new file mode 100644 index 0000000..fd99cb8 Binary files /dev/null and b/assets/airfoil.gif differ