def solvePressure(P_hat, U_hat):
global F_tmp, F_tmp2
U_tmp4[:] = 0
U_tmp3[:] = 0
F_tmp2[:] = 0
Ni = F_tmp2
# dudx = 0 from continuity equation. Use Shen Dirichlet basis
# Use regular Chebyshev basis for dvdx and dwdx
F_tmp[0] = Cm.matvec(U_hat[0])
F_tmp[0, u_slice] = SFTc.TDMA_3D(a0, b0, bc, c0, F_tmp[0, u_slice])
dudx = U_tmp4[0] = ifst(F_tmp[0], U_tmp4[0], ST)
SFTc.Mult_DPhidT_3D(N[0], U_hat[1], U_hat[2], F_tmp[1], F_tmp[2])
dvdx = U_tmp4[1] = ifct(F_tmp[1], U_tmp4[1])
dwdx = U_tmp4[2] = ifct(F_tmp[2], U_tmp4[2])
dudy_h = 1j * K[1] * U_hat[0]
dudy = U_tmp3[0] = ifst(dudy_h, U_tmp3[0], ST)
dudz_h = 1j * K[2] * U_hat[0]
dudz = U_tmp3[1] = ifst(dudz_h, U_tmp3[1], ST)
Ni[0] = fst(U0[0] * dudx + U0[1] * dudy + U0[2] * dudz, Ni[0], ST)
U_tmp3[:] = 0
dvdy_h = 1j * K[1] * U_hat[1]
dvdy = U_tmp3[0] = ifst(dvdy_h, U_tmp3[0], ST)
dvdz_h = 1j * K[2] * U_hat[1]
dvdz = U_tmp3[1] = ifst(dvdz_h, U_tmp3[1], ST)
Ni[1] = fst(U0[0] * dvdx + U0[1] * dvdy + U0[2] * dvdz, Ni[1], ST)
U_tmp3[:] = 0
dwdy_h = 1j * K[1] * U_hat[2]
dwdy = U_tmp3[0] = ifst(dwdy_h, U_tmp3[0], ST)
dwdz_h = 1j * K[2] * U_hat[2]
dwdz = U_tmp3[1] = ifst(dwdz_h, U_tmp3[1], ST)
Ni[2] = fst(U0[0] * dwdx + U0[1] * dwdy + U0[2] * dwdz, Ni[2], ST)
F_tmp[0] = 0
SFTc.Mult_Div_3D(N[0], K[1, 0], K[2, 0], Ni[0, u_slice], Ni[1, u_slice],
Ni[2, u_slice], F_tmp[0, p_slice])
SFTc.Solve_Helmholtz_3D_complex(N[0], 1, F_tmp[0, p_slice], P_hat[p_slice],
u0N, u1N, u2N, LN)
return P_hat
def steps():
global t, tstep, e0, dU, U_hat, P_hat, Pcorr, U_hat1, U_hat0, P
while t < T - 1e-8:
#plt.pause(2)
t += dt
tstep += 1
#print "tstep ", tstep
# Tentative momentum solve
for jj in range(velocity_pressure_iters):
dU[:] = 0
dU = ComputeRHS(dU, jj)
SFTc.Solve_Helmholtz_3D_complex(N[0], 0, dU[0, u_slice],
U_hat[0, u_slice], u0, u1, u2, L0)
SFTc.Solve_Helmholtz_3D_complex(N[0], 0, dU[1, u_slice],
U_hat[1, u_slice], u0, u1, u2, L0)
SFTc.Solve_Helmholtz_3D_complex(N[0], 0, dU[2, u_slice],
U_hat[2, u_slice], u0, u1, u2, L0)
#SFTc.Solve_Helmholtz_3Dall_complex(N[0], 0, dU[:, u_slice], U_hat[:, u_slice], u0, u1, u2, L0)
# Pressure correction
dU[3] = pressurerhs(U_hat, dU[3])
SFTc.Solve_Helmholtz_3D_complex(N[0], 1, dU[3, p_slice],
Pcorr[p_slice], u0N, u1N, u2N, LN)
# Update pressure
#dU[3] *= (-dt) # Get div(u) in dU[3]
#dU[3, p_slice] = SFTc.TDMA_3D_complex(a0N, b0N, bcN, c0N, dU[3, p_slice])
#P_hat[p_slice] += (Pcorr[p_slice] - nu*dU[3, p_slice])
P_hat[p_slice] += Pcorr[p_slice]
#if jj == 0:
#print " Divergence error"
#print " Pressure correction norm %2.6e" %(linalg.norm(Pcorr))
# Update velocity
dU[:] = 0
pressuregrad(Pcorr, dU)
dU[0, u_slice] = SFTc.TDMA_3D(a0, b0, bc, c0, dU[0, u_slice])
dU[1, u_slice] = SFTc.TDMA_3D(a0, b0, bc, c0, dU[1, u_slice])
dU[2, u_slice] = SFTc.TDMA_3D(a0, b0, bc, c0, dU[2, u_slice])
U_hat[:3,
u_slice] += dt * dU[:3,
u_slice] # + since pressuregrad computes negative pressure gradient
for i in range(3):
U[i] = ifst(U_hat[i], U[i], ST)
# Rotate velocities
U_hat1[:] = U_hat0
U_hat0[:] = U_hat
U0[:] = U
P = ifst(P_hat, P, SN)
conv1[:] = conv0
timer()
if tstep % check_step == 0:
if case == "OS":
pert = (U[1] - (1 - X[0]**2))**2 + U[0]**2
e1 = 0.5 * energy(pert)
exact = exp(2 * imag(OS.eigval) * t)
if rank == 0:
print "Time %2.5f Norms %2.12e %2.12e %2.12e" % (
t, e1 / e0, exact, e1 / e0 - exact)
elif case == "MKK":
e0 = energy(U0[0]**2)
e1 = energy(U0[2]**2)
if rank == 0:
print "Time %2.5f Energy %2.12e %2.12e " % (t, e0, e1)
#if tstep == 100:
#Source[2, :] = 0
#Sk[2] = fss(Source[2], Sk[2], ST)
if tstep % plot_step == 0 and enable_plotting:
if case == "OS":
im1.ax.clear()
im1.ax.contourf(X[1, :, :, 0], X[0, :, :, 0],
U[1, :, :, 0] - (1 - X[0, :, :, 0]**2), 100)
im1.autoscale()
im2.ax.clear()
im2.ax.contourf(X[1, :, :, 0], X[0, :, :, 0], U[0, :, :, 0],
100)
im2.autoscale()
im3.ax.clear()
im3.ax.contourf(X[1, :, :, 0], X[0, :, :, 0], P[:, :, 0], 100)
im3.autoscale()
im4.set_UVC(U[1, :, :, 0] - (1 - X[0, :, :, 0]**2), U[0, :, :,
0])
elif case == "MKK":
im1.ax.clear()
im1.ax.contourf(X[1, :, :, 0], X[0, :, :, 0], U[1, :, :, 0],
100)
im1.autoscale()
im2.ax.clear()
im2.ax.contourf(X[1, :, :, 0], X[0, :, :, 0], U[0, :, :, 0],
100)
im2.autoscale()
im3.ax.clear()
im3.ax.contourf(X[1, :, :, 0], X[0, :, :, 0], P[:, :, 0], 100)
im3.autoscale()
plt.pause(1e-6)
