Packing of interior facet data in parallel with ghost facets

[jorgensd]

As raised by @adajel, the compute_interface_data is not sufficient for sensible packing in parallel.
This is due to the fact that there might exists facets (ghosted on proc), that only has a single cell on this process( i.e. the facet is part of a ghost cell). This means that facet interpolation on interior facets might be wrong when relying on the fact that facets (cells in the submesh) is connected to two cells on the parent process.

The below code changes the code such that one pack per side (given by a set of cell tags) individually, and handle the 1 and 2 facet connection case in a unified way.

from mpi4py import MPI
import numpy as np
import numpy.typing as npt
import dolfinx
from typing import Iterable

def extract_facet_integration_data(
    cell_tags: dolfinx.mesh.MeshTags,
    cell_values: Iterable[int],
    facet_indices: npt.NDArray[np.int32],
) -> npt.NDArray[np.int32]:
    """
    Given a facet in `facet_indices`, find the cells tagged with a given set of `cell_values`,
    and compute the integration data for those cells.

    """
    if isinstance(cell_values, int):
        cell_values = (cell_values,)

    if len(facet_indices) == 0:
        return np.empty((0, 2), dtype=np.int32)
    cell_tags.topology.create_connectivity(cell_tags.topology.dim - 1, cell_tags.topology.dim)
    cell_tags.topology.create_connectivity(cell_tags.topology.dim, cell_tags.topology.dim-1)

    f_to_c = cell_tags.topology.connectivity(cell_tags.topology.dim - 1, cell_tags.topology.dim)
    c_to_f = cell_tags.topology.connectivity(cell_tags.topology.dim, cell_tags.topology.dim - 1)

    # Extract the cells connected to each facet.
    # Assumption is that there can only be two cells per facet, and should always be
    # two cells per facet.
    exterior_facet_indices = dolfinx.cpp.mesh.exterior_facet_indices(cell_tags.topology)
 
    is_exterior = np.isin(facet_indices, exterior_facet_indices)
    if is_exterior.any():
        raise RuntimeError("Facet is assumed to be an interior facet")
 
    num_cells_per_facet = f_to_c.offsets[facet_indices + 1] - f_to_c.offsets[facet_indices]


    num_facets = 0
    num_facets_per_cell = c_to_f.offsets[1:] - c_to_f.offsets[:-1]
  
    idata = np.zeros((0, 2), dtype=np.int32)
    for num_connections in [1,2]:
        n_sided_interior = np.flatnonzero(num_cells_per_facet == num_connections)
        if len(n_sided_interior) == 0:
            continue
        sub_facets = facet_indices[n_sided_interior]
        num_facets += len(n_sided_interior)
        
        facet_pos = np.vstack([f_to_c.offsets[sub_facets]+i for i in range(num_connections)]).T
        cells = f_to_c.array[facet_pos].flatten()
        assert all(
        num_facets_per_cell[cells.flatten()] == num_facets_per_cell[cells.flatten()[0]]
    ), "All connected cells must have the same number of facets (not mixed topology)."
    
        facets = np.vstack(
            [
                c_to_f.array[c_to_f.offsets[cells] + i]
                for i in range(num_facets_per_cell[cells[0]])
            ]
        ).T
    
        # Repeat facet indices twice to be able to do vectorized match
        rep_fi = np.repeat(sub_facets, num_connections)
        indicator = (facets == rep_fi[:, None])
        _row, local_pos = np.nonzero(indicator)
        assert np.unique(_row).shape[0] == len(_row)
        _idata = np.vstack([cells, local_pos]).T.reshape(-1, 2)
        idata = np.vstack([idata, _idata])

    assert num_facets == len(num_cells_per_facet), "Facets shared by more than two cells are not supported"

    is_marked = cell_tags.indices[np.isin(cell_tags.values, np.asarray(cell_values))]

    return idata[np.isin(idata[:, 0], is_marked)]


mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10, ghost_mode=dolfinx.mesh.GhostMode.shared_facet)


def locater(x):
    return x[0] <= x[1]

def locater2(x):
    return locater(x) & (x[1] < 0.7)


tdim = mesh.topology.dim
num_cells_local = mesh.topology.index_map(tdim).size_local + mesh.topology.index_map(tdim).num_ghosts
all_cells = np.full(num_cells_local, 1, dtype=np.int32)
all_cells[dolfinx.mesh.locate_entities(mesh, tdim, locater)] = 2
all_cells[dolfinx.mesh.locate_entities(mesh, tdim, locater2)] = 3

ct = dolfinx.mesh.meshtags(mesh, tdim, np.arange(num_cells_local, dtype=np.int32), all_cells)

interface = dolfinx.mesh.locate_entities(mesh, tdim-1, lambda x: np.isclose(x[0], x[1]))
ft = dolfinx.mesh.meshtags(mesh, tdim-1, interface, np.full_like(interface, 6, dtype=np.int32))
ft.name = "facet tags"

with dolfinx.io.XDMFFile(mesh.comm, "tags.xdmf", "w") as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_meshtags(ct, mesh.geometry)
    mesh.topology.create_connectivity(tdim-1,tdim)
    xdmf.write_meshtags(ft, mesh.geometry)


data_1 = extract_facet_integration_data(ct, (2, 3), interface)
data_2 = extract_facet_integration_data(ct, 1, interface)

print(data_1.shape, data_2.shape)