\section{ForwardNet Architecture}
\label{sec:forwardnet_architecture}
The architecture of the encoder is shown in~\Cref{tab:encoder}, where conv stands for a convolutional layer. Each convolutional layer consists of a convolutional, a batch normalization~\cite{ioffe2015batch}, and a leaky ReLU activation function~\cite{nair2010rectified} with slope 0.2. Of note, the kernel size is fixed at 3 for each convolution. The columns in\_channel and out\_channel represent the number of channels applied in each layer. Also, we show the changes in data size throughout the layers in the last two columns. As a result, the latent representations have the dimension 512*1*1. 
\begin{table}[!ht]
  \centering
  \caption{Encoder}
    \begin{tabular}{ccccccc}
    layer & in\_channel & out\_channel & stride & padding  & in\_size & out\_size \\
    \midrule
    conv  & 3     & 32    & 2     & 1     & 70*70 & 35*35 \\
    conv  & 32    & 32    & 1     & 1     & 35*35 & 35*35 \\
    conv  & 32    & 64    & 2     & 1     & 35*35 & 18*18 \\
    conv  & 64    & 64    & 1     & 1     & 18*18 & 18*18 \\
    conv  & 64    & 128   & 2     & 1     & 18*18 & 9*9 \\
    conv  & 128   & 128   & 1     & 1     & 9*9   & 9*9 \\
    conv  & 128   & 256   & 2     & 1     & 9*9   & 5*5 \\
    conv  & 256   & 256   & 1     & 1     & 5*5   & 5*5 \\
    conv  & 256   & 512   & 2     & 1     & 5*5   & 3*3 \\
    conv  & 512   & 512   & 1     & 1     & 3*3   & 3*3 \\
    conv  & 512   & 512   & 2     & 0     & 3*3   & 1*1 \\
    \end{tabular}%
  \label{tab:encoder}%
\end{table}%

On the other hand, \Cref{tab:decoder} reveals the details of the decoder structure. The aforementioned convolutional layer is also a part of the building block in the decoder. In this table, the kernel is set in size 3 as well. The stride and the padding are hidden since they are fixed at 1 and 1, respectively, in all convolutions. Moreover, up\_conv stands for the combination of an up-sampling operation and a convolutional layer. To avoid the ``checkerboard effect", ForwardNet employs up-sampling instead of deconvolution (transposed convolution). The up-sampling operation is implemented with the nearest-value mode and has a parameter scale that controls the output size.  Besides up\_conv and conv, the crop layer removes the surrounding part of the intermediate representation and retrieves the middle component to obtain the 70*70 size that matches the wavefield. Finally, the output channel is 2 because the wavefield is complex-valued. The two channels correspond to the real part and the imaginary part. 
\begin{table}[!ht]
  \centering
  \caption{Decoder}
    \begin{tabular}{cccccc}
    layer & in\_channel & out\_channel & scale & in\_size & out\_size \\
    \midrule
    up\_conv & 512   & 512   & 5     & 1*1   & 5*5 \\
    conv  & 512   & 512   &       & 5*5   & 5*5 \\
    up\_conv & 512   & 256   & 2     & 5*5   & 10*10 \\
    conv  & 256   & 256   &       & 10*10 & 10*10 \\
    up\_conv & 256   & 128   & 2     & 10*10 & 20*20 \\
    conv  & 128   & 128   &       & 20*20 & 20*20 \\
    up\_conv & 128   & 64    & 2     & 20*20 & 40*40 \\
    conv  & 64    & 64    &       & 40*40 & 40*40 \\
    up\_conv & 64    & 32    & 2     & 40*40 & 80*80 \\
    conv  & 32    & 32    &       & 80*80 & 80*80 \\
    crop  &       &       &       & 80*80 & 70*70 \\
    conv  & 32    & 2     &       & 70*70 & 70*70 \\
    \end{tabular}%
  \label{tab:decoder}%
\end{table}%


\section{Background}
\label{sec:background}

\subsection{Helmholtz equation}
\label{subsec:helmholtz}
Prior to an in-depth discussion of the proposed method, we first briefly introduce the basics of the acoustic wave equation. Let $(x,~z) \in \mathbb{R} \times \mathbb{R}$ denote a point in a two-dimensional spatial domain. Specifically, $x$ is the horizontal distance, and $z$ is the depth. Then, the two-dimensional acoustic wave equation with a constant density is  
\begin{align}
    \nabla^2 \bm p(t, x, z) - \frac{1}{v(x, z)^2} \frac{\partial^2 \bm p(t, x, z)}{\partial t^2} = s(t, x_s, z_s),
    \label{eqn:time_wave}
\end{align}
where $\bm p$ denotes the acoustic pressure wavefield, $v(x,~z)$ is the velocity at location $(x,~z)$, and $s$ is the point source at location $(x_s,~z_s)$. The symbol $\nabla^2 := \frac{\partial^2 }{\partial x^2} + \frac{\partial^2 }{\partial z^2}$ is the Laplacian. \Cref{eqn:time_wave} describes the spatiotemporal pattern of how seismic energy spreads out from the source. However, it is challenging to learn the time domain wavefield, given complicated velocity models. Particularly, the subsurface structures reflected by the velocity models cause highly nonlinear scattering effects, making the pattern of the time domain wavefields exceedingly complicated. Moreover, the temporal derivative induces the dependence between snapshots of the wavefield at different time steps, which hinders the training of a neural network. Therefore, we focus on solving the Helmholtz equation:
\begin{align}
    \nabla^2 \bm u(x, z; \omega) + k^2 \bm u(x, z; \omega) = f(x_s, z_s; \omega),
    \label{eqn:helmholtz}
\end{align}
where $\omega$ is the angular frequency, and $k = \frac{\omega}{v}$ is the wavenumber. The quantify of interest is the frequency domain pressure wavefield $\bm u$. This equation can be viewed as the result of applying the Fourier transform to the time component while preserving the spatial information. Unlike $\cref{eqn:time_wave}$, the Helmholtz equation does not have a temporal derivative, but only revolves around frequency $\omega$. Hence, the wavefields of different frequencies are independent of each other.


