Learning Dominant Wave Directions For Plane Wave Methods For High-Frequency Helmholtz Equations

We present a ray-based finite element method (ray-FEM) by learning basis adaptive to the underlying high-frequency Helmholtz equation in smooth media. Based on the geometric optics ansatz of the wave field, we learn local dominant ray directions by probing the medium using low-frequency waves with the same source. Once local ray directions are extracted, they are incorporated into the finite element basis to solve the high-frequency Helmholtz equation. This process can be continued to further improve approximations for both local ray directions and the high frequency wave field iteratively. The method requires a fixed number of grid points per wavelength to represent the wave field and achieves an asymptotic convergence as the frequency $\omega\rightarrow \infty$ without the pollution effect. A fast solver is developed for the resulting linear system with an empirical complexity $\mathcal{O}(\omega^d)$ up to a poly-logarithmic factor. Numerical examples in 2D are presented to corroborate the claims.


Introduction
We consider the Helmholtz equation: plus boundary (or radiation) conditions, where ω is the frequency, c(x) > 0 is the wave speed, and f (x) is the source distribution. The numerical solution of the Helmholtz equation (1) in the high-frequency regime, i.e., ω 1, is notoriously hard to compute. From Shannon's sampling principle [25], to resolve a general wave field oscillating at frequency ω, a mesh size h = O(ω −1 ) is necessary and sufficient. Hence the intrinsic degrees of freedom is O(ω d ), implying that the theoretical optimal overall complexity to solve (1) is O(ω d ). In general, an overall complexity of optimal order is difficult to achieve in practice due to two typical challenges: how to design a discretization method that can achieve both accuracy and stability with the minimum order of degrees of freedom; and how to solve the linear system issued from the discretization with an almost linear complexity as the frequency becomes large.
Numerically, it is difficult to design a sparse discretization that can achieve both accuracy and stability under the condition ωh = O(1) as ω becomes large. This is usually called pollution effect in error estimates for finite element methods [2,1,14], i.e., the ratio between numerical error and best approximation error from a discrete finite element space is ω dependent.
A popular strategy to solve the high-frequency Helmholtz equation (1) is to incorporate appropriate oscillatory behaviors into the basis for Galerkin methods in order to represent the oscillatory wave field governed by the Helmholtz equation more accurately and efficiently. The methods based on this strategy are usually called wave based methods. The key issue for this strategy is how to design the oscillatory basis. Several approaches have been developed along this direction. One approach is to incorporate plane waves with a predetermined even distribution in directions into the basis. For example, products of plane waves with local finite element basis are used in the generalized finite element method [1,27]. Trefftz-type methods use local solutions of the Helmholtz equation [12] as the basis functions, which in the case of piece-wise constant media are plane waves.
A challenging issue in these approaches is how to choose the number of plane wave directionsà priori. It is well known that in order to achieve a good accuracy, a fine, ω dependent, resolution in angle space is required. However, since many of these directions may not be relevant to a specific problem setup, it will not only increase the degrees of freedom significantly but also make the resulting linear system ill-conditioned due to the numerical loss of orthogonality of the elements of the basis. For heterogeneous media, the construction of local solutions of the Helmholtz equation becomes more challenging, and some recent papers have focused on how to construct such local solution: a generalized plane wave method was developed in [15]; and a two-scale method using local numerical solutions of the Helmholtz equation on a fine mesh as the basis function on a coarse mesh was developed in [9].
Another approach to solve the high-frequency Helmholtz equation is based on the geometric optics ansatz of the wave field. In the ansatz, the high-frequency wave field is approximated by a superposition of a finite number, that we denote by N , of dominant wave fronts of the form at each point; see [18] for examples. The main advantage is that both the amplitude functions A n (x) and phase functions φ n (x) are independent of the frequency ω and hence are non-oscillatory and smooth except at a measure zero set, e.g., focus points, caustics, corners in a smooth medium. However, both the amplitude and phase functions depend on the medium, source distribution and domain boundary, i.e., they are global and problem specific. Once the phase functions of the wave fronts are available, the oscillatory pattern of the wave field is known and can be incorporated into the numerical approximation to improve both stability and accuracy. Instead of incorporating many problem independent plane waves to the basis, the main idea of phase based numerical methods is to explicitly incorporate the phase function into the basis functions [10,20]. The main issue in phase-based methods is how to compute the phase functions, which are global unknown functions depending on the medium and the source distribution. Only in simple geometry and homogenous medium the phase function can be computed analytically; otherwise, it has to be computed numerically [18]. Usual approaches include Lagrangian methods, such as ray tracing by solving ordinary differential equations (ODE), and Eulerian methods based on partial differential equation (PDE), e.g., the Eikonal equation, satisfied by the phase function. Both these methods face numerical difficulties for general media in practice. For example, convergence and divergence of rays cause difficulties to obtain an accurate and uniform ray field. Since the solution to the PDE is single valued, it can only capture the phase function corresponding to the first arrival time of the wave front. For a general medium with varying speed, computing the appropriate global phase functions φ n (x) is a challenging task since different phase functions may be defined in different regions the boundaries of which are difficult to determine. In other words, the number of terms in the geometric optics ansatz varies from point to point; see [18] for examples. Moreover, since the phase functions are precomputed once for all and fed into the computation of the solution for the Helmholtz equation in these approaches, the phase function has to be computed extremely accurate because the numerical error is amplified by ω. Instead of incorporating the global phase function into the basis, one can use local ray directions, d n (x) = ∇φn(x) |∇φn(x)| , as local dominant plane wave direction of e iω dn·x and incorporate them into the local basis. Compared to the usual plane wave methods that typically require a large number of evenly distributed wave directions independent of the problem, only directions of dominant wave fronts relevant to the problem are involved. Hence the degrees of freedom can be kept minimal and ill-conditioning of the resulting linear system due to redundancy can be reduced. This combines the advantages of plane wave methods and phase-based methods. As is shown in [7], using dominant wave directions in plane wave methods can significantly improve efficiency and accuracy for solving the highfrequency Helmholtz equation in heterogeneous medium. Moreover, it was shown that there is some tolerance in the approximation of the dominant wave directions.
For those examples in [7], a simple ray tracing method was used to find the rays and the corresponding dominant wave directions at each point for a very simple setup in a homogeneous medium and the exact phase function was given for a heterogeneous medium.
Although incorporating problem specific dominant wave directions into basis functions has shown very promising results, a key issue in practice is how to find the dominant wave directions in a stable, efficient and systematic way for general media.
In this work we propose an approach that is based on learning the dominant wave directions specific to the medium and source distribution. In particular, we probe the same medium by the same source, i.e., solving the Helmholtz equation (1) with the same c(x), f (x), for a relative low frequency ω ∼ √ ω. The computed wave field is post-processed by numerical micro-local analysis (NMLA) or other signal processing tools to estimate dominant wave directions. These estimated wave directions are then used in plane wave methods to solve the original high-frequency Helmholtz equation. In our approach, global phase functions are not needed. Instead, dominant wave directions are computed locally where both the number of dominant wave directions and the directions can vary from point to point. It provides the flexibility to deal with general media. Moreover, once a more accurate wave field is computed, it can be used to get a better estimation of the dominant wave directions and then used to improve the high-frequency wave field again. In other words, both the dominant wave directions and the high-frequency wave field can be computed and improved in an iterative way. In this paper, we develop a simple ray-based finite element (ray-FEM) method in 2D for smooth media as a proof of concept study of our proposed approach. We start with a finite element mesh with mesh size h satisfying wh = O(1), i.e., a few points per wavelength. First, the low frequency Helmholtz equation with ω ∼ √ ω is solved on the mesh with quasi-optimality since ω 2 h = O(1). Then NMLA [3,5] (see Section 3) is applied to the computed low-frequency wave field to estimate local dominant wave directions. In the second step we incorporate these estimated dominant wave directions into the local finite element basis as the generalized finite element method to solve the high-frequency Helmholtz equation on the same mesh. If necessary, the solution for the high-frequency Helmholtz equation can also be processed by NMLA to improve the estimate of local dominant wave directions which can be used to further improve the high-frequency solution. We also develop an iterative fast solver for the discretized linear system based on polarized traces which can achieve O(N ) complexity with a possible poly-logarithmic factor for smooth media, where N is the total number of unknowns. We use numerical examples to show that our approach can achieve the following goals: (1) a stable discretization with degrees of freedom of the minimal order O(ω d ), (2) a fast solver with approximate linear complexity for the discretized linear systems, (3) an asymptotic convergence rate of O(ω − 1 2 ) as ω → ∞. Here is an outline of this paper. We first describe the ray-FEM using geometric optics ansatz as the motivation and study its approximation property in Section 2. In Section 3 we introduce the NMLA with its stability and local ray direction error analyzed in Appendix A and B. Section 4 provides the full presentation of the numerical algorithm whose complexity is given in Section 5. Numerical results are presented in Section 6. Conclusion and future works are summarized in Section 7.

