[问题讨论]使用Python学习CFD初级理论系列一二维Burger方程（10/10）

nx = 41
ny = 41
nt = 120
c = 1
dx = 2 / (nx - 1)
dy = 2 / (ny - 1)
sigma = .0009
nu = 0.01
dt = sigma * dx * dy / nu

x = numpy.linspace(0, 2, nx)
y = numpy.linspace(0, 2, ny)

u = numpy.ones((ny, nx))  
v = numpy.ones((ny, nx))
un = numpy.ones((ny, nx))
vn = numpy.ones((ny, nx))
comb = numpy.ones((ny, nx))

# 指定u初始化条件u(.5<=x<=1 && .5<=y<=1 )= 2
u[int(.5 / dy):int(1 / dy + 1),int(.5 / dx):int(1 / dx + 1)] = 2
# 指定v初始化条件u(.5<=x<=1 && .5<=y<=1 )= 2
v[int(.5 / dy):int(1 / dy + 1),int(.5 / dx):int(1 / dx + 1)] = 2

# 绘制初始条件
fig = pyplot.figure(figsize=(11, 7), dpi=100)
ax = fig.gca(projection='3d')
X, Y = numpy.meshgrid(x, y)
ax.plot_surface(X, Y, u[:], cmap=cm.viridis, rstride=1, cstride=1)
ax.plot_surface(X, Y, v[:], cmap=cm.viridis, rstride=1, cstride=1)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')

for n in range(nt + 1): # 时间迭代
   un = u.copy()
   vn = v.copy()

   u[1:-1, 1:-1] = (un[1:-1, 1:-1] -
                    dt / dx * un[1:-1, 1:-1] *
                    (un[1:-1, 1:-1] - un[1:-1, 0:-2]) -
                    dt / dy * vn[1:-1, 1:-1] *
                    (un[1:-1, 1:-1] - un[0:-2, 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]))

   v[1:-1, 1:-1] = (vn[1:-1, 1:-1] -
                    dt / dx * un[1:-1, 1:-1] *
                    (vn[1:-1, 1:-1] - vn[1:-1, 0:-2]) -
                    dt / dy * vn[1:-1, 1:-1] *
                   (vn[1:-1, 1:-1] - vn[0:-2, 1:-1]) +
                    nu * dt / dx**2 *
                    (vn[1:-1, 2:] - 2 * vn[1:-1, 1:-1] + vn[1:-1, 0:-2]) +
                    nu * dt / dy**2 *
                    (vn[2:, 1:-1] - 2 * vn[1:-1, 1:-1] + vn[0:-2, 1:-1]))

   u[0, :] = 1
   u[-1, :] = 1
   u[:, 0] = 1
   u[:, -1] = 1

   v[0, :] = 1
   v[-1, :] = 1
   v[:, 0] = 1
   v[:, -1] = 1

fig = pyplot.figure(figsize=(11, 7), dpi=100)
ax = fig.gca(projection='3d')
X, Y = numpy.meshgrid(x, y)
ax.plot_surface(X, Y, u, cmap=cm.viridis, rstride=1, cstride=1)
ax.plot_surface(X, Y, v, cmap=cm.viridis, rstride=1, cstride=1)
ax.set_xlabel('$x$')
ax.set_ylabel('$y$')