2D Schrodinger's Equation - Finite Difference

• 二维薛定谔方程初边值问题:
  二维薛定谔方程如下,
  \begin{equation}
  \mathrm{i}\hbar\frac{\partial\psi}{\partial t} = -\frac{\hbar^2}{2m}\left( \frac{\partial^2}{\partial x^2} + \frac{\partial^2}{\partial y^2} \right)\psi + V(x, y)\psi, \quad (x, y)\in \Omega = [-L_x, L_x]\times [-L_y, L_y], t \in [0, T]
  \label{eq_1}
  \end{equation}
  初始条件如下,
  \begin{equation}
  \psi(x, y, 0) = \psi_0(x, y), \quad (x, y) \in \Omega
  \label{eq_2}
  \end{equation}
  边界条件如下,
  \begin{equation}
  \psi(x, y, t) = 0, \quad (x, y) \in \partial\Omega, t \in [0, T]
  \label{eq_3}
  \end{equation}
• 连续型系统的离散化:
  本文拟采用中心差分离散化式$\eqref{eq_1}$,
  \begin{equation}
  \mathrm{i}\hbar\frac{\partial \psi_{i, j}}{\partial t} = -\frac{\hbar^2}{2m}\left( \frac{\psi_{i+1, j} + \psi_{i-1, j} - 2\psi_{i,j}}{h_x^2} + \frac{\psi_{i, j+1} + \psi_{i, j-1} - 2\psi_{i, j}}{h_y^2} \right) + V_{i, j}\psi_{i, j}
  \label{eq_4}
  \end{equation}
  令,
  \begin{equation*}
  U = \begin{bmatrix}
  \psi_{0, 0} & \cdots & \psi_{0, n_x}   \\
  \vdots & \ddots & \vdots   \\
  \psi_{n_y, 0} & \cdots & \psi_{n_y, n_x}   \\
  \end{bmatrix}\quad
  V = \begin{bmatrix}
  V_{0, 0} & \cdots & V_{0, n_x}   \\
  \vdots & \ddots & \vdots   \\
  V_{n_y, 0} & \cdots & V_{n_y, n_x}   \\
  \end{bmatrix}\quad
  A_x = \begin{bmatrix}
  -2 & 1 & & &   \\
  1 & -2 & 1 & &   \\
  & \ddots & \ddots & \ddots &   \\
  & & 1 & -2 & 1   \\
  & & & 1 & -2   \\
  \end{bmatrix}\quad
  A_y = \begin{bmatrix}
  -2 & 1 & & &   \\
  1 & -2 & 1 & &   \\
  & \ddots & \ddots & \ddots &   \\
  & & 1 & -2 & 1   \\
  & & & 1 & -2   \\
  \end{bmatrix}
  \end{equation*}
  式$\eqref{eq_4}$转换如下,
  \begin{equation}
  \mathrm{i}\hbar\frac{\partial U}{\partial t} = -\frac{\hbar^2}{2m}\left( \frac{A_yU}{h_y^2} + \frac{UA_x}{h_x^2} \right) + V\odot U
  \label{eq_5}
  \end{equation}
  注意, 上式$\eqref{eq_5}$中$U$之边界值需以边界条件式$\eqref{eq_3}$填充. 同时, 将上式$\eqref{eq_5}$转换如下,
  \begin{equation}
  \frac{\partial U}{\partial t} = \frac{\mathrm{i}\hbar}{2m}\left( \frac{A_yU}{h_y^2} + \frac{UA_x}{h_x^2} \right) -\frac{\mathrm{i}}{\hbar}V\odot U
  \label{eq_6}
  \end{equation}
  本文拟采用Runge-Kutta方法求解上式$\eqref{eq_6}$, 则有,
  \begin{equation}
  U^{k+1} = U^k + \frac{1}{6}(K_1 + 2K_2 + 2K_3 + K_4)h_t
  \label{eq_7}
  \end{equation}
  注意, 上式$\eqref{eq_7}$中$U$之初始值$U^0$需以初始条件式$\eqref{eq_2}$填充. 其中,
  \begin{gather*}
  F(U) = \frac{\mathrm{i}\hbar}{2m}\left( \frac{A_yU}{h_y^2} + \frac{UA_x}{h_x^2} \right) -\frac{\mathrm{i}}{\hbar}V\odot U \\
  K_1 = F(U)   \\
  K_2 = F(U + \frac{h_t}{2}K_1)     \\
  K_3 = F(U + \frac{h_t}{2}K_2)     \\
  K_4 = F(U + h_tK_3)     \\
  \end{gather*}
  由此, 根据式$\eqref{eq_7}$即可完成式$\eqref{eq_1}$二维薛定谔方程的数值求解.