The Ray-FEM Method
In this section we describe the ray-FEM method for the Helmholtz equation and its rationale based on the geometric optics ansatz. We explain briefly the ansatz and approximate it locally via a superposition of plane waves propagating in dominant directions. These plane waves with their associated dominant directions are incorporated into the finite element basis functions to improve both stability and accuracy for the computation of the solution to the high-frequency Helmholtz equation.
In this section we suppose that the dominant directions are known exactly. However, in Section 3 we will describe how to learn the dominant wave directions by probing the medium using low-frequency waves.
We use the following boundary value problem in 2D for the illustration of our method, −∆u − k 2 (x)u = f, in Ω, where Ω is a bounded Lipschitz domain in R 2 , and k(x) = ω/c(x) is the inhomogeneous wave number. (3) is usually refered to as the Helmholtz equation with impedance boundary conditions. This equation was chosen in order to easily impose another types of boundary conditions by modifying the coefficient β. Specifically, the Dirichlet boundary condition corresponds to β = ∞ and the first order absorbing boundary condition to β = ±1. Moreover, it is easy to extend (3) to incorporate absorbing boundary conditions implemented via PML [6], as it will be performed in the numerical experiments in Section 6.3.

Geometric optics ansatz and local plane wave approximation
The standard derivation of geometric optics ansatz is the use of WKJB approximation [23,16] (or the Lüneberg-Kline expansion [17]) for the solution to the Helmholtz equation (1): By taking ω → ∞ and considering only the first term one has where A is usually called the amplitude and φ the phase. The key features of the geometric optics ansatz are: • A and φ are independent of the frequency ω; • A and φ depend on the medium, c(x); and the source distribution, f (x).
Moreover, except for a small set of points, e.g., source/focus points, caustics, and discontinuities of the medium, they are smooth functions satisfying the following PDE system: In general, the phase function, φ, and the amplitude function, A, are multi-valued functions corresponding to multiple arrivals of wave fronts. Hence one can further decompose the geometric optics ansatz into a superposition of a number of wave fronts in the form: where N (x) is the number of fronts/rays passing through x, and the phases φ n and amplitudes A n are single valued functions satisfying the eikonal/transport equations (6), each defined in a suitable domain with suitable boundary conditions [4]. Based on the above geometric optics ansatz, one can derive a local plane wave approximation at any point where φ n and A n are smooth with variations on a O(1) scale. Indeed, using Taylor expansions on a small neighborhood around an observation point x 0 , we have for |x − x 0 | < h 1. Define as the ray directions of the wave fronts at x 0 , k(x 0 ) = ω/c(x 0 ), and the affine complex amplitude. By replacing (9) and (10) in (8) we have for |x − x 0 | < h 1. From (11) we have that u can be approximated locally by a superposition of plane waves propagating in certain directions with affine complex amplitudes. Moreover, as ω → ∞, such that ωh = O(1), the asymptotic error for the local plane wave approximation (11) is O(ω −1 ), which is of the same order as the asymptotic error for the original geometric ansatz (7). We use (11) as the motivation to construct local finite element basis with mesh size h = O(ω −1 ), in which an affine function is multiplied by plane waves oscillating in those ray directions, resulting in local approximations similar to (11).