\subsection{Fourier neural operator}
\label{subsec:fno}
The idea of the Fourier neural operator (FNO) stems from the Green's function of the linear differential operator $L$. A linear PDE can be written in the form $L u = f$, where $u$ is the solution to the PDE, and $f$ is the external force.  The corresponding Green's function $G(\cdot,~\cdot)$ then gives the solution as an integral,
\begin{align}
    u(\bm{{x}}) = \int_{\mathcal{D}} G(\bm{{x}}, \bm{{y}}) f(\bm{{y}}) \mathrm{d} \bm{{y}},
    \label{eqn:green_integral}
\end{align}
where $\bm{{x}}, \bm{{y}} \in \mathcal{D} \subset \mathbb{R}^d$. If there exists a function $g$ such that $G(\bm{{x}},~\bm{{y}}) = g(\bm{{x}} - \bm{{y}})$, \cref{eqn:green_integral} reduces to a convolution operation, and $G$ serves as the convolution kernel. For a parametric PDE, the Green's function also depends on the parameter. For example, the velocity $v$ is a parameter in the wave equation or the Helmholtz equation. In this framework, the parameter is assumed to be a function of $\bm{{x}}$. Therefore, \cite{li2020fourier} defines the kernel integral operator by 
\begin{align}
    \big(K_{\theta}(v) h_i \big)(\bm{{x}}) = \int_{\mathcal{D}} \kappa_{\theta}(\bm{{x}}, \bm{{y}}, v(\bm{{x}}), v(\bm{{y}})) \, h_i(\bm{{y}}) \mathrm{d} \bm{{y}},
    \label{eqn:kernel_integral_operator}
\end{align}
where $v$ is the PDE parameter, and the kernel integral operator is parameterized with a neural architecture with trainable parameters $\theta$ in the parameter space $\Theta$. Moreover, the intermediate representation $h$ is updated by the following,
\begin{align*}
    h_{i+1}(\bm{{x}}) = \sigma\left(W h_i(\bm{{x}}) + \big(K_{\theta}(v) h_i \big)(\bm{{x}})\right).
\end{align*}
Here, $W$ denotes a linear transform. Importantly, the neural operator consists of a sequence of non-local kernel integral operators combined with the activation function $\sigma$ so that it is capable of approximating highly nonlinear operators. In this sense, the neural operator can be considered as a generalized Green's function. However, it requires more effort to approach the convolution integral in \cref{eqn:kernel_integral_operator}. Unlike~\cite{li2020neural} in which message passing in graph neural networks is employed, FNO leverages the convolution theorem which states convolution is equivalent to multiplication in the Fourier space. Therefore, FNO directly parameterizes the kernel function $\kappa_{\theta}$ in the Fourier space. Then, it Fourier transforms the latent representations, performs the multiplication, and inverse Fourier transforms the results back to the original space. In the implementation, the fast Fourier transform (FFT) preserves computational efficiency. 

As FNO maps the parameter $v$ to the solution $u$, it succeeds in solving PDEs with variable parameters. For comparison, frameworks such as~\cite{raissi2019physics,sun2020theory} only solve one instance of the PDE after fixing the parameter. Another advantage of FNO is it is mesh-independent. This is because FNO takes the coordinate $\bm{{x}}$ as part of the input. When the resolution of the discretization changes, FNO generalizes to unseen locations. Our work leverages FNO to simulate the Helmholtz operator $\nabla^2 + k^2$ in \cref{eqn:helmholtz}, mapping velocity $v$ to the solution.




\section{Conclusion}
In this work, the Fourier neural operator (FNO) is applied to efficiently solve the Helmholtz equation under the context of variable velocity. Furthermore, we extend FNO to the multi-source and multi-frequency setting and propose the paralleled Fourier neural operator (PFNO) that consists of a group of FNOs. The paralleled structure distributes the data of different frequencies to different FNOs such that each FNO concentrates on a single frequency but multiple sources. PFNO achieves high accuracy even on the complicated high-frequency wavefields while FNO has less impressive performance. In the numerical experiments on the $6$ OpenFWI datasets, we demonstrate the accuracy of PFNO with wavefields visualization and show that it outperforms FNO and ForwardNet on all the $6$ datasets. Moreover, a trained PFNO is able to generalize to other datasets provided that the velocity models share similar features. Also, PFNO is robust in the case of noisy data. This strength makes PFNO reliable against measurement errors in practice. Finally, the ``optimal" 9-point method is compared with PFNO in terms of computational efficiency. PFNO embraces fast inferences as expected, taking milliseconds to solve for the wavefield corresponding to a velocity model. Even when the training time is incorporated, PFNO still gains an edge over the ``optimal" 9-point method given a large number of velocity models. With the strengths discussed above, FNO or PFNO can be a building block of an accurate and fast numerical solver for the Helmholtz equation when the model approaches a large number of velocity models. 




\section{Introduction}
\label{sec:introduction}
The last decade has witnessed the fast development of deep learning, thanks to the enormous growth in the computational power as well as increasingly ample data that is accessible in many fields. By the universal approximation theorem~\cite{barron1994approximation,cybenko1989approximation,hornik1989multilayer}, neural networks are capable of approximating any nonlinear function satisfying certain smoothness conditions. Therefore, their high expressiveness allows them to capture highly sophisticated and latent patterns from data. Furthermore, neural architectures can be exploited to learn the mappings between functions in the infinite-dimensional space. As shown in the work~\cite{chen1995universal}, neural networks can effectively approximate continuous nonlinear operators. This universal approximation theorem for operators leads to the development of DeepONet~\cite{lu2019deeponet,lu2021learning}. Under the dome of operator learning, another string of research defines the neural operator by a series of integral operators and activation functions~\cite{li2020fourier,li2020neural}.

