Consistent orientation of bifurcation integrals

script.py

# Consistently orient jump integrals on a mesh with a bifurcation
#
# The idea of this demo is to say that the jump at a bifurcation with three cells should be:
# u(0) - u(1) - u(2) where u(0) is taken from the cell with the lowest value, while u(1) and u(2) are taken from the two other cells.
#
#  \         /
#   \ u1    / u2
#    \     /
#     \   /
#      \ /
#       |
#       | u0
#       |
#
# Author : Jørgen S. Dokken
# SPDX-License-Identifier: MIT


from mpi4py import MPI
import dolfinx
import ufl
import numpy as np
from pathlib import Path
import numpy.typing as npt


def unroll_dofmap(dofs: npt.NDArray[np.int32], bs: int) -> npt.NDArray[np.int32]:
    """
    Given a two-dimensional dofmap of size `(num_cells, num_dofs_per_cell)`
    Expand the dofmap by its block size such that the resulting array
    is of size `(num_cells, bs*num_dofs_per_cell)`
    """
    num_cells, num_dofs_per_cell = dofs.shape
    unrolled_dofmap = np.repeat(dofs, bs).reshape(
        num_cells, num_dofs_per_cell * bs) * bs
    unrolled_dofmap += np.tile(np.arange(bs), num_dofs_per_cell)
    return unrolled_dofmap


# Read mesh from file
filename = Path("mesh_double_Y_bifurcation/mesh.xdmf")
with dolfinx.io.XDMFFile(MPI.COMM_WORLD, filename, "r") as xdmf:
    mesh = xdmf.read_mesh(name="mesh")
    ct = xdmf.read_meshtags(mesh, name="Cell tags")


# Create connectivities that should be used to go from `facet_index` to `(cell_index, local_facet_index)`
mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim - 1)
f_to_c = mesh.topology.connectivity(mesh.topology.dim - 1, mesh.topology.dim)
c_to_f = mesh.topology.connectivity(mesh.topology.dim, mesh.topology.dim - 1)
num_facets = mesh.topology.index_map(mesh.topology.dim - 1).size_local

# Compute bifurcation facets and their cell, local_facet_index representation
bifurcation_facets = []
facet_as_integration_entities = []
for facet in range(num_facets):
    cells = f_to_c.links(facet)
    if len(cells) > 2:
        assert len(cells) == 3
        bifurcation_facets.append(facet)
        local_list = []
        for cell in cells:
            facets = c_to_f.links(cell)
            local_pos = np.flatnonzero(facets == facet)[0]
            local_list.extend([cell, local_pos])
        facet_as_integration_entities.append(local_list)
bifurcation_facets = np.array(bifurcation_facets)
facet_as_integration_entities = np.array(facet_as_integration_entities)

# Get the cell marker for the cells
marked_cells = facet_as_integration_entities[:, ::2]
marked_values = ct.values[marked_cells]
sorted_order_mv = np.argsort(marked_values, axis=1)

# Make tuples (low, high_0), (low, high_1)
# Unroll for integration entity tuple
sorted_unrolled = unroll_dofmap(sorted_order_mv, 2)
stacked_unrolled = np.hstack([sorted_unrolled, sorted_unrolled[:, :2]])[
    :, [0, 1, 2, 3, 6, 7, 4, 5]]
sorted_bf = np.take_along_axis(
    facet_as_integration_entities, stacked_unrolled, axis=1).flatten()

# Create verification test
V = dolfinx.fem.functionspace(mesh, ("DG", 0))
u = dolfinx.fem.Function(V)
u.x.array[:] = ct.values

s_id = 72
dB = ufl.Measure("dS", domain=mesh, subdomain_data=[(s_id, sorted_bf)])

J = (0.5*u("+")-u("-")) * dB(s_id)

val = dolfinx.fem.assemble_scalar(dolfinx.fem.form(J))

sorted_values = np.take_along_axis(marked_values, sorted_order_mv, axis=1)
reference = np.sum(
    sorted_values[:, 0] - sorted_values[:, 1] - sorted_values[:, 2])

assert np.isclose(reference, val)