Open
Description
Describe new/missing feature
The main difference between PETSc NEST and PETSc block matrices seems to be the location of the ghosts (blocked matrix puts all ghosts at the end, while nest matrix puts them at the end of each block).
@schnellerhase , @michalhabera and I tried to use a direct solver with a nest matrix, and it seems to work.
Is there any performance study that indicates that our blocked matrices has a performance benefit compared to the nest matrices? If so, we should of course keep them.
from packaging.version import Version
from mpi4py import MPI
from petsc4py import PETSc
from dolfinx.cpp.la.petsc import scatter_local_vectors, get_local_vectors
import dolfinx.fem.petsc
import numpy as np
from scifem import create_real_functionspace, assemble_scalar
from scifem.petsc import apply_lifting_and_set_bc, zero_petsc_vector
import ufl
M = 40
mesh = dolfinx.mesh.create_unit_square(
MPI.COMM_WORLD, M, M, dolfinx.mesh.CellType.triangle, dtype=np.float64
)
V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
def u_exact(x):
return ufl.sin(2 * ufl.pi * x[0])
x = ufl.SpatialCoordinate(mesh)
n = ufl.FacetNormal(mesh)
g = ufl.dot(ufl.grad(u_exact(x)), n)
f = -ufl.div(ufl.grad(u_exact(x)))
h = assemble_scalar(u_exact(x) * ufl.dx)
# We then create the Lagrange multiplier space
R = create_real_functionspace(mesh)
# Next, we can create a mixed-function space for our problem
W = ufl.MixedFunctionSpace(V, R)
u, lmbda = ufl.TrialFunctions(W)
du, dl = ufl.TestFunctions(W)
# We can now define the variational problem
zero = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0))
a00 = ufl.inner(ufl.grad(u), ufl.grad(du)) * ufl.dx
a01 = ufl.inner(lmbda, du) * ufl.dx
a10 = ufl.inner(u, dl) * ufl.dx
L0 = ufl.inner(f, du) * ufl.dx + ufl.inner(g, du) * ufl.ds
L1 = ufl.inner(zero, dl) * ufl.dx
a = [[a00, a01], [a10, None]]
L = [L0, L1]
a_compiled = dolfinx.fem.form(a)
L_compiled = dolfinx.fem.form(L)
# Note that we have defined the variational form in a block form, and
# that we have not included $h$ in the variational form. We will enforce this
# once we have assembled the right hand side vector.
# We can now assemble the matrix and vector
A = dolfinx.fem.petsc.assemble_matrix(a_compiled, kind="nest", bcs=[])
A.assemble()
# On the main branch of DOLFINx, the `assemble_vector` function for blocked spaces has been rewritten to reflect how
# it works for standard assembly and `nest` assembly. This means that lifting is applied manually.
# In this case, with no Dirichlet BC, we could skip those steps.
# However, for clarity we include them here.
bcs = []
b = dolfinx.fem.petsc.assemble_vector(L_compiled, kind="nest")
apply_lifting_and_set_bc(b, a_compiled, bcs=bcs)
# Next, we modify the second part of the block to contain `h`
# We start by enforcing the multiplier constraint $h$ by modifying the right hand side vector.
# On the main branch, this is greatly simplified
uh = dolfinx.fem.Function(V, name="u")
po = {
"ksp_type":"preonly",
"pc_type":"lu",
"pc_factor_mat_solver_type": "mumps",
"ksp_error_if_not_converged": True,
}
problem = dolfinx.fem.petsc.LinearProblem(a, L, bcs=bcs, petsc_options=po,
kind="nest"
)
uh, rh = problem.solve()
print(problem.solver.getConvergedReason())
diff = uh - u_exact(x)
error = dolfinx.fem.form(ufl.inner(diff, diff) * ufl.dx)
print(f"L2 error: {np.sqrt(assemble_scalar(error)):.2e}")