The solution of a partial differential equation (PDE) is a type of complicated mapping that usually does not have an explicit form. In this work, we intend to solve the acoustic wave equation as wave equation forward modeling is the engine of many fields of study that are related to wave phenomena. We take the full waveform inversion (FWI) as an example, which is a powerful geophysical technique to reconstruct high-resolution subsurface seismic velocity models from seismic records~\cite{tarantola1984inversion}. It has been proven in the past decade that FWI is capable of meeting most of the industrial production demands~\cite{virieux2009overview}. However, the core of FWI involves numerically solving second-order PDEs (wave equations) twice in one iteration to calculate the gradient regarding the velocity update. The wavefield modeling step in FWI has a computational complexity of $O(N^3)$ for 2D cases and $O(N^4)$ for 3D cases. Usually, hundreds of such iterations would be needed before the algorithm converges to a satisfactory velocity model. Thus, a successful FWI induces tremendous computational costs in real-data implementations, and such costs will even dramatically increase when dealing with large-scale datasets. As more and more seismic inversion frameworks~\cite{jin2021unsupervised,sun2021physics} perform forward modeling and the inversion concurrently, it is increasingly crucial to solve the forward problem accurately and efficiently.

It has been an emerging direction in which deep neural networks are applied to efficiently solve PDEs~\cite{berg2018unified,han2018solving,long2018pde,malek2006numerical}. This is motivated by numerous cases where the PDE admits no straightforward numerical methods, or where traditional solvers are computationally expensive. Under such a context, the solution to the PDE, which is a function of time and location, is approximated by neural networks. PDE-constrained frameworks~\cite{raissi2019physics,wang2020deep,zhang2019quantifying,zhu2019physics} impose the governing equation(s) as a constraint in the training, since purely data-driven models may yield solutions that violate the governing equation. Nevertheless, these frameworks require that the form of the PDE be explicitly known. As mentioned above, the wave equation is of particular interest as it governs the behavior of seismic waves. As an instance, the physics-informed neural network (PINN) was applied to solve the wave equation in the time domain~\cite{rasht2022physics}. Most importantly, the frameworks above only solve the wave equation on a fixed velocity model. This is an obvious limitation since abundant velocity models present in reality, and we expect a model to solve the equation given different velocity models. Another direction aims to embed the PDE information into the network architecture. For example, \cite{sun2021physics,sun2020theory} harness a recurrent neural network (RNN) structure based on the finite difference of the wave equation. Particularly, the sequence of interest is formed by the wavefield values that are consecutive in time so that the time evolution of the wave directly depends on preceding states. However, the PDE information is not strictly imposed during the training. As a consequence, it remains an open question whether the model can generalize to new datasets. Also, this framework relies on a regular mesh because it is built on the finite difference. In contrast, FNO is mesh-independent since its input includes the locations (coordinates). 

Motivated to build a mesh-independent model  incorporating variable velocity models, we hereby turn to operator-learning-based neural architectures for parametric PDEs. More specifically, the wave equation is considered a parametric PDE which takes the velocity as a variable parameter, and our goal is to map different velocity models to their corresponding wavefields, governed by the wave equation. Importantly, this work solves the wave equation in the frequency domain. Due to the scattering effects caused by perplexing subsurface structures (reflected by the perplexity of the velocity model) and the causal relationship between time steps, the wavefields in the time domain can be extremely sophisticated. In comparison, the frequency domain wavefields exhibit benign properties. Specifically, the wavefields induced by different velocity models share common structures: they are composed of layers of ``circles" centered at the seismic source position. The deformed circles and other irregularities in a wavefield directly correspond to the features of the velocity model. In the frequency domain, the wave propagation is governed by the Helmholtz equation, which is nothing but the Fourier transformed wave equation along the time axis. The Helmholtz equation has been solved with PINN in a series of papers~\cite{alkhalifah2021wavefield,song2021wavefield,song2021solving}. However, they only solved the equation on a single velocity model. 

Inspired by the seminal work~\cite{li2020fourier}, we apply the Fourier neural operator (FNO) to learn the wavefields based on variable velocity models. There are previous works  that intend to solve the Helmholtz equation using FNO~\cite{konuk2021physics,song2022high}, but these frameworks focus on other variable parameters, e.g. source location and frequency, rather than velocity. The most related work to ours is ~\cite{yang2021seismic}. This work applies FNO to perform time evolution of the wave equation based on multiple velocity models generated by the Gaussian random field, which may not yield geologically realistic features  such as layered structures and fault zones. Additionally, the relatively low frequency source wavelet has little sensitivity to the high wavenumber random scatters of the velocity models in their experiments. Hence, scattered events do not appear in the corresponding time domain wavefields under such a configuration, making the wavefield smooth and easier to predict than reality. Our work not only adapts FNO to approaching multiple realistic velocity models but also extends FNO so that it is capable of solving the Helmholtz equation in a multi-source-location and multi-frequency setting. To this end, we propose the paralleled Fourier neural operator (PFNO) where a group of FNOs approach data of different frequencies in parallel. This structure simplifies the mapping by avoiding taking the frequency as another variable parameter. The resulting wavefields predicted by PFNO tend to capture irregularities in the data of high frequency. In other words, PFNO preserves the prediction accuracy while incorporating multiple sources and frequencies. 

