Dolfinx version of "Plot unit normal vectors of facets (interior boundary) in a 2D mesh"

This question has already been answered at:

Here is an updated version to DOLFINx v0.9, which also includes the project step.

import time

import dolfinx.fem.petsc
import numpy as np
import ufl
from dolfinx import cpp, fem
from mpi4py import MPI
from petsc4py import PETSc

N = 10
degree = 1
mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N)
V = dolfinx.fem.functionspace(mesh, ("CG", degree,(mesh.geometry.dim,)))
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
print(
    f"N={N}, Num dofs global {V.dofmap.index_map.size_global*V.dofmap.index_map_bs}")
# form is defined on the boundary => matrix wil have many zero diagonal entries
a_bnd = ufl.inner(u, v) * ufl.ds
form = fem.form(a_bnd)

# Find all dofs associated to the facets we are integrating over
mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
boundary_blocks = dolfinx.fem.locate_dofs_topological(
    V, V.mesh.topology.dim - 1, exterior_facets)

# Set diagonal for everything not associated with a facet in sparsity pattern
imap = V.dofmap.index_map
all_blocks = np.arange(imap.size_local+imap.num_ghosts, dtype=np.int32)
zero_blocks = np.delete(all_blocks, boundary_blocks)
start = time.perf_counter()
sp = fem.create_sparsity_pattern(form)
sp.insert_diagonal(zero_blocks)
sp.finalize()
end = time.perf_counter()
print(f"Create sparsity pattern {end-start}")

start = time.perf_counter()
A = cpp.la.petsc.create_matrix(mesh.comm, sp)
A.setOption(A.Option.NEW_NONZERO_ALLOCATION_ERR, 1)
fem.petsc.assemble_matrix(A, form)
A.assemble(A.AssemblyType.FLUSH)
end = time.perf_counter()
print(f"Assemble boundary {end-start}")


def bc_diag(A):
    bc = dolfinx.fem.dirichletbc(
        dolfinx.fem.Constant(mesh, np.zeros(u.ufl_shape, dtype=PETSc.ScalarType)), zero_blocks, V)
    start = time.perf_counter()
    dolfinx.cpp.fem.petsc.insert_diagonal(A, V._cpp_object, [bc._cpp_object], 1.)
    end = time.perf_counter()
    print(f"BC insert diag {end-start}s")
    return bc

bc = bc_diag(A)
start = time.perf_counter()
A.assemble()
end = time.perf_counter()
print(f"Matrix communication {end-start}")
#
print(f"min(diag) = {A.getDiagonal().array.min()}")
print(f"number of diagonal zeros = {sum(A.getDiagonal().array == 0)}")
print(f"norm(A) = {A.norm()}")



b = dolfinx.fem.form(ufl.inner(ufl.FacetNormal(mesh), v) * ufl.ds)
b_vec = dolfinx.fem.petsc.assemble_vector(b)
dolfinx.fem.petsc.apply_lifting(b_vec, [form], [[bc]])
b_vec.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
dolfinx.fem.petsc.set_bc(b_vec, [bc])

ksp = PETSc.KSP().create(mesh.comm)
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.setErrorIfNotConverged(True)
uh = dolfinx.fem.Function(V)

start = time.perf_counter()
ksp.solve(b_vec, uh.x.petsc_vec)
ksp.destroy()
b_vec.destroy()
end = time.perf_counter()
uh.x.scatter_forward()
print(f"CG solve {end-start}s")

with dolfinx.io.VTXWriter(mesh.comm, "normal.bp", [uh]) as writer:
    writer.write(0.0)