Learning time-dependent PDE via graph neural networks and deep operator network for robust accuracy on irregular grids

Authors: sung woong cho
jae yong lee
hyung ju hwang
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-06 08:01:20

Abstract: There has been growing interest in models that learn the operator from the parameters of a partial differential equation (PDE) to the corresponding solutions. Deep Operator Network (DeepONet) and Fourier Neural operator, among other models, have been designed with structures suitable for handling functions as inputs and outputs, enabling real-time predictions as surrogate models for solution operators. There has also been significant progress in the research on surrogate models based on graph neural networks (GNNs), specifically targeting the dynamics in time-dependent PDEs. In this paper, we propose GraphDeepONet, an autoregressive model based on GNNs, to effectively adapt DeepONet, which is well-known for successful operator learning. GraphDeepONet exhibits robust accuracy in predicting solutions compared to existing GNN-based PDE solver models. It maintains consistent performance even on irregular grids, leveraging the advantages inherited from DeepONet and enabling predictions on arbitrary grids. Additionally, unlike traditional DeepONet and its variants, GraphDeepONet enables time extrapolation for time-dependent PDE solutions.

Main paper: 
 introduction: 
 In recent years, operator learning frameworks have gained significant attention in the field of artificial intelligence. The primary goal of operator learning is to employ neural networks to learn the mapping from the parameters (external force, initial, and boundary condition) of a PDE to its corresponding solution operator. To accomplish this, researchers are exploring diverse models and methods, such as the deep operator network (DeepONet) (Lu et al., 2019) and Fourier neural operator (FNO) (Li et al., 2020) , to effectively handle functions as inputs and outputs of neural networks. These frameworks present promising approaches to solving PDEs by directly learning the underlying operators from available data. Several studies (Lu et al., 2022; Goswami et al., 2022) have conducted comparisons between DeepONet and FNO, and with theoretical analyses (Lanthaler et al., 2022; Kovachki et al., 2021a) have been performed to understand their universality and approximation bounds.In the field of operator learning, there is an active research focus on predicting time-evolving physical quantities. The DeepONet can be applied to simulate time-dependent PDEs by incorporating a time variable, denoted as t, as an additional input with spatial variables, denoted as x. However, Figure 1 : Framework of the proposed GraphDeepONet the use of both t and x as inputs at once to the DeepONet can only predict solutions within a fixed time domain and they should be treated differently from a coefficient and basis perspective. FNO (Li et al., 2020; Kovachki et al., 2021b ) also introduces two methods specifically designed for this purpose: FNO-2d, which utilizes an autoregressive model, and FNO-3d. However, a drawback of FNO is its reliance on a fixed uniform grid. To address this concern, recent studies have explored the modified FNO (Lingsch et al., 2023; Lin et al., 2022) , such as geo-FNO (Li et al., 2022) and F-FNO (Tran et al., 2023) .To overcome this limitation, researchers have explored the application of GNNs and message passing methods (Scarselli et al., 2008; Battaglia et al., 2018; Gilmer et al., 2017; Sanchez-Gonzalez et al., 2020; Pfaff et al., 2021; Lienen & Günnemann, 2022) to learn time-dependent PDE solutions. In particular, Brandstetter et al. (2022) and Boussif et al. (2022) focused on solving the time-dependent PDE based on GNNs. Brandstetter et al. (2022) proposed a Message-Passing Neural PDE Solver (MP-PDE) that utilizes message passing to enable the learning of the solution operator for PDEs, even on irregular domains. However, a limitation of their approach is that it can only predict the solution operator on the same irregular grid used as input, which poses challenges for practical simulation applications. To address this limitation, Boussif et al. (2022) introduced the Mesh Agnostic Neural PDE solver (MAgNet), which employs a network for interpolation in the feature space. This approach allows for more versatile predictions and overcomes the constraints of using the same irregular grid for both input and solution operator prediction. We aim to employ the DeepONet model, which learns the basis of the target function's spatial domain, to directly acquire the continuous space solution operator of time-dependent PDEs without requiring additional interpolation steps. By doing so, we seek to achieve more accurate predictions at all spatial positions without relying on separate interpolation processes. Our main contributions can be summarized as follows:• By effectively incorporating time information into the branch net using a GNN, GraphDeepONet enables time extrapolation prediction for PDE solutions, a task that is challenging for traditional DeepONet and its variants.• Our method exhibits robust accuracy in predicting the solution operator at arbitrary positions of the input on irregular grids compared to other graph-based PDE solver approaches. The solution obtained through GraphDeepONet is a continuous solution in the spatial domain. 
