写一个有限差分法模拟三维地震波的python代码

以下是一个简单的有限差分法模拟三维地震波的Python代码示例：

import numpy as np
import matplotlib.pyplot as plt

# 地震波速度模型
vp = np.zeros((100, 100, 100)) + 2000
vp[50:70, 50:70, 50:70] = 2200

# 初始化网格和时间步长
nx, ny, nz = vp.shape
dx, dy, dz = 10, 10, 10
dt = 0.0001

# 初始化波场
u = np.zeros((nx, ny, nz))
v = np.zeros((nx, ny, nz))
w = np.zeros((nx, ny, nz))

# 循环迭代
for n in range(5000):
    # 计算二阶空间导数
    uxx = np.gradient(np.gradient(u, axis=0), axis=0) / dx**2
    uyy = np.gradient(np.gradient(u, axis=1), axis=1) / dy**2
    uzz = np.gradient(np.gradient(u, axis=2), axis=2) / dz**2
    
    # 计算速度和应力
    vel = vp * vp
    sxx = vel * (uxx + uyy + uzz)
    
    # 更新波场
    u = u + dt * v
    v = v + dt * ((1 / vp**2) * sxx - 2 * v)
    w = w + dt * ((1 / vp**2) * sxx - 2 * w)
    
    # 绘制波场快照
    if n % 100 == 0:
        plt.imshow(u[:, :, 50], cmap='gray')
        plt.title("Time Step: {}".format(n))
        plt.show()

在此示例中，我们使用了一个简单的地震波速度模型，并使用了有限差分法来计算波场的更新。我们对波场进行了5000次迭代，每100次迭代绘制了一个波场的快照。此代码可能需要根据不同的地震波速度模型进行修改。