Coarse-Gridded Simulation of the Nonlinear Schrödinger Equation with Machine Learning

Abstract

A numerical method for evolving the nonlinear Schrödinger equation on a coarse spatial grid is developed. This trains a neural network to generate the optimal stencil weights to discretize the second derivative of solutions to the nonlinear Schrödinger equation. The neural network is embedded in a symmetric matrix to control the scheme’s eigenvalues, ensuring stability. The machine-learned method can outperform both its parent finite difference method and a Fourier spectral method. The trained scheme has the same asymptotic operation cost as its parent finite difference method after training. Unlike traditional methods, the performance depends on how close the initial data are to the training set.

1. Introduction

A numerical time-stepping algorithm is developed to approximate the solution of a partial differential equation—the cubic nonlinear Schrödinger equation (NLS) in the coarse discretization regime. Recently, machine learning-based methods have been proposed as a method for creating PDE solvers when classic methods are under-resolved [1,2,3,4]. In this work, a machine learning approach is developed to create a numerical scheme for solving the nonlinear Schrödinger equation (NLS). The method uses a neural network to generate spatial differencing weights in an approximation of the second derivative. The network is trained using the truncation error of the discretization (using a second-order implicit-explicit [IMEX2] procedure for temporal stepping) as an objective function. The training set is generated using a Fourier spectral method on a fine grid, which is then downsampled to a coarse grid before training (the method could also be trained on exact solutions). The result is a numerical scheme which can outperform both Fourier collocation and its parent finite difference scheme when used on coarse grids and on initial data near the training set.

The equation simulated in this work is the focusing-type cubic nonlinear Schrödinger equation (NLS). NLS is a common model used when the envelope of a wave packet is a quantity of interest [5,6,7,8,9]. NLS is fully integrable [10], can be solved using the inverse scattering transform [11], and contains solitary wave solutions. None of these properties are required for the machine-learned algorithm presented here. NLS was chosen as it requires the extension of [1] to a complex-valued solution in an equation with a symmetric linear operator and for its importance in laser propagation [9,12]. The resulting scheme generalizes to other equations where symmetric differential operators are discretized, for example, Poisson, Navier–Stokes, and the heat equation. These other equations arise in important applications, such as numerical weather prediction, computational fluid dynamics, and aerospace [13]. The nonlinear Schrödinger equation [14] takes the form

$$i{u}_t+u_{xx}+2{\left|u\right|}^{2}u=0.$$

(1)

This is the focusing case of (1), which has closed-form solitary wave solutions

$$\tilde{u}(x,t)=\alpha \sec \mathrm{h}\left(\alpha (x-2\beta t)\right)exp\left(i(\beta x+({\alpha }^2-{\beta }^2)t)\right),$$

(2)

where $\alpha$ is the amplitude and $\beta$ is the speed [14].

Machine learning methods are a popular tool for studying PDEs [15,16,17]. They can be used to find the PDE underlying a set of data [18,19,20,21,22,23]. They have been used to improve the accuracy of solutions to PDEs [3,18,24,25,26,27,28,29,30]. Recently, neural networks have been used to create new PDE solvers [3,15,18,31,32,33]. These works use a variety of layer depths, ranging from 3 to 22 layers. To avoid the known dangers of deep networks, including long training time and/or overfitting [32,33], this work uses a single-layer fully connected neural network.

The machine learning model in this work serves as a spatial discretization of the second derivative term, with fixed stencil width (see [3] for an example with variable stencil). The spatial discretization is paired with an IMEX2 scheme for time-stepping. IMEX2 was chosen for its unbounded stability region [34]. Fourth-order Runge–Kutta (RK4) [18,26,28,29,35] is a popular competitor as a time-stepping algorithm, but the finite stability region of explicit methods poses a stability risk from a trained spatial differencer. In addition to choosing a time-stepper with an unbounded stability region, the machine learning network is built in such a way as to preserve the spatial operator’s symmetry, guaranteeing eigenvalues which live in the stability region of the scheme. The resulting scheme is stable for all time regardless of spatial or temporal grid spacing, a challenge for machine-learned time-steppers, as noted in [3]. This novel approach both gives a stable scheme for the nonlinear Schrödinger equation in this work, but generalizes to other partial differential equations with symmetric differential operators (Poisson, heat equation, Navier–Stokes, etc.).

