Fourier Neural Solver for Large Sparse Linear Algebraic Systems

Abstract

Large sparse linear algebraic systems can be found in a variety of scientific and engineering fields and many scientists strive to solve them in an efficient and robust manner. In this paper, we propose an interpretable neural solver, the Fourier neural solver (FNS), to address them. FNS is based on deep learning and a fast Fourier transform. Because the error between the iterative solution and the ground truth involves a wide range of frequency modes, the FNS combines a stationary iterative method and frequency space correction to eliminate different components of the error. Local Fourier analysis shows that the FNS can pick up on the error components in frequency space that are challenging to eliminate with stationary methods. Numerical experiments on the anisotropic diffusion equation, convection–diffusion equation, and Helmholtz equation show that the FNS is more efficient and more robust than the state-of-the-art neural solver.

1. Introduction

Large sparse linear algebraic systems are ubiquitous in scientific and engineering computation, such as discretization of partial differential equations (PDE) and linearization of non-linear problems. Designing efficient, robust, and adaptive numerical methods for solving them is a long-term challenge. Iterative methods are an effective way to resolve this issue. They can be classified into single-level and multi-level methods. There are two types of single-level methods: stationary and non-stationary [1]. Due to sluggish convergence, stationary methods, such as weighted Jacobi, Gauss–Seidel and successive over-relaxation methods [2] are frequently utilized as smoothers in multi-level approaches or as preconditioners. Non-stationary methods typically refer to Krylov subspace methods, such as conjugate gradient (CG) and generalized minimal residual (GMRES) methods [3,4], whose convergence rate is heavily influenced by certain factors, such as the initial value. Multi-level methods mainly comprise the geometric multigrid (GMG) method [5,6,7] and the algebraic multigrid (AMG) method [8,9]. They are both affected by many factors, such as smoother and coarse grid correction, which heavily affect convergence. Identifying these factors for a concrete problem is an art that requires extensive analysis, innovation, and trial.

In recent years, the technique of automatically picking parameters for Krylov and multi-level methods and constructing a learnable iterative scheme based on deep learning has attracted much interest. Many neural solvers have achieved satisfactory results for second-order elliptic equations with smooth coefficients. Hsieh et al. [10] utilized a convolutional neural network (CNN) to accelerate convergence of the Jacobi method. Luna et al. [11] accelerated the convergence of GMRES with a learned initial value. Zhang et al. [12] combined standard relaxation methods and the DeepONet [13] to target distinct regions in the spectrum of eigenmodes. Significant efforts have also been made in the development of multigrid solvers, such as the learning smoother, the transfer operator [14,15] and coarse-fine splitting [16].

Huang et al. [17] exploited a CNN to design a more sensible smoother for anisotropy elliptic equations. The results showed that the magnitude of the learned smoother was dispersed along the anisotropic direction. Wang et al. [18] introduced a learning-based local weighted least square method for the AMG interpolation operator and applied it to random diffusion equations and one-dimensional small wavenumber Helmholtz equations. Fanaskov [19] produced the learned smoother and transfer operator of GMG in a neural network form.

When the anisotropic strength is mild (within two orders of magnitude), the studies referred to evidence considerable acceleration. Chen et al. [20] proposed the Meta-MgNet to learn a basis vector of Krylov subspace as the smoother of GMG for strong anisotropic cases. However, the convergence rate was still sensitive to the anisotropic strength. For convection–diffusion equations, Katrutsa et al. [21] trained the weighted Jacobi smoother and transfer operator of GMG, which had a positive effect on the upwind discretization system and was also applied to solve a one-dimensional Helmholtz equation. For second-order elliptic equations with random diffusion coefficients, Greenfeld et al. [22] employed a residual network to construct the prolongation operator of AMG for uniform grids. Luz et al. [23] extended it to non-uniform grids using graph neural networks, which outperformed classical AMG methods. For jumping coefficient problems, Antonietti et al. [24] presented a neural network to forecast the strong connection parameter to speed up AMG and used it as a preconditioner for CG. For the Helmholtz equation, Stanziola et al. [25] constructed a fully learnable neural solver, the helmnet, which was built on U-net and a recurrent neural network [26]. Azulay et al. [27] developed a preconditioner based on U-net and shift-Laplacian MG [28] and applied the flexible GMRES [29] to solve the discrete system. For solid and fluid mechanics equations, several neural solver methods for associated discrete systems have been proposed, such as, but not limited to, learning initial values [30,31], constructing preconditioners [32], learning the search directions of CG [33], and learning the parameters of GMG [34,35].

