1 # 2D Wave equation之求解 - 有限差分法
  2 
  3 import numpy
  4 from matplotlib import pyplot as plt
  5 from matplotlib import animation
  6 
  7 
  8 class WaveEq(object):
  9     
 10     def __init__(self, nx=300, ny=300, nt=1000, Lx=1, Ly=1, Lt=6):
 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.__D = 0.1
 18         
 19         self.__hx = 2 * Lx / nx
 20         self.__hy = 2 * Ly / ny
 21         self.__ht = Lt / nt
 22         
 23         self.__X, self.__Y = self.__build_gridPoints()
 24         self.__T = numpy.linspace(0, Lt, nt + 1)
 25         self.__Ax, self.__Ay = self.__build_2ndOrdMat()
 26         
 27         
 28     def get_solu(self):
 29         '''
 30         数值求解
 31         '''
 32         UList = list()
 33         
 34         U0 = self.__get_initial_U0()
 35         V0 = self.__get_initial_V0()
 36         self.__fill_boundary_U(U0)
 37         UList.append(U0)
 38         for t in self.__T[:-1]:
 39             # print(t, numpy.max(U0))
 40             Uk, Vk = self.__calc_Uk_Vk(t, U0, V0)
 41             UList.append(Uk)
 42             U0, V0 = Uk, Vk
 43         
 44         return self.__X, self.__Y, self.__T, UList
 45             
 46     
 47     def __calc_Uk_Vk(self, t, U0, V0):
 48         K1 = self.__calc_F(U0, V0)
 49         tmpU, tmpV = U0 + self.__ht / 2 * K1[0], V0 + self.__ht / 2 * K1[1]
 50         self.__fill_boundary_U(tmpU)
 51         
 52         K2 = self.__calc_F(tmpU, tmpV)
 53         tmpU, tmpV = U0 + self.__ht / 2 * K2[0], V0 + self.__ht / 2 * K2[1]
 54         self.__fill_boundary_U(tmpU)
 55         
 56         K3 = self.__calc_F(tmpU, tmpV)
 57         tmpU, tmpV = U0 + self.__ht * K3[0], V0 + self.__ht * K3[1]
 58         self.__fill_boundary_U(tmpU)
 59         
 60         K4 = self.__calc_F(tmpU, tmpV)
 61         
 62         Uk = U0 + (K1[0] + 2 * K2[0] + 2 * K3[0] + K4[0]) / 6 * self.__ht
 63         Vk = V0 + (K1[1] + 2 * K2[1] + 2 * K3[1] + K4[1]) / 6 * self.__ht
 64         self.__fill_boundary_U(Uk)
 65         return Uk, Vk
 66         
 67         
 68     def __calc_F(self, U0, V0):
 69         F1 = V0
 70         term0 = numpy.matmul(U0, self.__Ax) / self.__hx ** 2
 71         term1 = numpy.matmul(self.__Ay, U0) / self.__hy ** 2
 72         F2 = self.__D * (term0 + term1)
 73         return F1, F2
 74     
 75         
 76     def __fill_boundary_U(self, U):
 77         '''
 78         填充边界条件
 79         '''
 80         U[0, :] = 0
 81         U[-1, :] = 0
 82         U[:, 0] = 0
 83         U[:, -1] = 0
 84         
 85         
 86     def __get_initial_V0(self):
 87         '''
 88         获取V之初始条件
 89         '''
 90         V0 = numpy.zeros(self.__X.shape)
 91         return V0
 92         
 93         
 94     def __get_initial_U0(self):
 95         '''
 96         获取U之初始条件
 97         '''
 98         xc = 0
 99         yc = 0
100         U0 = 0.1 * numpy.exp(-((self.__X - xc) ** 2 + (self.__Y - yc) ** 2) / 0.0005)
101         return U0
102         
103         
104     def __build_2ndOrdMat(self):
105         '''
106         构造2阶微分算子的矩阵形式
107         '''
108         Ax = self.__build_AMat(self.__nx + 1)
109         Ay = self.__build_AMat(self.__ny + 1)
110         return Ax, Ay
111         
112         
113     def __build_AMat(self, n):
114         term0 = numpy.identity(n) * -2
115         term1 = numpy.eye(n, n, 1)
116         term2 = numpy.eye(n, n, -1)
117         AMat = term0 + term1 + term2
118         return AMat
119         
120         
121     def __build_gridPoints(self):
122         '''
123         构造网格节点
124         '''
125         X = numpy.linspace(-self.__Lx, self.__Lx, self.__nx + 1)
126         Y = numpy.linspace(-self.__Ly, self.__Ly, self.__ny + 1)
127         X, Y = numpy.meshgrid(X, Y)
128         return X, Y
129         
130         
131 class WavePlot(object):
132     
133     fig = None
134     ax1 = None
135     line = None
136     X, Y, T, Z = None, None, None, None
137     
138     @classmethod
139     def plot_ani(cls, waveObj):
140         cls.X, cls.Y, cls.T, cls.Z = waveObj.get_solu()
141         
142         cls.fig = plt.figure(figsize=(10, 10), dpi=40)
143         cls.ax1 = plt.subplot()
144         cls.line = cls.ax1.pcolormesh(cls.X, cls.Y, cls.Z[0][:-1, :-1], cmap="jet")
145         
146         ani = animation.FuncAnimation(cls.fig, cls.update, frames=range(1, len(cls.T), 25), 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.Z[frame][:-1, :-1])
155         cls.line.set_norm(norm=None)
156         return cls.line,
157         
158         
159 if __name__ == "__main__":
160     nx = 300
161     ny = 300
162     nt = 4000
163     
164     Lx = 0.5
165     Ly = 0.5
166     Lt = 10
167     waveObj = WaveEq(nx, ny, nt, Lx, Ly, Lt)
168     # waveObj.get_solu()
169     
170     WavePlot.plot_ani(waveObj)