将2D系统Split Operator方法扩展至3D的坐标网格构建咨询

原2D代码中存在一处明显错误：初始化波函数ψ_0 = np.zeros((N,My*My),dtype = 'complex')里的第二维应为My*Mz而非My*My，扩展到3D前建议先修复该问题保证2D版本运行正确。

核心疑问解答：3D位置/动量网格构建 #

你原有2D版本中tile+转置的逻辑本质是让1D坐标向量广播匹配为对应维度的网格，扩展到3D时不能直接照搬2D的转置写法，需要匹配三维维度顺序。更推荐用np.meshgrid实现，比手动tile+转置的出错概率低：

# 第一步新增x维度基础参数
Mx = 128*2
Lx = 10
LxT = Lx*2
x0 = np.linspace(-Lx, Lx, Mx)
k0x = np.linspace(-Mx*np.pi/LxT, Mx*np.pi/LxT - 2*np.pi/LxT, Mx)

# 构建3D位置网格，indexing='ij'保证维度顺序为(x,y,z)，和你原有维度习惯对齐
x0op, y0op, z0op = np.meshgrid(x0, y0, z0, indexing='ij')
# 3D动量网格同理
k0xop, k0yop, k0zop = np.meshgrid(k0x, k0y, k0z, indexing='ij')

如果你坚持用手动tile的写法，正确实现如下：

# x维度在最外层，两次tile扩展到y、z维度
x0op = np.tile(x0.reshape(Mx,1,1), (1, My, Mz))
# y维度在中间层，tile扩展到x、z维度
y0op = np.tile(y0.reshape(1, My, 1), (Mx, 1, Mz))
# z维度在最内层，tile扩展到x、y维度
z0op = np.tile(z0.reshape(1,1,Mz), (Mx, My, 1))
# 动量网格构建逻辑和位置网格完全一致
k0xop = np.tile(k0x.reshape(Mx,1,1), (1, My, Mz))
k0yop = np.tile(k0y.reshape(1, My,1), (Mx,1,Mz))
k0zop = np.tile(k0z.reshape(1,1,Mz), (Mx,My,1))

其余3D扩展关键修改点 #

• 波函数初始化调整：将原2D的维度乘积修改为3D版本

  ψ_0 = np.zeros((N, Mx*My*Mz), dtype='complex')
  

• 初始高斯波包调整：新增x维度的高斯项，利用广播机制相乘得到3D波包

  # 先更新基础参数数组为3维：ω = np.array([1,1,1]), gs = np.array([0,0,0])
  x0c = gs[2]
  kIx = 0
  σ_x = np.sqrt(2/ω[0])
  temp_x = np.exp(-((x0 - x0c)/σ_x)**2) * np.exp(1j * kIx * x0)
  # 三个1D高斯广播相乘得到3D波包
  temp_1 = temp_x.reshape(Mx,1,1) * temp_y.reshape(1,My,1) * temp.reshape(1,1,Mz)
  ψ_0[iS,:] = temp_1.reshape(1, Mx*My*Mz)
  

• 动能传播子与FFT逻辑调整：动能项新增x维度贡献，FFT切换为3D版本

  # 3D动能算符与传播子
  T = (k0xop**2)/2 + (k0yop**2)/2 + (k0zop**2)/2
  TP = np.exp(-1j * T * dt)
  
  # 循环内的传播逻辑修改为3D适配
  for n in range(N):
      temp2 = ψ[n,:]
      temp3 = temp2.reshape(Mx, My, Mz)
      temp4 = fft.fftshift(fft.fftn(temp3)) # 3D傅里叶变换
      ek[t,n] = np.real(np.sum(np.conj(temp4)*T*temp4) / np.sum(np.conj(temp4)*temp4))
      temp5 = temp4 * TP
      temp6 = fft.ifftn(fft.ifftshift(temp5)) # 3D逆傅里叶变换
      ψ[n,:] = temp6.reshape(1, Mx*My*Mz)
  

• 两能级适配：如果要扩展为两能级系统，保持N=2，按原有逻辑添加势能项的非对角耦合部分即可，无势场景下范数与能量会自动守恒。

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