Coarse Graining with Neural Operators for Simulating Chaotic Systems

TL;DR

This work shows that traditional closure-based coarse-graining struggles with chaotic systems due to non-uniqueness of the reduced-force mapping and slow convergence of empirical measures. It introduces a physics-informed neural operator (PINO) with a multi-fidelity training strategy that first learns coarse-grained dynamics on cheap data and then refines on limited high-fidelity data with PDE losses, providing resolution-invariant long-term statistics. Theoretical guarantees are given via a functional Liouville-flow framework, demonstrating convergence of MF-PINO to the projected invariant measure under reasonable conditions. Empirically, MF-PINO attains substantial speedups (up to 330x) and accurate long-term statistics (relative error around 10%) on KS and NS chaotic flows, outperforming classical closures and learning-based closures.

Abstract

Accurately predicting the long-term behavior of chaotic systems is crucial for various applications such as climate modeling. However, achieving such predictions typically requires iterative computations over a dense spatiotemporal grid to account for the unstable nature of chaotic systems, which is expensive and impractical in many real-world situations. An alternative approach to such a full-resolved simulation is using a coarse grid and then correcting its errors through a \textit{closure model}, which approximates the overall information from fine scales not captured in the coarse-grid simulation. Recently, ML approaches have been used for closure modeling, but they typically require a large number of training samples from expensive fully-resolved simulations (FRS). In this work, we prove an even more fundamental limitation, i.e., the standard approach to learning closure models suffers from a large approximation error for generic problems, no matter how large the model is, and it stems from the non-uniqueness of the mapping. We propose an alternative end-to-end learning approach using a physics-informed neural operator (PINO) that overcomes this limitation by not using a closure model or a coarse-grid solver. We first train the PINO model on data from a coarse-grid solver and then fine-tune it with (a small amount of) FRS and physics-based losses on a fine grid. The discretization-free nature of neural operators means that they do not suffer from the restriction of a coarse grid that closure models face, and they can provably approximate the long-term statistics of chaotic systems. In our experiments, our PINO model achieves a 330x speedup compared to FRS with a relative error $\sim 10\%$. In contrast, the closure model coupled with a coarse-grid solver is $60$x slower than PINO while having a much higher error $\sim186\%$ when the closure model is trained on the same FRS dataset.

Figures (12)

