Poisson submesh DG coupling

dg.py

from mpi4py import MPI
import dolfinx
import dolfinx.fem.petsc
import dolfinx.nls.petsc
import numpy
import argparse
from ufl import (
    Circumradius,
    FacetNormal,
    SpatialCoordinate,
    TrialFunction,
    TestFunction,
    derivative,
    div,
    dx,
    ds,
    dS,
    grad,
    inner,
    grad,
    avg,
    jump,
    dot,
    system,
)
from petsc4py import PETSc


class NewtonSolver:
    max_iterations: int
    bcs: list[dolfinx.fem.DirichletBC]
    A: PETSc.Mat
    b: PETSc.Vec
    J: dolfinx.fem.Form
    b: dolfinx.fem.Form
    dx: PETSc.Vec

    def __init__(
        self,
        F: list[dolfinx.fem.form],
        J: list[list[dolfinx.fem.form]],
        w: list[dolfinx.fem.Function],
        bcs: list[dolfinx.fem.DirichletBC] | None = None,
        max_iterations: int = 5,
        petsc_options: dict[str, str | float | int | None] = None,
    ):
        self.max_iterations = max_iterations
        self.bcs = [] if bcs is None else bcs
        self.b = dolfinx.fem.petsc.create_vector_block(F)
        self.F = F
        self.J = J
        self.A = dolfinx.fem.petsc.create_matrix_block(J)
        self.dx = self.A.createVecLeft()
        self.w = w
        self.x = dolfinx.fem.petsc.create_vector_block(F)

        # Set PETSc options
        opts = PETSc.Options()
        if petsc_options is not None:
            for k, v in petsc_options.items():
                opts[k] = v

        # Define KSP solver
        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
        self._solver.setOperators(self.A)
        self._solver.setFromOptions()

        # Set matrix and vector PETSc options
        self.A.setFromOptions()
        self.b.setFromOptions()

    def solve(self, tol=1e-6, beta=1.0):
        i = 0

        while i < self.max_iterations:
            dolfinx.cpp.la.petsc.scatter_local_vectors(
                self.x,
                [si.x.petsc_vec.array_r for si in self.w],
                [
                    (
                        si.function_space.dofmap.index_map,
                        si.function_space.dofmap.index_map_bs,
                    )
                    for si in self.w
                ],
            )
            self.x.ghostUpdate(
                addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
            )

            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
            with self.b.localForm() as b_local:
                b_local.set(0.0)

            dolfinx.fem.petsc.assemble_vector_block(
                self.b, self.F, self.J, bcs=self.bcs, x0=self.x, scale=-1.0
            )
            self.b.ghostUpdate(
                PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
            )
            # Assemble Jacobian
            self.A.zeroEntries()
            dolfinx.fem.petsc.assemble_matrix_block(self.A, self.J, bcs=self.bcs)
            self.A.assemble()
            self._solver.solve(self.b, self.dx)
            # self._solver.view()
            assert (
                self._solver.getConvergedReason() > 0
            ), "Linear solver did not converge"
            offset_start = 0
            for s in self.w:
                num_sub_dofs = (
                    s.function_space.dofmap.index_map.size_local
                    * s.function_space.dofmap.index_map_bs
                )
                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
                )
                s.x.petsc_vec.ghostUpdate(
                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
                )
                offset_start += num_sub_dofs
            # Compute norm of update

            correction_norm = self.dx.norm(0)
            print(f"Iteration {i}: Correction norm {correction_norm}")
            if correction_norm < tol:
                break
            i += 1


parser = argparse.ArgumentParser()
parser.add_argument(
    "--mixed",
    "-m",
    action="store_true",
    default=False,
    help="Used mixed expression where products are split by test function",
)
subparser = parser.add_subparsers(dest="solver")
linear_parser = subparser.add_parser("linear", help="Solve linear problem")
nonlinear_parser = subparser.add_parser("nonlinear", help="Solve nonlinear problem")
nonlinear_parser.add_argument(
    "--custom",
    "-c",
    action="store_true",
    default=False,
    help="Use custom Newton solver",
)

