Here are the code. I really appreciate your help. Thank you.
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import pyvista
from dolfinx import fem, mesh, io, plot
from dolfinx.fem.petsc import assemble_vector, assemble_matrix, create_vector, apply_lifting, set_bc
import ufl
import numpy as np
#define the constants:
lenght = 2
width = 0.5
mu = 1e-2
rho = 1
#create the rectangular mesh
domain = mesh.create_rectangle(MPI.COMM_WORLD, [np.array([0, 0]), np.array([lenght, width])],
[41, 6], cell_type=mesh.CellType.triangle)
#time steps
t = 0
Tend = 5
num_steps = 5000
dt = Tend / num_steps
#create the function space
from basix.ufl import element
v_cg2 = element("Lagrange", domain.topology.cell_name(), 2, shape=(domain.geometry.dim, ))
s_cg1 = element("Lagrange", domain.topology.cell_name(), 1)
V = fem.functionspace(domain, v_cg2)
Q = fem.functionspace(domain, s_cg1)
#create the trial and test functions
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
p, q = ufl.TrialFunction(Q), ufl.TestFunction(Q)
#define the boundary conditions
gdim = domain.geometry.dim
fdim = domain.topology.dim - 1
#walls
def walls(x):
return np.logical_or(np.isclose(x[1], 0), np.isclose(x[1], width))
w_b_facets = mesh.locate_entities_boundary(domain, fdim, walls)
u_noslip = np.array((0,) * domain.geometry.dim, dtype=PETSc.ScalarType)
bcu_noslip = fem.dirichletbc(u_noslip, fem.locate_dofs_topological(V, fdim, w_b_facets), V)
# Inlet
def inflow(x):
return np.isclose(x[0], 0)
class InletVelocity():
def __init__(self, t):
self.t = t
def __call__(self, x):
values = np.zeros((gdim, x.shape[1]), dtype=PETSc.ScalarType)
values[0, :] = 2
return values
in_b_facets = mesh.locate_entities_boundary(domain, fdim, inflow)
u_inlet = fem.Function(V)
inlet_velocity = InletVelocity(t)
u_inlet.interpolate(inlet_velocity)
bcu_inflow = fem.dirichletbc(u_inlet, fem.locate_dofs_topological(V, fdim, in_b_facets))
#outlet
def outflow(x):
return np.isclose(x[0], lenght)
out_b_facets = mesh.locate_entities_boundary(domain, fdim, outflow)
bcp_outlet = fem.dirichletbc(PETSc.ScalarType(0), fem.locate_dofs_topological(Q, fdim, out_b_facets), Q)
bcu = [bcu_noslip, bcu_inflow]
bcp = [bcp_outlet]
#define the variational problem
u_n = fem.Function(V)
u_n.name = "u_n"
U = 0.5 * (u_n + u)
n = ufl.FacetNormal(domain)
f = fem.Constant(domain, PETSc.ScalarType((0, 0)))
k = fem.Constant(domain, PETSc.ScalarType(dt))
mu = fem.Constant(domain, PETSc.ScalarType(mu))
rho = fem.Constant(domain, PETSc.ScalarType(rho))
# Define strain-rate tensor
def epsilon(u):
return ufl.sym(ufl.nabla_grad(u))
# Define stress tensor
def sigma(u, p):
return 2 * mu * epsilon(u) - p * ufl.Identity(len(u))
# Define the variational problem for the first step
p_n = fem.Function(Q)
p_n.name = "p_n"
F1 = rho * ufl.dot((u - u_n) / k, v) * ufl.dx
F1 += rho * ufl.dot(ufl.dot(u_n, ufl.nabla_grad(u_n)), v) * ufl.dx
F1 += ufl.inner(sigma(U, p_n), epsilon(v)) * ufl.dx
F1 += ufl.dot(p_n * n, v) * ufl.ds - ufl.dot(mu * ufl.nabla_grad(U) * n, v) * ufl.ds
F1 -= ufl.dot(f, v) * ufl.dx
a1 = fem.form(ufl.lhs(F1))
L1 = fem.form(ufl.rhs(F1))
A1 = assemble_matrix(a1, bcs=bcu)
A1.assemble()
b1 = create_vector(L1)
# Define variational problem for step 2
u_ = fem.Function(V)
a2 = fem.form(ufl.dot(ufl.nabla_grad(p), ufl.nabla_grad(q)) * ufl.dx)
L2 = fem.form(ufl.dot(ufl.nabla_grad(p_n), ufl.nabla_grad(q)) * ufl.dx - (rho / k) * ufl.div(u_) * q * ufl.dx)
A2 = assemble_matrix(a2, bcs=bcp)
A2.assemble()
b2 = create_vector(L2)
# Define variational problem for step 3
p_ = fem.Function(Q)
a3 = fem.form(rho * ufl.dot(u, v) * ufl.dx)
L3 = fem.form(rho * ufl.dot(u_, v) * ufl.dx - k * ufl.dot(ufl.nabla_grad(p_ - p_n), v) * ufl.dx)
A3 = assemble_matrix(a3)
A3.assemble()
b3 = create_vector(L3)
# Solver for step 1
solver1 = PETSc.KSP().create(domain.comm)
solver1.setOperators(A1)
solver1.setType(PETSc.KSP.Type.BCGS)
pc1 = solver1.getPC()
pc1.setType(PETSc.PC.Type.HYPRE)
pc1.setHYPREType("boomeramg")
# Solver for step 2
solver2 = PETSc.KSP().create(domain.comm)
solver2.setOperators(A2)
solver2.setType(PETSc.KSP.Type.BCGS)
pc2 = solver2.getPC()
pc2.setType(PETSc.PC.Type.HYPRE)
pc2.setHYPREType("boomeramg")
# Solver for step 3
solver3 = PETSc.KSP().create(domain.comm)
solver3.setOperators(A3)
solver3.setType(PETSc.KSP.Type.CG)
pc3 = solver3.getPC()
pc3.setType(PETSc.PC.Type.SOR)
from pathlib import Path
folder = Path("results")
folder.mkdir(exist_ok=True, parents=True)
vtx_u = io.VTXWriter(domain.comm, folder / "u.bp", u_n, engine="BP4")
vtx_p = io.VTXWriter(domain.comm, folder / "p.bp", p_n, engine="BP4")
vtx_u.write(t)
vtx_p.write(t)
for i in range(num_steps):
# Update current time step
t += dt
# Step 1: Tentative veolcity step
with b1.localForm() as loc_1:
loc_1.set(0)
assemble_vector(b1, L1)
apply_lifting(b1, [a1], [bcu])
b1.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
set_bc(b1, bcu)
solver1.solve(b1, u_.vector)
u_.x.scatter_forward()
# Step 2: Pressure corrrection step
with b2.localForm() as loc_2:
loc_2.set(0)
assemble_vector(b2, L2)
apply_lifting(b2, [a2], [bcp])
b2.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
set_bc(b2, bcp)
solver2.solve(b2, p_.vector)
p_.x.scatter_forward()
# Step 3: Velocity correction step
with b3.localForm() as loc_3:
loc_3.set(0)
assemble_vector(b3, L3)
b3.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
solver3.solve(b3, u_.vector)
u_.x.scatter_forward()
# Update variable with solution form this time step
u_n.x.array[:] = u_.x.array[:]
p_n.x.array[:] = p_.x.array[:]
# Write solutions to file
vtx_u.write(t)
vtx_p.write(t)
# Close xmdf file
vtx_u.close()
vtx_p.close()
b1.destroy()
b2.destroy()
b3.destroy()
solver1.destroy()
solver2.destroy()
solver3.destroy()