Ray-based FEM formulation
We use a finite element method to compute the solution to (3) whose standard weak formulation is given by where The domain, Ω, is discretized with a standard regular triangulated mesh, with mesh size h, which we denote by T h = {K}, where K represents a triangle of the mesh. Using the aforementioned mesh we define two approximation spaces for the variational formulation (12): • Standard FEM (S-FEM), we use low order P1 finite elements, i.e., piece-wise bilinear functions, • Ray-FEM, we use P1 finite elements multiplied by plane waves as in (11).
For a given element K ∈ T h , we denote by V j and x j , j = 1, 2, 3, the vertices of K and their coordinates respectively. Moreover, we denote by {ϕ j (x)} 3 j=1 a partition of unity consisting of piecewise bilinear functions satisfying ϕ j (x i ) = δ ij , i, j = 1, 2, 3, where δ ij is the Kronecker delta. The basis given by {ϕ j (x)} 3 j=1 is usually called the nodal basis for Lagrange P1 finite elements. The standard local approximation space is given by and the global P1 finite element space To define the ray-FEM we enrich the P1 finite elements by incorporating the ray information. Let { d j,l } n j l=1 be n j ray directions at the vertex V j , define the ray-based local approximation space by and the global ray-FEM space by We can define the Standard FEM method by Analogously, we define the ray-FEM method by

Approximation property of ray-FEM with exact ray information
We provide a simple computation to estimate the approximation error of the ray-FEM space. In particular, we compute an asymptotic bound on inf u h ∈V Ray (T h ) ||u − u h || L 2 (Ω) , which is achieved by estimating the interpolating error using V Ray (K) as a basis.
In the computation we assume that the ray direction, which is the gradient of the phase function φ, and the phase function itself, are exactly known. For simplicity, we assume N = 1 for the asymptotic formula in (7), i.e., that only one ray crosses each point of the domain. In addition we suppose that f , the source, is zero inside the domain. Under those circumstances A and φ are smooth; given that the source is outside the domain there is no singularity in the amplitude, and given that only one ray crosses each point in the domain, no caustic occurs. From the geometric optics ansatz, we have Let Ω be our domain of interest, T h a triangle mesh of Ω with width h. We denote by N h the total number of vertices on the mesh are the coordinates of all mesh nodes and their corresponding nodal basis functions for standard P1 element. We note that e iω[φ(x j )−∇φ(x j )·x j ] is a constant for the nodal basis associated to x j in an element K. From this observation we can easily deduce that the local ray-FEM space can be rewritten as Hence the nodal interpolation of the solution can be written as which, by construction lies within the global ray-FEM space V Ray (T h ). Let S j be the support of ϕ j (x), and |S j | ∼ O(h 2 ) be the area of S j . Then using the triangular inequality and the smoothness assumptions we have If the exact rays are known and the mesh size follows h ∼ ω −1 , then we have i.e. that the approximation error decays linearly with 1 ω , without oversampling.
Remark 1. The ray information can be incorporated into other Galerkin basis in the same fashion. For example, in the hybrid numerical asymptotic method of [11], the basis functions are constructed by multiplying nodal piece-wise bilinear function to oscillating functions with phase factors; the plane wave DG method of [7] employs the products of small degree polynomials and dominant plane waves as basis functions; the phase-based hybridizable DG method of [20] considers basis functions as products of polynomials and phase-based oscillating functions. However, the phase or ray information in these methods is obtained from solving the eikonal equation with ray tracing and related techniques.