args = parser.parse_args()
N = 5
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, N, N)
V = dolfinx.fem.functionspace(mesh, ("DG", 1))
a = 0.8
c = 1

x = SpatialCoordinate(mesh)
u_ex = 1 + a * x[0] ** 2 + c * x[1] ** 2
f = -div(grad(u_ex))
n = FacetNormal(mesh)
h = 2 * Circumradius(mesh)
alpha = 10
gamma = 10
h_avg = avg(h)


def mixed_term(u, v, n):
    return dot(grad(u), n) * v


def exterior_nitsche(u, v, n, h, alpha):
    F = inner(grad(u), grad(v)) * dx - inner(n, grad(u)) * v * ds

    F += -inner(n, grad(v)) * u * ds + alpha / h * inner(u, v) * ds
    F -= inner(f, v) * dx
    F -= -inner(n, grad(v)) * u_ex * ds + alpha / h * inner(u_ex, v) * ds
    return F


def standard_ip(u, v, n, h_avg, gamma):
    F = -inner(jump(v, n), avg(grad(u))) * dS
    F += -inner(avg(grad(v)), jump(u, n)) * dS
    F += +gamma / h_avg * inner(jump(v, n), jump(u, n)) * dS
    return F


if args.solver == "linear":
    print("Solving linear problem")
    u = TrialFunction(V)
    v = TestFunction(V)
    F = exterior_nitsche(u, v, n, h, alpha)
    a, L = system(F)
    # Add DG/IP terms
    if args.mixed:
        print("Solvng with mixed form")
        u_b = u("+")
        u_t = u("-")
        v_b = v("+")
        v_t = v("-")
        n_b = n("+")

        h_b = h("+")
        h_t = h("-")

        a += -0.5 * mixed_term(v_b, (u_b - u_t), n_b) * dS
        a += -0.5 * mixed_term(v_t, (u_b - u_t), n_b) * dS

        a += -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dS
        a += +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dS

        a += 2 * gamma / (h_b + h_t) * (u_b - u_t) * v_b * dS
        a += -2 * gamma / (h_b + h_t) * (u_b - u_t) * v_t * dS

    else:
        print("Solving with standard form")
        a += standard_ip(u, v, n, h_avg, gamma)

    problem = dolfinx.fem.petsc.LinearProblem(
        a,
        L,
        petsc_options={
            "ksp_type": "preonly",
            "pc_type": "lu",
            "pc_factor_mat_solver_type": "mumps",
        },
    )
    uh = problem.solve()
else:
    print("Solving nonlinear problem")
    uh = dolfinx.fem.Function(V)
    v = TestFunction(V)
    F = exterior_nitsche(uh, v, n, h, alpha)
    if args.mixed:
        print("Solvng with mixed form")
        u_b = uh("+")
        u_t = uh("-")
        v_b = v("+")
        v_t = v("-")
        n_b = n("+")

        h_b = h("+")
        h_t = h("-")

        F += -0.5 * mixed_term(v_b, (u_b - u_t), n_b) * dS
        F += -0.5 * mixed_term(v_t, (u_b - u_t), n_b) * dS

        F += -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dS
        F += +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dS

        F += 2 * gamma / (h_b + h_t) * (u_b - u_t) * v_b * dS
        F += -2 * gamma / (h_b + h_t) * (u_b - u_t) * v_t * dS

    else:
        F += standard_ip(uh, v, n, h_avg, gamma)

    if args.custom:
        J = derivative(F, uh)

        solver = NewtonSolver(
            [dolfinx.fem.form(F)],
            [[dolfinx.fem.form(J)]],
            [uh],
            petsc_options={
                "ksp_type": "preonly",
                "pc_type": "lu",
                "pc_factor_mat_solver_type": "mumps",
            },
        )
        solver.solve()
    else:
        problem = dolfinx.fem.petsc.NonlinearProblem(F, uh)
        solver = dolfinx.nls.petsc.NewtonSolver(mesh.comm, problem)
        solver.krylov_solver.setType("preonly")
        solver.krylov_solver.getPC().setType("lu")
        dolfinx.log.set_log_level(dolfinx.log.LogLevel.INFO)
        solver.solve(uh)

