Python中的split step Fourier方法

Question

很抱歉这篇文章很长但是我试图用分步法在python中模拟二维薛定谔方程。

这个等式的一维问题已在这篇文章中解释过： https://jakevdp.github.io/blog/2012/09/05/quantum-python/

我尝试添加另一个维度并按照上面的帖子修改了运算符，但我很困惑如何绘制psi函数现在我添加了另一个维度，这里是我通过添加和额外维度修改的类方程，我也修改了辅助函数现在我添加了额外的维度：

import numpy as np
from scipy.fftpack import fft, ifft


class Schrodinger(object):
    def __init__(self, x,y,psi_x0,psi_y0,V_x,V_y,k0=None, hbar=1, m=1, t0=0.0):

        self.x,self.y,psi_x0,psi_y0,self.V_x,V_y=map(np.asarray(x,y,psi_x0,psi_y0,V_x,V_y))
        N=self.x.size
        assert self.x.shape==(N,)
        assert self.y.shape==(N,)
        assert psi_x0.shape==(N,)
        assert self.V_x.shape==(N,)
        assert self.V_y.shape==(N,)

        self.hbar = hbar
        self.m = m
        self.t = t0
        self.dt_ = None
        self.N = len(x)
        self.dx=self.x[1]-self.x[0]
        self.dy=self.y[1]-self.y[0]
        self.dk = 2 * np.pi / (self.N * self.dx)

        if k0 == None:
            self.k0 = -0.5 * self.N * self.dk
        else:
            self.k0 = k0
        self.k = self.k0 + self.dk * np.arange(self.N)

        self.psi_x = psi_x0
        self.psi_y=psi_y0
        self.compute_k_from_x()
        self.compute_k_from_y()

        self.x_evolve_half = None
        self.x_evolve = None
        self.y_evolve_half = None
        self.y_evolve = None
        self.k_evolve = None

        self.psi_x_line = None
        self.psi_y_line = None
        self.psi_k_line = None
        self.V_x_line = None
        self.V_y_line = None

    def _set_psi_x(self,psi_x):
        self.psi_mode_x=(psi_x*np.exp(-1j*self.k[0]*self.x)*self.dx/np.sqrt(2*np.pi))
    def _get_psi_x(self):
        return (self.psi_mode_x*np.exp(1j*self.k[0])*np.sqrt(2*np.pi)/self.dx)

    def _set_psi_y(self,psi_y):
        self.psi_mode_y=(psi_y*np.exp(-1j*self.k[0]*self.y)*self.dy/np.sqrt(2*np.pi))
    def _get_psi_y(self):
        return (self.psi_mode_y*np.exp(1j*self.k[0])*np.sqrt(2*np.pi)/self.dy)

    def _set_psi_k_x(self,psi_k_x):
        self.psi_mode_k_x=psi_k_x*np.exp(1j*self.x[0]*self.dk*np.arange(self.N))
    def _get_psi_k_x(self):
        return self.psi_mode_k_x*np.exp(-1j*self.x[0]*self.dk*np.arange(self.N))

    def _set_psi_k_y(self,psi_k_y):
        self.psi_mode_k_y=psi_k_y*np.exp(1j*self.y[0]*self.dk*np.arange(self.N))
    def _get_psi_k_y(self):
        return self.psi_mode_k_y*np.exp(-1j*self.y[0]*self.dk*np.arange(self.N))

    def _get_dt(self):
        return self.dt_
    def _set_dt(self,dt):
        if dt!=self.dt_:
            self.dt_=dt
            self.x_evolve_half=np.exp(-0.5*1j*self.V_x/self.hbar*dt)
            self.y_evolve_half=np.exp(-0.5*1j*self.V_y/self.hbar*dt)
            self.x_evolve=self.x_evolve_half*self.x_evolve_half
            self.y_evolve=self.y_evolve_half*self.y_evolve_half
            self.k_evolve=np.exp(-0.5*1j*self.hbar/self.m*(self.k*self.k)*dt)

    psi_x=property(_get_psi_x,_set_psi_x)
    psi_y=property(_get_psi_y,_set_psi_y)
    psi_k_x=property(_get_psi_k_x,_set_psi_k_x)
    psi_k_y=property(_get_psi_k_y,_set_psi_k_y)
    dz=property(_get_dt,_set_dt)

    def compute_k_from_x(self):
        self.psi_mode_k_x=fft(self.psi_mode_x)
    def compute_k_from_y(self):
        self.psi_mode_k_y=fft(self.psi_mode_y)

    def compute_x_from_k(self):
        self.psi_mode_x=ifft(self.psi_mode_k_x)
    def compute_y_from_k(self):
        self.psi_mode_y=ifft(self.psi_mode_k_y)

    def time_step(self,dt,Nsteps=1):
        self.dt=dt

        if Nsteps>0:
            self.psi_mode_x*=self.x_evolve_half
            self.psi_mode_y*=self.y_evolve_half

        for i in range(Nsteps-1):
            self.compute_k_from_x()
            self.compute_k_from_y()
            self.psi_mode_k_x*=self.k_evolve
            self.psi_mode_k_y*=self.k_evolve
            self.compute_x_from_k()
            self.compute_y_from_k()
            self.psi_mode_x*=self.x_evolve
            self.psi_mode_y*=self.y_evolve

        self.compute_k_from_x()
        self.compute_k_from_y()
        self.psi_mode_k_x*=self.k_evolve
        self.psi_mode_k_y*=self.k_evolve

        self.compute_x_from_k()
        self.psi_mode_x*=self.x_evolve_half
        self.compute_y_from_k()
        self.psi_mode_y*=self.y_evolve_half

        self.compute_k_from_x()
        self.compute_k_from_y()

        self.t+=dt*Nsteps


