Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Interface discontinuity #719

Open
RemDelaporteMathurin opened this issue Mar 26, 2024 · 13 comments
Open

Interface discontinuity #719

RemDelaporteMathurin opened this issue Mar 26, 2024 · 13 comments
Labels
fenicsx Issue that is related to the fenicsx support help wanted Extra attention is needed

Comments

@RemDelaporteMathurin
Copy link
Collaborator

RemDelaporteMathurin commented Mar 26, 2024

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!

@RemDelaporteMathurin RemDelaporteMathurin added help wanted Extra attention is needed fenicsx Issue that is related to the fenicsx support labels Mar 26, 2024
@jpdean
Copy link

jpdean commented Mar 26, 2024

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
Copy link
Collaborator Author

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
Copy link

jpdean commented Mar 26, 2024

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
Copy link
Collaborator Author

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
Copy link
Collaborator Author

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

image

@RemDelaporteMathurin
Copy link
Collaborator Author

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
Copy link
Collaborator Author

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
Copy link

@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
Copy link
Collaborator Author

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
Copy link
Collaborator Author

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
Copy link

Could you try with the nightly docker image?

@jorgensd
Copy link

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: FEniCS/dolfinx#3260
which is the last place I got segfaults for submesh assembly.

@RemDelaporteMathurin
Copy link
Collaborator Author

@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:

image

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
fenicsx Issue that is related to the fenicsx support help wanted Extra attention is needed
Projects
Status: In Progress
Development

No branches or pull requests

3 participants