Time integration schemes based on neural networks for solving partial
  differential equations on coarse grids
The accuracy of solving partial differential equations (PDEs) on coarse grids
is greatly affected by the choice of discretization schemes. In this work, we
propose to learn time integration schemes based on neural networks which
satisfy three distinct sets of mathematical constraints, i.e., unconstrained,
semi-constrained with the root condition, and fully-constrained with both root
and consistency conditions. We focus on the learning of 3-step linear multistep
methods, which we subsequently applied to solve three model PDEs, i.e., the
one-dimensional heat equation, the one-dimensional wave equation, and the
one-dimensional Burgers' equation. The results show that the prediction error
of the learned fully-constrained scheme is close to that of the Runge-Kutta
method and Adams-Bashforth method. Compared to the traditional methods, the
learned unconstrained and semi-constrained schemes significantly reduce the
prediction error on coarse grids. On a grid that is 4 times coarser than the
reference grid, the mean square error shows a reduction of up to an order of
magnitude for some of the heat equation cases, and a substantial improvement in
phase prediction for the wave equation. On a 32 times coarser grid, the mean
square error for the Burgers' equation can be reduced by up to 35% to 40%.
Xinxin Yan, Zhideng Zhou, Xiaohan Cheng, Xiaolei Yang
Time integration schemes based on neural networks for solving partial differential equations on coarse grids

###### Abstract

The accuracy of solving partial differential equations (PDEs) on coarse grids is greatly affected by the choice of discretization schemes. In this work, we propose to learn time integration schemes based on neural networks which satisfy three distinct sets of mathematical constraints, i.e., unconstrained, semi-constrained with the root condition, and fully-constrained with both root and consistency conditions. We focus on the learning of 3-step linear multistep methods, which we subsequently applied to solve three model PDEs, i.e., the one-dimensional heat equation, the one-dimensional wave equation, and the one-dimensional Burgers' equation. The results show that the prediction error of the learned fully-constrained scheme is close to that of the Runge-Kutta method and Adams-Bashforth method. Compared to the traditional methods, the learned unconstrained and semi-constrained schemes significantly reduce the prediction error on coarse grids. On a grid that is \(4\times\) coarser than the reference grid, the mean square error shows a reduction of up to an order of magnitude for some of the heat equation cases, and a substantial improvement in phase prediction for the wave equation. On a \(32\times\) coarser grid, the mean square error for the Burgers' equation can be reduced by up to \(35\%\) to \(40\%\).

**Keywords:** Time integration scheme, Partial differential equation, Numerical simulation on coarse grids, Neural networks, Mathematical constraints

## 1 Introduction

Many environmental and engineering problems, encompassing a broad spectrum of spatial and temporal scales, are described by partial differential equations (PDEs) of high dimension. Since resolving all the scales often demands exceedingly extensive computational resources, solving the high-dimensional PDEs poses a great challenge to the current computing systems. For example, direct numerical simulations of high-Reynolds number turbulent flows can be particularly challenging in this regard. Solving the PDEs on coarse grids, reducing the dimension of the problem, can substantially reduce the computational costs, while meanwhile introduces discretization errors. Discretization errors depend not only on the spatial discretization schemes but also on the time integration schemes. In this study, we propose to learn time integration schemes to reduce the error of solving PDEs on coarse grids.

The traditional approach to reduce the prediction error on coarse grid is to directly account for the effect of the physics of the unresolved scales. One way is to solve the spatially filtered PDEs instead of the original PDE. However, the filtering procedure on the nonlinear term often introduces a new unclored term, i.e., the subgrid term. The subgrid term can be modelled explicitly by establishing its relation with the resolved scales. For instance, the eddy viscosity model for large-eddy simulation (LES) of turbulent flows governed by the Navier-Stokes (NS) equations [1, 2]. The subgrid term can also be modelled implicitly by changing the properties of the schemes for discretizing the spatial derivatives. One example is the Implicit Large-Eddy Simulation (ILES) method for solving turbulent flows [3]. The success of such approaches largely depends on our understanding of the physics of the unresolved scales. Moreover, the developed models are often valid for only a range of unresolved scales. When the dominant physics of the unresolved scales change from one to the other as the cut-off scale changes, a grey region appears where neither of the models in
each regime works. For instance, the grey region in the microscale and mesoscale simulations of the atmospheric flow in the meteorology research [4, 5].

The machine learning methods are revolutionizing the numerical methods for solving PDEs [6, 7, 8, 9, 10, 11]. Such approaches approximate the solution based on machine learning models [12, 13], e.g., neural networks [14], avoiding the need for a grid to discretize the computational domain. In the review paper by [15], three types of the methods based on neural networks were reviewed: 1) the physics-informed neural networks (PINN) [6, 16, 17], in which the PDEs with initial and boundary conditions are approximately enforced via loss functions; 2) the methods based on the Feynman-Kac formula [18], which approximate the solution of a PDE as the expectation of a stochastic process; and 3) the methods based on the solution of backward stochastic differential equations [19, 20], in which deep neural networks (DNN) are employed for computing the gradient of solutions. As noted in the review [15], the last two methods showed promising results for high-dimensional linear and semilinear systems. The methods based on PINN are capable of handling complex nonlinear PDEs and inverse problems with incomplete models and imperfect data, while they are not competitive for solving well-posed, high-dimensional forward problems when compared with the traditional grid-based numerical methods [12, 15].

Improving the predictive capability of grid-based numerical methods using the machine learning methods received lots of attention as well [21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]. Some efforts are particularly focused on improving the traditional subgrid models for solving PDEs on coarse grids, such as models for turbulent flow simulations [21, 22, 23, 24, 25, 26, 27, 28]. Other attempts were made to improve the performance of the spatial discretization schemes for PDEs [29, 30, 31]. Particularly, the spatial discretization schemes for solving PDEs on coarse grids were learned using neural networks in the work by Bar-Sinai et al. [31]. The results of the one-dimensional Burgers' equation showed that the learned discretization schemes can reduce the prediction error and extend the median valid simulation time for very coarse grids.