Learning Local Dominant Ray Directions
In Section 2 we use geometric optics to provide the motivation for the ray-FEM by building an adaptive approximation space that incorporates ray information specific to the underlying Helmholtz equation. However, the ray directions, which depend on the medium and source distribution, are unknown quantities themselves, hence they need to be computed or estimated. One way is to compute the global phase function, by either ray tracing or solving the eikonal equation, and take its gradient. As discussed in the introduction, computing the global phase function in a general varying medium can be difficult.
In the present paper we propose a totally different approach. This novel approach is based on learning the dominant ray directions by probing the same medium with the same source using a relative low-frequency wave. To be more specific, we first solve the Helmholtz equation (1) with the same speed function c(x), right hand side f (x) and boundary conditions but with a relative low-frequency ω ∼ √ ω on a mesh with size h = O( ω −2 ) = O(ω −1 ) with a standard finite element method, which is quasi-optimal in that regime. Then the local dominant ray directions are estimated based on the computed low-frequency wave field. The key point is that the low-frequency wave has probed the medium specific to the problem globally and only local dominant ray directions need to be learned, which allows us to handle multiple arrivals of wave fronts seamlessly. In particular, we use numerical microlocal analysis (NMLA), which is simple and robust, in this work to extract the dominant ray directions locally. However, this is a signal processing task that can be accomplished using other methods such as Prony's method [8], Pisarenko's method [22], MUSIC [24], matrix pencil [13], among many others.

NMLA
In this subsection, for the sake of completness, we provide a brief introduction to NMLA developed in [3,5]. If we suppose that a wave field is a weighted superposition of plane waves with the same wave number, but propagating in different directions, then the aim of NMLA is to extract from samples of the wave field, the directions and the weights. In the sequel we use a 2D example to illustrate the method.
Suppose that a wave field, denoted by u(x), is composed of N plane waves around an observation point x 0 , We suppose that we can sample the wave field, u(x), and its derivative on a circle S r (x 0 ) centered at x 0 with radius r. The wave field can be written under the model assumption in (24) as Furthermore, define implicitly the angle variable θ = θ( s); accordingly we denote θ n = θ( d n ) the angles to recover, and x(θ) = x 0 + r s(θ). Using the angle based notation we sample the impedance quantity which removes any possible ambiguity due to resonance, and improves the robustness to noise for solutions to the Helmholtz equation [5], on the circle S r (x 0 ). Then we apply the filtering operator B to the impedance quantity where , J l is the Bessel functions of order l, J l is their derivatives and is the l-th Fourier coefficient of U . It is shown in [5] that where Then it is possible to obtain the directions and the amplitudes by picking the peaks in the filtered data in (29). However, for applications, the measured data is never a perfect superposition of plane waves; therefore, we provide, for completeness, stability and error estimates for NMLA from [5] in Appendix A. In principle, as long as the perturbation is relative small with respect to the true signal, the estimation error is O( 1 kr ). In other words, the larger the radius of the circle compared to wavelength the more accurate the estimation is.
In our application, our data is the numerical solution of the Helmholtz equation. In addition to noises and numerical errors, there are perturbations due to two model errors : • the geometric optics ansatz has an asymptotic error of order O(ω −1 ) (see (7)); • in the geometric ansatz the wave field at a point is a superposition of curved wave fronts. In particular, the curvature of the wave fronts results in a compromise in the choice of the radius of the sampling circle to be of order O(ω − 1 2 ) for the NMLA in order to achieve the stability and the minimal error of order Detailed analysis is provided in Appendix B. Below is a summary of the NMLA algorithm.

Approximation property of numerical ray-FEM
In this section we incorporate the errors in the estimation of the ray directions into the approximation error for the ray-FEM method, in which ray directions are first estimated by using Algorithm 1 to the solution of the Helmholtz equation with a relative low-frequency. The estimated ray directions are then used to generate the approximation space. With the same assumptions as in Section 2.3, we estimate an upper bound on inf when the ray-FEM space, V h Ray (T h ), is constructed using the estimated ray directions from high-frequency waves by NMLA.
From Appendix B, the error estimation of dominant ray directions is We denote by the nodal interpolation of the solution in V h Ray (T h ). Then we have Comparing with (22) and (34), the error in the estimation of dominant ray directions due to NMLA leads to the extra term ω 1/2 h, which is the leading order in the high-frequency regime. Specifically, if ωh = O(1), then we have

