基于分裂算符法的1D薛定谔方程模拟：波函数静态问题排查

本人为编程新手，目前尝试编写代码，采用分裂算符法求解一维Schrödinger方程，模拟初始波函数的演化过程。但运行代码后发现波函数始终处于静态状态，代码如下：

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

# Define problem parameters
L = 10  # Box size
N = 1000  # Number of points on the grid
dx = L/N  # Spacing between points
x = np.linspace(0, L, N)  # Position vector

# Define the potential V(x)
def potential(x):
    return 0.5*x**2

# Define time step and number of iterations
dt = 0.01
num_iterations = 1000

# Define the initial wave function
psi = np.exp(-0.5*(x-5)**2)

# Create the figure and axis
fig, ax = plt.subplots()

# Define the plot of the squared magnitude of the wave function
line, = ax.plot(x, np.abs(psi)**2, color='blue', lw=2)

# Set axis labels
ax.set_xlabel('x')
ax.set_ylabel('$|\psi(x)|^2$')

# Define the update function
def update_func(i):
    # Evolve the wave function in imaginary time
    for j in range(50):
        psi = psi * np.exp(-1j*potential(x)*dt/2)
        psi = np.fft.fft(psi)
        psi = psi * np.exp(-1j*(np.pi/L)**2*dt)
        psi = np.fft.ifft(psi)
        psi = psi * np.exp(-1j*potential(x)*dt/2)
    
    # Update the plot of the squared magnitude of the wave function
    line.set_ydata(np.abs(psi)**2)
    ax.set_title('Iteration %d' % i)
    
    return line,

# Create the animation
animation = animation.FuncAnimation(fig, update_func, frames=num_iterations, interval=50)

# Save the animation as an mp4 file
animation.save('animation.mp4', fps=30, extra_args=['-vcodec', 'libx264'])

# Show the animation
plt.show()

关键修改点 #

• 修正变量作用域：update_func内的psi默认是局部变量，修改后不会影响全局的波函数状态，导致每次迭代都从初始值重新计算。需要在函数开头添加global psi声明，让函数可以修改全局的psi变量。
• 修正动量空间演化算符：原代码中动量项的计算完全错误，必须用np.fft.fftfreq生成符合网格特征的波矢k，动能演化算符应为exp(-1j * (k²/2) * dt)（对应动能算符p²/(2m)，取m=1）。
• 添加波函数归一化：数值计算中误差会导致波函数概率守恒被破坏，每次演化后重新归一化能避免概率分布失真。
• 固定y轴范围：设置固定的y轴上下限，防止动画过程中轴范围频繁跳动影响观测。

修改后的完整代码 #

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

# Define problem parameters
L = 10  # Box size
N = 1000  # Number of points on the grid
dx = L/N  # Spacing between points
x = np.linspace(0, L, N)  # Position vector

# Generate correct wave vector k for momentum space
k = 2 * np.pi * np.fft.fftfreq(N, d=dx)

# Define the potential V(x)
def potential(x):
    return 0.5*x**2

# Define time step and number of iterations
dt = 0.01
num_iterations = 1000

# Define and normalize initial wave function
psi = np.exp(-0.5*(x-5)**2)
psi /= np.sqrt(np.sum(np.abs(psi)**2 * dx))  # Normalize to total probability 1

# Create the figure and axis
fig, ax = plt.subplots()
ax.set_ylim(0, 0.6)  # Fixed y-limit for stable animation

# Define the plot of the squared magnitude of the wave function
line, = ax.plot(x, np.abs(psi)**2, color='blue', lw=2)

# Set axis labels
ax.set_xlabel('x')
ax.set_ylabel('$|\psi(x)|^2$')

# Define the update function
def update_func(i):
    global psi
    # Split-operator time evolution
    # Potential half-step
    psi *= np.exp(-1j * potential(x) * dt / 2)
    # Momentum full-step
    psi_fft = np.fft.fft(psi)
    psi_fft *= np.exp(-1j * (k**2 / 2) * dt)
    psi = np.fft.ifft(psi_fft)
    # Potential half-step
    psi *= np.exp(-1j * potential(x) * dt / 2)
    
    # Normalize to prevent numerical drift
    psi /= np.sqrt(np.sum(np.abs(psi)**2 * dx))
    
    # Update plot data and title
    line.set_ydata(np.abs(psi)**2)
    ax.set_title(f'Iteration {i}, Time = {i*dt:.2f}')
    
    return line,

# Create the animation
ani = animation.FuncAnimation(fig, update_func, frames=num_iterations, interval=50)

# Save the animation as an mp4 file
ani.save('animation.mp4', fps=30, extra_args=['-vcodec', 'libx264'])

# Show the animation
plt.show()

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