Modify matrix diagonal -- dolfinx version for A.ident_zeros()

General
3

I would suggest the following:

import time

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

N = 75
degree = 1
mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N)
V = dolfinx.fem.FunctionSpace(mesh, ("CG", degree))
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
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, 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.assemble()
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, PETSc.ScalarType(1.)), zero_blocks, V)
    start = time.perf_counter()
    dolfinx.cpp.fem.petsc.insert_diagonal(A, V._cpp_object, [bc], 1.)
    end = time.perf_counter()
    print(f"BC insert diag {end-start}s")


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()}")

This is leveraging that you actually know what degrees of freedom that will give contributions to your assembly.
With output:

N=75, Num dofs global 438976
Create sparsity pattern 0.28873315098462626
Assemble boundary 0.22457787103485316
BC insert diag 0.08238964097108692s
Matrix communication 0.030512886005453765
min(diag) = 5.925925925925912e-05
number of diagonal zeros = 0
norm(A) = 636.5720700068672
3 Likes