Algorithms
In this section we provide the full algorithm for ray-FEM including a fast iterative solver based on a modification of the method of polarized traces for the resulting linear systems. In order to streamline the presentation and to make the algorithm easier to understand, we introduce several subroutines, which the main algorithm relies on.
We can separate the full algorithm into three conceptual stages: 1. probing the medium by solving a relative low-frequency Helmholtz equation with the standard FEM, 2. learning the dominant ray directions from low-frequency probing wave field by NMLA, 3. solving the high-frequency Helmholtz equation in ray-FEM space.
If necessary the second stage can be applied to the high-frequency wave field computed in stage 3 to improve the estimation of dominant ray directions and then repeat stage 3.
We remind the reader that the ultimate objective of the algorithm presented in this paper (i.e., Algorithm 5) is to solve the Helmholtz equation (1) at frequency ω with a total O(ω d ) (up to poly-logarithmic factors) computational complexity. In order to achieve this objective, we discretize the PDE with a mesh size h = O(ω −1 ), which leads to a total of O(ω d ) number of degrees of freedom and a sparse linear system with O(ω d ) number of nonzeros. Then we develop a fast iterative solver with quasi-linear complexity to solve the resulting linear system after discretization. Below is a more detailed description of the three stages. Finally, following the notation defined in the prequel, we denote the triangular mesh by T h .

Probing
We first solve the Helmholtz equation with frequency ω ∼ √ ω in the same medium and with the same source on T h . The low-frequency problem is solved using the standard finite element method (S-FEM) with linear elements as prescribed by Algorithm 2.

