Stokes equations with Taylor-Hood elements¶

This demo show how to solve the Stokes problem using Taylor-Hood elements with a range of different linear solvers.

Equation and problem definition¶

Strong formulation¶

\[ \begin{align}\begin{aligned}- \nabla \cdot (\nabla u + p I) &= f \quad {\rm in} \ \Omega,\\\nabla \cdot u &= 0 \quad {\rm in} \ \Omega.\end{aligned}\end{align} \]

Note

The sign of the pressure has been flipped from the classical definition. This is done in order to have a symmetric system of equations rather than a non-symmetric system of equations.

A typical set of boundary conditions on the boundary \(\partial \Omega = \Gamma_{D} \cup \Gamma_{N}\) can be:

\[ \begin{align}\begin{aligned}u &= u_0 \quad {\rm on} \ \Gamma_{D},\\\nabla u \cdot n + p n &= g \, \quad\;\; {\rm on} \ \Gamma_{N}.\end{aligned}\end{align} \]

Weak formulation¶

We formulate the Stokes equations mixed variational form; that is, a form where the two variables, the velocity and the pressure, are approximated. We have the problem: find \((u, p) \in W\) such that

\[a((u, p), (v, q)) = L((v, q))\]

for all \((v, q) \in W\), where

\[ \begin{align}\begin{aligned}a((u, p), (v, q)) &:= \int_{\Omega} \nabla u \cdot \nabla v - \nabla \cdot v \ p + \nabla \cdot u \ q \, {\rm d} x,\\L((v, q)) &:= \int_{\Omega} f \cdot v \, {\rm d} x + \int_{\partial \Omega_N} g \cdot v \, {\rm d} s.\end{aligned}\end{align} \]

The space \(W\) is mixed (product) function space \(W = V \times Q\), such that \(u \in V\) and \(q \in Q\).

Domain and boundary conditions¶

We shall the lid-driven cavity problem with the following definitions domain and boundary conditions:

\(\Omega = [0,1]\times[0,1]\) (a unit square)
\(\Gamma_D = \partial \Omega\)
\(u_0 = (1, 0)^\top\) at \(x_1 = 1\) and \(u_0 = (0, 0)^\top\) otherwise
\(f = (0, 0)^\top\)

Implementation¶

We first import the modules and function that the program uses:

import dolfinx
import numpy as np
import ufl
from dolfinx import DirichletBC, Function, FunctionSpace, RectangleMesh
from dolfinx.cpp.mesh import CellType
from dolfinx.fem import locate_dofs_geometrical, locate_dofs_topological
from dolfinx.io import XDMFFile
from dolfinx.mesh import locate_entities_boundary
from mpi4py import MPI
from petsc4py import PETSc
from ufl import div, dx, grad, inner

We create a Mesh and attach a coordinate map to the mesh:

# Create mesh
mesh = RectangleMesh(MPI.COMM_WORLD,
                     [np.array([0, 0, 0]), np.array([1, 1, 0])],
                     [32, 32],
                     CellType.triangle, dolfinx.cpp.mesh.GhostMode.none)


# Function to mark x = 0, x = 1 and y = 0
def noslip_boundary(x):
    return np.logical_or(np.logical_or(np.isclose(x[0], 0.0),
                                       np.isclose(x[0], 1.0)),
                         np.isclose(x[1], 0.0))


# Function to mark the lid (y = 1)
def lid(x):
    return np.isclose(x[1], 1.0)


# Lid velocity
def lid_velocity_expression(x):
    return np.stack((np.ones(x.shape[1]), np.zeros(x.shape[1])))

We define two FunctionSpace instances with different finite elements. P2 corresponds to piecewise quadratics for the velocity field and P1 to continuous piecewise linears for the pressure field:

P2 = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 2)
P1 = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
V, Q = FunctionSpace(mesh, P2), FunctionSpace(mesh, P1)

We can define boundary conditions:

# No-slip boundary condition for velocity field (`V`) on boundaries
# where x = 0, x = 1, and y = 0
noslip = Function(V)
with noslip.vector.localForm() as bc_local:
    bc_local.set(0.0)

facets = locate_entities_boundary(mesh, 1, noslip_boundary)
bc0 = DirichletBC(noslip, locate_dofs_topological(V, 1, facets))

