latticebmjax

About

'latticebmjax' is a Python package designed for hardware-accelerated simulation of fluid flows using the Lattice Boltzmann Method (LBM) with JAX as the computational backend. It is capable of:

• Solving 2D incompressible flows using the D2Q9 BGK-LBM scheme
• Applying Zou-He and bounce-back boundary conditions
• Running on CPU with JIT compilation via JAX
• Validating results against the Ghia et al. (1982) benchmark data
• Plotting velocity fields, streamlines and centreline profiles
• Benchmarking JIT vs non-JIT performance across multiple solvers

Features

• D2Q9 lattice with BGK single-relaxation-time collision operator
• Zou-He velocity boundary condition for the moving lid
• Bounce-back no-slip boundary conditions for stationary walls
• JIT-compiled time loop via jax.jit for high performance
• Lid-driven cavity solver validated against Ghia et al. (1982)
• Poiseuille solver validated against the analytical parabolic profile
• Centreline velocity profiles compared to benchmark data (Re=100, 400, 1000)

Examples

• 'jit_benchmark.py' — JIT vs no-JIT performance comparison across all solvers
• 'lid_cavity.py' — lid-driven cavity with live plotting and Ghia comparison
• 'poiseuille.py' — Poiseuille channel flow with analytical validation

Installation

Dependencies

Ensure you have the required dependencies installed, which are listed in requirements.txt. You can install them using:

pip install -r requirements.txt

User installation

Clone the repository and install locally:

git clone https://github.com/veracfernanda/latticebmjax.git
cd latticebmjax
pip install -e .

To install the library, you can use the 'pip' command:

pip install latticebmjax

Folder structure

latticebmjax
├─── README.md
├─── LICENSE
├─── pyproject.toml
├─── requirements.txt
├─── benchmarks
│   ├─── convergence.py
│   └─── ghia_data.py
├─── examples
│   ├─── poiseuille.py
│   ├─── jit_benchmark.py
│   └─── lid_cavity.py
├─── latticebmjax
│   ├─── __init__.py
│   ├─── collision.py
│   ├─── equilibrium.py
│   ├─── lattice.py
│   └─── streaming.py
├─── solvers
│   ├─── __init__.py
│   ├─── base_solver.py
│   ├─── boundary.py
│   ├─── lid_cavity.py
│   └─── poiseuille.py
├─── testcase_data
│   ├─── ghia_re100.csv
│   ├─── ghia_re400.csv
│   └─── ghia_re1000.csv
└─── tests
    ├─── test_collision.py
    ├─── test_lid_cavity.py
    └─── test_streaming.py

Examples

Some usage examples of the library can be found in the examples folder

License

This software is distributed under the MIT License and further information about the license can be found in the LICENSE file

Third-party libraries

The library currently uses the following third-party libraries:

• jax — JIT compilation and GPU acceleration
• numpy — array utilities and CSV loading
• matplotlib — velocity field and centreline plots
• pandas — benchmark data loading
• pytest — unit and integration testing

Contributing

The project is in the early development phase. Upon completion and delivery, contributions will be welcome! Please fork the repository and submit a pull request for any improvements or bug fixes.