Solve a problem of parts of a unit square mesh

submesh_problem.py

import dolfinx
from mpi4py import MPI
import ufl
import numpy as np
from petsc4py.PETSc import ScalarType


def transfer_submesh_data(u_parent: dolfinx.fem.Function, u_sub: dolfinx.fem.Function,
                          sub_to_parent_cells: np.ndarray, inverse: bool = False):
    """
    Transfer data between a function from the parent mesh and a function from the sub mesh.
    Both functions has to share the same element dof layout

    Args:
        u_parent: Function on parent mesh
        u_sub: Function on sub mesh
        sub_to_parent_cells: Map from sub mesh (local index) to parent mesh (local index)
        inverse: If true map from u_sub->u_parent else u_parent->u_sub
    """
    V_parent = u_parent.function_space
    V_sub = u_sub.function_space
    # FIXME: In C++ check elementlayout for equality
    if inverse:
        for i, cell in enumerate(sub_to_parent_cells):
            bs = V_parent.dofmap.bs
            bs_sub = V_sub.dofmap.bs
            assert(bs == bs_sub)
            parent_dofs = V.dofmap.cell_dofs(cell)
            sub_dofs = V_sub.dofmap.cell_dofs(i)
            for p_dof, s_dof in zip(parent_dofs, sub_dofs):
                for j in range(bs):
                    u_parent.x.array[p_dof * bs + j] = u_sub.x.array[s_dof * bs + j]
    else:
        for i, cell in enumerate(sub_to_parent_cells):
            bs = V_parent.dofmap.bs
            bs_sub = V_sub.dofmap.bs
            assert(bs == bs_sub)
            parent_dofs = V.dofmap.cell_dofs(cell)
            sub_dofs = V_sub.dofmap.cell_dofs(i)
            for p_dof, s_dof in zip(parent_dofs, sub_dofs):
                for j in range(bs):
                    u_sub.x.array[s_dof * bs + j] = u_parent.x.array[p_dof * bs + j]


# Create parent mesh
N = 8
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, N, N, ghost_mode=dolfinx.mesh.GhostMode.shared_facet)
V = dolfinx.fem.FunctionSpace(mesh, ("CG", 1))

# Create source function on parent mesh
f = dolfinx.fem.Function(V)
f.interpolate(lambda x: 2 * np.pi**2 * (np.sin(x[0] * np.pi) * np.sin(np.pi * x[1])))

# Create sub mesh consisting of M elements in x direction
M = 5
assert(M <= N)
x_b = M / N
left_entities = dolfinx.mesh.locate_entities(mesh, mesh.topology.dim, lambda x: x[0] <= x_b + 5e-16)
local_left_entities = np.sort(left_entities)[left_entities < mesh.topology.index_map(mesh.topology.dim).size_local]
sub_mesh, cell_map, vertex_map, geometry_map = dolfinx.mesh.create_submesh(mesh, mesh.topology.dim, local_left_entities)

# Compute all boundary facets of sub mesh
sub_mesh.topology.create_connectivity(sub_mesh.topology.dim - 1, sub_mesh.topology.dim)
boundary_indicator = dolfinx.mesh.compute_boundary_facets(sub_mesh.topology)
boundary_facets = np.flatnonzero(boundary_indicator)

# Transfer source function from parent mesh
V_sub = dolfinx.fem.FunctionSpace(sub_mesh, ("CG", 1))
f_sub = dolfinx.fem.Function(V_sub)
transfer_submesh_data(f, f_sub, cell_map)

# Compute boundary facets of sub mesh for Neumann condition
mid_facets = dolfinx.mesh.locate_entities_boundary(sub_mesh, sub_mesh.topology.dim - 1, lambda x: np.isclose(x[0], x_b))
mid_facets = np.sort(mid_facets)

