Mid9ts
大学生
Mid9ts的博客

[算法]-一个笨拙的二维薛定谔方程数值模拟

[算法]-一个笨拙的二维薛定谔方程数值模拟

使用了Crank-Nicolson算法。知识欠缺,只能用大量的计算来填补。

二维薛定谔方程如下

\[ i\hbar \frac{\partial }{\partial t}\varPsi = \underbrace{\left[ -\frac{\hbar^2}{2m} \left( \frac{\partial ^2}{\partial x^2} + \frac{\partial ^2}{\partial y^2} \right) + V(x,y) \right]}_{H}\varPsi \]

也就是

\[ i\hbar \frac{\partial }{\partial t} \varPsi = H \varPsi \]

可将其视为关于时间的一阶偏微分方程,分别采用向前差分,有

$$ \begin{gather*} \frac{\varPsi^{n+1} – \varPsi^n}{\Delta t} = -\frac{iH^n}{\hbar} \varPsi^n \\ \frac{\varPsi^{n+1} – \varPsi^n}{\Delta t} = -\frac{iH^{n+1}}{\hbar} \varPsi^{n+1} \end{gather*} $$

变换后得到

\[ \left( 1+\frac{iH_{i,j}^{n+1}\Delta t}{2\hbar} \right)\varPsi^{n+1}_{i,j} = \left( 1-\frac{iH_{i,j}^n\Delta t}{2\hbar} \right) \varPsi^{n}_{i,j} \]

然后进一步对其展开

$$ \begin{align*} &\varPsi_{i,j}^{n+1} + \frac{i\Delta t}{2\hbar} \left[ \frac{\hbar^2}{2m}\left( \frac{\varPsi_{i,j-1}^{n+1}-2\varPsi_{i,j}^{n+1}+\varPsi_{i,j+1}^{n+1}}{(\Delta x)^2}+\frac{\varPsi_{i-1,j}^{n+1}-2\varPsi_{i,j}^{n+1}+\varPsi_{i+1,j}^{n+1}}{(\Delta y)^2} \right) + V^{n+1}_{i,j}\varPsi^{n+1}_{i,j} \right] \\ =&\varPsi_{i,j}^{n} – \frac{i\Delta t}{2\hbar} \left[ \frac{\hbar^2}{2m}\left( \frac{\varPsi_{i,j-1}^{n}-2\varPsi_{i,j}^{n}+\varPsi_{i,j+1}^{n}}{(\Delta x)^2}+\frac{\varPsi_{i-1,j}^{n}-2\varPsi_{i,j}^{n}+\varPsi_{i+1,j}^{n}}{(\Delta y)^2} \right) + V^n_{i,j}\varPsi^{n}_{i,j} \right] \end{align*} $$

可以看到,当前有一个挑战便是,如果依然将$\varPsi$记为一个$N+1\times N+1$的矩阵,将难以用矩阵乘法描述$V_{i,j}\varPsi^{n}_{i,j}$这一项。因此选择将$\varPsi$展开为向量$\boldsymbol{\Psi}=[\varPsi_{0,0},\cdots ,\varPsi_{N,0},\varPsi_{0,1},\cdots ,\varPsi_{N,1},\cdots ,\varPsi_{N,N}]^{\mathrm{T}}$,这样得到了一个从标号$(i,j)$到标号$k$的变换:$k=i+Nj+j$