To validate our models, comprehensive numerical tests have been implemented over $6$ datasets from $3$ families in the OpenFWI datasets~\cite{deng2021openfwi}. The first objective is to verify whether FNO-based solvers make high-accuracy predictions on the wavefields. In the single source location and single frequency setting, an FNO gives satisfactory results from data of low frequency to data of high frequency. However, it is not as impressive when multiple sources and multiple frequencies are involved. In addition, a ForwardNet, which is based on an encoder-decoder architecture, is added for comparison. In this case, PFNO outperforms FNO and ForwardNet in terms of both the mean-squared error and the wavefield visualization. Generalization is another important factor for model evaluation. To test the generalization of PFNO, we train a PFNO for each of the $6$ datasets, and test its performance on the other $5$. This cross-dataset generalization test shows that PFNO is able to generalize to unseen velocity models that share common features with the velocity models in the training set. Moreover, the wavefields are perturbed by random noise to investigate the robustness of the model. PFNO gives consistent performances under different levels of uncertainty. The last part of the numerical experiments address the computational efficiency of PFNO. We show that PFNO can predict the wavefield corresponding a velocity model within milliseconds. Even after incorporating the training time, PFNO dominates the ``optimal" 9-point finite difference method~\cite{jo1996optimal}, a standard traditional numerical method for the Helmholtz equation, in the case of large-scale inferences. 

The paper is organized as follows. We first give some background on the Helmholtz equation and FNO in \Cref{sec:background}. \Cref{sec:method} follows to describe the FNO-based solver of the Helmholtz equation with variable velocity. Then, the performance of the models is discussed in detail in \Cref{sec:numerical}. Finally, we arrive at the conclusion in the last section. 




\section{Acknowledgment}


This work was funded by the Los Alamos National Laboratory~(LANL) - Laboratory Directed Research and Development program under project number 20210542MFR.

\vskip 0.2in
\bibliographystyle{plain}

\section{Method}
\label{sec:method}

Given multiple velocity models, we have the data denoted by $\{\tensor V^i,~\bm u^i \}_i^N$, wherein $\tensor V^i \in \mathbb{R}^{n_z \times n_x}$ is a velocity model sample on a $n_z \times n_x$ grid. Of note, the location is represented by $\bm{{x}} = (x,~z)$. In particular, $x$ and $z$ denote the  horizontal offset and the depth, respectively. The label $\bm u^i \in \mathbb{C}^{n_z \times n_x}$ is the frequency domain pressure wavefield. Our goal is to solve the Helmholtz equation numerically by learning the wavefield with FNO. For convenience, we denote the shape of a data sample by $[n_z,~n_x]$. The velocity values are given on coordinates $(x,~z)$, $x$ and $z$ each giving a coordinate sample of the shape $[n_z,~n_x]$. Since FNO incorporates the coordinates into its input, a complete input is a 3-dimensional tensor of the shape $[n_z,~n_x,~3]$, as illustrated in \Cref{fig:data_shape}. It is worth mentioning that the grid can be irregular as long as the velocity values match the coordinates. The last dimension contains $v,~x,~z$ values, making it an analog to the dimension of channels as in image data. Specifically, we have $3$ ``channels" here. 
\begin{figure}[!ht]
    \centering
    \includegraphics[width=0.75\textwidth]{figure/data_shape.png}
    \caption{Data shape: an example of $n_z = n_x = 3$ with $\Delta x = \Delta z = 100$.}
    \label{fig:data_shape}
\end{figure}

The forward modeling is formulated by the following optimization problem:
\begin{align*}
    \min_{\theta \in \Theta} \quad \frac{1}{N}\sum_{i=1}^N \left\Vert \bm \hat{\bm u}_\theta^i -  \bm u^i \right\Vert_2,
\end{align*}
with $\hat{\bm u}_\theta^i = \text{FNO}(\tensor T^i; \theta)$ and $\tensor T^i = \tensor V^i \oplus \tensor X^i \oplus \tensor Z^i$, where $\oplus$ denotes concatenation. 

In real-life applications, multiple sources at different locations are usually required. Seismic data simulated by multiple sources has a better representation (or illumination) of the velocity model compared to the single source case because different source locations provide various incident waves that illuminate the same spatial position. Besides, multi-frequency simulation is essential, as well. Lower frequency components in the data are only sensitive to the longer wavelength parts of the velocity model, which represents the smoothed background. In comparison, higher frequencies are more sensitive to the shorter wavelengths, which represents the sharp scatters in the velocity model. Thus, it requires the model to be capable of solving the Helmholtz equation with multiple source locations and multiple seismic frequencies in order to have an accurate seismic data representation of the velocity model. Accordingly, FNO incorporates $\bm x_s,~\omega$ as another two variable parameters. Regarding the source location, we use binary encoding: the source locations are encoded by $1$ whereas all other locations are $0$. On the other hand, the frequency is set as a constant across the spatial domain, since it relies on neither the coordinates nor the velocity. As a result, a data sample ends up in shape $[n_s,~n_\omega,~n_z,~n_x,~5]$. A naive way of training is to feed these 5-dimensional tensors into a single FNO. Nevertheless, different source locations and different frequencies significantly increase the complexity of the mapping between the inputs and the wavefields. Hence, we hereby introduce the paralleled Fourier neural operator (PFNO), a structure in which we stack multiple FNOs in parallel, and let each specializes in only one frequency in order to partially counteract the increase in the mapping complexity. In other words, each FNO only takes two variable parameters, the source location, and the velocity, instead of three. \Cref{fig:PFNO} depicts the data flow through PFNO. 

\begin{figure}[!ht]
    \centering
    \includegraphics[width=\textwidth]{figure/PFNO.png}
    \caption{PFNO: each FNO receives data of only one frequency.}
    \label{fig:PFNO}
\end{figure}


\section{Numerical Experiments}
\label{sec:numerical}