graphdeeponet for time-dependent pdes: 
 For a fixed set of positional sensors x i (0 ≤ i ≤ N -1), we formulate a graph G = (V, E), where each node i belongs to V and each edge (i, j) to E. The nodes represent grid cells, and the edges signify local neighborhoods.Encoder ϵ. The encoder maps node embeddings from the function space to the latent space. For a given node i, it maps the last solution values at node position x i , denoted as u 0 i := u 0 (x i ), to the latent embedding vector. Formally, the encoding function ϵ : R 1+d → R dlat produces the node embedding vector f 0 i as follows:EQUATION) where ϵ is multilayer perceptron (MLP). It is noteworthy that the sampling method, which includes both the number of sensors N and their respective locations x i for 0 ≤ i ≤ N -1, can differ for each input. EQUATIONEQUATIONfor m = 0, 1, ..., M -1 with h 0 i = f 0 i , where N (i) denotes the neighboring nodes of node i. Both ϕ and ψ are implemented as MLPs. The use of relative positions, i.e., x j -x i , capitalizes on the translational symmetry inherent in the considered PDEs. After the M iterations of message passing, the processor emits a vector h M i for each node i. This is used to update the latent vector f 0 i as follows:f 1 i = f 0 i + h M i , 0 ≤ i ≤ N -1. (4) The updated latent vector f 1 0:N -1 := {f 1 i } N -1 i=0is used to predict the next time step solution u 1 (x).Decoder1 -Soft attention aggregation ω.We first predict the p-coefficients for each next timestep. Here, we use the soft attention aggregation layer with the feature-level gating described by Li et al. (2019) . The soft attention aggregation ν : R dlat×N → R p consists of two neural networks to calculate the attention scores and latent vectors as follows:EQUATIONwhere ⊙ represents the element-wise product, and ω gate : R dlat+d → R p and ω feature : R dlat+1 → R p are MLPs. Note that ν is well-defined for any number of sensors N ∈ N.Decoder2 -Inner product of coefficients and basis τ . The final output is reconstructed using the p-coefficients ν[f 1 0:N -1 , ∆t] and trained global basis via trunk net τ (x) = [τ 1 (x), ..., τ p (x)] with τ j : R d → R. The next timestep is predicted asEQUATIONwhere ν[f 1 0:N -1 , ∆t] := [ν 1 , ν 2 , ..., ν p ] ∈ R p . The GraphDeepONet is trained using the mean square error Loss (1) = MSE( u 1 (x), u 1 (x)). Since the GraphDeepONet use the trunk net to learn the global basis, it offers a significant advantage in enforcing the boundary condition B[u] = 0 as hard constraints. The GraphDeepONet can enforce periodic boundaries, unlike other graph-based methods, which often struggle to ensure such precise boundary conditions (See Appendix B.6). One significant advantage of models based on Graph Neural Networks (GNNs), such as MP-PDE, MAgNet, and GraphDeepONet (ours), compared to traditional numerical methods for solving timedependent PDEs, lies in their efficiency during inference. In traditional numerical methods, solving PDEs for different initial conditions requires recalculating the entire PDE, and in real-time weather prediction scenarios (Kurth et al., 2022) , where numerous PDEs with different initial conditions must be solved simultaneously, this can result in a substantial computational burden. On the other hand, models based on GNNs (MP-PDE, MAgNet, GraphDeepONet), including the process of learning the operator, require data for a few frames of PDE. However, after training, they enable rapid inference, allowing real-time PDE solving. More details on advantage using operator learning model compared to traiditional numerical method is explained in many studies (Goswami et al., 2022; Kovachki et al., 2021b) .Table 4 presents a computational time comparison between our proposed GraphDeepONet and other GNN-based models. Due to its incorporation of global interaction using equation 5 for a better understanding of irregular grids, GraphDeepONet takes longer during both training and inference compared to MP-PDE. However, the MAgNet model, which requires separate interpolation for irregular grids, takes even more time than MP-PDE and GraphDeepONet. This illustrates that our GraphDeepONet model exhibits a trade-off, demonstrating a stable accuracy for irregular grids compared to MP-PDE, while requiring less time than MAgNet. 