# Driving velocity condition u = (1, 0) on top boundary (y = 1)
lid_velocity = Function(V)
lid_velocity.interpolate(lid_velocity_expression)
facets = locate_entities_boundary(mesh, 1, lid)
bc1 = DirichletBC(lid_velocity, locate_dofs_topological(V, 1, facets))

# Collect Dirichlet boundary conditions
bcs = [bc0, bc1]

We now define the bilinear and linear forms corresponding to the weak mixed formulation of the Stokes equations in a blocked structure:

# Define variational problem
(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)
(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)
f = dolfinx.Constant(mesh, (0, 0))

a = [[inner(grad(u), grad(v)) * dx, inner(p, div(v)) * dx],
     [inner(div(u), q) * dx, None]]

L = [inner(f, v) * dx,
     inner(dolfinx.Constant(mesh, 0), q) * dx]

We will use a block-diagonal preconditioner to solve this problem:

a_p11 = inner(p, q) * dx
a_p = [[a[0][0], None],
       [None, a_p11]]

Nested matrix solver¶

We now assemble the bilinear form into a nested matrix A, and call the assemble() method to communicate shared entries in parallel. Rows and columns in A that correspond to degrees-of-freedom with Dirichlet boundary conditions are zeroed and a value of 1 is set on the diagonal.

A = dolfinx.fem.assemble_matrix_nest(a, bcs)
A.assemble()

We create a nested matrix P to use as the preconditioner. The top-left block of P is shared with the top-left block of A. The bottom-right diagonal entry is assembled from the form a_p11:

P11 = dolfinx.fem.assemble_matrix(a_p11, [])
P = PETSc.Mat().createNest([[A.getNestSubMatrix(0, 0), None], [None, P11]])
P.assemble()

Next, the right-hand side vector is assembled and then modified to account for non-homogeneous Dirichlet boundary conditions:

b = dolfinx.fem.assemble.assemble_vector_nest(L)

# Modify ('lift') the RHS for Dirichlet boundary conditions
dolfinx.fem.assemble.apply_lifting_nest(b, a, bcs)

# Sum contributions from ghost entries on the owner
for b_sub in b.getNestSubVecs():
    b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)

# Set Dirichlet boundary condition values in the RHS
bcs0 = dolfinx.cpp.fem.bcs_rows(dolfinx.fem.assemble._create_cpp_form(L), bcs)
dolfinx.fem.assemble.set_bc_nest(b, bcs0)

Ths pressure field for this problem is determined only up to a constant. We can supply the vector that spans the nullspace and any component of the solution in this direction will be eliminated during the iterative linear solution process.

# Create nullspace vector
null_vec = dolfinx.fem.create_vector_nest(L)

# Set velocity part to zero and the pressure part to a non-zero constant
null_vecs = null_vec.getNestSubVecs()
null_vecs[0].set(0.0), null_vecs[1].set(1.0)

# Normalize the vector, create a nullspace object, and attach it to the
# matrix
null_vec.normalize()
nsp = PETSc.NullSpace().create(vectors=[null_vec])
assert nsp.test(A)
A.setNullSpace(nsp)

Now we create a Krylov Subspace Solver ksp. We configure it to use the MINRES method, and a block-diagonal preconditioner using PETSc’s additive fieldsplit type preconditioner:

ksp = PETSc.KSP().create(mesh.mpi_comm())
ksp.setOperators(A, P)
ksp.setType("minres")
ksp.setTolerances(rtol=1e-8)
ksp.getPC().setType("fieldsplit")
ksp.getPC().setFieldSplitType(PETSc.PC.CompositeType.ADDITIVE)

# Define the matrix blocks in the preconditioner with the velocity and
# pressure matrix index sets
nested_IS = P.getNestISs()
ksp.getPC().setFieldSplitIS(
    ("u", nested_IS[0][0]),
    ("p", nested_IS[0][1]))

# Set the preconditioners for each block
ksp_u, ksp_p = ksp.getPC().getFieldSplitSubKSP()
ksp_u.setType("preonly")
ksp_u.getPC().setType("gamg")
ksp_p.setType("preonly")
ksp_p.getPC().setType("jacobi")

# Monitor the convergence of the KSP
ksp.setFromOptions()

To compute the solution, we create finite element Function for the velocity (on the space V) and for the pressure (on the space Q). The vectors for u and p are combined to form a nested vector and the system is solved:

u, p = Function(V), Function(Q)
x = PETSc.Vec().createNest([u.vector, p.vector])
ksp.solve(b, x)

Norms of the solution vectors are computed:

