Add create_entity_markers function

.gitignore

@@ -162,3 +162,8 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Output files
+*.xdmf
+*.h5
+*.bp

README.md

@@ -14,6 +14,7 @@ The package is still in its early stages and many functionalities are still miss
 - Point sources for usage in DOLFINx (>=v0.8.0)
   - Point sources in vector spaces are only supported on v0.9.0, post [DOLFINx PR 3429](https://github.com/FEniCS/dolfinx/pull/3429).
     For older versions, apply one point source in each sub space.
+- Simplified wrapper to create MeshTags based on a list of tags and corresponding locator functions.
 - Maps between degrees of freedom and vertices: `vertex_to_dofmap` and `dof_to_vertex`
 
 ## Installation

_toc.yml

@@ -9,6 +9,7 @@ parts:
     - file: "examples/real_function_space.py"
     - file: "examples/point_source.py"
     - file: "examples/xdmf_point_cloud.py"
+    - file: "examples/mark_entities.py"
   - caption: Reference
     chapters:
     - file: "docs/api.rst"

docs/api.rst

@@ -2,8 +2,5 @@
 API reference
 #############
 
-
-scifem
--------
 .. automodule:: scifem
     :members:

examples/mark_entities.py

@@ -0,0 +1,102 @@
+# # Entity markers
+#
+# Author: Jørgen S. Dokken
+#
+# SPDX-License-Identifier: MIT
+#
+# DOLFINx gives you full control for marking entities.
+# However, sometimes this can feel a bit repetative.
+# In this example we will show how to use {py:func}`scifem.create_entity_markers`.
+
+from mpi4py import MPI
+import dolfinx
+import numpy as np
+
+# We start by creating a simple unit square
+
+mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 60, 60)
+
+# Next, we want to mark some of the cells in our domain.
+# We create three marker functions below.
+# Each of them takes in a ``(3, num_points)`` array, and returns a boolean array of size ``num_points``.
+
+
+def left(x):
+    return x[0] < 0.2
+
+
+def right(x):
+    return x[0] > 0.9
+
+
+def inner(x):
+    # We use numpy bit operator `&` for "and"
+    return (x[0] > 0.3) & (x[0] < 0.7)
+
+
+# We want to mark these entities with  with unique integers 1, 3 and 7.
+
+from scifem import create_entity_markers
+
+cell_tag = create_entity_markers(mesh, mesh.topology.dim, [(1, left), (3, right), (7, inner)])
+
+# Next we can plot these marked entities
+
+import pyvista
+
+pyvista.start_xvfb()
+vtk_grid = dolfinx.plot.vtk_mesh(mesh, cell_tag.dim, cell_tag.indices)
+grid = pyvista.UnstructuredGrid(*vtk_grid)
+grid.cell_data["Marker"] = cell_tag.values
+
+# Create plotter
+
+plotter = pyvista.Plotter()
+plotter.add_mesh(grid)
+plotter.view_xy()
+if not pyvista.OFF_SCREEN:
+    plotter.show()
+
+
+# We can also mark lower order entities, such as facets
+
+
+def circle(x):
+    return x[0] ** 2 + x[1] ** 2 <= 0.16**2
+
+
+def top(x):
+    return x[1] > 0.9
+
+
+facet_tags = create_entity_markers(mesh, mesh.topology.dim - 1, [(2, top), (7, circle)])
+
+
+facet_grid = dolfinx.plot.vtk_mesh(mesh, facet_tags.dim, facet_tags.indices)
+
+fgrid = pyvista.UnstructuredGrid(*facet_grid)
+fgrid.cell_data["Marker"] = facet_tags.values
+
+fplotter = pyvista.Plotter()
+fplotter.add_mesh(fgrid)
+fplotter.view_xy()
+if not pyvista.OFF_SCREEN:
+    fplotter.show()
+
+# We can also exclude interior facets by adding `on_boundary: True` (by default this is set to False).
+
+boundary_facet_tags = create_entity_markers(
+    mesh, mesh.topology.dim - 1, [(2, top, True), (7, circle, False)]
+)
+
+
+boundary_grid = dolfinx.plot.vtk_mesh(mesh, boundary_facet_tags.dim, boundary_facet_tags.indices)
+
+bfgrid = pyvista.UnstructuredGrid(*boundary_grid)
+bfgrid.cell_data["Marker"] = boundary_facet_tags.values
+
+bfplotter = pyvista.Plotter()
+bfplotter.add_mesh(bfgrid)
+bfplotter.view_xy()
+if not pyvista.OFF_SCREEN:
+    bfplotter.show()