# We compute the error of the computation by comparing it to the analytical solution

error_form = dolfinx.fem.form(inner(uh - u_ex, uh - u_ex) * dx)
errorL2 = numpy.sqrt(dolfinx.fem.assemble_scalar(error_form))
print(rf"$L^2$-error: {errorL2:.2e}")

with dolfinx.io.VTXWriter(mesh.comm, "uh.bp", [uh], engine="BP4") as bp:
    bp.write(0.0)

problem.py

# SPDX-License-Identifier: MIT

from mpi4py import MPI
import dolfinx
import dolfinx.fem.petsc
import ufl
import numpy as np
from petsc4py import PETSc
import numpy.typing as npt
from packaging.version import Version

def compute_interface_data(
    cell_tags: dolfinx.mesh.MeshTags, facet_indices: npt.NDArray[np.int32]
) -> npt.NDArray[np.int32]:
    """
    Compute interior facet integrals that are consistently ordered according to the `cell_tags`,
    such that the data `(cell0, facet_idx0, cell1, facet_idx1)` is ordered such that
    `cell_tags[cell0]`<`cell_tags[cell1]`, i.e the cell with the lowest cell marker is considered the
    "+" restriction".

    Args:
        cell_tags: MeshTags that must contain an integer marker for all cells adjacent to the `facet_indices`
        facet_indices: List of facets (local index) that are on the interface.
    Returns:
        The integration data.
    """
    # Future compatibilty check
    integration_args: tuple[int] | tuple
    if Version("0.10.0") <= Version(dolfinx.__version__):
        integration_args = ()
    else:
        fdim = cell_tags.dim - 1
        integration_args = (fdim,)
    idata = dolfinx.cpp.fem.compute_integration_domains(
        dolfinx.fem.IntegralType.interior_facet,
        cell_tags.topology,
        facet_indices,
        *integration_args,
    )
    ordered_idata = idata.reshape(-1, 4).copy()
    switch = cell_tags.values[ordered_idata[:, 0]] > cell_tags.values[ordered_idata[:, 2]]
    if True in switch:
        ordered_idata[switch, :] = ordered_idata[switch][:, [2, 3, 0, 1]]
    return ordered_idata
def communicate_indices(
    index_map: dolfinx.common.IndexMap, local_indices: npt.NDArray[np.int32]
) -> npt.NDArray[np.int32]:
    """
    Given a set of local indices find ghosted indices marked by any
    other other process (other ghosted processes or owner process).
    Args:
        index_map: Map describing the ownership structure of the
            entities
        local_indices: List of indices local to process that should
            be distributed
    Returns:
        All indices (local index) on process (including ghosts) that
        have been marked by this or and other process.
    """
    index_accumulator = dolfinx.la.vector(index_map, 1)
    index_accumulator.array[:] = 0
    index_accumulator.array[local_indices] = 1
    index_accumulator.scatter_reverse(dolfinx.la.InsertMode.add)
    return np.flatnonzero(index_accumulator.array).astype(np.int32)


