Steady Stokes equation (3D) - dolfinx

General
4

It is quite straightforward to port demos from DOLFIN to DOLFINx.
For instance, take:
https://fenicsproject.org/olddocs/dolfin/latest/python/demos/stokes-iterative/demo_stokes-iterative.py.html
This can be converted as follows:

# Demo solving Stokes problem in 3D in DOLFINx using MINRES and AMG
# Copyright: 2022
# Author: Jørgen S. Dokken
# License: MIT

import numpy as np
import dolfinx
import ufl
from petsc4py import PETSc
from mpi4py import MPI

N = 16
mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N,
                                     cell_type=dolfinx.mesh.CellType.hexahedron)
P2 = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 2)
P1 = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
TH = ufl.MixedElement([P2, P1])
W = dolfinx.fem.FunctionSpace(mesh, TH)


def right(x, tol=1e-14): return x[0] > 1 - tol


def top_bottom(x, tol=1e-14):
    return np.logical_or(
        np.isclose(x[1], 0, atol=tol), np.isclose(x[1], 1, atol=tol))


V, V_to_W = W.sub(0).collapse()

# No-slip boundary condition for velocity
nonslip = dolfinx.fem.Function(V)
nonslip.x.set(0.0)
nonslip_facets = dolfinx.mesh.locate_entities_boundary(
    mesh, mesh.topology.dim-1, top_bottom)
nonslip_dofs = dolfinx.fem.locate_dofs_topological(
    (W.sub(0), V), mesh.topology.dim-1, nonslip_facets)
bc0 = dolfinx.fem.dirichletbc(nonslip, nonslip_dofs, W.sub(0))

# Inflow boundary condition for velocity


def inflow(x):
    zero_comp = np.zeros(x.shape[1], dtype=np.float64)
    return (-np.sin(np.pi*x[1]), zero_comp, zero_comp)


inflow_facets = dolfinx.mesh.locate_entities_boundary(mesh, mesh.topology.dim-1,
                                                      right)
inflow_dofs = dolfinx.fem.locate_dofs_topological((W.sub(0), V), mesh.topology.dim-1,
                                                  inflow_facets)
u_bc = dolfinx.fem.Function(V)
u_bc.interpolate(inflow)
bc1 = dolfinx.fem.dirichletbc(u_bc, inflow_dofs, W.sub(0))

bcs = [bc0, bc1]

(u, p) = ufl.TrialFunctions(W)
(v, q) = ufl.TestFunctions(W)
f = dolfinx.fem.Constant(mesh, (0., 0., 0.))

dx = ufl.Measure("dx", domain=mesh)
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx + \
    ufl.div(v)*p*dx + q*ufl.div(u)*dx
L = ufl.inner(f, v) * dx

a_f = dolfinx.fem.form(a)
L_f = dolfinx.fem.form(L)
A = dolfinx.fem.petsc.assemble_matrix(a_f, bcs=bcs)
A.assemble()
b = dolfinx.fem.petsc.assemble_vector(L_f)
dolfinx.fem.petsc.apply_lifting(b, [a_f], [bcs])
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
dolfinx.fem.petsc.set_bc(b, bcs)
b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

# Preconditioner matrix form
Pf = dolfinx.fem.form(ufl.inner(ufl.grad(u), ufl.grad(v)) * dx + p * q * dx)
P = dolfinx.fem.petsc.assemble_matrix(Pf, bcs=bcs)
P.assemble()

ksp = PETSc.KSP().create(mesh.comm)
ksp.setOperators(A, P)
ksp.setType(PETSc.KSP.Type.MINRES)
pc = ksp.getPC()
pc.setType("hypre")
pc.setHYPREType("boomeramg")
pc.setUp()

# To view Preconditioner info
# pc.view()


wh = dolfinx.fem.Function(W)
ksp.solve(b, wh.vector)

# To view solver info
# ksp.view()

uh = wh.sub(0).collapse()
p = wh.sub(1).collapse()

print(f"Converged Reason {ksp.getConvergedReason()}" +
      f"\nNumber of iterations {ksp.getIterationNumber()}")

with dolfinx.io.VTXWriter(mesh.comm, "u.bp", [uh]) as vtx:
    vtx.write(0.)

with dolfinx.io.VTXWriter(mesh.comm, "p.bp", [uh]) as vtx:
    vtx.write(0.)
3 Likes