The machine-learned scheme is tested on a family of initial data and compared in performance against a Fourier spectral and second-order finite difference discretization. The machine-learned scheme gives smaller errors than either of these methods for data sufficiently close to the training set but loses accuracy when the initial data are far from the training set.

There are many known strategies employed for creating training and test data. For applications in PDEs, training data can be generated from highly resolved classic methods (e.g., finite difference or Fourier spectral) [18,19,21,25,26,28,36]. A single short time interval can be used for training [1], or a random ensemble of long time data can be used for training [3]. In this work, a single solution trajectory is used for training, and testing is conducted on the time dynamics of solutions with nearby initial data.

The remainder of the paper is organized as follows: Section 2 discusses the machine learning model, the numerical scheme, and the training procedure. Section 3 presents the results of the training scheme and comparisons against classic methods, including both cases where the machine learning method outperforms, and those where it is outperformed by, its competition. Section 4 provides conclusions and avenues for future research.

2. Methodology

The machine learning method uses a neural network to improve a second-order centered finite difference scheme approximation of the second derivative,

$$u_{xx}={\displaystyle \frac{u_{j+1}-2u_j+u_{j-1}}{h^2}}+O\left(h^2\right),$$

(3)

and IMEX2 time-stepper,

$${\displaystyle \frac{u^{n+1}-u^{n-1}}{2k}}={\displaystyle \frac{1}{2}}\left(g\left(u^{n+1}\right)+g\left(u^{n-1}\right)\right)+f\left(u^n\right).$$

(4)

In (4), the function $g\left(u\right)$ is the discretization of the linear term in Equation (1) and $f\left(u\right)$ a discretization of the cubic nonlinearity. Two discretizations of the nonlinearity are considered,

$$f_j^{\left(1\right)}\left(\mathbf{u}\right)=|u_j{|}^2{u}_j,\mathrm{and}{f}_j^{\left(2\right)}\left(\mathbf{u}\right)={\left|u\right|}_j^2(u_{j+1}+u_{j-1})/2.$$

The latter has been shown to outperform the former on coarse grids [37].

This scheme has the linear stability region of the trapezoid rule, and thus is unconditionally stable when the spectrum of the linear operator is pure imaginary. While the scheme is stable for any time and space step pair, very coarse grids make the scheme useless due to large truncation errors. The goal of the machine-learned scheme developed in this work is to extend the range over which the scheme is useful, by decreasing the truncation error for large spatial step sizes.

The machine-learned model replaces the stencil weights of the centered difference approximation (3) with a neural network. The neural network stencil weights are restricted to preserve the symmetry of the operator, choosing the coefficients on the subdiagonal of the differencing matrix to equal those on the superdiagonal. The weights $a_1$ (on the superdiagonal of the differencing matrix) and $a_2$ (on the diagonal of the differencing matrix) are functions of u, defined as

$$\mathbf{a}{\left({\mathbf{u}}^n\right)}_j=\mathbf{b}+{\mathbf{W}}_1\sigma \left({\mathbf{W}}_2\left(\begin{array}{c}\mathrm{real}\left({\mathbf{X}}_j{\mathbf{u}}^n\right)\\ \mathrm{imag}\left({\mathbf{X}}_j{\mathbf{u}}^n\right)\end{array}\right)\right).$$

(5)

