---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.19.6
  main_language: python
---

# Mixed Poisson on a branching surface

This demo solves the Poisson equation in mixed form on three sheets
meeting along a line, a surface that branches. It illustrates how to

* give each branch its own $H(\mathrm{div})$-conforming space, using
  submeshes,
* couple the branches through a Lagrange multiplier on the joint, a
  mesh of lower dimension,
* tag cells and facets while creating a mesh on one process, and
  transfer the facet tags to a submesh,
* integrate over the facets of chosen cells, from the side of each
  cell,
* interpolate functions from the submeshes into a single function on
  the parent mesh, for visualisation.

```{admonition} Download sources
:class: download
* {download}`Python script <./demo_branching_manifold.py>`
* {download}`Jupyter notebook <./demo_branching_manifold.ipynb>`
```

## Equation and problem definition

The surface $\Omega$ consists of three branches $\Omega_b$,
$b = 0, 1, 2$, joined along a curve $\Gamma$. On each branch the
Poisson equation in mixed form reads

$$
\begin{aligned}
  \sigma_b - \nabla_\Omega u_b &= 0, \\
  \nabla_\Omega \cdot \sigma_b &= g,
\end{aligned}
$$

with $u_b = 0$ on the outer boundary. At the joint, the potential is
single-valued and the flux is conserved:

$$
\begin{aligned}
  u_b &= u_\Gamma \quad \text{on } \Gamma \text{ for every } b, \\
  \sum_b \sigma_b \cdot \nu_b &= 0 \quad \text{on } \Gamma,
\end{aligned}
$$

where $\nu_b$ is the outward unit normal to $\Gamma$ that lies in
$\Omega_b$, and $u_\Gamma$ is the common value of the potential there.
The second condition is Kirchhoff's law: whatever flows out of one
branch flows into the others. These are the interface conditions of
{cite:t}`branching-hansbo2017composite`, who solve the primal problem on
such composite surfaces with a Nitsche method.

### Discretization

Each branch therefore gets its own Raviart-Thomas space, conforming
up to the joint, and the conditions at the joint are imposed weakly
with $u_\Gamma$ as a Lagrange multiplier. This is the partially
hybridized mixed method {cite}`branching-arnold1985hybrid`: $u_\Gamma$
lives in the normal-trace space of the Raviart-Thomas elements on
$\Gamma$. {cite:t}`branching-arbogast2000mortar` use it to couple
blocks that meet in pairs; here three meet on one interface.

```{admonition} Networks of curves
:class: note dropdown
On networks of curves, where the joint is a point, the same
construction couples the segments at each bifurcation
{cite}`branching-gjerde2022graphnics`, and
{cite:t}`branching-dokken2025networks` implements it in DOLFINx in the
same way as here. A point has no extent, so there any flux has a
scalar trace; on a curve the multiplier must instead lie in the
normal-trace space, on meshes that match along $\Gamma$.
```

The unknowns live in three kinds of space:

* the potential $u \in V(\Omega)$, a discontinuous Lagrange space on
  the whole surface, since $u$ needs no continuity, not even across
  $\Gamma$,
* the flux $\sigma_b \in Q_b = Q(\Omega_b)$ on each branch, a
  Raviart-Thomas space,
* the joint potential $u_\Gamma \in Z(\Gamma)$, a discontinuous
  Lagrange space on the joint.

```{admonition} Why one space is not enough
:class: note dropdown
The degrees-of-freedom of a Raviart-Thomas space on a facet are shared
by the cells of that facet and fix the normal flux across it. On a
surface without branches each edge has two cells, and
{doc}`demo_mobius` shows how signing the facets makes that flux
continuous. An edge of $\Gamma$, however, has three cells. Their
three fluxes satisfy one condition, Kirchhoff's law, which leaves two
independent fluxes, more than one set of facet degrees-of-freedom can
carry. No choice of signs helps:
{py:meth}`create_facet_orientations
<dolfinx.mesh.Topology.create_facet_orientations>` refuses such a mesh.
```

Testing the first equation with $\tau_b \in Q_b$ and integrating by
parts on each branch moves the gradient from $u_b$ onto $\tau_b$, and
leaves the boundary term
$\int_\Gamma u_b \, \tau_b \cdot \nu_b ~\mathrm{d}s$, in which $u_b$
is replaced by $u_\Gamma$. The weak form is: find $u \in V(\Omega)$,
$\sigma_b \in Q_b$ and $u_\Gamma \in Z(\Gamma)$ such that

