如何修改2D波动方程代码：移除Pygame并实现RK系列时间步长方案

一、移除Pygame可视化，改用Matplotlib轻量替代 #

为简化代码，Matplotlib的动画模块足够实现基础波动效果可视化，且依赖更轻量化。

步骤1：替换导入模块 #

• 删除所有import pygame相关代码
• 添加核心依赖：

  import numpy as np
  import matplotlib.pyplot as plt
  from matplotlib.animation import FuncAnimation
  

步骤2：重构可视化逻辑 #

• 移除Pygame的窗口初始化、事件循环、像素绘制代码
• 用Matplotlib创建3D画布，通过FuncAnimation实现帧更新逻辑

二、替换Euler为Runge-Kutta系列时间步长方案 #

2D波动方程可拆分为一阶微分方程组，适配RK类方法求解。假设波动方程为：

$u_{tt} = c^2 (u_{xx} + u_{yy})$

将其拆解为两个一阶方程：

• $v = u_t$（速度为位移的时间导数）
• $v_t = c^2 (\nabla^2 u)$（加速度由拉普拉斯算子决定）

通用前置函数实现 #

首先定义有限差分法计算拉普拉斯算子：

def laplacian(u, dx, dy):
    # 四阶精度有限差分计算空间二阶导数
    lap = (np.roll(u, 1, axis=0) + np.roll(u, -1, axis=0) - 2*u)/dx**2 + \
          (np.roll(u, 1, axis=1) + np.roll(u, -1, axis=1) - 2*u)/dy**2
    # 周期性边界处理（可按需改为固定边界）
    lap[0,:] = lap[1,:]
    lap[-1,:] = lap[-2,:]
    lap[:,0] = lap[:,1]
    lap[:,-1] = lap[:,-2]
    return lap

然后定义微分方程组右端项函数：

def rhs(state, c, dx, dy, t):
    u, v = state
    lap_u = laplacian(u, dx, dy)
    du_dt = v
    dv_dt = c**2 * lap_u
    return np.array([du_dt, dv_dt])

1. RK4实现 #

def rk4_step(state, dt, c, dx, dy, t):
    k1 = rhs(state, c, dx, dy, t)
    k2 = rhs(state + 0.5*dt*k1, c, dx, dy, t + 0.5*dt)
    k3 = rhs(state + 0.5*dt*k2, c, dx, dy, t + 0.5*dt)
    k4 = rhs(state + dt*k3, c, dx, dy, t + dt)
    return state + (dt/6)*(k1 + 2*k2 + 2*k3 + k4)

2. RK5(Butcher)实现 #

def rk5_step(state, dt, c, dx, dy, t):
    k1 = rhs(state, c, dx, dy, t)
    k2 = rhs(state + dt*(1/5)*k1, c, dx, dy, f, t + dt/5)
    k3 = rhs(state + dt*(3/40)*k1 + dt*(9/40)*k2, c, dx, dy, t + 3*dt/10)
    k4 = rhs(state + dt*(44/45)*k1 - dt*(56/15)*k2 + dt*(32/9)*k3, c, dx, dy, t + 4*dt/5)
    k5 = rhs(state + dt*(19372/6561)*k1 - dt*(25360/2187)*k2 + dt*(64448/6561)*k3 - dt*(212/729)*k4, c, dx, dy, t + 8*dt/9)
    k6 = rhs(state + dt*(9017/3168)*k1 - dt*(355/33)*k2 + dt*(46732/5247)*k3 + dt*(49/176)*k4 - dt*(5103/18656)*k5, c, dx, dy, t + dt)
    return state + dt*(35/384)*k1 + dt*(500/1113)*k3 + dt*(125/192)*k4 - dt*(2187/6784)*k5 + dt*(11/84)*k6

3. RK7(Verner)实现 #

