JAX-Fluids 是一款完全可微分的计算流体力学（CFD）求解器，基于 JAX 编写，支持 CPU/GPU/TPU 加速与自动微分。它的子模块jaxfluids_nn专门负责"神经网络 + CFD 求解器"的耦合训练，让你可以端到端地用梯度优化流场代理模型、反演物理参数，是 ML 与 CFD 交叉研究的理想起点 🚀。

一、为什么需要 jaxfluids_nn？

传统 CFD 算例跑一次可能要数小时，而"参数 → 流场"的反复试算成本极高。JAX-Fluids 的核心优势在于：整条时间推进过程都可以反向传播梯度。这意味着你可以把神经网络嵌入求解器内部，让"模拟"本身成为损失函数的一部分，从而：

• 🧠 用神经网络学习未知源项、材料属性（代理模型）
• 🔍 通过可微分前向模拟反演单一物理参数（如粘度）
• ⚡ 用 GPU/TPU 并行训练，规模可达数千核

二、模块速览：jaxfluids_nn 里有什么

src/jaxfluids_nn/目录结构非常精简，四个核心文件各司其职：

Trainer内部用jax.jit加速每个train_step，每checkpoint_freq个 epoch 自动保存checkpoint_{epoch}.pkl（含参数、优化器状态、完整 history），训练曲线自动绘制为history.png——工程化程度接近主流深度学习框架。

三、官方实战三部曲：notebooks/advanced

仓库自带三个由浅入深的 Notebook，覆盖 jaxfluids_nn 的全部典型用法：

1️⃣ 用神经网络学习未知源项

01_nn_training.ipynb（配置：shear_flow.json、numerical_setup.json）

目标：给定目标速度剖面 u(y)，用一个两层 MLP 学习产生它的源项 S_x = NN(y)。流程分四步：

1. 生成训练数据：先用已知源项sin(2πy)跑一次正向模拟（配置文件里写死），把结果存为 HDF5
2. 嵌入网络：通过Callback的after_compute_rhs钩子，把 NN 输出的源项加到右端项 buffer 上
3. 定义 loss：调用SimulationManager._feed_forward做 200 步可微分前向模拟，取末端速度剖面与目标做均方误差
4. 梯度下降：jax.value_and_grad对整条模拟链路反传，Adam 优化 10000 步

关键机制在 ml_buffers.py：ParametersSetup/CallablesSetup把"网络权重"和"网络函数"分别注入求解器，base_callback.py 的after_compute_rhs负责在每步积分时读取网络参数。

2️⃣ 可微分反演单一参数：Poiseuille 粘度优化

02_poiseuille_optimization.ipynb 演示最经典的逆问题：先以参考粘度 μ=1.0 跑一遍生成目标剖面，再从 μ=0.5 出发，用 optax 的 Adam 优化器反传 125 步，让模拟出的速度剖面与目标贴合。这里不需要神经网络——被优化的就是标量粘度本身，通过DiffusiveFluxesSetup(dynamic_viscosity=...)传入求解器。配套脚本 run_optimization.py 可直接命令行复现。

3️⃣ 神经网络预测温度相关热导率 λ(T)

03_heat_equation_nn.ipynb 是最完整的"代理模型"案例：参考解用高斯型热导率生成，一个 3 层 MLP 学习 λ_θ(T)，输出层用1.0 + softplus(x)保证热导率恒大于 1（物理约束）。热导率通过CallablesSetup注入DiffusiveFluxesSetup.thermal_conductivity，优化器使用optax.exponential_decay阶梯学习率，300 步即可收敛。脚本版见 run_optimization.py。

四、动手三步走：从安装到第一次训练

第 1 步：安装

git clone https://gitcode.com/gh_mirrors/ja/JAXFLUIDS cd JAXFLUIDS pip install --upgrade "jax[cpu]" # GPU 用户请安装 jax[cuda12] pip install -e .

第 2 步：跑通正向模拟。先打开 notebooks/basics/ 里的 1D Sod 激波管教程，熟悉InputManager+SimulationManager的两文件配置范式（case_setup + numerical_setup）。

第 3 步：进入 advanced 目录做 NN 训练。三个 Notebook 都自带 json 配置文件，直接在各自的目录里按顺序运行即可；02 和 03 的 Notebook 也支持从仓库根目录运行（会自动回退路径）。

五、实用小贴士 💡

• 开启双精度：涉及优化的算例建议jax.config.update("jax_enable_x64", True)，01 号 Notebook 开头就有这一行
• jit 是关键：step_fn用jax.jit编译后首次运行会有编译开销，之后每步毫秒级（02 号算例日志里记录了 COMPILE TIME 与 Time/Batch）
• FeedForwardSetup 参数：outer_steps（总步数）与inner_steps（scan 内分块大小）控制可微分前向的长度，见 feed_forward/data_types.py
• 学习率调度：OptimizerSetup支持 Constant / Polynomial / Exponential / Piecewise_constant 四种调度，见 helper_functions.py 的initialize_optimizer

六、总结

jaxfluids_nn 把"CFD 求解器"变成了神经网络训练栈的一等公民：Callback 注入源项、ML Buffers 注入网络参数、Trainer 管理训练循环。配合 notebooks/advanced 下的三个完整案例，你可以在一小时内走通"数据生成 → 可微分前向 → 梯度优化"的完整闭环，把 JAX-Fluids 真正变成你的流体智能工具箱 🛠️。