experiments: 
 We conduct experiments comparing the proposed GraphDeepONet model with other benchmark models. Firstly, we explore the simulation of time-dependent PDEs by comparing the original DeepONet and VIDON with GraphDeepONet for regular and irregular sensor points. Specifically, we assess how well GraphDeepONet predicts in arbitrary positions, especially concerning irregular sensor points, compared to models such as MP-PDE, and MAgNet. Furthermore, we include FNO-2D, a well-established model known for operator learning, in our benchmark comparisons. Given the difficulty FNO faces in handling input functions with irregular sensor points, we also consider Factorized FNO (F-FNO) (Tran et al., 2023) , which extends FNO to irregular grids (See Appendix B.3). We consider the 1D Burgers' equation data from Brandstetter et al. (2022) , the 2D shallow water equation data from Takamoto et al. (2022) , and the 2D Navier-Stokes (N-S) equation data from Kovachki et al. (2021b) . For datasets with periodic boundaries, the GraphDeepONet leveraged the advantage of enforcing the condition (See Appendix B.6). The PyTorch Geometric library (Fey & Lenssen, 2019) is used for all experiments. The relative L 2 error by averaging the prediction solutions for all time is used for error estimate. See Appendix B for more details.Comparison with DeepONet and its variants The fourth and fifth columns in Comparison with GNN-based PDE-solvers The third, sixth, and seventh columns of Table 1 depict the accuracy of the FNO-2D and GNN-based models. While FNO outperformed the other models on a regular grid, unlike graph-based methods and our approach, it is not applicable to irregular sensor points, which is specifically designed for uniform grids. F-FNO also faces challenges when applied to the irregular grid. When compared to GNN-based models, with the exception of F-FNO, our model slightly outperformed MP-PDE and MAgNet, even on an irregular grid. Table 2 summarizes the results of our model along with other models, when applied to various irregular grids for 2D shallow water equation and 2D N-S equation, namely, Irregular I,II, and III for each equation.Remarkably, on one specific grid, MP-PDE outperformed our model. However, the MP-PDE has a significant inconsistency in the predicted performance. In contrast, our model consistently demonstrated high predictive accuracy across all grid cases. This is because, unlike other methods, the solution obtained through GraphDeepONet is continuous in the spatial domain. Figure 3 displays the time-evolution predictions of models trained on the shallow water equation for an initial condition. The GNN-based models are trained on fixed irregular sensors as seen in the second column and are only capable of predicting on the same grid, necessitating interpolation for prediction. In contrast, GraphDeepONet leverages the trunk net, enabling predictions at arbitrary grids, resulting in more accurate predictions. 
