Interface discontinuity

[RemDelaporteMathurin]

For multi-material simulations, interface discontinuities are needed.

Let $c$ be the concentration of hydrogen, we have

$$ \nabla \cdot ( D \nabla c) = -f \quad \text{on} \ \Omega $$

and

$$ \frac{c^-}{K^-} = \frac{c^+}{K^+} \quad \text{on \ interface} $$

@SamueleMeschini started working on a DG method implementation in pure fenics.

However, we need to find a way to determine what materials correspond to the "-" and "+" sides of the interface to have it generic.

Ideally, we would only have DG elements on the interface and CG elements everywhere else. @jpdean and @jorgensd have helped us a lot on this topic thank you!

[jpdean]

You can do this by specifying the facets to integrate over manually, see this test. You'd need to tag the facet on the interface, get the cells connected to them via the facet to cell connectivity, and pass a list of ((cell, local_facet_index), (cell, local_facet_index)) pairs when the integration measure is created. The first entry is "+" and the second is "-". Admittedly, this isn't the most intuitive interface, so I'll have a think about how to tidy it up. It is now possible to create CG spaces over each subdomain and couple across the interface with Nitsche's methods in FEniCSx main. I'm planning to add a demo on how to do all of this soon :)

[RemDelaporteMathurin]

You can do this by specifying the facets to integrate over manually

Means we couldn't use mesh.locate_entities and mesh.meshtags for this but have to tag them manually?

pass a list of ((cell, local_facet_index), (cell, local_facet_index)) pairs

in your example it looks more like a list of [cell, local_facet_index, cell, local_facet_index, ...] than [(cell, local_facet_index), (cell, local_facet_index)], correct?

It is now possible to create CG spaces over each subdomain and couple across the interface with Nitsche's methods in FEniCSx main

I think @jorgensd mentioned it to me yes! I'm trying to figure it out but a small demo would be great!

[jpdean]

Means we couldn't use mesh.locate_entities and mesh.meshtags for this but have to tag them manually?

You can still use mesh.locate_entities and mesh.meshtags, it's just you'd need to get the (cell, local_facet_index) pairs from the facet numbers via the mesh connectivity and pass them to to the integration measure.

in your example it looks more like a list of [cell, local_facet_index, cell, local_facet_index, ...] than [(cell, local_facet_index), (cell, local_facet_index)], correct?

The list is flattened, but it represents a list of ((cell, local_facet_index), (cell, local_facet_index)) pairs, where each pair identifies a given interior facet from each cell sharing it.

When I add a demo, I'll add some helper functions / tidy the interface so that this is more intuitive

[RemDelaporteMathurin]

I've tried adapting the code in the test to @SamueleMeschini 's example but I get an error because when creating the NonLinearProblem. I understand it's because the subdomain data for each integral isn't the same.

Does that mean I cannot do:

int_facet_domain = []
for f in interface_facets:
   if f >= facet_imap.size_local or len(f_to_c.links(f)) != 2:
      continue
   c_0, c_1 = f_to_c.links(f)[0], f_to_c.links(f)[1]
   local_f_0 = np.where(c_to_f.links(c_0) == f)[0][0]
   local_f_1 = np.where(c_to_f.links(c_1) == f)[0][0]
   int_facet_domain.append(c_0)
   int_facet_domain.append(local_f_0)
   int_facet_domain.append(c_1)
   int_facet_domain.append(local_f_1)
int_facet_domains = [(2, int_facet_domain)]

ds = ufl.Measure("ds", domain=msh, subdomain_data=mesh_tags_facets)
dS = ufl.Measure("dS", domain=msh, subdomain_data=mesh_tags_facets)

dS_manual = ufl.Measure("dS", subdomain_data=int_facet_domains, domain=msh)

and use a combination of ds, dx, dS and dS_manual in the form?

[RemDelaporteMathurin]

@jpdean @jorgensd I managed to get what I want

