频散曲线绘制

作者：Sheaping
链接：https://www.zhihu.com/question/458713107/answer/1878851102
来源：知乎
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。

import numpy as np
import matplotlib.pyplot as plt
from obspy import read

# 读取面波数据并画图。
st = read('MASW_DATA/Sample_Data/*.SAC')
dt = st[0].stats.delta
data = []
scale = 0.05
dx = 2
plt.figure(figsize=(8, 6))
for i, tr in enumerate(st):
    d = tr.data
    data.append(d)
    t = np.arange(len(d)) * dt
    plt.plot(t, d*scale+(i+1)*dx, lw=1, color='b')
plt.xlabel('Time (s)')
plt.ylabel('Offset (m)')
plt.tight_layout()
plt.savefig('Surface_wave.png')
plt.show()

# 二维FFT。
d = np.array(data)
n = len(d[0])
# m为空间方向的采样点数，m增大可以让FK谱光滑一点，以达到插值效果。
m = len(d[:, 0]) * 5
D = np.zeros((m, n))
D[:len(d[:, 0])] = d
# 时间采样率。
fs = 1 / dt
# 空间采样率
xs = 1 / dx
# 频率 (赫兹)。
f = np.arange(-n//2, n//2) * fs / (n-1)
# 波数 (每米)。
k = 2 * np.pi * np.arange(-m//2, m//2) * xs / (m-1)
# 二维FFT。
fk = np.fft.fft2(D)

# 作图。
pmin = -10
P = abs(np.fft.fftshift(fk)); P /= P.max(); P = 10 * np.log10(P)
P2 = abs(fk); P2 /= P2.max(); P2 = 10 * np.log10(P2)

plt.figure(figsize=(11, 8))
plt.subplot(221)
plt.pcolormesh(f, k, P2, cmap='magma', vmin=pmin, vmax=0)
plt.xlabel('Frequency (s$^{-1}$)')
plt.ylabel('Wave number (2$\pi$m$^{-1}$)')

plt.subplot(222)
plt.pcolormesh(f, k, P, cmap='magma', vmin=pmin, vmax=0)
plt.plot([f[n//2], f[-1], f[-1], f[n//2], f[n//2]],
        [k[0], k[0], k[m//2], k[m//2], k[0]],
         lw=2, ls='--', color='r')
plt.xlabel('Frequency (s$^{-1}$)')
plt.ylabel('Wave number (m$^{-1}$)')

plt.subplot(223)
plt.pcolormesh(f[n//2:], k[:m//2], P[:m//2, n//2:],
               cmap='magma', vmin=pmin, vmax=0)
plt.xlabel('Frequency (s$^{-1}$)')
plt.ylabel('Wave number (m$^{-1}$)')

plt.subplot(224)
plt.pcolormesh(f[n//2:], abs(k[:m//2][::-1]), P[:m//2, n//2:][::-1],
               cmap='magma', vmin=pmin, vmax=0)
cbar = plt.colorbar()
cbar.set_label(r'FK spectra (dB)')
plt.xlim(0, 100)
plt.xlabel('Frequency (s$^{-1}$)')
plt.ylabel('Wave number (m$^{-1}$)')
plt.tight_layout()
plt.show()