first commit

.gitignore

@@ -0,0 +1,143 @@
+# Editors
+.vscode/
+.idea/
+
+# Outputs
+*vti
+*vtr
+*vtk
+*jpg
+*zraw
+*mhd
+*png
+# Exclusions
+!jax-lbm.png
+!airfoil.png
+!car.png
+# Vagrant
+.vagrant/
+
+# Meshes
+*.obj
+*.stl
+*.ply
+*.off
+
+# Mac/OSX
+.DS_Store
+
+# Windows
+Thumbs.db
+
+# Source for the following rules: https://raw.githubusercontent.com/github/gitignore/master/Python.gitignore
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json

AUTHORS

@@ -0,0 +1,7 @@
+# This is the list of XLB's significant contributors.
+# This compilation acknowledges the principal contributors to XLB but does not include every individual who has contributed code. 
+# Due to the collective efforts of numerous employees from various corporations, some contributors may not be listed here. 
+# For a comprehensive view of all contributors, please refer to the revision history in the source control.
+
+Mehdi Ataei (Autodesk Inc)
+Hesam Saleipour (Autodesk Inc)

CODE_OF_CONDUCT.md

@@ -0,0 +1,23 @@
+# Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our community a safe, inclusive and harassment-free experience for everyone. We pledge to act and interact in ways that contribute to an open, welcoming, diverse, inclusive and healthy community.
+
+## Our Standards
+
+Our Open Source Community works to:
+
++ Be kind towards other people which enables us to be empathic to each other
++ Be respectful of differing opinions, viewpoints, and experiences
++ Give and gracefully accept constructive feedback
++ Accept responsibility and apologize to those affected by our mistakes, and learning from the experience
++ Focus on what is best not just for us as individuals, but for the overall community
+
+We will not tolerate the following behaviors:
++ Violent threats or language
++ The use of sexualized language or imagery, and sexual attention or advances of any kind
++ Trolling, insulting or derogatory comments, and personal or political attacks
++ Public or private harassment
++ Publishing others’ private information, such as a physical or email address, without their explicit permission
++ Other conduct which could reasonably be considered inappropriate in a professional setting

CONTRIBUTING.md

