Parallel Sparse PETSc Matrix Generation with dofs as rows and columns

4

The key here is that you are only access local (owned) indices of the vector (when you call get_local(), not the ghosted values. This means that you are ignoring several contributions to your matrix.

I do not know of an easy way of getting those ghost values in dolfin, however, here is a code that does the same thing for dolfinx.

from petsc4py import PETSc
from mpi4py import MPI
import dolfinx
import dolfinx.geometry
import basix
import numpy as np
import ufl
el = 30

mesh = dolfinx.UnitIntervalMesh(MPI.COMM_WORLD, el, dolfinx.cpp.mesh.GhostMode.shared_facet)
# mesh = dolfinx.UnitSquareMesh(MPI.COMM_WORLD, el, el)
#mesh = dolfinx.UnitCubeMesh(MPI.COMM_WORLD, el, el, el)


def fl_subdomain_func(x, eps=1e-16):
    x = x[0]
    x_f = 0.25
    a_f = 0.025
    return np.logical_and(x_f - a_f - eps <= x, x <= x_f + a_f + eps)

tdim = mesh.topology.dim
marked_cells = dolfinx.mesh.locate_entities(mesh, tdim, fl_subdomain_func)
fl = 1
subdomain = dolfinx.MeshTags(mesh, tdim, marked_cells, np.full(len(marked_cells), fl, dtype=np.int32))

degree = 1
CG = ufl.FiniteElement("CG", mesh.ufl_cell(), degree)
W = dolfinx.FunctionSpace(mesh, ufl.MixedElement([CG, CG]))

#This function returns a's nonzero dofs and their values as a list
v_1, v_2 = ufl.TestFunctions(W)
dx = ufl.Measure("dx", subdomain_data=subdomain)

V_fl = MPI.COMM_WORLD.allreduce(dolfinx.fem.assemble_scalar(dolfinx.Constant(mesh, 1)*dx(fl)), op=MPI.SUM)
b_1 = dolfinx.Function(W)
b_1.x.array[:] = 0
a_1 = dolfinx.fem.assemble_vector(b_1.vector, v_1/V_fl*dx(fl))
b_1.vector.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
b_2 = dolfinx.Function(W)
b_2.x.array[:] = 0
a_2 = dolfinx.fem.assemble_vector(b_2.vector, v_2/V_fl*dx(fl))
b_2.vector.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)

# NOTE: This includes ghosted dofs
aa1 = b_1.x.array
aa2 = b_2.x.array
indices1 = np.flatnonzero(aa1)
indices2 = np.flatnonzero(aa2)

aa1 = list(zip(indices1, aa1[indices1]))
aa2 = list(zip(indices2, aa2[indices2]))
a =  np.concatenate((aa1, aa2))

print(MPI.COMM_WORLD.rank, a, W.dofmap.index_map.size_local)

point = np.array([[0.2, 0, 0]])
v = np.array([[0, 0, 1]]).T
if tdim == 1:
    v = np.array([[1]])
elif tdim == 2:
    v = np.array([[1, 0]]).T

# Finds the basis function's derivative at point x
# and returns the relevant dof and derivative as a list
bb_tree = dolfinx.geometry.BoundingBoxTree(mesh, tdim)
cell_candidates = dolfinx.geometry.compute_collisions_point(bb_tree, point[0])
# Choose one of the cells that contains the point
cell = dolfinx.geometry.select_colliding_cells(mesh, cell_candidates, point[0], 1)

# Data required for pull back of coordinate
gdim = mesh.geometry.dim
num_local_cells = mesh.topology.index_map(tdim).size_local
num_dofs_x = mesh.geometry.dofmap.links(0).size  # NOTE: Assumes same cell geometry in whole mesh
t_imap = mesh.topology.index_map(tdim)
num_cells = t_imap.size_local + t_imap.num_ghosts
x = mesh.geometry.x
x_dofs = mesh.geometry.dofmap.array.reshape(num_cells, num_dofs_x)
cell_geometry = np.zeros((num_dofs_x, gdim), dtype=np.float64)
points_ref = np.zeros((1, tdim))

# Data required for evaluation of derivative
ct = dolfinx.cpp.mesh.to_string(mesh.topology.cell_type)
element = basix.create_element(basix.finite_element.string_to_family(
        "Lagrange", ct), basix.cell.string_to_type(ct), degree, basix.LagrangeVariant.equispaced)
dofmaps = [W.sub(i).dofmap for i  in range(2)]
coordinate_element = basix.create_element(basix.finite_element.string_to_family(
        "Lagrange", ct), basix.cell.string_to_type(ct), 1, basix.LagrangeVariant.equispaced)

point_ref = None
B = []
if len(cell) > 0:
    cell = cell[0]
    # Only add contribution if cell is owned
    if cell < num_local_cells:
        # Map point in cell back to reference element
        cell_geometry[:] = x[x_dofs[cell], :gdim]
        point_ref = mesh.geometry.cmap.pull_back(point[:,:gdim], cell_geometry)
        dphi = coordinate_element.tabulate(1, point_ref)[1:,0,:]
        J = np.dot(cell_geometry.T, dphi.T)
        Jinv = np.linalg.inv(J)  
        for j in range(2):
            cell_dofs = dofmaps[j].cell_dofs(cell)
            global_dofs = dofmaps[j].index_map.local_to_global(cell_dofs)
            # Compute gradient on physical element by multiplying by J^(-T)
            d_dx = (Jinv.T @ element.tabulate(1, point_ref)[1:, 0, :]).T
            d_dv = np.dot(d_dx, v)[:, 0]
            for i in range(len(cell_dofs)):
                B.append([global_dofs[i], d_dv[i]])
    else:
        print(MPI.COMM_WORLD.rank, "Ghost", cell)

print(np.asarray(B))

Regarding you question about mesh partitioning, there is functionality to control this in dolfinx, and @IgorBaratta can probably provide you with an example of how to control it.

3 Likes