Physics-informed Neural Networks with Fourier Features for Seismic Wavefield Simulation in Time-Domain Nonsmooth Complex Media
Yi Ding Su Chen Hiroe Miyake Xiaojun Li







Preprint submitted to IEEE
Shell et al.: A Sample Article Using IEEEtran.cls for IEEE Journals

0000–0000/00$00.00  2024 IEEE



Physics-informed Neural Networks with Fourier Features for Seismic Wavefield Simulation in Time-Domain Nonsmooth Complex Media
    Yi Ding, Su Chen, Hiroe Miyake, Xiaojun Li
        
Manuscript received ****, *** 2024; This work was supported in part by National Natural Science Foundation of China (Grant Nos. 52192675, U1839202), and in part by the China Scholarship Council 202306540044 (Corresponding author: Su Chen.)
        
Yi Ding is with the Key Laboratory of Urban Security and Disaster Engineering of the Ministry of Education, Beijing University of Technology, Beijing 100124, China, and also with the Earthquake Research Institute, The University of Tokyo, Tokyo, 113-0032, Japan (e-mail: dingyi18@mails.ucas.edu.cn).
        
Su Chen and Xiaojun Li are with the Key Laboratory of Urban Security and Disaster Engineering of the Ministry of Education, Beijing University of Technology, Beijing 100124, China, and Xiaojun Li is also with the Institute of Geophysics, China Earthquake Administration, Beijing 100081, China (e-mail: chensuchina@126.com; beerli@vip.sina.com).
        
Hiroe Miyake is with the Earthquake Research Institute, The University of Tokyo, Tokyo, 113-0032, Japan (e-mail: hiroe@eri.u-tokyo.ac.jp).
    5 September 2024
==================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



§ ABSTRACT 
Physics-informed neural networks (PINNs) have great potential for flexibility and effectiveness in forward modeling and inversion of seismic waves. However, coordinate-based neural networks (NNs) commonly suffer from the "spectral bias" pathology, which greatly limits their ability to model high-frequency wave propagation in sharp and complex media. We propose a unified framework of Fourier feature physics-informed neural networks (FF-PINNs) for solving the time-domain wave equations. The proposed framework combines the stochastic gradient descent (SGD) strategy with a pre-trained wave velocity surrogate model to mitigate the singularity at the point source. The performance of the activation functions and gradient descent strategies are discussed through ablation experiments. In addition, we evaluate the accuracy comparison of Fourier feature mappings sampled from different families of distributions (Gaussian, Laplace, and uniform). The second-order paraxial approximation-based boundary conditions are incorporated into the loss function as a soft regularizer to eliminate spurious boundary reflections. Through the non-smooth Marmousi and Overthrust model cases, we emphasized the necessity of the absorbing boundary conditions (ABCs) constraints. The results of a series of numerical experiments demonstrate the accuracy and effectiveness of the proposed method for modeling high-frequency wave propagation in sharp and complex media.



Physics-informed neural networks (PINNs), Seismic wave propagation simulation, Fourier feature neural networks, Absorbing boundary conditions (ABCs), Spectral bias.




§ INTRODUCTION

Improving the accuracy and stability of the inner domain and artificial boundaries in the simulation of complex wave propagation problems is a fundamental and common requirement for important inversion techniques <cit.>. For example, full waveform inversion (FWI) is widely used in oil and gas exploration to detect underground structures from artificial earthquake seismograms. By minimizing the differences between observed data and simulated data from the wave equations, FWI can provide high-resolution estimates of subsurface parameters. With the successful application of machine learning in many fields, there is increasing interest in utilizing the approximation capabilities of machine learning techniques for seismic inversion <cit.>. However, the lack of labeled training datasets is a common challenge in applying deep learning to solve most geoscience problems. Semi-supervised or unsupervised learning stands as a prospective avenue to surmount this obstacle, leveraging prior constraints to formulate unsupervised loss functions for the training of deep neural networks <cit.>.



With the universal approximation theorem of neural networks <cit.> and advances in automatic differentiation, deep learning tools are introducing a new trend, offering an alternative approach for solving partial differential equations (PDEs) through the combination of hidden layers to provide nonlinear approximations. The physics-informed neural networks (PINNs) proposed by Raissi et al. <cit.> seamlessly integrate observation and governing physics and train the neural network by minimizing a physics-informed loss function constructed based on PDE residuals and initial/boundary conditions (I/BCs). PINNs provide new perspectives and ideas for addressing the challenges faced by traditional numerical methods in complex domains <cit.>, and solving inverse problems <cit.>. In the geophysical forward and inverse problems, PINNs have been successfully applied to the Eikonal equation <cit.>, the Maxwell equation <cit.>, wave equations for isotropic and anisotropic media in both time <cit.>, and frequency domains <cit.>.

