Numerical Solution of the Wave Equation: K-Space Method

Method

Assume we are given the following wave equation.

$$ \partial_{t}^{2} u(\mathbf{x}, t) = \Delta_{\mathbf{x}} u (\mathbf{x}, t), \qquad (\mathbf{x}, t) \in \mathbb{R}^{n} \times \mathbb{R} $$

By taking the Fourier transform of both sides with respect to the variable $\mathbf{x}$, we obtain the following.

$$ \partial_{t}^{2} \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) = \mathcal{F}_{\mathbf{x}}[\Delta u (\cdot, t)]u(\mathbf{k}, t) = -|\mathbf{k}|^{2} \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) $$

Additionally, the left-hand side of the above equation can be approximated using the finite difference method as follows.

$$ \partial_{t}^{2} \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) \approx \dfrac{\mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t + \Delta t) - 2\mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) + \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t - \Delta t)}{(\Delta t)^{2}} $$

From the above two equations, we obtain the following.

$$ \dfrac{\mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t + \Delta t) - 2\mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) + \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t - \Delta t)}{(\Delta t)^{2}} = -|\mathbf{k}|^{2} \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) $$

$$ \implies \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t + \Delta t) = 2\mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) - \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t - \Delta t) - (\Delta t)^{2}|\mathbf{k}|^{2} \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) $$

Taking the inverse Fourier transform on both sides,

$$ u(\mathbf{x}, t + \Delta t) = 2u(\mathbf{x}, t) - u(\mathbf{x}, t - \Delta t) - \mathcal{F}_{\mathbf{k}}^{-1}\Big[ (\Delta t)^{2}|\mathbf{k}|^{2} \mathcal{F}_{\mathbf{x}}u(\mathbf{k}, t) \Big](\mathbf{x}, t) \tag{1} $$

Explanation

$k$-space method refers to the technique of calculating the propagation of waves using Fourier transforms as shown above. The name derives from the fact that taking the Fourier transform converts the spatial variable $\mathbf{x}$ into the frequency variable $\mathbf{k}$. When solving the wave equation using the finite element method, the ratio of the spatial grid to the time grid is crucial, but in the $k$-space method, such considerations are unnecessary.

Additionally, in the last term of $(1)$, replacing $(\Delta t |\mathbf{k}|)^{2}$ with $4 \sin^{2} \left( \dfrac{|\mathbf{k}|\Delta t}{2} \right)$ yields better results.^{1 2}

Code

Below is the code for the $k$-space method written in Julia, and the execution result is shown as a GIF.

using Plots
using FFTW

# Initialization
U = zeros(Float64, 256, 256, 256)

# x ∈ [-1, 1]
# t ∈ [0, 1]
Δt = 1/256

for i ∈ 1:256, j ∈ 1:256
    if 10^2 < (i-80-64)^2 + (j-80-64)^2 < 12^2
        U[i, j, 1] = 1
        U[i, j, 2] = 1
    end
    if (i-45-64)^2 + (j-45-64)^2 < 4^2
        U[i, j, 1] = 1
        U[i, j, 2] = 1
    end
end

# make k-space grid
k = LinRange(-128., 128, 256) * π
k1 = k * fill!(similar(k'), 1)
k2 = fill!(similar(k), 1) * k'

K = sqrt.(k1.^2 + k2.^2)
K = 4 * (sin.(Δt * K/2)).^2
K = ifftshift(K)

for i ∈ 3:256
    U[:, :, i] = 2U[:, :, i-1] - U[:, :, i-2] - real(ifft(K .* fft(U[:, :, i-1])))
end

w2b = cgrad([:white, :black])
anim = @animate for i ∈ 1:256
    heatmap(U[:,:,i], c=w2b, clim=(0,1), title="$i")
end

gif(anim)