def rk7_step(state, dt, c, dx, dy, t):
    k1 = rhs(state, c, dx, dy, t)
    k2 = rhs(state + dt*(2/27)*k1, c, dx, dy, t + 2*dt/27)
    k3 = rhs(state + dt*(1/36)*k1 + dt*(1/12)*k2, c, dx, dy, t + dt/9)
    k4 = rhs(state + dt*(1/24)*k1 + dt*(0)*k2 + dt*(1/8)*k3, c, dx, dy, t + dt/6)
    k5 = rhs(state + dt*(5/12)*k1 + dt*(0)*k2 + dt*(-25/16)*k3 + dt*(25/16)*k4, c, dx, dy, t + 5*dt/12)
    k6 = rhs(state + dt*(1/20)*k1 + dt*(0)*k2 + dt*(0)*k3 + dt*(1/4)*k4 + dt*(1/5)*k5, c, dx, dy, t + dt/2)
    k7 = rhs(state + dt*(-25/108)*k1 + dt*(0)*k2 + dt*(125/108)*k3 + dt*(-65/27)*k4 + dt*(125/54)*k5 + dt*(0)*k6, c, dx, dy, t + 5*dt/6)
    k8 = rhs(state + dt*(31/300)*k1 + dt*(0)*k2 + dt*(0)*k3 + dt*(61/225)*k4 + dt*(-2/9)*k5 + dt*(13/900)*k6 + dt*(0)*k7, c, dx, dy, t + dt)
    k9 = rhs(state + dt*(2)*k1 + dt*(0)*k2 + dt*(-53/6)*k3 + dt*(704/45)*k4 + dt*(-107/9)*k5 + dt*(67/90)*k6 + dt*(3)*k7 + dt*(0)*k8, c, dx, dy, t + 2*dt)
    k10 = rhs(state + dt*(-91/108)*k1 + dt*(0)*k2 + dt*(23/108)*k3 + dt*(-976/135)*k4 + dt*(311/54)*k5 + dt*(-19/60)*k6 + dt*(17/6)*k7 + dt*(0)*k8 + dt*(-1/12)*k9, c, dx, dy, t + dt)
    return state + dt*(13/160)*k1 + dt*(0)*k2 + dt*(0)*k3 + dt*(2375/5984)*k4 + dt*(0)*k5 + dt*(512/1865)*k6 + dt*(125/1992)*k7 + dt*(2187/6784)*k8 + dt*(0)*k9 + dt*(11/84)*k10

三、整合完整代码示例 #

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# 仿真参数
nx, ny = 100, 100
dx, dy = 0.1, 0.1
c = 1.0  # 波速
dt = 0.005
t_total = 10.0

# 初始化状态：[位移u, 速度v]
x = np.linspace(0, (nx-1)*dx, nx)
y = np.linspace(0, (ny-1)*dy, ny)
X, Y = np.meshgrid(x, y)
u = np.exp(-((X - nx*dx/2)**2 + (Y - ny*dy/2)**2)/(2*(2*dx)**2))  # 初始高斯脉冲
v = np.zeros((nx, ny))
state = np.array([u, v])

def laplacian(u, dx, dy):
    lap = (np.roll(u, 1, axis=0) + np.roll(u, -1, axis=0) - 2*u)/dx**2 + \
          (np.roll(u, 1, axis=1) + np.roll(u, -1, axis=1) - 2*u)/dy**2
    lap[0,:] = lap[1,:]
    lap[-1,:] = lap[-2,:]
    lap[:,0] = lap[:,1]
    lap[:,-1] = lap[:,-2]
    return lap

def rhs(state, c, dx, dy, t):
    u, v = state
    lap_u = laplacian(u, dx, dy)
    du_dt = v
    dv_dt = c**2 * lap_u
    return np.array([du_dt, dv_dt])

def rk4_step(state, dt, c, dx, dy, t):
    k1 = rhs(state, c, dx, dy, t)
    k2 = rhs(state + 0.5*dt*k1, c, dx, dy, t + 0.5*dt)
    k3 = rhs(state + 0.5*dt*k2, c, dx, dy, t + 0.5*dt)
    k4 = rhs(state + dt*k3, c, dx, dy, t + dt)
    return state + (dt/6)*(k1 + 2*k2 + 2*k3 + k4)

# 可视化逻辑
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
ax.set_zlim(-1, 1)

def update(frame):
    global state
    t = frame * dt
    state = rk4_step(state, dt, c, dx, dy, t)
    ax.clear()
    ax.plot_surface(X, Y, state[0], cmap='viridis')
    ax.set_zlim(-1, 1)
    return ax,

ani = FuncAnimation(fig, update, frames=int(t_total/dt), interval=10, blit=False)
plt.show()

切换RK方法说明 #

只需将update函数中的rk4_step替换为rk5_step或rk7_step即可，确保函数参数完全匹配。

四、关键注意事项 #

• 稳定性：所有RK方法需满足CFL条件：$dt \leq \frac{dx}{c\sqrt{2}}$，高阶方法可适当放宽但需测试
• 边界条件：示例用周期性边界，若需固定边界，可将边界点的位移设为0
• 性能：RK7计算量最大，若追求速度优先选RK4；大网格下可降低帧更新频率或缩小网格分辨率

内容的提问来源于stack exchange，提问作者donet3088