Elastic wave equation implementation on a staggered grid

This is a first attempt at implemenenting the elastic wave equation as described in:

[1] Jean Virieux (1986). ”P-SV wave propagation in heterogeneous media: Velocity‐stress finite‐difference method.” GEOPHYSICS, 51(4), 889-901. https://doi.org/10.1190/1.1442147

The current version actually attempts to mirror the FDELMODC implementation by Jan Thorbecke:

[2] https://janth.home.xs4all.nl/Software/fdelmodcManual.pdf

Explosive source

We will first attempt to replicate the explosive source test case described in [1], Figure 4. We start by defining the source signature $g(t)$, the derivative of a Gaussian pulse, given by Eq 4:

$$g(t) = -2 \alpha(t - t_0)e^{-\alpha(t-t_0)^2}$$

In [1]:
from devito import *
from examples.seismic.source import WaveletSource, RickerSource, GaborSource, TimeAxis
from examples.seismic import plot_image
import numpy as np

from sympy import init_printing, latex
init_printing(use_latex=True)
In [2]:
# Initial grid: 1km x 1km, with spacing 100m
extent = (2000., 2000.)
shape = (201, 201)
x = SpaceDimension(name='x', spacing=Constant(name='h_x', value=extent[0]/(shape[0]-1)))
z = SpaceDimension(name='z', spacing=Constant(name='h_z', value=extent[1]/(shape[1]-1)))
grid = Grid(extent=extent, shape=shape, dimensions=(x, z))
In [3]:
class DGaussSource(WaveletSource):
    
    def wavelet(self, f0, t):
        a = 0.004
        return -2.*a*(t - 1/f0) * np.exp(-a * (t - 1/f0)**2)

# Timestep size from Eq. 7 with V_p=6000. and dx=100
t0, tn = 0., 600.
dt = (10. / np.sqrt(2.)) / 6.
time_range = TimeAxis(start=t0, stop=tn, step=dt)

src = RickerSource(name='src', grid=grid, f0=0.01, time_range=time_range)
src.coordinates.data[:] = [1000., 1000.]
src.show()
src_coords(p_src, d) None (0, 0)
src_coords(p_src, d) {p_src: left, d: left} <function generic_derivative at 0x7fac72a1a488>
In [4]:
# Now we create the velocity and pressure fields
so = 2
vx= TimeFunction(name='vx', grid=grid, staggered=x, space_order=so)
vz = TimeFunction(name='vz', grid=grid, staggered=z, space_order=so)
txx = TimeFunction(name='txx', grid=grid, staggered=NODE, space_order=so)
tzz = TimeFunction(name='tzz', grid=grid, staggered=NODE, space_order=so)
txz = TimeFunction(name='txz', grid=grid, staggered=(x, z), space_order=so)
vx(t, x, z) x (0, 1, 0)
vx(t, x, z) {t: left, x: right, z: left} <function staggered_diff at 0x7fac72a1a620>
vz(t, x, z) z (0, 0, 1)
vz(t, x, z) {t: left, x: left, z: right} <function staggered_diff at 0x7fac72a1a620>
txx(t, x, z) node (0, 0, 0)
txx(t, x, z) {t: left, x: left, z: left} <function staggered_diff at 0x7fac72a1a620>
tzz(t, x, z) node (0, 0, 0)
tzz(t, x, z) {t: left, x: left, z: left} <function staggered_diff at 0x7fac72a1a620>
txz(t, x, z) (x, z) (0, 1, 1)
txz(t, x, z) {t: left, x: right, z: right} <function staggered_diff at 0x7fac72a1a620>
In [5]:
# Now let's try and create the staggered updates
t = grid.stepping_dim
time = grid.time_dim

# We need some initial conditions
V_p = 4.0
V_s = 1.0
density = 3.

# The source injection term
src_xx = src.inject(field=txx.forward, expr=src)
src_zz = src.inject(field=tzz.forward, expr=src)

#c1 = 9.0/8.0;
#c2 = -1.0/24.0;

# Thorbecke's parameter notation
cp2 = V_p*V_p
cs2 = V_s*V_s
ro = 1/density

mu = cs2*ro
l = (cp2*ro - 2*mu)

# fdelmodc reference implementation
u_vx = Eq(vx.forward, vx - dt*ro*(txx.dx + txz.dz))

u_vz = Eq(vz.forward, vz - ro*dt*(txz.dx + tzz.dz))