\subsection{Datasets}
In this section, we demonstrate the performance of FNO and PFNO using $6$ datasets from the dataset collection OpenFWI, namely \emph{CurveFault-A} \& \emph{CurveFault-B}, \emph{Style-A} \& \emph{Style-B}, and \emph{FlatVel-A} \& \emph{FlatVel-B}. More specifically, \emph{CurveFault-A} and \emph{CurveFault-B} contain discontinuities (geological faults) in the velocity models. These faults are induced by shifting the rock layers. Of note, \emph{CurveFault-B} exhibits more complicated changes of the velocities than \emph{CurveFault-A}. In contrast, \emph{FlatVel-A} and \emph{FlatVel-B} have flat layer interfaces without faults, but the velocity values are randomly distributed across different depths in \emph{FlatVel-B}, which increases the complexity. The velocity models in \emph{Style-A} and \emph{Style-B} are generated with style transfer, where the Marmousi model~\cite{brougois1990marmousi} serves as the style image while the content images are drawn from the COCO dataset~\cite{lin2014microsoft}. The difference between them is that the velocity models in \emph{Style-A} are smoother. In the above datasets, the spatial domain is defined on a $70 \times 70$ grid with $\Delta x = \Delta z = 10$ m, and all velocity values fall in the interval $[1500 \text{ m/s}, 4500 \text{ m/s}]$. \Cref{fig:velocity} displays examples from each of the datasets.
\begin{figure}[!ht]
  \centering
  \subfloat[]{
  \includegraphics[width=0.33\textwidth]{figure/velocity_examples/CurveFault_A_vel_34_0.png}}
  \subfloat[]{
  \includegraphics[width=0.33\textwidth]{figure/velocity_examples/Style_A_vel_34_0.png}}
  \subfloat[]{
  \includegraphics[width=0.33\textwidth]{figure/velocity_examples/FlatVel_A_vel_34_0.png}} \\
  \subfloat[]{
  \includegraphics[width=0.33\textwidth]{figure/velocity_examples/CurveFault_B_vel_34_0.png}}
  \subfloat[]{
  \includegraphics[width=0.33\textwidth]{figure/velocity_examples/Style_B_vel_34_0.png}}
  \subfloat[]{
  \includegraphics[width=0.33\textwidth]{figure/velocity_examples/FlatVel_B_vel_34_0.png}}
  \caption{Velocity models: (a) \emph{CurveFault-A}; (b) \emph{Style-A}; (c) \emph{FlatVel-A}; (d) \emph{CurveFault-B}; (e) \emph{Style-B}; (f) \emph{FlatVel-B}. } 
  \label{fig:velocity}
\end{figure}

Based on the velocity models, we solve \cref{eqn:time_wave} by a high-order finite difference scheme~\cite{jin2021unsupervised}. Let $p_{x, z}^t := p(t, x, z)$. The temporal derivative approximated with the second-order central difference 
\begin{align*}
    \frac{\partial^2 p}{\partial t^2} \approx \frac{p_{x,z}^{t+\Delta t} - 2p_{x,z}^{t} + p_{x,z}^{t-\Delta t}}{\Delta t^2}.
\end{align*}
To guarantee the accuracy, the Laplacian is discretized with the fourth-order central difference,
\begin{align*}
	\nabla^2 p 
	= \frac{\partial^2 p}{\partial x^2} + \frac{\partial^2 p}{\partial z^2} \nonumber  \approx \frac{1}{\Delta x^2} \sum_{i=-2}^{2} c_i p_{x+i, z}^t + \frac{1}{\Delta z^2} \sum_{i=-2}^{2} c_i p_{x, z+i}^t,
\end{align*}
where $c_{-1} = c_1 = -\frac{1}{12}$, $c_{-2} = c_2 = \frac{4}{3}$, $c_0 = -\frac{5}{2}$ are the coefficients. For convenience, we write $x+i := x + i \Delta x$ and $z+i := z + i \Delta z$. With $v_{x, z} := v(x, z), s_{x, z}^t := s(t, x, z)$, \cref{eqn:time_wave} can be discretized as follows, 
\begin{align*}
	\frac{1}{\Delta x^2} \sum_{i=-2}^{2} c_i p_{x+i, z}^t + \frac{1}{\Delta z^2} \sum_{i=-2}^{2} c_i p_{x, z+i}^t 
	- \frac{p_{x,z}^{t+\Delta t} - 2p_{x,z}^{t} + p_{x,z}^{t-\Delta t}}{v_{x, z}^2 \Delta t^2} = s_{x,z}^t.
\end{align*}
To satisfy the Courant-Friedrichs-Lewy (CFL) condition for numerical stability, we take $\Delta t = 0.001$ and let the finite difference solve to $t = 1$. The source is generated by the Ricker wavelet whose peak frequency is set at $15$Hz. Moreover, $120$ absorbing layers with damping parameters are applied around the spatial domain to avoid reflection artifacts. Lastly, we apply the Fourier transform to the $t$ dimension and convert the time domain wavefields to the frequency domain wavefields which will be the labels. Notably, the frequency domain wavefields are complex-valued. When training FNOs, we split the complex numbers into their real part and imaginary part. The training loss is then obtained by averaging the losses from the two parts.    

As mentioned in \Cref{sec:method}, we intend to incorporate multiple source locations and multiple frequencies. Particularly, the datasets have $5$ sources at the depth of $10$m below the surface. Horizontally, they are distributed evenly from $0$m to $690$m. On the other hand, the selection of the frequencies is slightly more sophisticated. With a fast Fourier transform (FFT) from the time domain, there are $1,000$ frequencies corresponding to the $1,000$ time steps. However, it is computationally prohibitive to learn the wavefields of all the frequencies. Hence, a common technique in the field of signal processing is adopted in this work. Specifically, we attend to the magnitude of the wavefields since it reflects the energy level of the frequency domain signal. Then, we only focus on the spectrum that contributes to the majority of the total energy and filter out the rest. Based on \emph{CurveFault-A}, we plot the magnitude against the frequencies in \Cref{fig:spectrum}. Over $90\%$ of the energy is contributed by the signals of the spectrum $1$Hz $\sim$ $30$Hz. Therefore, it is chosen to be the spectrum of interest. 
\begin{figure}[!ht]
    \centering
    \includegraphics[width= 0.7\textwidth]{figure/spectrum.png}
    \caption{Spectrum of interest: the magnitude is averaged over the dimension of space, source, and sample.}
    \label{fig:spectrum}
\end{figure}

