Physics-Informed Neural Operator for  Learning Partial Differential Equations

1 INTRODUCTION . Machine learning-based methods are starting to show promise in scientific computing and especially in solving partial differential equations ( PDEs ) . They have demonstrated advantages in both efficiency and accuracy compared to conventional solvers . They are even able to tackle previously intractable problems such as higher-dimensional , multi-scale , high-contrast , and chaotic PDE systems ( Um et al. , 2020 ; Brunton et al. , 2020 ; Fan et al. , 2018 ; Long et al. , 2018 ; Han et al. , 2018 ; Bruno et al. , 2021 ) . Broadly , ML-based approaches for PDEs can be divided into two categories : optimizing to solve for a specific solution function of PDE vs. learning the solution operator over a family of PDEs . Optimization of solution function and PINN . Most ML-based methods , as well as the conventional solvers , fall into this category . Conventional solvers such as FDM and FEM usually discretize the domain into a grid and optimize/approximate the solution function on the grid , which imposes a truncation error . The Physics-Informed Neural Network ( PINN ) -type methods are proposed to overcome the discretization issue ( Raissi et al. , 2019 ) . They use a neural network as the ansatz of the solution function and take advantage of auto-differentiation to compute the exact , mesh-free derivatives . Recently , researchers have developed numerous variations of PINN with promising results on inverse problems and partially observed tasks ( Lu et al. , 2021a ; Zhu et al. , 2019 ; Smith et al. , 2021 ) . However , compared to conventional solvers , PINNs face several optimization issues : ( 1 ) the challenging optimization landscape from soft physics or PDE constraints ( Wang et al. , 2021a ) , ( 2 ) the difficulty to propagate information from the initial or boundary conditions to unseen parts of the interior or to future times ( Dwivedi & Srinivasan , 2020 ) , and ( 3 ) the sensitivity to hyper-parameters selection ( Sun et al. , 2020 ) . As a result , PINNs are still unable to compete with conventional solvers in most cases , and they often fail to converge on high-frequency or multi-scale PDEs ( Wang et al. , 2020b ; Fuks & Tchelepi , 2020 ; Raissi et al. , 2020 ) . In this work , we propose to overcome these optimization challenges by integrating operator learning with PINN . ( a ) learn the solution operator from a family of equations . ( b ) use the learned operator as ansatz to solve for a specific instance . Operator learning and neural operators . A recent alternative approach is to learn the solution operator of a family of PDEs , defined by the map from the input–initial conditions and boundary conditions , to the output–solution functions . In this case , usually , a dataset of input-output pairs from an existing solver is given . There are two main aspects to consider ( a ) model : to design models for learning highly complicated PDE solution operators , and ( b ) data : to be data-efficient and to improve generalization . Recent advances in operator learning replace traditional convolutional neural networks and U-Nets from computer vision with operator-based model tailored to PDEs with greatly improved model expressiveness ( Li et al. , 2020c ; Lu et al. , 2019 ; Patel et al. , 2021 ; Wang et al. , 2020a ; Duvall et al. , 2021 ) . Specifically , the neural operator generalizes the neural network to the operator setting where the input and output spaces are infinite-dimensional . The framework has shown success in learning resolution-invariant solution operators for highly non-linear problems such as turbulence flow ( Li et al. , 2020b ; a ) . However , the data challenges remain : ( 1 ) the need for training data , which assumes an existing solver or experimental setup , ( 2 ) the non-negligible generalization error , and ( 3 ) extrapolation to unseen conditions . These issues can be addressed by adding physics or PDE constraints to operator learning ( Zhu et al. , 2019 ; Wang et al. , 2021b ; Zhang et al. , 2021 ) . Our contributions . To overcome the shortcomings of both physics-informed optimization and data-driven operator learning , we propose the physics-informed neural operator ( PINO ) that combines operator learning with equation solving ( test-time optimization ) . It requires fewer or no data points to learn the operator and generalizes better . In PINO , we use the pre-trained operator as the ansatz to optimize for the solution function at test time , which reduces the generalization error . Compared to PINN , PINO has a much better optimization landscape and representation space , and hence , PINO converges faster and more accurately . Our contributions can be summarized as follows : • We propose the physics-informed neural operator ( PINO ) , combining the operator-learning and physics-informed settings . We introduce the pre-training and test-time optimization schemes that utilize both the data and equation constraints ( whichever are available ) . We develop an efficient method to compute the exact gradient for neural operators to incorporate the equation constraints . • By utilizing pre-trained operator ansatz , PINO overcomes the challenge of propagating information from the initial condition to future time steps with ( soft ) physics constraints . It can solve the 2d transient flow over an extremely long time period , where PINN and DeepONet ( Lu et al. , 2019 ) fail . Even without any pre-training and using only PDE constraints for the given instance , PINO still outperforms PINN by 20x smaller error and 25x speedup on the chaotic Kolmogorov flow , demonstrating superior expressivity of the neural operator over standard neural networks . • By utilizing the equation constraints , PINO requires fewer or no training data and generalizes better compared to FNO ( Li et al. , 2020c ) . On average it has 7 % smaller error on the transient and Kolmogorov flows , while matching the speedup of FNO ( 400x ) compared to the GPU-based pseudo-spectral solver ( He & Sun , 2007 ) , matching FNO . Further , the pre-trained PINO model on the Navier Stokes equation can be easily transferred to different Reynolds numbers ranging from 100 to 500 using test-time optimization . • We propose the forward and backward PINO models for inverse problems . Our approach accurately recovers the coefficient function in the Darcy flow which is 3000x faster than the conventional solvers using accelerated MCMC ( Cotter et al. , 2013 ) . Our major novelty and contributions are to use the pre-trained operator ansatz with instance-wise fine-tuning to overcome the optimization challenges in PINN and the generalization challenges in operator learning . Previous works such as PINN-DeepONet ( Wang et al. , 2021b ) and Physicsconstrained modeling ( Zhu et al. , 2019 ) use the PDE constraints in operator learning , like we do during the pre-training phase in PINO . However , we propose several methodological advances as well as extensive experiments to understand the optimization and generalization challenges . Our methodological advances include : ( 1 ) Instance-wise fine-tuning at test-time to further improve the fidelity of the operator ansatz . ( 2 ) Efficient Fourier-space methods for computing derivatives present in the PDE loss . ( 3 ) Efficient learning through the design of data augmentation and loss functions . ( 4 ) Novel formulation for inverse problems that results in accurate recovery as well as good speedups . PINN vs. PINO : pointwise vs. function-wise optimization . The neural operator ansatz in PINO has an easier optimization landscape and a more expressive representation space compared to the neural networks ansatz in PINN . The neural operator parameterizes the solution function as an aggregation of basis functions , and hence , the optimization is in the function space . This is easier than just optimizing a single function as in PINN . Further , we can learn these basis functions in the pre-training phase which makes the test-time optimization on the querying instance even easier . In this case , PINO does not need to solve from scratch . It just fine-tunes the solution function parameterized by the solution operator . Thus , PINO is much faster and more accurate compared to PINN . 2 PRELIMINARIES AND PROBLEM SETTINGS . 2.1 PROBLEM SETTINGS . We consider two natural class of PDEs . In the first , we consider the stationary system P ( u , a ) = 0 , in D ⊂ Rd u = g , in ∂D ( 1 ) where D is a bounded domain , a ∈ A ⊆ V is a PDE coefficient/parameter , u ∈ U is the unknown , and P : U ×A → F is a possibly non-linear partial differential operator with ( U , V , F ) a triplet of Banach spaces . Usually the function g is a fixed boundary condition ( potentially can be entered as a parameter ) . This formulation gives rise to the solution operator G† : A → U defined by a 7→ u . A prototypical example is the second-order elliptic equation P ( u , a ) = −∇ · ( a∇u ) + f . In the second setting , we consider the dynamical system du dt = R ( u ) , in D × ( 0 , ∞ ) u = g , in ∂D × ( 0 , ∞ ) u = a in D̄ × { 0 } ( 2 ) where a = u ( 0 ) ∈ A ⊆ V is the initial condition , u ( t ) ∈ U for t > 0 is the unknown , and R is a possibly non-linear partial differential operator with U , and V Banach spaces . As before , we take g to be a known boundary condition . We assume that u exists and is bounded for all time and for every u0 ∈ U . This formulation gives rise to the solution operator G† : A → C ( ( 0 , T ] ; U ) defined by a 7→ u. Prototypical examples include the Burgers ’ equation and the Navier-Stokes equation . 2.2 SOLVING EQUATION USING THE PHYSICS-INFORMED LOSS ( PINN ) . Given an instance a and a solution operator G† defined by equations ( 1 ) or ( 2 ) , we denote by u† = G† ( a ) the unique ground truth . The equation solving task is to approximate u† . This setting consists of the ML-enhanced conventional solvers such as learned finite element , finite difference , and multigrid solvers ( Kochkov et al. , 2021 ; Pathak et al. , 2021 ; Greenfeld et al. , 2019 ) , as well as purely neural network-based solvers such as the Physics-Informed Neural Networks ( PINNs ) , Deep Galerkin Method , and Deep Ritz Method ( Raissi et al. , 2019 ; Sirignano & Spiliopoulos , 2018 ; Weinan & Yu , 2018 ) . Especially , these PINN-type methods use a neural network uθ with parameters θ as the the ansatz to approximate the solution function u† . The parameters θ are found by minimizing the physics-informed loss with exact derivatives computed using automatic-differentiation ( autograd ) . In the stationary case , the physics-informed loss is defined by minimizing the l.h.s . of equation ( 1 ) in the squared norm of F . A typical choice is F = L2 ( D ) , giving the loss function Lpde ( a , uθ ) = ‖P ( a , uθ ) ‖2L2 ( D ) + α‖uθ|∂D − g‖ 2 L2 ( ∂D ) = ∫ D |P ( uθ ( x ) , a ( x ) ) |2dx+ α ∫ ∂D |uθ ( x ) − g ( x ) |2dx ( 3 ) In the case of a dynamical system , it minimizes the residual of equation ( 2 ) in some natural norm up to a fixed final time T > 0 . A typical choice is the L2 ( ( 0 , T ] ; L2 ( D ) ) norm , yielding Lpde ( a , uθ ) = ∫ T 0 ∫ D |duθ dt ( t , x ) −R ( uθ ) ( t , x ) |2dxdt+ α ∫ T 0 ∫ ∂D |uθ ( t , x ) − g ( t , x ) |2dxdt + β ∫ D |uθ ( 0 , x ) − a ( x ) |2dx ( 4 ) The PDE loss consists of the physics loss in the interior and the data loss on the boundary and initial conditions , with hyper-parameters α , β > 0 . It can be generalized to variational form as in ( Weinan & Yu , 2018 ) . Challenges of PINN PINNs take advantage of the universal approximability of neural networks , but , in return , suffer from the low-frequency induced bias . Empirically , PINNs often fail to solve challenging PDEs when the solution exhibits high-frequency or multi-scale structure ( Wang et al. , 2021a ; 2020b ; Fuks & Tchelepi , 2020 ; Raissi et al. , 2020 ) . Further , as an iterative solver , PINNs have difficulty propagating information from the initial condition or boundary condition to unseen parts of the interior or to future times ( Dwivedi & Srinivasan , 2020 ) . For example , in challenging problems such as turbulence , PINNs are only able to solve the PDE on a relatively small domain ( Jin et al. , 2021 ) , or otherwise , require extra observational data which is not always available in practice ( Raissi et al. , 2020 ; Cai et al. , 2021 ) . In this work , we propose to overcome the challenges posed by the optimization by integrating operator learning with PINNs .

Summary: The authors introduce physics-informed neural operator for learning partial differential equations (PDEs) from data with deep learning. This approach combines two recent and popular methods in the field of machine learning and PDEs, the physics-informed neural network (PINNs) for solving PDEs and Fourier neural operator (FNO) for learning solution operators. The paper takes advantages of both approaches to overcome convergence issues arising in PINNs by incorporating PDE constraints into FNO. Then, the authors apply their method to a number of challenging test cases and obtain better results than PINNs on PDE solving problems and FNO for learning solution operators.