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: The authors propose to model advection-diffusion partial differential equations as a composition of multiple neural networks. According to the authors, this leads to better generalization over different initial and boundary conditions for the physics system, and also the ability to learn different factors of the process as modeled by the terms in the PDE. Extensive experimental results are presented to support the author’s claims that their framework performs better than state-of-the-art.