The stability and accuracy of solving PDEs on coarse grids depend on both temporal and spatial discretization schemes. Empowering time integration schemes through machine learning methods, however, has been overlooked in existing studies. When solving PDEs on coarse grids, errors in the solution due to coarse spatial resolution accumulate over time. For example, this can lead to phase errors when solving the wave equation on coarse grids. In this study, we employ neural networks to learn time integration schemes to reduce the error of solving PDEs on coarse grids. Specifically, we propose three learning approaches with different constraints and test the learned schemes using the one-dimensional heat equation, wave equation, and Burgers' equation.

In the rest of the paper, the three learning approaches are first described and applied to learn a 3-step linear multistep method in section 2. Then, the test results for the one-dimension heat equation, wave equation, and Burgers' equation are presented sequentially in section 3. Lastly, the conclusions are drawn in section 4.

## 2 Learning of time integration schemes

In this section, we describe the approach for developing data-driven time integration schemes. The 3-step linear multistep method is taken as an example. Similar ideas can be applied to other time integration schemes. The general formulation of the linear multistep method is given in section 2.1 focusing on the constraints for the stability and consistency of the method. The three different models for imposing the constraints are then presented in section 2.2. The definition of the loss function is given in section 2.3. Lastly, the details for the application in a 3-step linear multistep method is given in section 2.4.

### General formulation of the linear multistep method

We consider the following ordinary differential equation (ODE) with independent variable \(t\)

\[\frac{\partial v}{\partial t}=F(x,t), \tag{1}\]

where \(x\) is a parameter (which can be the spatial coordinate). The linear multistep method of \(k\)-step for discretizing the above ODE can be expressed as

\[\alpha_{k}v^{n+1}=-\sum_{i=0}^{k-1}\alpha_{i}v^{n+i-k+1}+\Delta t \sum_{i=0}^{k}\beta_{i}F_{n+i-k+1}, \tag{2}\]

where \(\alpha_{i}\) and \(\beta_{i}\) are real coefficients, \(v^{n}\) represent the approximate value of \(v(t_{n})\) at \(t=t_{n}\), \(\Delta t\) is the size of time step and \(F_{n+i}=F(x,t_{n+i})\), \(t_{n+i}=t_{n}+i\Delta t\). For the sake of simplicity, we assume that \(\alpha_{k}=1\), \(\alpha_{0}^{2}+\beta_{0}^{2}\neq 0\).

The generating polynomials of the multistep method, which can be written as

\[\rho(\chi)=\alpha_{k}\chi^{k}+\alpha_{k-1}\chi^{k-1}+\cdots+ \alpha_{0}, \tag{3}\] \[\sigma(\chi)=\beta_{k}\chi^{k}+\beta_{k-1}\chi^{k-1}+\cdots+\beta _{0}, \tag{4}\]
are employed to obtain the consistency and stability constraints [32]. The obtained consistency condition is

\[\rho(1)=0,\quad\rho^{\prime}(1)=\sigma(1), \tag{5}\]

which is also the condition that the linear multistep method should satisfy to be of first-order accuracy 1.

Footnote 1: According to Taylor expansion, the necessary and sufficient condition for a linear multistep method of \(k\)-step to be of order \(P\) is \(\sum_{i=0}^{k}\alpha_{i}=0\), \(\sum_{i=0}^{k}\alpha_{i}i^{x}=s\sum_{i=0}^{k}\beta_{i}i^{x-1},s=1,2,\cdots P\).

The stability condition we employ here is called zero-stability, that a general method is stable when \(\Delta t\to 0\). For a linear multistep method, it is stable if \(\rho(\chi)\) satisfies the root condition [32], i.e., the moduli of the roots of \(\rho(\chi)\) are less than or equal to 1, and if one root's modulus is 1, the root must be single. That is to say, the roots of \(\rho(\chi)\) are on or within the unit circle, and the roots on the unit circle are single. It is worth mentioning that the convergence and stability of the linear multistep method are equivalent under the premise of the consistency condition.

### Three models for learning time integration schemes

The stability and consistency constraints discussed in the last subsection are imposed during the training of the neural network model, with the aim to increase the interpretability and stability of the learned time integration schemes. Three different strategies for applying the constraints are employed, i.e., 1) the BP (backpropagation) neural network model with no constraints (the unconstrained model, Un-con); 2) the BP neural network model with root condition (stability condition) enforced (the semi-constrained model, Semi-con); 3) the BP neural network model with both root condition and consistency condition enforced (the fully-constrained model, Full-con). In the following, the root condition and consistency condition will be reformulated in a way to facilitate their implementation as constraints for training the neural network models. As the results from the explicit Runge-Kutta method of order 3(2) with adaptive time-stepping will be employed as the reference for examining the accuracy, in the following we focus on deriving the constraints for the data-driven explicit linear multistep method of 3-step (\(k=3\) in Eq. (2)).

The consistency condition, which is equivalent to imposing the first-order constraint to a third-order explicit linear multistep method, is implemented in a way that some parameters are the direct outputs of the neural network with the rest of the parameters (i.e., \(\alpha_{i}\) and \(\beta_{i}\) in Eq. (2)) adjusted to satisfy the constraint (i.e., Eq. (5)).

In regard to the root condition (stability condition), i.e., the moduli of the roots of \(\rho(\chi)\) are less than or equal to 1 (\(\|\chi\|\leq 1\)), and if one root's modulus is 1, the root must be single, we employ the transform \(\chi=\frac{1+z}{1-z}\) to map the above constraint to the one that the real parts of all roots of a real coefficient polynomial with respect to \(z\) are negative. With \(Re(z)<0\), the original root condition is replaced by a stricter one, \(\|\chi\|<1\). Therefore, the \(\rho(\chi)\) with the root condition promised can be turned into a Hurwitz Polynomial \(\psi(z)\)[33]. Specifically, the Hurwitz Polynomials are in the following form,

