基于 JAX 的 GPU 上的可微流体力学及端到端优化
流体动力学
2024-07-01 v1
摘要
本项目旨在推进面向多孔介质的高超声速耦合流动的可微流体力学，展示了基于自动微分（AD）的优化用于端到端解决方案的潜力。高效地利用 AD 处理高维优化问题，提供了传统方法的灵活替代方案。我们利用 JAX-Fluids—a 基于 JAX 框架开发的求解器，结合了 autograd 和 TensorFlow 的 XLA。在配备 NVIDIA A100 GPU 的 HAWK-AI 节点上编译的 JAX-Fluids 显示出与 FLEXI 等高阶代码相当的计算性能。通过压缩湍流通道流 DNS 案例的验证显示出色，同时成功测试了用于模型化多孔介质的新边界条件。我们的研究的后续步骤仍在计划中。
关键词
引用
@article{arxiv.2406.19494,
title = {JAX-based differentiable fluid dynamics on GPU and end-to-end optimization},
author = {Wenkang Wang and Xuanwei Zhang and Deniz Bezgin and Aaron Buhendwa and Xu Chu and Bernhard Weigand},
journal= {arXiv preprint arXiv:2406.19494},
year = {2024}
}