Current research on PINNs methods for modeling wave propagation behavior has received widespread attention, but there have also been certain challenges in applying vanilla PINNs to seismic wave propagation problems. Firstly, some studies have reported that PINNs suffer from point-source singularity when solving wave equations in the frequency and time domains. Alkhalifah et al. <cit.> addressed the scatter wavefield of the Helmholtz equation based on analytical background wavefields to avoid the singularity arising from high sparsity. To solve the time-domain wave equations, the early initial wavefields obtained by analytical <cit.> or conventional numerical methods <cit.> can be utilized to provide information about the source. This scheme leads to additional initial conditions loss components, which can be adjusted using a loss weight adaptive algorithm <cit.> to balance the backpropagation gradients between different loss terms.

On the other hand, neural networks tend to learn low-frequency functions and encounter difficulties in learning high-frequency functions, a phenomenon known as "spectral bias" <cit.>. Tancik et al. <cit.> introduced the Fourier feature network, which employs Fourier feature mapping to transform the input coordinates into a series of high-frequency sinusoidal waves, enhancing the neural network's capability to learn high-frequency functions. Subsequently, the Fourier feature neural networks have been introduced into physics-informed learning <cit.>, then utilized to solve the multi-frequency Helmholtz equation <cit.> and the acoustic wave equation <cit.>.

In conventional discrete numerical algorithms, the discrete equations of motion of the nodes need to be integrated step-by-step using appropriate time integration methods to ensure that the solutions are obtained in temporal order. In contrast, PINNs simultaneously predict the solution over the entire spatio-temporal domain, which may violate temporal causality and tend to converge to erroneous solutions <cit.>. The inherent implicit bias in PINNs hinders the learning process associated with long-term complex dynamic processes. A sequential training strategy via time domain decomposition has been introduced to improve the ability of PINNs to model complex wave propagation behavior <cit.>.

Specifying proper initial conditions (ICs) and boundary conditions (BCs) for a PDE is essential to have a well-posed problem. The original formulation of PINNs utilized a “soft constraints” manner to incorporate I/BCs. Emerging studies have rigorously guaranteed the implementation of hard-embedding I/BCs and avoided the issue of imbalanced gradients in multi-component loss functions <cit.>. Absorbing boundary conditions (ABCs) are commonly used to model wave propagation in infinite or semi-infinite media under classical numerical methods. At present, a consensus has not been reached regarding the necessity of including ABCs in the solution of the wave equation using PINNs. When addressing the Helmholtz equation in non-smooth media, the incorporation of perfect matching layer conditions into the loss function is proposed to enhance the coupling of the real and imaginary components of the wavefield <cit.>. Additionally, for modeling seismic wave propagation in 2D elastic media, an absorbing boundary condition is integrated into the network as a soft regularizer to address truncated boundaries <cit.>. However, the potential of PINNs for modeling seismic wavefields in the time domain has not been fully explored, especially concerning complex media and high-frequency wave behavior in the time domain.

Motivated by the foregoing developments, this work proposes a novel unified PINNs architecture for time-domain seismic wavefield simulation in complex media. To alleviate the singularity issue of point sources, we first inject the source using the smoothed Gaussian spatial mapping to approximate the Dirac delta function. Additionally, a fast wave velocity model prediction module is proposed to combine with the stochastic gradient descent (SGD) training strategy, randomly sampling training points at each training step. This allows PINNs to effectively capture information about the seismic source.

To ensure an accurate implementation of the zero-initial condition, the hard embedding scheme proposed by Sethi et al. <cit.> is introduced. We carry out a series of ablation experiments to compare the impact of activation functions and training strategies on the proposed framework, as well as the performance of different families of distributions of Fourier feature parameters. Taking the Marmousi and Overthrust models as examples, we show that non-smooth complex medium velocity models lead to inaccurate wavefields when ABCs are not considered. The paraxial approximation technique of the acoustic wave equations is introduced into the loss function as a soft regularization to solve the wave propagation problem in the infinite domain. Combining ABCs and time-domain decomposition training strategies, we demonstrate the potential and versatility of the proposed framework.



§ THEORY



 §.§ Problem Setting

In this paper, we aim to investigate the potential of Fourier feature PINNs for modeling wave propagation in complex velocity media. The propagation of seismic waves can be simulated in the time domain by solving the 2D acoustic wave equation as follows:

    ∂^2 u(𝐱, t)/∂ t^2    = c^2(𝐱) ∇^2 u(𝐱, t)+s(t) G(𝐱, 𝐱_s) 
    
    s(t)     = M_0(1-2(π f_0(t-t_0))^2) exp((π f_0(t-t_0))^2) 
    
    G(𝐱, 𝐱_s)     =exp(-1/2𝐱-𝐱_s/α_2^2)

where 𝐱={x,z} is a vector of spatial coordinates for 2D media, 𝐱_s={x_s,z_s} denotes the coordinate of the source. c(𝐱) is the velocity model and u(𝐱,t) is the displacement wavefield in the time domain. The source time function s(t) is injected via a smooth space dependent field G(𝐱,𝐱_s) of Gaussian shape. The kernel width α controls the spatial mapping range of the point source, M_0 is the amplitude. A Ricker wavelet is used as the source time function, with a center frequency of f_0 and a delay time of t_0=1/f_0.

