Merge pull request #152 from scientificcomputing/dokken/general_interpolation_matrix

Interpolate cell submesh function onto an facet submesh

Author: jorgensd

src/scifem/interpolation.py

@@ -1,10 +1,11 @@
+from packaging.version import Version
 import dolfinx
 import ufl
 import numpy as np
 import numpy.typing as npt
 from .utils import unroll_dofmap
 
-__all__ = ["interpolation_matrix", "prepare_interpolation_data"]
+__all__ = ["interpolation_matrix", "prepare_interpolation_data", "interpolate_to_surface_submesh"]
 
 if dolfinx.has_petsc4py:
     from petsc4py import PETSc
@@ -255,3 +256,65 @@ def scatter(
         )
         A.assemble()
         return A
+
+
+def interpolate_to_surface_submesh(
+    u_volume: dolfinx.fem.Function,
+    u_surface: dolfinx.fem.Function,
+    submesh_facets: npt.NDArray[np.int32],
+    integration_entities: npt.NDArray[np.int32],
+):
+    """
+    Interpolate a function `u_volume` into the function `u_surface`.
+    Note:
+        Does not work for DG as no dofs are associated with the facets
+    Args:
+        u_volume: Function to interpolate data from
+        u_surface: Function to interpolate data to
+        submesh_facets: Cells in facet mesh
+        integration_entities: Integration entities on the parent mesh
+            corresponding to the facets in `submesh_facets`
+    """
+    if Version(dolfinx.__version__) < Version("0.10.0"):
+        raise RuntimeError("interpolate_to_submesh requires dolfinx version 0.10.0 or higher")
+    V_vol = u_volume.function_space
+    mesh = V_vol.mesh
+    V_surf = u_surface.function_space
+    submesh = V_surf.mesh
+    ip = V_surf.element.interpolation_points
+
+    expr = dolfinx.fem.Expression(u_volume, ip)
+
+    mesh.topology.create_connectivity(mesh.topology.dim, submesh.topology.dim)
+    mesh.topology.create_connectivity(submesh.topology.dim, mesh.topology.dim)
+
+    data = expr.eval(mesh, integration_entities)
+
+    submesh.topology.create_entity_permutations()
+    mesh.topology.create_entity_permutations()
+    cell_info = mesh.topology.get_cell_permutation_info()
+    ft = V_surf.element.basix_element.cell_type
+
+    for i in range(integration_entities.shape[0]):
+        perm = np.arange(data.shape[1], dtype=np.int32)
+        V_vol.element.basix_element.permute_subentity_closure_inv(
+            perm,
+            cell_info[integration_entities[i, 0]],
+            ft,
+            int(integration_entities[i, 1]),
+        )
+        data[i] = data[i][perm]
+
+    if len(data.shape) == 3:
+        # Data is now (num_cells, value_size,num_points)
+        data = data.swapaxes(1, 2)
+        # Data is now (value_size, num_cells, num_points)
+        data = data.swapaxes(0, 1)
+
+    if expr.value_size == 1:
+        shaped_data = data.flatten()
+    else:
+        shaped_data = data.reshape(expr.value_size, -1)
+
+    u_surface._cpp_object.interpolate(shaped_data, submesh_facets)
+    u_surface.x.scatter_forward()

src/scifem/mesh.py

@@ -9,6 +9,7 @@
 from typing import Protocol
 from packaging.version import Version
 