class NewtonSolver:
    max_iterations: int
    bcs: list[dolfinx.fem.DirichletBC]
    A: PETSc.Mat
    b: PETSc.Vec
    J: dolfinx.fem.Form
    b: dolfinx.fem.Form
    dx: PETSc.Vec

    def __init__(
        self,
        F: list[dolfinx.fem.form],
        J: list[list[dolfinx.fem.form]],
        w: list[dolfinx.fem.Function],
        bcs: list[dolfinx.fem.DirichletBC] | None = None,
        max_iterations: int = 5,
        petsc_options: dict[str, str | float | int | None] = None,
        problem_prefix="newton",
    ):
        self.max_iterations = max_iterations
        self.bcs = [] if bcs is None else bcs
        self.b = dolfinx.fem.petsc.create_vector_block(F)
        self.F = F
        self.J = J
        self.A = dolfinx.fem.petsc.create_matrix_block(J)
        self.dx = self.A.createVecLeft()
        self.w = w
        self.x = dolfinx.fem.petsc.create_vector_block(F)

        # Set PETSc options
        opts = PETSc.Options()
        if petsc_options is not None:
            for k, v in petsc_options.items():
                opts[k] = v

        # Define KSP solver
        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
        self._solver.setOperators(self.A)
        self._solver.setFromOptions()

        # Set matrix and vector PETSc options
        self.A.setFromOptions()
        self.b.setFromOptions()

    def solve(self, tol=1e-6, beta=1.0):
        i = 0

        while i < self.max_iterations:
            dolfinx.cpp.la.petsc.scatter_local_vectors(
                self.x,
                [si.x.petsc_vec.array_r for si in self.w],
                [
                    (
                        si.function_space.dofmap.index_map,
                        si.function_space.dofmap.index_map_bs,
                    )
                    for si in self.w
                ],
            )
            self.x.ghostUpdate(
                addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
            )

            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
            with self.b.localForm() as b_local:
                b_local.set(0.0)
            if Version(dolfinx.__version__) < Version("0.9.0"):
                dolfinx.fem.petsc.assemble_vector_block(
                    self.b, self.F, self.J, bcs=self.bcs, x0=self.x, scale=-1.0
                )
            else:
                dolfinx.fem.petsc.assemble_vector_block(
                    self.b, self.F, self.J, bcs=self.bcs, x0=self.x, alpha=-1.0
                )
            self.b.ghostUpdate(
                PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD
            )


            # Assemble Jacobian
            self.A.zeroEntries()
            dolfinx.fem.petsc.assemble_matrix_block(self.A, self.J, bcs=self.bcs)
            self.A.assemble()

            self._solver.solve(self.b, self.dx)
            # self._solver.view()
            assert (
                self._solver.getConvergedReason() > 0
            ), "Linear solver did not converge"
            offset_start = 0
            for s in self.w:
                num_sub_dofs = (
                    s.function_space.dofmap.index_map.size_local
                    * s.function_space.dofmap.index_map_bs
                )
                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
                )
                s.x.petsc_vec.ghostUpdate(
                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
                )
                offset_start += num_sub_dofs
            # Compute norm of update

            correction_norm = self.dx.norm(0)
            print(f"Iteration {i}: Correction norm {correction_norm}")
            if correction_norm < tol:
                break
            i += 1

    def __del__(self):
        self.A.destroy()
        self.b.destroy()
        self.dx.destroy()
        self._solver.destroy()
        self.x.destroy()


def transfer_meshtags_to_submesh(
    mesh, entity_tag, submesh, sub_vertex_to_parent, sub_cell_to_parent
):
    """
    Transfer a meshtag from a parent mesh to a sub-mesh.
    """

    tdim = mesh.topology.dim
    cell_imap =mesh.topology.index_map(tdim)
    num_cells = cell_imap.size_local + cell_imap.num_ghosts
    mesh_to_submesh = np.full(num_cells, -1)
    mesh_to_submesh[sub_cell_to_parent] = np.arange(
        len(sub_cell_to_parent), dtype=np.int32
    )
    sub_vertex_to_parent = np.asarray(sub_vertex_to_parent)

    submesh.topology.create_connectivity(entity_tag.dim, 0)

    num_child_entities = (
        submesh.topology.index_map(entity_tag.dim).size_local
        + submesh.topology.index_map(entity_tag.dim).num_ghosts
    )
    submesh.topology.create_connectivity(submesh.topology.dim, entity_tag.dim)

    c_c_to_e = submesh.topology.connectivity(submesh.topology.dim, entity_tag.dim)
    c_e_to_v = submesh.topology.connectivity(entity_tag.dim, 0)

    child_markers = np.full(num_child_entities, 0, dtype=np.int32)

    mesh.topology.create_connectivity(entity_tag.dim, 0)
    mesh.topology.create_connectivity(entity_tag.dim, mesh.topology.dim)
    p_f_to_v = mesh.topology.connectivity(entity_tag.dim, 0)
    p_f_to_c = mesh.topology.connectivity(entity_tag.dim, mesh.topology.dim)
    sub_to_parent_entity_map = np.full(num_child_entities, -1, dtype=np.int32)
    for facet, value in zip(entity_tag.indices, entity_tag.values):
        facet_found = False
        for cell in p_f_to_c.links(facet):
            if facet_found:
                break
            if (child_cell := mesh_to_submesh[cell]) != -1:
                for child_facet in c_c_to_e.links(child_cell):
                    child_vertices = c_e_to_v.links(child_facet)
                    child_vertices_as_parent = sub_vertex_to_parent[child_vertices]
                    is_facet = np.isin(
                        child_vertices_as_parent, p_f_to_v.links(facet)
                    ).all()
                    if is_facet:
                        child_markers[child_facet] = value
                        facet_found = True
                        sub_to_parent_entity_map[child_facet] = facet
    tags = dolfinx.mesh.meshtags(
        submesh,
        entity_tag.dim,
        np.arange(num_child_entities, dtype=np.int32),
        child_markers,
    )
    tags.name = entity_tag.name
    return tags, sub_to_parent_entity_map


