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

以下是一个简单的有限差分法模拟三维地震波的Python代码，其中采用了PML吸收边界条件。

import numpy as np
import matplotlib.pyplot as plt

# Parameters
nx = 101  # Number of grid points in x-direction
ny = 101  # Number of grid points in y-direction
nz = 101  # Number of grid points in z-direction
dx = 10  # Grid spacing in x-direction (m)
dy = 10  # Grid spacing in y-direction (m)
dz = 10  # Grid spacing in z-direction (m)
dt = 0.001  # Time step (s)
nt = 500  # Number of time steps
fc = 20  # Centre frequency of the source (Hz)
src = np.zeros((nt,))  # Source time function
src[np.arange(1, nt+1) * dt * fc * 2 * np.pi <= np.pi] = np.sin(np.arange(1, nt+1) * dt * fc * 2 * np.pi / np.pi / 2) ** 2  # Ricker wavelet
pml_width = 20  # Width of PML layer (grid points)
pml_alpha_max = 0.01  # Maximum PML absorption coefficient
rho = 2700  # Density of medium (kg/m^3)
vp = 6000  # P-wave velocity of medium (m/s)
vs = 3464  # S-wave velocity of medium (m/s)
mu = vs ** 2 * rho  # Shear modulus of medium (Pa)

# Grid
x = np.arange(0, nx * dx, dx)
y = np.arange(0, ny * dy, dy)
z = np.arange(0, nz * dz, dz)
X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
dx2 = dx ** 2
dy2 = dy ** 2
dz2 = dz ** 2
dt2 = dt ** 2

# PML coefficients
pml_x = np.zeros((nx,))
pml_y = np.zeros((ny,))
pml_z = np.zeros((nz,))
for i in range(pml_width):
    pml_x[i] = pml_alpha_max * ((pml_width - i) / pml_width) ** 2
    pml_x[-i-1] = pml_alpha_max * ((pml_width - i) / pml_width) ** 2
    pml_y[i] = pml_alpha_max * ((pml_width - i) / pml_width) ** 2
    pml_y[-i-1] = pml_alpha_max * ((pml_width - i) / pml_width) ** 2
    pml_z[i] = pml_alpha_max * ((pml_width - i) / pml_width) ** 2
    pml_z[-i-1] = pml_alpha_max * ((pml_width - i) / pml_width) ** 2

# Coefficients for finite difference scheme
c1 = (dt ** 2) / (dx2 * rho)
c2 = (dt ** 2) / (dy2 * rho)
c3 = (dt ** 2) / (dz2 * rho)
c4 = (dt ** 2) / (dx * dy * mu)
c5 = (dt ** 2) / (dy * dz * mu)
c6 = (dt ** 2) / (dx * dz * mu)

# Initial conditions
p = np.zeros((nx, ny, nz))
vx = np.zeros((nx, ny, nz))
vy = np.zeros((nx, ny, nz))
vz = np.zeros((nx, ny, nz))

