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.)