For each interface (assuming it's only shared between 2 volumes) I look at the first facet and determine what is the tag of subdomain '+' and subdomain '-'. From that, I can know what are the correct Ks.

cell_tag_to_K = {
    1: fem.Constant(msh, PETSc.ScalarType(2.0)),
    2: fem.Constant(msh, PETSc.ScalarType(8.0)),
    3: fem.Constant(msh, PETSc.ScalarType(4.0)),
}

for interface_id, interface_facets in zip([2, 3, 4], [interface_facets_1, interface_facets_2, interface_facets_3]):
    # look at the first facet on interface
    # and get the two cells that are connected to it
    # and get the material properties of these cells
    first_facet_interface = interface_facets[0]
    c_plus, c_minus = f_to_c.links(first_facet_interface)[0], f_to_c.links(first_facet_interface)[1]
    id_minus, id_plus = values_cells[c_plus], values_cells[c_minus]
    # TODO ensure that id_minus and id_plus are the same for all facets on the interface

    K1 = cell_tag_to_K[id_minus]
    K2 = cell_tag_to_K[id_plus]

Complete code:

from dolfinx import mesh, fem, plot, geometry
import dolfinx.fem.petsc
import dolfinx.nls.petsc
from mpi4py import MPI
import ufl
from petsc4py import PETSc
from ufl import grad, dot, jump, avg
import numpy as np
import matplotlib.pyplot as plt
import pyvista

msh = mesh.create_unit_square(MPI.COMM_WORLD, 100, 100)
V = fem.FunctionSpace(msh, ("DG", 1))
uD = fem.Function(V)
uD.interpolate(lambda x: np.full(x[0].shape, 0.0))

# create mesh tags
tol = 1e-16
def marker_interface_1(x):
    return np.logical_and(np.isclose(x[0], 0.5), [x[1] > 0.5 - tol])

def marker_interface_2(x):
    return np.logical_and(np.isclose(x[0], 0.5), [x[1] < 0.5 + tol])

def marker_interface_3(x):
    return np.logical_and([x[0] > 0.5 - tol], [np.isclose(x[1], 0.5)])

tdim = msh.topology.dim
fdim = tdim - 1
msh.topology.create_connectivity(fdim, tdim)
c_to_f = msh.topology.connectivity(tdim, tdim - 1)
f_to_c = msh.topology.connectivity(tdim - 1, tdim)
facet_map = msh.topology.index_map(tdim - 1)
cell_imap = msh.topology.index_map(tdim)
boundary_facets = mesh.exterior_facet_indices(msh.topology)
interface_facets_1 = mesh.locate_entities(msh, tdim - 1, marker_interface_1)
interface_facets_2 = mesh.locate_entities(msh, tdim - 1, marker_interface_2)
interface_facets_3 = mesh.locate_entities(msh, tdim - 1, marker_interface_3)
num_facets = facet_map.size_local + facet_map.num_ghosts
indices = np.arange(0, num_facets)
# values = np.arange(0, num_facets, dtype=np.intc)
values = np.zeros(indices.shape, dtype=np.intc)  # all facets are tagged with zero

values[boundary_facets] = 1
values[interface_facets_1] = 2
values[interface_facets_2] = 3
values[interface_facets_3] = 4

mesh_tags_facets = mesh.meshtags(msh, tdim - 1, indices, values)

# tag cells
def left_domain(x):
    tol = 1e-16
    return x[0] < 0.5 + tol

def top_right_domain(x):
    tol = 1e-16
    return np.logical_and([x[0] > 0.5 - tol], [x[1] > 0.5 - tol])

def bottom_right_domain(x):
    tol = 1e-16
    return np.logical_and([x[0] > 0.5 - tol], [x[1] < 0.5 + tol])

left_cells = mesh.locate_entities(msh, tdim, left_domain)
top_right_cells = mesh.locate_entities(msh, tdim, top_right_domain)
bottom_right_cells = mesh.locate_entities(msh, tdim, bottom_right_domain)
num_cells = cell_imap.size_local
indices = np.arange(0, num_cells)
values_cells = np.zeros(indices.shape, dtype=np.intc)
values_cells[left_cells] = 1
values_cells[top_right_cells] = 2
values_cells[bottom_right_cells] = 3

mesh_tags_cells = mesh.meshtags(msh, tdim, indices, values_cells)

# Create measures
ds = ufl.Measure("ds", domain=msh, subdomain_data=mesh_tags_facets)
dS = ufl.Measure("dS", domain=msh, subdomain_data=mesh_tags_facets)
dx = ufl.Measure("dx", domain=msh, subdomain_data=mesh_tags_cells)

u = fem.Function(V)
u_n = fem.Function(V)
v = ufl.TestFunction(V)

h = ufl.CellDiameter(msh)
n = ufl.FacetNormal(msh)

# Define parameters
alpha = 1000

# Simulation constants
f = fem.Constant(msh, PETSc.ScalarType(2.0))

cell_tag_to_K = {
    1: fem.Constant(msh, PETSc.ScalarType(2.0)),
    2: fem.Constant(msh, PETSc.ScalarType(8.0)),
    3: fem.Constant(msh, PETSc.ScalarType(4.0)),
}

# Define variational problem
F = 0
F += (
    dot(grad(v), grad(u)) * dx
    - dot(v * n, grad(u)) * ds
    - dot(avg(grad(v)), jump(u, n)) * dS(0)
    - dot(jump(v, n), avg(grad(u))) * dS(0)
    + alpha / avg(h) * dot(jump(v, n), jump(u, n)) * dS(0)
    + alpha / h * v * u * ds
)

# source
F += -v * f * dx

# Dirichlet BC
F += -dot(grad(v), u * n) * ds + uD * dot(grad(v), n) * ds - alpha / h * uD * v * ds

for interface_id, interface_facets in zip([2, 3, 4], [interface_facets_1, interface_facets_2, interface_facets_3]):
    # look at the first facet on interface
    # and get the two cells that are connected to it
    # and get the material properties of these cells
    first_facet_interface = interface_facets[0]
    c_plus, c_minus = f_to_c.links(first_facet_interface)[0], f_to_c.links(first_facet_interface)[1]
    id_minus, id_plus = values_cells[c_plus], values_cells[c_minus]
    # TODO ensure that id_minus and id_plus are the same for all facets on the interface

    K1 = cell_tag_to_K[id_minus]
    K2 = cell_tag_to_K[id_plus]

    F += -dot(avg(grad(v)), n("-")) * (u("-") * (K1 / K2 - 1)) * dS(interface_id)
    F += (
        alpha
        / avg(h)
        * dot(jump(v, n), n("-"))
        * (u("-") * (K1 / K2 - 1))
        * dS(interface_id)
    )
    # symmetry
    F += -dot(avg(grad(v)), jump(u, n)) * dS(interface_id)

    # coercivity
    F += +alpha / avg(h) * dot(jump(v, n), jump(u, n)) * dS(interface_id)

problem = dolfinx.fem.petsc.NonlinearProblem(F, u)
solver = dolfinx.nls.petsc.NewtonSolver(MPI.COMM_WORLD, problem)
solver.solve(u)

[RemDelaporteMathurin]

Update for dolfinx 0.8.0

from dolfinx import mesh, fem, plot, geometry
import dolfinx.fem.petsc
import dolfinx.nls.petsc
from mpi4py import MPI
import ufl
from petsc4py import PETSc
from ufl import grad, dot, jump, avg
import numpy as np
import matplotlib.pyplot as plt
import pyvista

msh = mesh.create_unit_square(MPI.COMM_WORLD, 100, 100)
V = fem.functionspace(msh, ("DG", 1))
uD = fem.Function(V)
uD.interpolate(lambda x: np.full(x[0].shape, 0.0))

# create mesh tags
tol = 1e-16

def marker_interface_1(x):
    return np.logical_and(np.isclose(x[0], 0.5), x[1] > 0.5 - tol)

def marker_interface_2(x):
    return np.logical_and(np.isclose(x[0], 0.5), x[1] < 0.5 + tol)

def marker_interface_3(x):
    return np.logical_and(x[0] > 0.5 - tol, np.isclose(x[1], 0.5))

tdim = msh.topology.dim
fdim = tdim - 1
msh.topology.create_connectivity(fdim, tdim)
c_to_f = msh.topology.connectivity(tdim, tdim - 1)
f_to_c = msh.topology.connectivity(tdim - 1, tdim)
facet_map = msh.topology.index_map(tdim - 1)
cell_imap = msh.topology.index_map(tdim)
boundary_facets = mesh.exterior_facet_indices(msh.topology)
interface_facets_1 = mesh.locate_entities(msh, fdim, marker_interface_1)
interface_facets_2 = mesh.locate_entities(msh, fdim, marker_interface_2)
interface_facets_3 = mesh.locate_entities(msh, fdim, marker_interface_3)
num_facets = facet_map.size_local + facet_map.num_ghosts
indices = np.arange(0, num_facets)
# values = np.arange(0, num_facets, dtype=np.intc)
values = np.zeros(indices.shape, dtype=np.intc)  # all facets are tagged with zero

values[boundary_facets] = 1
values[interface_facets_1] = 2
values[interface_facets_2] = 3
values[interface_facets_3] = 4

mesh_tags_facets = mesh.meshtags(msh, fdim, indices, values)

# tag cells
def left_domain(x):
    tol = 1e-16
    return x[0] < 0.5 + tol

def top_right_domain(x):
    tol = 1e-16
    return np.logical_and(x[0] > 0.5 - tol, x[1] > 0.5 - tol)

def bottom_right_domain(x):
    tol = 1e-16
    return np.logical_and(x[0] > 0.5 - tol, x[1] < 0.5 + tol)

left_cells = mesh.locate_entities(msh, tdim, left_domain)
top_right_cells = mesh.locate_entities(msh, tdim, top_right_domain)
bottom_right_cells = mesh.locate_entities(msh, tdim, bottom_right_domain)
num_cells = cell_imap.size_local
indices = np.arange(0, num_cells)
values_cells = np.zeros(indices.shape, dtype=np.intc)
values_cells[left_cells] = 1
values_cells[top_right_cells] = 2
values_cells[bottom_right_cells] = 3

mesh_tags_cells = mesh.meshtags(msh, tdim, indices, values_cells)

# Create measures
ds = ufl.Measure("ds", domain=msh, subdomain_data=mesh_tags_facets)
dS = ufl.Measure("dS", domain=msh, subdomain_data=mesh_tags_facets)
dx = ufl.Measure("dx", domain=msh, subdomain_data=mesh_tags_cells)

u = fem.Function(V)
u_n = fem.Function(V)
v = ufl.TestFunction(V)

h = ufl.CellDiameter(msh)
n = ufl.FacetNormal(msh)

# Define parameters
alpha = 1000

# Simulation constants
f = fem.Constant(msh, PETSc.ScalarType(2.0))

cell_tag_to_K = {
    1: fem.Constant(msh, PETSc.ScalarType(2.0)),
    2: fem.Constant(msh, PETSc.ScalarType(8.0)),
    3: fem.Constant(msh, PETSc.ScalarType(4.0)),
}

# Define variational problem
F = 0
F += (
    dot(grad(v), grad(u)) * dx
    - dot(v * n, grad(u)) * ds
    - dot(avg(grad(v)), jump(u, n)) * dS(0)
    - dot(jump(v, n), avg(grad(u))) * dS(0)
    + alpha / avg(h) * dot(jump(v, n), jump(u, n)) * dS(0)
    + alpha / h * v * u * ds
)

# source
F += -v * f * dx

# Dirichlet BC
F += -dot(grad(v), u * n) * ds + uD * dot(grad(v), n) * ds - alpha / h * uD * v * ds

for interface_id, interface_facets in zip(
    [2, 3, 4], [interface_facets_1, interface_facets_2, interface_facets_3]
):
    # look at the first facet on interface
    # and get the two cells that are connected to it
    # and get the material properties of these cells
    first_facet_interface = interface_facets[0]
    c_plus, c_minus = (
        f_to_c.links(first_facet_interface)[0],
        f_to_c.links(first_facet_interface)[1],
    )
    id_minus, id_plus = values_cells[c_plus], values_cells[c_minus]
    # TODO ensure that id_minus and id_plus are the same for all facets on the interface

    K1 = cell_tag_to_K[id_minus]
    K2 = cell_tag_to_K[id_plus]

    F += -dot(avg(grad(v)), n("-")) * (u("-") * (K1 / K2 - 1)) * dS(interface_id)
    F += (
        alpha
        / avg(h)
        * dot(jump(v, n), n("-"))
        * (u("-") * (K1 / K2 - 1))
        * dS(interface_id)
    )
    # symmetry
    F += -dot(avg(grad(v)), jump(u, n)) * dS(interface_id)

    # coercivity
    F += +alpha / avg(h) * dot(jump(v, n), jump(u, n)) * dS(interface_id)

problem = dolfinx.fem.petsc.NonlinearProblem(F, u)
solver = dolfinx.nls.petsc.NewtonSolver(MPI.COMM_WORLD, problem)
solver.solve(u)

pyvista.OFF_SCREEN = True

pyvista.start_xvfb()

u_topology, u_cell_types, u_geometry = plot.vtk_mesh(V)
u_grid = pyvista.UnstructuredGrid(u_topology, u_cell_types, u_geometry)
u_grid.point_data["u"] = u.x.array.real
u_grid.set_active_scalars("u")
u_plotter = pyvista.Plotter()
u_plotter.add_mesh(u_grid, show_edges=False)
u_plotter.view_xy()
if not pyvista.OFF_SCREEN:
    u_plotter.show()
else:
    figure = u_plotter.screenshot("DG.png")

[RemDelaporteMathurin]

Thanks to @dokken and @jpdean we now have a prototype for interfaces!

# docker run -ti -v $(pwd):/root/shared -w /root/shared jpdean/mixed_domain

# 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

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 = 10.0

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

K_b = K_0(b_res)
K_t = K_1(t_res)

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()

This runs on dolfinx 0.8

[jorgensd]

@jpdean has a «simplified»/slightly different version of this for problem (But assuming linear problem) at https://github.com/jpdean/mixed_domain_demos/blob/main/poisson_domain_decomp.py It is slightly out of date, but illustrates the same strategy, just for problems that one want to solve without a newton solver.

[RemDelaporteMathurin]

I think our problems will all be non-linear, I'm happy to help streamline the newton solver integration into dolfinx if I can! 😄

[RemDelaporteMathurin]

Next step is to make this work only on a component of a vector function space:

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)

    degree = 1
    element_CG = basix.ufl.element(
        basix.ElementFamily.P,
        submesh.basix_cell(),
        degree,
        basix.LagrangeVariant.equispaced,
    )
    element = basix.ufl.mixed_element([element_CG, element_CG])
    V = dolfinx.fem.functionspace(submesh, element)
    u = dolfinx.fem.Function(V)
    us = list(ufl.split(u))
    vs = list(ufl.TestFunctions(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(us[0]), ufl.grad(vs[0])) * dx_r - val * vs[0] * dx_r
    k = 0.05
    p = 0.1
    n = 1.0
    F += k * us[0] * (n - us[1]) * vs[1] * dx_r - p * us[1] * vs[1] * dx_r
    return u, vs, F, mesh_to_submesh

u_0, v_0s, F_00, m_to_b = define_interior_eq(mesh, 2, submesh_b, submesh_b_to_mesh, 0.0)
u_1, v_1s, 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 = v_0s[0](b_res)
v_t = v_1s[0](t_res)

u_bs = list(ufl.split(u_0))
u_ts = list(ufl.split(u_1))
u_b = u_bs[0](b_res)
u_t = u_ts[0](t_res)

But this doesn't work

[0]PETSC ERROR: ------------------------------------------------------------------------
[0]PETSC ERROR: Caught signal number 11 SEGV: Segmentation Violation, probably memory access out of range
[0]PETSC ERROR: Try option -start_in_debugger or -on_error_attach_debugger
[0]PETSC ERROR: or see https://petsc.org/release/faq/#valgrind and https://petsc.org/release/faq/
[0]PETSC ERROR: configure using --with-debugging=yes, recompile, link, and run 
[0]PETSC ERROR: to get more information on the crash.
[0]PETSC ERROR: Run with -malloc_debug to check if memory corruption is causing the crash.
Abort(59) on node 0 (rank 0 in comm 0): application called MPI_Abort(MPI_COMM_WORLD, 59) - process 0

[jorgensd]

Could you try with the nightly docker image?

[jorgensd]

The following code does not segfault on the nightly branch:

from mpi4py import MPI
import dolfinx
import dolfinx.fem.petsc
import ufl
import numpy as np
from petsc4py import PETSc
import basix

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: parent_to_sub_b,
    submesh_t: 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)

    degree = 1
    element_CG = basix.ufl.element(
        basix.ElementFamily.P,
        submesh.basix_cell(),
        degree,
        basix.LagrangeVariant.equispaced,
    )
    element = basix.ufl.mixed_element([element_CG, element_CG])
    V = dolfinx.fem.functionspace(submesh, element)
    u = dolfinx.fem.Function(V)
    us = list(ufl.split(u))
    vs = list(ufl.TestFunctions(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(us[0]), ufl.grad(vs[0])) * dx_r - val * vs[0] * dx_r
    k = 0.05
    p = 0.1
    n = 1.0
    F += k * us[0] * (n - us[1]) * vs[1] * dx_r - p * us[1] * vs[1] * dx_r
    return u, vs, F, mesh_to_submesh

u_0, v_0s, F_00, m_to_b = define_interior_eq(mesh, 2, submesh_b, submesh_b_to_mesh, 0.0)
u_1, v_1s, 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 = v_0s[0](b_res)
v_t = v_1s[0](t_res)

u_bs = list(ufl.split(u_0))
u_ts = list(ufl.split(u_1))
u_b = u_bs[0](b_res)
u_t = u_ts[0](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 = 10.0

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

K_b = K_0(b_res)
K_t = K_1(t_res)

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=10,
    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.sub(0).collapse()], engine="BP4")
