DeepONet for QG Model

Problem setup for beta model

The overarching goal of this study is to train a DeepONet aided neural operator to approximate a map from the \(\psi(\bar{x},\bar{y},t)\) to \(\Psi(x,y,t)\), where \(\psi\) and \(\Psi\) represent low and high resolution Quasi-Geostrophic (QG) flow fields in space. Therefore, our aim is to formulate a DeepONet cite{lu2021learning} architecture to learn the operator mapping \(\mathcal{G}\) from the functional space \(\psi\) to the function space \(\Psi\), which is expressed as

\[\mathcal{G}: \psi(\bar{x},\bar{y},t) \xrightarrow[]{} \Psi(x,y,t).\]

To train the DeepONet based neural operator, we generated the data by solving the two layer QG system for a very long time interval.

Using a variable-separation approach (in space-time), we postulate that the QG flow field can be expressed as

\[\psi = \phi(x, y)\zeta(t).\]

Motivated by the above equation, we construct a novel DeepONet architecture by pairing DeepONet with an LSTM architecture in a single framework in the following manner:

  1. The first step is to train a DeepONet to approximate the operator \(\mathcal{G}\) in space by mapping the low resolution data \(\psi(\bar{x},\bar{y},t)\) \({\in \mathbb{R}^{24 \times 24}}\) to high resolution data \(\Psi(x,y,t)\) \({\in \mathbb{R}^{24 \times 24} \approx \mathcal{P}\left( \mathbb{R}^{128 \times 128}\right)}\), where \(\mathcal{P}\) is high order and numerically constructed projection operator. In other words, we want to approximate \(\phi(x, y)\) using a DeepONet.

  2. The second step is to use the high resolution output approximated by DeepONet and incorporate the memory of system by using a Long short-term memory (LSTM) network, which is approximating \(\zeta(t)\) using a sequence-to-sequence mapping. However, the solution of QG system lies in a high dimensional space \((\mathbb{R}^{24\times24})\) space, which is also the dimension of the feature space for training and testing of LSTM, resulting in a computationally taxing process for the training of LSTM. To circumvent this issue instead of vanilla LSTM, we use a 2D Convolutional LSTM, which replaces the binary operation in vanilla LSTM with convolution.

Alternative text
Alternative text

Results for beta model

The figure below shows the prediction of the pdfs of streamfunctions \(\psi_1\) and \(\psi_2\), and the amplitude of wavenumbers \((1,0), (1,1)\) of the barotropic streamfunction for \(\beta = 2\) and \(r = 0.2\), while training took place for \(\beta = 2\) and \(r = 0.1\).

The objective of this task is to correct the statistics in space and time by utilizing the the low resolution data for a prescribed measurable and forecast in time.

Alternative text
Alternative text
Alternative text
Alternative text
Alternative text
Alternative text