In this paper, we propose a Fourier neural solver (FNS), a deep learning and fast Fourier transform (FFT)-based [36] neural solver. The FNS is made up of two modules: a stationary method and a frequency space correction. Since stationary methods, such as the weighted Jacobi method, have difficulty eliminating low-frequency error, the FNS uses FFT and CNN to learn these modes in the frequency space. Local Fourier analysis (LFA) [5] has shown that the FNS can pick up on the error components in frequency space that are challenging to eradicate using stationary methods. The FNS builds a complementary relationship with the stationary method and CNN to eliminate error. With the help of FFT, the single-step iteration of the FNS has a $O\left(N{log}_{2}N\right)$ computational complexity. All matrix-vector products are implemented using convolution, which is both storage-efficient and straightforward to parallelize. We investigated the effectiveness and robustness of the FNS on three types of convection–diffusion–reaction equations. For anisotropic diffusion equations, numerical experiments showed that the FNS was able to reduce the number of iterations by nearly 10-times compared to the state-of-the-art Meta-MgNet when the anisotropic strength was relatively strong. For non-symmetric systems arising from the convection–diffusion equation discretized by the central difference method, the FNS can converge, while MG and CG methods diverge. In addition, FNS is faster than other algorithms, such as GMRES and BiCGSTAB(ℓ) [37]. For indefinite systems arising from the Helmholtz equation, the FNS outperforms GMRES and BiCGSTAB for medium wavenumbers. In this paper, we apply the FNS to the above three PDE systems. However, the principles employed by the FNS indicate that the FNS has the potential to be useful for a broad range of sparse linear algebraic systems.

The remainder of this paper is organized as follows: Section 2 proposes a general form of linear convection–diffusion–reaction equation and describes the motivation for designing the FNS. Section 3 presents the FNS algorithm. Section 4 examines the performance of the FNS with respect to anisotropy, convection–diffusion, and the Helmholtz equations. Finally, Section 5 describes the conclusions and potential future work.

2. Motivation

We consider the general linear convection–diffusion–reaction equation with a Dirichlet boundary condition