u_txx = Eq(txx.forward, txx - (l+2*mu)*dt * vx.forward.dx -  l*dt * vz.forward.dz)
u_tzz = Eq(tzz.forward, tzz - (l+2*mu)*dt * vz.forward.dz -  l*dt * vx.forward.dx)

u_txz = Eq(txz.forward, txz - mu*dt * (vx.forward.dz + vz.forward.dx))
In [6]:
op = Operator([u_vx, u_vz, u_txx, u_tzz, u_txz] + src_xx + src_zz)
In [7]:
# Reset the fields
vx.data[:] = 0.
vz.data[:] = 0.
txx.data[:] = 0.
tzz.data[:] = 0.
txz.data[:] = 0.

op()
GNUCompiler: cache hit `/tmp/devito-jitcache-uid944902/0bc1050f82b25c37ce3a012bb323a18caf275737.c` [0.01 s]
=========================================================================================
section0<<510,201,201>,<510,201,201>> with OI=1.17 computed in 0.056 s [17.44 GFlops/s, 1.85 GPts/s]
section1<<510,1>,<510,1>,<510,1>,<510,1>,<510,1>,<510,1>,<510,1>> with OI=9.67 computed in 0.000 s [4.73 GFlops/s, 0.16 GPts/s]
=========================================================================================
In [8]:
# Let's see what we got....
plot_image(vx.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(vz.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(txx.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(tzz.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(txz.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
In [9]:
# Now that looks pretty! But let's do it again with a higher order...
so = 12
vx= TimeFunction(name='vx', grid=grid, staggered=x, space_order=so)
vz = TimeFunction(name='vz', grid=grid, staggered=z, space_order=so)
txx = TimeFunction(name='txx', grid=grid, staggered=NODE, space_order=so)
tzz = TimeFunction(name='tzz', grid=grid, staggered=NODE, space_order=so)
txz = TimeFunction(name='txz', grid=grid, staggered=(x, z), space_order=so)

# fdelmodc reference implementation
u_vx = Eq(vx.forward, vx - dt*ro*(txx.dx + txz.dz))

u_vz = Eq(vz.forward, vz - ro*dt*(txz.dx + tzz.dz))

u_txx = Eq(txx.forward, txx - (l+2*mu)*dt * vx.forward.dx -  l*dt * vz.forward.dz)
u_tzz = Eq(tzz.forward, tzz - (l+2*mu)*dt * vz.forward.dz -  l*dt * vx.forward.dx)

u_txz = Eq(txz.forward, txz - mu*dt * (vx.forward.dz + vz.forward.dx))

op = Operator([u_vx, u_vz, u_txx, u_tzz, u_txz] + src_xx + src_zz)

# Reset the fields
vx.data[:] = 0.
vz.data[:] = 0.
txx.data[:] = 0.
tzz.data[:] = 0.
txz.data[:] = 0.

op()

plot_image(vx.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(vz.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(txx.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(tzz.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
plot_image(txz.data[0], vmin=-.5*1e-2, vmax=.5*1e-2, cmap="seismic")
vx(t, x, z) x (0, 1, 0)
vx(t, x, z) {t: left, x: right, z: left} <function staggered_diff at 0x7fac72a1a620>
vz(t, x, z) z (0, 0, 1)
vz(t, x, z) {t: left, x: left, z: right} <function staggered_diff at 0x7fac72a1a620>
txx(t, x, z) node (0, 0, 0)
txx(t, x, z) {t: left, x: left, z: left} <function staggered_diff at 0x7fac72a1a620>
tzz(t, x, z) node (0, 0, 0)
tzz(t, x, z) {t: left, x: left, z: left} <function staggered_diff at 0x7fac72a1a620>
txz(t, x, z) (x, z) (0, 1, 1)
txz(t, x, z) {t: left, x: right, z: right} <function staggered_diff at 0x7fac72a1a620>
GNUCompiler: compiled `/tmp/devito-jitcache-uid944902/c9839d148f7ed3b1bc27d4bb3219351ae6a60553.c` [0.37 s]
=========================================================================================
section0<<510,201,201>,<510,201,201>> with OI=8.07 computed in 0.207 s [33.59 GFlops/s, 0.50 GPts/s]
section1<<510,1>,<510,1>,<510,1>,<510,1>,<510,1>,<510,1>,<510,1>> with OI=9.67 computed in 0.000 s [4.73 GFlops/s, 0.16 GPts/s]
=========================================================================================