The initiation of a point-like source at a single grid point can be readily accomplished utilizing the finite differences method. However, this is no longer suitable within the PINNs method. Similar issues are observed with pseudo-spectral methods called the Gibbs phenomenon, as the Fourier transform of a spike-like function creates oscillations that damage the accuracy of the solution <cit.>. We use the Gaussian function in (<ref>) to define the space-dependent part of the source, injecting it with a smooth spatial distribution of Gaussian shape to alleviate singularity issues associated with point sources.



 §.§ Hard embedding initial conditions in PINNs

Differing from data-driven neural networks that require a large amount of data to establish the mapping between input and output, the optimization process of PINNs relies only on known physical principles, without the need for any labeled data. Considering the 2D acoustic wave equations described by (1), the residual ℛ_r(t_r^i,𝐱_r^i) and loss function ℒ_r(θ) corresponding to the PDEs are


    ℛ_r(t_r, 𝐱_r)    =∂^2û_θ/∂ t_r^2-c^2∇^2û_θ+s(t_r) G(𝐱_r, 𝐱_s), 
    ℒ_r(θ)    =1/N_r∑_i=1^N_r|ℛ_r(t_r^i, 𝐱_r^i)|^2,

where û_θ is the unknown potential solution approximated by the deep neural networks. {t_r^i, 𝐱_r^i}_i=1^N_r denotes the collocation points of the PDEs residual, with the number being N_r.

In the vanilla PINNs framework, I/BCs are imposed as loss terms in the target function in a "soft constraint" manner. The gradient imbalance problem of different loss terms may lead to the failure of neural network training, and the weight parameters of different loss terms must be carefully determined. Meanwhile, the optimization process cannot guarantee that the soft constraint loss term is completely zero. The output of NNs is directly the displacement solution of the second order wave equations, which can directly satisfy the initial and boundary conditions related to u. To ensure a unique solution to the wave equations, it is imperative to rigorously define the initial conditions. Before the seismic source excitation, the system is in a quiescent state with zero initial conditions. The neural network output f_θ(t,𝐱) can be transformed as

    û_θ(t, 𝐱)=f_θ(t, 𝐱) · t^2,

which guarantees that û_θ=∂û_θ / ∂ t=0 at t=0, thus realizing hard embedding with zero initial conditions. Here t^2 is used as the approximate distance function (ADF) and is not unique. The employed ADF must satisfy two required criteria: 1) The ADF must be equal to zero at the initial moment t=0, and 2) for any t>0, the ADF and its gradient must be smooth and non-zero.



 §.§ Fourier Feature Mapping

With the Neural Tangent Kernel (NTK) theory <cit.>, the "spectral bias" pathology of neural networks can be rigorously analyzed and elucidated. The neural network is biased to first learn the target function along the eigendirections of NTK with larger eigenvalues, and then learn the remaining components corresponding to the smaller eigenvalues <cit.>. However, NTK theory suggests that this is because standard coordinate-based multilayer perceptrons (MLPs) correspond to kernels with a rapid frequency falloff, which effectively prevents them from learning the high-frequency functions <cit.>.

In PINNs methods, MLPs are generally used as universal approximators to learn the solutions of PDEs. The Fourier feature NNs accept the coordinates as input and map them to a feature space with a Fourier feature mapping before passing them to the hidden layer. Let v={t,𝐱}∈ℝ^d denote the input coordinates. The Fourier feature mapping γ(v) performs a high-dimensional frequency mapping on the input coordinates, which is then fed into the MLPs. Following the original formulation of Tancik et al. <cit.>, the Fourier feature neural network is

    γ(v)     =[[ cos (2 π𝐁v); sin (2 π𝐁v) ]], 
    H^(1)    =ϕ(W^(1)·γ(v)+b^(1)), 
    H^(l)    =ϕ(W^(l)·H^(l-1)+b^(l)),    l=2, ⋯, L, 
    f_θ(v)     =W^(L+1)·H^(L)+b^(l),

where W^(l)∈ℝ^d_l× d_l-1, 𝐛^(l)∈ℝ^d_l are the trainable weight and bias of the lth layer, respectively. ϕ denotes the nonlinear activation function. Fourier basis frequency 𝐁∈ℝ^m× d can be sampled from different distribution families (Gaussian, Laplacian, uniform), leading to different ways of embedding Fourier features. For example, in a Gaussian Fourier feature NNs, each entry of 𝐁 is i.i.d. sampled from a normal distribution 𝒩(0,σ^2) with the standard deviation σ>0. m is the number of basis functions in the set of Fourier basis functions. m is the number of basis functions in the Fourier basis set. During the network training process, 𝐁 can be kept constant or optimized alongside the weights and biases as the trainable parameter θ. In all the cases considered in this work, we set 𝐁 as a fixed parameter and let m=256.

The standard deviation σ is a hyperparameter that directly governs the scale of the distribution of the Fourier base frequencies and the resulting eigenspace of the NTK. It is usually specified empirically or requires an expensive hyperparameter search to determine. The performance of FF-PINNs is strongly influenced by the hyperparameter σ, and optimal performance is only achieved when σ is tuned to a certain range. However, for specific problems, we can obtain prior frequency information for the specific problem from existing boundary conditions, source time functions, or observed data to avoid costly wide-range hyperparameter searches. For specific problems, prior frequency information can be obtained from available I/BCs, source time functions, or observations to avoid extensive and expensive hyperparameter searches.



 §.§ Absorbing boundary conditions constraints