def gauss_x(x, a, x0, k0):
    return ((a*np.sqrt(np.pi))**(-0.5)* np.exp(-0.5*((x-x0)* 1./a)**2 +1j*x*k0))
def gauss_y(y, a, y0, k0):
    return ((a*np.sqrt(np.pi))**(-0.5)* np.exp(-0.5*((y-y0)* 1./a)**2 +1j*y*k0))

def gauss_k_x(k,a,x0,k0):
    return ((a/np.sqrt(np.pi))**0.5* np.exp(-0.5*(a*(k- k0))** 2- 1j*(k- k0)*x0))
def gauss_k_y(k,a,y0,k0):
    return ((a/np.sqrt(np.pi))**0.5* np.exp(-0.5*(a*(k- k0))** 2- 1j*(k- k0)*y0))


def theta(x):
    x = np.asarray(x)
    y = np.zeros(x.shape)
    y[x > 0] = 1.0
    return y


def square_barrier(x, width, height):
    return height * (theta(x) - theta(x - width))

dt = 0.01
N_steps = 50
t_max = 120
frames = int(t_max / float(N_steps * dt))

hbar = 1.0   
m = 1.9      


N = 2 ** 11
dx = 0.1
dy=0.1
x = dx * (np.arange(N) - 0.5 * N)
y = dy * (np.arange(N) - 0.5 * N)

V0 = 1.5
L = hbar / np.sqrt(2 * m * V0)
a = 3 * L
x0 = -60 * L
y0 = -60 * L
V_x = square_barrier(x, a, V0)
V_y = square_barrier(y, a, V0)
V_x[x < -98] = 1E6
V_x[x > 98] = 1E6
V_y[y < -98] = 1E6
V_y[y > 98] = 1E6


p0 = np.sqrt(2 * m * 0.2 * V0)
dp2 = p0 * p0 * 1./80
d = hbar / np.sqrt(2 * dp2)

k0 = p0 / hbar
v0 = p0 / m
psi_x0 = gauss_x(x, d, x0, k0)
psi_y0 = gauss_y(y, d, y0, k0)


S =Schrodinger(x=x,y=y,psi_x0=psi_x0,psi_y0=psi_y0,V_x=V,V_y=V_y,hbar=hbar,m=m,k0=-28)

一维情节的代码和动画如下：

from matplotlib import pyplot as pl
from matplotlib import animation
fig = pl.figure()

xlim = (-100, 100)
klim = (-5, 5)

ymin = 0
ymax = V0
ax1 = fig.add_subplot(211, xlim=xlim,
                      ylim=(ymin - 0.2 * (ymax - ymin),
                            ymax + 0.2 * (ymax - ymin)))
psi_x_line, = ax1.plot([], [], c='r', label=r'$|\psi(x)|$')
V_x_line, = ax1.plot([], [], c='k', label=r'$V(x)$')
center_line = ax1.axvline(0, c='k', ls=':',
                          label = r"$x_0 + v_0t$")

title = ax1.set_title("")
ax1.legend(prop=dict(size=12))
ax1.set_xlabel('$x$')
ax1.set_ylabel(r'$|\psi(x)|$')

ymin = abs(S.psi_k).min()
ymax = abs(S.psi_k).max()
ax2 = fig.add_subplot(212, xlim=klim,
                      ylim=(ymin - 0.2 * (ymax - ymin),
                            ymax + 0.2 * (ymax - ymin)))
psi_k_line, = ax2.plot([], [], c='r', label=r'$|\psi(k)|$')

p0_line1 = ax2.axvline(-p0 / hbar, c='k', ls=':', label=r'$\pm p_0$')
p0_line2 = ax2.axvline(p0 / hbar, c='k', ls=':')
mV_line = ax2.axvline(np.sqrt(2 * V0) / hbar, c='k', ls='--',
                      label=r'$\sqrt{2mV_0}$')
ax2.legend(prop=dict(size=12))
ax2.set_xlabel('$k$')
ax2.set_ylabel(r'$|\psi(k)|$')

V_x_line.set_data(S.x, S.V_x)


# Animate plot
def init():
    psi_x_line.set_data([], [])
    V_x_line.set_data([], [])
    center_line.set_data([], [])

    psi_k_line.set_data([], [])
    title.set_text("")
    return (psi_x_line, V_x_line, center_line, psi_k_line, title)

def animate(i):
    S.time_step(dt, N_steps)
    psi_x_line.set_data(S.x, 4 * abs(S.psi_x))
    V_x_line.set_data(S.x, S.V_x)
    center_line.set_data(2 * [x0 + S.t * p0 / m], [0, 1])

    psi_k_line.set_data(S.k, abs(S.psi_k))
    title.set_text("t = %.2f" % S.t)
    return (psi_x_line, V_x_line, center_line, psi_k_line, title)


anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=frames, interval=30, blit=True)



pl.show()