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)
2D Schrodinger's Equation - Finite Difference
原文:https://www.cnblogs.com/xxhbdk/p/14497468.html