$$\left\{\begin{array}{cc}-\epsilon \nabla ·\left(\mathbf{\alpha }\right(x)\nabla u)+\nabla ·\left(\mathbf{\beta }\right(x\left)u\right)+\gamma u=f\left(x\right)\hfill & \mathrm{in}\Omega \hfill \\ u\left(x\right)=g\left(x\right)\hfill & \mathrm{on}\partial \Omega \hfill \end{array}\right.$$

(1)

where $\Omega \subseteq {\mathbb{R}}^d$ is an open and bounded domain. $\mathbf{\alpha }\left(x\right)$ is the $d\times d-$ order diffusion coefficient matrix. $\mathbf{\beta }\left(x\right)$ is the $d\times 1$ velocity field that the quantity is moving with. $\gamma$ is the reaction coefficient. f is the source term.

We can obtain a linear algebraic system once we discretize Equation (13) by the finite element method (FEM) or finite difference method (FDM)

$$\mathbf{A}\mathbf{u}=\mathbf{f},$$

(2)

where $\mathbf{A}\in {\mathbb{R}}^{N\times N}$, $\mathbf{f}\in {\mathbb{R}}^N$ and N is the spatial discrete degrees of freedom.

Classical stationary iterative methods, such as Gauss–Seidel and weighted Jacobi methods, have the generic form

$${\mathbf{u}}^{k+1}={\mathbf{u}}^k+\mathbf{B}\left(\mathbf{f}-\mathbf{A}{u}^k\right),$$

(3)

where $\mathbf{B}$ is an easily computed operator, such as the inverse of the diagonal matrix (Jacobi method) or the inverse of the lower triangle matrix (Gauss–Seidel method). However, the convergence rate of such methods is relatively low. As an example, we utilize the weighted Jacobi method to solve a special case of Equation (1) and use LFA to analyze the rationale.

Taking $\epsilon =1,\mathbf{\alpha }\left(x\right)=\left(\begin{array}{cc}1\hfill & 0\hfill \\ 0\hfill & 1\hfill \end{array}\right)$, $\mathbf{\beta }\left(x\right)=\left(\begin{array}{c}0\hfill \\ 0\hfill \end{array}\right)$ and $\gamma =0$, Equation (1) becomes the Poisson equation. With a linear FEM discretization, in stencil notation, the resulting discrete operator reads

$$\left[\begin{array}{ccc}& -1& \\ -1& 4& -1\\ & -1\end{array}\right].$$

(4)

In the weighted Jacobi method, $\mathbf{B}=\omega \mathbf{I}/4$, where $\omega \in (0,1]$ and $\mathbf{I}$ is the identity matrix. Equation (3) can be written in the pointwise form

$$u_{ij}^{k+1}=u_{ij}^k+\frac{\omega }{4}\left(f_{ij}-(4u_{ij}^k-u_{i-1,j}^k-u_{i+1,j}^k-u_{i,j-1}^k-u_{i,j+1}^k)\right).$$

(5)

Let $u_{ij}$ be the true solution and define error $e_{ij}^k=u_{ij}-u_{ij}^k$. Then, we have

$$e_{ij}^{k+1}=e_{ij}^k-\frac{\omega }{4}(4e_{ij}^k-e_{i-1,j}^k-e_{i+1,j}^k-e_{i,j-1}^k-e_{i,j+1}^k).$$

(6)

Expanding error in a Fourier series $e_{ij}^k={\sum }_{p_1,p_2}{v}^k{e}^{i2\pi \left(p_1{x}_i+p_2{y}_j\right)}$, substituting the general term $v^k{e}^{i2\pi \left(p_1{x}_i+p_2{y}_j\right)}$, $p_1,p_2\in [-N/2,N/2)$ into Equation (6), we have

$$\begin{array}{cc}\hfill v^{k+1}& =v^k\left(1-\frac{\omega }{4}\left(4-e^{-i2\pi p_{1}h}-e^{i2\pi p_{1}h}-e^{-i2\pi p_{2}h}-e^{-i2\pi p_{2}h}\right)\right)\hfill \\ \hfill & =v^k\left(1-\frac{\omega }{4}\left(4-2cos\left(2\pi p_{1}h\right)-2cos\left(2\pi p_{2}h\right)\right)\right).\hfill \end{array}$$

The convergence factor of the weighted Jacobi method (also known as the smoother factor in the MG framework [7]) is

$${\mu }_{\mathrm{loc}}:=\left|\frac{v^{k+1}}{v^k}\right|=\left|1-\omega +\frac{\omega }{2}\left(\cos \left(2\pi p_{1}h\right)+cos\left(2\pi p_{2}h\right)\right)\right|.$$

(7)

Figure 1a shows the distribution of the convergence factor ${\mu }_{\mathrm{loc}}$ of weighted Jacobi ($\omega =2/3$) in solving a linear system for the Poisson equation. For a better understanding, let ${\theta }_1=2\pi p_{1}h$, ${\theta }_2=2\pi p_{2}h$, $\mathbf{\theta }=({\theta }_1,{\theta }_2)\in {[-\pi ,\pi )}^2$. Define the high and low-frequency regions

$$\begin{array}{cc}\hfill & T^{\mathrm{low}}:={\left[-\frac{\pi }{2},\frac{\pi }{2}\right)}^2,\hfill \\ \hfill & T^{\mathrm{high}}:={[-\pi ,\pi )}^2\backslash {\left[-\frac{\pi }{2},\frac{\pi }{2}\right)}^2,\hfill \end{array}$$

(8)

as shown in Figure 1b. It can be seen that, in the high-frequency region, ${\mu }_{\mathrm{loc}}$ is approximately zero, whereas, in the low-frequency region, ${\mu }_{\mathrm{loc}}$ is close to one. As a result, the weighted Jacobi method attenuates high-frequency errors quickly but is mostly ineffective for low-frequency errors.

The explanation has two aspects. Firstly, the high-frequency reflects the local oscillation, while the low-frequency reflects the global pattern. Since $\mathbf{A}$ is sparse and the basic operation $\mathbf{Au}$ of the weighted Jacobi method is a local operation, it is challenging to remove low-frequency global error components. Secondly, $\mathbf{A}$ is sparse and ${\mathbf{A}}^{-1}$ is commonly dense, which means that the mapping $\mathbf{f}\to {\mathbf{A}}^{-1}\mathbf{f}$ is non-local, making local operations difficult to approximate.

Therefore, we should seek the solution in another space to obtain an effective approximation of the non-local mapping. For example, the Krylov method approximates the solution in a subspace spanned by a basis set. The MG generates a coarse space to broaden the receptive field of the local operation. However, as mentioned in Section 1, these methods require the careful design of various parameters. In this paper, we propose the FNS, a generic solver that uses FFT to learn solutions in frequency space, with the parameters obtained automatically in a data-driven manner.

3. Fourier Neural Solver

Denote stationary iterative methods of (3) in an operator form

$${\mathbf{v}}^{k+1}=\Phi \left({\mathbf{u}}^k\right),$$

(9)

and the $k-$th step residual

$${\mathbf{r}}^k:=\mathbf{f}-\mathbf{A}{\mathbf{v}}^{k+1},$$

(10)

then the $k-$th step error ${\mathbf{e}}^k:=\mathbf{u}-{\mathbf{v}}^{k+1}$ satisfies residual equation

$$\mathbf{A}{\mathbf{e}}^k={\mathbf{r}}^k.$$

(11)

As shown in the preceding section, the slow convergence rate of stationary methods is due to the difficulty of reducing low-frequency errors. Even high-frequency errors might not be effectively eliminated by $\Phi$ in many cases. We employ stationary methods to rapidly erase some components of the error and use FFT to convert the remaining error components to frequency space. The resulting solver is the Fourier neural solver.

Figure 2 shows a flowchart of the $k-$th step for the FNS. The module for solving the residual equation in frequency space is denoted as $\mathcal{H}$. It consists of three steps: FFT→Hadamard product→IFFT. The parameter $\mathbf{\vartheta }$ of $\mathcal{H}$ is the output of the hyper-neural network (HyperNN). The input $\eta$ of the HyperNN are the PDE parameters corresponding to the discrete systems. During training, the only parameter $\mathbf{\theta }$ of the HyperNN serves as the optimization parameter.

The three-step operation of $\mathcal{H}$ was inspired by the fast Poisson solver [38]. Let eigenvalues and eigenvectors of $\mathbf{A}$ be ${\lambda }_1,\dots ,{\lambda }_N$ and ${\mathbf{q}}_1,\dots ,{\mathbf{q}}_N$, respectively. Solving Equation (2) entails three steps:

• Expand $\mathbf{f}$ as a combination $\mathbf{f}=a_1{\mathbf{q}}_1+\cdots +a_N{\mathbf{q}}_N$ of the eigenvectors
• Divide each $a_k$ by ${\lambda }_k$
• Recombine eigenvectors into $\mathbf{u}=\left(a_1/{\lambda }_1\right){\mathbf{q}}_1+\cdots +\left(a_N/{\lambda }_N\right){\mathbf{q}}_N$.

In particular, when $\mathbf{A}$ is the system generated by a five-point stencil (4), its eigenvector ${\mathbf{q}}_k$ is the sine function. The first and third steps above can now be performed with a computational complexity of $O\left(N{log}_{2}N\right)$ using fast sine transform (based on the FFT). The computational complexity of each iteration of the FNS is $O\left(N{log}_{2}N\right)$.

It is worth noting that, although $\Phi$ can smooth some components of the error, the components that are removed are indeterminate. As a result, instead of filtering high-frequency modes in frequency space, $\mathcal{H}$ learns the error components that $\Phi$ cannot easily eliminate. For $\Phi$ with a fixed stencil, we can use LFA to demonstrate that the learned $\mathcal{H}$ is complementary to $\Phi$.

The loss function used here for training is the relative residual

$$\mathcal{L}=\sum _{i=1}^{N_b}\frac{\parallel {\mathbf{f}}_i-{\mathbf{A}}_i{\mathbf{u}}_i^K{\parallel }_2}{\parallel {\mathbf{f}}_i{\parallel }_2},$$

(12)

where $\{{\mathbf{A}}_i,{\mathbf{f}}_i\}$ are the training data. $N_b$ is the batch size. K indicates that the $K-$th step solution ${\mathbf{u}}^K$ is used to calculate the loss. These specific values will be given in the next section. The training and testing algorithms of the FNS are summarized in Algorithms 1 and 2, respectively.

Algorithm 1: FNS offline traning.
	Data: PDE parameters ${\left\{{\eta }_i\right\}}_{i=1}^{N_{train}}$ and corresponding discrete systems ${\{{\mathbf{A}}_i,{\mathbf{f}}_i\}}_{i=1}^{N_{train}}$
	Input: $\Phi$, HyperNN($\theta$), K and Epochs
1	forepoch$=1,\dots ,$Epochsdo serial
2			for$i=1,\dots ,N_{train}$do parallel
3					${\vartheta }_i=$ HyperNN $({\eta }_i,\theta )$
4					${\mathbf{u}}_i^0=$ zeros like ${\mathbf{f}}_i$
5					for$k=0,\dots ,K-1$do serial
6							${\mathbf{v}}_i^{k+1}=\Phi \left({\mathbf{u}}_i^k\right)$
7							${\mathbf{r}}_i^k={\mathbf{f}}_i-{\mathbf{A}}_i{\mathbf{v}}_i^{k+1}$
8							${\widehat{\mathbf{r}}}_i^k=\mathcal{F}\left({\mathbf{r}}_i^k\right)$
9							${\widehat{\mathbf{e}}}_i^k={\widehat{\mathbf{r}}}_i^k\circ {\vartheta }_i$
10							${\mathbf{e}}_i^k={\mathcal{F}}^{-1}\left({\mathbf{e}}_i^k\right)$
11							${\mathbf{u}}_i^{k+1}={\mathbf{v}}_i^{k+1}+{\mathbf{e}}_i^k$
12					end
13			end
14			Compute loss function (12)
15			Apply Adam algorithm  [39] to update parameters $\theta$
16	end
17	returnlearned FNS

Algorithm 2: FNS online testing.

4. Numerical Experiments

We used the anisotropic diffusion equation, the convection–diffusion equation, and the Helmholtz equation as examples to demonstrate the performance of the FNS. In all experiments, the matrix-vector products were implemented by CNN based on the Pytorch [40] platform. All the code can be found at https://github.com/cuichen1996/FourierNeuralSolver (accessed on 13 September 2022).

4.1. Anisotropic Diffusion Equation

Consider the anisotropic diffusion equation

$$\left\{\begin{array}{cc}\hfill -\nabla ·(C\nabla u)& =f,\mathrm{in}\Omega ,\hfill \\ \hfill u& =0,\mathrm{on}\partial \Omega ,\hfill \end{array}\right.$$

(13)

the diffusion coefficient matrix

$$C=C(\xi ,\theta )=\left(\begin{array}{cc}cos\theta & -sin\theta \\ sin\theta & cos\theta \end{array}\right)\left(\begin{array}{cc}1\hfill & 0\hfill \\ 0\hfill & \xi \hfill \end{array}\right)\left(\begin{array}{cc}cos\theta & sin\theta \\ -sin\theta & cos\theta \end{array}\right),$$

(14)

$0<\xi <1$ is the anisotropic strength and $\theta \in [0,\pi ]$ is the anisotropic direction, $\Omega ={(0,1)}^2$. We use bilinear FEM to discretize (13) with a uniform $n\times n$ quadrilateral mesh. The associated discrete system is shown in (2), where $N=(n-1)\times (n-1)$. We performed experiments for the two cases described below.

4.1.1. Case 1: Generalization Ability of Anisotropic Strength with Fixed Direction

In this case, we use the same training and testing data as [20]. For fixed $\theta =0$, we randomly sample 20 distinct parameters $\xi$ from the distribution ${log}_{10}\frac{1}{\xi }\sim U[0,5]$ and obtain ${\left\{{\mathbf{A}}_i\right\}}_{i=1}^{20}$ by discretizing (13) using bilinear FEM with $n=256$. We randomly select 100 right-hand functions for each ${\mathbf{A}}_i$. Each entry of $\mathbf{f}$ is sampled from the Gaussian distribution $N(0,1)$. Therefore, there are $N_{train}=2000$ training data. The hyperparameters used for training, including batch size, learning rate, K in the loss function, and the concrete network structure of HyperNN are listed in Appendix A.1.

The FNS can take various kinds of $\Phi$. In this case, we use weighted Jacobi, Chebyshev semi-iterative (Cheby-semi) and Krylov methods. The weight of the weighted Jacobi method is $2/3$. We use a DenseNet [41] to give the basis vector of the Krylov subspace, then approximate the solution in this subspace by least squares as in [20]. For the Chebyshev semi-iterative method, we provide a brief summary here. More details can be obtained in [42,43].

If we vary the parameter of the Richardson iteration at each step

$${\mathbf{u}}^{k+1}={\mathbf{u}}^k+{\tau }_k\left(\mathbf{f}-\mathbf{A}{u}^k\right),$$

(15)

and the maximum and minimum eigenvalues of $\mathbf{A}$ are known, then ${\tau }_k$ can be determined as

$${\tau }_k=\frac{2}{{\lambda }_{max}+{\lambda }_{min}-\left({\lambda }_{min}-{\lambda }_{max}\right)x_k},k=0,\dots ,m-1,$$

(16)

where

$$x_k=cos\frac{\pi (2k+1)}{2n},k=0,\dots ,m-1,$$

(17)

are the roots of an $m-$order Chebyshev polynomial. Here, ${\lambda }_{max}$ is obtained by the power method [44], but calculating ${\lambda }_{min}$ often incurs an expensive computational cost. Therefore, we use ${\lambda }_{max}/\mathbf{\alpha }$ to replace ${\lambda }_{min}$. The resulting method is referred to as a Chebyshev semi-iteration. We take $m=10,\mathbf{\alpha }=3$. Figure 3 shows the convergence factor obtained by LFA. It can be seen that the smooth effect improves as m increases. However, the high-frequency error along the y direction is also difficult to eliminate. When $\Phi$ is the Jacobi method, we implement $\Phi$ five-times and transform the residual to frequency space to correct errors. This is because employing the stationary method many times can enhance its smoothing effect.

After training, we choose $\theta =0,\xi ={10}^{-1},\dots ,{10}^{-6}$ for testing and generate 10 random right-hand functions for each parameter. The iteration is terminated when the relative residual is less than ${10}^{-6}$. We use “mean ± std” to show the mean and standard deviation of the iterations over the test set as in [20].

Table 1. Mean and standard deviation of the number of iterations required to achieve the stopping criterion over all tests for the anisotropic diffusion equation case 1. “−” means that it cannot converge within 10,000 steps, and “ ” means that [20] does not provide test results for this parameter.

$\mathit{\xi }$	FNS (Cheby-Semi)	FNS (Jacobi)	FNS (Krylov)	Meta-MgNet (Krylov)  [20]	MG (Jacobi)	MG (Line-Jacobi)
$\xi ={10}^{-1}$	$67.9\pm 3.81$	$138.9\pm 11.18$	$30.0\pm 4.58$	$7.5\pm 0.50$	$90.2\pm 0.98$	$13.0\pm 0.00$
$\xi ={10}^{-2}$	$101.6\pm 8.72$	$167.8\pm 13.81$	$38.5\pm 3.83$	$35.1\pm 1.04$	$752.8\pm 12.23$	$13.0\pm 0.00$
$\xi ={10}^{-3}$	$151.0\pm 7.24$	$221.7\pm 11.56$	$48.6\pm 3.26$	$171.6\pm 6.34$	$5600\pm 119.42$	$13.0\pm 0.00$
$\xi ={10}^{-4}$	$233.2\pm 5.67$	$330.1\pm 9.16$	$65.5\pm 2.80$	$375.2\pm 5.88$	−	$11.0\pm 0.00$
$\xi ={10}^{-5}$	$340.1\pm 9.43$	$466.2\pm 13.47$	$80.7\pm 7.21$	$797.8\pm 12.76$	−	$11.0\pm 0.00$
$\xi ={10}^{-6}$	$348.1\pm 11.15$	$477.9\pm 16.10$	$85.9\pm 7.52$		−	$11.0\pm 0.00$

shows the test results of different solvers. The iteration steps of all the solvers grow as anisotropic strength increases except for MG (line-Jacobi). The growth of FNS is substantially lower than Meta-MgNet and MG. When the FNS employs the same $\Phi$ as Meta-MgNet, the number of iterations is nearly 10-times lower than that of Meta-MgNet at $\epsilon ={10}^{-5}$ with the same computational complexity for a single-step iteration. The line-Jacobi smoother can only be applied to several specific $\theta$, i.e., $0,\pi /4,\pi /2,3\pi /4,\pi$. However, the FNS can be available for arbitrary $\theta$.

We use $\xi ={10}^{-1},{10}^{-6},n=64$ and Cheby-semi ($m=10$) as examples to illustrate the error that $\mathcal{H}$ learned. Figure 4a shows the convergence factor obtained by LFA of Cheby-semi ($m=10$) for solving the system with $\xi ={10}^{-1},\theta =0$. It can effectively eliminate the error components except for low-frequency errors. Figure 4b shows the distribution of errors before correction in frequency space. The result is consistent with the guidance of LFA, i.e., the error is concentrated in the low-frequency modes. Figure 4c shows the distribution of errors learned by $\mathcal{H}$ in the frequency space at this time. Its distribution is largely similar to that of Figure 4b. Figure 4d–f show the corresponding situation for $\xi ={10}^{-6},\theta =0$. In this case, Cheby-semi ($m=10$) is unable to eliminate the error along the y direction. However, $\mathcal{H}$ is still capable of learning.

4.1.2. Case 2: Generalization Ability of Anisotropic Direction with Fixed Strength

We randomly sample 20 parameters $\theta$ according to the distribution $\theta \sim U[0,\pi ]$ with fixed $\xi ={10}^{-6}$. The training and testing data are generated in a similar manner as in Section 4.1.1.

Table 2. Mean and standard deviation of the number of iterations required to achieve the stopping criterion over all tests for Equation (13) case 2.

$\mathit{\theta }$	$0.1\mathit{\pi }$	$0.2\mathit{\pi }$	$0.3\mathit{\pi }$	$0.4\mathit{\pi }$	$0.6\mathit{\pi }$	$0.7\mathit{\pi }$	$0.8\mathit{\pi }$	$0.9\mathit{\pi }$
FNS (Jacobi)	$300.2\pm 23.95$	$252.4\pm 34.81$	$269.3\pm 36.70$	$356.4\pm 39.69$	$338.4\pm 32.07$	$265.0\pm 29.96$	$266.5\pm 25.07$	$316.7\pm 17.26$
FNS (Krylov)	$58.4\pm 4.45$	$46.5\pm 4.84$	$45.1\pm 2.30$	$64.0\pm 7.01$	$54.4\pm 5.75$	$41.3\pm 7.11$	$43.0\pm 2.83$	$60.3\pm 3.93$

shows the test results. It can be seen that whether $\Phi$ is Jacobi or Krylov, the FNS can maintain robust performance in all situations, while the line-smoother is not available for these cases.

Take $\theta =j\pi /10,j=1,\dots ,4,6,\dots ,9$, $\xi ={10}^{-6},n=64$ and $\Phi$ is the weighted Jacobi method with weight $2/3$. Figure 5 shows the test results. The first row shows the weighted Jacobi method convergence factor ${\mu }_{\mathrm{loc}}$ for each $\theta$ which is computed by LFA. The region of ${\mu }_{\mathrm{loc}}\sim 1$ means that the error is difficult to eliminate. These error components are distributed along the anisotropic direction. The second row shows the error distribution in frequency space before correction, which is consistent with the results obtained by the LFA. The third row shows the distribution of the error learned by $\mathcal{H}$. It can be seen that $\mathcal{H}$ can automatically learn the error components that $\Phi$ has difficulty eliminating. The line-smoother is not feasible for these $\theta$.

4.2. Convection–Diffusion Equation

Consider the convection–diffusion equation

$$\left\{\begin{array}{cc}-\epsilon \Delta u+u_x+u_y=0,\hfill & \Omega ={(0,1)}^2,\hfill \\ u=0,\hfill & x=0,0\le y<1\mathrm{and}y=0,0\le x<1,\hfill \\ u=1,\hfill & x=1,0\le y<1\mathrm{and}y=1,0\le x\le 1,\hfill \end{array}\right.$$

(18)

We use the central difference method to discretize (18) on a uniform mesh with spatial size h in both x and y directions, which yields a non-symmetric stencil

$$\frac{1}{h^2}\left[\begin{array}{ccc}& h/2-\epsilon & \\ -h/2-\epsilon & 4\epsilon & h/2-\epsilon \\ & -h/2-\epsilon \end{array}\right].$$

(19)

To meet the stability requirement, the central difference scheme needs to satisfy the Peclet condition

$$Pe:=\frac{h}{\epsilon }max\left(\right|a|,|b\left|\right)\le 2,$$

(20)

which means that the central difference method cannot approximate the PDE solution when $\epsilon$ is extremely small. However, here, we only take into account the solver of the linear system. Thus, we continue to use this discretization method to demonstrate the performance of the FNS. In the following experiments, we explore diffusion- and convection-dominant cases, respectively.

4.2.1. Case 1: $\epsilon \in (0.01,1)$

We utilize the weighted Jacobi method as $\Phi$ in this case. Taking $\epsilon =0.1,h=1/64$ as an example, Figure 6a illustrates the convergence factor obtained by LFA of the weighted Jacobi method ($\omega =4/5$) for solving the system. We use five-times the consecutive weighted Jacobi method as $\Phi$. Figure 6b–e show that such a $\Phi$ is a good smoother. Figure 6f shows the distribution of error learned by $\mathcal{H}$ in the frequency space. It can be seen that this is essentially complementary to $\Phi$.

Figure 7 uses $\epsilon =0.5,0.1,0.05$ as examples to show the relative residual of the FNS and the weighted Jacobi method. It can be seen that the FNS has acceleration and the weighted Jacobi method ramps up as $\epsilon$ decreases. This is because when $\epsilon$ declines, the diagonal element $4\epsilon$ becomes small, and the weight along the gradient direction increases. However, the Jacobi and other gradient descent algorithms will diverge as long as $\epsilon$ continues to decrease unless the weight is drastically reduced.

4.2.2. Case 2: $\epsilon \in [{10}^{-6},{10}^{-3}]$

In this case, since the diagonal elements of the discrete system are notably less than the off-diagonal elements, the system is non-symmetric. Many methods, such as Jaocbi, CG, and MG (Jacobi), might diverge. Figure 8 shows the convergence factor of the weighted Jacobi ($\omega =4/5$) for solving the system when $\epsilon ={10}^{-2},{10}^{-3},{10}^{-6}$. It can be seen that, when $\epsilon$ is small, the convergence factor of the weighted Jacobi method for most frequency modes is larger than 1, which causes this iterative method to diverge and to be unsuitable as a smoother.

Consequently, we learn the $\Phi$ in Equation (3), where $\mathbf{B}$ is a two-layer linear CNN with channels 1→8→1, and the kernel size is $3\times 3$. The $\Phi$ is trained together with $\mathcal{H}$; the training hyperparameters are listed in Appendix A.2. Figure 9a illustrates how the learned FNS is able to solve the linear system when $\epsilon ={10}^{-3},{10}^{-4},{10}^{-5},{10}^{-6}$. It is evident that the FNS converges rapidly. Figure 9b shows the change in the relative residual for the FNS, GMRES, and BiCGSTAB(ℓ) ($\ell =15$) with $\epsilon ={10}^{-6}$. It is clear that the FNS has the fastest convergence rate.

4.3. Helmholtz Equation

The Helmholtz equation we consider here is

$$-\Delta u\left(x\right)-{\kappa }^{2}u\left(x\right)=g\left(x\right),x\in \Omega ,$$

(21)

where $\Omega ={(0,1)}^2$, $\kappa$ is the wavenumber. We currently only take into account the zero Dirichlet boundary condition. We use the second-order FDM to discretize (21) on a uniform mesh with spatial size h. The corresponding stencil reads

$$\frac{1}{h^2}\left[\begin{array}{ccc}0& -1& 0\\ -1& 4-{\kappa }^2{h}^2& -1\\ 0& -1& 0\end{array}\right].$$

(22)

We examine the FNS performance at a low wavenumber ($\kappa =25$) and medium wavenumber ($\kappa =125$). For $\kappa =25$, we take $h=1/64$ and $h=1/256$ for $\kappa =125$. Using Krylov in [20] as $\Phi$ and $g\left(x\right)=1$; the training hyperparameters are listed in Appendix A.3. Figure 10 shows how the relative residual decreases with different solvers. For $\kappa =25$, the FNS performs best for the first 300 steps; however, BiCGSTAB performs better at the end. For $\kappa =125$, the FNS outperforms BiCGSTAB. The GMRES results were too poor to display for this case.

5. Conclusions and Future Work

This paper proposes an interpretable FNS to solve large sparse linear systems. It is composed of a stationary method and a frequency correction, which are used to eliminate errors in different Fourier modes. Numerical experiments undertaken showed that the FNS was more effective and robust than other solvers in solving the anisotropic diffusion equation, the convection–diffusion equation and the Helmholtz equation. The core concepts discussed here are relevant to a broad range of systems.

There is still a great deal of work to do. First, we only considered uniform mesh in this paper. We intend to generalize the FNS to non-uniform grids by exploiting geometric deep learning tools, such as graph neural networks and graph Fourier transform. Secondly, as previously discussed, the stationary method converges slowly or diverges in some situations, which has prompted researchers to approximate solutions in other transform space. This is true for almost all advanced iterative methods, including MG, Krylov subspace methods and the FNS. This specified space, however, may not always be the best choice. In the future, we will investigate additional potential transforms, such as Chebyshev, Legendre transforms, and potentially learnable transforms based on data.