Wavelet neural operator: a neural operator for parametric partial differential equations

TL;DR

A novel operator learning algorithm referred to as the Wavelet Neural Operator (WNO) that blends integral kernel with wavelet transformation and is used to build a digital twin capable of predicting Earth's air temperature based on available historical data.

Abstract

With massive advancements in sensor technologies and Internet-of-things, we now have access to terabytes of historical data; however, there is a lack of clarity in how to best exploit the data to predict future events. One possible alternative in this context is to utilize operator learning algorithm that directly learn nonlinear mapping between two functional spaces; this facilitates real-time prediction of naturally arising complex evolutionary dynamics. In this work, we introduce a novel operator learning algorithm referred to as the Wavelet Neural Operator (WNO) that blends integral kernel with wavelet transformation. WNO harnesses the superiority of the wavelets in time-frequency localization of the functions and enables accurate tracking of patterns in spatial domain and effective learning of the functional mappings. Since the wavelets are localized in both time/space and frequency, WNO can provide high spatial and frequency resolution. This offers learning of the finer details of the parametric dependencies in the solution for complex problems. The efficacy and robustness of the proposed WNO are illustrated on a wide array of problems involving Burger's equation, Darcy flow, Navier-Stokes equation, Allen-Cahn equation, and Wave advection equation. Comparative study with respect to existing operator learning frameworks are presented. Finally, the proposed approach is used to build a digital twin capable of predicting Earth's air temperature based on available historical data.

Figures (6)

• Figure 1: Illustration of the architecture of the wavelet neural operator (WNO). (a) Schematic of the proposed neural operator. First lift the inputs to a higher dimension by a local transformation $P(\cdot)$. Then pass the lifted inputs to a series of wavelet kernel integral layer. Transform back the final output of the wavelet integral layer using a local transformation $Q(\cdot)$. Finally, activate the output of $Q(\cdot)$, which will provide the solution $u(x)$. In the wavelet kernel integral layer, the multilevel wavelet decomposition of the inputs are performed first, as a result of which the horizontal, vertical and diagonal coefficients at different levels are obtained. The convolution between neural network weights and the subband coefficients of last level are then performed. The inverse wavelet transform is performed on the convoluted coefficients to transform the reduced inputs to original dimension. Suitable activation is done on the output of the wavelet kernel integration layer. (b) A simple WNO with one wavelet kernel integral layer. The inputs contain the initial parameters and space information. The local transformations $P(\cdot)$ and $Q(\cdot)$ are modelled as shallow fully connected neural network. The output of $P(\cdot)$ is fed into the wavelet integral layer. The integral layer consists of two separate branches. In the first branch the wavelet decomposition of inputs followed by parameterization of integral kernel is done. In the second branch a convolution neural network (CNN) with kernel size 1 is constructed. The outputs of the two branches are then summed and activation's are performed. Then the outputs are passed through the transformation $Q(\cdot)$, which provides the target solution $u(x)$. In similar manner, a WNO with arbitrary number of wavelet integral layers can be be constructed.
• Figure 2: Operator learning of flow problems. First row: Burgers' PDE with 1024 spatial resolution. (a) the initial conditions. (b) The solution to the 1D Burgers flow equation at $t=1$. The aim was to learn the differential operator of Burgers' PDE. For given input functions, the predictions are shown and it can be seen that the predictions of proposed WNO matches with the true solution almost exactly. Darcy flow in the rectangular domain with spatial resolution of ${\bf{85 \times 85}}$. The aim is to learn the differential operator of Darcy equation. The prediction of output pressure fields from two different input permeability conditions are shown in the figure. The errors indicate that the predictions match the true solutions almost exactly. Darcy flow in a triangular domain with a notch. The goal here is to test the performance of proposed WNO for learning nonlinear operators in complex geometric conditions. For the given boundary conditions, the pressure fields obtained using the proposed WNO are shown in the figure. The errors against each of the predictions depict that the proposed WNO can learn operators in complex geometric conditions also. Navier-Stokes equation with spatial resolution of ${\bf{64 \times 64}}$. In this example, the aim is to use WNO for learning of differential operators of 2D time dependent problems. The initial vorticity (top row) and the predictions at $t=11$s (second row first column) and $t=20$s (second row second column) are shown. The errors indicates that the proposed WNO can learn highly nonlinear operators.
• Figure 3: Operator learning of Allen-Cahn reactor-diffusion equation and wave advection equation. Wave advection PDE: For solution, a spatio-temporal resolution of $40 \times 40$ is used. The proposed WNO is used here to learn the operator that can predict the wave propagation through continuous medium. Three independent initial conditions (IC) were taken, for which the true, prediction and corresponding errors are plotted. The error plots demonstrate that the prediction results of proposed WNO is very accurate. Allen-Cahn PDE: For three different initial conditions, the final solutions at $t=20$s are plotted. This PDE has chaotic nature, i.e., the solutions diverges significantly for small perturbations in initial conditions. The aim here is to test the ability of the proposed WNO for prediction of chaotic type solutions. The solutions are obtained on a spatial resolution of $43 \times 43$. Three different IC are taken and for each IC the true and WNO predicted solutions are shown. The errors against each case demonstrate that the proposed WNO could predict the true solutions almost exactly.
• Figure 4: Methodology for the short-term weekly forecast of 2m air temperature on a resolution of ${\mathbf{2^{\circ} \times 2^{\circ}}}$. The aim of the short-term forecast is to train the WNO using previous 7 days data and then predict the 2m temperature for next 7 days. To do this, a recurrent neural network (RNN) structure is employed within the WNO. The RNN framework takes the previous 7 days data as input and then performs prediction for Day 8. For prediction the parameters are learned through convolution using the proposed WNO architecture. The Day 8 prediction data is then appended in the previous 7 day database. To maintain the recurrent structure i.e., to transfer the information from Day 8 to future predictions the information of Day 1 from previous data is discarded and the process is continued until the next 7 day predictions are made. This results in a highly efficient and accurate.
• Figure 5: WNO for forecast of the weekly 2m air temperature on a resolution of ${\mathbf{2^{\circ} \times 2^{\circ}}}$. The predictions along with their corresponding errors for weekly forecast of 2m temperature at 12.00 UTC between $10^{th}$ to $15^{th}$ April are shown in the above figure. The errors are obtained with respect to the forecast made by Integrated Forecast System (IFS) of ECMWF. The predictions using proposed WNO are obtained with an error of the magnitude of $<1 \%$ that is a proof of the nearly accurate match with the actual forecast data. Once the training is done the predictions for the entire week are made within 3 seconds on the given GPU, which is evident of the computational advantage of the proposed WNO framework.