Mixed formulation of the Poisson equation with a block-preconditioner/solver # noqa
This demo illustrates how to solve the Poisson equation using a mixed (two-field) formulation and a block-preconditioned iterative solver. In particular, it illustrates how to
Use mixed and non-continuous finite element spaces.
Set essential boundary conditions for subspaces and
spaces.Construct a blocked linear system.
Construct a block-preconditioned iterative linear solver using PETSc/petsc4y.
Construct a Hypre Auxiliary Maxwell Space (AMS) preconditioner for
problems in two-dimensions.
Download sources
Equation and problem definition
An alternative formulation of Poisson equation can be formulated by
introducing an additional (vector) variable, namely the (negative)
flux:
with boundary conditions
where
where the variational forms
and
To discretize the above formulation, two discrete function spaces
To solve the linear system for the mixed problem, we will use am iterative method with a block-diagonal preconditioner that is based on the Riesz map, see for example this paper.
Implementation
Import the required modules:
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
import dolfinx
import ufl
from basix.ufl import element
from dolfinx import fem, mesh
from dolfinx.fem.petsc import discrete_gradient, interpolation_matrix
from dolfinx.mesh import CellType, create_unit_square
# Solution scalar (e.g., float32, complex128) and geometry (float32/64)
# types
dtype = dolfinx.default_scalar_type
xdtype = dolfinx.default_real_type
Create a two-dimensional mesh. The iterative solver constructed later requires special construction that is specific to two dimensions. Application in three-dimensions would require a number of changes to the linear solver.
msh = create_unit_square(MPI.COMM_WORLD, 96, 96, CellType.triangle, dtype=xdtype)
Here we construct compatible function spaces for the mixed Poisson
problem. The V
Raviart-Thomas (W
space is a
space of discontinuous Lagrange function of degree k
.
Note
The
In the lowest-order case
k = 1
V = fem.functionspace(msh, element("RT", msh.basix_cell(), k, dtype=xdtype))
W = fem.functionspace(msh, element("Discontinuous Lagrange", msh.basix_cell(), k - 1, dtype=xdtype))
Q = ufl.MixedFunctionSpace(V, W)
Trial functions for
sigma, u = ufl.TrialFunctions(Q)
tau, v = ufl.TestFunctions(Q)
# The source function is set to be $f = 10\exp(-((x_{0} - 0.5)^2 +
# (x_{1} - 0.5)^2) / 0.02)$:
x = ufl.SpatialCoordinate(msh)
f = 10 * ufl.exp(-((x[0] - 0.5) * (x[0] - 0.5) + (x[1] - 0.5) * (x[1] - 0.5)) / 0.02)
We now declare the blocked bilinear and linear forms. We use
ufl.extract_blocks
to extract the block structure of the bilinear
and linear form. For the first block of the right-hand side, we provide
a form that efficiently is 0. We do this to preserve knowledge of the
test space in the block. Note that the defined L
corresponds to
dx = ufl.Measure("dx", msh)
a = ufl.extract_blocks(
ufl.inner(sigma, tau) * dx + ufl.inner(u, ufl.div(tau)) * dx + ufl.inner(ufl.div(sigma), v) * dx
)
L = [ufl.ZeroBaseForm((tau,)), -ufl.inner(f, v) * dx]
In preparation for Dirichlet boundary conditions, we use the function
locate_entities_boundary
to locate mesh entities (facets) with which
degree-of-freedoms to be constrained are associated with, and then use
locate_dofs_topological
to get the degree-of-freedom indices. Below
we identify the degree-of-freedom in V
on the (i) top (
fdim = msh.topology.dim - 1
facets_top = mesh.locate_entities_boundary(msh, fdim, lambda x: np.isclose(x[1], 1.0))
facets_bottom = mesh.locate_entities_boundary(msh, fdim, lambda x: np.isclose(x[1], 0.0))
dofs_top = fem.locate_dofs_topological(V, fdim, facets_top)
dofs_bottom = fem.locate_dofs_topological(V, fdim, facets_bottom)
Now, we create Dirichlet boundary objects for the condition
cells_top_ = mesh.compute_incident_entities(msh.topology, facets_top, fdim, fdim + 1)
cells_bottom = mesh.compute_incident_entities(msh.topology, facets_bottom, fdim, fdim + 1)
g = fem.Function(V, dtype=dtype)
g.interpolate(lambda x: np.vstack((np.zeros_like(x[0]), np.sin(5 * x[0]))), cells0=cells_top_)
g.interpolate(lambda x: np.vstack((np.zeros_like(x[0]), -np.sin(5 * x[0]))), cells0=cells_bottom)
bcs = [fem.dirichletbc(g, dofs_top), fem.dirichletbc(g, dofs_bottom)]
Rather than solving the linear system
For this problem, we introduce the preconditioner
$P
:
a_p = ufl.extract_blocks(
ufl.inner(sigma, tau) * dx + ufl.inner(ufl.div(sigma), ufl.div(tau)) * dx + ufl.inner(u, v) * dx
)
We create finite element functions that will hold the
sigma, u = fem.Function(V, name="sigma", dtype=dtype), fem.Function(W, name="u", dtype=dtype)
We now create a linear problem solver for the mixed problem.
As we will use different preconditions for the individual blocks of
the saddle point problem, we specify the matrix kind to be “nest”,
so that we can use # fieldsplit
(block) type and set the ‘splits’ between the
problem = fem.petsc.LinearProblem(
a,
L,
u=[sigma, u],
P=a_p,
kind="nest",
bcs=bcs,
petsc_options_prefix="demo_mixed_poisson_",
petsc_options={
"ksp_type": "gmres",
"pc_type": "fieldsplit",
"pc_fieldsplit_type": "additive",
"ksp_rtol": 1e-8,
"ksp_gmres_restart": 100,
"ksp_view": "",
},
)
ksp = problem.solver
ksp.setMonitor(
lambda _, its, rnorm: PETSc.Sys.Print(f"Iteration: {its:>4d}, residual: {rnorm:.3e}")
)
For the
If PETSc has been configured with Hypre, we use the Hypre Auxiliary Maxwell Space
(AMS) algebraic multigrid preconditioner. We can use
AMS for this
ksp_sigma, ksp_u = ksp.getPC().getFieldSplitSubKSP()
pc_sigma = ksp_sigma.getPC()
if PETSc.Sys().hasExternalPackage("hypre") and not np.issubdtype(dtype, np.complexfloating):
pc_sigma.setType("hypre")
pc_sigma.setHYPREType("ams")
opts = PETSc.Options()
opts[f"{ksp_sigma.prefix}pc_hypre_ams_cycle_type"] = 7
opts[f"{ksp_sigma.prefix}pc_hypre_ams_relax_times"] = 2
# Construct and set the 'discrete gradient' operator, which maps
# grad H1 -> H(curl), i.e. the gradient of a scalar Lagrange space
# to a H(curl) space
V_H1 = fem.functionspace(msh, element("Lagrange", msh.basix_cell(), k, dtype=xdtype))
V_curl = fem.functionspace(msh, element("N1curl", msh.basix_cell(), k, dtype=xdtype))
G = discrete_gradient(V_H1, V_curl)
G.assemble()
pc_sigma.setHYPREDiscreteGradient(G)
assert k > 0, "Element degree must be at least 1."
if k == 1:
# For the lowest order base (k=1), we can supply interpolation
# of the '1' vectors in the space V. Hypre can then construct
# the required operators from G and the '1' vectors.
cvec0, cvec1 = fem.Function(V), fem.Function(V)
cvec0.interpolate(lambda x: np.vstack((np.ones_like(x[0]), np.zeros_like(x[1]))))
cvec1.interpolate(lambda x: np.vstack((np.zeros_like(x[0]), np.ones_like(x[1]))))
pc_sigma.setHYPRESetEdgeConstantVectors(cvec0.x.petsc_vec, cvec1.x.petsc_vec, None)
else:
# For high-order spaces, we must provide the (H1)^d -> H(div)
# interpolation operator/matrix
V_H1d = fem.functionspace(msh, ("Lagrange", k, (msh.geometry.dim,)))
Pi = interpolation_matrix(V_H1d, V) # (H1)^d -> H(div)
Pi.assemble()
pc_sigma.setHYPRESetInterpolations(msh.geometry.dim, None, None, Pi, None)
# High-order elements generally converge less well than the
# lowest-order case with algebraic multigrid, so we perform
# extra work at the multigrid stage
opts[f"{ksp_sigma.prefix}pc_hypre_ams_tol"] = 1e-12
opts[f"{ksp_sigma.prefix}pc_hypre_ams_max_iter"] = 3
ksp_sigma.setFromOptions()
else:
# If Hypre is not available, use LU factorisation on the $P_{00}$
# block
pc_sigma.setType("lu")
use_superlu = PETSc.IntType == np.int64
if PETSc.Sys().hasExternalPackage("mumps") and not use_superlu:
pc_sigma.setFactorSolverType("mumps")
elif PETSc.Sys().hasExternalPackage("superlu_dist"):
pc_sigma.setFactorSolverType("superlu_dist")
Once we have set the preconditioners for the two blocks, we can
solve the linear system. The LinearProblem
class will
automatically assemble the linear system, apply the boundary
conditions, call the Krylov solver and update the solution
vectors u
and sigma
.
problem.solve()
converged_reason = problem.solver.getConvergedReason()
assert converged_reason > 0, f"Krylov solver has not converged, reason: {converged_reason}."
We save the solution u
in VTX format:
try:
from dolfinx.io import VTXWriter
u.name = "u"
with VTXWriter(msh.comm, "output_mixed_poisson.bp", u, "bp4") as f:
f.write(0.0)
except ImportError:
print("ADIOS2 required for VTX output.")