JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows

• URL: http://arxiv.org/abs/2402.05193v1
• Date: Wed, 7 Feb 2024 19:05:27 GMT
• Title: JAX-Fluids 2.0: Towards HPC for Differentiable CFD of Compressible Two-phase Flows
• Authors: Deniz A. Bezgin, Aaron B. Buhendwa, Nikolaus A. Adams
• Abstract summary: JAX-Fluids is a Python-based fully-differentiable CFD solver designed for compressible single- and two-phase flows. 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. The new code version offers enhanced two-phase flow modeling capabilities.
• Score: 0.0
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• Abstract: 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.

Related papers

• DualFluidNet: an Attention-based Dual-pipeline Network for FLuid Simulation [4.694954114339147]
  We propose an innovative approach for 3D fluid simulations utilizing an Attention-based Dual-pipeline Network. We find a way to achieve a better balance between global fluid control and physical law constraints. We also propose a new dataset, Tank3D, to further explore the network's ability to handle more complicated scenes.
  arXiv  Detail & Related papers  (2023-12-28T07:37:11Z)
• Vertical Federated Learning over Cloud-RAN: Convergence Analysis and System Optimization [82.12796238714589]
  We propose a novel cloud radio access network (Cloud-RAN) based vertical FL system to enable fast and accurate model aggregation. We characterize the convergence behavior of the vertical FL algorithm considering both uplink and downlink transmissions. We establish a system optimization framework by joint transceiver and fronthaul quantization design, for which successive convex approximation and alternate convex search based system optimization algorithms are developed.
  arXiv  Detail & Related papers  (2023-05-04T09:26:03Z)
• AMT: All-Pairs Multi-Field Transforms for Efficient Frame Interpolation [80.33846577924363]
  We present All-Pairs Multi-Field Transforms (AMT), a new network architecture for video framegithub. It is based on two essential designs. First, we build bidirectional volumes for all pairs of pixels, and use the predicted bilateral flows to retrieve correlations. Second, we derive multiple groups of fine-grained flow fields from one pair of updated coarse flows for performing backward warping on the input frames separately.
  arXiv  Detail & Related papers  (2023-04-19T16:18:47Z)
• Machine Learning model for gas-liquid interface reconstruction in CFD numerical simulations [59.84561168501493]
  The volume of fluid (VoF) method is widely used in multi-phase flow simulations to track and locate the interface between two immiscible fluids. A major bottleneck of the VoF method is the interface reconstruction step due to its high computational cost and low accuracy on unstructured grids. We propose a machine learning enhanced VoF method based on Graph Neural Networks (GNN) to accelerate the interface reconstruction on general unstructured meshes.
  arXiv  Detail & Related papers  (2022-07-12T17:07:46Z)
• Tutel: Adaptive Mixture-of-Experts at Scale [20.036168971435306]
  Sparsely-gated mixture-of-experts (MoE) has been widely adopted to scale deep learning models to trillion-plus parameters with fixed computational cost. We present Flex, a highly scalable stack design and implementation for MoE with dynamically adaptive parallelism and pipelining. Our evaluation shows that Flex efficiently and effectively runs a real-world MoE-based model named SwinV2-MoE, built upon Swin Transformer V2, a state-of-the-art computer vision architecture.
  arXiv  Detail & Related papers  (2022-06-07T15:20:20Z)
• JAX-FLUIDS: A fully-differentiable high-order computational fluid dynamics solver for compressible two-phase flows [0.0]
  We propose JAX-FLUIDS: a comprehensive Python CFD solver for compressible two-phase flows. JAX-FLUIDS allows the simulation of complex fluid dynamics with phenomena like three-dimensional turbulence, compressibility effects, and two-phase flows.
  arXiv  Detail & Related papers  (2022-03-25T16:38:24Z)
• Accelerating Distributed K-FAC with Smart Parallelism of Computing and Communication Tasks [13.552262050816616]
  Kronecker-Factored Approximate Curvature (KFAC) is one of the most efficient approximation algorithms for training deep models. Yet, when leveraging GPU clusters to train models with KFAC, it incurs extensive computation as well as introduces extra communications during each iteration. We propose D-KFAC with smart parallelism of computing and communication tasks to reduce the iteration time.
  arXiv  Detail & Related papers  (2021-07-14T08:01:07Z)
• DiffPD: Differentiable Projective Dynamics with Contact [65.88720481593118]
  We present DiffPD, an efficient differentiable soft-body simulator with implicit time integration. We evaluate the performance of DiffPD and observe a speedup of 4-19 times compared to the standard Newton's method in various applications.
  arXiv  Detail & Related papers  (2021-01-15T00:13:33Z)
• Combining Differentiable PDE Solvers and Graph Neural Networks for Fluid Flow Prediction [79.81193813215872]
  We develop a hybrid (graph) neural network that combines a traditional graph convolutional network with an embedded differentiable fluid dynamics simulator inside the network itself. We show that we can both generalize well to new situations and benefit from the substantial speedup of neural network CFD predictions.
  arXiv  Detail & Related papers  (2020-07-08T21:23:19Z)
• ASFD: Automatic and Scalable Face Detector [129.82350993748258]
  We propose a novel Automatic and Scalable Face Detector (ASFD) ASFD is based on a combination of neural architecture search techniques as well as a new loss design. Our ASFD-D6 outperforms the prior strong competitors, and our lightweight ASFD-D0 runs at more than 120 FPS with Mobilenet for VGA-resolution images.
  arXiv  Detail & Related papers  (2020-03-25T06:00:47Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.