• Figure 1: (A-B):Conceptual Illustration of Different Methods (training and inference)(A) During training, learning-based closure models can only leverage information captured by the coarse grids, though it is derived from fully-resolved simulations. In contrast, neural operators can directly learn the fine-scale information through either supervised or physics-informed training with fine-grid inputs. (B) (i) Fully-resolved simulations are conducted on fine grids. They provide gold-standard estimations of statistics but are extremely costly. (ii) Coarse-graining framework with closure model ($\mathcal{A}\overline{u}+clos(\overline{u};\theta)$), which adopts a closure model in conjunction with a coarse-grid numerical solver. (iii) Coarse graining with neural operator. After training with fine-grid inputs in (A), neural operators seamlessly support coarse-grid inputs. fine-scale information is implicitly incorporated into coarse-grid simulations through the interaction between inputs and model parameters. (C): Non-Unique Issue of Closure Model. The arrows represent the evolving direction of trajectories (i.e. the time derivative of the governing dynamics). In chaotic systems, trajectories with slightly different initializations rapidly diverge from one another, but ultimately they all converge to the attractor (if one exists) after long-term evolution. Coarse graining can be viewed as designing a vector field that drives the dynamics in the reduced space. In closure model framework, the assigned vector field is $\mathcal{A}\overline{u}+clos(\overline{u};\theta)$. Multiple points (e.g., $u_1$ and $u_2$) of the ground-truth attractor (i.e., equilibrium state) map to the same filtered value $\overline{u}$, precluding the closure model from identifying the correct dynamics ($\mathcal{F}(\mathcal{A}u_1)$ and $\mathcal{F}(\mathcal{A} u_2)$) in the filtered space $\mathcal{F}(\mathcal{H})$. By minimizing the loss function, the model learns to predict the average of these multiple choices (green arrow), leading the simulation to wrongly diverge from the filtered attractor. (D): Prediction of energy spectrum at the attractor. 'Closure Model' represents a single-state learning-based closure. (E):Total variation distance from ground-truth invariant measure versus computation cost. 'Fully-Resolved': gold-standard fully-resolved simulations, providing ground-truth statistics. 'CGS': coarse-grid simulation without closure model. 'Smag.': Smagorinsky closure model.'ML-Closure': learning-based single-state model. 'DSM': Dynamical Smagorinsky closure model. 'RNN': history-aware closure model with recurrent neural network. 'DM': stochastic closure model based on diffusion model. 'MF-FNO': Multi-fidelity FNO. 'MF-PINO': multi-fidelity physics-informed neural operator. All learning-based methods are trained with same amount of fully-resolved data. Our methods with neural operators (MF-FNO and MF-PINO) are the fastest. MF-PINO is the closest to ground truth among all methods that take coarse-grid inputs, beating those methods using an explicit closure model.
• Figure 2: Total Variation (TV) error results for Navier-Stokes with $Re=1.6\times 10^4$. The $(k_x,k_y)$-element represents the TV error compared with ground-truth fully-resolved simulations regarding the distribution of the mode length of the component for $(k_x,k_y)$- Fourier basis $e^{i\frac{2\pi}{L}(k_xx+k_yy)}$. Total variation ranges within $[0,1]$. The smaller the TV error, the closer the distribution is to the ground truth. From left to right: coarse-grid simulation without closure model (CGS), Smagorinsky model (Smag.), Dynamic Smagorinsky Model (DSM), single-state learning-based model (Single), history-aware closure model based on recurrent neural network (RNN), stochastic closure model based on diffusion model (DM), Multi-fidelity FNO (MFF), and multi-fidelity physics-informed neural operator (MF-PINO). Our approach with MF-PINO performs the best among all methods that simulate on coarse grid systems.
• Figure 3: Experiment Results of Some Statistics for NS Equation with $Re=1.6\times 10^4$. 'Ground Truth'(blue curve) refers to fully-resolved simulation. 'DSM': dynamical Smagorinsky model. 'Single': learning-based single-state closure model. More results for other statistics and all methods can be found in the Appendix. Our method (red) is the closest to ground truth among all methods with coarse-grid inputs.
• Figure 4: Illustration of the model architecture.Left: The neural operator model takes the initial condition (a spatial function) as input, e.g. $u_0(x,y)$ in 2D case. Its output is the trajectory $u(x,y,t)$ in time interval $t\in [0,h]$, where $h$ is a model parameter. Right: Architecture of neural operator. The input spatial function is first lifted into a spatial-temporal function by replicating $u_0$ in the temporal dimension and equipping it with temporal position embedding. Then it is transformed with point-wise operator $P$ which is usually a multi-layer perceptron (MLP) that increases the feature dimension. The Fourier layers consist of spectral convolution implemented based on Fourier transform $\mathcal{F}$, and point-wise operation. In the figure, $W$ denote the weight matrix in MLP and $\sigma$ denotes nonlinear activation functions. After several Fourier layers, the final output is obtained with a point-wise operator $Q$.
• Figure 5: Visualization of Dataset These figures are snapshots of one FRS trajectory for the Navier-Stokes equation with $Re=1.6\times 10^4$. From the plots, the fluid field appears to approach a dynamical equilibrium after $t=50$.

Theorems & Definitions (59)

• Theorem 2.1
• Remark 2.2
• Theorem 2.3
• Remark 2.4
• Remark 2.5
• Remark 2.6
• Definition A.1
• Definition A.2
• Lemma A.3: Shadowing Lemma
• Remark A.4