Learning
Once the low-frequency problem has been solved, we extract the dominant ray directions from u ω,h using NMLA as described in Section 3.1 around each mesh node. We utilize the smoothness of the phase functions, and hence the smoothness of the ray directions field to reduce the computational cost. The reduction is achieved by restricting the learning of the dominant ray directions to vertices of a coarse mesh down-sampled from T h . Such coarse mesh is denoted by T hc , where h c = O( √ h). The resulting dominant ray directions are then interpolated onto the fine mesh T h .
Note that at each vertex of T hc , the wave field u ω,h on the fine mesh T h is used for NMLA to estimate the dominant ray directions. There are three sources of error in the learning stage we need to be aware of: • numerical error of u ω,h , • model error in geometric optics ansatz, • interpolation error.
The numerical error for u ω,h by Algorithm 2 in L 2 norm [28] is O( ωh 2 + ω 2 h 2 ) = O(ω −1 ), which is negligible with respect to the model error in the geometric optics ansatz. The error introduced by the geometric optics approximation is O( ω − 1 2 ) as shown in Section 3.1 and Appendix B. The error due to linear interpolation on T hc to get ray direction estimation at every vertex on , which is much smaller than the model error in geometric optics ansatz. Hence the overall error in ray direction estimation based on NMLA on u ω,h and interpolation is O( ω − 1 2

High-frequency solver
Once the dominant ray directions on T h have been computed, we can construct the ray-FEM space V Ray (T h ) and solve the high-frequency Helmholtz equations following (18), which is implemented in Algorithm 4. In general, the more accurate of the dominant ray directions learned by NMLA, the more accurate ray-FEM high-frequency solution can be computed by Algorithm 4. From section 4.2, the accuracy order of learning stage from low-frequency wave field is O( ω − 1 2 ), and following the error analysis of section 3.2, the consequent ray-FEM solution has the same order of accuracy. However, we can apply the learning stage to the computed high-frequency wave field to improve approximation for dominant ray directions. Moreover, the improved ray information can be used to improve the approximation for high-frequency wave field computed by ray-FEM. This process can be continued iteratively to further improve the approximations of the ray directions and the high-frequency wave field. Therefore, we develop an iterative ray-FEM Helmholtz solver for high-frequency ω in Algorithm 5. ω

Algorithm 4 Ray-FEM Helmholtz Solver
end while 13: end function Remark 2. Extensive numerical experiments and Appendix A suggest that the NMLA process in learning dominant ray directions is remarkably stable even for noisy plane wave data. Hence, the iterative process in Algorithm 5 usually needs very few iterations to reach the desired accuracy. Typically, we only need 1 or 2 iterations in our numerical tests.

Fast linear solver
To achieve the overall complexity mentioned in the introduction, it is necessary to solve the linear system resulting from both standard finite element methods and ray-FEM, which we write in a generic form as in linear complexity (up to poly-logarithmic factors). This solver is, in fact, the bottleneck of Algorithms 2 and 4.
For a smooth medium, this can be achieved by modifying the method of polarized traces [30], of which we provide a brief review here. For further details we refer the interested readers to [30]. The method of polarized traces is a domain decomposition method that encompasses the following aspects • layered domain decomposition; • absorbing boundary conditions between subdomains implemented via PML [6]; • transmission conditions issued from a discrete Green's representation formula; • efficient preconditioner arising from localization of the waves via an incomplete Green's formula.
The first two aspects can be effortlessly implemented. Consider a layered partition of Ω into L slabs, or layers {Ω } L =1 . Define f as the restriction of f to Ω , i.e., f = f χ Ω ; and define the local Helmholtz operators as with absorbing boundary conditions implemented via PML between slabs. The method of polarized traces aims to solve the global linear system in (36) by solving the local systems H , which are the discrete version of (37).
In order to solve the global system, or in this case to find a good approximate solution, we need to "glue" the subdomains together, which is achieved via a discrete Green's integral formula deduced by imposing discontinuous solutions.
In the original formulation of the method of polarized traces [30], the Green's representation formula was used to build a global surface integral equation (SIE) at the interfaces between slabs. The SIE was solved using an efficient preconditioner coupled with a multilevel compression of the discrete kernels to accelerate the on-line stage of the algorithm. The original algorithm had superlinear off-line complexity which was amortized among a large number of right-hand sides, which represents a typical situation in explorations geophysics.
In the context of the present paper, the linear systems issued from the raybased FEM depends on the source distribution, making it impossible to amortize a super-linear off-line cost. In order to reduce the off-line cost we use a matrixfree formulation (see Chapter 2 in [29]) with a domain decomposition in thin layers. In this case, the cost per iteration is linear on the number of degrees of freedom, depending on the grow or the auxiliary degrees of freedom corresponding to the PML's. Finally the convergence is normally achieved in O(log ω) iterations, as it will be shown in the sequel.

Complexity
In this section we provide an overall computational complexity count of our algorithm for the high-frequency Helmholtz equation (1) in terms of ω. Our count includes learning ray directions by NMLA and the linear solver for the discretized systems from both standard FEM and ray FEM for low-frequency and high-frequency Helmholtz equation respectively. The summary of the complexity of the method is given in Table 2.

Ray direction learning
As described in Section 4.2, Algorithm 3 applies NMLA to computed wave fields with low-frequency ω ∼ √ ω or high-frequency ω and first estimate ray directions at vertices on a downsampled coarse mesh T hc and then interpolate the ray directions to the vertices on a fine mesh T h . We remind the reader the following scalings: . These scalings allow us to strike a balance among the number of observation points at which NMLA is used to estimate ray directions, the radius of sampling circle, and the corresponding number of sampling points on the circle to resolve the wave field to reach the optimal accuracy of NMLA with desired total computational complexity.
We estimate the dominant computational cost which is applying NMLA to the high-frequency wave field to get ray direction estimation on T h in 2D as an illustration of Table 1. It is shown in Appendix B, the least error that can be achieved by NMLA is O(ω − 1 2 ) when radius r of the sampling circle centered at an observation point is O(ω − 1 2 ). Hence the number of points sampled on the circle to resolve the wave field of frequency ω is M ω = O(ω 1 2 ). Since NMLA is a linear filter based on a Fourier transform in angle space, the corresponding computational complexity is O(M ω log M ω ) [3]. The number of observation points we need to perform NMLA is the number of vertices on the coarse mesh which is O(h −2 c ) = O(ω). Hence the computation cost to obtain ray directions at all vertices on the coarse mesh by NMLA is O(ω 3 2 log ω). Finally the ray directions estimated at the vertices on the coarse mesh by NMLA are interpolated to the fine mesh T h . Interpolation is a linear operation and hence its computation complexity is O(ω 2 ). Table 1 provides the complexity for each operation of NMLA, where d is the dimension and C N M LA , C ray,hc , C Int , and C ray,h is the computation complexity of NMLA at a single vertex, NMLA on the under sampled coarse mesh, interpolation of local ray directions to the fine mesh and full algorithm for learning local ray directions at frequency ω on mesh T h respectively.

Helmholtz Solver
The most computationally intensive component in the whole ray-FEM algorithm is solving the linear systems after discretization of the Helmholtz equation. Algorithm 5 solves both u ω,h = S-FEM( ω, h, c, f, g) by S-FEM and u dω,h = Ray-FEM(ω, h, c, f, g, {d j ω,h } N h j=1 ) by ray-FEM on the same mesh T h . Each solver is composed of three steps: the assembling step, the setup step, and the iterative solve step.
Since the basis functions are locally supported, the resulting matrix is sparse. The complexity of the assembling step is of the same order as the degrees of freedom In the setup stage, the computational domain is decomposed into subdomains of thin layers whose width is comparable to the characteristic wavelength. The local problems in each subdomain are factorized using multifrontal method in O(  N h log N h )) complexity per iteration. Extensive numerical experiments suggest that the number of iterations to converge is O(log N h ) for the both high-and low-frequency solves for smooth media. Hence, the empirical overall complexity is O(N h log N h ) for the high-frequency solve and O(N h log 3 N h ) for the low-frequency one, which is summarized in the following table 2:

Numerical Experiments
In this section we provide several numerical experiments to test the proposed ray-FEM and corroborate our claims. For all cases, the domain of interest is Ω = [−1/2, 1/2] 2 with different source terms and boundary conditions. Ω is discretized using a standard traingular mesh. The mass and stiffness matrices are assembled using a high-order Gaussian quadrature rule to compute the integrals numerically 1 .