$$
\begin{aligned}
  \int_{\Omega_b} \sigma_b \cdot \tau_b ~\mathrm{d}x
    + \int_{\Omega_b} u \, \nabla_\Omega \cdot \tau_b ~\mathrm{d}x
    - \int_\Gamma u_\Gamma \, \tau_b \cdot \nu_b ~\mathrm{d}s
    &= 0 && \forall \tau_b \in Q_b, \quad b = 0, 1, 2, \\
  \sum_b \int_{\Omega_b} (\nabla_\Omega \cdot \sigma_b) \, v
    ~\mathrm{d}x &= \int_\Omega g \, v ~\mathrm{d}x
    && \forall v \in V(\Omega), \\
  -\sum_b \int_\Gamma \sigma_b \cdot \nu_b \, v_\Gamma ~\mathrm{d}s
    &= 0 && \forall v_\Gamma \in Z(\Gamma).
\end{aligned}
$$

The last equation is Kirchhoff's law, tested against $v_\Gamma$.

### Method of manufactured solutions

Each branch is a unit square, rotated about the $z$-axis by
$2\pi b / 3$, with $\Gamma$ the segment $0 \le z \le 1$ of that axis.
With $r$ the distance from the axis, take

$$
  u_b = (1 - r)(1 + a_b r) \sin(\pi z).
$$

It vanishes on the outer boundary, and on $\Gamma$, where $r = 0$, it
equals $\sin(\pi z)$ on every branch. The flux out of branch $b$ into
the joint is $\sigma_b \cdot \nu_b = -\partial_r u_b = (1 - a_b)
\sin(\pi z)$, so Kirchhoff's law holds when $\sum_b a_b = 3$. We take
$a = (3, 1, -1)$: branch $0$ drains the joint, branch $2$ feeds it,
and nothing crosses into branch $1$. The branches are flat, so
$g = \Delta u_b$ is the usual Laplacian in $(r, z)$.

## Implementation

```python
from mpi4py import MPI

import numpy as np

import basix.ufl
import dolfinx
import dolfinx.fem.petsc
import ufl
```

The branches are at angles `angles`, the exact solution has slopes
`slopes`, `order` is the degree of the Raviart-Thomas spaces, and
`joint_marker` tags the edges of the joint.

```python
angles = 2 * np.pi * np.arange(3) / 3
slopes = (3.0, 1.0, -1.0)
order = 2
joint_marker = 3
```

### Mesh-creation

The surface is built by hand: each branch is a grid of `n` by `m`
squares, cut into triangles, with the column of vertices on the axis
shared by all three. Every edge on the axis then has three cells,
which `max_facet_to_cell_links=3` allows.

Each input cell is marked with its branch, and each edge of the joint
with `joint_marker`. Like the cells themselves, the markers exist only
on the process that creates the mesh, so they are sent to the
processes that hold the entities, owned or ghost, with
{py:func}`distribute_entity_data <dolfinx.io.distribute_entity_data>`
and turned into {py:class}`MeshTags <dolfinx.mesh.MeshTags>`. An edge
is given by its two vertices, in the numbering of the input.

```python
def branching_mesh(n, m):
    """Three unit squares joined along the z-axis, tagged by branch."""
    if MPI.COMM_WORLD.rank == 0:
        z = np.linspace(0, 1, m + 1)
        r = np.linspace(0, 1, n + 1)
        points = [np.column_stack([np.zeros_like(z), np.zeros_like(z), z])]
        cells = []
        for b, theta in enumerate(angles):
            rr, zz = np.meshgrid(r[1:], z, indexing="ij")
            rr, zz = rr.ravel(), zz.ravel()
            points.append(np.column_stack([rr * np.cos(theta), rr * np.sin(theta), zz]))

            # Column 0 is the axis, numbered first and shared
            offset = (m + 1) * (1 + b * n)

            def node(i, j, offset=offset):
                return j if i == 0 else offset + (i - 1) * (m + 1) + j

            for i in range(n):
                for j in range(m):
                    cells.append([node(i, j), node(i + 1, j), node(i, j + 1)])
                    cells.append([node(i + 1, j), node(i + 1, j + 1), node(i, j + 1)])
        x = np.vstack(points)
        cells = np.array(cells, dtype=np.int64)
        markers = np.repeat(np.arange(3, dtype=np.int32), 2 * n * m)
        # The joint edges join consecutive vertices on the axis
        joint_edges = np.column_stack([np.arange(m), np.arange(1, m + 1)]).astype(np.int64)
    else:
        x = np.zeros((0, 3))
        cells = np.zeros((0, 3), dtype=np.int64)
        markers = np.zeros(0, dtype=np.int32)
        joint_edges = np.zeros((0, 2), dtype=np.int64)
    domain = ufl.Mesh(
        basix.ufl.element("Lagrange", "triangle", 1, shape=(3,), dtype=dolfinx.default_real_type)
    )
    msh = dolfinx.mesh.create_mesh(
        MPI.COMM_WORLD,
        cells,
        domain,
        x.astype(dolfinx.default_real_type),
        max_facet_to_cell_links=3,
    )
    tdim = msh.topology.dim
    local_cells, local_markers = dolfinx.io.distribute_entity_data(msh, tdim, cells, markers)
    cell_tags = dolfinx.mesh.meshtags_from_entities(
        msh, tdim, dolfinx.graph.adjacencylist(local_cells), local_markers
    )
    local_edges, local_values = dolfinx.io.distribute_entity_data(
        msh, tdim - 1, joint_edges, np.full(len(joint_edges), joint_marker, dtype=np.int32)
    )
    msh.topology.create_connectivity(tdim - 1, tdim)
    facet_tags = dolfinx.mesh.meshtags_from_entities(
        msh, tdim - 1, dolfinx.graph.adjacencylist(local_edges), local_values
    )
    return msh, cell_tags, facet_tags
```