ABCs play a crucial role in the numerical simulation of wave propagation problems. To achieve the goal that all outgoing waves are not reflected by boundaries, the computed boundary motions should take into account all the complexities of the outgoing waves. The paraxial boundary conditions are derived by rational approximation of the 2D acoustic-wave dispersion relation <cit.>. For example, ck_z / ω=-√(1-s^2) with respect to the variable s=ck_x=ω in the negative z direction. Thus, the paraxial approximation technique can be strictly used for constructing ABCs for 2D acoustic wave equations.

In PINNs, it is straightforward to impose paraxial-approximation-based ABCs at the boundary in a soft constraint manner. ABCs in continuous form are treated directly in PINNs, the relevant derivatives are computed by automatic differentiation without the need for numerical discretization. In this study, a second-order expression of the Clayton-Engquist acoustic paraxial approximation boundary <cit.> is considered. The computational model is a rectangular region, the coordinate origin is located in the upper left corner of the model, and the boundary directions are parallel to the x and z coordinate axes, respectively. Then the residuals of the second-order Clayton-Engquist boundary conditions can be expressed as 

x -axis positive:

    ℛ_x l=∂^2û_θ/∂ x_l∂ t_c+1/c∂^2û_θ/∂ t_c^2-c/2∂^2û_θ/∂ z_c^2,

x -axis negative:

    ℛ_x b=∂^2û_θ/∂ x_b∂ t_c-1/c∂^2û_θ/∂ t_c^2+c/2∂^2û_θ/∂ z_c^2,

z -axis positive:

    ℛ_z l=∂^2û_θ/∂ z_l∂ t_c+1/c∂^2û_θ/∂ t_c^2-c/2∂^2û_θ/∂ x_c^2,

z -axis negative:

    ℛ_z b=∂^2û_θ/∂ z_b∂ t_c-1/c∂^2û_θ/∂ t_c^2+c/2∂^2û_θ/∂ x_c^2 .


The loss function now consists of the PDEs residual loss defined in (3) and the ABCs residual loss ℒ_ce,

    ℒ(θ)     =λ_rℒ_r(θ)+λ_c eℒ_c e(θ)=λ_r/N_r∑_i=1^N_r|ℛ_r|^2
        +λ_c e/N_c e∑_i=1^N_c e (|ℛ_x l|^2+|ℛ_x b|^2+|ℛ_z l|^2+|ℛ_z b|^2)

where {x_b,z_b,x_l,z_l} denotes the collocation points from the spatial boundaries with the number being N_ce, the subscript b is the lower bound and l is the upper bound. x_c, z_c and t_c are randomly sampled collocation points across the x, z and t axes, respectively. Fig. <ref> illustrates the distribution of boundary collocation points for a rectangular domain with ABCs imposed. λ_r and λ_ce are the weight hyperparameters of the corresponding loss function terms, which are set to 1 in the subsequent cases. Algorithm <ref> demonstrates the training pipeline of the proposed framework.









 §.§ Experimental Setup

Activation functions. The choice of activation functions has a significant effect on the training dynamics and task performance. Popular activation Rectified Linear Unit (ReLU) is deficient for high-order PDEs since its second-order derivative is zero. The utilization of the hyperbolic tangent (Tanh) activation function has been prevalent in previous research pertaining to PINNs. However, it has been observed that Tanh encounters challenges when tasked with learning high-frequency and periodic functions <cit.>. In the field of Implicit Neural Representation (INR), various activation functions have been suggested to help coordinate-MLPs encode high-frequency signals, including periodic (sinusoidal <cit.>) and nonperiodic (e.g., Gaussian <cit.>) functions. In addition, Swish <cit.>, an activation function discovered by combining exhaustive and reinforcement learning-based search, achieves impressive performance in many challenging tasks. We carry out a detailed experimental comparison in Section <ref> to provide guidance on the choice of activation functions.

Random sampling. In comparison to batch gradient descent (BGD), stochastic gradient descent (SGD) <cit.> significantly reduces the memory requirements and the computational cost of each iteration. Moreover, according to our observation, the use of random sampling plays an important role in the training efficiency and model performance of the selected points, whereas the use of full-batch gradient descent to train PINNs may lead to overfitting of PDE residuals. The continuously changing random sampling across the entire domain ensures that the neural network gains sufficient understanding of the point source region, even if it is sparse relative to the entire spatial domain, enabling PINNs to fully capture information about the seismic source. Section <ref> introduces the velocity surrogate model combined with SGD algorithm. A detailed comparison of the performance between SGD and BGD algorithms will be conducted in the ablation study in Section <ref>.