\[\psi_{2}(z)=(1-z)^{2}\rho_{2}\left(\frac{1+z}{1-z}\right)=(1-\alpha_{1}+ \alpha_{0})z^{2}+2(1-\alpha_{0})z+(1+\alpha_{1}+\alpha_{0}), \tag{6}\]

\[\psi_{3}(z)=(1-z)^{3}\rho_{3}\left(\frac{1+z}{1-z}\right)=(1-\alpha_{2}+ \alpha_{1}-\alpha_{0})z^{3}+(3-\alpha_{2}-\alpha_{1}+3\alpha_{0})z^{2}+(3+ \alpha_{2}-\alpha_{1}-3\alpha_{0})z+(1+\alpha_{2}+\alpha_{1}+\alpha_{0}), \tag{7}\]

which correspond to the following quadratic and cubic \(\rho(\chi)\) polynomials \(\rho_{2}(\chi)=\chi^{2}+\alpha_{1}\chi+\alpha_{0}\) and \(\rho_{3}(\chi)=\chi^{3}+\alpha_{2}\chi^{2}+\alpha_{1}\chi+\alpha_{0}\), respectively.

According to the Routh-Hurwitz criterion [33], the constraint that the real parts of the roots of \(\psi_{2}(z)\) and \(\psi_{3}(z)\) are negative is equivalent to the condition that their polynomial coefficients satisfy the following inequalities,

