#%% import libraries
import cupy as cp
import matplotlib.pyplot as plt
#%% define parameters
l_cr = 1 # characteristic length
h = l_cr / 100 # grid spacing
dt = 0.001 # time step
L = 1 # domain length
D = 1 # domain depth
W = 1 # domain width
Nx = round(L / h) # grid points in x-axis
Ny = round(D / h) # grid points in y-axis
Nz = round(W / h) # grid points in z-axis
nu = 1 / 1000 # kinematic viscosity
Uinf = 1 # free stream velocity / inlet velocity / lid velocity
cfl = dt * Uinf / h # cfl number
travel = 10 # times the disturbance travels entire length of computational domain
TT = travel * L / Uinf # total time
ns = int(TT / dt) # number of time steps
Re = round(l_cr * Uinf / nu) # Reynolds number
#%% intialization
u = cp.zeros((Nx, Ny, Nz)) # x-velocity
v = cp.zeros((Nx, Ny, Nz)) # y-velocity
w = cp.zeros((Nx, Ny, Nz)) # z-velocity
p = cp.zeros((Nx, Ny, Nz)) # pressure
X, Y, Z = cp.meshgrid(cp.linspace(0, L, Nx), cp.linspace(0, D, Ny), cp.linspace(0, W, Nz), indexing = 'ij') # spatial grid
#%% pre calculate for speed
P1 = 1 / (2 * h * dt)
P2 = 1 / (4 * h * h)
P3 = h**2
P4 = 1 / 6
P5 = (2 / Re) * dt / h**2
P6 = dt / h
P7 = 1 - (6 * P5)
P8 = 0.75 # under relaxation
P9 = 1 - P8
#%% solve 3D Navier-Stokes equations
for nt in range(ns):
dudx = u[2:, 1:-1, 1:-1] - u[:-2, 1:-1, 1:-1]
dvdy = v[1:-1, 2:, 1:-1] - v[1:-1, :-2, 1:-1]
dwdz = w[1:-1, 1:-1, 2:] - w[1:-1, 1:-1, :-2]
pn = p.copy()
b = P1 * (dudx + dvdy + dwdz) - P2 * (dudx**2 + dvdy**2 + dwdz**2 + 2 * ((u[1:-1, 2:, 1:-1] - u[1:-1, :-2, 1:-1]) * (v[2:, 1:-1, 1:-1] - v[:-2, 1:-1, 1:-1]) + (u[1:-1, 1:-1, 2:] - u[1:-1, 1:-1, :-2]) * (w[2:, 1:-1, 1:-1] - w[:-2, 1:-1, 1:-1]) + (v[1:-1, 1:-1, 2:] - v[1:-1, 1:-1, :-2]) * (w[1:-1, 2:, 1:-1] - w[1:-1, :-2, 1:-1]))) # divergence
p[1:-1, 1:-1, 1:-1] = P4 * (pn[2:, 1:-1, 1:-1] + pn[:-2, 1:-1, 1:-1] + pn[1:-1, 2:, 1:-1] + pn[1:-1, :-2, 1:-1] + pn[1:-1, 1:-1, 2:] + pn[1:-1, 1:-1, :-2] - P3 * b) # mass
p[0, :, :] = p[1, :, :] # dp/dx = 0 at x = 0
p[-1, :, :] = p[-2, :, :] # dp/dx = 0 at x = L
p[:, 0, :] = p[:, 1, :] # dp/dy = 0 at y = 0
p[:, -1, :] = p[:, -2, :] # dp/dy = 0 at y = D
p[:, :, 0] = p[:, :, 1] # dp/dz = 0 at z = 0
p[:, :, -1] = p[:, :, -2] # dp/dz = 0 at z = H
p = P8 * p + P9 * pn
un = u.copy()
vn = v.copy()
wn = w.copy()
u[1:-1, 1:-1, 1:-1] = un[1:-1, 1:-1, 1:-1] * P7 - P6 * (un[1:-1, 1:-1, 1:-1] * (un[2:, 1:-1, 1:-1] - un[:-2, 1:-1, 1:-1]) + vn[1:-1, 1:-1, 1:-1] * (un[1:-1, 2:, 1:-1] - un[1:-1, :-2, 1:-1]) + wn[1:-1, 1:-1, 1:-1] * (un[1:-1, 1:-1, 2:] - un[1:-1, 1:-1, :-2]) + p[2:, 1:-1, 1:-1] - p[:-2, 1:-1, 1:-1]) + P5 * (un[2:, 1:-1, 1:-1] + un[:-2, 1:-1, 1:-1] + un[1:-1, 2:, 1:-1] + un[1:-1, :-2, 1:-1] + un[1:-1, 1:-1, 2:] + un[1:-1, 1:-1, :-2]) # x momentum
u[0, :, :] = 0 # u = 0 at x = 0
u[-1, :, :] = 0 # u = 0 at # x = L
u[:, 0, :] = 0 # u = 0 at y = 0
u[:, -1, :] = 0 # u = 0 at y = D
u[:, :, 0] = 0 # u = 0 at z = 0
u[:, :, -1] = Uinf # u = Uinf at z = W
u = P8 * u + P9 * un
v[1:-1, 1:-1, 1:-1] = vn[1:-1, 1:-1, 1:-1] * P7 - P6 * (un[1:-1, 1:-1, 1:-1] * (vn[2:, 1:-1, 1:-1] - vn[:-2, 1:-1, 1:-1]) + vn[1:-1, 1:-1, 1:-1] * (vn[1:-1, 2:, 1:-1] - vn[1:-1, :-2, 1:-1]) + wn[1:-1, 1:-1, 1:-1] * (vn[1:-1, 1:-1, 2:] - vn[1:-1, 1:-1, :-2]) + p[1:-1, 2:, 1:-1] - p[1:-1, :-2, 1:-1]) + P5 * (vn[2:, 1:-1, 1:-1] + vn[:-2, 1:-1, 1:-1] + vn[1:-1, 2:, 1:-1] + vn[1:-1, :-2, 1:-1] + vn[1:-1, 1:-1, 2:] + vn[1:-1, 1:-1, :-2]) # y momentum
v[0, :, :] = 0 # v = 0 at x = 0
v[-1, :, :] = 0 # v = 0 at # x = L
v[:, 0, :] = 0 # v = 0 at y = 0
v[:, -1, :] = 0 # v = 0 at y = D
v[:, :, 0] = 0 # v = 0 at z = 0
v[:, :, -1] = 0 # v = 0 at z = W
v = P8 * v + P9 * vn
w[1:-1, 1:-1, 1:-1] = wn[1:-1, 1:-1, 1:-1] * P7 - P6 * (un[1:-1, 1:-1, 1:-1] * (wn[2:, 1:-1, 1:-1] - wn[:-2, 1:-1, 1:-1]) + vn[1:-1, 1:-1, 1:-1] * (wn[1:-1, 2:, 1:-1] - wn[1:-1, :-2, 1:-1]) + wn[1:-1, 1:-1, 1:-1] * (wn[1:-1, 1:-1, 2:] - wn[1:-1, 1:-1, :-2]) + p[1:-1, 1:-1, 2:] - p[1:-1, 1:-1, :-2]) + P5 * (wn[2:, 1:-1, 1:-1] + wn[:-2, 1:-1, 1:-1] + wn[1:-1, 2:, 1:-1] + wn[1:-1, :-2, 1:-1] + wn[1:-1, 1:-1, 2:] + wn[1:-1, 1:-1, :-2]) # z momentum
w[0, :, :] = 0 # w = 0 at x = 0
w[-1, :, :] = 0 # w = 0 at # x = L
w[:, 0, :] = 0 # w = 0 at y = 0
w[:, -1, :] = 0 # w = 0 at y = D
w[:, :, 0] = 0 # w = 0 at z = 0
w[:, :, -1] = 0 # w = 0 at z = W
w = P8 * w + P9 * wn
if nt % 1000 == 0:
print("% Complete", 100 * nt / ns)
#%% post process
fig = plt.figure(dpi = 500)
ax = fig.add_subplot(111, projection = '3d')
ax.contourf(X[:, Ny // 2, :].get(), u[:, Ny // 2, :].get(), Z[:, Ny // 2, :].get(), zdir = 'y', offset = D / 2, levels = 128, cmap='jet', alpha = 0.5) # vertical plane
ax.set_xlim(0, L)
ax.set_ylim(0, D)
ax.set_zlim(0, W)
ax.set_xticks([0, L])
ax.set_yticks([0, D])
ax.set_zticks([0, W])
ax.tick_params(axis='x', pad = -2)
ax.tick_params(axis='y', pad = -2)
ax.tick_params(axis='z', pad = -2)
ax.set_xlabel('x [m]', labelpad = -15)
ax.set_ylabel('y [m]', labelpad = -15)
ax.set_zlabel('z [m]', labelpad = -15)
plt.gca().set_aspect('equal')
ax.view_init(elev = 22.5, azim = -45)
plt.show()
The results from post processing are shown in Fig. 1. Within Fig. 1, the u, v and w components of velocities are shown along the plane of flow. The resulting pressure field and the velocity streamlines are shown within Fig. 2.