• 代码实现:
  Part Ⅰ:
  现以如下势能函数及初始条件为例进行算法实施,
  \begin{gather*}
  V(x, y) = \frac{1}{2}m\omega_x^2x^2 + \frac{1}{2}m\omega_y^2y^2   \\
  \psi_0(x, y) = \mathrm{e}^{-\alpha(x^2 + y^2)} \cdot \mathrm{e}^{\mathrm{i}\beta x}
  \end{gather*}
  Part Ⅱ:
  利用finite difference求解2D Schrodinger‘s equation, 实现如下,
  
  

    1 # Schrodinger‘s equation之实现 - 有限差分法
    2 
    3 import numpy
    4 from matplotlib import pyplot as plt
    5 from matplotlib import animation
    6 
    7 
    8 class SchrodingerEq(object):
    9     
   10     def __init__(self, nx=150, ny=100, nt=10000, Lx=1.5, Ly=1, Lt=1):
   11         self.__nx = nx                                 # x轴子区间划分数
   12         self.__ny = ny                                 # y轴子区间划分数
   13         self.__nt = nt                                 # t轴子区间划分数
   14         self.__Lx = Lx                                 # x轴半长
   15         self.__Ly = Ly                                 # y轴半长
   16         self.__Lt = Lt                                 # t轴全长
   17         self.__hbar = 1                                # 普朗克常数取值
   18         self.__m = 1                                   # 质量取值
   19         
   20         self.__hx = 2 * Lx / nx
   21         self.__hy = 2 * Ly / ny
   22         self.__ht = Lt / nt
   23         
   24         self.__X, self.__Y = self.__build_gridPoints()
   25         self.__T = numpy.linspace(0, Lt, nt + 1)
   26         self.__Ax, self.__Ay = self.__build_2ndOrdMat()
   27         self.__V = self.__get_V()
   28         
   29         
   30     def get_solu(self):
   31         ‘‘‘
   32         数值求解
   33         ‘‘‘
   34         UList = list()
   35         
   36         U0 = self.__get_initial_U0()
   37         self.__fill_boundary_U(U0)
   38         UList.append(U0)
   39         for t in self.__T[:-1]:
   40             Uk = self.__calc_Uk(t, U0)
   41             UList.append(Uk)
   42             U0 = Uk
   43         
   44         return self.__X, self.__Y, self.__T, UList
   45     
   46     
   47     def __calc_Uk(self, t, U0):
   48         K1 = self.__calc_F(U0)
   49         self.__fill_boundary_U(K1)
   50         
   51         K2 = self.__calc_F(U0 + self.__ht / 2 * K1)
   52         self.__fill_boundary_U(K2)
   53         
   54         K3 = self.__calc_F(U0 + self.__ht / 2 * K2)
   55         self.__fill_boundary_U(K3)
   56         
   57         K4 = self.__calc_F(U0 + self.__ht * K3)
   58         self.__fill_boundary_U(K4)
   59         
   60         Uk = U0 + (K1 + 2 * K2 + 2 * K3 + K4) / 6 * self.__ht
   61         return Uk
   62         
   63         
   64     def __calc_F(self, U):
   65         term0 = numpy.matmul(self.__Ay, U) / self.__hy ** 2
   66         term1 = numpy.matmul(U, self.__Ax) / self.__hx ** 2
   67         term2 = -1j / self.__hbar * self.__V * U
   68         FVal = 1j * self.__hbar / 2 / self.__m * (term0 + term1) + term2
   69         return FVal
   70         
   71         
   72     def __get_initial_U0(self):
   73         ‘‘‘
   74         获取初始条件
   75         ‘‘‘
   76         alpha = 50
   77         beta = 5
   78         U0 = numpy.exp(-alpha * (self.__X ** 2 + self.__Y ** 2)) * numpy.exp(1j * beta * self.__X)
   79         return U0
   80         
   81         
   82     def __fill_boundary_U(self, U):
   83         ‘‘‘
   84         填充边界条件
   85         ‘‘‘
   86         U[0, :] = 0
   87         U[-1, :] = 0
   88         U[:, 0] = 0
   89         U[:, -1] = 0
   90         
   91         
   92     def __get_V(self):
   93         ‘‘‘
   94         获取势能函数
   95         ‘‘‘
   96         omegax = 1
   97         omegay = 2
   98         V = 0.5 * self.__m * omegax ** 2 * self.__X ** 2 + 0.5 * self.__m * omegay ** 2 * self.__Y ** 2
   99         return V
  100         
  101         
  102     def __build_2ndOrdMat(self):
  103         ‘‘‘
  104         构造2阶微分算子的矩阵形式
  105         ‘‘‘
  106         Ax = self.__build_AMat(self.__nx + 1)
  107         Ay = self.__build_AMat(self.__ny + 1)
  108         return Ax, Ay
  109         
  110         
  111     def __build_AMat(self, n):
  112         term0 = numpy.identity(n) * -2
  113         term1 = numpy.eye(n, n, 1)
  114         term2 = numpy.eye(n, n, -1)
  115         AMat = term0 + term1 + term2
  116         return AMat
  117         
  118         
  119     def __build_gridPoints(self):
  120         ‘‘‘
  121         构造网格节点
  122         ‘‘‘
  123         X = numpy.linspace(-self.__Lx, self.__Lx, self.__nx + 1)
  124         Y = numpy.linspace(-self.__Ly, self.__Ly, self.__ny + 1)
  125         X, Y = numpy.meshgrid(X, Y)
  126         return X, Y
  127         
  128         
  129 class SchrodingerPlot(object):
  130     
  131     fig = None
  132     ax1 = None
  133     line = None
  134     txt = None
  135     X, Y, T, Z = None, None, None, None
  136     
  137     @classmethod
  138     def plot_ani(cls, schObj):
  139         cls.X, cls.Y, cls.T, UList = schObj.get_solu()
  140         cls.ZList = list(numpy.abs(item) for item in UList)
  141     
  142         cls.fig = plt.figure(figsize=(6, 4))
  143         cls.ax1 = plt.subplot()
  144         cls.line = cls.ax1.pcolormesh(cls.X, cls.Y, cls.ZList[0][:-1, :-1], cmap="jet", vmin=0)
  145 
  146         ani = animation.FuncAnimation(cls.fig, cls.update, frames=numpy.arange(1, len(cls.ZList), 100), blit=True, interval=5, repeat=True)
  147         
  148         ani.save("plot_ani.gif", writer="PillowWriter", fps=200)
  149         plt.close()
  150         
  151         
  152     @classmethod
  153     def update(cls, frame):
  154         cls.line.set_array(cls.ZList[frame][:-1, :-1])
  155         return cls.line,
  156         
  157         
  158         
  159 if __name__ == "__main__":
  160     nx = 150
  161     ny = 100
  162     nt = 20000
  163     Lx = 1.5
  164     Ly = 1
  165     Lt = 1
  166     schObj = SchrodingerEq(nx, ny, nt, Lx, Ly, Lt)
  167     
  168     SchrodingerPlot.plot_ani(schObj)

  View Code
• 结果展示:
  注意, 以上为概率密度分布情况, 即波函数之模方.
• 使用建议:
  ①. 虚数单位在程序中可直接参与运算;
  ②. 判断数值精度是否达标: 增加网格密度, 当前后数值解重合较好时, 则精度达标(绘图查看即可).
• 参考文档:
  吴一东. 数值方法5:偏微分方程3:二维热传导方程和薛定谔方程. bilibili, 2020.

2D Schrodinger's Equation - Finite Difference

原文：https://www.cnblogs.com/xxhbdk/p/14497468.html