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, 实现如下,
    2D Schrodinger's Equation - Finite Difference
      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
  • 结果展示:
    2D Schrodinger's Equation - Finite Difference注意, 以上为概率密度分布情况, 即波函数之模方.
  • 使用建议:
    ①. 虚数单位在程序中可直接参与运算;
    ②. 判断数值精度是否达标: 增加网格密度, 当前后数值解重合较好时, 则精度达标(绘图查看即可).
  • 参考文档:
    吴一东. 数值方法5:偏微分方程3:二维热传导方程和薛定谔方程. bilibili, 2020.
上一篇:3.Android网络编程-http介绍


下一篇:Mybatis之if标签使用(Insert操作)