norm_u_0 = u.vector.norm()
norm_p_0 = p.vector.norm()
if MPI.COMM_WORLD.rank == 0:
    print("(A) Norm of velocity coefficient vector (nested, iterative): {}".format(norm_u_0))
    print("(A) Norm of pressure coefficient vector (nested, iterative): {}".format(norm_p_0))

The solution fields can be saved to file in XDMF format for visualization, e.g. with ParView. Before writing to file, ghost values are updated.

with XDMFFile(MPI.COMM_WORLD, "velocity.xdmf", "w") as ufile_xdmf:
    u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    ufile_xdmf.write_mesh(mesh)
    ufile_xdmf.write_function(u)

with XDMFFile(MPI.COMM_WORLD, "pressure.xdmf", "w") as pfile_xdmf:
    p.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    pfile_xdmf.write_mesh(mesh)
    pfile_xdmf.write_function(p)

Monolithic block iterative solver¶

Next, we solve same problem, but now with monolithic (non-nested) matrices and iterative solvers.

A = dolfinx.fem.assemble_matrix_block(a, bcs)
A.assemble()
P = dolfinx.fem.assemble_matrix_block(a_p, bcs)
P.assemble()
b = dolfinx.fem.assemble.assemble_vector_block(L, a, bcs)

# Set near null space for pressure
null_vec = A.createVecLeft()
offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
null_vec.array[offset:] = 1.0
null_vec.normalize()
nsp = PETSc.NullSpace().create(vectors=[null_vec])
assert nsp.test(A)
A.setNullSpace(nsp)

# Build IndexSets for each field (global dof indices for each field)
V_map = V.dofmap.index_map
Q_map = Q.dofmap.index_map
offset_u = V_map.local_range[0] * V.dofmap.index_map_bs + Q_map.local_range[0]
offset_p = offset_u + V_map.size_local * V.dofmap.index_map_bs
is_u = PETSc.IS().createStride(V_map.size_local * V.dofmap.index_map_bs, offset_u, 1, comm=PETSc.COMM_SELF)
is_p = PETSc.IS().createStride(Q_map.size_local, offset_p, 1, comm=PETSc.COMM_SELF)

# Create Krylov solver
ksp = PETSc.KSP().create(mesh.mpi_comm())
ksp.setOperators(A, P)
ksp.setTolerances(rtol=1e-8)
ksp.setType("minres")
ksp.getPC().setType("fieldsplit")
ksp.getPC().setFieldSplitType(PETSc.PC.CompositeType.ADDITIVE)
ksp.getPC().setFieldSplitIS(
    ("u", is_u),
    ("p", is_p))

# Configure velocity and pressure sub KSPs
ksp_u, ksp_p = ksp.getPC().getFieldSplitSubKSP()
ksp_u.setType("preonly")
ksp_u.getPC().setType("gamg")
ksp_p.setType("preonly")
ksp_p.getPC().setType("jacobi")

# Monitor the convergence of the KSP
opts = PETSc.Options()
opts["ksp_monitor"] = None
opts["ksp_view"] = None
ksp.setFromOptions()

We also need to create a block vector,``x``, to store the (full) solution, which we initialize using the block RHS form L.

# Compute solution
x = A.createVecRight()
ksp.solve(b, x)

# Create Functions and scatter x solution
u, p = Function(V), Function(Q)
offset = V_map.size_local * V.dofmap.index_map_bs
u.vector.array[:] = x.array_r[:offset]
p.vector.array[:] = x.array_r[offset:]

We can calculate the \(L^2\) norms of u and p as follows:

norm_u_1 = u.vector.norm()
norm_p_1 = p.vector.norm()
if MPI.COMM_WORLD.rank == 0:
    print("(B) Norm of velocity coefficient vector (blocked, iterative): {}".format(norm_u_1))
    print("(B) Norm of pressure coefficient vector (blocked, interative): {}".format(norm_p_1))
assert np.isclose(norm_u_1, norm_u_0)
assert np.isclose(norm_p_1, norm_p_0)

Monolithic block direct solver¶

Solve same problem, but now with monolithic matrices and a direct solver

# Create LU solver
ksp = PETSc.KSP().create(mesh.mpi_comm())
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("superlu_dist")

We also need to create a block vector,``x``, to store the (full) solution, which we initialize using the block RHS form L.

# Compute solution
x = A.createVecLeft()
ksp.solve(b, x)

