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)
View Code
注意, 以上为幅值分布情况, 为加强比对, 每一帧之色彩分布均是相对的.