google/jax-cfd

JAX‑CFD: JAX 中的可微分計算流體力學

是什麼 – 一個開源研究函式庫，使用 JAX 實作核心 CFD 求解器（有限體積法和偽譜法），因此程式碼可以自動求導，並能在 GPU/TPU 上執行。它還附帶了一些使用 Haiku 建立的機器學習增強模型。

為什麼重要 – 透過利用 JAX 的即時編譯（JIT）和自動微分功能，研究人員可以：

• 訓練與流體模擬互動的類神經網路（例如：代理模型、閉合模型）。
• 直接在 CFD 求解器上進行基於梯度的優化或數據同化。
• 在不更改程式碼的情況下，在 CPU、GPU 或 TPU 上執行相同的程式碼。

核心組件

• jax_cfd.base – 在交錯網格（Arakawa‑C/MAC）上的有限體積/差分求解器。
• jax_cfd.spectral – 具有抗混疊功能的偽譜渦度求解器。
• jax_cfd.ml – ML 增強的 CFD 模型（Haiku 類神經網路）和範例預訓練模型。
• jax_cfd.data – 用於載入、處理和視覺化模擬數據的工具（Xarray + Pillow）。

數值計算重點

• 非定常湍流。
• 空間離散化：交錯網格 FV/Diff 或偽譜渦度。
• 時間離散化：一階顯式平流，隱式/顯式擴散。
• 壓力求解：共軛梯度法或基於 FFT 的對角化（僅限週期性邊界條件）。
• 平流方案：二階 Van Leer。
• 湍流閉合：Smagorinsky 渦粘度模型。

如何開始

1. 安裝核心套件：

   pip install jax-cfd
   

2. 加入可選的額外組件以使用 ML 或數據工具：

   pip install jax-cfd[ml]   # Haiku + ML helpers
   pip install jax-cfd[data] # Xarray, Pillow, etc.
   

3. 執行 notebooks/ 資料夾中的現成 Colab notebook，以查看：
   • 2-D 交錯網格 FV 演示
   • 2-D 偽譜湍流演示
   • 通道流範例
   • 重現 PNAS 2021 論文的結果（數據分析與模型推論）

專案狀態 – 已標記為不再維護（2024 年 6 月）。README 指向了積極維護的替代方案，例如 JAX‑Fluids、PhiFlow 和 Exponax。

典型使用場景

• 可微分 CFD 的學術研究，例如：學習湍流閉合模型或逆向觀察算符。
• 建立需要透過流體求解器進行梯度的 ML‑CFD 混合模型原型。
• 在流體力學工作負載上測試 JAX 的效能。

引用 – 如果你使用有限體積法程式碼或 ML 模型，請引用 PNAS 2021 論文。對於譜法程式碼，請引用 arXiv 2022 pre‑print。

相關專案 – README 列出了其他可微分 CFD 函式庫（PhiFlow, Autograd fluid examples）以及更廣泛的 JAX 科學工具（JAX‑MD, JAX‑DFT, jax‑cosmo, Veros）。

---

以上所有細節均直接取自儲存庫的 README。