Investigate issue with following example

[jorgensd]

import dolfinx
import dolfinx_mpc
from dolfinx.fem import FunctionSpace, Function, Constant, assemble_scalar, form
from dolfinx.mesh import locate_entities_boundary, meshtags

from ufl import grad, div, inner, Measure, TrialFunctions, TestFunctions, FiniteElement, VectorElement
import ufl
from petsc4py import PETSc
import numpy as np
from mpi4py import MPI
# MPI initialization
comm = MPI.COMM_WORLD
rank = MPI.COMM_WORLD.rank

#############################################
#              Mesh generation              #
#############################################
domain = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 50, 50, dolfinx.mesh.CellType.quadrilateral)
gdim = 2
fdim = 1
##  cells for different subdomains
def cells_1(x):
    return np.logical_or(np.logical_or(x[0] > 0.6, x[0] < 0.4), np.logical_or(x[1] > 0.6, x[1] < 0.4))
def cells_2(x):
    return np.logical_and(np.logical_and(x[0] <= 0.6, x[1] <= 0.6), np.logical_and(x[0] >= 0.4, x[1] >= 0.4))

# locate entities
cells_1_dom = dolfinx.mesh.locate_entities(domain, domain.topology.dim, cells_1)
cells_2_dom = dolfinx.mesh.locate_entities(domain, domain.topology.dim, cells_2)

## create cell tags
marked_cells = np.hstack([cells_1_dom, cells_2_dom])
marked_values = np.hstack([np.full_like(cells_1_dom, 1), np.full_like(cells_2_dom, 2)])
sorted_cells = np.argsort(marked_cells)
ct = dolfinx.mesh.meshtags(domain, gdim, marked_cells[sorted_cells], marked_values[sorted_cells])

## create facet tag at interface
def create_interface(facet_marker: np.typing.NDArray[np.int32],
                     tag: int,
                     fa: np.typing.NDArray[np.int32],
                     fo: np.typing.NDArray[np.int32],
                     cell_marker: np.typing.NDArray[np.int32]):
    for f in range(len(facet_marker)):
        cells = fa[fo[f]:fo[f + 1]]
        c_val = cell_marker[cells]
        if len(np.unique(c_val)) == 2:
            facet_marker[f] = tag


# Cold start of function
create_interface(np.empty(0, dtype=np.int32), 0, np.empty(
    0, dtype=np.int32), np.empty(0, dtype=np.int32), np.empty(0, dtype=np.int32))

cell_tag_1 = np.unique(ct.values)[0]
cell_tag_2 = np.unique(ct.values)[1]

domain.topology.create_connectivity(domain.topology.dim - 1, domain.topology.dim)
f_to_c = domain.topology.connectivity(domain.topology.dim - 1, domain.topology.dim)
fa = f_to_c.array
fo = f_to_c.offsets

facet_map = domain.topology.index_map(domain.topology.dim - 1)
num_facets = facet_map.size_local + facet_map.num_ghosts
facet_marker = np.zeros(num_facets, dtype=np.int32)
cell_map = domain.topology.index_map(domain.topology.dim)
num_cells = cell_map.size_local + cell_map.num_ghosts
cell_marker = np.full(num_cells, cell_tag_2, dtype=np.int32)
cell_marker[ct.indices] = ct.values

interface_marker = 9
create_interface(facet_marker, interface_marker, fa, fo, cell_marker)
ft = meshtags(domain, domain.topology.dim - 1,
                     np.arange(num_facets), facet_marker)

dx = Measure('dx', domain=domain, subdomain_data=ct)
#############################################
#             Mesh generation finished      #
#############################################



# elements and funcionspace
###########################

# Taylor-Hood elements and Mixed Functionspace
P1 = FiniteElement("CG", domain.ufl_cell(), 1)
P2 = VectorElement("CG", domain.ufl_cell(), 2)
TH = ufl.MixedElement([P2, P1])
V = FunctionSpace(domain, TH)

# define function, trial function and test function
(ve, pr) = TrialFunctions(V)
(dve, dpr) = TestFunctions(V)

# define material parameter
S = FunctionSpace(domain, ("DG", 0))
vis = Function(S)
fluid_visc = ct.find(1)
vis.x.array[fluid_visc] = np.full_like(fluid_visc, 1, dtype=PETSc.ScalarType)
inclusion_visc = ct.find(2)
vis.x.array[inclusion_visc]  = np.full_like(inclusion_visc, 1e+5, dtype=PETSc.ScalarType)
vis.x.scatter_forward()

# boundary conditions
#####################
## set all fluid velocities at the interface to zero (no slip boundary condition)
# create collapsed functionspace
V_coll0, _ = V.sub(0).collapse()
v_bc = Function(V_coll0)
v_bc.x.array[:] = 0.0
v_bc.x.scatter_forward()
# dofs of interface and dirichlet bcs
dofs_v_interf = dolfinx.fem.locate_dofs_topological((V.sub(0), V_coll0)  , ft.dim, ft.find(9))  # interface tag: 9
bcs_v_interf = dolfinx.fem.dirichletbc(v_bc, dofs_v_interf, V.sub(0))

## set pressures at the corners to zero
def boundary_node(x):
    return np.logical_or(
        np.logical_and(
            np.isclose(x[0], 0.),
            np.logical_or(
                np.isclose(x[1], .0),
                np.isclose(x[1], 1.0),
            )
        ),
        np.logical_and(
            np.isclose(x[0], 1.0),
            np.logical_or(
                np.isclose(x[1], 0.),
                np.isclose(x[1], 1.0),
            )
        ),
    )