\subsection{Wavefield Prediction}
\label{subsec:wavefield_prediction}
The motivation for applying FNO or PFNO is to obtain a numerical solver that can map velocity models to their corresponding wavefields. In this part, we demonstrate the performance of the models on the aforementioned datasets. For the following experiments, we use $15,000$ velocity models for training and $3000$ for testing. Moreover, we allow $600$ epochs in the training with a batch size of $16$. The initial learning rate is set to $0.0016$ which is decayed by a factor of $0.5$ every $125$ epochs. For the optimization purpose, the AdamW optimizer is applied with $\beta_1 = 0.9,~\beta_2 = 0.999$ and weight decay $0.0001$. Following the setting in~\cite{li2020fourier}, we admit $4$ Fourier layers and only $12$ Fourier modes for truncation. The other hyperparameter in FNO is the width of each Fourier layer. In PFNO, the dimensions are specified by the rule: $1)$ $32$ if the frequency is in $[1\text{Hz},~15\text{Hz}]$; $2)$ 64 if the frequency is in $[16\text{Hz},~25\text{Hz}]$; $3)$ 96 if the frequency is in $[26\text{Hz},~30\text{Hz}]$. This is consistent with the fact that the higher the frequency the more complicated the wavefield. 

Starting with a simple case of only one source location and one frequency, we train two FNOs based on the  \emph{CurveFault-A} dataset. In this experiment, the source is located on the right edge of the domain. For comparison purposes, the two FNOs are trained on data of $10$Hz and $30$Hz, respectively. In addition, the width is $32$ for the $10$Hz data, whereas we use $96$ for $30$Hz data in that it is much more challenging to fit. As \Cref{fig:wavefield_freqs} verifies, the wavefields of $30$Hz data have smaller wavelengths as well as more prominent scattering effects. Consequently, the FNO trained on $10$Hz data outperforms its counterpart in terms of the relative misfit (absolute difference between the prediction and the label). Additionally, the imaginary part has highly similar patterns to the real part regardless of the data frequency, but they differ in magnitude. 
\begin{figure}[!ht]
  \centering
  \subfloat[]{
  \includegraphics[width= \textwidth, height= 0.35\textheight]{figure/wavefield/freqs/v9_s5_10Hz.png}}\\
  \subfloat[]{
  \includegraphics[width= \textwidth, height= 0.35\textheight]{figure/wavefield/freqs/v9_s5_30Hz.png}}
  \caption{Wavefield with a single source location and a single frequency: (a) $10$Hz; (b) $30$Hz; the first row is the real part and the second row is the imaginary part. } 
  \label{fig:wavefield_freqs}
\end{figure}

The mapping between the inputs and the wavefield becomes significantly more complicated when multiple source locations and multiple frequencies are incorporated. In our numerical experiments, we find that the frequency has a greater impact on the performance of FNO than the source location. Therefore, we propose the PFNO structure where multiple FNOs are paralleled. Each of the FNOs only approaches data of a single frequency (but still of multiple source locations). Hence, the source location and the velocity are the variable parameters for each FNO. In the numerical experiments, the datasets have the aforementioned 5 source locations and the following 10 frequencies: $\{1, 3, 5, 7, 9, 12, 15, 19, 25, 30 \}\text{Hz}$. We choose more lower than higher frequencies in the experiments mainly to serve seismic imaging problems like FWI, where the lower frequency components are more important than the higher ones. It is because the lower frequencies represent the fundamental low-wavenumber part of subsurface kinematics, and lack of such information leads to the higher frequencies cycle-skipping problems~\cite{brossier2009seismic}.

Figures \ref{fig:wavefields_CurveFault},~\ref{fig:wavefields_Style}, and~\ref{fig:wavefields_FlatVel} visualize the performance of PFNO on \emph{CurveFault} datasets, \emph{Style} datasets, and \emph{FlatVel} datasets, respectively. For demonstration purposes, only the real part of the $25$Hz wavefields is shown, since the imaginary part is similar. As depicted in \Cref{fig:wavefields_CurveFault}, the geological faults induce a number of scattering artifacts that form the cobwebs in some local areas. Given smooth features in the velocity models,  \Cref{fig:wavefields_Style_A} exhibits smooth curves in the wavefields as expected. In comparison, ruffled layers manifest in \Cref{fig:wavefields_Style_B}. The irregularities in \Cref{fig:wavefields_FlatVel} reveal the reflections of energy between layer interfaces, although these interfaces are flat. Since \emph{Style} and \emph{FlatVel} have relatively simpler wavefields, PFNO performs better on these two families than on \emph{CurveFault}. In addition, the resolution of the wavefield at a certain point is inverse proportional to its wavelength and the velocity value at that point if we fix the frequency. Thus, the blurry parts are caused by the small wavenumber at the current frequency. 
\begin{figure}[!ht]
  \centering
  \subfloat[]{\label{fig:wavefields_CurveFault_A}
  \includegraphics[width= \textwidth, height= 0.22\textheight]{figure/wavefield/performance_datasets/CurveFault_A_v0_s3_25Hz.png}}\\
  \subfloat[]{\label{fig:wavefields_CurveFault_B}
  \includegraphics[width= \textwidth, height= 0.22\textheight]{figure/wavefield/performance_datasets/CurveFault_B_v0_s3_25Hz.png}}
  \caption{Real part of 25Hz wavefields - \emph{CurveFault}: (a) \emph{CurveFault-A}; (b) \emph{CurveFault-B}.}
  \label{fig:wavefields_CurveFault}
\end{figure}

\begin{figure}[!ht]
  \centering
  \subfloat[]{\label{fig:wavefields_Style_A}
  \includegraphics[width= \textwidth, height= 0.22\textheight]{figure/wavefield/performance_datasets/Style_A_v9_s1_25Hz.png}}\\
  \subfloat[]{\label{fig:wavefields_Style_B}
  \includegraphics[width= \textwidth, height= 0.22\textheight]{figure/wavefield/performance_datasets/Style_B_v9_s1_25Hz.png}}
  \caption{Real part of 25Hz wavefields - \emph{Style}: (a) \emph{Style-A}; (b) \emph{Style-B}.}
  \label{fig:wavefields_Style}