Optimizer. In this study, only Adam <cit.> was used as the default optimizer, random sampling from the spatio-temporal domain at each epoch. We set the initial learning rate to 5e^-3 and an exponential decay with a decay rate of 0.9 for every 1000 decay epochs. Since the L-BFGS <cit.> strongly relies on the historical gradient to approximate the inverse Hessian matrix, when these gradients are computed using collocation points the process can be unstable. The stable implementation of quasi-Newton updating in a multi-batch setting continues to garner significant attention in the academic literature <cit.>. Finally, dense layers will be typically initialized using the Glorot scheme <cit.>.

Time domain decomposition strategy. In this study, we have employed a time-domain decomposition strategy to address the wave propagation problem in complex media. The time-domain decomposition decomposes the difficult optimization problem over the whole time period into several relatively simple problems, which significantly reduces the optimization difficulty of learning the complete evolution of the dynamical system. Specifically, if the entire simulation time domain is [0, T], the solution of PDEs is sequentially learned in the time period [0, t_i](i =1, 2, … and 0<t_1<t_2<⋯ T). The training parameters of the previous time period are passed to the next time period to provide good parameter initialization.

Evaluation. Predictions û_θ for PINNs/FF-PINNs and ĉ_θ for the wave velocity surrogate model are evaluated using the relative ℓ_2 error

    ϵ_ℓ_2=√(∑_i=1^n(ŷ_θ-y_ref)/∑_i=1^n y_ref)

where ŷ_θ is the prediction result of the neural networks. and y_ref represents the ground truth. n is the number of test sampling points.

In order to assess the accuracy of the proposed method for solving the wave equation, the results of the finite difference method (FDM) are used as reference data. We use unsplit convolutional perfectly matched layer (PML) boundary condition <cit.> for staggered-grid second-order finite difference modeling of the time-dependent 2D acoustic wave equation. For the infinite medium wave propagation problem, we add 10-grid PML boundaries to each of the four boundaries of the rectangular domain to prevent spurious boundary-reflected waves.



 §.§ The Fourier feature velocity surrogate model

Previous studies have avoided the tendency of neural networks to converge to a zero solution by ensuring that sufficient collocation points are added near the source location. However, this approach results in a significant amount of computation and may lead to erroneous attention mechanisms during the training process, as the NNs excessively focus on optimizing around the source while neglecting other areas. According to our experiments, random sampling in the spatio-temporal domain in each iteration and applying SGD can sufficiently learn the information of the source, thus avoiding additional sampling in the region around the source.

Due to the heterogeneity of the model, the wave velocity corresponding to the randomly sampled spatial coordinates must be obtained quickly in each epoch to participate in the calculation of the PDE residual loss. To solve this problem, we use a supervised neural network for fast prediction of wave velocity models. The velocity model prediction network is trained using coordinate-velocity labeling pairs before the simulation starts. The input to the network is the spatial coordinates 𝐱={x,z} and the output is the corresponding wave velocity ĉ_θ. Upon completion of the velocity network training, the prediction of velocity can be executed in each subsequent PINNs/FF-PINNs training epoch, rendering the impact on training time negligible.

The predictive efficiency and accuracy of the velocity model neural network are crucial to the overall workflow. Since neural networks tend to learn low-frequency functions, a reasonable conjecture is that they may struggle to learn the detailed structure of complex velocity models. We use a comparative experiment with the Overthrust velocity model to demonstrate the importance of Fourier feature mapping. The Fourier feature NNs and fully connected NNs followed the same hyperparameter settings. We trained 5-layer neural networks with 20 neurons per hidden layer for 100,000 epochs using the Adam optimizer, the activation function used is Swish. Scaling of the wave velocity and spatial coordinates by the same proportion was implemented according to the principle of wave velocity normalization.






Table <ref> shows the comparison of the relative ℓ_2 errors in the predictions of fully connected NNs and Fourier feature NNs with different σ. As can be seen from Table <ref>, the using of Fourier feature mapping improves the prediction accuracy by about an order of magnitude. The relative ℓ_2 error of the prediction of the fully connected NNs is 2.166e^-2, while that of the Fourier feature NNs with σ=15 is 1.483e^-2. A comparison of the absolute errors of the predictions of fully connected NNs and σ=15 Fourier feature NNs with the true velocity model is shown in Fig. <ref>. For complex velocity models with sharp interfaces, fully connected NNs are unable to learn high-frequency information and thus give blurred predictions because of spectral bias pathology. On the other hand, the Fourier feature NNs showed low sensitivity to sigma in the tests, achieving low errors (ϵ_ℓ_2<2e^-3) over a wide range (σ∈{5, 25}). The Fourier feature velocity surrogate model proposed in this section provides a robust and accurate prediction of the medium model for the next stage of the solution process.



§ ABLATION STUDY

We conduct a detailed ablation study through a homogeneous medium case, which is used to evaluate several aspects that researchers and practitioners should consider when solving the wave equation using FF-PINNs. The discussion highlights the importance of choosing appropriate activation functions, random sampling strategies, and scaling of Fourier feature mappings.