\[\begin{cases}1-\alpha_{1}+\alpha_{0}>0,\\ 1-\alpha_{0}>0,\\ 1+\alpha_{1}+\alpha_{0}>0,\end{cases}\quad\text{(a)}\quad\quad\quad\quad \begin{cases}1-\alpha_{2}+\alpha_{1}-\alpha_{0}>0,\\ 1+\alpha_{2}+\alpha_{1}+\alpha_{0}>0,\\ 1-\alpha_{1}+\alpha_{2}\alpha_{0}-\alpha_{0}^{2}>0,\\ 1-\alpha_{0}>0,\\ 1+\alpha_{0}>0.\end{cases}\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad\quad \
The root condition is then reduced to that the rest of the two roots of \(\rho_{2}(\chi)\) must lie within the unit circle. The coefficients of \(\chi^{2}+p\chi+q\) should meet the Routh-Hurwitz criterion, i.e., Eq. (8a). With \(p\), \(q\) given by the learned neural network model, the values of \(\alpha_{0}\), \(\alpha_{1}\) and \(\alpha_{2}\) can then be obtained.

### Loss function

The loss function consists of two parts, i.e., the part related to the error of the solution predicted by the learned time integration scheme, and the other part for the constraints imposed on the scheme coefficients given by the neural network model. The first part of the loss function (\(\text{MSE}(v_{pred},v_{true})\)) is defined as the mean squared error between the predictions on the coarse grid (\(v_{pred}\)) and those on the same grid produced by coarsening the solutions from the high resolution grid (\(v_{true}\)). The second part of the loss function is in the form of the internal penalty function, which originates from the inequalities of the root condition constraint. The internal penalty function, which is also called the barrier function, establishes a barrier at the boundary of the feasible region to prevent the iteration from leaving the area [34]. When the output values of the neural network are close to the boundary of the feasible region, the value of the internal penalty function tends to become infinity. In this work, the barrier function of the following reciprocal form,

\[B_{1}(\vec{\alpha})=\sum_{i}\frac{1}{|g_{i}(\vec{\alpha})|+\epsilon},\quad \text{(a)}\qquad\qquad B_{2}(p,q)=\sum_{i}\frac{1}{g_{i}(p,q)},\quad\text{(b)} \tag{10}\]

that the boundary of the feasible region is determined by \(g_{i}(\cdot)>0\), is employed. The barrier function for the semi-constrained model Eq. (10a) sets the constraints on the values of \(\alpha_{i}\) directly. Because the initial output value may be at the boundary of the feasible region, the application of absolute values and \(\epsilon\) here is to prevent the denominator from changing to 0 during training, resulting in subsequent training not being possible. In addition we need to test the output of the semi-constrained model after each training. If the output coefficients are not in the feasible domain, the model will be retrained with modified initial weights.

For the fully-constrained model, the barrier function Eq. (10b) is enforced on \(p\), \(q\) as shown in Eq. (9), in which the consistency condition \(\rho(1)=0\) is guaranteed. Due to the initial output of the fully constrained model being within the feasible domain, it will not run out of the feasible domain at small learning rates. There is no need for techniques to ensure that the denominator is not zero. The specific initial network settings are shown in subsection 2.4 and section 3. As a consequence, the loss function can be written as

\[\text{loss}=\text{MSE}(v_{pred},v_{true})+\gamma B, \tag{11}\]

where \(\gamma\) is an adjustable hyperparameter and set to 0 for the unconstrained model. The \(\epsilon\) in Eq. (10a) is a small quantity and should be adjusted with \(\gamma\) changes.

### Learning of 3-step linear multistep methods

The data preparation and model training of a 3-step linear multistep method are presented to show the procedure of developing data-driven time integration schemes based on the BP neural networks, which is then applied to the one-dimensional heat equation, wave equation and Burgers' equation.

The outputs of the neural network model are the coefficients \(\alpha_{i}\) and \(\beta_{i}\) for the linear multistep formulation (Eq. (2)). The inputs consist of coarse-grained solutions from the previous time step, i.e., \(v^{n}\) at different spatial locations (with the solutions \(v^{n-1}\), \(v^{n-2}\) for computing the right-hand-side term for the 3-step linear multistep method). Consequently, the learned coefficients are influenced by spatial and temporal variations in the solutions. Since the internal penalty function is employed, the initial values of the outputs of the neural network need to be specified in the feasible region. A small learning rate is employed, such that the step size of updating the hyperparameter of network is small, which ensures the outputs are located in the feasible region. Max-min normalization is used for inputs during training and testing to ensure valid constraints during testing.

To train the data-driven model for a 3-step linear multistep method, the coefficients of the third-order explicit Adams methods, which can be written as

\[v^{n+1}=v^{n}+\Delta t\left(\frac{23}{12}F_{n}-\frac{16}{12}F_{n-1}+\frac{5}{ 12}F_{n-2}\right), \tag{12}\]

are employed to set the initial output of the neural network. To prevent the training results from falling into local optimization, small perturbations are added to the coefficients. With the assumption that the optimal coefficients are close to the initial guess, the initial biases of the network are set manually, and the initial weights are set as small random numbers for realizing the optimization in small step sizes and ensuring that the output coefficients are still in the feasible region when the inputs are not in the training set.
A schematic showing the training process is shown in Figure 1. During each iteration of the model training, the solution from the previous time step is first given as the input. With the input, the scheme coefficients are then given by the NN model (neural network model). At last, the loss is computed as the mean square error of the predicted solution \(v_{pre}^{n+1}\) and the constraints for the coefficients (i.e., Eq. (11)), in which \(v_{pre}^{n+1}\) is computed by advancing the equation for one step using the output coefficients of the NN model, the solution at previous time steps, and the right-hand-side term.

The employed neural network is composed of three layers, i.e., the input layer with the number of neurons equals to the number of grid points in \(x\), a hidden layer with 20 neurons, and the output layer with 6 neurons for the unconstrained model and the semi-constrained model, and 4 neurons for the fully-constrained model. Due to the initial loss being below \(10^{-5}\), small learning rates, somewhere around \(10^{-7}\), are employed to adjust the initial guess of the coefficients. The three models are all trained using the Adam optimizer. The activation function is ReLU. As the number of biases in the first two layers is greater than the output, the values of the additional bias will default to the last number of the manually set bias for the output layer. The models are trained in the framework of TensorFlow 1.0.

The training data are generated by coarsening the numerical solution on the fine grid using the cell-averaging approach. The spatial discretizations based on the finite volume method are performed, which has good stability and conservation properties, being consistent with the cell-averaging approach employed for coarsening.

The procedure for using the learned data-driven 3-step time integration scheme is shown in Figure 2. In this figure, the process of advancing in time is shown by the green arrows. At beginning, the solutions \(v^{1}\) and \(v^{2}\) are calculated from the initial solution \(v^{0}\) using the Runge-Kutta method. The black arrows indicate the process of employing the data-driven time integration scheme, while the red arrows indicate the calculation of the right-hand-side terms.

## 3 Results

In this section, we present the results from the data-driven time integration schemes and compare them with those from the Runge-Kutta method of order 3(2) with adaptive time-stepping. In subsection 3.1, we apply the data-driven

Figure 1: A schematic for the training procedure of a data-driven 3-step linear multistep scheme.

Figure 2: A schematic for the procedure of using the learned 3-step linear multistep scheme.
time integration schemes with the finite volume method to solve 1-D heat equations with different thermal diffusivities on a grid \(4\times\) coarser than the reference grid. And then we consider first-order wave equations with varying wave speeds in subsection 3.2, and analyze why data-driven time integration schemes can achieve highly accurate results. In subsection 3.3, we test the data-driven temporal schemes on the Burgers' equation on a grid \(32\times\) coarser than the reference grid, with the spatial derivatives approximated using the data-driven schemes proposed in [31].

### 1-D Heat Equation

The heat equation considered here is

\[\begin{cases}\frac{\partial v}{\partial t}=\frac{\partial}{\partial x}\left( \lambda\frac{\partial v}{\partial x}\right),&0\leqslant x\leqslant 1,\ 0\leqslant t\leqslant 1,\\ v(x,0)=\sin(2\pi x),&0\leqslant x\leqslant 1,\end{cases} \tag{13}\]

where \(\lambda\) is thermal diffusivity which expresses the ability of the a system tending to a uniform temperature in heating or cooling [35]. This equation describes the heat conduction or diffusion in one-dimensional isotropic medium. Eq. (13) employs the periodic boundary condition with the domain of size \(L=1\). The exact solution is

\[v(x,t)=\mathrm{e}^{-4\pi^{2}\lambda t}\sin(2\pi x). \tag{14}\]

The thermal diffusivity \(\lambda\) here is set in the range of \(0.1\) to \(1\). We apply the data-driven time integration schemes trained under a particular thermal diffusivity to other thermal diffusivities in this range. The first-order finite-volume method is employed for spatial discretization during both training and testing.

The three different data-driven time integration schemes with/without constraints (i.e., the unconstrained, the semi-constrained, and the fully-constrained schemes) are trained on a training set with a fine grid solution coarsen by a factor of \(4\) for \(\lambda=0.5\), and then tested on \(\lambda\in\{0.1,0.2,0.3,0.4,0.6,0.7,0.8,0.9,1.0\}\). We assume that, for example, if time integration schemes have lower errors at \(\lambda=0.7\) and \(0.8\), then it is highly likely that schemes will perform better in \(\lambda\in[0.7,0.8]\). The fine grid solution has \(64\) grid points in the \(x\) direction for \(0\leqslant x\leqslant 1\), which means there are \(16\) grid points on the \(4\times\) coarse grid. The size of the time step is set to \(0.0001\).

Specifically, to learn the data-driven 3-step linear multistep model, the initial weights are uniformly distributed in the range from \(-0.0005\) to \(0.0005\) for the unconstrained model and the semi-constrained model, with the same initial biases \([0,0,-1,5/12,-4/3,23/12]\) for \(\alpha_{0,1,2}\) and \(\beta_{0,1,2}\). With the consistency condition enforced, the number of outputs in the fully-constrained model is \(4\), i.e., \(p\), \(q\), \(\beta_{0}\) and \(\beta_{1}\). Considering that \(p,q=0\) for the generating polynomial \(\rho(\lambda)\) of Eq. (12) when simplifying it as the form of Eq. (9), the initial biases of the fully-constrained model are set as \([0,0,5/12,-4/3]\), and the initial weights are uniformly distributed in the range from \(-0.005\) to \(0.005\). The learning rate is \(10^{-7}\) for both unconstrained model and semi-constrained model and is \(5\times 10^{-7}\) for fully-constrained model. The three models are trained by Adam optimizer for \(8000\) steps. The \(\gamma\) in loss function Eq. (11) is set to \(10^{-18}\) for the semi-constrained model and \(10^{-12}\) for the fully-constrained model. The purpose is that when the output coefficients are in the feasible domain, the penalty function will minimize its impact on the main part of the loss, i.e. the prediction error.

Figure 3 shows the results of the 1-D heat equation for \(\lambda=0.3,0.7\) and \(1.0\). Since the solution tends to be a constant after a time greater than \(0.5\), only the solution and error for t from \(0\) to \(0.5\) are plotted here. The results for a larger range of t and other \(\lambda\) can be found in table 4 in the Appendix. From Figure 3 and table 4, it can be seen that unconstrained model and semi-constrained model have lower values of mean square error (MSE) and mean absolute error (MAE) than the Runge-Kutta method. Little difference is observed between prediction from the fully-constrained model and Runge-Kutta method. We believe this is because the fully-constrained model needs to find coefficients satisfying both the root condition and the consistency condition, leaving almost no space for optimisation.

As can be seen from the table 4 and figure 3, while the unconstrained model works well for \(\lambda=0.3\) to \(1\) (especially for \(\lambda\) close to the training set around \(0.5\)), it performs the worst for \(\lambda=0.1\) and \(0.2\), with the MSE one to two orders of magnitude greater than the Runge-Kutta method. This is understandable as the unconstrained model has the space to learn the mechanism for error reduction, but with a narrow scope of generalisation as a result of not satisfying certain mathematical constraints.

In contrast, the generalisation ability of the semi-constrained model is well demonstrated, with its predictions closer to the exact solution when compared with the Runge-Kutta method for the cases with \(\lambda\) from \(0.2\) to \(1\). This can be explained by the fact that the root constraint enhances the stability of the learned scheme, and meanwhile leaves room for optimisation.
### 1-D Wave Equation

In this subsection one-dimensional first-order wave equation (i.e. linear convection equation [35]) with wave speed \(c\) is considered, which can be written as

\[\begin{cases}\frac{\partial\tau}{\partial t}+c\frac{\partial\tau}{\partial x}=0,\quad 0\leqslant x\leqslant 1,\ 0\leqslant t\leqslant 1,\\ v(x,0)=\sin(4\pi x),\quad 0\leqslant x\leqslant 1.\end{cases} \tag{15}\]

where the wave speed \(c>0\). The wave equation describes the propagation of a wave in a homogeneous medium with a velocity \(c\) in the \(x\) direction. Eq. (15) employs the periodic boundary condition with the domain of size \(L=1\). The exact solution is

\[v(x,t)=\sin\left[4\pi(x-ct)\right]. \tag{16}\]

The wave speed \(c\) here is set from \(0.1\) to \(1\). The spatial discretization scheme is the second order finite-volume method. The training set is generated by coarsening the numerical solution on the fine grid with \(c=0.5\). The test set has the values of \(c\in\{0.1,0.2,0.3,0.4,0.6,0.7,0.8,0.9,1.0\}\). The fine grid case has points of \(64\) in \(x\)-direction. The cell-averaging method is employed for coarsening find grid solution to obtain the training data. The size of the time step is set to \(0.0001\).

To train the model, the initial weights are set uniformly distributed in the range of \([-0.0005,0.0005]\) for the unconstrained model and \([0,0.0005]\) for the semi-constrained model while \([-0.005,0.005]\) for the fully-constrained

Figure 3: Test results of the 1-D heat equation for (A) \(\lambda=0.3\), (B) \(\lambda=0.7\) and (C) \(\lambda=1.0\) and their curves of error over time (D, E, F, G, H, I). The subgraph on the top of (A, B, C) is the exact solution obtained from Eq. (14) and then followed by realizations of solutions and absolute error distribution from different time integration schemes. The numbers in brackets on the right of (A, B, C) are the mean square error averaged over the whole domain. The error shown in (D, E, F) is the mean squared error obtained by averaging the error in space \(([0,1])\) and time \(([0,t])\). The error shown in (G, H, I) is the mean absolute error averaged over space \(([0,1])\) at \(t\) instant. (D, G) is the error curves for (A): \(\lambda=0.3\) and (E, H) for (B): \(\lambda=0.7\), (F, I) for (C): \(\lambda=1.0\).
model. The initial biases, learning rate, and optimizer are identical with the 1-D heat equation case but training for 10000 steps.

Figure 4 shows the results of the 1-D wave equation for \(c=0.2,0.7\) and 1.0. The MSE and the MAE for other \(c\) could be found in table 5 in the Appendix. Here we clearly see that the coarse-grained prediction of the unconstrained model for \(c\) from 0.1 to 1 is the closest to the exact solution among the considered methods, with the values of MSE and MAE reduced by at least one order of magnitude in comparison with the Runge-Kutta method and the third-order Adams methods. Significant and consistent improvements are also observed for the semi-constrained model, with at least about 90 percent reduction in MSE and 60 percent reduction in MAE for each \(c\). The fully-constrained model, on the other hand, predicts almost the same solutions as the Runge-Kutta method and the third-order Adams method. For the sake of comparison, the third-order Adams method will be used here as the baseline method.

In the following, we further analyze the error of the 1-D wave equation to provide some understanding of the improvement obtained from the learned time integration scheme. As shown from the Fourier analysis, the spatial discretization scheme introduces significant dispersion errors on coarse grids for the wave equation, here manifesting as a gradual lagging behind the phase during propagation. Figure 5 illustrates the solution of different models corresponding to different \(c\) at \(t=1\). It is evident that the phase error increases as \(c\) increases. Among them, the Runge-Kutta method, the third-order Adams method, and the fully-constrained model lag in phase the most.

Specifically, we aim to investigate how the coefficients of the unconstrained and semi-constrained models lead to error reductions. The third-order explicit linear multistep method combined with a second-order finite-volume method corresponds to a four-level explicit scheme. Assuming that the solution \(v^{n-2},v^{n-1}\) and \(v^{n}\) from the three previous time

Figure 4: Test results of wave equation for (A) \(c=0.2\), (B) \(c=0.7\) and (C) \(c=1.0\) and their curves of error over time (D, E, F, G, H, I). The subgraph on the top of (A, B, C) is the exact solution obtained from Eq. (16) and then followed by realizations of solutions and absolute error distribution from different time integration schemes. The numbers in brackets on the right of (A, B, C) are the mean square error averaged over the whole domain. The error shown in (D, E, F) is the mean squared error obtained by averaging the error in space \(([0,1])\) and time \(([0,t])\). The error shown in (G, H, I) is the mean absolute error averaged over space \(([0,1])\) at \(t\) instant. (D, G) is the error curves for (A): \(c=0.2\) and (E, H) for (B): \(c=0.7\), (F, I) for (C): \(c=1.0\).
steps are given by the exact solution Eq. (16), we intend to examine the phase error brought about by the discretization scheme after advancing for one time step to \(n+1\).

It can be proven that the phase displacement per one time step is a constant for the coefficients given by the linear multistep method for initial conditions in Eq (15). Figure 6 shows the phase displacement per one time step for cases with different \(c\). Since the coefficients from the data-driven models vary with the inputs, we have plotted the time-averaged phase displacements for these models. Ideally, the phase would move \(c\Delta t\) in one time step, with the exact slope being of \(\Delta t\) as shown in Figure 6. It is seen that phase displacement from the unconstrained model is the closest to the exact one for different values of \(c\). Improvements are also observed for the semi-constrained model. As for the fully-constrained model and third-order Adams method, however, phase displacements lower by approximately 10 percent are observed, being consistent with the observations in Figure 5. Overall, Figure 6 demonstrates that the learned time integration scheme is capable of correcting the numerical phase displacement, which reduces the dispersion error caused by the spatial discretizations. More details on the mathematical derivations can be found in Appendix C.

### 1-D Burgers' Equation

In this section, we present the results of Burgers' equation from the data-driven time integration schemes and compare them with those from the Runge-Kutta method of order 3(2). Three different test cases, including different forcing terms, long integration time and large computational domain are considered. The performance of the data-driven time integration schemes is evaluated by applying them to simulations on a coarse grid, which is \(32\times\) coarser than the fine grid simulations (512 grids points) employed for generating the training data. In the \(32\times\) coarse grid simulations, the data-driven spatial discretization schemes trained (the number of steps and the learning rate for training are slightly different from the reference) using the method in the reference [31] are employed.

The considered one-dimensional Burgers' equation is in the following form,

\[\frac{\partial v}{\partial t}+\frac{\partial}{\partial x}\left(\frac{v^{2}}{2} -\eta\frac{\partial v}{\partial x}\right)=f(x,t), \tag{17}\]

where \(\eta=0.01\) is the viscosity and \(f\) is the forcing term. The initial condition is \(v(x,t=0)=0\). The periodic boundary condition is applied in \(x\)-direction. The forcing term \(f\) is given as follows [31],

\[f(x,t)=\sum_{i=1}^{20}A_{i}\sin(\omega t+2\pi l_{i}x/L+\phi_{i}), \tag{18}\]

Figure 5: The snapshots of wave equation for \(c\in\{0.1,0.2,0.3,0.4,0.6,0.7,0.8,0.9,1.0\}\) at \(t=1\). The horizontal axis represents the spatial position, and the vertical axis represents the coarse-grained waves predicted by different methods at \(t=1\).
where \(A_{i}\in[-0.5,0.5]\), \(\alpha_{i}\in[-0.4,0.4]\), \(\phi_{i}\in[0,2\pi]\) and \(l_{i}\in\{3,4,5,6\}\). The initial weights are uniformly distributed in the range of \(-0.0001\) to \(0.0001\) and \(-0.001\) to \(0.001\) for the unconstrained model and the semi-constrained mode respectively, with the same initial biases \([0,0,-1,1/2,-4/3,23/12]\). The initial biases of the fully-constrained model are set as \([0,0,1/3,-4/3]\), and the initial weights are uniformly distributed in the range of \(-0.01\) to \(0.01\). The \(\gamma\) in loss function Eq. (11) is set to \(10^{-15}\) for the semi-constrained model and \(10^{-12}\) for the fully-constrained model. The three coarse-grained solutions of the Burgers' equation for \(t\in[0,20]\) with different forcing terms in the training set do not intersect with the test cases.

First, we examine the performance of the three data-driven time integration schemes using the test cases with the forcing term different from the training dataset. Figure 7 shows the results from two typical cases with (B, D with green-background title) and without (A, C with blue-background title) significant improvements, respectively. It is seen that the predictions from different time integration schemes are basically the same in the beginning of time. As further advancing in time, the unconstrained model and the semi-constrained model perform better than the other two models, which is measured using MSE and MAE. Little difference is observed between the predictions from the fully-constrained model and Runge-Kutta method, as no much space is left for further improvement when the constraints for deriving a third-order method are fully enforced when training the data-driven schemes.

A greater expectation of the proposed data-driven time integration scheme is their superior performance over an integration time longer than that of the training dataset. Figure 8 compares the predictions from the three learned time integration schemes trained on a temporal domain with \(t\in[0,20]\) with those from the explicit Runge-Kutta method on a temporal domain with \(t\in[0,100]\). It is clear that the unconstrained model and semi-constrained model outperform the fully-constrained model and the Runge-Kutta method as shown by the MSE and MAE for a long integration time. In order to systematically evaluate the performance of the data-driven time integration schemes, 20 cases with vastly different random forcing terms are carried out. The exact value of error and the percentage of the error reduction are shown in the AppendixD. As shown in Figure 9, the error of learned unconstrained (purple diamond) and semi-constrained (orange square) discretization schemes, which are close to each other, are less than that of the Runge-Kutta method (blue triangle) in terms of MSE and MAE for most cases. As for the fully-constrained discretization scheme, the values of the MSE and MAE (green pentagon) are approximately the same with those from the Runge-Kutta method.

Furthermore, the differences between the error from the three data-driven time integration schemes and the errors from the Runge-Kutta method are analysed by the paired t-tests. The P-values are shown in table 1 for the three schemes. At a significance level of \(0.01\), it is seen that the error from the learned unconstrained and semi-constrained time integration schemes is significantly different from (lower than) the error from the Runge-Kutta method.

Figure 6: Phase displacement per one time step for different time integration schemes for 1-D wave equation with different values of \(c\). The grey dashed lines from top to bottom are the lines fitted by the unconstrained, semi-constrained, fully-constrained, and third-order Adams methods, with the latter two almost overlapping. Two pink dashed lines with \(k=1.000e-4\) is the slope for the exact solution, which equals to the size of the time step. The values in the grey box correspond to the slopes of various lines in the figure. Since \(c=0.5\) is used to generate the training data, the corresponding results are not included.
\begin{tabular}{l c c c} \hline \hline  & Un-Con & Semi-Con & Full-Con \\ \hline Mean MSEs & 0.009832659 & 0.010041047 & 0.011981383 \\ MSE p-value & 0.000147995 & 0.000201141 & 0.087204851 \\ MAE p-value & 0.000111214 & 0.000180005 & 0.008416531 \\ \hline \hline \end{tabular}

Overall, the test results from these cases with different forcing terms and a long integration time (\(t\in[0,100]\) compared with \(t\in[0,20]\) for the training dataset) demonstrate that the data-driven unconstrained and semi-constrained time integration schemes can reduce the error caused by the spatial coarse-graining. However, the fully-constrained model can hardly reduce the error, for which we deduce that fully enforcing the constraint conditions, which leaves very little room for improvement when training the model, is the key reason.

So far, the size of the spatial domain (\([0,2\pi]\)) employed in the test cases is the same as that of the training dataset. In the following, we test the data-driven time integration schemes using test cases with a larger spatial domain, in which it will use part of the spatial points as input to determine the coefficients of time integration schemes. Specifically, a \(10\times\) domain is employed in these test cases, that the periodic boundary condition is applied for \(x\in[0,20\pi]\) and the forcing term \(f\) should be modified as reference [31]. This poses a challenge on selecting the grid points where the solutions are employed as the input of the data-driven model. In the tests of this work, the numerical solution on the first 16 grid points for \(x\in[0,2\pi]\) are fed into the data-driven time integration schemes. The length of the integration time is 40.

Figure 10 illustrates one experiment of the prediction on \(x\in[0,20\pi]\). It is seen that the unconstrained and the semi-constrained models still outperform the Runge-Kutta method and the fully-constrained model, keeping lower values of MSE and MAE as advancing in time. Figure 11 shows the error of 10 tests in the domain of \([0,20\pi]\times[0,40]\) and \([0,20\pi]\times[0,100]\), respectively. As seen, an overall improvement is obtained, although the performance deteriorates for several cases as the time advances further to 100.

Given the small initial weights of our neural network, the time-varying coefficients are approximately constant after retaining several decimal places. For data-driven unconstrained and semi-constrained time integration schemes, the constant optimized coefficients are used to calculate the above 20 samples on the domain \(t\in[0,100]\), \(x\in[0,2\pi]\). Paired t-tests similar to the above are carried out. The exact value of error are shown in the AppendixD and the p-values

\begin{table}
\begin{tabular}{l c c c} \hline \hline  & Un-Con & Semi-Con & Full-Con \\ \hline Mean MSEs & 0.009832659 & 0.010041047 & 0.011981383 \\ MSE p-value & 0.000147995 & 0.000201141 & 0.087204851 \\ MAE p-value & 0.000111214 & 0.000180005 & 0.008416531 \\ \hline \hline \end{tabular}
\end{table}
Table 1: The average of mean square errors and the P-values of paired t-test of the mean square error and the mean absolute error of three data-driven time integration schemes (The mean MSE is 0.011989284 for the Runge-Kutta method).

Figure 7: Test results for two distinct forcing terms for (A, B) the realizations of solutions and error distribution and (C, D) the corresponding mean square error curves and mean absolute error curves. The employed grid is 32 times coarser than the reference fine grid. The numbers in brackets below each subgraph in (A, B) are the mean square error averaged over the whole domain. The error shown in (C, D) is obtained by averaging the error in space \(([0,2\pi])\) and time \(([0,t])\)
are shown in table 2. It is seen that the constant optimized coefficients can produce results similar with the time-varying coefficients. Although the p-values are slightly changed, it still can be considered acceptable at the significance level of 0.01. From the perspective of computing efficiency, the obtained constant optimized coefficients can be employed.

Last but not least, we attempt to explore the reason for the improved performance of the data-driven time integration schemes. Can the improvements be obtained by increasing the scheme's order of accuracy? To probe into this, the

Figure 8: Test results for a long integration time with \(t\in[0,100]\) for (A) a realization of the solution and the corresponding error distribution in space and time, (B,D) the mean square error curves, and (C,E) the mean absolute error curves. The employed grid is 32 times coarser than the reference fine grid. The error shown in (B, C) is obtained by averaging the error in space (\([0,2\pi]\)) and time (\([0,t]\)). The error shown in (D,E) is the error averaged over space (\([0,2\pi]\)) at \(t\) instant.

Figure 9: Mean square error and mean absolute error of 20 samples for comparison with the time period of 40 and 100. (A) is the mean square error of these samples (\(0\leq t\leq 40\)) solved by four different time integration schemes. (B) is the mean absolute error of samples (\(0\leq t\leq 40\)) solved by four different time integration schemes. (C) and (D) are the same as (A) and (B), but for \(0\leq t\leq 100\).
Adams-Bashforth schemes of different orders of accuracy as shown in Table 3 are tested and compared with the constant optimized coefficients obtained from the unconstrained model.

As shown in the figure 12 for the results from a sample, improving the order of accuracy of the time integration schemes does not reduce the error during the advancement in time. The low spatial resolution employed in coarse grained simulations is the major source for error. The data-driven time integration scheme provides a new mechanism for canceling the error due to the coarse graining in space, at the price of destroying some properties or/and conditions employed for developing the conventional time integration schemes. Therefore, it is not surprising that imposing less or no constraints during model training, which leaves more space for optimizing the scheme coefficients, achieves an

\begin{table}
\begin{tabular}{l c c} \hline \hline  & Const-coef (Un-Con) & Const-coef (Semi-Con) \\ \hline Mean MSEs & 0.009864970 & 0.009423226 \\ MSE p-value & 0.000160390 & 0.000341949 \\ MAE p-value & 0.000120366 & 0.000350297 \\ \hline \hline \end{tabular}
\end{table}
Table 2: The average of mean square errors and the P-values of paired t-test of the mean square error and the mean absolute error of the constant optimized coefficients obtained from the unconstrained and semi-constrained models.

Figure 10: Test results on a 10\(\times\) larger spatial domain for (A) a realization of the solution, (B) the corresponding error distribution in space and time, (C,E) the mean square error curves, and (D,F) the mean absolute error curves. The spatial resolution is the same as the 32\(\times\) coarse test for spatial domain \([0,20\pi]\). The error shown in (C, D) is obtained by averaging the error in space (\([0,20\pi]\)) and time (\([0,t]\)).The error shown in (E,F) is the error averaged over space (\([0,20\pi]\)) at \(t\) instant.
overall better performance. However, it is likely that the data-driven time integration schemes learned for certain types of cases (e.g., one specific coarsening or one specific PDE) may not improve the performance for other cases, as the learned schemes are uniquely tuned to reduce the error of a specific grid resolution for a specific PDE discretized in space using a specific scheme. The worst case is that the data-driven schemes have to be learned case by case, which is difficult for problems requiring a large amount of computational resources, e.g., high Reynolds number turbulent flows. This issue of generalization ability exists for almost all data-driven models. Further investigations need to be carried out in future work.

## 4 Conclusions

In this work, we proposed to learn time integration schemes using neural networks for solving partial differential equations on coarse grids, and tested the learned 3-step linear multistep method using the one-dimensional heat equation, the one-dimensional wave equation, and the one-dimensional Burgers' equation. During the training of the model, mathematical constraints, i.e., the consistency condition and the root condition (stability condition), are enforced. A backpropagation neural network with three layers and a low learning rate was employed with the initial

\begin{table}
\begin{tabular}{l c c} \hline \hline k-step & Schemes & Truncation error \\ \hline \(k=3\) & \(v^{n+1}=v^{n}+\frac{\Delta t}{12}\left(23F_{n}-16F_{n-1}+5F_{n-2}\right)\) & \(\frac{3}{8}(\Delta t)^{4}\frac{d^{4}v}{d^{4}v}|_{t}=\zeta_{n}\) \\ \(k=4\) & \(v^{n+1}=v^{n}+\frac{\Delta t}{24}\left(55F_{n}-59F_{n-1}+37F_{n-2}-9F_{n-3}\right)\) & \(\frac{251}{720}(\Delta t)^{5}\frac{d^{5}v}{d^{5}v}|_{t}=\zeta_{n}\) \\ \(k=5\) & \(v^{n+1}=v^{n}+\frac{\Delta t}{720}\left(1901F_{n}-2774F_{n-1}+2616F_{n-2}-1274 F_{n-3}+251F_{n-4}\right)\) & \(\frac{95}{288}(\Delta t)^{6}\frac{d^{6}v}{d^{6}v}|_{t}=\zeta_{n}\) \\ \hline \hline \end{tabular}
\end{table}
Table 3: Adams-Bashforth schemes for \(k=3,4,5\) and the corresponding expression for truncation errors.

Figure 12: Test results of Adams-Bashforth (k=3,4,5) and constant optimized coefficients (Un-Con) for (A) a realization of the solution and the corresponding error distribution in space and time, (B,D) the mean square error curves, and (C,E) the mean absolute error curves. The meaning of subgraphs are the same as Figure 8.
values of the coefficients given as those of the conventional time integration schemes with perturbations. Three distinct time integration schemes were trained using the unconstrained model, the semi-constrained model with the root condition enforced, the fully-constrained model with both root and consistency conditions enforced.

The test results showed that the time integration schemes learned using the semi-constrained model and the unconstrained model are capable of reducing the mean square error (MSE) and the mean absolute error (MAE) for most cases, showing a reduction in error as high as an order of magnitude for the 1-D heat and the 1-D wave equation, and the most reduction in error in the range of 35% to 40% for the 1-D Burgers' equation. For the fully-constrained model, the prediction errors are close to those of the conventional time integration schemes.

Analysis of the 1-D wave equation case revealed that the learned scheme effectively mitigates the dispersion error induced by the coarse grid. Further analysis of results from the Burgers' equation indicates that, instead of the order of accuracy, the data-driven model learned a mechanism to offset the error due to low spatial resolutions. Such mechanism is not readily analyzed using the existing methods such as Taylor analysis or Fourier analysis. To develop discretization schemes that exhibit strong generalization and interpretability, future research should focus on developing mathematical theories and tools for analyzing and evaluating these learned schemes.

## Acknowledgment

This work was funded by the NSFC Basic Science Center Program for "Multiscale Problems in Nonlinear Mechanics" (NO. 11988102), National Natural Science Foundation of China (NO. 12172360), Institute of Mechanics CAS, and Chinese Academy of Sciences.