The vector $\mathbf{b}$ and matrices ${\mathbf{W}}_1$ and ${\mathbf{W}}_2$ contain the training parameters. Their dimensions are $b\in {\mathbb{R}}^2$, $W_1\in {\mathbb{R}}^{2\times 2}$, and $W_2\in {\mathbb{R}}^{2\times 6}$. The function $\sigma$ is an activation function, where $\sigma \left(x\right)=\tan \mathrm{h}\left(x\right)$, which acts on the input vector component-wise. No hyperparameters were used. The entire network is explicit in (5), which is fully connected, the matrices $W_1$ and $W_2$ matrices are full, and contains a single activation function. ${\mathbf{X}}_j$ is a $3\times N$ matrix with entries

$${\left({\mathbf{X}}_j\right)}_{i,l}={\delta }_{i,(l-j+2)\mathrm{mod}N},$$

which is used to restrict spatial vectors to the three-point stencil centered at point $x_j$. The modulo N in the second index imposes periodicity in the spatial dimension. The result is that ${\mathbf{X}}_j$ applied to $u^n$ restricts the support of the input vector to its three-point stencil,

$${\mathbf{X}}_j{\mathbf{u}}^n=\left(\begin{array}{c}{u}_{j-1}^n\\ u_j^n\\ u_{j+1}^n\end{array}\right).$$

(6)

Equation (6) is presented for $j\in [2,N-1]$; for $j=1$ or N, the periodicity in the j-index will be evident (and is explicit in the definition of $X_j$).

The outputs of the machine-learned model are stencil weights for the approximation of the second derivative, of the form ${\mathbf{u}}_{xx}\approx \mathbf{B}\left(\mathbf{u}\right)\mathbf{u}$, with

$$\mathbf{B}\left(u\right)=\left[\begin{array}{ccccccccc}{\left(a_2\right)}_1& {\left(a_1\right)}_1& 0& \cdots & & & \cdots & 0& {\left(a_1\right)}_N\\ {\left(a_1\right)}_1& {\left(a_2\right)}_2& {\left(a_1\right)}_2& 0& \cdots & & & \cdots & 0\\ 0& {\left(a_1\right)}_2& \ddots & \ddots & 0& \cdots & & & ⋮\\ ⋮& 0& \ddots & \ddots & \ddots & 0& \cdots & & \\ & ⋮& 0& \ddots & \ddots & \ddots & 0& \cdots & \\ & & ⋮& 0& \ddots & \ddots & \ddots & 0& \cdots \\ ⋮& & & ⋮& 0& \ddots & {\left(a_2\right)}_{N-2}& {\left(a_1\right)}_{N-2}& 0\\ 0& ⋮& & & ⋮& 0& {\left(a_1\right)}_{N-2}& {\left(a_2\right)}_{N-1}& {\left(a_1\right)}_{N-1}\\ {\left(a_1\right)}_N& 0& \cdots & & & ⋮& 0& {\left(a_1\right)}_{N-1}& {\left(a_2\right)}_N\end{array}\right].$$

Indicially, the second derivative operator is approximated

$${\left(u_{xx}\right)}_j\approx {\displaystyle \frac{{\left(a_1\right)}_{j-1}{u}_{j-1}+{\left(a_2\right)}_j{u}_j+{\left(a_1\right)}_j{u}_{j+1}}{h^2}}.$$

(7)

Notice that the weight on $u_{j-1}$ is ${\left(a_1\right)}_{j-1}$, rather than ${\left(a_1\right)}_j$. With this choice, the second derivative matrix is symmetric. The symmetry of the matrix, $\mathbf{B}\left(\mathbf{u}\right)$, ensures that this machine-learned operator’s spectrum lies in the stability region of IMEX2 regardless of the training results.

The combined time-stepping scheme to be trained is

$$i{\displaystyle \frac{{\mathbf{u}}^{n+1}-{\mathbf{u}}^{n-1}}{2k}}=-{\displaystyle \frac{1}{2h^2}}\mathbf{B}\left({\mathbf{u}}^n\right){\mathbf{u}}^{n+1}-{\displaystyle \frac{1}{2h^2}}\mathbf{B}\left({\mathbf{u}}^n\right){\mathbf{u}}^{n-1}+f\left({\mathbf{u}}^n\right).$$