# Use remaining boundary facets for DirichletBC
bc_facets = np.setdiff1d(np.sort(boundary_facets), mid_facets)
bc_dofs = dolfinx.fem.locate_dofs_topological(
    V_sub, sub_mesh.topology.dim - 1, bc_facets)
bc = dolfinx.fem.dirichletbc(ScalarType(0), bc_dofs, V_sub)

# Create meshtag on sub mesh
exterior_facets = np.hstack([bc_facets, mid_facets])
argsort = np.argsort(exterior_facets)
values = np.hstack([np.ones(len(bc_facets), dtype=np.int32), np.full(len(mid_facets), 2, dtype=np.int32)])
sub_tag = dolfinx.mesh.meshtags(sub_mesh, sub_mesh.topology.dim - 1, exterior_facets[argsort], values[argsort])
print(len(sub_tag.indices[sub_tag.values == 1]))
# Create variational form on sub mesh
ds = ufl.Measure("ds", domain=sub_mesh, subdomain_data=sub_tag, subdomain_id=2)
u = ufl.TrialFunction(V_sub)
v = ufl.TestFunction(V_sub)
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
x_sub = ufl.SpatialCoordinate(sub_mesh)
g = ufl.pi * ufl.cos(ufl.pi * x_sub[0]) * ufl.sin(ufl.pi * x_sub[1])
L = ufl.inner(f_sub, v) * ufl.dx + ufl.inner(g, v) * ds

# Solve linear problem on sub mesh
problem = dolfinx.fem.petsc.LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
u_sub = problem.solve()

with dolfinx.io.XDMFFile(sub_mesh.comm, "sub_problem.xdmf", "w") as xdmf:
    xdmf.write_mesh(sub_mesh)
    # xdmf.write_meshtags(sub_tag)
    xdmf.write_function(u_sub)

# Transfer solution to parent mesh
u = dolfinx.fem.Function(V)
u.x.array[:] = 0
transfer_submesh_data(u, u_sub, cell_map, inverse=True)
u.x.scatter_forward()
with dolfinx.io.XDMFFile(mesh.comm, "parent_problem.xdmf", "w") as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(u)

# Solve equivalent problem on the full mesh
mesh.topology.create_entities(mesh.topology.dim - 1)
mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
boundary_facets = np.flatnonzero(dolfinx.mesh.compute_boundary_facets(mesh.topology))
bc = dolfinx.fem.dirichletbc(ScalarType(0), dolfinx.fem.locate_dofs_topological(
    V, mesh.topology.dim - 1, boundary_facets), V)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
L = ufl.inner(f, v) * ufl.dx
problem = dolfinx.fem.petsc.LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
u_full = problem.solve()

with dolfinx.io.XDMFFile(mesh.comm, "full_problem.xdmf", "w") as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(u_full)

# Compute error between sub mesh solution and parent solution
u_full_sub = dolfinx.fem.Function(V_sub)
transfer_submesh_data(u_full, u_full_sub, cell_map)

L2_error = dolfinx.fem.form(ufl.inner(u_sub - u_full_sub, u_sub - u_full_sub) * ufl.dx)
error = sub_mesh.comm.allreduce(dolfinx.fem.assemble_scalar(L2_error))
print(f"L2 error: {error}")
x = ufl.SpatialCoordinate(mesh)
u_ex = ufl.sin(x[0] * ufl.pi) * ufl.sin(np.pi * x[1])
print("Full error: ", mesh.comm.allreduce(dolfinx.fem.assemble_scalar(
    dolfinx.fem.form(ufl.inner(u_full - u_ex, u_full - u_ex) * ufl.dx))))


u_ex = ufl.sin(x_sub[0] * ufl.pi) * ufl.sin(np.pi * x_sub[1])
print("Sub error : ", sub_mesh.comm.allreduce(dolfinx.fem.assemble_scalar(
    dolfinx.fem.form(ufl.inner(u_sub - u_ex, u_sub - u_ex) * ufl.dx))))