src/scifem/__init__.py

@@ -6,6 +6,7 @@
 from .point_source import PointSource
 from .assembly import assemble_scalar
 from . import xdmf
+from .mesh import create_entity_markers
 
 __all__ = [
     "PointSource",
@@ -17,6 +18,7 @@
     "xdmf",
     "vertex_to_dofmap",
     "dof_to_vertexmap",
+    "create_entity_markers",
 ]
 
 

src/scifem/mesh.py

@@ -0,0 +1,64 @@
+import dolfinx
+import typing
+import numpy as np
+import numpy.typing as npt
+
+__all__ = ["create_entity_markers"]
+
+
+# (tag, locator, on_boundary) where on_boundary is optional
+TaggedEntities = (
+    tuple[int, typing.Callable[[npt.NDArray[np.floating]], npt.NDArray[np.bool_]]]
+    | tuple[int, typing.Callable[[npt.NDArray[np.floating]], npt.NDArray[np.bool_]], bool]
+)
+
+
+def create_entity_markers(
+    domain: dolfinx.mesh.Mesh,
+    dim: int,
+    entities_list: list[TaggedEntities],
+) -> dolfinx.mesh.MeshTags:
+    """Mark entities of specified dimension according to a geometrical marker function.
+
+    Args:
+        domain: A ``dolfinx.mesh.Mesh`` object
+        dim: Dimension of the entities to mark
+        entities_list: A list of tuples with the following elements:
+
+            - ``index 0``: The tag to assign to the entities
+            - ``index 1``: A function that takes a point and returns a boolean array
+              indicating whether the point is inside the entity
+            - ``index 2``: Optional, if True, the entities will be marked on the boundary
+
+    Returns:
+        A ``dolfinx.mesh.MeshTags`` object with the corresponding entities marked.
+        If an entity satisfies multiple input marker functions,
+        it is not deterministic what value the entity gets.
+
+    Note:
+
+    """
+    # Create required connectivities
+    domain.topology.create_entities(dim)
+    domain.topology.create_connectivity(dim, 0)
+    domain.topology.create_connectivity(dim, domain.topology.dim)
+
+    # Create marker function
+    e_map = domain.topology.index_map(dim)
+    num_entities_local = e_map.size_local + e_map.num_ghosts
+    markers = np.full(num_entities_local, -1, dtype=np.int32)
+
+    locate_entities = (
+        lambda on_boundary: dolfinx.mesh.locate_entities_boundary
+        if on_boundary
+        else dolfinx.mesh.locate_entities
+    )
+
+    # Concatenate and sort the arrays based on indices
+    for tagged_entity in entities_list:
+        on_boundary = False if len(tagged_entity) == 2 else tagged_entity[2]
+        entities = locate_entities(on_boundary)(domain, dim, tagged_entity[1])
+        markers[entities] = tagged_entity[0]
+
+    facets = np.flatnonzero(markers != -1).astype(np.int32)
+    return dolfinx.mesh.meshtags(domain, dim, facets, markers[facets])

tests/test_mesh.py