(8)

Strictly speaking, the above scheme is no longer the IMEX2 in (4), since the operator g now depends on u at two different time steps. The new scheme, (8), is still second order, and has the same linear stability region as (4). The scheme is not self-starting; IMEX1 is used for the first time step. In Equation (8) and throughout, j indexes a spatial location, k is the temporal step size, h is the spatial step size, and n indexes the spatial location.

The scheme was trained using the truncation error against a highly resolved numerical simulation, $\tilde{u}$, downsampled to the coarse grid as the objective function,

$$H(b,W_1,W_2)={∥i{\displaystyle \frac{{\tilde{u}}^{n+1}-{\tilde{u}}^{n-1}}{2k}}+{\displaystyle \frac{1}{2}}\mathbf{B}\left({\tilde{u}}^n\right){\tilde{u}}^{n+1}+{\displaystyle \frac{1}{2h^2}}\mathbf{B}\left({\tilde{u}}^n\right){\tilde{u}}^{n-1}-f\left({\tilde{u}}^n\right)∥}_F^2.$$

(9)

The argument of the norm is the scheme at some choice of space and time points, a doubly indexed set. The argument is interpreted as a matrix, so the objective function is the square of the Frobenius norm of that matrix. The size of the training set comes from the number of spatial and temporal points this objective function is evaluated on, both of which are training parameters. The objective function, H in Equation (9), is minimized with classic steepest descent. Training is considered successful only if the objective function is decreased by a factor of ten from its initialization. The unknowns are initialized with classic centered finite difference as an initial guess, $\overrightarrow{b}={[1,-2]}^T,{\mathbf{W}}_1=\mathbf{0},$ and ${\mathbf{W}}_2=\mathbf{0}$.

In the testing phase of this procedure, both the truncation error (9) and the forward error are used. The forward error is calculated against highly resolved Fourier solutions, $u_F$, which are downsampled after their computation onto the coarse grid used by the machine-learned scheme. The forward error might be considered preferable as an objective function, but costs significantly more to evaluate, and we observe that the forward error and truncation error closely track one another.

In summation, a neural network is used for the weights in a differencing scheme (8). The training set is a coarse sampling of a highly resolved simulation (simulated on a fine grid) along a single solution trajectory. The testing set comprises other solution trajectories from different initial data. The machine-learned scheme has the same stencil breadth as the finite difference scheme. The memory and operational cost of using the machine-learned method after training are asymptotically identical to the finite difference, $O\left(n\right)$ flops per time step. The training procedure is a significant extra cost, but may be worth it when the trained method can achieve error thresholds which other methods cannot.

3. Results

The method was trained on a single solution trajectory, whose initial data are taken from

$$u(x,0)=D\left[\alpha \sec \mathrm{h}\left(\alpha x\right)exp\left(i\beta x\right)\right].$$

(10)

For some choices of $\alpha ,\beta ,D\in \mathbb{R}$. When $D=1$, this is a soliton solution of NLS [14] whose envelope travels as speed $2\beta$, with formula

$$u(x,t)=\alpha \sec \mathrm{h}\left(\alpha (x-2\beta t)\right)exp(i\beta x+i({\beta }^2-{\alpha }^2)t).$$

(11)

If $D\ne 1$, the initial data are a dilation of a solitary wave’s and the solution does not correspond to a single NLS soliton but has some other dynamic (e.g., including dispersive radiation [11]).

