Add equidistribution marking

cpp/dolfinx/refinement/mark.h

@@ -20,6 +20,34 @@
 namespace dolfinx::refinement
 {
 
+namespace impl
+{
+
+/// @brief Threshold marking helper
+///
+/// @param[in] marker Input marker
+/// @param[in] threshold Lower bound for values to mark
+///
+/// @returns indices i which satisfy e_i > threshold.
+template <std::floating_point T>
+std::vector<std::int32_t> mark_threshold(std::span<const T> marker, T threshold)
+{
+  auto mark = [=](T e) { return e > threshold; };
+
+  std::vector<std::int32_t> indices;
+  indices.reserve(std::ranges::count_if(marker, mark));
+
+  for (std::int32_t i = 0; i < static_cast<std::int32_t>(marker.size()); ++i)
+  {
+    if (mark(marker[i]))
+      indices.push_back(i);
+  }
+
+  return indices;
+}
+
+} // namespace impl
+
 /// @brief Maximum marking of a marker.
 ///
 /// @param[in] marker Input marker (local) - usually an error indicator per
@@ -39,18 +67,45 @@ std::vector<std::int32_t> mark_maximum(std::span<const T> marker, T theta,
                          : std::ranges::max(marker);
   MPI_Allreduce(MPI_IN_PLACE, &max, 1, dolfinx::MPI::mpi_t<T>, MPI_MAX, comm);
 
-  auto mark = [=](T e) { return e > theta * max; };
+  auto indices = impl::mark_threshold<T>(marker, theta * max);
 
-  std::vector<std::int32_t> indices;
-  indices.reserve(std::ranges::count_if(marker, mark));
+  spdlog::info("Marking (max) {} / {} (local) entities.", indices.size(),
+               marker.size());
 
-  for (std::int32_t i = 0; i < static_cast<std::int32_t>(marker.size()); ++i)
-  {
-    if (mark(marker[i]))
-      indices.push_back(i);
-  }
+  return indices;
+}
 
-  spdlog::info("Marking (max) {} / {} (local) entities.", indices.size(),
+/// @brief Equidistribution marking of a marker.
+///
+/// @param[in] marker Input marker (local) - usually an error indicator per
+/// entity
+/// @param[in] theta Parameter, 0 < θ < 1
+/// @param[in] comm Communicator over which the total marker is computed.
+/// @return Indices (local) of marker elements, which satisfy: marker_i ≥ θ
+/// (sum_i marker_i^2)^1/2 / N^1/2.
+template <std::floating_point T>
+std::vector<std::int32_t> mark_equidistribution(std::span<const T> marker,
+                                                T theta, MPI_Comm comm)
+{
+  if ((theta <= 0) || (theta >= 1))
+    throw std::invalid_argument("Theta needs to fullfill 0 < θ < 1.");
+
+  T norm{0};
+  for (T e : marker)
+    norm += std::pow(e, 2);
+
+  MPI_Allreduce(MPI_IN_PLACE, &norm, 1, dolfinx::MPI::mpi_t<T>, MPI_SUM, comm);
+
+  norm = std::sqrt(norm);
+
+  std::int32_t count = marker.size();
+  MPI_Allreduce(MPI_IN_PLACE, &count, 1, dolfinx::MPI::mpi_t<std::int32_t>,
+                MPI_SUM, comm);
+
+  auto indices
+      = impl::mark_threshold<T>(marker, theta * norm / std::sqrt(count));
+
+  spdlog::info("Marking (equi) {} / {} (local) entities.", indices.size(),
                marker.size());
 
   return indices;

cpp/test/mesh/refinement/mark.cpp

@@ -51,3 +51,44 @@ TEMPLATE_TEST_CASE("Mark maximum", "[refinement][mark][maximum]", double, float)
     CHECK(expect_marked == marked);
   }
 }
+
+TEMPLATE_TEST_CASE("Mark equidistribution empty",
+                   "[refinement][mark][equidistribution]", double, float)
+{
+  std::vector<TestType> marker;
+  auto indices = mark_equidistribution<TestType>(marker, .5, MPI_COMM_WORLD);
+  CHECK(indices.size() == 0);
+}
+
+TEMPLATE_TEST_CASE("Mark equidistribution",
+                   "[refinement][mark][equidistribution]", double, float)
+{
+  MPI_Comm comm = MPI_COMM_WORLD;
+
+  std::vector<TestType> marker;
+  marker.reserve(10);
+  for (std::size_t i = 0; i < 10; i++)
+    marker.push_back(10 * dolfinx::MPI::rank(comm) + i);
+
+  TestType theta = 0.5;
+  auto indices = mark_equidistribution<TestType>(marker, theta, comm);
+
+  CHECK(std::ranges::all_of(
+      indices, [&](auto e)
+      { return (0 <= e) && (e <= static_cast<std::int32_t>(marker.size())); }));
+
+  std::int32_t n = dolfinx::MPI::size(comm) * 10 - 1;
+  TestType norm = std::sqrt(n * (n + 1) * (2 * n + 1) / 6);
+
+  auto mark = [=](auto e) { return e > theta * norm / std::sqrt(n); };
+
+  CHECK(std::ranges::count_if(marker, mark)
+        == static_cast<std::int32_t>(indices.size()));
+
+  for (std::int32_t i = 0; i < static_cast<std::int32_t>(marker.size()); ++i)
+  {
+    bool expect_marked = mark(marker[i]);
+    bool marked = std::ranges::find(indices, i) != indices.end();
+    CHECK(expect_marked == marked);
+  }
+}

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 @@
     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)