# Main loop
for n in range(1, nt):
    # Update p
    p[1:-1, 1:-1, 1:-1] += c1 * (vx[1:-1, 1:-1, 1:-1] - vx[:-2, 1:-1, 1:-1] + vy[1:-1, 1:-1, 1:-1] - vy[1:-1, :-2, 1:-1] + vz[1:-1, 1:-1, 1:-1] - vz[1:-1, 1:-1, :-2])
    p[1:-1, 1:-1, 1:-1] -= c4 * (vx[1:-1, 1:-1, 1:-1] - vx[:-2, 1:-1, 1:-1])
    p[1:-1, 1:-1, 1:-1] -= c5 * (vy[1:-1, 1:-1, 1:-1] - vy[1:-1, :-2, 1:-1])
    p[1:-1, 1:-1, 1:-1] -= c6 * (vz[1:-1, 1:-1, 1:-1] - vz[1:-1, 1:-1, :-2])
    p[1:-1, 1:-1, 1:-1] *= (2 - pml_x[1:-1, np.newaxis, np.newaxis] - pml_y[np.newaxis, 1:-1, np.newaxis] - pml_z[np.newaxis, np.newaxis, 1:-1])
    p[0, :, :] = 0
    p[:, 0, :] = 0
    p[:, :, 0] = 0
    p[-1, :, :] = 0
    p[:, -1, :] = 0
    p[:, :, -1] = 0
    
    # Update vx
    vx[1:-1, 1:-1, 1:-1] += c4 * (p[1:-1, 1:-1, 1:-1] - p[:-2, 1:-1, 1:-1])
    vx[1:-1, 1:-1, 1:-1] -= c2 * (vy[1:-1, 1:-1, 1:-1] - vy[1:-1, :-2, 1:-1])
    vx[1:-1, 1:-1, 1:-1] -= c3 * (vz[1:-1, 1:-1, 1:-1] - vz[1:-1, 1:-1, :-2])
    vx[1:-1, 1:-1, 1:-1] *= (2 - pml_x[1:-1, np.newaxis, np.newaxis] - pml_y[np.newaxis, 1:-1, np.newaxis] - pml_z[np.newaxis, np.newaxis, 1:-1])
    vx[0, :, :] = 0
    vx[:, 0, :] = 0
    vx[:, :, 0] = 0
    vx[-1, :, :] = 0
    vx[:, -1, :] = 0
    vx[:, :, -1] = 0
    
    # Update vy
    vy[1:-1, 1:-1, 1:-1] += c4 * (p[1:-1, 1:-1, 1:-1] - p[1:-1, :-2, 1:-1])
    vy[1:-1, 1:-1, 1:-1] -= c1 * (vx[1:-1, 1:-1, 1:-1] - vx[:-2, 1:-1, 1:-1])
    vy[1:-1, 1:-1, 1:-1] -= c3 * (vz[1:-1, 1:-1, 1:-1] - vz[1:-1, 1:-1, :-2])
    vy[1:-1, 1:-1, 1:-1] *= (2 - pml_x[1:-1, np.newaxis, np.newaxis] - pml_y[np.newaxis, 1:-1, np.newaxis] - pml_z[np.newaxis, np.newaxis, 1:-1])
    vy[0, :, :] = 0
    vy[:, 0, :] = 0
    vy[:, :, 0] = 0
    vy[-1, :, :] = 0
    vy[:, -1, :] = 0
    vy[:, :, -1] = 0
    
    # Update vz
    vz[1:-1, 1:-1, 1:-1] += c4 * (p[1:-1, 1:-1, 1:-1] - p[1:-1, 1:-1, :-2])
    vz[1:-1, 1:-1, 1:-1] -= c2 * (vy[1:-1, 1:-1, 1:-1] - vy[1:-1, :-2, 1:-1])
    vz[1:-1, 1:-1, 1:-1] -= c1 * (vx[1:-1, 1:-1, 1:-1] - vx[:-2, 1:-1, 1:-1])
    vz[1:-1, 1:-1, 1:-1] *= (2 - pml_x[1:-1, np.newaxis, np.newaxis] - pml_y[np.newaxis, 1:-1, np.newaxis] - pml_z[np.newaxis, np.newaxis, 1:-1])
    vz[0, :, :] = 0
    vz[:, 0, :] = 0
    vz[:, :, 0] = 0
    vz[-1, :, :] = 0
    vz[:, -1, :] = 0
    vz[:, :, -1] = 0
    
    # Add source
    p[50, 50, 50] += src[n]
    
    # Visualize wavefield
    if n % 10 == 0:
        plt.clf()
        plt.imshow(p[:, :, 50], cmap='seismic', vmin=-0.5e-5, vmax=0.5e-5)
        plt.colorbar()
        plt.title('Time step: %d' % n)
        plt.pause(0.01)

请注意，此代码仅适用于学术用途，可能需要根据实际应用进行修改和优化。