\end{figure}

\begin{figure}[!ht]
  \centering
  \subfloat[]{\label{fig:wavefields_FlatVel_A}
  \includegraphics[width= \textwidth, height= 0.22\textheight]{figure/wavefield/performance_datasets/FlatVel_A_v9_s5_25Hz.png}}\\
  \subfloat[]{\label{fig:wavefields_FlatVel_B}
  \includegraphics[width= \textwidth, height= 0.22\textheight]{figure/wavefield/performance_datasets/FlatVel_B_v9_s5_25Hz.png}}
  \caption{Real part of 25Hz wavefields - \emph{FlatVel}: (a) \emph{FlatVel-A}; (b) \emph{FlatVel-B}.}
  \label{fig:wavefields_FlatVel}
\end{figure}

As a comparison, we also implement a single FNO with a width of $96$ that takes source location, frequency, and velocity into its inputs. In addition, an encoder-decoder architecture is introduced as an image-to-image method. This architecture is inspired by InversionNet~\cite{wu2019inversionnet}. With the purpose of solving the forward problem, an encoder of $11$ convolutional layers maps the velocity model to a latent space and then a decoder also of $11$ convolutional layers maps the latent representation to the wavefield. For convenience, it is named ForwardNet. The specific structure and hyperparameters of ForwardNet can be found in~\Cref{sec:forwardnet_architecture}.  Similarly, we examine the performance of FNO and ForwardNet on the $6$ datasets and compare them with PFNO. Here, the mean-squared error (MSE) is employed to measure how PFNO performs on the testing set. As \Cref{tab:comparison} shows, PFNO consistently outperforms the other two models on all datasets. To provide an intuitive comparison, \Cref{fig:model_comparison} visualizes the predicted wavefields and the ground truth using a $25$Hz example from \emph{CurveFault-A}. Highly agitated patterns of the layer edges are recognized by PFNO, whereas FNO and ForwardNet only capture the smooth part. Moreover, the wavefields predicted by FNO turns blurry as the wave spreads rightward, but PFNO still gives sharp scatterings as the ground truth suggests. 
\begin{table}[!ht]
  \centering
  \caption{Performance of the models, measured by MSE}
    \begin{tabular}{cccc}
          & FNO   & ForwardNet & PFNO \\
    \midrule
    \emph{CurveFault-A} & 4.231e-05 & 4.441e-05 & \textbf{3.357e-05} \\
    \emph{CurveFault-B} & 2.213e-04 & 2.588e-04 & \textbf{1.959e-04} \\
    \emph{Style-A} & 2.262e-05 & 3.304e-05 & \textbf{1.713e-05} \\
    \emph{Style-B} & 1.702e-05 & 2.575e-05 & \textbf{8.478e-06} \\
    \emph{FlatVel-A} & 8.280e-06 & 1.372e-05 & \textbf{3.209e-06} \\
    \emph{FlatVel-B} & 3.316e-05 & 3.591e-05 & \textbf{2.302e-05} \\
    \end{tabular}%
  \label{tab:comparison}%
\end{table}%


\begin{figure}[!ht]
    \centering
    \subfloat[]{\label{fig:model_comparison_FNO}
    \includegraphics[width= 0.45\textwidth, height= 0.22\textheight]{figure/model_comparison_vis/FNO_v9_s1_25Hz.png}}
    \subfloat[]{\label{fig:model_comparison_ForwardNet}
    \includegraphics[width= 0.45\textwidth, height= 0.22\textheight]{figure/model_comparison_vis/ForwardNet_v9_s1_25Hz.png}} \\
    \subfloat[]{\label{fig:model_comparison_PFNO}
    \includegraphics[width= 0.45\textwidth, height= 0.22\textheight]{figure/model_comparison_vis/PFNO_v9_s1_25Hz.png}}
    \subfloat[]{\label{fig:model_comparison_label}
    \includegraphics[width= 0.45\textwidth, height= 0.22\textheight]{figure/model_comparison_vis/label_v9_s1_25Hz.png}} \\
    \caption{Model performance comparison: (a) FNO; (b) ForwardNet; (c) PFNO; (d) Ground Truth.}
    \label{fig:model_comparison}
\end{figure}

