Take the intersection for nested lists of subdomain_ids

Author: pbrubeck

firedrake/mesh.py

@@ -25,7 +25,7 @@
 from pyop2.mpi import (
     MPI, COMM_WORLD, temp_internal_comm
 )
-from functools import cached_property
+from functools import cached_property, reduce
 from pyop2.utils import as_tuple
 import petsctools
 from petsctools import OptionsManager, get_external_packages
@@ -4809,6 +4809,7 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig
     subdomain_id : int | Sequence | None
         Subdomain ID representing the submesh.
         If multiple subdomain IDs are provided, their union is taken.
+        If nested lists of subdomain IDs are provided, their intersection is taken.
         If `None` the submesh will cover the entire domain,
         this is useful to obtain a codim-1 submesh over all facets or
         a submesh over a different communicator.
@@ -4878,20 +4879,33 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig
         elif subdim == dim - 1:
             label_name = dmcommon.FACE_SETS_LABEL
 
-    if isinstance(subdomain_id, (tuple, list)):
-        # A list of subdomain ids requires us to build an internal DM label with the union
-        iset = PETSc.IS().createGeneral([], comm=comm or mesh.comm)
+    # Parse non-integer subdomain_id
+    if subdomain_id == "on_boundary":
+        subdomain_id = tuple(mesh.exterior_facets.unique_markers)
+
+    if isinstance(subdomain_id, Sequence):
+        # Create a temporary DMLabel with the union of the labels in the list
+        icomm = comm or mesh.comm
+        iset = PETSc.IS().createGeneral([], comm=icomm)
         for sub in subdomain_id:
-            iset = iset.union(plex.getStratumIS(label_name, sub))
-        label_name = "temp_union"
+            if isinstance(sub, Sequence):
+                # Take the intersection of the (closure of the) labels from nested lists
+                ises = [plex.getStratumIS(label_name, subi) for subi in sub]
+                closure = [[plex.getTransitiveClosure(p)[0] for p in i.indices] for i in ises]
+                indices = reduce(np.intersect1d, closure)
+                cur = PETSc.IS().createGeneral(indices, comm=icomm)
+            else:
+                cur = plex.getStratumIS(label_name, sub)
+            iset = iset.union(cur)
+        label_name = "temp_label"
         subdomain_id = 1
         plex.createLabel(label_name)
         label = plex.getLabel(label_name)
         label.setStratumIS(subdomain_id, iset)
 
     subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm=comm)
 
-    if label_name == "temp_union":
+    if label_name == "temp_label":
         plex.removeLabel(label_name)
 
     comm = comm or mesh.comm