boundary_facets = locate_entities_boundary(domain, 0, boundary_node) # it is zero dimension for a point
dofs_corners = dolfinx.fem.locate_dofs_topological(V.sub(1), 0, boundary_facets)  # it is zero dimension for a point
bc_corners = dolfinx.fem.dirichletbc(PETSc.ScalarType(0.0), dofs_corners, V.sub(1))

bcs = [bcs_v_interf, bc_corners]

## periodic boundary conditions
pbc_directions = []
pbc_slave_tags = []
pbc_is_slave = []
pbc_is_master = []
pbc_meshtags = []
pbc_slave_to_master_maps = []

mpc = dolfinx_mpc.MultiPointConstraint(V)

def generate_pbc_slave_to_master_map(i):
    def pbc_slave_to_master_map(x):
        out_x = x.copy()
        out_x[i] = x[i] - 1.0
        return out_x
    return pbc_slave_to_master_map

def generate_pbc_is_slave(i):
    return lambda x: np.isclose(x[i], 1.0)

def generate_pbc_is_master(i):
    return lambda x: np.isclose(x[i], 0.0)

for i in range(gdim):

    pbc_directions.append(i)
    pbc_slave_tags.append(i + 2)
    pbc_is_slave.append(generate_pbc_is_slave(i))
    pbc_is_master.append(generate_pbc_is_master(i))
    pbc_slave_to_master_maps.append(generate_pbc_slave_to_master_map(i))

    facets = locate_entities_boundary(domain, fdim, pbc_is_slave[-1])
    arg_sort = np.argsort(facets)
    pbc_meshtags.append(meshtags(domain,
                                 fdim,
                                 facets[arg_sort],
                                 np.full(len(facets), pbc_slave_tags[-1], dtype=np.int32)))

N_pbc = len(pbc_directions)
for i in range(N_pbc):
    if N_pbc > 1:
        def pbc_slave_to_master_map(x):
            out_x = pbc_slave_to_master_maps[i](x)
            idx = pbc_is_slave[(i + 1) % N_pbc](x)
            out_x[pbc_directions[i]][idx] = np.nan
            return out_x
    else:
        pbc_slave_to_master_map = pbc_slave_to_master_maps[i]

    for ij in range(gdim):
        mpc.create_periodic_constraint_topological(V.sub(0).sub(ij), pbc_meshtags[i],
                                                   pbc_slave_tags[i],
                                                   pbc_slave_to_master_map,
                                                   bcs)

    mpc.create_periodic_constraint_topological(V.sub(1), pbc_meshtags[i],
                                                   pbc_slave_tags[i],
                                                   pbc_slave_to_master_map,
                                                   bcs)

if len(pbc_directions) > 1:
    def pbc_slave_to_master_map(x):
        out_x = x.copy()
        out_x[0] = x[0] - domain.geometry.x[:, 0].max()
        out_x[1] = x[1] - domain.geometry.x[:, 1].max()
        idx = np.logical_and(pbc_is_slave[0](x), pbc_is_slave[1](x))
        out_x[0][~idx] = np.nan
        out_x[1][~idx] = np.nan
        return out_x

    for ij in range(gdim):
        mpc.create_periodic_constraint_topological(V.sub(0).sub(ij), pbc_meshtags[1],
                                                   pbc_slave_tags[1],
                                                   pbc_slave_to_master_map,
                                                   bcs)

    mpc.create_periodic_constraint_topological(V.sub(1), pbc_meshtags[1],
                                               pbc_slave_tags[1],
                                               pbc_slave_to_master_map,
                                               bcs)

mpc.finalize()

# Variational problem
#####################
f_constant = Constant(domain, PETSc.ScalarType((1,0)))

## weak formulation
a_f = form(vis * inner(grad(ve), grad(dve)) * dx + pr * div(dve) * dx + div(ve) * dpr * dx)
L_f = form(inner(f_constant, dve) * dx)

# solution function
func_sol = Function(V)

## iterative solver
A = dolfinx_mpc.assemble_matrix(a_f, mpc, bcs=bcs)
A.assemble()
b = dolfinx_mpc.assemble_vector(L_f, mpc)
dolfinx_mpc.apply_lifting(b, [a_f], [bcs], mpc)
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
dolfinx.fem.petsc.set_bc(b, bcs)
b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)

# Preconditioner matrix form
Pf = dolfinx.fem.form(ufl.inner(ufl.grad(ve), ufl.grad(dve)) * dx + pr * dpr * dx)
P = dolfinx_mpc.assemble_matrix(Pf, mpc, bcs=bcs)
P.assemble()

# solver settings
ksp = PETSc.KSP().create(domain.comm)
ksp.setOperators(A, P)
ksp.setType(PETSc.KSP.Type.MINRES)
pc = ksp.getPC()
pc.setType("hypre")
pc.setHYPREType("boomeramg")
pc.setUp()

# solving and backsubstitution to obtain the values at the slave degrees of freedom
ksp.solve(b, func_sol.vector)
func_sol.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
mpc.backsubstitution(func_sol)

# collapse mixed functionspace
ve = func_sol.sub(0).collapse()
pr = func_sol.sub(1).collapse()

# average velocity
v_avg = domain.comm.allreduce(assemble_scalar(form(ve[0] * dx)), op=MPI.SUM)
if rank == 0:
    print("Average velocity in x direction:", v_avg)

# write results to file
with dolfinx.io.VTXWriter(MPI.COMM_WORLD, "results_2D_square/pressure.bp", [pr], engine="BP4") as vtx:
    vtx.write(0.0)
with dolfinx.io.VTXWriter(MPI.COMM_WORLD, "results_2D_square/velocity.bp", [ve], engine="BP4") as vtx:
    vtx.write(0.0)

context at https://fenicsproject.discourse.group/t/mpc-backsubstitution-in-parallel/12401/9