Merge pull request #26 from mehdiataei/main

Updated README after paper submission

.gitignore

@@ -6,6 +6,7 @@
 *vti
 *vtr
 *vtk
+*mp4
 *jpg
 *zraw
 *mhd
@@ -17,6 +18,7 @@
 !car.png
 !logo-transparent.png
 !building.png
+!XLB_diff.png
 # Vagrant
 .vagrant/
 

README.md

@@ -4,70 +4,112 @@
   <img src="assets/logo-transparent.png" alt="" width="700">
 </p>
 
-# 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:
 
 <p align="center">
-  <img src="assets/cavity.gif" alt="" width="500">
+  <img src="assets/airfoil.gif" width="800">
 </p>
 <p align="center">
-  Lid-driven Cavity flow at Re=100,000 (~25 million voxels)
+  On GPU in-situ rendering using <a href="https://github.com/loliverhennigh/PhantomGaze">PhantomGaze</a> library (no I/O). Flow over a NACA airfoil using KBC Lattice Boltzmann Simulation with ~10 million cells.
 </p>
 
+
 <p align="center">
-  <img src="assets/building.png" alt="" width="700">
+  <img src="assets/car.png" alt="" width="500">
 </p>
 <p align="center">
-  Airflow in to, out of, and within a building (~400 million voxels)
+<a href=https://www.epc.ed.tum.de/en/aer/research-groups/automotive/drivaer > DrivAer model </a> in a wind-tunnel using KBC Lattice Boltzmann Simulation with approx. 317 million cells
 </p>
 
+<p align="center">
+  <img src="assets/building.png" alt="" width="700">
+</p>
+<p align="center">
+  Airflow in to, out of, and within a building (~400 million cells)
+</p>
 
 <p align="center">
-  <img src="assets/car.png" alt="" width="500">
+  <img src="assets/XLB_diff.png" alt="" width="900">
 </p>
 <p align="center">
-<a href=https://www.epc.ed.tum.de/en/aer/research-groups/automotive/drivaer > DrivAer model </a> 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)
 </p>
 
+<br>
+
 <p align="center">
-  <img src="assets/airfoil.png" width="500">
+  <img src="assets/cavity.gif" alt="" width="500">
 </p>
 <p align="center">
-  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)
 </p>
 
 ## 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.