Learned discretizations for passive scalar advection in a two-dimensional turbulent flow

The computational cost of fluid simulations increases rapidly with grid resolution. This has given a hard limit on the ability of simulations to accurately resolve small-scale features of complex flows. Here we use a machine learning approach to learn a numerical discretization that retains high accuracy even when the solution is under-resolved with classical methods. We apply this approach to passive scalar advection in a two-dimensional turbulent flow. The method maintains the same accuracy as traditional high-order flux-limited advection solvers, while using $4\times$ lower grid resolution in each dimension. The machine learning component is tightly integrated with traditional finite-volume schemes and can be trained via an end-to-end differentiable programming framework. The solver can achieve near-peak hardware utilization on CPUs and accelerators via convolutional filters.

Figure 1

End-to-end learning framework with differential programming. During training, the model is optimized to predict future concentrations across multiple time steps, based on a precomputed dataset of snapshots from high-resolution simulations. During inference, the optimized model is repeatedly applied to predict time evolution. The neural network component contains a stack of 2D convolutional layers with ReLU activation functions (degraded to 1D convolution for 1D problems). Physical constraints are imposed before and after the convolutional layers (Sec. 2d). In the “Time-stepping” block, $H$ is the advection operator that computes the concentration update based on the machine learning estimate of spatial derivatives.

Figure 2

One test sample for 1D advection under constant velocity. The concentration field is advected by a half grid box every time step, and returns to the original position after every 64 time steps because the domain is periodic. Our neural network model is able to maintain the initial shape indefinitely, while traditional solvers accumulate numerical diffusion over time.

Figure 3

Error for 1D advection on test data. Here we only plot even time steps ($0,2,4,...$) for a smooth curve, because the error oscillates between odd and even time steps (a result of $\text{CFL}=0.5$).

Figure 4

Neural network prediction on out-of-sample data. The neural network model is only trained on square waves, but applied to Gaussian initial conditions. The model gradually turns Gaussian waves into squares, and then maintains the squares indefinitely.

Figure 5

Result on 2D deformational flow. The flow reverses at $t=2.5$ and returns to the initial condition at $t=5$. The neural network model is able to maintain a sharp gradient, while the baseline model incurs large numerical diffusion. The spatial domain is $[0,1]\times [0,1]$ (not plotted on axis).

Figure 6

Entropy for advection under 2D deformational flow. Entropy is conserved under pure advection and increases under diffusion. Traditional monotonic solvers are only allowed to increase entropy, while our neural network model is allowed to decrease entropy and thus minimizes diffusion error over a long time.

Figure 7

One test sample under 2D turbulent flow. The initial blobs (first column) are stretched into thin filaments under the turbulent flow (last column), illustrated by the vorticity $\omega ={\partial }_x{u}_y-{\partial }_y{u}_v$ (rightmost column). The baseline solver (second-order Van Leer scheme) can resolve such filaments on the fine-resolution grid, but incurs large numerical diffusion on the coarse grid. The neural network model can preserve the sharp features on the coarse grid. The spatial domain is $[0,2\pi ]\times [0,2\pi ]$ (not plotted on axis).

Figure 8

Error for 2D turbulent advection on test data. The neural network model achieves the same accuracy as the second-order Van Leer baseline scheme at $4\times$ resolution, and entropy similar to the baseline at $8\times$ resolution.

Figure 9

Limitations of long-time stability under 2D turbulent flow. (a) Ten randomly chosen examples of concentration fields from the neural net model after 1024 time steps. One field (first row, fourth column) is entirely covered by “checkerboard” artifacts, and two others (top left and bottom right) show checkerboard artifacts in limited regions. (b) Empirical distribution function for absolute error across all models after 256 and 1024 time steps. The neural net performance suffers significantly, with about 10% of solutions having an absolute error greater than 1.

Figure 10

Interpolation coefficients for 2D advection. Illustration of how prediction of the interpolation coefficients changes for different combinations of concentration (top row) and velocity field (left column) inputs. Concentration values represented by color; velocity field has unit magnitude and changes direction as shown in the plot. The target location for the interpolation is marked by a red bar on the concentration plots. The model predominantly interpolates along the velocity field and concentration edges, rediscovering the upwindinglike methods at the corner cases of the facet. While the model learned some general symmetries, we expect even better results for models that incorporate symmetries by construction.

Figure 11

Accuracy-speed trade-off for neural network model. Each data point is a neural network model with different hyperparamters (detailed in Fig. 12). The performance of the baseline solver at intermediate grid resolutions ($64\times 64$ and $128\times 128$) is overlaid for comparison. The model accuracy is evaluated on the 2D turbulence case after 256 time steps (Sec. 3c), and the run time is measured on a single Nvidia V100 GPU. The $x$ axis shows the wall-clock time per advection time step on the coarse grid, which requires two or four time steps for the $64\times 64$ or $128\times 128$ baseline due to the CFL condition.

Figure 12

Hyperparameter effect on neural network model performance. Same as Fig. 11, but here the data points are grouped by different numbers of convolutional layers and the number of filters, with symbol shape denoting the size of the stencil. Duplicate points correspond to identical models trained with different random seeds.

Figure 13

Sample results for advection under 2D turbulent flow. Concentration fields after 256 time steps are illustrated for six different randomly initialized concentration and velocity fields. The models are the same as those illustrated in Fig. 7.