The region of the computational model is x∈[0, 600] m, z∈[0, 600] m, with c=500 m/s. The total computation time is 0.9 s. Since the velocity spatial independence in the case of homogeneous media, it can be explicitly defined during the iterative process, obviating the need for the Fourier feature velocity surrogate model mentioned in Section <ref>. We do not consider ABCs in this case, thus the loss function contains only the PDE loss term. In the vanilla PINNs method and FF-PINNs, a fully connected neural network with 5 hidden layers and 50 neurons per layer was used. The number of collocation points for the PDEs residual is N_r=3,0000. The results are averaged from 5 independent trials optimized by the Adam optimizer for 10, 000 iterations.



 §.§ Activation functions and random sampling strategies

The central frequency of the Ricker source time function considered is f_0=10 Hz, and α=0.02. We used the SGD algorithm to evaluate the performance of some of the activation functions discussed in Section <ref>, and show all the results in Table <ref>. It is worth noting that the commonly used Tanh activation function performs poorly (ϵ_ℓ_2=0.6212) even when sigma is set to the proper value. In addition, Swish outperforms Gaussian and Sin in all cases, with Gaussian slightly outperforming Sin. The best performance (ϵ_ℓ_2=0.0398) was achieved with Swish activation when σ=1 for the Gaussian Fourier feature network. As discussed in Section <ref>, the BGD strategy cannot be used in this problem, probably due to insufficient learning of information about the point sources. The results in Table <ref> validate this finding. Fig. <ref> shows a comparison of the wavefield snapshots of the predicted solutions of FF-PINNs with the FDM results for five moments. All subsequent experiments herein employ Swish activation and SGD training strategy.





 §.§ Fourier mapping

The parameter 𝐁 determines the preferred frequency range learned by the Fourier feature network described by (5), and thus the reasonable selection of the distribution of 𝐁 is a prerequisite for accurate prediction by FF-PINNs. In Gaussian Fourier feature neural networks, the scale factor σ controls the range of the 𝐁 distribution. Previous works that utilized FF-PINNs to simulate the wave equation in the frequency <cit.> and time domains <cit.> sampled 𝐁 from a uniform distribution with zero as the axis of symmetry. We use Fourier feature mappings sampled from different distribution families (Gaussian, Laplacian, and uniform) and sweep over the standard deviation of each distribution. The Gaussian distribution is 1/√(2π)σe^-(x-μ)^2/2σ^2, where σ is the standard deviation. The Laplacian distribution is 1/2be^-|x-μ|/b, where b>0 is a scale parameter, which is sometimes referred to as the “diversity”. In the Laplacian distribution, the standard deviation σ=√(2)b. The standard deviation of a uniform distribution with distribution range [-b_max,b_max] is b_max / √(3). Building on the previous section, a source with f_0=10 Hz, α=0.02 was considered in this section.

Fig. <ref> shows the results of hyperparameter sweeps. The scatter plots of the results induced by three distribution families exhibit a characteristic "U-shaped" pattern. Both excessively small and excessively large values of σ lead to comparatively large prediction errors. Lower σ may make it difficult to learn information about point sources, while higher coding frequencies can introduce salt-and-pepper artifacts. Consequently, both excessively small and large values result in comparatively large prediction errors. Ideally, an appropriate σ should be selected such that the bandwidth of NTK matches that of the target signals. This not only accelerates the training convergence but also improves the prediction accuracy. However, the spectral information of the solution may not be accessible when solving forward PDEs. In this work, we perform a proper hyperparameter search starting from σ=1. From the ablation study in this section, it can be observed that the exact sampling distribution family is much less important than the distribution’s scale (standard deviation), consistent with the findings of Tancik et al. <cit.> in their experiments on 1D target signals. In the subsequent experiments, we only use Gaussian Fourier features PINNs to model the wave forward propagation.








§ EXPERIMENTAL RESULTS

In this section, we evaluate the representational capacity of the designed FF-PINNs architecture in modeling the propagation of acoustic waves through complex infinite media. Before conducting FF-PINN simulations, we construct a velocity surrogate model using coordinate-velocity label data. We train Gaussian Fourier feature NNs with 5 hidden layers, each containing 20 neurons, employing an Adam optimizer with a learning rate of 5e^-3 for 100,000 epochs, where σ=15. All the numerical implementations are coded in Jax <cit.> and performed on an NVIDIA Tesla V100 GPU card (32G) in a standard workstation.

The accuracy impact due to ABCs is shown in Table <ref>, where the number of the ABCs collocation points sampled randomly from every boundary in each training iteration is N_ce=2000. The term "w/ ABCs" in Table <ref> refers to the second-order Clayton-Engquist acoustic paraxial approximation boundary applied on all four boundaries of the rectangular domain.






 §.§ 4-layer velocity model

A 4-layer velocity model case is designed to assess the ability of FF-PINNs to learn the wavefield solutions in horizontal layered media. As shown in Fig. <ref>a, the computational model is a rectangular domain with a width of 1.2 km and a height of 1.2 km. We have set up four layers distributed evenly along the depth, with velocities of 0.6 km/s, 0.8 km/s, 1.0 km/s, and 1.4 km/s respectively. A Ricker source excitation with f_0=5 Hz and α= 0.01 is applied at {x_s,z_s}={0.6, 0.24} km in the spatial domain. The total calculated time for wave propagation is 1.2 s.



