SPFNO: Spectral operator learning for PDEs with Dirichlet and Neumann boundary conditions

The SOL is a kind of specifically designed neural architecture [41], which involves learnable spectral operators that are induced by suitable basis function and implemented through correspoding transformation. Simultaneously, it emphases the concept of strictly satisfying BCs in deep surrogate models for PDEs. This architecture allows for the utilization of various spectral methods techniques in deep learning methods and leads to improved accuracy and physical validity of the neural operator models.

We now briefly introduce this architecture through the task of prediction of time-dependent PDEs. However, its application is not limited to this scenario, and other cases can be found in the numerical experiments and Fig. 1. Consider the following PDE

with inital–boundary conditions

where $\mathcal{N}$ is a continuous operator and $\mathcal{B}$ is an operator corresponding to specific BCs. The task is to learn the solution operator $\mathcal{S}_{\tau}$ which evolves the initial condition $u_{0}$ to the solution at $\tau$, namely, $\mathcal{S}_{\tau}(u_{0}(x))=u(x,t=\tau)$. We let $\left\{u_{k}(x)\right\}_{k\in\mathbb{N}}$ and $\mathcal{T}$ be a set of basis functions satisfying the BC (1) and the linear transform induced by such a basis, respectively. Then the SOL model for $\mathcal{S}(\tau)$ consists of a stack of neural spectral layers $\mathcal{L}^{(l)}$ that are induced by $\left\{u_{k}\right\}$ and $\mathcal{T}$, and are in the form of

where $\sigma$ is the nonlinear activation function, $W_{l}$ is an auxiliary pointwise shallow neural network, and $A_{l}$ is a learnable spectrum-wise matrix.

The form of architecture (2) is first demonstrated by Li et al. for FNO [27], and then adopted by multiple neural operators, such as the multiwavelet-based neural operator (MWT-NO, [48]), wavelet neural operator (WNO, [49]), spectral neural operator (SNO, [50]), orthogonal polynomial neural operator (OPNO, [41]), multi-channel IAE-Net [51] and so on. Some Examples are listed in Tab. 1. While matching the BCs by selecting a suitable underlying basis is a fundamental criterion in spectral methods, this criterion was first introduced in [41] to operator learning, and we specifically refer to the those approach that inherently satisfies the BCs of the problem as spectral operator learning. So when solving PDEs with periodic BCs, FNO also falls into the category of SOL.

In this subsection, the domain $\Omega$ is limited to one dimensional interval $[0,1]$ for convenience, while the conclusions can be easily generalized to separable multi-dimensional domains. When focusing on the cases of Dirichlet or Neumann BCs, the basis of specified trigonometric polynomials with “semi-period” is a competitive choice due to its fast transformation algorithms. More concretely, for Dirichlet BCs, it represents the Fourier sine polynomials $u_{k}(x)=\sin k\pi x,x\in[0,1]$ and (discrete) sine transform; and the basis of cosine polynomials $u_{k}(x)=\cos k\pi x,x\in[0,1]$ and (discrete) cosine transform for Neumann BCs.

The spectral methods of the aforementioned specified trigonometric polynomials have been investigated in [52, 53]. The SPFNO model is a natural derivative of these methods under the framework of SOL, with $\mathcal{T}$ substituted by the sine or cosine transform.

The discrete sine/cosine transforms (DST/DCT) of the 1st kind involve taking the discrete Fourier transform on the odd/even extensions of the given input, respectively. With a slight abuse of notation, we denote $\hat{f}$ as the odd extension of function $f\in C[0,1]$ if $f$ satisfies Dirichlet BCs or the even extension if $f$ satisfies Neumann BCs:

So the solutions are discretized at uniform grids and BCs are imposed the on boundary points, see Fig 2b. Then, the sine or cosine polynomials form a basis for $\hat{f}$.

The proof of Theorem 2.1 will be presented in Section 3.2. Theorem 2.1 ensures the effectiveness of DST and DCT as decomposition transforms in solving PDEs with corresponding specific BCs. It also leads to the following conclusion:

Finally, the structure of the SPFNO is given in Fig. 2. The projection filter illustrated in Fig. 2c is an optional component that ultimately eliminates the tiny errors of output with respect to the BCs. Generally, its influence on performance is minimal. For example, we simultaneously present the results of SPFNO with and without it in Sec. 3.5. And it is worth noting that a FNO with a similar DFT projection filter satisfies the periodic BCs exactly.

Since $f$ is typically a strong solution to PDEs in the implementation of SPFNO, it is reasonable to assume that $f$ possesses sufficient smoothness, such as belonging to $f\in H^{1}([0,1])$. But for now we only need $f\in C[0,1]$, or the differentiability on the boundary additionally for the following proof of Neumann BCs.

