import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

nx = 31  # 定義x方向節(jié)點(diǎn)數(shù)
ny = 31  # 定義y方向節(jié)點(diǎn)數(shù)
nt = 17  # 定義時(shí)間步長(zhǎng)
nu  = 0.05  # 定義擴(kuò)散系數(shù)
dx = 2/(nx-1)  # 定義x方向網(wǎng)格尺寸
dy = 2/(ny-1)  # 定義y方向網(wǎng)格尺寸
sigma = 0.25   # 定義庫(kù)朗數(shù)
dt = sigma * dx * dy /nu

x = np.linspace(0,2,nx)   # x取值范圍
y = np.linspace(0,2,ny)   # y取值范圍
# 定義邊界條件
u = np.ones((ny,nx))
un = np.ones((ny,nx))
# 定義初始條件
u[int(.5 / dy):int(1 / dy + 1),int(.5 / dx):int(1 / dx + 1)] = 2

# 繪制初始條件
fig = plt.figure()
ax = fig.gca(projection='3d')
X, Y = np.meshgrid(x, y)
surf = ax.plot_surface(X, Y, u, rstride=1, cstride=1, cmap=cm.viridis,
       linewidth=0, antialiased=False)
ax.set_xlim(0, 2)
ax.set_ylim(0, 2)
ax.set_zlim(1, 2.5)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')
plt.show()

# 定義函數(shù)進(jìn)行求解
def diffuse(nt):
   u[int(.5 / dy):int(1 / dy + 1),int(.5 / dx):int(1 / dx + 1)] = 2

   for n in range(nt + 1):
       un = u.copy()
       u[1:-1, 1:-1] = (un[1:-1,1:-1] +
                       nu * dt / dx**2 *
                       (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, 0:-2]) +
                       nu * dt / dy**2 *
                       (un[2:,1: -1] - 2 * un[1:-1, 1:-1] + un[0:-2, 1:-1]))
       u[0, :] = 1
       u[-1, :] = 1
       u[:, 0] = 1
       u[:, -1] = 1


   fig = plt.figure()
   ax = fig.gca(projection='3d')
   surf = ax.plot_surface(x, y, u[:], rstride=1, cstride=1, cmap=cm.viridis,
       linewidth=0, antialiased=True)
   ax.set_zlim(1, 2.5)
   ax.set_xlabel('$x$')
   ax.set_ylabel('$y$')
   plt.show()