COUPLED MULTIWAVELET NEURAL OPERATOR LEARN-ING FOR COUPLED PARTIAL DIFFERENTIAL EQUATIONS

Abstract

Coupled partial differential equations (PDEs) are key tasks in modeling the complex dynamics of many physical processes. Recently, neural operators have shown the ability to solve PDEs by learning the integral kernel directly in Fourier/Wavelet space, so the difficulty for solving the coupled PDEs depends on dealing with the coupled mappings between the functions. Towards this end, we propose a coupled multiwavelets neural operator (CMWNO) learning scheme by decoupling the coupled integral kernels during the multiwavelet decomposition and reconstruction procedures in the Wavelet space. The proposed model achieves significantly higher accuracy compared to previous learning-based solvers in solving the coupled PDEs including Gray-Scott (GS) equations and the non-local mean field game (MFG) problem. According to our experimental results, the proposed model exhibits a 2ˆ" 4ˆimprovement relative L2 error compared to the best results from the state-of-the-art models.

1. INTRODUCTION

Human perception relies on detecting and processing waves. While our eyes detect waves of electromagnetic radiation, our ears detect waves of compression in the surrounding air. Going beyond waves, from complex dynamics of blood flow to sustain tissue growth and life, to navigating underwater, ground and aerial vehicles at high speeds requires discovering, learning and controlling the partial differential equations (PDEs) governing individual or webs of biological, physical and chemical phenomena (Lacasse et al., 2007; Henriquez, 1993; Laval & Leclercq, 2013; Ghanavati et al., 2017; Radmanesh et al., 2020) . Within this context, neural operators have been successfully used to learn and solve various PDEs. By representing the integral kernel termed as Green's function in the Fourier or Wavelet spaces, the fourier neural operator (Li et al., 2020b) and the multiwaveletbased neural operator (Gupta et al., 2021b; a) ) exhibit significant improvements on solving PDEs compared with previous work. However, when it comes to coupled systems characterized by coupled differential equations such as mean field games (MFGs) (Lasry & Lions, 2007; Achdou & Capuzzo-Dolcetta, 2010) , analysis of coupled cyber-physical systems (Xue & Bogdan (2017) , or analysis of the surface currents in the tropical Pacific Ocean Bonjean & Lagerloef (2002) , the interactions between the variables/functions need to be considered to decouple the system. Without the knowledge of underlying PDEs, the complex interactions can be hardly represented in the data-driven model. To build a data-driven model that can give a general representation of the interactions to efficiently solve coupled differential equations, we propose the coupled multiwavelets neural operator (CMWNO). Neural Operators. Neural operators (Li et al., 2020b; c; a; Gupta et al., 2021b; Bhattacharya et al., 2020; Patel et al., 2021) focus on learning the mapping between infinite-dimensional spaces of functions. The critical feature for neural operators is to model the integral operator namely the Green's function through various neural network architectures. The graph neural operators (Li et al., 2020b; c) use the graph kernel to model the integral operator inspired by graph neural networks; the ˚Equal Contribution Fourier neural operator (Li et al., 2020b) uses an iterative architecture to learn the integral operator in Fourier space. The multiwavelet neural operators (Gupta et al., 2021b; a) utilize the non-standard form of the multiwavelets to represent the integral operator through 4 neural networks in the Wavelet space. The neural operators are completely data-driven and resolution independent by learning the mapping between the functions directly, which can achieve the state-of-the-art performance on solving PDEs and initial value problems (IVPs). To deal with coupled PDEs in the coupled system and be data-efficient, we aim to decode the various interaction scenarios inside the neural operators. Multiwavelet Transform. In contrast to wavelets, multiwavalets (refer to Appendix C) use more than one scaling functions which are orthogonal. The multiwavelets exploit the advantages of wavelets, such as (i) the vanishing moments, (ii) the orthogonality, and (iii) the compact support. Along the essence of wavelet transform, a series of wavelet bases are introduced with scaled/shifted versions in multiwavelets to construct the basis of the coarsest scale polynomial subspace. The multiwavelet bases have been proved to be successful for representing integral operators as shown in (Alpert et al., 1993) (the discrete version of multiwavelets) and (Alpert, 1993b) . In our proposed model, to develop compactly supported multiwavelets, we use the Legendre polynomials (Appendix D) which are non-zero only over a finite interval as the basis. The differential (B{Bt) and the integral ( ť Ω ) operators can be represented by the first-order multiwavelet coefficients (s and d) of orthogonal bases via decomposition in the Wavelet space. Mean Field Games (MFGs). As a representative problem for coupled systems in the real world, MFGs gains raising attentions in various areas, including economics (Achdou et al., 2014; 2022) , finance (Guéant et al., 2011; Huang et al., 2019) and engineering (De Paola et al., 2019; Gomes et al., 2021) , etc. Building on statistical mechanics principles and infusing them into the study of strategic decision making, MFGs investigate the dynamics of a large population of interacting agents seen as particles in a thermodynamic gas. Simply speaking, MFGs consist of (i) a Fokker-Planck equation (or related PDE) that describes the dynamics of the aggregate distribution of agents, which is coupled to (ii) a Hamilton-Jacobi-Bellman equation (another PDE) prescribing the optimal control of an individual (Lasry & Lions, 2006; 2007; Huang et al., 2006; 2007) . Among different types of MFGs, the class of non-potential MFGs system with mixed couplings is particularly important as it is more reflective of the real world with a continuum of agents in a differential game.

Solutions on MFGs.

Previous works either only restrict to systems without non-local coupling, such as alternating direction method of multipliers (ADMM) (Benamou & Carlier, 2015; Benamou et al., 2017) and primal-dual hybrid gradient (PDHG) algorithm (Briceno-Arias et al., 2019; 2018) or use general purpose numerical methods for solving the MFG problems (Achdou et al., 2013a; b; Achdou & Capuzzo-Dolcetta, 2010) , which misses specific information from the target structure. In addition, the aforementioned works are not parallelizable with linear computational cost under the coupled MFGs settings. Recently, (Liu & Nurbekyan, 2020) considers dual variables of nonlocal couplings in Fourier or feature space. Furthermore, (Liu et al., 2021) expands the feature-space in the kernel-based representations of machine learning methods and uses expansion coefficients to decouple the mean field interactions. However, both dual variables and expansion coefficients need to bound the interactions of coupled system in a reasonable interval with prior knowledge. In our work, we first introduce the neural operator into coupled MFG fields, which can decouple the various interactions inside the multiwavelet domain. Novel Contributions. The main novel contributions of our work are summarized as follows: • For coupled differential equations, we propose a coupled neural operator learning scheme, named CMWNO. To the best of our knowledge, CMWNO is the first neural operator work using pure data-driven method to decouple and then solve coupled differential equations. • Utilizing multiwavelet transform, CMWNO can deal with the interactions between the kernels of coupled differential equations in the Wavelet space. Specifically, we first yield the representation of coupled information during the decomposition process of multiwavelet transform. Then, the decoupled representation can interact separately to help the operators' reconstruction process. In addition, we propose a dice strategy to mimic the information interaction during the training process.

