我正在尝试用有限差分法在python中编写2D Schrödinger方程的曲线图

import numpy as np
import matplotlib.pyplot as plt

# 定义网格的大小和步长
N = 100   # 网格大小
L = 10    # 区域大小
dx = L/N  # 步长

# 定义时间步长和总时间
dt = 0.01
T = 1

# 定义势场函数
def potential(x, y):
    return x**2 + y**2

# 定义初始波函数
def wave_function(x, y):
    return np.exp(-(x**2 + y**2))

# 创建网格
x = np.linspace(-L/2, L/2, N)
y = np.linspace(-L/2, L/2, N)
X, Y = np.meshgrid(x, y)

# 初始化波函数
psi = wave_function(X, Y)

# 进行时间演化
t = 0
while t < T:
    # 计算哈密顿量
    H = -0.5 * (np.roll(psi, 1, axis=0) + np.roll(psi, -1, axis=0) +
                np.roll(psi, 1, axis=1) + np.roll(psi, -1, axis=1) - 4 * psi)/dx**2 + potential(X, Y) * psi

    # 使用欧拉方法更新波函数
    psi = psi + dt * H

    t += dt

# 绘制曲线图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, np.abs(psi)**2, cmap='viridis')

# 设置坐标轴标签
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel(r'$|\psi|^2$')

# 显示图形
plt.show()