The FF-PINNs contain 5 hidden layers, 80 neurons per layer, and a standard deviation σ=1. The PDE residual collocation points with N_r=80,000 were randomly sampled from the spatio-temporal domain in each training iteration. Fig. <ref> illustrates the comparison between the results from the staggered-grid finite difference method and FF-PINNs, after 30,000 training epochs using the Adam optimizer. Seismic waves propagating through layered velocity interfaces give rise to reflection and transmission phenomena, as well as compression and expansion of the waves. Under the same experimental setup, only considering the effect of ABCs soft regularization constraint. The relative ℓ_2 error of the FF-PINNs prediction with ABCs is 0.0614, while the relative ℓ_2 error without ABCs is 0.0763. Since the wave velocity model is not complex, FF-PINNs without considering ABCs can also learn the behavior of transmitted outgoing waves. Moreover, the additional ABCs loss term may slow down the optimization of the PDEs loss term in the early training epochs, which can explain the small difference in the errors. Fig. <ref> illustrates the comparison between the FF-PINNs predictions incorporating ABCs and the results obtained through FDM. Upon examining the wavefield snapshots presented in Fig. <ref>, it becomes evident that the FF-PINNs adeptly capture distinct physical phenomena occurring near the medium interfaces.





 §.§ Random mixed Gaussian velocity model

We further increase the complexity of the wavefield by using three source terms with different α values and considering a random mixed Gaussian velocity model. Based on a background wave velocity of c_0(𝐱)=1.0 km/s, the following random mixed Gaussian model is incorporated

    G_c(𝐱)=M_cexp(-1/2𝐱-𝐱_c/α_c_2^2),

where M_c represents the amplitude, which is set as a random number between [-3, 3]. The location 𝐱_c is uniformly distributed between [0, 1]. The α_c is set to 1. The resulting random mixed Gaussian velocity model (c_0(𝐱)+G_c(𝐱)) is illustrated in Fig. <ref>b.



We consider the Ricker wavelet time function with a duration of 1.0 s and a central frequency of 10 Hz. The kernel width α of the three sources {s_1,s_2,s_3} are 0.02, 0.03, and 0.025, and the source locations 𝐱_c={x_c,z_c} are {0.3,0.3} km, {0.5, 0.5} km and {0.65, 0.65} km, respectively. The FF-PINNs with standard deviation of σ=1 is used to approximate the displacement solution. The neural networks consist of 8 hidden layers with 80 neurons per layer. During each training iteration, the number of PDE residual collocation points randomly sampled from the spatio-temporal domain is N_r=80,000. After 30,000 epochs of training using the Adam optimizer, the relative ℓ_2 errors of the predicted solutions for FF-PINNs with and without ABCs are 0.0676 and 0.0976, respectively. Table <ref> shows the corresponding comparison. Fig. <ref> shows the comparison between the FF-PINNs prediction with ABCs and the FDM results. It is clear that FF-PINNs learn the complex wave propagation dynamics generated by the three sources, and the predicted solution matches the ground truth very well.



 §.§ 2D Marmousi model and Overthrust model

Next, we apply the proposed framework to the extracted Marmousi model and Overthrust model, as shown in Fig. <ref>. To test the capability of FF-PINNs in handling velocity interfaces with sharp variations, we did not apply any smoothing to the velocity models. For the extracted Marmousi model, the model size is 176 × 138, and the spatial sampling interval is 16 m in both x and z directions. The source time function is the Ricker wavelet at the center of the model with f_0=10 Hz and α= 0.01 located. For the extracted Overthrust model, the model size is 201×100 and the spatial sampling interval is 25 m. The center frequency f_0 of the Ricker wavelet is 15 Hz and α=0.02. The total computational time for wave propagation in both cases is 1.0 second.





The FF-PINNs with standard deviation of σ=1.5 is used to approximate the displacement solution. The neural networks consist of 5 hidden layers with 128 neurons per layer. The number of PDE residual collocation points randomly sampled from the spatio-temporal domain is N_r=80,000.

Due to the increased complexity of the Marmousi and Overthrust models under consideration, the sequential training strategy of time-domain decomposition introduced in Section <ref> is applied to both of these cases. This training strategy is consistent with the laws of physics for the time-dependent dynamical systems. We initially conducted 5,000 pre-training epochs using the Adam optimizer within the time range of [0, 0.3] s, followed by 10,000 epochs of training within the time ranges of [0, 0.4] s and [0, 0.5] s, and finally, 20,000 epochs within the range of [0, 0.6] s to obtain the solution for the entire time domain.

