Autodesk/XLB

XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML

XLB – 可微分格子波茲曼函式庫 (Python)

簡介 – XLB 是一個實作格子波茲曼方法 (LBM) 進行流體動力學模擬的 Python 函式庫,並使每一個步驟皆可微分。它透過三個後端在 CPU、單 GPU、多 GPU 和 TPU 硬體上執行:

  • JAX – 利用 JAX 的 XLA 編譯器,適用於 CPU、GPU 和 TPU。
  • NVIDIA Warp – 用於頂尖單 GPU 效能的自訂 CUDA 核心框架。
  • Neon – Autodesk 的多 GPU 程式設計模型,增加了網格細化 (多解析度) 與多 GPU 擴展能力。

由於核心是可微分的,XLB 可以嵌入機器學習管線中 (例如:基於梯度的設計最佳化、物理資訊神經網路,或基於伴隨的形狀最佳化)。


快速入門 (來自 README)

# 僅 CPU (無特殊硬體)
pip install xlb

# 使用 Warp 的單 GPU (最快的單 GPU 路徑)
pip install "xlb[warp]"

# 在支援 CUDA 的 GPU 上使用 JAX
pip install "xlb[cuda]"

# 在 TPU 上使用 JAX
pip install "xlb[tpu]"

# 使用 Neon 進行多 GPU / 多解析度 (需要 Linux 與 Python 3.11‑3.14)
git clone https://github.com/Autodesk/XLB.git
cd XLB
pip install -r requirements.txt
pip install '.[neon]'

請勿在同一環境中混合使用 Warp 與 Neon。


核心功能 (列於 README)

類別 功能
後端 JAX, NVIDIA Warp, Neon (多 GPU, 網格細化)
LBM 模型 BGK, KBC (在高雷諾數下穩定), Smagorinsky LES
格子集 D2Q9, D3Q19, D3Q27 (KBC 所需)
邊界條件 平衡, 全/半反彈, Zou/He, 正規化, 外推流出, 插值反彈, 混合等
可微分性 所有碰撞與串流核心皆與 JAX 相容,支援針對初始場、幾何形狀或材料參數的梯度計算
擴展性 單 GPU (Warp), 多 GPU (Neon), 透過 JAX shard‑map 進行分散式多 GPU/TPU
精度 混合精度 (儲存與計算分離)
I/O 與視覺化 透過 PyVista 進行 VTK (二進位/ASCII), 帶 gzip 的 HDF5/XDMF, 使用 PhantomGaze 進行原位 GPU 渲染, 使用 Orbax 進行檢查點儲存, 影像切片, 透過 trimesh 進行體素化
機器學習整合 與 Flax, Haiku, Optax 等協作;提供可微分邊界條件
路線圖亮點 核心外模擬, 無損 GPU 壓縮, 流體‑熱耦合, 伴隨形狀最佳化, ML 加速求解器, 降階模型

典型使用案例

  • 基於物理的 ML 研究 – 將流體模擬器嵌入神經網路損失函數中,並透過模擬進行反向傳播,以學習控制策略或反向設計。
  • 基於梯度的設計最佳化 – 透過對幾何形狀的阻力/升力目標進行微分,來最佳化形狀 (翼型、車體、建築佈局)。
  • 高效能 CFD – 在單一 GPU (Warp) 或多個 GPU (Neon/JAX) 上執行大規模 LBM 模擬 (數億個網格)。
  • 多解析度研究 – 使用 Neon 的網格細化功能,在需要的地方 (例如牆壁附近) 集中解析度,同時保持整體記憶體佔用在可控範圍內。
  • 教育 / 研究原型設計 – 純 Python API 搭配 NumPy 風格語法,無需編寫 CUDA 核心即可輕鬆進行實驗。

深入了解

  • 範例 – 儲存庫的 examples/ 資料夾包含風力渦輪機、NACA 翼型、DrivAer 車型、建築氣流、頂蓋驅動空腔及可微分最佳化演示的即用型腳本。
  • 論文XLB: A differentiable massively parallel lattice Boltzmann library in Python (Computer Physics Communications, 2024) 提供了基準測試數據與驗證細節。
  • 引用 – 若您發表使用 XLB 的研究,請引用主要論文;若使用網格細化功能,請同時引用 README 中列出的 Neon 相關會議論文。

授權

Apache 2.0 – 學術與商業用途皆免費。


總結 – XLB 是一個真正的生產級函式庫,架起了高效能流體模擬與現代可微分程式設計之間的橋樑。它非常適合需要快速、可擴展 CFD,且需針對機器學習驅動的設計或最佳化任務進行微分的研究人員與工程師。

相關

  • 專案
  • 專案
  • 專案
  • 專案