@@ -0,0 +1,40 @@
+# XLB's Community and How to Contribute
+
+The XLB team is deeply committed to the ethos of open-source. We are always open to, and greatly value, contributions from our community, which can take the form of suggesting new features, reporting issues, and contributing code. This document will guide you through the various processes you can follow to contribute to our project.
+
+## Providing Suggestions
+
+We believe that XLB should continually evolve in response to community feedback. As such, we highly value your suggestions on how to enhance the design or functionality of our platform. Please use the enhancement tag when submitting issues that are specifically suggestions, as this will help us categorize and respond appropriately.
+
+## Filing Bugs
+
+Despite our best efforts, like any software, XLB may occasionally have bugs. If you encounter any, please report them as regular issues on our GitHub page. We are continuously monitoring these issues, and we will prioritize and schedule fixes accordingly.
+
+The most effective bug reports provide a detailed method for reliably reproducing the issue and, if possible, a working example demonstrating the problem.
+
+## Contributing Code
+
+Contributing your code to our project involves three main steps: signing a Contributor License Agreement, discussing your goals with the community, adhering to XLB's coding standards when writing your code, and finally, submitting a pull request.
+
+
+### Contributor License Agreement (CLA)
+
+Before you can contribute any code to this project, we kindly request you to sign a Contributor License Agreement (CLA). We are unable to accept any pull request without a signed CLA.
+
+- If you are contributing as an individual, the process of signing the CLA is integrated into the pull request procedure.
+
+- If you are contributing on behalf of your employer, please sign our [**Corporate Contributor License Agreement**](https://github.com/Autodesk/autodesk.github.io/releases/download/1.0/ADSK.Form.Corp.Contrib.Agmt.for.Open.Source.docx). The document includes instructions on where to send the completed forms to. Once a signed form has been received, we can happily review and accept your pull requests.
+
+### Coordinate With the Community
+
+We strongly advise that you initiate your contribution process by opening an issue on GitHub to outline your objectives prior to beginning any coding. This proactive approach facilitates early feedback from the community and helps prevent potential overlaps in contributions.
+
+### Git Workflow
+
+We follow the [GitFlow](http://nvie.com/posts/a-successful-git-branching-model/) development model. 
+If you would like to contribute your code to XLB, you should:
+- Include your work in a feature branch created from the XLB `main` branch. The `main` branch contains the latest work in XLB. 
+- Then, create a pull request against the `main` branch.
+
+
+When you submit your code, please include relevant tests as part of the pull request, and ensure that your comments and coding style align with the rest of the project. You can refer to the existing code for examples of the testing and style practices that the project follows.

LICENSE

@@ -0,0 +1,13 @@
+Copyright 2023 Autodesk Inc
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

README.md

@@ -0,0 +1,97 @@
+<p align="center">
+  <img src="assets/logo-transparent.png" alt="" width="700">
+</p>
+
+# XLB: Hardware-Accelerated, Scalable, and Differentiable Lattice Boltzmann Simulation Framework based on JAX
+
+XLB (Accelerated LB) is a hardware-accelerated fully differentiable 2D/3D Lattice Boltzmann Method solver based on the JAX library. It is designed to solve fluid dynamics problems in a computationally efficient and differentiable manner.
+
+## Documentation
+Coming soon
+## Showcase
+
+The following examples showcase the capabilities of XLB:
+
+<p align="center">
+  <img src="assets/cavity.gif" alt="" width="500">
+</p>
+<p align="center">
+  Lid-driven Cavity flow at Re=100,000 (~25 million voxels)
+</p>
+
+<p align="center">
+  <img src="assets/car.png" alt="" width="500">
+</p>
+<p align="center">
+  Q-criterion over a DrivAer model
+</p>
+
+<p align="center">
+  <img src="assets/airfoil.png" width="500">
+</p>
+<p align="center">
+  Q-criterion over a NACA airfoil
+</p>
+
+## Capabilities 
+
+### LBM
+- BGK collision model (Standard LBM collision model)
+- KBC collision model (unconditionally stable for flows with high Reynolds number)
+
+### Lattice Models
+- D2Q9
+- D3Q19
+- D3Q27 (Must be used for KBC simulation runs)
+
+### Output
+- Binary VTK output
+- ASCII VTK output
+- Image Output (2D and 3D slice)
+- 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.
+
+- Half-Way Bounceback BC: Similar to the Full-Way Bounceback BC, in this boundary condition, the velocity of the fluid populations is partially reflected back to the fluid side of the boundary, resulting in a non-zero fluid velocity at the boundary.
+
+- Do Nothing BC: In this boundary condition, the fluid populations are allowed to pass through the boundary without any reflection or modification.
+
+- Zouhe BC: This boundary condition is used to impose a prescribed velocity or pressure profile at the boundary.
+- 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)
+
+## Installation Guide
+
+To install XLB, you can run the following commands:
+
+```bash
+pip install --upgrade pip
+
+# For CPU run
+pip install --upgrade "jax[cpu]"
+
+# For GPU run
+
+# 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
+
+# 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
+
+# Please refer to https://github.com/google/jax for the latest installation documentation
+
+# Run dependencies
+pip install jmp pyvista numpy matplotlib Rtree trimesh jmp
+```
+## Citing XLB
+Accompanying publication coming soon:
+
+**M. Ataei, H. Salehipour**. XLB: Hardware-Accelerated, Scalable, and Differentiable Lattice Boltzmann Simulation Framework based on JAX. TBA

examples/MLUPS2d.py

@@ -0,0 +1,59 @@
+import os
+
+from src.models import BGKSim
+from src.lattice import LatticeD2Q9
+
+#os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
+import jax.numpy as jnp
+import numpy as np
+from src.utils import *
+from jax.config import config
+from time import time
+import argparse
+from src.boundary_conditions import *
+
+#config.update('jax_disable_jit', True)
+# Use 8 CPU devices
+
+precision = 'f32/f32'
+
+
+class Cavity(BGKSim):
+
+    def set_boundary_conditions(self):
+        # concatenate the indices of the left, right, and bottom walls
+        walls = np.concatenate((self.boundingBoxIndices['left'], self.boundingBoxIndices['right'], self.boundingBoxIndices['bottom']))
+        # apply bounce back boundary condition to the walls
+        self.BCs.append(BounceBack(tuple(walls.T), self.grid_info, self.precision_policy))
+
+        # apply inlet equilibrium boundary condition to the top wall
+        moving_wall = self.boundingBoxIndices['top']
+
+        rho_wall = np.ones(moving_wall.shape[0], dtype=self.precision_policy.compute_dtype)
+        vel_wall = np.zeros(moving_wall.shape, dtype=self.precision_policy.compute_dtype)
+        vel_wall[:, 0] = u_wall
+        self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.grid_info, self.precision_policy, rho_wall, vel_wall))
+
+
+if __name__ == '__main__':
+
+    lattice = LatticeD2Q9(precision)
+
+    parser = argparse.ArgumentParser("simple_example")
+    parser.add_argument("N", help="The total number of voxels will be NxN", type=int)
+    parser.add_argument("timestep", help="Number of timesteps", type=int)
+    args = parser.parse_args()
+
+    N = args.N
+    max_iter = args.timestep
+    Re = 100.0
+    u_wall = 0.1
+    clength = N - 1
+
+    visc = u_wall * clength / Re
+    omega = 1.0 / (3. * visc + 0.5)
+    print('omega = ', omega)
+    assert omega < 2.0, "omega must be less than 2.0"
+    os.system('rm -rf ./*.vtk && rm -rf ./*.png')
+    sim = Cavity(lattice, omega, N, N, precision=precision)
+    sim.run(max_iter, MLUPS=True)