Composing Partial Differential Equations with Physics-Aware Neural Networks

1 INTRODUCTION . Artificial neural networks ( ANNs ) are considered universal function approximators ( Cybenko , 1989 ) . Their effective learning ability , however , greatly depends on domain and task-specific prestructuring and methodological modifications referred to as inductive biases ( Battaglia et al. , 2018 ) . Typically , inductive biases limit the space of possible models by reducing the opportunities for computational shortcuts , which can lead to erroneous implications derived from a potentially limited dataset ( overfitting ) . The recently evolving field of physics-informed machine learning employs physical knowledge as inductive bias providing huge generalization advantages in contrast to pure machine learning ( ML ) in physical domains ( Raissi et al. , 2019 ) . While numerous approaches have been introduced to augment ANNs with physical knowledge , these methods either do not allow the incorporation of explicitly defined physical equations ( Long et al. , 2018 ; Seo et al. , 2019 ; Guen & Thome , 2020 ; Li et al. , 2020a ; Sitzmann et al. , 2020 ) or can not generalize to other initial and boundary conditions than those encountered during training ( Raissi et al. , 2019 ) . In this work , we present the finite volume neural network ( FINN ) model—a physics-aware neural network structure adhering to the idea of spatial and temporal discretization in numerical simulation . FINN consists of multiple neural network modules that interact in a distributed , compositional manner ( Battaglia et al. , 2018 ; Lake et al. , 2017 ; Lake , 2019 ) . The modules are designed to account for specific parts of advection-diffusion equations , a class of partial differential equations ( PDEs ) . This modularization allows to combine two advantages that are not yet met by state-of-the-art models : the explicit incorporation of physical knowledge and the generalization over initial and boundary conditions . To the best of our knowledge , FINN ’ s ability to adjust to different initial and boundary conditions and to explicitly learn constitutive relationships and reaction terms is unique , yielding excellent out-of-distribution generalization . The core contributions of this work are : • Introduction of FINN , a physics-aware neural network model , explicitly designed to generalize over initial and boundary conditions , demonstrating excellent generalization ability . • Evaluation of state-of-the-art pure ML and physics-aware models , contrasted to FINN on one- and two-dimensional benchmarks , demonstrating benefits of explicit model design . • Application of FINN to a real-world contamination-diffusion problem , verifying its applicability to real , spatially and temporally constrained training data . 2 RELATED WORK . Non-physics-aware ANN architectures Pure ML models that are designed for spatiotemporal data processing can be separated into temporal convolution ( TCN , Kalchbrenner et al. , 2016 ) and recurrent neural networks . While the former perform convolutions over space and time , representatives of the latter , e.g. , convolutional LSTM ( ConvLSTM , Shi et al. , 2015 ) or DISTANA ( Karlbauer et al. , 2019 ) , aggregate spatial neighbor information to further process the temporal component with recurrent units . Since pure ML models do not adhere to physical principles , they require enormous amounts of training data and parameters in order to approximate a desired physical process ; but still are not guaranteed to behave consistently outside the regime of the training data . Physics-aware ANN architectures When designed to satisfy physical equations , ANNs are reported to have greater robustness in terms of physical plausibility . For example , the physicsinformed neural network ( PINN , Raissi et al. , 2019 ) consists of an MLP that satisfies an explicitly defined PDE with specific initial and boundary conditions , using automatic differentiation . However , the remarkable results beyond the time steps encountered during training are limited to the very particular PDE and its conditions . A trained PINN can actually not be applied to other different initial and boundary conditions , which limits its applicability in real-world scenarios . Other methods by Long et al . ( PDENet , 2018 ) , Guen & Thome ( PhyDNet , 2020 ) , or Sitzmann et al . ( SIREN , 2020 ) learn the first n derivatives to achieve a physically plausible behavior , but lack the option to include physical equations . The same limitation applies when operators are learned to approximate PDEs ( Li et al. , 2020a ; b ) , or when physics-aware graph neural networks are applied ( Seo et al. , 2019 ) . Yin et al . ( APHYNITY , 2020 ) propose to approximate equations with an appropriate physical model and to augment the result by an ANN , preventing the ANN to approximate a distinct part within the physical model . For more comparison to related work , please refer to subsection A.1 . In summary , none of the above methods can explicitly learn particular constitutive relationships or reaction terms while simultaneously generalizing beyond different initial and boundary conditions . 3 FINITE VOLUME NEURAL NETWORK ( FINN ) . Problem formulation Here , we focus on modeling spatiotemporal physical processes . Specifically , we consider systems governed by advection-diffusion type equations ( Smolarkiewicz , 1983 ) : ∂u ∂t = D ( u ) ∂2u ∂x2 − v ( u ) ∂u ∂x + q ( u ) , ( 1 ) where u is the quantity of interest , t is time , x is the spatial coordinate , D is the diffusion coefficient , v is the advection velocity , and q is the source/sink term . Eq . 1 can be partitioned into three parts : the storage term , the flux terms , and the source/sink term . The storage term ∂u∂t describes the change of the quantity u over time . The flux terms are the advective flux v ( u ) ∂u∂x and the diffusive flux D ( u ) ∂ 2u ∂x2 . Both calculate the amount of u exchanged between neighboring positions . The source/sink term q ( u ) describes the generation or elimination of the quantity u. Eq . 1 is a general form of PDE with up to second order spatial derivative , but it has a wide range of applicability due to the flexibility of defining D ( u ) , v ( u ) , and q ( u ) as different functions of u , as is shown by the numerical experiments in this work . The finite volume method ( FVM , Moukalled et al. , 2016 ) discretizes a simulation domain into control volumes ( i = 1 , . . . , Nx ) , where exchange fluxes are calculated using a surface integral ( Riley et al. , 2009 ) . In order to match this structure in FINN , we introduce two different kernels , which are ( spatially ) casted across the discretized control volumes : the flux kernel , modeling the flux terms , and the state kernel , modeling the source/sink term as well as the storage term . The overall FINN architecture is shown in Figure 1 . Flux kernel The flux kernel F approximates the surface integral for each control volume i with boundary Ω by a composition of multiple subkernels fj , each representing the flux through a discretized surface element j : Fi = Nsi∑ j=1 fj ≈ ∮ ω⊆Ω ( D ( u ) ∂2u ∂x2 − v ( u ) ∂u ∂x ) · n̂ dΓ , ( 2 ) where Nsi is the number of discrete surface elements of control volume i , ω is a continuous surface element ( a subset of Ω ) , fj are subkernels ( which are realized as feedforward network modules ) , and n̂ is the unit normal vector pointing outwards of ω . In our exemplary one-dimensional arrangement , two subkernels fi− and fi+ ( see Figure 1 ) contain the modules ϕD , ϕA , and ϕN ( ϕNi− or ϕNi+ for fi− and fi+ , respectively ) . The module ϕN is a linear layer with the purpose to learn the numerical FVM stencil . More specifically , the module ϕNi− takes the inputs ui and its neighbor ui−1 , while ϕNi+ takes the inputs ui and ui+1 , and output the approximation of the spatial derivative ∂u ∂x . This signifies that the weight of ϕN should amount to [ −1 , 1 ] with respect to [ ui , ui−1 ] and [ ui , ui+1 ] ( i.e . simple difference between neighboring control volumes ) in ideal one-dimensional problems . The other modules ϕA and ϕD are responsible for advective and diffusive fluxes , respectively . Both modules receive only ui as input . Moreover , the moduleR applies an upwind differencing scheme , which prevents numerical instability in the first order spatial derivative calculation ( Versteeg & Malalasekera , 1995 ) . It ensures that calculation of the advective fluxes using ϕA is performed only on one control volume surface ( either left or right , and not both at the same time ) , whereas calculation of the diffusive fluxes using ϕD is performed on both surfaces . This means , for the advective flux calculation , the summation of the numerical stencil from both fi− and fi+ will only lead to [ −1 , 1 ] , which corresponds to either [ ui , ui−1 ] when ϕA > 0 or [ ui , ui+1 ] when ϕA < 0 ( i.e . only first order spatial derivative ) . For the diffusive flux calculation , the summation of the numerical stencil will lead to the classical one-dimensional numerical Laplacian with [ 1 , −2 , 1 ] corresponding to [ ui−1 , ui , ui+1 ] , that is , the second order spatial derivative , because the calculation of ϕD is performed on both surfaces ( see subsection A.3 for more details ) . Accordingly , module R generates the inductive bias to make ϕA only approximate the advective , and ϕD the diffusive flux . For the advective flux calculation , if the velocity v is a function of u , the module ϕA approximates this dependence using a feedforward neural network , i.e . ϕA ( u ) ≈ v ( u ) . Otherwise , ϕA is a scalar value ϕA ≡ v , which can be also set as a learnable parameter . The output of ϕA is activated by the following case-sensitive ReLUR : R ( ϕA ) = { ReLU ( ϕA ) , on fi− , −ReLU ( −ϕA ) , on fi+ , ( 3 ) and then multiplied with the output of ϕN ( i.e . ∂u∂x ) to obtain the corresponding advective flux at the surface . Similarly for the diffusive fluxes , if the diffusion coefficient D depends on u , the module ϕD takes u as an input to learn the function ϕD ( u ) ≈ D ( u ) using a feedforward neural network . Otherwise , ϕD is a scalar value ϕD ≡ D , which can also be set as a learnable parameter . The output of ϕD is not activated by the module R , but is directly multiplied with the output of ϕN to obtain the corresponding diffusive flux at the surface . Consequently , integration of subkernels fj in Eq . 2 and substitution of Eq . 3 leads to the following representation : Fi = { ϕN ( ui , ui−1 ) [ ϕD − ϕA ] fi− + ϕN ( ui , ui+1 ) [ ϕD ] fi+ , if ϕA > 0 , ϕN ( ui , ui−1 ) [ ϕD ] fi− + ϕN ( ui , ui+1 ) [ ϕD − ϕA ] fi+ , otherwise . ( 4 ) Boundary conditions A means of applying boundary conditions in the model is essential when solving PDEs . Currently available models mostly adopt a convolutional structure for modelling spatiotemporal processes . However , a convolutional structure only allows a constant value to be padded at the domain boundaries ( e.g . zero-padding or mirror-padding ) , which is only appropriate for the implementation of Dirichlet or periodic boundary condition types . Other types of frequently used boundary conditions are Neumann and Cauchy . They are defined as a derivative of the quantity of interest , and hence can not be easily implemented in convolutional models . With certain predefined boundary condition types , the flux kernels at the boundaries are adjusted accordingly to allow for straightforward boundary condition implementation . For Dirichlet boundary condition , a constant value u = ub is set as the input ui−1 ( for the flux kernel fi− ) or ui+1 ( for fi+ ) at the corresponding boundary . For Neumann boundary condition ν , the output of the flux kernel fi− or fi+ at the corresponding boundary is set to be equal to ν . With Cauchy boundary condition , the solution-dependent derivative is calculated and set as ui−1 or ui+1 at the corresponding boundary . State kernel The state kernel S calculates the source/sink and storage terms of Eq . 1 . The source/sink ( if required ) is learned using the module Φψ ≈ q ( u ) , which takes u as input . The storage , ∂u∂t , is then calculated using the output of the flux kernel and module Φψ of the state kernel : Si = Fi + Φψ ( ui ) ≈ ∂ui ∂t . ( 5 ) By doing so , the PDE in Eq . 1 is now reduced to a system of coupled ordinary differential equations ( ODEs ) , which are functions of ui , ui−1 , ui+1 , and t. Thus , the solutions of the coupled ODE system can be found using a numerical integration over time . Since first order explicit approaches , such as the Euler method ( Butcher , 2008 ) , suffer from numerical instability ( Courant et al. , 1967 ; Isaacson & Keller , 1994 ) , we employ the neural ordinary differential equation method ( Neural ODE , Chen et al. , 2018 ) to reduce numerical instability via the Runge-Kutta adaptive time-stepping strategy . When using the Neural ODE method , the predicted values of u at time t are fed back into the network as direct inputs to predict u at time t + 1 . Therefore , the entire training is performed in closed-loop , improving stability and accuracy of the prediction compared to networks trained with teacher forcing , i.e . one-step-ahead prediction ( Praditia et al. , 2020 ) . The weight update is realized by applying backpropagation through time ( indicated by red arrows in Figure 1 ) . In short , FINN takes only the initial condition u at time t = 0 and propagates the dynamics forward .

Summary: This paper presents a compositional physics-aware neural network (FINN) for learning spatiotemporal advection-diffusion processes. It claims that the FINN outperforms pure machine learning and other state-of-the-art physics-aware models in all cases—often even by multiple orders of magnitude. However, the design of the network depends too much on the form of the equation, which leads to a very narrow application of the method.