We first consider the one-dimensional viscous Burgers equation as follows

subject to the initial-boundary conditions

where $\Omega=\left[-1,1\right]$. Burgers’ equation is a fundamental PDE with applications in modeling turbulence, nonlinear acoustics, and traffic flow. The complexity of the dynamical system it describes poses challenges for the learning of deep models, so it has been adopted as one of the most popular benchmark problems in the field of AI4Science.

The dataset used in this experiment is sourced from the paper on the OPNO [41], where $\nu$ is fixed to $0.1/\pi$ and the task is to learn the solution operator $S_{1}:u_{0}(x)\to u(x,t=1)$. Because the input functions in PDE training and test dataset are usually randomly sampled from the same distribution, in order to ensure the generalization of the trained model, the generator of samples should possess the ability to approximate any function on the input space and a sufficient number of degrees of freedom. In this dataset, the initial condition $u_{0}(x)$ is generated by sampling from a Gaussian random field according to $u_{0}\sim\mathcal{N}(0,625(-4\Delta+25I)^{-2})$ with Neumann BCs to maintain the complexity of the sample space.

As we have concluded in [41], the BC-satisfying property is crucial for enhancing the accuracy and credibility of surrogate models for PDEs. One can reasonably anticipate that other BC-satisfying neural operators should also demonstrate notable superiority, even when employed on different grids. This expectation arises from the fact that neural operators are specifically designed to learn the underlying operator, rather than the values at discrete grids of a function.

Our results first ruled out the suspicion that the advantage of SPFNO over FNO arises from a larger bandwidth rather than BC satisfaction. This is illustrated in Fig. 3, which also indicates that this technique is not suitable for FNO. Therefore, the bandwidth of FNO is maintained at 1.

The results presented in Tab. 3 and Fig. 3 show that the SPFNO achieves a comparable average accuracy and superior worst performance when compared to the previous SOL model OPNO. On the other hand, it significantly outperforms all non-BC-satisfying baseline models as well as the BC-satisfying BOON model. And although the FNO and LSM models have achieved acceptable global errors in some cases, their errors on the BCs can be several orders of magnitude larger.

Moreover, it is noteworthy that the BOON model also achieves better performance than any non-BC-satisfying models. This further supports our viewpoints on the BC-satisfication of neural operators, as it is mutually corroborated through different neural architectures.

With the computational domain set to $\Omega=[-1,1]^{2}$, the following 2D Burgers’ equation is considered in this experiment:

Nevertheless, the output of the target operator consists of solutions at multiple fixed time steps, i.e. $S_{\tau_{0},...,\tau_{n}}:u_{0}\mapsto\left\{u(\cdot,\tau_{1}),...,u(\cdot,\tau_{n})\right\}$, so that the time-dependent PDEs can be efficiently solved by only one forward propagation. As a result, both the operator and the task can be more complicated compared with the 1D case. We take the dataset from [41] and choose a subset of time discretization by fixing $\left\{\tau_{i}\right\}=\left\{0.2,0.6,1.0\right\}$. In this dataset, $u_{0}$ is generated according to the Gaussian random field $\mathcal{N}(0,16(\Delta+16I)^{-2})$. During the training, the bandwidth of the learnable matrix $A_{l}$ is also set as $4$. The results are illustrated in Tab. 4, where the results of different SOL models are also very close, and significantly outperform the baseline models.

The coupled reaction–diffusion (Allen Cahn) equations are formulated as follows

where

The nonlinearly coupled variables $u$ and $v$ represent the activator and inhibitor in the system, respectively, to which the Neumann BCs are imposed. The dataset is provided by PDEBench [56], a comprehensive set of benchmarks for scientific machine learning. Since the data are given on a staggered uniform grid, directly sub-sampling would yields an unexpected non-uniform grid. So we only perform the experiment with the original resolution of $128\times 128$.

In [56], models are trained and evaluated using an autoregressive approach, where the output at each time step serves as the input for the subsequent time step in a time series. This approach is often applied in predicting time-dependent PDEs [27] and global weather [36] using neural operators but may lead to training instability and GPU memory limitations. The results are shown in Tab. 5(a), where SPFNO acheives the lowest errors.

Furthermore, it should be noted that the difference in efficiency between SPFNO and FNO is mainly due to the constant factor arising from performing FFT on the odd/even extensions of length 2N for the DST/DCT. However, by employing optimized butterfly operations, it is possible to achieve the same efficiency as FFT.

Darcy’s law describes the flow of fluid through a porous medium. It has been widely implemented in various fields, including hydrogeology, petroleum engineering, and soil mechanics. In this experiment, the 2-D steady-state Darcy flow equations in a unit box are formulated as the following boundary value problem (BVP):