# Create Functions and scatter x solution
u, p = Function(V), Function(Q)
offset = V_map.size_local * V.dofmap.index_map_bs
u.vector.array[:] = x.array_r[:offset]
p.vector.array[:] = x.array_r[offset:]

We can calculate the \(L^2\) norms of u and p as follows:

norm_u_2 = u.vector.norm()
norm_p_2 = p.vector.norm()
if MPI.COMM_WORLD.rank == 0:
    print("(C) Norm of velocity coefficient vector (blocked, direct): {}".format(norm_u_2))
    print("(C) Norm of pressure coefficient vector (blocked, direct): {}".format(norm_p_2))
assert np.isclose(norm_u_2, norm_u_0)
assert np.isclose(norm_p_2, norm_p_0)

Non-blocked direct solver¶

Again, solve the same problem but this time with a non-blocked direct solver approach

# Create the function space
TH = P2 * P1
W = FunctionSpace(mesh, TH)
W0 = W.sub(0).collapse()

# No slip boundary condition
noslip = Function(V)
facets = locate_entities_boundary(mesh, 1, noslip_boundary)
dofs = locate_dofs_topological((W.sub(0), V), 1, facets)
bc0 = DirichletBC(noslip, dofs, W.sub(0))


# Driving velocity condition u = (1, 0) on top boundary (y = 1)
lid_velocity = Function(W0)
lid_velocity.interpolate(lid_velocity_expression)
facets = locate_entities_boundary(mesh, 1, lid)
dofs = locate_dofs_topological((W.sub(0), V), 1, facets)
bc1 = DirichletBC(lid_velocity, dofs, W.sub(0))


# Since for this problem the pressure is only determined up to a
# constant, we pin the pressure at the point (0, 0)
zero = Function(Q)
with zero.vector.localForm() as zero_local:
    zero_local.set(0.0)
dofs = locate_dofs_geometrical((W.sub(1), Q),
                               lambda x: np.isclose(x.T, [0, 0, 0]).all(axis=1))
bc2 = DirichletBC(zero, dofs, W.sub(1))

# Collect Dirichlet boundary conditions
bcs = [bc0, bc1, bc2]

# Define variational problem
(u, p) = ufl.TrialFunctions(W)
(v, q) = ufl.TestFunctions(W)
f = Function(W0)
zero = dolfinx.Constant(mesh, 0.0)
a = (inner(grad(u), grad(v)) + inner(p, div(v)) + inner(div(u), q)) * dx
L = inner(f, v) * dx

# Assemble LHS matrix and RHS vector
A = dolfinx.fem.assemble_matrix(a, bcs)
A.assemble()
b = dolfinx.fem.assemble.assemble_vector(L)

dolfinx.fem.assemble.apply_lifting(b, [a], [bcs])
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)

# Set Dirichlet boundary condition values in the RHS
dolfinx.fem.assemble.set_bc(b, bcs)

# Create and configure solver
ksp = PETSc.KSP().create(mesh.mpi_comm())
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("superlu_dist")

# Compute the solution
U = Function(W)
ksp.solve(b, U.vector)

# Split the mixed solution and collapse
u = U.sub(0).collapse()
p = U.sub(1).collapse()

# Compute norms
norm_u_3 = u.vector.norm()
norm_p_3 = p.vector.norm()
if MPI.COMM_WORLD.rank == 0:
    print("(D) Norm of velocity coefficient vector (monolithic, direct): {}".format(norm_u_3))
    print("(D) Norm of pressure coefficient vector (monolithic, direct): {}".format(norm_p_3))
assert np.isclose(norm_u_3, norm_u_0)

# Write the solution to file
with XDMFFile(MPI.COMM_WORLD, "new_velocity.xdmf", "w") as ufile_xdmf:
    u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    ufile_xdmf.write_mesh(mesh)
    ufile_xdmf.write_function(u)

with XDMFFile(MPI.COMM_WORLD, "my.xdmf", "w") as pfile_xdmf:
    p.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
    pfile_xdmf.write_mesh(mesh)
    pfile_xdmf.write_function(p)

Stokes equations with Taylor-Hood elements¶

Equation and problem definition¶

Strong formulation¶

Weak formulation¶

Domain and boundary conditions¶

Implementation¶

Nested matrix solver¶

Monolithic block iterative solver¶

Monolithic block direct solver¶

Non-blocked direct solver¶

Table of Contents

Previous topic

Next topic

This Page