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

以下是一个简单的三维地震波模拟代码，使用有限差分法进行计算。

import numpy as np
import matplotlib.pyplot as plt

# 地震波速度模型
vp = np.zeros((50,50,50))
vp[25:35, 25:35, 25:35] = 3000
vp[20:30, 20:30, 20:30] = 4000
vp[15:25, 15:25, 15:25] = 5000

# 时间步长和空间步长
dt = 0.001
dx = 10

# 空间网格和时间网格
nx, ny, nz = vp.shape
nt = int(1 / dt)

# 初始化波场
u = np.zeros((nx, ny, nz, 3))
u[:,:,:,0] = np.random.rand(nx, ny, nz) * 0.1

# 有限差分系数
c1 = 9/8
c2 = -1/24

# 循环计算
for n in range(1, nt):
    for i in range(1, nx-1):
        for j in range(1, ny-1):
            for k in range(1, nz-1):
                # 二阶中心差分
                d2udx2 = (u[i+1,j,k,0] - 2*u[i,j,k,0] + u[i-1,j,k,0]) / dx**2
                d2udy2 = (u[i,j+1,k,0] - 2*u[i,j,k,0] + u[i,j-1,k,0]) / dx**2
                d2udz2 = (u[i,j,k+1,0] - 2*u[i,j,k,0] + u[i,j,k-1,0]) / dx**2
                
                # 计算下一个时间步的波场
                u[i,j,k,1] = 2*u[i,j,k,0] - u[i,j,k,2] + (vp[i,j,k]**2 * dt**2) * (c1*d2udx2 + c1*d2udy2 + c1*d2udz2)
    
    # 更新波场
    u[:,:,:,2] = u[:,:,:,1]
    u[:,:,:,1] = u[:,:,:,0]
    u[:,:,:,0] = u[:,:,:,2]

# 绘制地震波传播动画
fig = plt.figure()
ax = fig.gca(projection='3d')

for i in range(0, nt, 10):
    ax.clear()
    ax.voxels(vp > 0, facecolors='g', edgecolor='k', alpha=0.5)
    ax.voxels(u[:,:,:,0] > 0, facecolors='r', edgecolor='k', alpha=0.5)
    ax.set_xlim([0, nx])
    ax.set_ylim([0, ny])
    ax.set_zlim([0, nz])
    ax.set_title('Time step: %d' % i)
    plt.pause(0.01)

plt.show()

该代码使用了三维数组来表示地震波速度模型和波场，使用了二阶中心差分来计算空间二阶导数，使用了有限差分系数来进行时间推进，最后使用了matplotlib库的3D绘图功能来展示地震波传播动画。