Moreover, the homogeneous Dirichlet BCs are imposed. The task is to learn the operator $\mathcal{G}(a)=u$ that maps the diffusion coefficient $a$ to the solution $u$, where the input $a$ is possibly discontinuous. This problem serves as another most commonly used benchmark for deep PDE solvers since the dataset is provided in [27]. In this dataset, the diffusion coefficient $a(x)$ is taken as a piecewise constant, while the reference solution is generated using a finite difference method. Under this premise, however, the nondifferentiable variable coefficient makes the 2nd order finite difference method unsuitable for solving the problem. So we utilize the 2D-Darcy dataset of PDEBench [56] with $f$ fixed as $100.0$ instead. Its input functions are also piecewise constant and its larger size ($10^{4}$ pieces of data compared to $10^{3}$ in [27]) contributes to the model achieving higher accuracy. The experiment is conducted using the original spatial resolution $128\times 128$.

Compared to the Neumann BC, the Dirichlet BC is much easier to learn because it does not involve any derivative. Additionally, the heterogeneity between the input and output functions leads to a much more complicated spectral structure of the mapping operator. Actually, multiple non-neural-operator method have been reported to surpass the performance of neural operators in solving Eq. (7), especially the transformers, among which the LSM acheives the highest accuracy that is known to us.

We adopt the ReduceLROnPlateau scheduler for SPFNO to accelerate the training. Similar scheduler is employed for baseline models if there is an improvement in accuracy. The results can be found in Tab. 5(b), where the SPFNO again obtains the lowest relative error. Notably, however, the LSM achieves the lowest maximum error. The main reason is that the input $a(x)$ in the dataset is discontinuous, and when the discontinuous input exhibits rapid variations, the accuracy of spectral methods will significantly decrease. Similarly, the worst performances of FNO and SPFNO are slightly inferior to that of U-Net. It might help address such issues to introduce local structures such as convolution layers or radial basis functions to capture the extremely high frequency information, but could potentially lead to a loss of flexibility or invariance in discretization [57].

We focus on the ability and flexibility of handling general geometries of our model. The incompressible flow through a pipe is considered in this experiment, of which the governing equation is incompressible Navier-Stokes equation

where $u$ is the velocity vector, $p$ is the pressure, and $\mu$ represents the viscosity. The Pipe dataset is given by the geometry-aware Fourier neural operator (Geo-FNO) [58], where the shape of the pipe is randomly generated, and the maximum velocity condition $u_{max}=[1;0]$ is imposed at the inlet, free boundary condition is imposed at the outlet, and no-slip boundary condition is imposed at the inner wall of the pipe. However, we have identified a few instances with artificial discontinuities at the upper edge of the outlet in the current Pipe dataset, see Fig. 4a. These discontinuities arise from bugs in the numerical solver when handling the free boundary condition. So we recomputed the reference solutions to address this issue. The updated dataset is now available at https://drive.google.com/drive/folders/1WPAs6bXttCPOWrDaUudC8B4dKPoju1OO. Notably, after excluding the errors caused by the dataset, the accuracy of all models has improved on the updated dataset.

In the experiment, we utilized $1000$ instances for training and $1000$ instances for test. The goal is to learn the operator that maps the mesh point locations to the horizonal fluid velocity on these mesh points. The experiment is conducted with a resolution of $129\times 129$.

The solution to the problem satisfies a fixed Dirichlet BC on the inlet and a Neumann BC on the outlet. To approximate the mixed BCs in the horizontal direction, the Fourier transform is applied on the dual Fourier extensions, as shown in Fig. 4b. When a continuous function $f\in C[0,1]$ satisfies $f(0)=0$ and $f^{\prime}(1)=0$, the Fourier expansion of its odd-even dual Fourier extension consists solely of odd-order sine series terms. The corresponding transform is also known as the “WAWS” transform [53], which is an abbreviation for the Whole-sample Antisymmetry on the left boundary and Whole-sample Symmetry on the right boundary. On the vertical direction, we use SPFNO with sine basis to impose the Dirichlet BCs. Finally, in this experiment, we implement a WAWS transformation in the horizontal direction for the mixed BCs, and a sine transform in the vertical direction for Dirichlet BCs.

In addition to the mesh deviating from the expected uniform grid in the design of SPFNO, the complexity of the problem is further compounded by the non-perpendicular alignment of the pipe’s artificial outlet with the longitudinal axis. This could lead to errors in the Neumann BC we imposed on the mesh at the outlet. To improve the flexibility of our model when dealing with general geometries, the projection filter is dropped. As shown in Tab. 6, the SPFNO models has a significant lower error compared to the baseline models. Besides, the projection filter only exhibits minimal impact on the performance. These results demonstrate the flexibility of SPFNO in handling general geometries.