In our effort to facilitate machine learning-assisted computational fluid
dynamics (CFD), we introduce the second iteration of JAX-Fluids. JAX-Fluids is
a Python-based fully-differentiable CFD solver designed for compressible
single- and two-phase flows. In this work, the first version is extended to
incorporate high-performance computing (HPC) capabilities. We introduce a
parallelization strategy utilizing JAX primitive operations that scales
efficiently on GPU (up to 512 NVIDIA A100 graphics cards) and TPU (up to 1024
TPU v3 cores) HPC systems. We further demonstrate the stable parallel
computation of automatic differentiation gradients across extended integration
trajectories. The new code version offers enhanced two-phase flow modeling
capabilities. In particular, a five-equation diffuse-interface model is
incorporated which complements the level-set sharp-interface model. Additional
algorithmic improvements include positivity-preserving limiters for increased
robustness, support for stretched Cartesian meshes, refactored I/O handling,
comprehensive post-processing routines, and an updated list of state-of-the-art
high-order numerical discretization schemes. We verify newly added numerical
models by showcasing simulation results for single- and two-phase flows,
including turbulent boundary layer and channel flows, air-helium shock bubble
interactions, and air-water shock drop interactions.

我们介绍了第二个版本的基于 Python 的完全可微分计算流体力学求解器 JAX-Fluids，该求解器设计用于可压缩的单相和两相流。我们通过引入能够在 GPU 和 TPU 上高效扩展的 JAX 原始操作并展示稳定的自动微分梯度计算， 实现了高性能计算能力的并行化策略。新版本还提供了增强的两相流模拟能力、正性保持限制器以增加鲁棒性、支持拉伸笛卡尔网格、重构的输入 / 输出处理、全面的后处理例程，以及最新的高阶数值离散方案。我们通过模拟单相和两相流的结果验证了新增的数值模型，包括湍流边界层和通道流、气 - 氦冲击泡交互和气 - 水冲击滴交互。