Conditionally Parameterized, Discretization-Aware Neural Networks for Mesh-Based Modeling of Physical Systems

1 Introduction . Numerical simulations of partial differential equations ( PDEs ) have become an indispensable tool in the study of complex physical systems . High-resolution simulations are , however , prohibitively expensive or intractable in many practical problems . Machine learning techniques have recently been explored to improve the efficiency and accuracy of traditional numerical methods . Successful applications include nonlinear model order reduction [ 1 , 2 , 3 ] , model augmentation [ 4 , 5 , 6 ] , and super-resolution [ 7 , 8 , 9 ] . Neural networks have also been used to replace traditional PDE-based solvers , and serve as a standalone prediction tool . Popular approaches include auto-regressive timeseries predictions [ 10 , 11 , 12 , 13 , 14 ] , Physics-Informed Neural Networks ( PINNs ) [ 15 , 16 , 17 ] . Despite promising results on canonical problems , commonly used network architectures such as autoencoders and CNNs have inherent limitations . An autoencoder generates a fixed mapping between the geometric coordinates and the encoded digits . This limits their portability for new geometries and dynamic patterns . A CNN requires interpolation of existing data to a structured , Euclidean space , introducing additional cost and error . Irregular geometry boundaries require constructs such as elliptic coordinate transformation [ 18 ] and Signed Distance Function ( SDF ) [ 1 ] . Moreover , models often ignore the hierarchical relations between heterogeneous features , and concatenate them into a single input vector , e.g . the common concatenation of the edge and node features in Graph 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Neural Networks ( GNNs ) . The learning of high-order terms remains mostly unguided – even simple quadratic terms are often fitted via a number of hidden units in a brute-force manner . With a focus on mesh-based modeling of physical systems , we use the idea of conditional parameterization ( CP ) to build the hierarchical relations between different physical quantities as well as numerical discretization information into the network architectures . The key contributions of our work are as follows 1 : 1 . We demonstrate that a drop-in CP modification can bring significant improvements for various existing models on several tasks essential to the modeling of physical systems . 2 . We propose a conditionally parameterized graph neural network ( CP-GNet ) , which effectively models complex physics such as chemical source terms , irregular mesh discretizations , and different types of boundary conditions . 3 . We conduct extensive numerical tests and demonstrate state-of-the-art performances on problems of different complexities , ranging from the basic viscous Burgers equation to a complex reacting flow . 2 Methodology . Conditional Parametrization : The idea of conditional parametrization ( CP ) is to use trainable functions of input parameters to generate the weights of a neural network . To demonstrate this , we start from a standard dense ( fully connected ) layer : h ( u ; W , b ) = σ ( Wu + b ) , ( 1 ) where u ∈ Rnx is the input feature vector , h ∈ Rnh is the output hidden state vector , W ∈ Rnh×nx and b ∈ Rnh are the trainable weights and bias , and σ is the activation function . It can be seen that in the evaluation stage , the values of W and b are fixed regardless of the inputs . Thus the performance of Eq . ( 1 ) is largely limited by the interpolation range of training data . By introducing a parameter vector p ∈ Rnp and a trainable function f ( p ) : Rnp → Rnh×nx that computes the weights W based on p , the conditionally parameterized version of Eq . ( 1 ) is given by : h ( u ; f ( p ) , b ) = σ ( f ( p ) u + b ) . ( 2 ) An easy way to incorporate the formulation into existing neural network models is by making f a single-layer MLP , the conditionally parameterized dense ( CP-Dense ) layer can be represented by : h ( u , p ; W , B , b ) = σ ( σ ( 〈W , p〉+ B ) u + b ) . ( 3 ) It should be noted that this would bring a change in the dimensions of weights and biases , which become W ∈ R ( nh×nu ) ×np , B ∈ Rnh×nu . When the layer width is kept the same , the total number of trainable parameters increases linearly with the parameter size np . In applications , p is not limited to an additionally-introduced parameter . When simply taking u as the parameter for itself , the quadratic terms will be introduced . High-order terms , which are prevalent in physical systems , can be easily modeled using multiple such layers . In Appendix B , we demonstrate how certain discretized PDE terms can be fitted exactly with simple conditionally parameterized layers . 2.1 CP-GNet for mesh-based modeling of physical systems . Graph representation of discretized systems : Consider a physical system governed by a set PDEs for a time-variant vector of variables q ( t ) . Using the popular finite volume discretization , the computational domain is divided into contiguous small cells , indexed by i . The discretized form of equation can be written as : dqi ( t ) dt = 1 Ωi ∑ j∈N ( i ) f ( qi , qj , nij ) Aij + s ( qi ) , ( 4 ) 1The source code is released to facilitate future research at https : //github.com/ davidxujiayang/cpnets where qi is the cell-centered value of cell i , Ωi is the volume ( 3D ) /area ( 2D ) of the cell , and N ( i ) is the neighborhood set of cells around i . Between a neighboring pair of cells i and j , Aij is area ( 3D ) /length ( 2D ) of the shared cell boundary , and nij = ( xi − xj ) /|xi − xj | is a vector between the cell center locations xi and xj . In the explicit numerical simulation of Eq . ( 2.1 ) , solutions are updated by computing the increment of ∆qki = q k+1 i − qki between discrete time steps indexed by k , which is determined by two terms . The flux term f computes the exchange of quantity between neighboring cells , which is a complex function involving both the cell values as well as the vector between them , e.g . [ 19 ] . The source term s computes physics that are local to the cell , such as the reaction of chemical species . In our setting , the discretized system is mapped to a graph G ( V , E ) , defined by nodes V of size |V | = nv connected by edges E ⊂ V × V of size |E| = ne . Each node vi is located at the corresponding cell center xi , and each edge ( i , j ) corresponds to a shared boundary between the finite volume cells . Denoting the sets of mapped quantities on all nodes and edges of G , Q = { qi , i ∈ V } , N = { nij , ( i , j ) ∈ E } , the target is to develop a graph neural network operator g that predicts the increment as ∆Qk = g ( Qk , N ) . CP-GNet architecture : The architecture for the proposed conditionally parameterized graph neural network , CP-GNet , can be written in an encoder-processor-decoder form . A schematic is provided in Fig . 1 . For clarity of different variables in the description of the network , we use ui for the latent variables on node i to distinguish from the physical variables qi , and use eij for latent variables on edge ( i , j ) to distinguish from the vector nij . Encoder : Numerical solution of PDEs ( e.g. , the compressible Navier–Stokes equations ) requires the processing of arbitrarily complex interactions between mesh elements . While large MLP architectures can represent this complexity , the data requirements to reliably train such networks might be large . In contrast , our proposed encoder takes two CP-Dense layers , taking the output from the previous layer as both the input and the conditional parameter . Through the encoder , high-order interactions can be easily extracted , allowing a degree of extrapolation by virtue of linearity . The CP-GNet uses two separate , but similarly constructed encoders to process the input node features qki and edge features nij , respectively . Processor : The flux term f in Eq . ( 2.1 ) can be effectively approximated by CP message-passing ( CP-MP ) between adjacent nodes on a graph . The source term s , on the other hand , can be modeled by CP-Dense layers . In CP-GNet the processor consists of multiple identical blocks with independent weights . Residual connections are added between the blocks . As shown in Fig . 1 , each block includes a CP-MP based section and a CP-Dense based section to address the two types of terms . Modified from the Edge Conditioned Convolution ( ECC ) [ 20 ] , the CP-MP computation is formulated as : Wij = σ ( 〈 W , eφij 〉 + B ) , ; hi = ∑ j∈N ( i ) wijσ ( 〈Wij , [ ui ; uj ] 〉 ) , ( 5 ) where eφij is the output from the edge encoder , ui and uj are the latent node features from the previous layer , hi is the nodal latent output , and wij = Aij/Ωi is the flux weight from Eq . ( 2.1 ) . Decoder : It is common in a PDE solver to use a Jacobian matrix J = ∂u/∂h to transform the variable increments ∆u = J∆h . The decoder in the GP-GNet serves a similar purpose – to convert hidden variables to the output on the physical space . Similar to the encoder , the decoder consists of three conditionally parameterized dense layers . The first two layers can actually be viewed as a dedicated “ encoder ” that is similar to the initial node encoder , taking qki as the input , but with independent weights . The purpose of this “ encoder ” is to extract a final conditional parameter , which is used in the third CP-Dense layer in the decoder to determine the weights for the output node feature uχi from the processor . The third decoder layer is also the final layer of the model , which outputs ∆qki ( with proper scaling ) . Except for the edge encoder and the last two layers in the decoder , all dense , CP-Dense , CP-MP layers are appended with LayerNormalization ( LN ) layers . Treatments for boundaries : The computational domain of a practical problem includes multiple types of boundaries , e.g . the case in Sec . 4.3 . In classic PDE solvers , they are treated with different boundary conditions , which define explicit formulations to compute relationships of the domain with the external world . However , these conditions and formulations are only defined for the physical quantities , thus can not be easily transferred for latent variables , especially when multiple messagepassing/convolution steps are used . To enable the GP-GNet to model different types of boundaries efficiently , special treatments are necessary . For boundaries with known inputs , such as the inlet and the outlet , their values are directly input to the corresponding nodes at every time step . For the boundaries imposing certain constraints , instead of a given physical value , such as Neumann and symmetry boundaries , ghost edges are introduced . For a cell i with a face lying on a boundary , we introduce a ghost edge vector nig , that points from the corresponding node i to the center of the boundary face . Ghost edges are processed together with the normal edges in the edge encoder . However , the CP-MP layer in the processor of the CP-GNet is slightly modified . More specifically , the concatenation [ ui ; uj ] in Eq . ( 5 ) is replaced with only ui . And for each type of boundary , the weights for the CP-MP layer are trained independently , to let the model learn different types of boundary condition for the latent variables . The effectiveness of this treatment is shown in Sec . 4.3 and further discussed in Appendix . A.3.3 .

Summary: This paper proposes methods for neural simulation of physical systems on graphs. In particular it proposes the use of hyper-network style layers (CP-dense) where the weights of the layer are determined by regressing some context parameter. These layers are combined in the CP-GNet architecture, which shares a similar encoder / process / decode structure to MeshGraphNets. The proposed layers are shown to yield performance improvements on simulation tasks relative to dense layers.