Discontinuity at interface using mixed domains

I think this is the variational form that you would like:

# SPDX-License-Identifier: MIT

from mpi4py import MPI
import dolfinx
import dolfinx.fem.petsc
import ufl
import numpy as np
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,
        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)
            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

    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)

# Split domain in half and set an interface tag of 5
gdim = mesh.geometry.dim
tdim = mesh.topology.dim
fdim = tdim - 1
top_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, top_boundary)
bottom_facets = 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

bottom_cells = dolfinx.mesh.locate_entities(mesh, tdim, half)
num_cells_local = (
    mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).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
)
all_b_facets = dolfinx.mesh.compute_incident_entities(
    mesh.topology, ct.find(3), tdim, fdim
)
all_t_facets = dolfinx.mesh.compute_incident_entities(
    mesh.topology, ct.find(4), tdim, fdim
)
interface = np.intersect1d(all_b_facets, all_t_facets)
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.find(3)
)[0:3]
submesh_t, submesh_t_to_mesh, t_v_map = dolfinx.mesh.create_submesh(
    mesh, tdim, ct.find(4)
)[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)

entity_maps = {submesh_b: parent_to_sub_b, submesh_t: parent_to_sub_t}
#entity_maps = {submesh_b._cpp_object: parent_to_sub_b, submesh_t._cpp_object: parent_to_sub_t}




def define_interior_eq(mesh,degree,  submesh, submesh_to_mesh, value):
    # Compute map from parent entity to submesh cell
    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)

    V = dolfinx.fem.functionspace(submesh, ("Lagrange", degree))
    u = dolfinx.fem.Function(V)
    v = ufl.TestFunction(V)
    ct_r = dolfinx.mesh.meshtags(mesh, mesh.topology.dim, submesh_to_mesh, np.full_like(submesh_to_mesh, 1, dtype=np.int32))
    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


u_0, F_00, m_to_b = define_interior_eq(mesh, 2, submesh_b, submesh_b_to_mesh, 0.0)
u_1, F_11, m_to_t = define_interior_eq(mesh, 1, submesh_t, submesh_t_to_mesh, 0.0)
u_0.name = "u_b"
u_1.name = "u_t"


# Add coupling term to the interface
# Get interface markers on submesh b
dInterface = ufl.Measure("dS", domain=mesh, subdomain_data=mt, subdomain_id=5)
b_res = "+"
t_res = "-"

v_b = ufl.TestFunction(u_0.function_space)(b_res)
v_t = ufl.TestFunction(u_1.function_space)(t_res)
u_b = u_0(b_res)
u_t = u_1(t_res)


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


W_0 = dolfinx.fem.functionspace(submesh_b, ("DG", 0))
K_0 = dolfinx.fem.Function(W_0)
K_0.x.array[:] = 1
W_1 = dolfinx.fem.functionspace(submesh_t, ("DG", 0))
K_1 = dolfinx.fem.Function(W_1)
K_1.x.array[:] = 1.5

n = ufl.FacetNormal(mesh)
n_b = n(b_res)
n_t = n(t_res)
K_b = K_0(b_res)
K_t = K_1(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/K_b - u_t/K_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/K_b - u_t/K_t), n_b) * dInterface
)
F_0 += 2 * gamma / (h_b + h_t) * (u_b/K_b - u_t/K_t) * v_b * dInterface
F_1 += -2 * gamma / (h_b + h_t) * (u_b/K_b - u_t/K_t) * v_t * dInterface

F_0 += F_00
F_1 += F_11

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),
]
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()
1 Like