\subsection{Generalization}
The model's ability to generalize to other datasets is of importance. If a trained model only approximates certain patterns within a dataset instead of the intrinsic properties of the Helmholtz operator, it needs to be retrained once the features of the velocity model change, which offsets the benefit of fast inference. To examine the generalization of PFNO, we hereby conduct a cross-dataset generalization test. More specifically, we test the performance of each trained PFNO in \cref{subsec:wavefield_prediction} on the other 5 datasets. In Table~\ref{tab:generalization}, the diagonal marks the intra-dataset test results. As expected, the smallest error for each testing set appears on the diagonal. Since \emph{CurveFault-A} and \emph{CurveFault-B} share similar features in their velocity models, PFNO trained on either one generalizes to the other. Furthermore, PFNOs trained on them can adapt to all other datasets except \emph{FlatVel-B}, because the \emph{CurveFault} family admits features seen in the other two families, such as curves and layered structures. Interestingly, the \emph{Style-A}-trained PFNO has a mild generalization to \emph{CurveFault-A} but it fails to generalize to \emph{CurveFault-B}. Also, it can generalize to \emph{Style-B} despite sharper features in it. In contrast, its \emph{Style-B}-trained counterpart only adapts to \emph{Style-A}. This is presumably caused by the unique sharp features in \emph{Style-B}. Among the $3$ families, \emph{FlatVel} seems to be biased to the other two distributions because it contains only flat layers and no lateral variation. Thus, PFNOs trained on them are unable to generalize to the others. Although the \emph{FlatVel-B}-trained PFNO makes accurate predictions on \emph{FlatVel-A}, the opposite direction does not apply. The reason roots in the randomly distributed velocities across layers in \emph{FlatVel-B}. 
\begin{table}[!ht]
  \centering
  \caption{Cross-dataset generalization: the rows are the training sets and the columns are the testing sets; errors are measured by MSE}
    \resizebox{\textwidth}{!}{\begin{tabular}{ccccccc}
          & \emph{CurveFault-A} & \emph{CurveFault-B} & \emph{Style-A} & \emph{Style-B} & \emph{FlatVel-A} & \emph{FlatVel-B} \\
    \midrule
    \emph{CurveFault-A} & \textbf{3.357e-05} & 8.873e-04 & 2.067e-04 & 1.269e-04 & 8.963e-05 & 3.451e-03 \\
    \emph{CurveFault-B} & 1.070e-04 & \textbf{1.959e-04} & 6.748e-05 & 4.835e-05 & 8.869e-05 & 2.034e-03 \\
    \emph{Style-A} & 8.221e-04 & 1.853e-03 & \textbf{1.713e-05} & 2.834e-05 & 2.851e-03 & 6.607e-03 \\
    \emph{Style-B} & 1.232e-02 & 2.768e-02 & 2.294e-04 & \textbf{8.478e-06} & 2.268e-02 & 3.550e-01 \\
    \emph{FlatVel-A} & 3.040e-03 & 4.793e-03 & 2.377e-03 & 1.070e-03 & \textbf{3.209e-06} & 1.185e-03 \\
    \emph{FlatVel-B} & 3.015e-03 & 4.704e-03 & 2.382e-03 & 1.054e-03 & 1.185e-05 & \textbf{2.302e-05} \\
    \end{tabular}}
  \label{tab:generalization}%
\end{table}%


\subsection{Robustness}
In reality, the measurement of the seismic pressure is prone to the errors of the receivers (sensors). This requires our model to perform consistently under a certain level of uncertainty. To test the robustness of the model, we consider injecting random noise into the labels and training PFNOs with these noisy labels. In particular, zero-mean Gaussian I.I.D. random variables are added to the labels in \emph{CurveFault-A}. To control the uncertainty level, we specify different standard deviations when generating the Gaussian random variables. To ensure that the data has a similar magnitude, we select these $5$ frequencies in the training: $[11, 13, 15, 17, 19]$Hz. 

The result is exhibited in \Cref{fig:robust}. It is not surprising to observe an upward trend in the training MSE. Since the random noise constantly perturb the labels, it is more challenging to learn the mapping. As a strong contrast, the testing MSE is barely impacted by the stochasticity up to $0.01$ standard deviation which is approximately $10\%$ of the average magnitudes. On the contrary to ascending, the testing MSE even decreases slightly. An explanation is that the stochasticity actually serves as a regularization. This benefits the PFNO because the random noises prevent the optimizer from driving towards sharp minima. Hence, PFNO generalizes slightly better. 
\begin{figure}[!ht]
    \centering
    \includegraphics[width= 0.8\textwidth]{figure/robust.png}
    \caption{Robustness}
    \label{fig:robust}
\end{figure}

\subsection{Computational Efficiency}
One major motivation for applying deep learning models over traditional methods is their fast inference. In the case of a large-scale dataset, i.e., tens of thousands of velocity models, it is time-consuming to solve the Helmholtz equation with traditional methods. In comparison, FNO and PFNO predict a wavefield within several milliseconds on average. Nevertheless, the apparent drawback is that they demand hours for training. Therefore, we take both the training and inference time into consideration when comparing the efficiency of FNO and PFNO with a traditional method, the ``optimal" 9-point finite difference method~\cite{jo1996optimal}.   

The comparison is conducted in two settings: $1)$ single source location and single frequency; $2)$ $5$ source locations and $10$ frequencies. In the first case, an FNO is trained with 1 GPU, while we exploit $4$ GPUs to train a PFNO in the second case. However, both cases use only $1$ GPU for inference. Again, there are $15,000$ velocity models for training and $3,000$ for inference. In the inference, the batch size of both FNO and PFNO is set to $64$. Not surprisingly, FNO outputs a wavefield in 4.23e-04 seconds on average in case $1$. On the other hand, the ``optimal" 9-point method needs about $2.30$ seconds to solve the equation based on a velocity model. The speed advantage of PFNO is more significant with multiple sources and multiple frequencies. For the multi-source and multi-frequency experiment setup on the same velocity model, PFNO averages 9.63e-03 seconds to predict the wavefield while in contrast, the ``optimal" 9-point method takes $145.65$ seconds on average. Considering that FNO and PFNO took $2.33$ hours and $12.02$ hours for training, respectively, we consolidate them with their inference time, which makes the comparison more reasonable. As \Cref{fig:efficiency_case1} illustrates, FNO has a negligible marginal increase in the wall time as the number of velocity models increases, leading to an almost flat line. Moreover, FNO is more efficient if more than $3,648$ velocity models are to be approached, even though the training time is incorporated. A similar conclusion can be drawn from \Cref{fig:efficiency_case2}, only that PFNO is even more preferable. Provided that we have more than $298$ velocity models in the dataset, PFNO prevails over the traditional method in terms of speed. From the evidence above, one may verify the efficiency of FNO and PFNO in large-scale inferences. 
\begin{figure}[!ht]
    \centering
    \subfloat[]{\label{fig:efficiency_case1}
    \includegraphics[width= 0.495\textwidth]{figure/efficiency/case1.png}}
    \subfloat[]{\label{fig:efficiency_case2}
    \includegraphics[width= 0.495\textwidth]{figure/efficiency/case2.png}} 
    \caption{Efficiency comparison: (a) case 1; (2) case 2; the training time of both FNO and PFNO corresponds to $15,000$ training samples of velocity model.}
    \label{fig:efficiency}
\end{figure}