Convergence tests
In the first test, the exact solution to the Helmholtz equation with Robin boundary condition is the wave field (normalized by the frequency ω) corresponding to a point source outside the domain. It is given by Numerically we solve the Helmholtz equation (1) with c(x) ≡ 1, source f (x) ≡ 0 and exact impedance boundary data with mesh size such that the number of points per wave length (NPW) is 6 for different ω and test convergence for both the ray direction estimation by NMLA and the final numerical solution by ray-FEM. First, a probing wave with low-frequency ω = √ ω is solved by standard FEM. Then NMLA is applied to the low-frequency probing wave to get an estimation of the local dominant ray directions d w . Instead of using the regular NMLA for plane wave decomposition, we use NMLA with a curvature correction version [5] to estimate the normal direction of a circular wave front. The local ray direction information is used in ray-FEM to produce the first numerical solution to the high frequency Helmholtz equation u d w .
We employ one more iteration in the framework of iterative ray-FEM by applying NMLA to u d w to get an improved local ray direction estimation d w and then use it again in ray-FEM to get a more accurate numerical solution to the high-frequency Helmholtz equation u dw .
Due to the curvature correction in NMLA, it can be shown that the error in ray direction estimation ) by using an analysis similar to that in [5] and Appendix B. Using the estimate in Section 3.2, one can show that the approximation error for numerical ray-FEM space is at least  Figure 1 show that the asymptotic error, for both ray estimation and numerical solution by ray-FEM, decreases as the frequency increases. Moreover, they show that the ray-FEM algorithm is stable and quasioptimal with fixed NPW, i.e., ωh = O(1). Quasi-optimality holds if the ratio of numerical solution error to the best approximation error is bounded by a constant which is independent of frequency ω. Here we assume the best approximation error in ray-FEM space inf u h ∈V h Ray (T h ) u − u h is of the same order as the interpolation error u − u h I . So we can define the estimated quasi-optimality constant by , and it is shown in Figure 2. Also it shows that one more iteration using iterative ray-FEM can significantly improve the final numerical solution to the order of O(ω −1 ), which is of the same order when exact ray direction d ex is used in ray-FEM, due to the asymptotic error for geometric ansatz. Note, however, that NMLA with curvature correction is only valid for a single point source in homogeneous media. u dex − u ex L 2 2.97e-05 1.49e-05 7.47e-06 3.74e-06 Table 3: Errors of one point source problem for fixed NPW = 6. θ ex is the exact ray angle, θ(d ω ) and θ(d ω ) are ray angle estimations using low and high frequency waves, respectively; u d ω , u dω and u dex are ray-FEM solutions using low-frequency ray estimation d ω , high frequency ray estimation d ω , and exact ray d ex , respectively.
Next we show that our method can handle multiple wave fronts by probing the whole domain and extracting dominant ray directions locally. The setup is exactly as above except that there are four point sources. The exact solution is given by The main difficulty of this example compared to the one above, is that the lowfrequency wave solution by the standard FEM contains multiple wave fronts at each point due to the interference of multiple sources. The numerical results are shown in right column of Figure 1. In this case, NMLA with curvature correction does not apply so we use the the standard NMLA version for plane wave decomposition described in Section 3.1 to estimate local dominant ray directions. As analyzed in Section 3.2 and Appendix B, the expected error for ray direction estimation and numerical solution is of order O(ω −1/2 ) due to the curved wave fronts. The numerical results show that the ray-FEM meets the expectation without pollution as the frequency increases.

Phase errors
Here we show that by incorporating the estimated ray directions, ray-FEM can capture the phase much more accurately. We test our algorithm with a source inside the domain. In particular, we use a point source, given its importance in many practical applications, in particular, in geophysics, in which the sources are often modeled as point sources. Moreover, in applications oriented towards inverse and imaging problems, having a numerical method that produces the correct phase in the far field is of great importance in order to properly locate features in the image.
In this experiments we focus our attention on the far-field since our current method can not deal with singularities in amplitude and phase at source points. We start by solving the Helmholtz equation using a slight modification of Algorithm 5 and using standard finite elements at the source point. In this case the source term is located inside the domain f = δ(x − (−0.4, −0.4)), but we will use the associated column of the mass matrix (normalized by mesh size h) as the right hand side. For vertices near the source, we use the exact ray direction, the radial direction, in our ray-FEM. And we find the ray directions by NMLA for vertices away from the source; see Figure 3 left part for the ray direction field.
To demonstrate the phase errors in numerical solutions, we plot the computed wave field on a 90 degree part of an annulus [26], with the radial coordinate varying on an interval of about two wavelengths; see Figure 3 right part.
In this case the frequency ω = 250π is fixed, but we increase the number of grid points per wavelength. Figure 4 depicts the behavior of both ray-FEM solution and standard FEM solution. From the figure we can easily observe the superiority of the ray-FEM on minimizing the phase error, even using relatively coarse meshes.
We then solve the Helmholtz equation using a heterogeneous medium given by Figure 5 with the source located inside. We also provide an experiment where we show the ability of the method in this paper to handle wave field with caustics; see Figure 6. Again radial directions are used for local ray directions near the source point.