We show in Table <ref> the accuracy comparison between the Marmousi case and the Overthrust case with and without ABCs. With ABCs, the relative ℓ_2 errors are 0.0781 and 0.0813 for the two cases, respectively. However, without ABCs, the errors significantly increase to 0.1806 and 0.2128, respectively. Figs. <ref> and <ref> show the snapshots of the predictions of FF-PINNs (with and without ABCs) compared to the FDM-simulated wavefields for the 5 moments of the Marmousi and Overthrust models. Compared to the incident waves, the refracted and reflected waves are much smaller in scale. This poses a significant challenge for fully-connected neural networks as approximators in PINNs. As can be observed, the continuous fine-grained time-domain decomposition training strategy and the soft regularization of ABCs have enabled FF-PINNs to capture the details of waves, yielding high-fidelity wavefield simulation results. When the ABCs soft regularization constraint is not imposed, spurious reflected waves are found to be generated around the boundaries. This leads to very large errors near the boundaries, which further extend to the full spatial domain, severely affecting the accuracy performance of FF-PINNs. On the contrary, the proper application of ABCs can effectively attenuate or even eliminate unreasonable errors at the boundaries, leading to a clear and accurate wavefield solution. This provides strong evidence supporting the necessity of imposing ABCs within the proposed framework.





§ DISCUSSION

Boundary conditions. This work focuses on the necessity and applicability of ABCs in solving 2D acoustic equations in infinite complex media using FF-PINNs. Free surface boundary conditions are not imposed in this work, mainly because of the possible failure of soft-constrained zero-stress boundary conditions <cit.>, which is still an open problem to be solved by the community. Hard-embedded boundary conditions are a reasonable and promising approach. For second-order wave equations, hard embedding of homogeneous Dirichlet boundary conditions (fixed boundaries) is straightforward for PINNs. However, hard embedding of the homogeneous Neumann boundary conditions cannot be achieved by directly modifying the neural network output. Developments in mesh-free methods and the deep learning community provide theoretical support for potential implementations <cit.>.

Efficacy of sharp velocity models. When dealing with non-smooth complex media, PINNs will tend to give smooth wavefields. This may be attributed to two reasons: 1) “spectral bias” pathologies that make learning high frequencies challenging, and 2) commonly used activation functions, such as Tanh, tend to produce smooth output results <cit.>. The FF-PINNs and Swish activation function used in this work significantly enhanced the efficacy and accuracy of neural networks in handling sharp velocity models. It is worth mentioning that updating the randomly sampled collocation points (SGD strategy and velocity surrogate model) at each training epoch is crucial for mitigating point source singularity and learning the details of small scattered and reflected waves. Thus, the proposed approach is an extensible, unified framework with high resolution and reasonable accuracy. Extension to time-domain elastic wave equations solving is a natural progression, and we will present related research in future work.

Computational efficiency. The second-order ABCs in our proposed framework involve additional second-order derivative calculations, which will increase the training cost to some extent. Nevertheless, the previous numerical case comparisons have provided evidence for the necessity of imposing ABCs. Recent studies have added source location as an additional input parameter to the network <cit.> and explored transfer learning <cit.> to improve simulation efficiency. However, computational efficiency remains a major limitation for the application of PINNs methods to large-scale computational scenarios. PINNs method requires the evaluation of PDEs residuals over a large number of sampling points on the internal domain, and the computation of higher-order derivatives through automatic differentiation can be time-consuming. Discrete learning methods <cit.> and neural spectral methods that learn PDE solutions in the spectral domain <cit.> have the potential to alleviate the computational burden of learning the propagation of seismic waves on a large scale.

Better frequency manipulation. The conventional initialization method for Fourier feature NNs based on empirical or hyperparameter sweep should still be used, which is a major shortcoming of this work. If observations are available, the initial values of the Fourier feature parameters can be manipulated based on the spectrum information of the target function. Furthermore, our framework embeds the spatio-temporal coordinates into the high frequency space with the same scale factor σ. However, the "frequencies" in the temporal and spatial domains may be different, or even the gap between them may be significant. In such cases, spatio-temporal multi-scale Fourier feature embedding <cit.> may be a promising idea. The frequency of the Fourier feature mapping determines the frequency of the NTK feature vectors, which can provide a more fundamental guidance of the training dynamics through the perspective of the NTK theory <cit.>. In conclusion, how to rigorously establish a general theory of the relationship between frequency and Fourier features is an open problem that requires theoretical support from more future research.



§ CONCLUSION

In this study, we present a unified framework of Fourier feature PINNs for modeling high-frequency wave propagation in sharp and complex media. The accuracy comparison of the activation functions and sampling strategies in dealing with the point source problem is discussed in detail. The combination of the Swish activation function and the SGD strategy achieves the best performance in all the experiments. The pre-trained wave velocity surrogate model provides fast and accurate wave velocity predictions in each training epoch. Fourier feature neural networks sampled from different families of distributions (Gaussian, Laplace, and uniform) have similar error distribution patterns. The standard deviation of the distribution family determines the range of coordinate frequency mapping and is an important parameter affecting the accuracy of FF-PINNs. For the case of non-smooth complex media, the imposition of ABCs avoids spurious boundary-reflected waves and significantly improves the ability of NNs to learn detailed waves. Realistic assessments demonstrate the efficacy of the proposed framework compared to the vanilla PINNs, particularly in scenarios involving high-frequencies and non-smooth, heterogeneous real-world models.



§ ACKNOWLEDGMENTS

The authors are grateful for the research funding provided by the National Natural Science Foundation of China (NSFC, Grant Nos. 52192675, U1839202).


IEEEtran