The training data were generated from a finely discretized, highly resolved numerical simulation. This simulation used Fourier collocation for the spatial derivatives and a second-order IMEX time-stepper for time-stepping. The resolutions used to generate the training set were $N_x=512,L=40$, and $k=0.002$. This same temporal step size was used for all simulations (training and testing), making the truncation error due to the temporal step size approximately four orders of magnitude smaller than the spatial truncation error. The highly resolved simulation was then downsampled to a coarse grid, with $N_x=16$ and $h=2.5$. The simulation time step was chosen so that the truncation errors in time are negligible compared to the truncation errors in space. The smallest truncation errors achieved by simulations on the coarse grid are ≈${10}^{-3}$. These are presumed to be almost entirely spatial based, since a second-order time-stepper is being used with a time step $k=2\times {10}^{-3}$. All increases in accuracy of the method are attributed to the increased accuracy of the spatial discretization due to training on the coarse grid. All simulations were developed and conducted in MATLAB.

A single training set was used to find the neural network used to generate the results of this section. This network was trained on the trajectory of a single NLS traveling solitary wave, with $D=1,\alpha =0.75$ and $\beta =0.2$ in Equation (11). The objective function for training was the truncation error (9), evaluated at 50 equally spaced times in the interval $0\le t\le 4$. The objective was minimized using the steepest descent, until the truncation error reached a threshold of $1/10$th of the truncation error of its parent finite difference method. The parent finite difference method takes $b={[1,-2]}^T$, ${\mathbf{W}}_1=\mathbf{0}$, and ${\mathbf{W}}_2=\mathbf{0}$, and was used as an initial guess for the minimization procedure.

Because the method is trained using the truncation error as the objective function, the resulting neural network is guaranteed to have truncation error, which is 10× smaller than its parent finite difference when evaluated on the training set. Since both schemes are linearly stable, one might expect that the forward error would also be 10× smaller. Unfortunately, there is no such guarantee for the forward error; it is possible for one scheme to have truncation errors that cancel out rather than accumulate, even though the schemes have the same stability. Thus, the ordering of the truncation errors does not imply an ordering on the forward errors. Such cancellation was not observed in these simulations, and the achieved forward errors were 10× smaller for the machine-learned scheme than the parent finite difference method on the training set.

A series of experiments were conducted to test the machine-learned numerical method. The simplest was comparison of solution profiles during the propagation of non-solitary initial data. The chosen initial data come from (10) with $D=1.1$ ($D=1$ would be solitary), $\beta =0.22$, and $\alpha =0.825$; all three parameters were changed by 10% from the trained values. The evolution is presented to $t=6$ (50% longer than the training time interval). Figure 1 depicts the results. The machine-learned differencing method best approximates the exact solution at all times. It outperforms its parent finite difference method but also defeats Fourier collocation. Thus, for this resolution and initial data, the machine learning method is the superior scheme. The method is guaranteed to be outperform its parent for within the training interval. Longer time simulations are in Figure 2; at the simulated resolution eventually all three schemes have $O\left(1\right)$ error.

The accuracy of the machine-learned differencing scheme depends on the distance between the solution and the training set. This dependence was probed with a series of test simulations. In the examples presented here, the spatial discretization uses 16 points, so the initial data are in ${\mathbb{C}}^{16}$. Rather than attempt to visualize the regions close to the training set in sixteen dimensional space, cross-sections of the initial data space were taken, using the parameter values of $D,\alpha$, and $\beta$ from (10).

The metric used to observe the performance of the machine-learned differencing scheme as the initial data get farther from the training set was the ratio of the forward error of the machine-learned method to the forward error of a competing method,

$$\mathrm{Forward}\mathrm{Error}\mathrm{Ratio}={\displaystyle \frac{\parallel \tilde{u}-u_{ML}\parallel }{\parallel \tilde{u}-u_C\parallel }}.$$

(12)

In Equation (12), $u_{ML}$ is the solution generated by the machine-learned differencing scheme, $\tilde{u}$ is a solution from a highly resolved Fourier collocation method (a proxy for an exact solution), and $u_C$ is the solution computed by a competing method (for example, the parent finite difference scheme or Fourier collocation conducted on a coarse grid).