+
 __all__ = [
     "create_entity_markers",
     "transfer_meshtags_to_submesh",
@@ -195,14 +196,9 @@ def extract_submesh(
     # Accumulate all entities, including ghosts, for the specfic set of tagged entities
     edim = entity_tag.dim
     mesh.topology.create_connectivity(edim, mesh.topology.dim)
-    emap = mesh.topology.index_map(entity_tag.dim)
-    marker = dolfinx.la.vector(emap)
     tags_as_arr = np.asarray(tags, dtype=entity_tag.values.dtype)
     all_tagged_indices = np.isin(entity_tag.values, tags_as_arr)
-    marker.array[entity_tag.indices[all_tagged_indices]] = 1
-    marker.scatter_reverse(dolfinx.la.InsertMode.add)
-    marker.scatter_forward()
-    entities = np.flatnonzero(marker.array)
+    entities = entity_tag.indices[all_tagged_indices]
     # Extract submesh
     submesh, cell_map, vertex_map, node_map = dolfinx.mesh.create_submesh(mesh, edim, entities)
 
@@ -261,7 +257,14 @@ def find_interface(
 
     # Compute intersecting facets
     interface_facets = np.intersect1d(facets0, facets1)
-    return reverse_mark_entities(facet_map, interface_facets)
+
+    reverse_mark_entities(facet_map, interface_facets)
+
+    topology.create_connectivity(tdim - 1, tdim)
+    f_to_c = topology.connectivity(tdim - 1, tdim)
+    num_cells_per_facet = f_to_c.offsets[interface_facets + 1] - f_to_c.offsets[interface_facets]
+    is_interface = interface_facets[num_cells_per_facet == 2]
+    return is_interface
 
 
 def compute_subdomain_exterior_facets(
@@ -324,7 +327,9 @@ def compute_subdomain_exterior_facets(
 
 
 def compute_interface_data(
-    cell_tags: dolfinx.mesh.MeshTags, facet_indices: npt.NDArray[np.int32]
+    cell_tags: dolfinx.mesh.MeshTags,
+    facet_indices: npt.NDArray[np.int32],
+    include_ghosts: bool = False,
 ) -> npt.NDArray[np.int32]:
     """
     Compute interior facet integrals that are consistently ordered according to the `cell_tags`,
@@ -336,6 +341,8 @@ def compute_interface_data(
         cell_tags: MeshTags that must contain an integer marker for all cells adjacent
             to the `facet_indices`
         facet_indices: List of facets (local index) that are on the interface.
+        include_ghosts: If `True` integration entities will include facets that are ghosts on
+        the current process. This is for instance useful for interpolation on interior facets.
     Returns:
         The integration data.
     """
@@ -346,12 +353,44 @@ def compute_interface_data(
     else:
         fdim = cell_tags.dim - 1
         integration_args = (fdim,)
-    idata = dolfinx.cpp.fem.compute_integration_domains(
-        dolfinx.fem.IntegralType.interior_facet,
-        cell_tags.topology,
-        facet_indices,
-        *integration_args,
-    )
+    if include_ghosts:
+        if len(facet_indices) == 0:
+            return np.empty((0, 4), dtype=np.int32)
+        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.
+        num_cells_per_facet = f_to_c.offsets[facet_indices + 1] - f_to_c.offsets[facet_indices]
+        assert np.all(num_cells_per_facet == 2), "All facets must be interior facets."
+        facet_pos = np.vstack([f_to_c.offsets[facet_indices], f_to_c.offsets[facet_indices] + 1]).T
+        cells = f_to_c.array[facet_pos].flatten()
+        # Extract facets connected to all cells
+        # Assumption is that all cells have the same number of facets
+        num_facets_per_cell = c_to_f.offsets[1:] - c_to_f.offsets[:-1]
+        assert all(
+            num_facets_per_cell[cells.flatten()] == num_facets_per_cell[cells.flatten()[0]]
+        ), "Cells must have facets."
+        facets = np.vstack(
+            [
+                c_to_f.array[c_to_f.offsets[cells.flatten()] + i]
+                for i in range(num_facets_per_cell[cells.flatten()[0]])
+            ]
+        ).T
+        # Repeat facet indices twice to be able to do vectorized match
+        rep_fi = np.repeat(facet_indices, 2)
+        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, 4)
+    else:
+        idata = dolfinx.cpp.fem.compute_integration_domains(
+            dolfinx.fem.IntegralType.interior_facet,
+            cell_tags.topology,
+            facet_indices,
+            *integration_args,
+        )
     ordered_idata = idata.reshape(-1, 4).copy()
     switch = cell_tags.values[ordered_idata[:, 0]] > cell_tags.values[ordered_idata[:, 2]]
     if True in switch:

tests/test_interpolation.py

@@ -1,4 +1,5 @@
 from mpi4py import MPI
+from packaging.version import Version
 import dolfinx
 import scifem.interpolation
 import pytest
@@ -203,3 +204,114 @@ def test_discrete_curl(degree, use_petsc, cell_type):
     w_ref.x.scatter_forward()
 
     np.testing.assert_allclose(w.x.array, w_ref.x.array, rtol=1e-10, atol=1e-11)
+
+
+@pytest.mark.parametrize("degree", [1, 2, 3])
+@pytest.mark.parametrize("family", ["Lagrange"])
+@pytest.mark.skipif(
+    Version(dolfinx.__version__) < Version("0.10.0"), reason="Requires DOLFINx version >0.10.0"
+)
+def test_interpolate_to_interface_submesh(family, degree):
+    # Create a unit square
+    comm = MPI.COMM_WORLD
+    domain = dolfinx.mesh.create_unit_square(
+        comm, 48, 48, ghost_mode=dolfinx.mesh.GhostMode.shared_facet
+    )
+
+    # Split unit square in two subdomains
+    cell_map = domain.topology.index_map(domain.topology.dim)
+    num_cells_local = cell_map.size_local + cell_map.num_ghosts
+    markers = np.full(num_cells_local, 1, dtype=np.int32)
+    markers[
+        dolfinx.mesh.locate_entities(domain, domain.topology.dim, lambda x: x[0] <= 0.5 + 1e-14)
+    ] = 2
+    ct = dolfinx.mesh.meshtags(
+        domain, domain.topology.dim, np.arange(num_cells_local, dtype=np.int32), markers
+    )
+
+    # Create submesh for each subdomain
+    omega_e, e_to_parent, _, _, _ = scifem.mesh.extract_submesh(domain, ct, (1,))
+    omega_i, i_to_parent, _, _, _ = scifem.mesh.extract_submesh(domain, ct, (2,))
+
+    # Compute submesh for the interface between omega_e and omega_i
+    interface_facets = scifem.mesh.find_interface(ct, (1,), (2,))
+    ft = dolfinx.mesh.meshtags(
+        domain,
+        domain.topology.dim - 1,
+        interface_facets,
+        np.full(interface_facets.shape, 1, dtype=np.int32),
+    )
+
+    gamma, gamma_to_parent, _, _, _ = scifem.mesh.extract_submesh(domain, ft, 1)
+
+    num_facets_local = (
+        gamma.topology.index_map(gamma.topology.dim).size_local
+        + gamma.topology.index_map(gamma.topology.dim).num_ghosts
+    )
+    gamma_to_parent_map = gamma_to_parent.sub_topology_to_topology(
+        np.arange(num_facets_local, dtype=np.int32), inverse=False
+    )
+
+    # Create functions on each subdomain
+    def fe(x):
+        return x[0] + x[1] ** degree
+
+    def fi(x):
+        return np.sin(x[0]) + np.cos(x[1])
+
+    Ve = dolfinx.fem.functionspace(omega_e, (family, degree))
+    ue = dolfinx.fem.Function(Ve)
+    ue.interpolate(fe)
+    ue.x.scatter_forward()
+    Vi = dolfinx.fem.functionspace(omega_i, (family, degree))
+    ui = dolfinx.fem.Function(Vi)
+    ui.interpolate(fi)
+    ui.x.scatter_forward()
+
+    # Compute ordered integration entities on the interface
+    interface_integration_entities = scifem.compute_interface_data(
+        ct, facet_indices=gamma_to_parent_map, include_ghosts=True
+    )
+    mapped_entities = interface_integration_entities.copy()
+
+    # For each submesh, get the relevant integration entities
+    parent_to_e = e_to_parent.sub_topology_to_topology(
+        np.arange(num_cells_local, dtype=np.int32), inverse=True
+    )
+    parent_to_i = i_to_parent.sub_topology_to_topology(
+        np.arange(num_cells_local, dtype=np.int32), inverse=True
+    )
+    mapped_entities[:, 0] = parent_to_e[interface_integration_entities[:, 0]]
+    mapped_entities[:, 2] = parent_to_i[interface_integration_entities[:, 2]]
+    assert np.all(mapped_entities[:, 0] >= 0)
+    assert np.all(mapped_entities[:, 2] >= 0)
+
+    # Create two functions on the interface submesh
+    Q = dolfinx.fem.functionspace(gamma, (family, degree))
+    qe = dolfinx.fem.Function(Q, name="qe")
+    qi = dolfinx.fem.Function(Q, name="qi")
+
+    # Interpolate volume functions (on submesh) onto all cells of the interface submesh
+    scifem.interpolation.interpolate_to_surface_submesh(
+        ue, qe, np.arange(len(gamma_to_parent_map), dtype=np.int32), mapped_entities[:, :2]
+    )
+    qe.x.scatter_forward()
+    scifem.interpolation.interpolate_to_surface_submesh(
+        ui, qi, np.arange(len(gamma_to_parent_map), dtype=np.int32), mapped_entities[:, 2:]
+    )
+    qi.x.scatter_forward()
+
+    # Compute the difference between the two interpolated functions
+    I = dolfinx.fem.Function(Q, name="i")
+    I.x.array[:] = qe.x.array - qi.x.array
+
+    reference = dolfinx.fem.Function(Q)
+    reference.interpolate(lambda x: fe(x) - fi(x))
+
+    qe_ref = dolfinx.fem.Function(Q)
+    qe_ref.interpolate(fe)
+    qi_ref = dolfinx.fem.Function(Q)
+    qi_ref.interpolate(fi)
+    np.testing.assert_allclose(qe.x.array, qe_ref.x.array)
+    np.testing.assert_allclose(qi.x.array, qi_ref.x.array)
+    np.testing.assert_allclose(I.x.array, reference.x.array, rtol=1e-14, atol=1e-14)