Python实现二维薛定谔方程球坐标波包动画及共振观测

问题背景 #

已完成一维含时薛定谔方程的有限差分求解代码，现需：

• 在球坐标系下以二维形式绘制波函数
• 制作动画观测3.48 Hartree能量下的波包共振现象
• 实现波包从$r=15$向$r=0$传播（已在初始波函数$\psi_0$的指数项添加负号）

已知条件：$E=k^{2/2$，已构建能量相关薛定谔方程，给定势场$V(x)=7.5x}2e^{-x}$与$\Delta k=0.2$。

现有代码问题分析 #

提供的代码存在两处未定义变量：D2（二阶差分矩阵）和V（势场变量），且当前为一维径向计算，未适配球坐标的动能算符，初始波包位置设为$x0=2$，不符合从$r=15$出发的要求。

核心实现思路 #

1. 球坐标系径向薛定谔方程适配 #

球对称势场下，径向含时薛定谔方程为：
$$i\hbar\frac{\partial\psi_r}{\partial t} = \left(-\frac{\hbar^{2}{2m}\left(\frac{d}2}{dr^2} + \frac{2}{r}\frac{d}{dr}\right) + V(r)\right)\psi_r$$
其中$\hbar=1, m=1$，所以动能项需修正为：
$$-\frac{1}{2}\left(D2 + \frac{2}{r}D1\right)\psi_r$$

• D2：稀疏二阶差分矩阵（有限差分法构建）
• D1：稀疏一阶差分矩阵

2. 初始波包位置修正 #

将初始波包中心$x0$改为15，确保波包从$r=15$处出发：
$$\psi_0(r) = \sqrt{A}\exp\left(-\frac{(r-15)^2}{2\Delta k^2}\right)\exp(-ikr)$$
指数项的负号-ikr表示波矢方向指向原点，波包会向$r=0$传播。

3. 二维波函数绘制（球对称转笛卡尔坐标） #

利用球对称的旋转对称性，将径向波函数扩展为二维分布：

• 生成极坐标网格$(r, \theta)$，转换为笛卡尔坐标$(x,y)=(r\cos\theta, r\sin\theta)$
• 对每个$r$，将径向波函数的模平方$|\psi_r(r)|^2$赋值给所有对应$(x,y)$位置，得到二维密度分布

4. 动画逻辑调整 #

将一维折线动画改为二维热力图动画：

• 初始化时绘制势场的径向分布（可选）和二维坐标轴
• 每一帧更新时，将当前时刻的径向波函数转换为二维密度图，更新图像并显示当前时间

修改后的完整代码示例 #

import numpy as np
from scipy import sparse
import matplotlib.pyplot as plt
import scipy.integrate as integrate
from matplotlib.animation import FuncAnimation
from IPython import display

# 空间网格参数
dx = 0.02
r_max = 15
r = np.arange(0.1, r_max, dx)  # 径向网格，避免r=0的奇点
N = len(r)

# 波包与能量参数
deltak = 0.2
E = 3.48
r0 = 15  # 初始波包位置改为15
k = np.sqrt(2 * E)
A = 1.0 / (deltak * np.sqrt(np.pi))  # 归一化常数

# 势场定义
V = 7.5 * r**2 * np.exp(-r)

# 构建稀疏差分矩阵（球坐标动能算符）
# 一阶差分矩阵D1（中心差分）
D1 = sparse.diags([-1, 1], [-1, 1], shape=(N, N)) / (2 * dx)
# 二阶差分矩阵D2（中心差分）
D2 = sparse.diags([1, -2, 1], [-1, 0, 1], shape=(N, N)) / (dx**2)
# 球坐标径向动能算符：-0.5*(D2 + 2/r * D1)
kinetic_op = -0.5 * (D2 + sparse.diags(2/r, 0) @ D1)

# 时间参数
dt = 0.1
t0 = 0.0
tf = 10.0  # 延长时间以观测共振
t_eval = np.arange(t0, tf, dt)

# 薛定谔方程右端函数
def psi_t(t, psi):
    return -1j * (kinetic_op.dot(psi) + V * psi)

# 初始波函数：从r=15向原点传播
psi0 = np.sqrt(A) * np.exp(-(r - r0)**2 / (2 * deltak**2)) * np.exp(-1j * k * r)

# 求解初值问题
sol = integrate.solve_ivp(psi_t, t_span=[t0, tf], y0=psi0, t_eval=t_eval, method="RK23")

# 二维绘图准备：生成笛卡尔网格
theta = np.linspace(0, 2*np.pi, 100)
R, Theta = np.meshgrid(r, theta)
X = R * np.cos(Theta)
Y = R * np.sin(Theta)

# 动画设置
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
# 左图：一维径向概率密度+势场
ax1.set_xlim(0, r_max)
ax1.set_ylim(0, 5)
ax1.plot(r, V, "k--", label="势场V(r)")
line_radial, = ax1.plot([], [], label="|ψ(r,t)|²")
ax1.legend()
# 右图：二维概率密度分布
im = ax2.pcolormesh(X, Y, np.zeros_like(X), vmin=0, vmax=5, cmap="viridis")
ax2.set_aspect("equal")
ax2.set_xlim(-r_max, r_max)
ax2.set_ylim(-r_max, r_max)
fig.colorbar(im, ax=ax2, label="概率密度")
title = fig.suptitle('')

# 初始化函数
def init():
    line_radial.set_data([], [])
    im.set_array(np.zeros_like(X))
    return line_radial, im

# 帧更新函数
def animate(i):
    # 更新一维图
    prob_radial = np.abs(sol.y[:, i])**2
    line_radial.set_data(r, prob_radial)
    # 更新二维图：将径向概率密度映射到极坐标网格
    prob_2d = prob_radial[np.newaxis, :].repeat(len(theta), axis=0)
    im.set_array(prob_2d.ravel())
    # 更新标题
    title.set_text(f'Time = {sol.t[i]:1.3f}')
    return line_radial, im

# 生成动画
anim = FuncAnimation(fig, animate, init_func=init, frames=len(sol.t), interval=50, blit=True)

# 在Jupyter中显示动画
video = anim.to_html5_video()
html = display.HTML(video)
display.display(html)
plt.close()

关键细节解释 #

1. 球坐标动能算符修正：由于球坐标系的径向动能包含$\frac{2}{r}\frac{d}{dr}$项，必须在差分矩阵中加入该一阶导数项，否则计算结果会偏离真实物理行为。
2. 初始波包的传播方向：指数项$\exp(-ikr)$对应波矢$k$指向$-r$方向，波包会向$r=0$移动；若改为$\exp(ikr)$则波包向外传播。
3. 二维绘制逻辑：利用球对称波函数的旋转不变性，只需将径向概率密度复制到所有角度方向，即可得到二维平面上的轴对称分布，直观展示波包的径向运动与共振现象。
4. 共振现象观测：延长模拟时间（示例中改为tf=10.0），当波包到达势场区域时，会与势场发生相互作用，部分波被反射，部分波穿透，形成共振时会观测到概率密度在势场区域内的周期性振荡。

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