In the first set of experiments, the machine-learned scheme is compared to its parent finite difference method. In all cases, the machine-learned scheme uses the weights for a classic centered difference scheme as an initial guess for training. Two different discretizations of the nonlinearity are compared. The first discretizes the nonlinearity at a single point and is known to have chaotic dynamics on coarse grids [37]. The second discretizes the nonlinearity at three points and is absent chaos. Figure 3 and Figure 4 compare the relative forward errors of the parent finite difference schemes and the machine-learned children at a variety of initial data, from (10). Figure 3 presents the relative accuracy of the parent and child schemes as functions of speed and amplitude for traveling solitary wave initial data. Both discretizations of the nonlinearity include regions near the training point where the machine-learned scheme has forward error, which is 10× smaller than the parent finite difference scheme. Figure 3 presents the relative errors in the case when the initial data have non-trivial phase ($\beta =0.2$) whose propagation is non-solitary ($D\ne 1$), varying in breadth and amplitude. Again both discretizations include regions near the training set where the forward error has decreased by a factor of 10.

The machine-learned difference methods were trained via optimization beginning with the parent finite difference method as the initialization. Although not guaranteed, it is not surprising that machine-learned methods outperform their parents near the training set. In Figure 5, the relative accuracy of the machine-learned differencing scheme with the three-point nonlinear discretization of [37] is compared to the parent finite difference method (left panel) as well as to the Fourier collocation calculated on the same coarse grid ($\Delta x=2.5$). The initial data used in this comparison are non-solitary and have a variety of initial phases (D and $\beta$ are varied in (10)). As in all comparisons with the parent data finite difference, the left panel shows an improvement by a factor of 10× near the training set. In the right panel, the method is better than Fourier collocation at the training data by only a factor of 2. There is, however, a region where the machine-learned difference scheme is 10× better than the Fourier collocation. That this region does not include the training data suggests that data in this region have a particularly large error in Fourier collocation, rather than a particularly small error for the machine-learned difference method.

In [37], the three-point discretization of the nonlinearity is observed to avoid chaotic trajectories that the coarse discretization with the single-point discretization is subject to (neither method displays chaotic trajectories if the grid is sufficiently refined). The initial data that were used to observe the chaotic trajectories were designed to trigger the Benjamin–Feir instability; a plane wave and a sideband were used on a small domain with sideband frequency within the Benjamin–Feir instability region, e.g.,

$$u(x,0)=a\left(1+\delta cos\left(\mu x\right)\right).$$

Attempts were made to train the single-point discretization of the nonlinearity on these data. Unfortunately, the optimization procedure did not converge to the prescribed threshold (a tenfold decrease in truncation error). The truncation errors of the parent finite difference scheme are also quite large on these initial data, a natural consequence of the chaotic trajectories. It appears the simple neural network does not have the flexibility required to overcome the size of the errors on this data set. It is possible that a broader and/or deeper network could be used to build a scheme which could suppress the chaos of the parent finite difference scheme in this regime. The multi-point discretization on the other hand, does not have chaotic trajectories in this region, and could be used as a parent with the current network.

4. Conclusions

A machine-learned differencing scheme for the nonlinear Schrödinger equation in the focusing regime is developed. The scheme is designed to have a pure imaginary spectrum, and thus to be linearly stable regardless of training. Two discretizations of the nonlinearity are considered, based on the observed chaos in coarse-gridded finite difference simulations of NLS in [37]. The methods are observed to be 10× more accurate than their parent finite difference schemes for data near the training set. The neural network is embedded into a symmetric matrix discreteization of the second derivative; thus, when combined with an IMEX time-stepper, it has guaranteed stability. This procedure generalizes to other PDEs with symmetric differential operators, for example, the other Schrödinger equations, the heat equation, or Poisson’s equation. The neural network comes with a training cost that traditional methods do not, but can be worth the cost when increased accuracy is necessary.