Complexity tests
In this subsection we test the computational complexity for ray-FEM. A key step is solving the sparse linear systems generated by ray-FEM using iterative methods with a performant preconditioner, e.g., domain decomposition techniques coupled with high-quality absorbing/transmission boundary conditions. In our tests, we use a modification of the method of polarized traces to solve the linear systems resulting from both standard FEM and ray-FEM as described in Section 4.4.
We use the numerical experiments to demonstrate the overall computational complexity of our ray-FEM method. In particular, we solve the Helmholtz equation with a point source in both homogeneous medium and heterogeneous medium. We compute for many different frequencies, using Algorithm 5 with only one iteration of ray-FEM, the solution to the Helmholtz equation posed on Ω = [−0.5, 0.5] × [−0.5, 0.5] with absorbing boundary conditions implemented via PML. For each frequency we report the execution time of the low and high frequency problems and the time spent in processing the data using NMLA to extract the dominant ray information.
As explained in Section 4, in order to process the data using NMLA we need to solve the low-frequency problem in a slightly larger domain. The size of the larger domain is given by the sampling radius of the NMLA. For the sake of simplicity, we use a low-frequency subdomain, Ω low = [−1, 1] × [−1, 1], i.e., four times bigger than the original domain. The size can be reduced in order to lower computational cost for the low-frequency problem.
The main issue with the low-frequency solver in our case are the PML's, given that each thin slab contains less than a wavelength across, the PML may not be very effective. In order to decrease the number of iterations to converge, we increase the PML points logarithmically with the frequency. This implies a slightly more expensive setup cost and solve cost as shown in Figures 7 left and 8 left. Figure 7 shows the runtime for solving the Helmholtz equation with a point source inside a homogeneous medium. We can observe that the overall cost is O(N ) up to poly-logarithmic factors as shown in our complexity study. The low-frequency solver has a slightly higher asymptotic cost in this case, given the ratio between the width of the PML and the characteristic wavelength inside the domain. Figure 8 shows the runtime for solving the Helmholtz equation with a point source inside a heterogenous medium. We can observe the same scaling as before, albeit with slightly larger constants. Error

Conclusion
In this work we present a numerical method, the ray-FEM, for the high frequency Helmholtz equation in smooth media based on learning problem specific basis functions to represent the wave field. The key information, local ray directions, is extracted from a relative low frequency wave field that has probed the whole domain. These local ray directions are then incorporated into the basis to improve both stability and accuracy in the computation for high frequency wave field. Moreover, both local ray directions and the high frequency wave field can be further improved through more iterations. Numerical tests suggest that our method only requires a fixed number of points per wave length without pollution effect as frequency becomes large. By designing a fast solver for the discretized linear systems an overall complexity of order O(ω d log ω) is achieved.
However, our ray-FEM can not handle singularities of both the amplitude and phase on a mesh. We will develop a hybrid method that combines local asymptotic expansion near the source and the ray-FEM away from the source in our future work.
where B * ≤ 0.89 is a pure constant independent of ω and B 1 is the complex amplitude of the plane wave. Then the error in the angle estimation is given by Similar results can be derived for multiple waves N > 1. We remark that 1 4B * 0.28, which implies that if the relative noise level does not surpass 28% the angle will be detected within an error of order O( 1 kr ). In Benamou's work [5], an analysis of a point source shows that |θ 1 − θ * | decreases like O(ω −1/2 ) when the point x 0 is far away from the source and the radius of the observation circle is chosen like r ∼ ω −1/2 for large ω. We obtain similar accuracy order for general noisy plane waves under some smoothness conditions, see details in Appendix B.
Appendix B. Error analysis of wave-field as a perturbed plane wave data As introduced in Section 3.1, NMLA is a tool to process a signal that is (approximately) a superposition of plane waves with frequency ω and to extract each plan wave component by sampling the signal on a circle/sphere with radius r around a reference point. As shown in Appendix A, provided that the perturbation of the signal is relative small compared to the signal, the estimation of the plane wave directions converges and the error is O( 1 ωr ). In this application, we use NMLA to process wave-field data, which is the numerical solution to the Helmholtz equation, to extract the directions of dominant wave fronts based on the geometric optics ansatz (7) in the high-frequency regime. Hence it is important to study the wave field data as a perturbation of plane wave data locally and estimate the error in the ray directions obtained from NMLA. In particular, this analysis allows us to find the optimal choice of the radius of the sampling circle/sphere, in order to achieve the minimal asymptotic error for the ray direction estimation in terms of the frequency ω of the Helmholtz equation which generates the wave-field data. The result is crucial for both error analysis and implementation of ray-FEM. Since the wave field data in our application is the numerical solution to the Helmholtz equation, its perturbation can be composed as the sum of three components: 1. numerical error in solving the Helmholtz equation and interpolation error in obtaining data on the sampling circle/sphere for NMLA from the numerical solution on a fixed mesh, 2. the asymptotic error in the geometric optics ansatz, 3. the local deviation of a smooth curved wave front from a planar wave front.
On a mesh with mesh size h = O(ω −1 ), the last components, which we call the phase error, is the dominant factor among the three. We present below an analysis of the phase error, in which, for simplicity, we only consider one wave front.
.  As shown in Appendix A, on one hand δU has to be small compared to U . On the other hand, the error in direction estimate from NMLA is O( 1 wr ). Assuming the smoothness of A(x) and φ(x), i.e., boundedness of ∇A(x), A(x) and ∇ 2 φ(x), the leading term in δU is ωr 2 |A(x 0 )| s T ∇ 2 φ(x 0 ) s as ω → ∞, where s T ∇ 2 φ(x 0 ) s is the curvature of the wave front. Hence the radius of the sampling circle can at most be chosen r ∼ O( 1 √ ω ) as ω → ∞. Let