@@ -4900,7 +4914,7 @@ def Submesh(mesh, subdim=None, subdomain_id=None, label_name=None, name=None, ig
     if subplex.getDimension() != subdim:
         raise RuntimeError(f"Found subplex dim ({subplex.getDimension()}) != expected ({subdim})")
     if reorder is None:
-        # Ideally we should set perm_is = mesh.dm_reordering[label_indices]
+        # Ideally we should set perm_is = mesh._dm_renumbering[label_indices]
         reorder = mesh._did_reordering
 
     submesh = Mesh(

tests/firedrake/submesh/test_submesh_facet.py

@@ -134,18 +134,3 @@ def test_submesh_facet_all_facets():
     rmesh = RelabeledMesh(mesh, [facet_function], [facet_value])
     submesh2 = Submesh(rmesh, mesh.topological_dimension - 1, facet_value)
     assert submesh2.cell_set.size == submesh1.cell_set.size
-
-
-def test_submesh_facet_subdomain_id_tuple():
-    mesh = UnitCubeMesh(2, 2, 2)
-    subdomain_id = (1, 3, 6)
-    submesh1 = Submesh(mesh, mesh.topological_dimension - 1, subdomain_id=subdomain_id)
-    assert abs(assemble(1*dx(domain=submesh1)) - len(subdomain_id)) < 1E-12
-
-    V = FunctionSpace(mesh, "HDiv Trace", 0)
-    facet_function = Function(V)
-    DirichletBC(V, 1, subdomain_id).apply(facet_function)
-    facet_value = 999
-    rmesh = RelabeledMesh(mesh, [facet_function], [facet_value])
-    submesh2 = Submesh(rmesh, mesh.topological_dimension - 1, facet_value)
-    assert submesh2.cell_set.size == submesh1.cell_set.size

tests/firedrake/submesh/test_submesh_interface.py

@@ -0,0 +1,93 @@
+import pytest
+import numpy as np
+from firedrake import *
+
+
+def test_submesh_subdomain_id_tuple():
+    mesh = UnitSquareMesh(4, 4)
+    x, y = SpatialCoordinate(mesh)
+    M = FunctionSpace(mesh, "DG", 0)
+    m1 = Function(M).interpolate(conditional(lt(x, 0.5), 1, 0))
+    m2 = Function(M).interpolate(conditional(lt(y, 0.5), 1, 0))
+    mesh.mark_entities(m1, 111)
+    mesh.mark_entities(m2, 222)
+
+    subdomain_id = [111, 222]
+    submesh1 = Submesh(mesh, mesh.topological_dimension, subdomain_id=subdomain_id)
+
+    m3 = Function(M).interpolate(m1 + m2 - m1 * m2)
+    expected = assemble(m3*dx)
+    assert abs(assemble(1*dx(domain=submesh1)) - expected) < 1E-12
+
+    mesh.mark_entities(m3, 333)
+    submesh2 = Submesh(mesh, mesh.topological_dimension, 333)
+    assert submesh2.cell_set.size == submesh1.cell_set.size
+    assert np.allclose(submesh2.coordinates.dat.data, submesh1.coordinates.dat.data)
+
+
+def test_submesh_subdomain_id_nested_tuple():
+    mesh = UnitSquareMesh(4, 4)
+    x, y = SpatialCoordinate(mesh)
+    M = FunctionSpace(mesh, "DG", 0)
+    m1 = Function(M).interpolate(conditional(lt(x, 0.5), 1, 0))
+    m2 = Function(M).interpolate(conditional(lt(y, 0.5), 1, 0))
+    mesh.mark_entities(m1, 111)
+    mesh.mark_entities(m2, 222)
+
+    subdomain_id = [(111, 222)]
+    submesh1 = Submesh(mesh, mesh.topological_dimension, subdomain_id=subdomain_id)
+
+    m3 = Function(M).interpolate(m1 * m2)
+    expected = assemble(m3*dx)
+    assert abs(assemble(1*dx(domain=submesh1)) - expected) < 1E-12
+
+    mesh.mark_entities(m3, 333)
+    submesh2 = Submesh(mesh, mesh.topological_dimension, 333)
+    assert submesh2.cell_set.size == submesh1.cell_set.size
+    assert np.allclose(submesh2.coordinates.dat.data, submesh1.coordinates.dat.data)
+
+
+@pytest.mark.parametrize("subdomain_id", ["on_boundary", (1, 3, 6)])
+def test_submesh_facet_subdomain_id_tuple(subdomain_id):
+    mesh = UnitCubeMesh(2, 2, 2)
+    submesh1 = Submesh(mesh, mesh.topological_dimension - 1, subdomain_id=subdomain_id)
+    if subdomain_id == "on_boundary":
+        area = assemble(1*ds(domain=mesh))
+    else:
+        area = assemble(1*ds(subdomain_id, domain=mesh))
+    assert abs(assemble(1*dx(domain=submesh1)) - area) < 1E-12
+
+    V = FunctionSpace(mesh, "HDiv Trace", 0)
+    facet_function = Function(V)
+    DirichletBC(V, 1, subdomain_id).apply(facet_function)
+    facet_value = 999
+    rmesh = RelabeledMesh(mesh, [facet_function], [facet_value])
+    submesh2 = Submesh(rmesh, mesh.topological_dimension - 1, facet_value)
+    assert submesh2.cell_set.size == submesh1.cell_set.size
+    assert np.allclose(submesh2.coordinates.dat.data, submesh1.coordinates.dat.data)
+
+
+def test_submesh_facet_subdomain_id_nested_tuple():
+    mesh = UnitSquareMesh(4, 4)
+    x, y = SpatialCoordinate(mesh)
+    M = FunctionSpace(mesh, "DG", 0)
+    m1 = Function(M).interpolate(conditional(lt(x, 0.5), 1, 0))
+    m2 = Function(M).interpolate(conditional(lt(x, 0.5), 0, 1))
+    mesh.mark_entities(m1, 111)
+    mesh.mark_entities(m2, 222)
+
+    subdomain_id = [(111, 222)]
+    submesh1 = Submesh(mesh, mesh.topological_dimension - 1, subdomain_id=subdomain_id, label_name="Cell Sets")
+
+    expected = 1
+    assert abs(assemble(1*dx(domain=submesh1)) - expected) < 1E-12
+
+    x, y = SpatialCoordinate(mesh)
+    V = FunctionSpace(mesh, "HDiv Trace", 0)
+    facet_function = Function(V)
+    facet_function.interpolate(conditional(lt(abs(x-0.5), 1E-8), 1, 0))
+    facet_value = 999
+    rmesh = RelabeledMesh(mesh, [facet_function], [facet_value])
+    submesh2 = Submesh(rmesh, mesh.topological_dimension - 1, facet_value)
+    assert submesh2.cell_set.size == submesh1.cell_set.size
+    assert np.allclose(submesh2.coordinates.dat.data, submesh1.coordinates.dat.data)