import numpy as np
import matplotlib.pyplot as plt
nx = 41 # 網(wǎng)格節(jié)點(diǎn)數(shù)
dx = 2/(nx - 1) # 網(wǎng)格尺寸
nt = 40       # 時(shí)間步數(shù)
nu = 0.3 # 擴(kuò)散系數(shù)
sigma = 0.2  #定義一個(gè)中間變量，用于確定時(shí)間步長(zhǎng)dt
dt = sigma * dx**2 / nu #時(shí)間步長(zhǎng)
# 定義初始條件
u = np.ones(nx)
u[int(0.5 / dx):int(1 / dx + 1)] = 2
un = np.ones(nx) # 定義一個(gè)定義數(shù)組
plt.ion()
plt.show()
for n in range(nt): #時(shí)間循環(huán)，計(jì)算每一個(gè)時(shí)間點(diǎn)上網(wǎng)格點(diǎn)上的數(shù)據(jù)
   plt.cla()
   un = u.copy()
   for i in range(1,nx-1): #空間循環(huán)，計(jì)算u[i]
       u[i] = un[i]+ nu *dt /dx**2 * (un[i+1] -2 * un[i] + un[i-1])
   # 下面為圖形繪制代碼
   plt.plot(np.linspace(0,2,nx) , u , 'm' , lw =3 )  #繪制速度隨時(shí)間變化曲線
   plt.ylim(0.9,2.1)          # 確定y軸上下限
   plt.title("times: %.4f s"% (n*dt)) # 將時(shí)間打印到圖形上
   plt.pause(0.1) #控制圖形顯示時(shí)間
plt.ioff()
plt.show()