Reconstructing 3D Flow from 2D Data with Diffusion Transformer

TL;DR

The paper tackles reconstructing 3D flow fields from 2D observations, addressing the high cost of 3D CFD and PIV by learning a data-driven prior via a Diffusion Transformer. It introduces plane position embeddings and window/plane attention to enable reconstruction from arbitrary sets of 2D planes while keeping computation tractable. Empirical results on INS/CNS DNS data show state-of-the-art or competitive performance (nRMSE, PSNR, SSIM), with ablations confirming efficiency gains and the importance of plane embeddings. This approach enables flexible, scalable 3D flow reconstruction from partial 2D measurements, with potential impact on experimental fluid dynamics and rapid flow analysis, though future work could incorporate physical constraints to further improve fidelity.

Abstract

Fluid flow is a widely applied physical problem, crucial in various fields. Due to the highly nonlinear and chaotic nature of fluids, analyzing fluid-related problems is exceptionally challenging. Computational fluid dynamics (CFD) is the best tool for this analysis but involves significant computational resources, especially for 3D simulations, which are slow and resource-intensive. In experimental fluid dynamics, PIV cost increases with dimensionality. Reconstructing 3D flow fields from 2D PIV data could reduce costs and expand application scenarios. Here, We propose a Diffusion Transformer-based method for reconstructing 3D flow fields from 2D flow data. By embedding the positional information of 2D planes into the model, we enable the reconstruction of 3D flow fields from any combination of 2D slices, enhancing flexibility. We replace global attention with window and plane attention to reduce computational costs associated with higher dimensions without compromising performance. Our experiments demonstrate that our model can efficiently and accurately reconstruct 3D flow fields from 2D data, producing realistic results.

Figures (7)

• Figure 1: Pipeline and architecture. We use the Diffuse Transformer to reconstruct 3D flow from 2D flow. Top: The 2D flow is fed into the model through three paths: 1) it is padded and concatenated with the input; 2) features are extracted using CLIP, combined with the timestep embedding and plane position embedding, and used as conditioning; and 3) the final layer output of the CLIP image encoder is fed into the model via cross-attention. Middle: The structure of the Diffuse Transformer. The input flow field undergoes patchification, passes through $N$ transformer layers, and is unpatchified to produce the predicted noise in the flow field. Global attention is replaced with more efficient window attention and plane attention within the transformer. Bottom: The structure of different attention mechanisms. Both window attention and plane attention are implemented through reshaping operations.
• Figure 2: Visual explanation of different attention. Each grid of the cube represents an input token, and only tokens of the same color can attend to each other. (a) Global attention. (b) Window attention. (c) Plane attention: from left to right, $yOz$ plane attention, $xOz$ plane attention, and $xOy$ plane attention.
• Figure 3: Visualization of results from different models on the INS(INT) (top left), INS(EXT) (top right), CNS (bottom left), and CNS $xOy$ planes where the reference $xOy$ plane is not provided. Note that the displayed 2D flow fields are not the inputs to the models.
• Figure 4: Relative positions to the center of the 3D flow field versus three test metrics on INS(INT).
• Figure 5: Relative positions to the center of the 3D flow field versus three test metrics on CNS.