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
注意, 以上为概率密度分布情况, 即波函数之模方.