def bottom_boundary(x):
    return np.isclose(x[1], 0.0)


def top_boundary(x):
    return np.isclose(x[1], 1.0)


def half(x):
    return x[1] <= 0.5 + 1e-14

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10, dolfinx.mesh.CellType.triangle, ghost_mode=dolfinx.mesh.GhostMode.shared_facet)

# Split domain in half and set an interface tag of 5
gdim = mesh.geometry.dim
tdim = mesh.topology.dim
fdim = tdim - 1
mesh.topology.create_connectivity(fdim, tdim)
facet_map = mesh.topology.index_map(fdim)

top_facets = communicate_indices(facet_map, dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary))
bottom_facets = communicate_indices(facet_map, dolfinx.mesh.locate_entities_boundary(mesh, fdim, bottom_boundary))
num_facets_local = (
    mesh.topology.index_map(fdim).size_local + mesh.topology.index_map(fdim).num_ghosts
)
facets = np.arange(num_facets_local, dtype=np.int32)
values = np.full_like(facets, 0, dtype=np.int32)
values[top_facets] = 1
values[bottom_facets] = 2

cell_map =  mesh.topology.index_map(tdim)
bottom_cells = dolfinx.mesh.locate_entities(mesh, tdim, half)
num_cells_local = (
   cell_map.size_local +cell_map.num_ghosts
)
cells = np.full(num_cells_local, 4, dtype=np.int32)
cells[bottom_cells] = 3
ct = dolfinx.mesh.meshtags(
    mesh, tdim, np.arange(num_cells_local, dtype=np.int32), cells
)
ct_b = communicate_indices(cell_map, ct.find(3))
ct_t = communicate_indices(cell_map, ct.find(4))

all_b_facets = dolfinx.mesh.compute_incident_entities(
    mesh.topology, ct_b, tdim, fdim
)
all_t_facets = dolfinx.mesh.compute_incident_entities(
    mesh.topology, ct_t, tdim, fdim
)
abf = communicate_indices(facet_map, all_b_facets)
atf = communicate_indices(facet_map, all_t_facets)
interface = np.intersect1d(abf, atf)
    

values[interface] = 5

mt = dolfinx.mesh.meshtags(mesh, mesh.topology.dim - 1, facets, values)

submesh_b, submesh_b_to_mesh, b_v_map = dolfinx.mesh.create_submesh(
    mesh, tdim, ct_b
)[0:3]
submesh_t, submesh_t_to_mesh, t_v_map = dolfinx.mesh.create_submesh(
    mesh, tdim, ct_t
)[0:3]
parent_to_sub_b = np.full(num_facets_local, -1, dtype=np.int32)
parent_to_sub_b[submesh_b_to_mesh] = np.arange(len(submesh_b_to_mesh), dtype=np.int32)
parent_to_sub_t = np.full(num_facets_local, -1, dtype=np.int32)
parent_to_sub_t[submesh_t_to_mesh] = np.arange(len(submesh_t_to_mesh), dtype=np.int32)    

