Add equidistribution marking (py)

Author: schnellerhase

python/dolfinx/mesh.py

@@ -35,6 +35,7 @@
 from dolfinx.cpp.refinement import (
     IdentityPartitionerPlaceholder,
     RefinementOption,
+    mark_equidistribution,
     mark_maximum,
 )
 from dolfinx.cpp.refinement import (
@@ -73,6 +74,7 @@
     "exterior_facet_indices",
     "locate_entities",
     "locate_entities_boundary",
+    "mark_equidistribution",
     "mark_maximum",
     "meshtags",
     "meshtags_from_entities",

python/dolfinx/wrappers/dolfinx_wrappers/refinement.h

@@ -134,6 +134,17 @@ void declare_refinement(nanobind::module_& m)
             std::span(marker.data(), marker.size()), theta, comm.get()));
       },
       nb::arg("marker"), nb::arg("theta"), nb::arg("comm"));
+
+  m.def(
+      "mark_equidistribution",
+      [](nb::ndarray<const T, nb::ndim<1>, nb::c_contig> marker, T theta,
+         MPICommWrapper comm)
+      {
+        return dolfinx_wrappers::as_nbarray(
+            dolfinx::refinement::mark_equidistribution(
+                std::span(marker.data(), marker.size()), theta, comm.get()));
+      },
+      nb::arg("marker"), nb::arg("theta"), nb::arg("comm"));
 }
 
 } // namespace dolfinx_wrappers

python/test/unit/refinement/test_mark.py

@@ -31,3 +31,26 @@ def test_mark_maximum(theta: float, dtype: np.dtype) -> None:
     msh.topology.create_entities(1)
     marked_edges = mesh.compute_incident_entities(msh.topology, marked_cells, tdim, 1)
     mesh.refine(msh, marked_edges)
+
+
+@pytest.mark.parametrize("theta", [0.2, 0.4, 0.6, 0.8])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_mark_equidistribution(theta: float, dtype: np.dtype) -> None:
+    msh = mesh.create_unit_square(comm := MPI.COMM_WORLD, n := 10, n, dtype=dtype)
+
+    tdim = msh.topology.dim
+    cell_count = (cell_im := msh.topology.index_map(tdim)).size_local + cell_im.num_ghosts
+    marker = np.random.default_rng(0).random(cell_count)
+
+    marked_cells = mesh.mark_equidistribution(marker, theta, comm)
+
+    norm = np.sqrt(comm.allreduce(np.sum(marker**2), MPI.SUM))
+    count = comm.allreduce(marker.size)
+    assert np.allclose(
+        marked_cells,
+        np.argwhere(marker > theta * norm / np.sqrt(count)).flatten(),
+    )
+
+    msh.topology.create_entities(1)
+    marked_edges = mesh.compute_incident_entities(msh.topology, marked_cells, tdim, 1)
+    mesh.refine(msh, marked_edges)