bp.write(0)
bp.close()
bp = dolfinx.io.VTXWriter(mesh.comm, "u_t.bp", [u_1.sub(0).collapse()], engine="BP4")
bp.write(0)
bp.close()

I'm not sure this is what you had in mind. I think it could be related to: https://github.com/FEniCS/dolfinx/pull/3260 which is the last place I got segfaults for submesh assembly.

[RemDelaporteMathurin]

@jorgensd your code runs without segfaults on the nightly docker image.

I fixed a few things (forgot to set the Dirichlet BC on the first component only) and it now produces the expected result:

Left is mobile particle concentration and right is trapped concentration.

The code:

from mpi4py import MPI
import dolfinx
import dolfinx.fem.petsc
import ufl
import numpy as np
from petsc4py import PETSc
import basix

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, 20, 20, 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: parent_to_sub_b,
    submesh_t: 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)

    degree = 1
    element_CG = basix.ufl.element(
        basix.ElementFamily.P,
        submesh.basix_cell(),
        degree,
        basix.LagrangeVariant.equispaced,
    )
    element = basix.ufl.mixed_element([element_CG, element_CG])
    V = dolfinx.fem.functionspace(submesh, element)
    u = dolfinx.fem.Function(V)
    us = list(ufl.split(u))
    vs = list(ufl.TestFunctions(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(us[0]), ufl.grad(vs[0])) * dx_r - val * vs[0] * dx_r
    k = 2
    p = 0.1
    n = 0.5
    F += k * us[0] * (n - us[1]) * vs[1] * dx_r - p * us[1] * vs[1] * dx_r
    return u, vs, F, mesh_to_submesh

u_0, v_0s, F_00, m_to_b = define_interior_eq(mesh, 2, submesh_b, submesh_b_to_mesh, 0.0)
u_1, v_1s, 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 = v_0s[0](b_res)
v_t = v_1s[0](t_res)

u_bs = list(ufl.split(u_0))
u_ts = list(ufl.split(u_1))
u_b = u_bs[0](b_res)
u_t = u_ts[0](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 = 10.0

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

K_b = K_0(b_res)
K_t = K_1(t_res)

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.sub(0), 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.sub(0), fdim, ft_t.find(1))
)
bcs = [bc_b, bc_t]

solver = NewtonSolver(
    F,
    J,
    [u_0, u_1],
    bcs=bcs,
    max_iterations=10,
    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.sub(0).collapse()], engine="BP4")
bp = dolfinx.io.VTXWriter(mesh.comm, "u_b_0.bp", [u_0.sub(0).collapse()], engine="BP4")
bp.write(0)
bp.close()
bp = dolfinx.io.VTXWriter(mesh.comm, "u_t_0.bp", [u_1.sub(0).collapse()], engine="BP4")
bp.write(0)
bp.close()
bp = dolfinx.io.VTXWriter(mesh.comm, "u_b_1.bp", [u_0.sub(1).collapse()], engine="BP4")
bp.write(0)
bp.close()
bp = dolfinx.io.VTXWriter(mesh.comm, "u_t_1.bp", [u_1.sub(1).collapse()], engine="BP4")
bp.write(0)
bp.close()