### Submesh creation of the joint and branches

Each branch is extracted from its cell tags as a submesh with
{py:func}`create_submesh <dolfinx.mesh.create_submesh>`. A branch is
an ordinary surface with a boundary, so its facets can be signed. The
joint is a submesh of the tagged edges.
Each submesh comes with an {py:class}`EntityMap
<dolfinx.mesh.EntityMap>` relating its entities to those of the whole
surface, which assembly needs to combine them.

The coupling terms are integrals over $\Gamma$ seen from one branch at
a time, so that $\nu_b$ is the normal of that branch. They are
integrals over cell facets, identified as `(cell, local facet)` pairs,
one integration domain per branch. On a branch the joint is part of
the boundary, so the facet tags are moved to the branch with
{py:func}`transfer_meshtags_to_submesh
<dolfinx.mesh.transfer_meshtags_to_submesh>` and the pairs found with
{py:func}`compute_integration_domains
<dolfinx.fem.compute_integration_domains>`, as for any boundary. The
integrals are taken over the whole surface, which the joint mesh is
related to, so the cells are then mapped back to it with the entity
map of the branch.

```python
def solve(n, m):
    """Solve on the surface with ``n`` by ``m`` squares per branch."""
    msh, cell_tags, facet_tags = branching_mesh(n, m)
    tdim = msh.topology.dim

    branches, entity_maps, joint_facets = [], [], []
    for b in range(3):
        sub, sub_map, vertex_map = dolfinx.mesh.create_submesh(msh, tdim, cell_tags.find(b))[:3]
        # Required for a Raviart-Thomas space on a surface. Here it signs
        # no facets, as the cells of each branch share an orientation.
        sub.topology.create_facet_orientations()
        branches.append(sub)
        entity_maps.append(sub_map)

        # The joint is part of the boundary of the branch: find its
        # (cell, local facet) pairs there, and map the cells back to the
        # whole surface, where the local facet index is the same
        sub_tags = dolfinx.mesh.transfer_meshtags_to_submesh(facet_tags, sub, sub_map, vertex_map)
        sub.topology.create_connectivity(tdim - 1, tdim)
        pairs = dolfinx.fem.compute_integration_domains(
            dolfinx.fem.IntegralType.exterior_facet, sub.topology, sub_tags.find(joint_marker)
        ).reshape(-1, 2)
        pairs[:, 0] = sub_map.sub_topology_to_topology(pairs[:, 0], inverse=False)
        joint_facets.append((b, pairs.reshape(-1)))

    on_axis = facet_tags.find(joint_marker)
    joint, joint_map = dolfinx.mesh.create_submesh(msh, tdim - 1, on_axis)[:2]
    entity_maps.append(joint_map)

    # Spaces, with one Raviart-Thomas space per branch
    Q = [dolfinx.fem.functionspace(sub, ("RT", order)) for sub in branches]
    V = dolfinx.fem.functionspace(msh, ("DG", order - 1))
    Z = dolfinx.fem.functionspace(joint, ("DG", order - 1))
    W = ufl.MixedFunctionSpace(*Q, V, Z)
    *sigma, u, u_Gamma = ufl.TrialFunctions(W)
    *tau, v, v_Gamma = ufl.TestFunctions(W)

    # dx(b) is branch b, ds(b) the joint seen from branch b, whose
    # outward normal there is nu
    dx = ufl.Measure("dx", domain=msh, subdomain_data=cell_tags)
    ds = ufl.Measure("ds", domain=msh, subdomain_data=joint_facets)
    nu = ufl.FacetNormal(msh)

    x = ufl.SpatialCoordinate(msh)

    def u_exact(b):
        """The exact solution on branch ``b``."""
        r = x[0] * np.cos(angles[b]) + x[1] * np.sin(angles[b])
        return (1 - r) * (1 + slopes[b] * r) * ufl.sin(ufl.pi * x[2])

    a = 0
    for b in range(3):
        a += ufl.inner(sigma[b], tau[b]) * dx(b)
        a += ufl.inner(u, ufl.div(tau[b])) * dx(b) + ufl.inner(ufl.div(sigma[b]), v) * dx(b)
        a -= ufl.inner(u_Gamma, ufl.dot(tau[b], nu)) * ds(b)
        a -= ufl.inner(ufl.dot(sigma[b], nu), v_Gamma) * ds(b)
    g = [ufl.div(ufl.grad(u_exact(b))) for b in range(3)]
    L = ufl.inner(g[0], v) * dx(0) + ufl.inner(g[1], v) * dx(1) + ufl.inner(g[2], v) * dx(2)

    # Only the equation tested with v has a source. The others get an
    # explicit zero, so that the right-hand side has a block per space.
    a_blocked: list[list[ufl.Form | None]] = ufl.extract_blocks(a)  # type: ignore[assignment]
    L_blocked: list[ufl.Form | ufl.ZeroBaseForm] = [
        *(ufl.ZeroBaseForm((t,)) for t in tau),
        L,
        ufl.ZeroBaseForm((v_Gamma,)),
    ]

    problem = dolfinx.fem.petsc.LinearProblem(
        a_blocked,
        L_blocked,
        entity_maps=entity_maps,
        petsc_options={
            "ksp_type": "preonly",
            "pc_type": "lu",
            "pc_factor_mat_solver_type": "mumps",
            "ksp_error_if_not_converged": True,
        },
        petsc_options_prefix="branching_",
        kind="mpi",
    )
    *sigma_h, u_h, u_Gamma_h = problem.solve()

    def norm(expr, measure):
        """The L2 norm of ``expr`` over ``measure``."""
        # UFL estimates the degree of sin as that of its argument plus
        # two, too low to measure the error of a high-degree space: at
        # the few points it gives, the multiplier superconverges.
        integral = dolfinx.fem.form(
            ufl.inner(expr, expr) * measure(metadata={"quadrature_degree": 2 * order + 6}),
            entity_maps=entity_maps,
        )
        return np.sqrt(abs(msh.comm.allreduce(dolfinx.fem.assemble_scalar(integral))))

    x_joint = ufl.SpatialCoordinate(joint)
    errors = {
        "L2(u)": np.hypot.reduce([norm(u_h - u_exact(b), dx(b)) for b in range(3)]),
        "L2(sigma)": np.hypot.reduce(
            [norm(sigma_h[b] - ufl.grad(u_exact(b)), dx(b)) for b in range(3)]
        ),
        "L2(u_Gamma)": norm(u_Gamma_h - ufl.sin(ufl.pi * x_joint[2]), ufl.dx(domain=joint)),
    }
    return msh, entity_maps, sigma_h, u_h, errors
```

