Add region tags

meshmode/interop/firedrake/mesh.py

@@ -248,22 +248,22 @@ def _get_firedrake_facial_adjacency_groups(fdrake_mesh_topology,
     # groups
     bdy_tags = _get_firedrake_boundary_tags(
         top, tag_induced_boundary=cells_to_use is not None)
-    boundary_tag_to_index = {bdy_tag: i for i, bdy_tag in enumerate(bdy_tags)}
+    from meshmode.mesh import make_tag_to_index, get_tag_bit
+    boundary_tag_to_index = make_tag_to_index(bdy_tags)
     marker_to_neighbor_value = {}
-    from meshmode.mesh import _boundary_tag_bit
     # for convenience,
     # None maps to the boundary tag for a boundary facet with no marker
     marker_to_neighbor_value[None] = \
-        -(_boundary_tag_bit(bdy_tags, boundary_tag_to_index, BTAG_REALLY_ALL)
-          | _boundary_tag_bit(bdy_tags, boundary_tag_to_index, BTAG_ALL))
+        -(get_tag_bit(boundary_tag_to_index, BTAG_REALLY_ALL)
+          | get_tag_bit(boundary_tag_to_index, BTAG_ALL))
     # firedrake exterior facets with no marker are assigned the
     # a dummy marker
     from firedrake.mesh import unmarked as fd_unmarked
     marker_to_neighbor_value[fd_unmarked] = marker_to_neighbor_value[None]
     # Now figure out the appropriate tags for each firedrake markers
     for marker in top.exterior_facets.unique_markers:
         marker_to_neighbor_value[marker] = \
-            -(_boundary_tag_bit(bdy_tags, boundary_tag_to_index, marker)
+            -(get_tag_bit(boundary_tag_to_index, marker)
               | -marker_to_neighbor_value[None])
 
     # {{{ build the FacialAdjacencyGroup for internal connectivity
@@ -314,12 +314,10 @@ def _get_firedrake_facial_adjacency_groups(fdrake_mesh_topology,
                                               newly_created_exterior_facs)
         new_ext_elements = int_elements[newly_created_exterior_facs]
         new_ext_element_faces = int_element_faces[newly_created_exterior_facs]
-        new_ext_neighbor_tag = -(_boundary_tag_bit(bdy_tags,
-                                                  boundary_tag_to_index,
-                                                  BTAG_REALLY_ALL)
-                                | _boundary_tag_bit(bdy_tags,
-                                                    boundary_tag_to_index,
-                                                    BTAG_INDUCED_BOUNDARY))
+        new_ext_neighbor_tag = -(get_tag_bit(boundary_tag_to_index,
+                                             BTAG_REALLY_ALL)
+                                | get_tag_bit(boundary_tag_to_index,
+                                              BTAG_INDUCED_BOUNDARY))
         new_ext_neighbors = np.full(new_ext_elements.shape,
                                     new_ext_neighbor_tag,
                                     dtype=IntType)

meshmode/mesh/__init__.py

@@ -38,10 +38,21 @@
 .. autoclass:: FacialAdjacencyGroup
 .. autoclass:: InterPartitionAdjacencyGroup
 
+.. autofunction:: make_region_tags
+.. autofunction:: make_boundary_tags
+.. autofunction:: make_tag_to_index
+.. autofunction:: get_tag_bit
+
 .. autofunction:: as_python
 .. autofunction:: check_bc_coverage
 .. autofunction:: is_boundary_tag_empty
 
+Predefined Region tags
+----------------------
+
+.. autoclass:: RTAG_NONE
+.. autoclass:: RTAG_ALL
+
 Predefined Boundary tags
 ------------------------
 
@@ -56,12 +67,22 @@
 
 # {{{ element tags
 
+class RTAG_NONE:  # noqa: N801
+    """A region tag representing an empty region."""
+    pass
+
+
+class RTAG_ALL:  # noqa: N801
+    """A region tag representing all regions."""
+    pass
+
+
 class BTAG_NONE:  # noqa: N801
-    """A boundary tag representing an empty boundary or volume."""
+    """A boundary tag representing an empty boundary."""
 
 
 class BTAG_ALL:  # noqa: N801
-    """A boundary tag representing the entire boundary or volume.
+    """A boundary tag representing the entire boundary.
 
     In the case of the boundary, :class:`BTAG_ALL` does not include rank boundaries,
     or, more generally, anything tagged with :class:`BTAG_NO_BOUNDARY`.
@@ -134,7 +155,10 @@ class BTAG_INDUCED_BOUNDARY(BTAG_NO_BOUNDARY):  # noqa: N801
     # firedrakeproject.org seems to reject connections from Github.
 
 
-SYSTEM_TAGS = {BTAG_NONE, BTAG_ALL, BTAG_REALLY_ALL, BTAG_NO_BOUNDARY,
+SYSTEM_RTAGS = {RTAG_NONE, RTAG_ALL}
+
+
+SYSTEM_BTAGS = {BTAG_NONE, BTAG_ALL, BTAG_REALLY_ALL, BTAG_NO_BOUNDARY,
                    BTAG_PARTITION, BTAG_INDUCED_BOUNDARY}
 
 # }}}
@@ -163,6 +187,11 @@ class MeshElementGroup(Record):
 
         *(dim, nunit_nodes)*
 
+    .. attribute:: regions
+
+        An array *(nelements)* of integers, with the bits of ``regions[i]``
+        indicating the mesh regions that contain element ``i``.
+
     .. attribute:: element_nr_base
 
         Lowest element number in this element group.
@@ -192,7 +221,7 @@ class MeshElementGroup(Record):
     """
 
     def __init__(self, order, vertex_indices, nodes,
-            element_nr_base=None, node_nr_base=None,
+            regions=None, element_nr_base=None, node_nr_base=None,
             unit_nodes=None, dim=None, **kwargs):
         """
         :arg order: the maximum total degree used for interpolation.
@@ -210,6 +239,7 @@ def __init__(self, order, vertex_indices, nodes,
             vertex_indices=vertex_indices,
             nodes=nodes,
             unit_nodes=unit_nodes,
+            regions=regions,
             element_nr_base=element_nr_base, node_nr_base=node_nr_base,
             **kwargs)
 
@@ -289,7 +319,7 @@ def __ne__(self, other):
 
 class _ModepyElementGroup(MeshElementGroup):
     def __init__(self, order, vertex_indices, nodes,
-            element_nr_base=None, node_nr_base=None,
+            regions=None, element_nr_base=None, node_nr_base=None,
             unit_nodes=None, dim=None, **kwargs):
         """
         :arg order: the maximum total degree used for interpolation.
@@ -341,6 +371,7 @@ def __init__(self, order, vertex_indices, nodes,
                         f" got {vertex_indices.shape[-1]}")
 
         super().__init__(order, vertex_indices, nodes,
+                regions=regions,
                 element_nr_base=element_nr_base,
                 node_nr_base=node_nr_base,
                 unit_nodes=unit_nodes,
@@ -630,6 +661,15 @@ class Mesh(Record):
         (Note that element groups are not necessarily geometrically contiguous
         like the figure may suggest.)
 
+    .. attribute:: region_tags
+
+        A list of region tag identifiers. :class:`RTAG_ALL` is guaranteed to exist.
+
+    .. attribute:: rtag_to_index
+
+        A mapping that maps region tag identifiers to their
+        corresponding index.
+
     .. attribute:: boundary_tags
 
         A list of boundary tag identifiers. :class:`BTAG_ALL` and
@@ -640,11 +680,6 @@ class Mesh(Record):
         A mapping that maps boundary tag identifiers to their
         corresponding index.
 
-        .. note::
-
-            Elements of :attr:`boundary_tags` that do not cover any
-            part of the boundary will not be keys in this dictionary.
-
     .. attribute:: vertex_id_dtype
 
     .. attribute:: element_id_dtype
@@ -667,6 +702,7 @@ def __init__(self, vertices, groups, *, skip_tests=False,
             skip_element_orientation_test=False,
             nodal_adjacency=None,
             facial_adjacency_groups=None,
+            region_tags=None,
             boundary_tags=None,
             vertex_id_dtype=np.int32,
             element_id_dtype=np.int32,
@@ -710,29 +746,31 @@ def __init__(self, vertices, groups, *, skip_tests=False,
             el_nr += ng.nelements
             node_nr += ng.nnodes
 
-        # {{{ boundary tags
+        # {{{ region tags
 
-        if boundary_tags is None:
-            boundary_tags = []
-        else:
-            boundary_tags = boundary_tags[:]
+        region_tags = make_region_tags(user_tags=region_tags)
+
+        max_region_tag_count = int(
+                np.log(np.iinfo(element_id_dtype).max)/np.log(2))
+        if len(region_tags) > max_region_tag_count:
+            raise ValueError("too few bits in element_id_dtype to represent all "
+                    "region tags")
+
+        rtag_to_index = make_tag_to_index(region_tags)
+
+        # }}}
+
+        # {{{ boundary tags
 
-        if BTAG_NONE in boundary_tags:
-            raise ValueError("BTAG_NONE is not allowed to be part of "
-                    "boundary_tags")
-        if BTAG_ALL not in boundary_tags:
-            boundary_tags.append(BTAG_ALL)
-        if BTAG_REALLY_ALL not in boundary_tags:
-            boundary_tags.append(BTAG_REALLY_ALL)
+        boundary_tags = make_boundary_tags(user_tags=boundary_tags)
 
         max_boundary_tag_count = int(
                 np.log(np.iinfo(element_id_dtype).max)/np.log(2))
         if len(boundary_tags) > max_boundary_tag_count:
             raise ValueError("too few bits in element_id_dtype to represent all "
                     "boundary tags")
 
-        btag_to_index = {
-                btag: i for i, btag in enumerate(boundary_tags)}
+        btag_to_index = make_tag_to_index(boundary_tags)
 
         # }}}
 
@@ -759,6 +797,8 @@ def __init__(self, vertices, groups, *, skip_tests=False,
                 self, vertices=vertices, groups=new_groups,
                 _nodal_adjacency=nodal_adjacency,
                 _facial_adjacency_groups=facial_adjacency_groups,
+                region_tags=region_tags,
+                rtag_to_index=rtag_to_index,
                 boundary_tags=boundary_tags,
                 btag_to_index=btag_to_index,
                 vertex_id_dtype=np.dtype(vertex_id_dtype),
@@ -822,10 +862,10 @@ def set_if_not_present(name, from_name=None):
         set_if_not_present("vertices")
         if "groups" not in kwargs:
             kwargs["groups"] = [group.copy() for group in self.groups]
+        set_if_not_present("region_tags")
         set_if_not_present("boundary_tags")
         set_if_not_present("nodal_adjacency", "_nodal_adjacency")
         set_if_not_present("facial_adjacency_groups", "_facial_adjacency_groups")
-        set_if_not_present("boundary_tags")
         set_if_not_present("vertex_id_dtype")
         set_if_not_present("element_id_dtype")
         set_if_not_present("is_conforming")
@@ -899,9 +939,11 @@ def facial_adjacency_groups(self):
 
         return self._facial_adjacency_groups
 
+    def region_tag_bit(self, region_tag):
+        return get_tag_bit(self.rtag_to_index, region_tag)
+
     def boundary_tag_bit(self, boundary_tag):
-        return _boundary_tag_bit(self.boundary_tags, self.btag_to_index,
-                        boundary_tag)
+        return get_tag_bit(self.btag_to_index, boundary_tag)
 
     def __eq__(self, other):
         return (
@@ -912,6 +954,7 @@ def __eq__(self, other):
                 and self.element_id_dtype == other.element_id_dtype
                 and self._nodal_adjacency == other._nodal_adjacency
                 and self._facial_adjacency_groups == other._facial_adjacency_groups
+                and self.region_tags == other.region_tags
                 and self.boundary_tags == other.boundary_tags
                 and self.is_conforming == other.is_conforming)
 
@@ -1029,17 +1072,55 @@ def _compute_nodal_adjacency_from_vertices(mesh):
 # }}}
 
 
-# {{{ boundary tag to bit
+# {{{ tags
+
+def make_region_tags(user_tags=None):
+    """Create a region tag list, optionally including extra *user_tags*."""
+    region_tags = []
+
+    if user_tags is not None:
+        if RTAG_NONE in user_tags:
+            raise ValueError("RTAG_NONE is not allowed to be part of region_tags")
+        region_tags.extend(user_tags)
+
+    if RTAG_ALL not in region_tags:
+        region_tags.append(RTAG_ALL)
+
+    return region_tags
+
+
+def make_boundary_tags(user_tags=None):
+    """Create a boundary tag list, optionally including extra *user_tags*."""
+    boundary_tags = []
+
+    if user_tags is not None:
+        if BTAG_NONE in user_tags:
+            raise ValueError("BTAG_NONE is not allowed to be part of boundary_tags")
+        boundary_tags.extend(user_tags)
+
+    if BTAG_ALL not in boundary_tags:
+        boundary_tags.append(BTAG_ALL)
+    if BTAG_REALLY_ALL not in boundary_tags:
+        boundary_tags.append(BTAG_REALLY_ALL)
+
+    return boundary_tags
+
+
+def make_tag_to_index(tags):
+    """Create a dict that maps tags to their respective index in the tag list."""
+    return {tag: i for i, tag in enumerate(tags)}
 
-def _boundary_tag_bit(boundary_tags, btag_to_index, boundary_tag):
-    if boundary_tag is BTAG_NONE:
+
+def get_tag_bit(tag_to_index, tag):
+    """Get the bit in a tag bitfield that corresponds to *tag*."""
+    if tag is RTAG_NONE or tag is BTAG_NONE:
         return 0
 
-    if boundary_tag not in boundary_tags:
-        raise ValueError("boundary tag '%s' is not known" % boundary_tag)
+    if tag not in tag_to_index:
+        raise ValueError("tag '%s' is not known" % tag)
 
     try:
-        return 1 << btag_to_index[boundary_tag]
+        return 1 << tag_to_index[tag]
     except KeyError:
         return 0
 
@@ -1135,10 +1216,10 @@ def _compute_facial_adjacency_from_vertices(groups, boundary_tags,
     if not groups:
         return None
 
-    boundary_tag_to_index = {tag: i for i, tag in enumerate(boundary_tags)}
+    boundary_tag_to_index = make_tag_to_index(boundary_tags)
 
     def boundary_tag_bit(boundary_tag):
-        return _boundary_tag_bit(boundary_tags, boundary_tag_to_index, boundary_tag)
+        return get_tag_bit(boundary_tag_to_index, boundary_tag)
 
     # Match up adjacent faces according to their vertex indices
 
@@ -1411,13 +1492,27 @@ def check_bc_coverage(mesh, boundary_tags, incomplete_ok=False,
 # }}}
 
 
+# {{{ is_region_tag_empty
+
+def is_region_tag_empty(mesh, region_tag):
+    """:returns: *True* if *region_tag* does not occur as part of *mesh*."""
+    rtag_bit = mesh.region_tag_bit(region_tag)
+    if not rtag_bit:
+        return True
+
+    for grp in mesh.groups:
+        if (grp.regions & rtag_bit).any():
+            return False
+
+    return True
+
+# }}}
+
+
 # {{{ is_boundary_tag_empty
 
 def is_boundary_tag_empty(mesh, boundary_tag):
-    """Return *True* if the corresponding boundary tag does not occur as part of
-    *mesh*.
-    """
-
+    """:returns: *True* if *boundary_tag* does not occur as part of *mesh*."""
     btag_bit = mesh.boundary_tag_bit(boundary_tag)
     if not btag_bit:
         return True