conclusion and discussion: 
 The proposed GraphDeepONet represents a significant advancement in the realm of PDE solution prediction. Its unique incorporation of time information through a GNN in the branch net allows for precise time extrapolation, a task that has long challenged traditional DeepONet and its variants. Additionally, our method outperforms other graph-based PDE solvers, particularly on irregular grids, providing continuous spatial solutions. Furthermore, GraphDeepONet offers theoretical assurance, demonstrating its universal capability to approximate continuous operators across arbitrary time intervals. Altogether, these innovations position GraphDeepONet as a powerful and versatile tool for solving PDEs, especially in scenarios involving irregular grids. While our GraphDeepONet model has demonstrated promising results, one notable limitation is its current performance on regular grids, where it is outperformed by FNO. Addressing this performance gap on regular grids remains an area for future improvement. As we have employed the temporal bundling method in our approach, one of our future endeavors includes exploring other techniques utilized in DeepONetrelated models and GNN-based PDE solver models to incorporate them into our model. Furthermore, exploring the extension of GraphDeepONet to handle more complex 2D time-dependent PDEs or the Navier-Stokes equations, could provide valuable insights for future works and applications. 
 
 Appendix: The notations in the paper is summarized in Table 3 . i (i = 0, 1, ..., N -1)the N -fixed sensor point in the spatial domain ∆t the discretized time K frame + 1 the number of frames in one solution trajectory K the number of grouping frames for temporal bundling method u k (x) (k = 0, 1, ..., K frame ) the solution at time t = k∆t ūk (x) (k = 0, 1, ..., K frame ) the values of solution at time t = k∆t in fixed sensor points u k (x) (k = 0, 1, ..., K frame ) the approximated solution at time t = k∆t G (k) the operator from the initial condition to the solution at time k∆t G GDON the approximated operator using GraphDeepONet G graph the approximated operator using other graph-based PDE solver p the number of basis (or coefficients) in DeepONet ν the branch net (or decoder) in DeepONet (or GraphDeepONet) τ the trunk net in DeepONet (or GraphDeepONet) ϵ the encoder in GraphDeepONet ϕ, ψ the neural networks of processor in GraphDeepONet ω the neural network of decoder in GraphDeepONet f i (i = 0, 1, ..., N -1) the feature vector at node iThe edges (i, j) ∈ E are constructed based on the proximity of node positions, connecting nodes within a specified distance. In actual experiments, we considered nodes as grids with given initial conditions. There are broadly two methods for defining edges. One approach involves setting a threshold based on the distances between grids in the domain, connecting edges if the distance between these grids is either greater or smaller than the specified threshold value. Another method involves utilizing classification techniques, such as the k-nearest neighbors (k-NN) algorithm, to determine whether to establish an edge connection. We determined whether to connect edges based on the k-NN algorithm with k =6 for 1D, k = 8 for 2D. Therefore, the processing of ϕ and ψ takes place based on these edges. The crucial point here is that once the Graph G = (V, E) is constructed according to a predetermined criterion, even with a different set of sensor points, ϕ and ψ remain unchanged as processor networks applied to the respective nodes and their connecting edges.Similar to other graph-based PDE solver studies (Brandstetter et al., 2022; Boussif et al., 2022) , we consider the 1D Burgers' equation asEQUATIONwhere δ(t, x) is randomly generated asδ(t, x) = 5 j=1 A j sin(a j t + b j x + ϕ j ) (8)Figure 4 : Three types of irregular grid (Irregular I, Irregular II, and Irregular III) used to train the models in shallow water eqaution where a j , b j and ϕ j are uniformly sampled asEQUATION)We conducted a direct comparison with the models using the data E1, E2, and E3 as provided in Brandstetter et al. (2022) ; Boussif et al. (2022) . For a more detailed understanding of the data, refer to those studies.Also, we take the 2D shallow water equation data from Takamoto et al. (2022) . The shallow water equations, which stem from the general Navier-Stokes equations, provide a suitable framework for the modeling of free-surface flow problems. In two dimensions, it can be expressed as the following system of hyperbolic PDEsEQUATIONwhere h(t, x, y) is the height of water with horizontal and vertical velocity (u, v) and g is the gravitational acceleration. We generate the random samples of initial conditions similar to the setting of Takamoto et al. (2022) . The initial condition is generated byh 0 (x, y) = 2.0, for r < x 2 + y 2 1.0, for r ≥ x 2 + y 2 (11)where the radius r is uniformly sampled from [0.3, 0.7].We utilize the same Navier-Stokes data employed in Li et al. (2020) . The dynamics of a viscous fluid are described by the Navier-Stokes equation. In the vorticity formulation, the incompressible Navier-Stokes equation on the unit torus can be represented as follows:EQUATIONHere, w, u, ν, and f represent the vorticity, velocity field, viscosity, and external force, respectively. 2023), which is state-of-the-art on regular grids and applicable to irregular grids. As shown in Table 2 , the F-FNO is applicable to irregular grids data, but it generally exhibits higher errors compared to GraphDeepONet. This is attributed to the limited capacity for the number of input features. We used the F-FNO model, which is built for Point Cloud data, to predict how solutions will evolve over time. At first, the model was designed to process dozens of input features, which made it difficult to include all the initial values from a two-dimensional grid.We trained various models, including DeepONet and VIDON, following the architecture and sizes as well as the training hyperparameters outlined in Prasthofer et al. (2022) . Additionally, MP-PDE and MAgNet utilized parameter settings as provided in Boussif et al. (2022) without modification. We trained our model, the GraphDeepONet, using the Adam optimizer, starting with an initial learning rate of 0.0005. This learning rate is reduced by 20In the small architecture, the encoder was set up with a width of 128 and a depth of 2 for epsilon. The processor components, ϕ, and ψ, each had a width of 128 and a depth of 2. We employed distinct ϕ and ψ for each of the three message-passing steps. In the decoder, we assigned ω gate and ω feature for aggregation to the neural network, which had a width of 128 and a depth of 3. The trunk net, τ , was configured with a width of 128 and a depth of 3.For the large architecture, the width of all neural networks was set to 128, and the depth was set to 3, except for the trunk net. The trunk net's depth was set to 5. The number of message-passing steps was set to 3. For more specific details, refer to the code. .. , naturally leading to a solution u(t, x) (x ∈ Ω) that satisfies the |Ω|-periodicity. As depicted in Figure 5 , the results reveal that while other models fail to perfectly match the periodic boundary conditions, GraphDeepONet successfully aligns with the boundary conditions.While our experiment primarily focuses on periodic boundary conditions, it is feasible to handle Dirichlet boundaries as well using ansatz extension as discussed in Choudhary et al. (2020) ; Horie & Mitsume (2022) . If we aim to enforce the solution u(t, x) = g(x) at x ∈ ∂Ω, we can construct the following solution:u(t, x) = g(x) + l(x) p j=1 ν j [f 1 0:N -1 , ∆t]τ j (x)where l(x) satisfies l(x) = 0, x ∈ ∂Ω, l(x) > 0, others.By constructing g(x) and l(x) appropriately, as described, we can effectively enforce Dirichlet boundary conditions as well. While the expressivity of the solution using neural networks may be somewhat reduced, there is a trade-off between enforcing boundaries and expressivity.For Burgers' equation, we generate the uniform grid of 50 points in [0, 16]. We divided the time interval from 0 to 4 seconds uniformly to create 250 time steps. We started with 25 initial values for each segment, then predicted the values for the next 25 instances, and so on. The total number of For the 2D shallow water equation, we generate the grid of 1024 = 32 2 points for the regular setting.For irregular data, we selected an equal number of points from a uniform distribution over 128 2 points within the rectangle [-2.5, 2.5] 2 and made predictions on a fixed grid. Figure 4 illustrates how we set up irregular sensor points for training GNN-based models and our model.We evenly divided the time interval from 0 to 1 second uniformly to create 101 time steps. We started with 10 initial values for each segment, then predicted the values for the next 10 instances, and so on. The total number of prediction steps is 9, calculated by dividing 101-1=100 by 10. We remark that the values at t = 1 were excluded from the data set. In all experiments, we used a batch size of 4. For both regular data and irregular data, the training data consisted of 600 samples, while both the validation and test samples contained 200 samples each. Note that the MAgNet has the capability to interpolate values using the neural implicit neural representation technique. However, we did not utilize this technique when generating Figure 3 , which assesses the interpolation ability for irregular data. For clarity, we've provided Figure 7 the predictions on the original irregular grid prior to interpolation.In reference to the 2D Navier-Stokes equation, we apply the data from Li et al. (2020) with a viscosity of 0.001. For regular data, we generate the grid of 1024 = 32 2 points. For irregular data, we selected an equal number of points from a uniform distribution over 64 2 points within the rectangle [0, 1] 2 and made predictions on a fixed grid. We evenly divided the time interval from 1 to 50 seconds uniformly to create 50 time steps. We started with 10 initial values for each segment, then predicted the values for the next 10 instances, and so on. The total number of prediction steps is 4, calculated by dividing 50-10=40 by 10. In all experiments, we used a batch size of 4. For both regular data and irregular data, the training data consisted of 600 samples, while both the validation and test samples contained 200 samples each. 


Decision: This paper presents a GNN method on surrogate models. The paper needs to clarify on the motivation, which can be address during the camera ready version.

Reviews: review: title: title: Nice work, but some details on the relevance of the work are missing
main_review: main_review: in the introduction, the authors say "Message Passing Neural PDE Solver (MO-PDE) ... a limitation of their approach is that it can only predict the solution ... on the same grid used as input"

The authors should better clarify:
1. why this is a problem for practical simulation applications
2. hHow the GraphDeepONet proposed approach addresses and solves this limitation.