### Results

The errors converge at the optimal rate, $O(h^k)$ for Raviart-Thomas
of degree $k$, in the flux, the potential and the multiplier alike.
Each error is followed by its rate, computed from the previous mesh.

```python
resolutions = [4, 8, 16]
columns = ["L2(u)", "L2(sigma)", "L2(u_Gamma)"]
results = [solve(n, n) for n in resolutions]
if MPI.COMM_WORLD.rank == 0:
    print(f"{'n':>4}" + "".join(f"{c:>14}{'rate':>6}" for c in columns))
    for i, (n, (*_, errors)) in enumerate(zip(resolutions, results, strict=True)):
        row = f"{n:>4}"
        for c in columns:
            if i == 0:
                rate = ""
            else:
                ratio = results[i - 1][-1][c] / errors[c]
                rate = f"{np.log(ratio) / np.log(n / resolutions[i - 1]):.2f}"
            row += f"{errors[c]:>14.3e}{rate:>6}"
        print(row)
```

### Visualisation: interpolation from submeshes to the parent mesh

The fluxes live on the branches. To view them together they are
interpolated into one discontinuous Lagrange space on the whole
surface, branch by branch, mapping each branch's cells to those of
the surface with its entity map.

```python
if dolfinx.has_adios2:
    msh, entity_maps, sigma_h, u_h, _ = results[-1]
    S = dolfinx.fem.functionspace(msh, ("DG", order - 1, (3,)))
    sigma_out = dolfinx.fem.Function(S, name="sigma")
    for sigma_b, sub_map in zip(sigma_h, entity_maps[:3], strict=True):
        sub_cell_map = sigma_b.function_space.mesh.topology.index_map(2)
        sub_cells = np.arange(sub_cell_map.size_local + sub_cell_map.num_ghosts, dtype=np.int32)
        sigma_out.interpolate(
            sigma_b,
            cells0=sub_cells,
            cells1=sub_map.sub_topology_to_topology(sub_cells, inverse=False),
        )
    with dolfinx.io.VTXWriter(msh.comm, "branching_manifold.bp", [sigma_out, u_h]) as writer:
        writer.write(0.0)
```

## References
```{bibliography}
   :filter: cited
   :labelprefix:
   :keyprefix: branching-
```