基于RungeKutta法的谐振子能级求解与波函数绘图问题排查

原代码存在的核心问题 #

• 初始条件错误：设置谐振子波函数在左边界x=-d/2处为0，这不符合谐振子波函数的行为（基态波函数在所有有限位置均为正，仅当x→±∞时趋近于0）。
• 函数参数冲突：RungeKutta2d函数定义中的x参数未被使用，循环中直接调用全局变量tpoints，且变量名x与循环变量重名，导致逻辑混乱。
• 冗余函数嵌套：V(x, potential_function)函数完全冗余，直接调用势能函数H(x)即可，无需额外封装。

修复后的完整代码 #

%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
from scipy.special import hermite

# 物理常数
m = 9.109383702 * 10**-31  # kg, 电子质量
hbar = 1.054571817 * 10**-34  # J·s, 约化普朗克常数
e = 1.602176634 * 10**-19  # C, 电子电荷
d = 5 * 10**-9  # m, 量子点边长

# 计算范围与采样点
x_start = -d/2
x_end = d/2
N = 2000
h = (x_end - x_start) / N
x_points = np.arange(x_start, x_end, h)

# 定义谐振子势能
V0 = 700 * e
def harmonic_potential(x):
    return (V0 * x**2) / ((d/2)**2)

# 薛定谔方程的一阶微分形式
def schrodinger_eq(r, x, E, potential):
    psi, psi_prime = r
    d_psi = psi_prime
    d_psi_prime = (2*m / hbar**2) * (potential(x) - E) * psi
    return np.array([d_psi, d_psi_prime])

# 四阶龙格-库塔积分函数
def runge_kutta_2d(initial_r, x_vals, eq_func, E, potential):
    r = np.copy(initial_r)
    psi_vals = []
    psi_prime_vals = []
    
    for x in x_vals:
        psi_vals.append(r[0])
        psi_prime_vals.append(r[1])
        
        # 计算RK4增量
        k1 = h * eq_func(r, x, E, potential)
        k2 = h * eq_func(r + 0.5*k1, x + 0.5*h, E, potential)
        k3 = h * eq_func(r + 0.5*k2, x + 0.5*h, E, potential)
        k4 = h * eq_func(r + k3, x + h, E, potential)
        
        r += (k1 + 2*k2 + 2*k3 + k4) / 6
    
    # 添加最后一步的结果
    psi_vals.append(r[0])
    psi_prime_vals.append(r[1])
    
    return np.array([psi_vals, psi_prime_vals])

# 设置能级n（可修改为任意非负整数）
n = 0

# 谐振子角频率与初始能级猜测
omega = np.sqrt(8 * V0 / (m * d**2))
E_guess = (n + 0.5) * hbar * omega
E1 = E_guess - 3e-19
E2 = E_guess

# 初始条件：左边界处的波函数与导数（使用渐近解设置）
alpha = np.sqrt(m * omega / hbar)
psi_initial = np.exp(-0.5 * alpha * x_start**2)
psi_prime_initial = -alpha * x_start * psi_initial
initial_r = np.array([psi_initial, psi_prime_initial])

# 割线法求解本征能级
tolerance = e / 100000
# 计算初始两个猜测的波函数终点值
soln1 = runge_kutta_2d(initial_r, x_points, schrodinger_eq, E1, harmonic_potential)
psi1 = soln1[0][-1]

soln2 = runge_kutta_2d(initial_r, x_points, schrodinger_eq, E2, harmonic_potential)
psi2 = soln2[0][-1]

while abs(E2 - E1) > tolerance:
    E3 = E2 - psi2 * (E2 - E1) / (psi2 - psi1)
    # 更新能级猜测
    E1, E2 = E2, E3
    # 重新计算波函数终点值
    psi1 = runge_kutta_2d(initial_r, x_points, schrodinger_eq, E1, harmonic_potential)[0][-1]
    psi2 = runge_kutta_2d(initial_r, x_points, schrodinger_eq, E2, harmonic_potential)[0][-1]

print(f"n={n} 对应的能级为 {E3} J，即 {E3/e} eV")

# 求解归一化的波函数
soln = runge_kutta_2d(initial_r, x_points, schrodinger_eq, E3, harmonic_potential)
psi = soln[0][:-1]  # 匹配x_points的长度

# 梯形法归一化波函数
integral = np.trapz(psi**2, x=x_points)
psi_norm = psi / np.sqrt(integral)

# 验证归一化
norm_check = np.trapz(psi_norm**2, x=x_points)
print(f"归一化积分值：{norm_check}（接近1表示正确）")

# 绘图
plt.figure(figsize=(8, 5))
plt.axvline(x=-d/2, c='#5f5f5f', ls='-', lw=2.5, label='边界')
plt.axvline(x=d/2, c='#5f5f5f', ls='-', lw=2.5)
plt.plot(x_points, psi_norm, label=f'n={n} 波函数')
plt.xlabel('位置 (m)')
plt.ylabel('归一化波函数')
plt.title(f'谐振子n={n}态波函数')
plt.legend()
plt.grid(alpha=0.3)
plt.show()

关键修复说明 #

1. 修正初始条件：使用谐振子波函数的渐近形式exp(-αx²/2)设置左边界的波函数值与导数，符合谐振子波函数的实际行为。
2. 优化函数逻辑：重命名变量避免混淆（如tpoints改为x_points），移除冗余的V函数，让runge_kutta_2d函数依赖传入的参数而非全局变量，提升代码可读性与可维护性。
3. 简化归一化计算：使用np.trapz替代手动梯形求和，代码更简洁且不易出错。
4. 增强可视化：添加网格、图例，优化绘图尺寸，让结果更直观。

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