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.
posted @ 2021-03-14 12:25  LOGAN_XIONG  阅读(956)  评论(4)  编辑  收藏  举报