# We need to modify the cell maps, as for `dS` integrals of interfaces between submeshes, there is no entity to map to.
# We use the entity on the same side to fix this (as all restrictions are one-sided)

# Transfer meshtags to submesh
ft_b, b_facet_to_parent = transfer_meshtags_to_submesh(
    mesh, mt, submesh_b, b_v_map, submesh_b_to_mesh
)
ft_t, t_facet_to_parent = transfer_meshtags_to_submesh(
    mesh, mt, submesh_t, t_v_map, submesh_t_to_mesh
)

t_parent_to_facet = np.full(num_facets_local, -1)
t_parent_to_facet[t_facet_to_parent] = np.arange(len(t_facet_to_parent), dtype=np.int32)


# Hack, as we use one-sided restrictions, pad dS integral with the same entity from the same cell on both sides
mesh.topology.create_connectivity(fdim, tdim)
f_to_c = mesh.topology.connectivity(fdim, tdim)
for facet in mt.find(5):
    cells = f_to_c.links(facet)
    assert len(cells) == 2
    b_map = parent_to_sub_b[cells]
    t_map = parent_to_sub_t[cells]
    parent_to_sub_b[cells] = max(b_map)
    parent_to_sub_t[cells] = max(t_map)

if Version(dolfinx.__version__) < Version("0.9.0"):
    entity_maps = {submesh_b._cpp_object: parent_to_sub_b, submesh_t._cpp_object: parent_to_sub_t}

else:
    entity_maps = {submesh_b: parent_to_sub_b, submesh_t: parent_to_sub_t}




def define_interior_eq(v, mesh, submesh_to_mesh, value):
    # Compute map from parent entity to submesh cell
    submesh = v.ufl_function_space().mesh

    codim = mesh.topology.dim - submesh.topology.dim
    ptdim = mesh.topology.dim - codim
    num_entities = (
        mesh.topology.index_map(ptdim).size_local
        + mesh.topology.index_map(ptdim).num_ghosts
    )
    mesh_to_submesh = np.full(num_entities, -1)
    mesh_to_submesh[submesh_to_mesh] = np.arange(len(submesh_to_mesh), dtype=np.int32)

    u = dolfinx.fem.Function(v.ufl_function_space())
    sort_order = np.argsort(submesh_to_mesh)
    ct_r = dolfinx.mesh.meshtags(mesh, mesh.topology.dim, submesh_to_mesh[sort_order], np.full_like(submesh_to_mesh, 1, dtype=np.int32)[sort_order])
    val = dolfinx.fem.Constant(submesh, value)
    dx_r = ufl.Measure("dx", domain=mesh, subdomain_data=ct_r, subdomain_id=1)
    F = ufl.inner(ufl.grad(u), ufl.grad(v)) * dx_r - val * v * dx_r
    return u, F, mesh_to_submesh

b_res = "+"
t_res = "-"
V0 = dolfinx.fem.functionspace(submesh_b, ("Lagrange", 2))
V1 = dolfinx.fem.functionspace(submesh_t, ("Lagrange", 1))
if Version(dolfinx.__version__) < Version("0.9.0"):
    v_0 = ufl.TestFunction(V0)
    v_1 = ufl.TestFunction(V1)
else:
    W = ufl.MixedFunctionSpace(V0, V1)
    v_0, v_1 = ufl.TestFunctions(W)

u_0, F_00, m_to_b = define_interior_eq(v_0, mesh, submesh_b_to_mesh, 3.0)
u_1, F_11, m_to_t = define_interior_eq(v_1, mesh, submesh_t_to_mesh, 1.0)
u_0.name = "u_b"
u_1.name = "u_t"