@@ -0,0 +1,109 @@
+from mpi4py import MPI
+import numpy as np
+import dolfinx
+import scifem
+
+import pytest
+
+
+@pytest.mark.parametrize(
+    "entities_list, set_values",
+    [
+        ([(1, lambda x: x[0] <= 1.0)], {1}),
+        ([(2, lambda x: x[0] <= 1.0)], {2}),
+        (
+            [
+                (1, lambda x: x[0] <= 1.0),
+                (2, lambda x: x[0] >= 0.0),
+            ],
+            {2},
+        ),
+    ],
+)
+def test_create_celltags_celltags_all(entities_list, set_values):
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
+    cell_tag = scifem.create_entity_markers(mesh, mesh.topology.dim, entities_list=entities_list)
+
+    im = mesh.topology.index_map(mesh.topology.dim)
+    assert cell_tag.dim == mesh.topology.dim
+    assert cell_tag.indices.shape[0] == im.size_local + im.num_ghosts
+    assert cell_tag.values.shape[0] == im.size_local + im.num_ghosts
+    assert cell_tag.values.dtype == np.int32
+    assert set(cell_tag.values) == set_values
+
+
+@pytest.mark.parametrize(
+    "entities_list, set_values",
+    [
+        ([(1, lambda x: x[0] <= 1.0, False)], {1}),
+        ([(2, lambda x: x[0] <= 1.0, False)], {2}),
+        (
+            [
+                (1, lambda x: x[0] <= 1.0, False),
+                (2, lambda x: x[0] >= 0.0, False),
+            ],
+            {2},
+        ),
+    ],
+)
+def test_create_facet_tags_all_on_boundary_False(entities_list, set_values):
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
+    facet_tag = scifem.create_entity_markers(
+        mesh, mesh.topology.dim - 1, entities_list=entities_list
+    )
+    im = mesh.topology.index_map(mesh.topology.dim - 1)
+
+    assert facet_tag.dim == mesh.topology.dim - 1
+    assert facet_tag.indices.shape[0] == im.size_local + im.num_ghosts
+    assert facet_tag.values.shape[0] == im.size_local + im.num_ghosts
+    assert facet_tag.values.dtype == np.int32
+    assert set(facet_tag.values) == set_values
+
+
+@pytest.mark.parametrize(
+    "entities_list, set_values",
+    [
+        ([(1, lambda x: x[0] <= 1.0, True)], {1}),
+        ([(2, lambda x: x[0] <= 1.0, True)], {2}),
+        (
+            [
+                (1, lambda x: x[0] <= 1.0, True),
+                (2, lambda x: x[0] >= 0.0, True),
+            ],
+            {2},
+        ),
+    ],
+)
+def test_create_facet_tags_all_on_boundary_True(entities_list, set_values):
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
+    facet_tag = scifem.create_entity_markers(
+        mesh, mesh.topology.dim - 1, entities_list=entities_list
+    )
+
+    mesh.topology.create_entities(1)
+    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
+    facet_indices = dolfinx.mesh.exterior_facet_indices(mesh.topology)
+    assert facet_tag.dim == mesh.topology.dim - 1
+    assert facet_tag.indices.shape[0] == len(facet_indices)
+    assert facet_tag.values.shape[0] == len(facet_indices)
+    assert facet_tag.values.dtype == np.int32
+    assert set(facet_tag.values) == set_values
+
+
+@pytest.mark.parametrize(
+    "entities_list",
+    [
+        ([(1, lambda x: x[0] < 0.0)]),
+        ([(1, lambda x: x[0] < 0.0)]),
+        ([]),
+    ],
+)
+def test_create_celltags_empty(entities_list):
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
+    cell_tag = scifem.create_entity_markers(mesh, mesh.topology.dim, entities_list=entities_list)
+
+    assert cell_tag.dim == mesh.topology.dim
+    assert cell_tag.indices.shape[0] == 0
+    assert cell_tag.values.shape[0] == 0
+    assert cell_tag.values.dtype == np.int32
+    assert set(cell_tag.values) == set()