PICT -- A Differentiable, GPU-Accelerated Multi-Block PISO Solver for Simulation-Coupled Learning Tasks in Fluid Dynamics

TL;DR

PICT tackles the challenge of differentiable fluid simulation by delivering a GPU-accelerated, multi-block PISO solver with end-to-end differentiability implemented in PyTorch. It combines a DtO-based forward-discrete gradient flow with OtD-style backward solves, enabling learned turbulence models to be trained via solver gradients while preserving divergence-free constraints. The framework is validated on 2D and 3D flows, showing that CNN-based subgrid-scale models can match or exceed high-fidelity references with substantial runtime savings, and that gradient-path ablation offers practical speedups for short rollouts. The work provides strong evidence that solver-gradient-based learning can yield stable, accurate, and efficient turbulence modeling, supported by open-source software and extensive experimentation across canonical benchmarks.

Abstract

Despite decades of advancements, the simulation of fluids remains one of the most challenging areas of in scientific computing. Supported by the necessity of gradient information in deep learning, differentiable simulators have emerged as an effective tool for optimization and learning in physics simulations. In this work, we present our fluid simulator PICT, a differentiable pressure-implicit solver coded in PyTorch with Graphics-processing-unit (GPU) support. We first verify the accuracy of both the forward simulation and our derived gradients in various established benchmarks like lid-driven cavities and turbulent channel flows before we show that the gradients provided by our solver can be used to learn complicated turbulence models in 2D and 3D. We apply both supervised and unsupervised training regimes using physical priors to match flow statistics. In particular, we learn a stable sub-grid scale (SGS) model for a 3D turbulent channel flow purely based on reference statistics. The low-resolution corrector trained with our solver runs substantially faster than the highly resolved references, while keeping or even surpassing their accuracy. Finally, we give additional insights into the physical interpretation of different solver gradients, and motivate a physically informed regularization technique. To ensure that the full potential of PICT can be leveraged, it is published as open source: https://github.com/tum-pbs/PICT.

Figures (26)

• Figure 1: A flow chart showing the high level computational graph of our solver. Viscosity, boundaries, velocity sources, transformations, and non-orthogonal correction have been omitted for clarity. The dashed line represents the pressure correction loop, and both linear solves (Adv and P) are highlighted in blue. All shown paths are differentiable.
• Figure 2: Three examples of transformed multi-block meshes that can be handled by our simulator. From left to right: a channel flow grid refined towards the walls, a ring grid with a round obstacle at the center, a refined C-grid around a NACA 0012 airfoil. The meshes have been coarsened for improved visibility and colors indicate different blocks.
• Figure 3: Velocity profiles for the 3D lid-driven cavity with $\text{Re}=1000$ for increasing resolutions. The left image in each subfigure is the $u$-velocity on the vertical center line, and the right is the $v$-velocity on the horizontal center line. (b) uses a grid that was refined towards all boundaries. The velocities are normalized with the Reynolds number. The reference is a high-resolution DNS Lid3D.
• Figure 4: Turbulence statistics for a 3D TCF with $\text{Re}_\tau = 550$. The statistics are averaged over time and stream- and span-wise direction, normalized with the average $u_\tau$ of the corresponding simulation, and plotted against the wall-normal direction. The solid lines show results from our solver, while the dashed lines show those of OpenFOAM's PISO implementation using the same computational mesh. Dotted lines indicate the spectral reference from Hoyas and Jiménez TCF_2008_10, while fine dotted lines in the $U^+$ plot are the log-law and law of the wall.
• Figure 5: Qualitative visualizations of the pressure and velocity components of our TCF benchmark. Only the flow close to the bottom wall in the computational space is shown. The color gradient and iso-surface values are chosen individually to highlight the detail of the flow structures. Top: $u$ and $v$, bottom: $w$ and $p$.