# Add coupling term to the interface
# Get interface markers on submesh b
ordered_integration_data = compute_interface_data(ct, interface)
# Pad entity maps for sparsity pattern
parent_cells_plus = ordered_integration_data[:, 0]
parent_cells_minus = ordered_integration_data[:, 2]
entity_maps[submesh_b][parent_cells_minus] = entity_maps[submesh_b][parent_cells_plus]
entity_maps[submesh_t][parent_cells_plus] = entity_maps[submesh_t][parent_cells_minus]
ordered_integration_data = ordered_integration_data.flatten()
integral_data_interface = [(13, ordered_integration_data)]
dInterface = ufl.Measure(
    "dS",
    domain=mesh,
    subdomain_data=integral_data_interface,
    subdomain_id=13,
)


u_b = u_0(b_res)
u_t = u_1(t_res)
v_b = v_0(b_res)
v_t = v_1(t_res)

def mixed_term(u, v, n):
    return ufl.dot(ufl.grad(u), n) * v


n = ufl.FacetNormal(mesh)
n_b = n(b_res)
n_t = n(t_res)
cr = ufl.Circumradius(mesh)
h_b = 2 * cr(b_res)
h_t = 2 * cr(t_res)
gamma = 50.0


F_0 = (
    -0.5 * mixed_term((u_b + u_t), v_b, n_b) * dInterface
    - 0.5 * mixed_term(v_b, (u_b - u_t), n_b) * dInterface
)

F_1 = (
    +0.5 * mixed_term((u_b + u_t), v_t, n_b) * dInterface
    - 0.5 * mixed_term(v_t, (u_b - u_t), n_b) * dInterface
)
F_0 += 2 * gamma / (h_b + h_t) * (u_b - u_t) * v_b * dInterface
F_1 += -2 * gamma / (h_b + h_t) * (u_b - u_t) * v_t * dInterface

F_0 += F_00
F_1 += F_11

if Version(dolfinx.__version__) < Version("0.9.0"):
    jac00 = ufl.derivative(F_0, u_0)

    jac01 = ufl.derivative(F_0, u_1)

    jac10 = ufl.derivative(F_1, u_0)
    jac11 = ufl.derivative(F_1, u_1)
    J00 = dolfinx.fem.form(jac00, entity_maps=entity_maps)


    J01 = dolfinx.fem.form(jac01, entity_maps=entity_maps)
    J10 = dolfinx.fem.form(jac10, entity_maps=entity_maps)
    J11 = dolfinx.fem.form(jac11, entity_maps=entity_maps)
    J = [[J00, J01], [J10, J11]]
    F = [
        dolfinx.fem.form(F_0, entity_maps=entity_maps),
        dolfinx.fem.form(F_1, entity_maps=entity_maps),
    ]
else:
    F = dolfinx.fem.form(ufl.extract_blocks(F_0 + F_1),
                         entity_maps=entity_maps)
    J = dolfinx.fem.form(ufl.extract_blocks(ufl.derivative(F_0+F_1, [u_0,u_1], ufl.TrialFunctions(W))),
                         entity_maps=entity_maps)

b_bc = dolfinx.fem.Function(u_0.function_space)
b_bc.x.array[:] = 0.2
submesh_b.topology.create_connectivity(
    submesh_b.topology.dim - 1, submesh_b.topology.dim
)
bc_b = dolfinx.fem.dirichletbc(
    b_bc, dolfinx.fem.locate_dofs_topological(u_0.function_space, fdim, ft_b.find(2))
)


t_bc = dolfinx.fem.Function(u_1.function_space)
t_bc.x.array[:] = 0.05
submesh_t.topology.create_connectivity(
    submesh_t.topology.dim - 1, submesh_t.topology.dim
)
bc_t = dolfinx.fem.dirichletbc(
    t_bc, dolfinx.fem.locate_dofs_topological(u_1.function_space, fdim, ft_t.find(1))
)
bcs = [bc_b, bc_t]


solver = NewtonSolver(
    F,
    J,
    [u_0, u_1],
    bcs=bcs,
    max_iterations=2,
    petsc_options={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
    },
)
solver.solve(1e-5)

bp = dolfinx.io.VTXWriter(mesh.comm, "u_b.bp", [u_0], engine="BP4")
bp.write(0)
bp.close()
bp = dolfinx.io.VTXWriter(mesh.comm, "u_t.bp", [u_1], engine="BP4")
bp.write(0)
bp.close()