$ show Autodesk/XLB
XLB is a differentiable, massively parallel Lattice Boltzmann Method (LBM) library written in Python, designed for physics-based machine learning and computational fluid dynamics (CFD). It provides a modular operator-based architecture with abstract base classes and backend-specific implementations for JAX (CPU/GPU/TPU), NVIDIA Warp (single-GPU CUDA), and Autodesk Neon (multi-GPU, multi-resolution). The library supports 2D and 3D simulations with multiple velocity sets (D2Q9, D3Q19, D3Q27), collision models (BGK, KBC, Smagorinsky LES), boundary conditions (equilibrium, bounce-back, regularized, Zou-He, hybrid, extrapolation outflow), and immersed boundary method (IBM) for fluid-structure interaction. A key differentiator is its native differentiability via JAX and Warp autodiff, enabling gradient-based optimization of initial conditions, geometry, or control parameters. The codebase also includes experimental out-of-core support for simulations exceeding GPU memory, multi-resolution grid refinement, and distributed computation across multiple devices. While primarily a software library for CFD practitioners and ML researchers, it is accompanied by research validation data and academic publications, giving it a hybrid application/research character.
The LLM identified 6 files worth highlighting for reuse:
Defines the abstract `Operator` base class with a `@register_backend` decorator that allows backend-specific implementations (JAX, Warp, Neon) to be registered and dispatched at runtime. This pattern cleanly separates physics logic from compute backend, making it easy to add new backends or operators without modifying existing code. The decorator-based registration is a non-obvious design choice that avoids large if-else chains and enables compile-time specialization.
Provides the `Grid` abstract class and factory functions (`grid_factory`, `multires_grid_factory`) that create backend-specific grid instances (WarpGrid, JaxGrid, NeonGrid, NeonMultiresGrid). The grid abstraction encapsulates field creation, indexing, and memory layout, allowing operators to work uniformly across backends. The factory pattern centralizes grid creation logic and supports multi-resolution grids with nested cuboid structures.
Abstract base class for velocity sets (D2Q9, D3Q19, D3Q27) that provides lattice vectors, weights, opposite indices, and backend-specific constant arrays. The design uses class methods and cached properties to compute constants once and reuse them across all operators. This is a clean, reusable pattern for any LBM implementation, and the backend-specific constant generation (e.g., Warp `wp.constant`) optimizes GPU kernel performance.
Defines `PrecisionPolicy` and `Precision` enums that map logical precision modes (e.g., FP32FP32, FP64FP64, FP32FP16) to actual JAX and Warp dtypes. This component centralizes precision configuration and ensures consistent type handling across all operators and grids. It is self-contained and can be extracted for any numerical simulation library that needs mixed-precision support.
Implements a singleton registry that assigns unique integer IDs to boundary condition instances at initialization. This is used to encode boundary condition types into the simulation grid for efficient GPU kernel dispatch. The registry pattern avoids runtime type checks and enables compact storage of BC information in field arrays. It is a non-obvious optimization for LBM boundary handling.
Defines an enum `MeshVoxelizationMethod` with options RAY, AABB, AABB_CLOSE, WINDING, each representing a different algorithm for converting triangle meshes to voxel masks. The enum also includes a `morphological_close` flag for smoothing. This component encapsulates the choice of voxelization strategy and can be reused in any project requiring mesh-to-grid conversion, independent of the LBM context.