Autodesk · hsalehipour · Oct 22, 2024 · Oct 22, 2024
diff --git a/.github/workflows/mkdocs.yml b/.github/workflows/mkdocs.yml
@@ -2,7 +2,6 @@ name: mkdocs-deployment
 on:
   push:
     branches:
-      - master 
       - main
     # paths:
     #   - 'docs/**'

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 - _No changes yet_  <!-- Placeholder for future changes -->
 
+### Fixed
+-  mkdocs is now configured correctly for the new project structure
+
 ## [0.2.0] - 2024-10-22
 
 ### Added

diff --git a/docs/assets/airfoil.png b/docs/assets/airfoil.png
diff --git a/docs/assets/car.png b/docs/assets/car.png
diff --git a/docs/assets/cavity.gif b/docs/assets/cavity.gif
diff --git a/docs/assets/logo-transparent.png b/docs/assets/logo-transparent.png
diff --git a/docs/assets/logo.svg b/docs/assets/logo.svg
diff --git a/docs/base.md b/docs/base.md
diff --git a/docs/boundary_conditions.md b/docs/boundary_conditions.md
diff --git a/docs/index.md b/docs/index.md
@@ -1,63 +1,139 @@
+[![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![GitHub star chart](https://img.shields.io/github/stars/Autodesk/XLB?style=social)](https://star-history.com/#Autodesk/XLB)
 <p align="center">
-  <img src="assets/logo-transparent.png" alt="" width="700">
+  <img src="https://raw.githubusercontent.com/autodesk/xlb/main/assets/logo-transparent.png" alt="" width="300">
 </p>
 
-# XLB: A Hardware-Accelerated Differentiable Lattice Boltzmann Simulation Framework based on JAX for Physics-based Machine Learning
+# XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based 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.
+🎉 **Exciting News!** 🎉 XLB version 0.2.0 has been released, featuring a complete rewrite of the library and introducing support for the NVIDIA Warp backend! 
+XLB can now be installed via pip: `pip install xlb`.
+
+XLB is a fully differentiable 2D/3D Lattice Boltzmann Method (LBM) library that leverages hardware acceleration. It supports [JAX](https://github.com/google/jax) and [NVIDIA Warp](https://github.com/NVIDIA/warp) backends, 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. With the new Warp backend, XLB now offers state-of-the-art performance for even faster simulations.
+
+## Getting Started
+To get started with XLB, you can install it using pip:
+```bash
+pip install xlb
+```
+
+To install the latest development version from source:
+
+```bash
+pip install git+https://github.com/Autodesk/XLB.git
+```
+
+The changelog for the releases can be found [here](https://github.com/Autodesk/XLB/blob/main/CHANGELOG.md).
+
+For examples to get you started please refer to the [examples](https://github.com/Autodesk/XLB/tree/main/examples) folder.
+
+## Accompanying Paper
+
+Please refer to the [accompanying paper](https://doi.org/10.1016/j.cpc.2024.109187) for benchmarks, validation, and more details about the library.
+
+## Citing XLB
+
+If you use XLB in your research, please cite the following paper:
+
+```
+@article{ataei2024xlb,
+  title={{XLB}: A differentiable massively parallel lattice {Boltzmann} library in {Python}},
+  author={Ataei, Mohammadmehdi and Salehipour, Hesam},
+  journal={Computer Physics Communications},
+  volume={300},
+  pages={109187},
+  year={2024},
+  publisher={Elsevier}
+}
+```
 
 ## 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 with billions of voxels.
+- **Multiple Backend Support:** XLB now includes support for multiple backends including JAX and NVIDIA Warp, providing *state-of-the-art* performance for lattice Boltzmann simulations. Currently, only single GPU is supported for the Warp backend.
+- **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 using the JAX backend, 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 on GPU 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="https://raw.githubusercontent.com/autodesk/xlb/main/assets/airfoil.gif" width="800">
+</p>
+<p align="center">
+  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="https://raw.githubusercontent.com/autodesk/xlb/main/assets/car.png" alt="" width="500">
+</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 cells
+</p>
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/autodesk/xlb/main/assets/building.png" alt="" width="700">
 </p>
 <p align="center">
-  Lid-driven Cavity flow at Re=100,000 (~25 million voxels)
+  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="https://raw.githubusercontent.com/autodesk/xlb/main/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="https://raw.githubusercontent.com/autodesk/xlb/main/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
+- Single GPU support for the Warp backend with state-of-the-art performance
+- Distributed Multi-GPU support using the JAX backend
+- 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.
@@ -70,43 +146,46 @@ 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)
+- **Interpolated Bounceback BC:** Interpolated bounce-back boundary condition for representing curved boundaries.
 
-## Installation Guide
+## Roadmap
 
-To use XLB, you must first install JAX and other dependencies using the following commands:
+### Work in Progress (WIP)
+*Note: Some of the work-in-progress features can be found in the branches of the XLB repository. For contributions to these features, please reach out.*
 
-```bash
-# Please refer to https://github.com/google/jax for the latest installation documentation
+ - 🌐 **Grid Refinement:** Implementing adaptive mesh refinement techniques for enhanced simulation accuracy.
 
-pip install --upgrade pip
+ - 💾 **Out-of-Core Computations:** Enabling simulations that exceed available GPU memory, suitable for CPU+GPU coherent memory models such as NVIDIA's Grace Superchips (coming soon).
 
-# For CPU run
-pip install --upgrade "jax[cpu]"
 
-# For GPU run
+- ⚡ **Multi-GPU Acceleration using [Neon](https://github.com/Autodesk/Neon) + Warp:** Using Neon's data structure for improved scaling.
 
-# CUDA 12 and cuDNN 8.8 or newer.
-pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+- 🗜️ **GPU Accelerated Lossless Compression and Decompression**: Implementing high-performance lossless compression and decompression techniques for larger-scale simulations and improved performance.
 
-# CUDA 11 and cuDNN 8.6 or newer.
-pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+- 🌡️ **Fluid-Thermal Simulation Capabilities:** Incorporating heat transfer and thermal effects into fluid simulations.
 
-# Run dependencies
-pip install jmp pyvista numpy matplotlib Rtree trimesh jmp
-```
+- 🎯 **Adjoint-based Shape and Topology Optimization:** Implementing gradient-based optimization techniques for design optimization.
 
-Run an example:
-```bash
-git clone https://github.com/Autodesk/XLB
-cd XLB
-export PYTHONPATH=.
-python3 examples/cavity2d.py
-```
-## Citing XLB
-Accompanying publication coming soon:
+- 🧠 **Machine Learning Accelerated Simulations:** Leveraging machine learning to speed up simulations and improve accuracy.
+
+- 📉 **Reduced Order Modeling using Machine Learning:** Developing data-driven reduced-order models for efficient and accurate simulations.
+
+
+### Wishlist
+*Contributions to these features are welcome. Please submit PRs for the Wishlist items.*
+
+- 🌊 **Free Surface Flows:** Simulating flows with free surfaces, such as water waves and droplets.
+
+- 📡 **Electromagnetic Wave Propagation:** Simulating the propagation of electromagnetic waves.
+
+- 🛩️ **Supersonic Flows:** Simulating supersonic flows.
+
+- 🌊🧱 **Fluid-Solid Interaction:** Modeling the interaction between fluids and solid objects.
+
+- 🧩 **Multiphase Flow Simulation:** Simulating flows with multiple immiscible fluids.
+
+- 🔥 **Combustion:** Simulating combustion processes and reactive flows.
+
+- 🪨 **Particle Flows and Discrete Element Method:** Incorporating particle-based methods for granular and particulate flows.
 
-**M. Ataei, H. Salehipour**. XLB: Hardware-Accelerated, Scalable, and Differentiable Lattice Boltzmann Simulation Framework based on JAX. TBA
+- 🔧 **Better Geometry Processing Pipelines:** Improving the handling and preprocessing of complex geometries for simulations.
diff --git a/docs/lattice.md b/docs/lattice.md
diff --git a/docs/models.md b/docs/models.md
diff --git a/docs/utils.md b/docs/utils.md
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -30,7 +30,7 @@ repo_url: https://github.com/Autodesk/XLB
 edit_uri: edit/master/docs/
 
 watch:
-- ./src/
+- ./xlb/
 
 # Can be used to add a meta tag to the HTML header
 site_description: 'Documentation for project XLB'
@@ -147,10 +147,4 @@ extra_javascript:
   - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js
 
 nav:
-  - XLB's home: index.md
-  - XLB API:
-    - XLB base: base.md
-    - XLB models: models.md
-    - XLB lattice: lattice.md
-    - XLB utils: utils.md
-    - XLB boundary conditions: boundary_conditions.md
+  - XLB's home: index.md
-Original file line number
+Diff line change
@@ Expand Up / @@ -2,7 +2,6 @@ name: mkdocs-deployment @@
     on:
       push:
         branches:
-          - master
           - main
         # paths:
         #   - 'docs/**'
@@ Expand Down @@