\[ \frac{\partial ^2}{\partial x^2}\boldsymbol{\Psi}_k = \left\{ \begin{aligned} &-2 \boldsymbol{\Psi}_k + \boldsymbol{\Psi}_{k+N+1} & 0\leq k \leq N \\ &-2 \boldsymbol{\Psi}_k + \boldsymbol{\Psi}_{k-N-1} & N^2+N \leq k \leq N^2+2N \\ &\boldsymbol{\Psi}_{k-N-1}-2\boldsymbol{\Psi}_{k} + \boldsymbol{\Psi}_{k+N+1} & \mathrm{else} \end{aligned} \right. \] \[ \frac{\partial ^2}{\partial y^2}\boldsymbol{\Psi}_k = \left\{ \begin{aligned} &-2 \boldsymbol{\Psi}_k + \boldsymbol{\Psi}_{k+1} & k=(N+1)m \\ &-2 \boldsymbol{\Psi}_k + \boldsymbol{\Psi}_{k-1} & k=(N+1)m+N \\ & \boldsymbol{\Psi}_{k-1}-2\boldsymbol{\Psi}_{k} + \boldsymbol{\Psi}_{k+1}& \mathrm{else} \\ &m\in \mathbb{Z} \end{aligned} \right. \] $$ \begin{align} &\boldsymbol{\Psi}^{n+1} + \frac{i\Delta t}{2\hbar} \left[ \frac{\hbar^2}{2m}\left( \frac{1}{(\Delta x)^2}\boldsymbol{A}+\frac{1}{(\Delta y)^2}\boldsymbol{B} \right) + \boldsymbol{V}^{n+1} \right] \boldsymbol{\Psi}^{n+1} \notag \\ =&\boldsymbol{\Psi}^{n} – \frac{i\Delta t}{2\hbar} \left[ \frac{\hbar^2}{2m}\left( \frac{1}{(\Delta x)^2}\boldsymbol{A}+\frac{1}{(\Delta y)^2}\boldsymbol{B} \right) + \boldsymbol{V}^n \right] \boldsymbol{\Psi}^{n} \end{align} $$

其中矩阵$\boldsymbol{V}$形如

\[ \boldsymbol{V}= \begin{bmatrix} V_0 & &0 \\ & \ddots & \\ 0& & V_N \\ \end{bmatrix} ,\, V_{k} = \begin{bmatrix} V_{0,k} & & 0 \\ &\ddots & \\ 0 & & V_{N,k} \end{bmatrix} \]

变形后得到

$$ \begin{align*} \boldsymbol{\Psi}^{n+1} =& \left( \boldsymbol{I} + \frac{i\Delta t}{2\hbar} \left[ \frac{\hbar^2}{2m}\left( \frac{1}{(\Delta x)^2}\boldsymbol{A}+\frac{1}{(\Delta y)^2}\boldsymbol{B} \right) – \boldsymbol{V}^{n+1} \right] \right)^{-1} \\ &\left( \boldsymbol{I} – \frac{i\Delta t}{2\hbar} \left[ \frac{\hbar^2}{2m}\left( \frac{1}{(\Delta x)^2}\boldsymbol{A}+\frac{1}{(\Delta y)^2}\boldsymbol{B} \right) – \boldsymbol{V}^n \right] \right) \boldsymbol{\Psi}^{n} \end{align*}$$

为了简化,令$\hbar = 1 , \, m=1$。初态为

\[    \varPsi_{t=0}=e^{-\frac{(x-0.5)^2 + (y-0.5)^2}{0.05^2}}e^{40000xi} \]

边界条件为

\[\varPsi|_{x=0}=\varPsi|_{x=1}=\varPsi|_{y=0}=\varPsi|_{y=1} = 0\]

让势场的矩阵形似一个有两个孔的挡板,然后就可以计算了。这个算法可以很快地让CPU吃满,然后迟迟不吐出结果。

模拟的结果,可以看到图像右边有干涉条纹出现

代码

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams
import math
import matplotlib.patches as mpatches
import matplotlib.transforms as mtransforms
plt.rcParams['font.sans-serif']=['SimSun']
plt.rcParams['axes.unicode_minus']=False

config = {
    "font.family": 'serif',
    "font.size": 12,
    "mathtext.fontset": 'stix',
    "font.serif": ['SimSun'],
}
rcParams.update(config)

class schrodingerEq(object):
    def __init__(self,L,n,Lt,nt):
        self.__L = L
        self.__Lt = Lt
        self.__n = n
        self.__nt = nt
        self.__m = 1
        self.__hbar = 1

        self.__dx = self.__L / self.__n
        self.__dt = self.__Lt / self.__nt

        self.__A, self.__B = self.__build_difmat()
        self.__X, self.__Y = self.__build_grid()

    def solve(self):
        """
        求解
        """
        varPsi0 = self.__get_initial_varPsi0()
        Psi0 = self.__Mat2Vector(varPsi0)

        varPsi_List = list()

        time_interval = np.ceil(self.__nt / 17)
        A, B = self.__build_difmat()
        I = np.identity((self.__n+1)**2)
        for i in range(0,self.__nt):
            t = i * self.__dt
            V = self.__get_V(t)
            term0 = self.__hbar**2/(2*self.__m) * 1/(self.__dx**2) * (A + B)
            term01 = I + 1j*self.__dt/(2*self.__hbar) * (term0 - V)
            term02 = np.linalg.inv(term01)

            term11 = I - 1j*self.__dt/(2*self.__hbar) * (term0 - V)
            term2 = term02.dot(term11)
            Psi1 = term2.dot(Psi0)

            varPsi1 = self.__Vector2Mat(Psi1)
            self.__boundary(varPsi1)
            Psi1 = self.__Mat2Vector(varPsi1)

            if i % time_interval == 0:
                varPsi_List.append(varPsi1)

            Psi0 = Psi1
        return varPsi_List

    def __build_difmat(self):
        N = self.__n
        A = np.zeros(((N + 1) ** 2, (N + 1) ** 2))
        B = A
        for k in range(0, (N + 1) ** 2):
            A[k, k] = -2
            if k <= N:
                A[k, k + N + 1] = 1
            elif k >= N ** 2 + N:
                A[k, k - N - 1] = 1
            else:
                A[k, k - N - 1] = 1
                A[k, k + N + 1] = 1
        for k in range(0, (N + 1) ** 2):
            B[k, k] = -2
            if k % (N + 1) == 0:
                B[k, k + 1] = 1
            elif (k - N) % (N + 1) == 0:
                B[k, k - 1] = 1
            else:
                B[k, k + 1] = 1
                B[k, k - 1] = 1
        return A,B

    def __build_grid(self):
        x = np.linspace(0,self.__L,self.__n+1)
        y = x
        X, Y = np.meshgrid(x,y)
        return X, Y

    def __get_initial_varPsi0(self):
        N = self.__n
        a = 0.05
        k = 40000
        x = np.linspace(0, 1, N + 1)
        y = np.linspace(0, 1, N + 1)
        X, Y = np.meshgrid(x, y)
        varPsi = np.power(np.e, -((X - 0.2) ** 2 + (Y - 0.5) ** 2) / a ** 2) * np.power(np.e, k * X * 1j)
        return varPsi

    def __get_V(self,t):
        N = self.__n
        V = np.zeros((N + 1, N + 1))
        l2 = int(np.ceil(0.6 * N))
        V[0:N + 1, l2] = 10000
        V[int(np.ceil(N / 2)) + 1:int(np.ceil(N / 2)) + 9, l2] = 0
        V[int(np.ceil(N / 2)) - 9:int(np.ceil(N / 2)) - 1, l2] = 0
        V0 = V.reshape(1, (N + 1) ** 2, order='F')[0]
        v = np.zeros(((N + 1) ** 2, (N + 1) ** 2))
        for i in range(0, (N + 1) ** 2):
            v[i, i] = V0[i]
        return v

    def __Mat2Vector(self,Mat):
        N = len(Mat[0])
        Vector = Mat.reshape(1, N ** 2, order='F')[0]
        return Vector

    def __Vector2Mat(self,Vector):
        N = np.sqrt(len(Vector))
        N = int(N)
        Mat = Vector.reshape(N, N,order='F')
        return Mat

    def __boundary(self,varPsi):
        varPsi[0,:] = 0
        varPsi[-1,:] = 0
        varPsi[:,0] = 0
        varPsi[:,-1] = 0


def data_plt(data, tl):
    """绘图程序"""
    fig, ax = plt.subplots(4, 4, figsize=(10,8), sharey='all', sharex='all')
    fig.subplots_adjust(hspace=0.02, wspace=0.02)
    n = 0
    for i in range(0, 4):
        for j in range(0, 4):
            im = ax[i, j].imshow(data[n], cmap='inferno')
            n = n + 1
            ax[i, j].tick_params(direction='in')
            ax[i, j].tick_params(length=4)
            ax[i, j].tick_params(axis='x', colors='white')
            ax[i, j].tick_params(axis='y', colors='white')
            for ll in ax[i, j].get_yticklabels():
                ll.set_color('k')
            for ll in ax[i, j].get_xticklabels():
                ll.set_color('k')
            ax[i, j].text(5,25, r'$t=%f$' % tl[n], c='white',
                          fontdict={'family': 'Times New Roman', 'size': 10})
    cbar = fig.colorbar(im, ax=ax,pad=0.01,aspect=50)
    fig.text(0.065,0.5,r"$32 \times 32$"
             , va='center', rotation='vertical',
             fontdict={'family': 'Times New Roman', 'size': 12})
    plt.savefig(r'E:\Latex\Luo\baogao4\codes\q0.pdf', bbox_inches='tight')
    plt.show()

def proj(L,n,Lt,nt):
    proj1 = schrodingerEq(L,n,Lt,nt)
    data0 = proj1.solve()
    data = list()
    for i in range(0, 16):
        d = data0[i]
        d1 = np.conjugate(d)
        d2 = np.abs(d * d1)
        data.append(d2)
    time_list = np.linspace(0,Lt,16)
    time_list = np.hstack(([0],time_list))
    return data, time_list

L = 1
n = 31
Lt = 0.03
nt = 600
data, tl = proj(L,n,Lt,nt)
data_plt(data,tl)

# #
首页      科研      代码      [算法]-一个笨拙的二维薛定谔方程数值模拟

发表评论

textsms
account_circle
email

Mid9ts的博客

[算法]-一个笨拙的二维薛定谔方程数值模拟
使用了Crank-Nicolson算法。知识欠缺,只能用大量的计算来填补。 二维薛定谔方程如下 \[ i\hbar \frac{\partial }{\partial t}\varPsi = \underbrace{\left[ -\frac{\hbar^2}{2m} \…
扫描二维码继续阅读
2022-05-25