Add equidistribution marking

cpp/dolfinx/refinement/mark.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2026 Paul T. Kühner
+// Copyright (C) 2026 Paul T. Kühner, Jack S. Hale
 //
 // This file is part of DOLFINX (https://www.fenicsproject.org)
 //
@@ -8,50 +8,171 @@
 
 #include <algorithm>
 #include <cassert>
+#include <cmath>
 #include <concepts>
 #include <cstdint>
 #include <limits>
 #include <mpi.h>
+#include <numeric>
+#include <span>
 #include <spdlog/spdlog.h>
+#include <stdexcept>
 #include <vector>
 
 #include "dolfinx/common/MPI.h"
 
 namespace dolfinx::refinement
 {
 
-/// @brief Maximum marking of a marker.
+namespace impl
+{
+
+/// @brief Computes local threshold marking of indicators.
+///
+/// Helper for other marking routines.
+///
+/// Returns the indices \f$ i \f$ of the indicators \f$ \eta_i \f$ that satisfy
+/// the threshold: \f$ \eta_i > \text{threshold} \f$.
 ///
-/// @param[in] marker Input marker (local) - usually an error indicator per
-/// entity
-/// @param[in] theta Cut off parameter, 0 < θ < 1
-/// @param[in] comm Communicator over which the maximum is computed.
-/// @return Indices (local) of marker elements, which satisfy: marker_i ≥ θ
-/// max(marker).
+/// @param[in] indicators Indicators (local) \f$ \eta_i \f$ -
+/// usually an error indicator associated with mesh entity \f$ i \f$.
+/// @param[in] threshold Threshold value; indicators greater than this are
+/// marked.
+/// @return Local indices of marked entities.
 template <std::floating_point T>
-std::vector<std::int32_t> mark_maximum(std::span<const T> marker, T theta,
-                                       MPI_Comm comm)
+std::vector<std::int32_t> mark_threshold(std::span<const T> indicators,
+                                         T threshold)
 {
-  if ((theta <= 0) || (theta >= 1))
-    throw std::invalid_argument("Theta needs to fullfill 0 < θ < 1.");
-
-  T max = marker.empty() ? std::numeric_limits<T>::lowest()
-                         : 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 mark = [threshold](T e) { return e > threshold; };
 
   std::vector<std::int32_t> indices;
-  indices.reserve(std::ranges::count_if(marker, mark));
+  indices.reserve(std::ranges::count_if(indicators, mark));
 
-  for (std::int32_t i = 0; i < static_cast<std::int32_t>(marker.size()); ++i)
+  for (std::int32_t i = 0; i < static_cast<std::int32_t>(indicators.size());
+       ++i)
   {
-    if (mark(marker[i]))
+    if (mark(indicators[i]))
       indices.push_back(i);
   }
 
-  spdlog::info("Marking (max) {} / {} (local) entities.", indices.size(),
-               marker.size());
+  return indices;
+}
+
+} // namespace impl
+
+/// @brief Computes maximum-based marking of indicators.
+///
+/// Returns the indices \f$ i \f$ of the indicators \f$ \eta_i \f$ that satisfy
+/// the maximum threshold: \f$ \eta_i > \theta \max_j \eta_j \f$.
+///
+/// @param[in] comm Communicator to compute the maximum over.
+/// @param[in] indicators Indicators (local) \f$ \eta_i \f$ -
+/// usually an error indicator associated with mesh entity \f$ i \f$.
+/// @param[in] theta Parameter, \f$ 0 < \theta < 1 \f$.
+/// @return Local indices of marked entities.
+template <std::floating_point T>
+std::vector<std::int32_t> mark_maximum(MPI_Comm comm,
+                                       std::span<const T> indicators, T theta)
+{
+  if ((theta <= 0) || (theta >= 1))
+    throw std::invalid_argument("theta must fulfill 0 < theta < 1.");
+
+  T max = indicators.empty() ? std::numeric_limits<T>::lowest()
+                             : std::ranges::max(indicators);
+  MPI_Allreduce(MPI_IN_PLACE, &max, 1, dolfinx::MPI::mpi_t<T>, MPI_MAX, comm);
+
+  std::vector<std::int32_t> indices
+      = impl::mark_threshold<T>(indicators, theta * max);
+
+  spdlog::info("Marking (maximum) {} / {} (local) entities.", indices.size(),
+               indicators.size());
+
+  return indices;
+}
+
+/// @brief Computes equidistribution threshold marking of indicators.
+///
+/// Returns the indices \f$ i \f$ of the indicators \f$ \eta_i \f$ that satisfy
+/// the equidistribution threshold: \f$\eta_i > \theta
+/// \frac{||\eta||}{\sqrt{N}} \f$ where \f$ N \f$ is the (global) number of
+/// indicators.
+///
+/// @warning `indicators` must contain owned entities only. Ghost values cause
+/// double-counting in the global sum reductions.
+///
+/// @param[in] comm Communicator over which the global equidistribution
+/// threshold is computed.
+/// @param[in] indicators Indicators for owned local entities \f$ \eta_i \f$
+/// - usually associated with mesh entity \f$ i \f$.
+/// @param[in] theta Parameter, \f$ 0 < \theta < 1 \f$.
+/// @return Local indices of indicators that satisfy the threshold.
+template <std::floating_point T>
+std::vector<std::int32_t>
+mark_equidistribution(MPI_Comm comm, std::span<const T> indicators, T theta)
+{
+  if ((theta <= 0) || (theta >= 1))
+    throw std::invalid_argument("theta must fulfill 0 < theta < 1.");
+
+  T norm = std::inner_product(indicators.begin(), indicators.end(),
+                              indicators.begin(), T{0});
+
+  MPI_Allreduce(MPI_IN_PLACE, &norm, 1, dolfinx::MPI::mpi_t<T>, MPI_SUM, comm);
+
+  T sqrt_norm = std::sqrt(norm);
+
+  // int64_t gives headroom for global sum across ranks.
+  std::int64_t count = indicators.size();
+  MPI_Allreduce(MPI_IN_PLACE, &count, 1, dolfinx::MPI::mpi_t<std::int64_t>,
+                MPI_SUM, comm);
+
+  std::vector<std::int32_t> indices = impl::mark_threshold<T>(
+      indicators, theta * sqrt_norm / std::sqrt(static_cast<T>(count)));
+
+  spdlog::info("Marking (equidistribution) {} / {} (local) entities.",
+               indices.size(), indicators.size());
+
+  return indices;
+}
+
+/// @brief Computes equidistribution threshold marking of a squared indicator.
+///
+/// Returns the indices \f$i\f$ of the squared indicators \f$ \eta_i^2 \f$ that
+/// satisfy the equidistribution threshold: \f$ \eta_i^2 > \theta^2
+/// \frac{||\eta||^2}{N} \f$ where \f$ N \f$ is the (global) number of
+/// indicators.
+///
+/// @warning `squared_indicators` must contain owned entities only. Ghost
+/// values cause double-counting in the global sum reductions.
+///
+/// @param[in] comm Communicator over which the global equidistribution
+/// threshold is computed.
+/// @param[in] squared_indicators Input squared indicators for owned local
+/// entities \f$ \eta^2_i \f$ - usually associated with mesh entity \f$ i \f$.
+/// @param[in] theta Parameter, \f$ 0 < \theta < 1 \f$.
+/// @return Local indices of squared indicators that satisfy the threshold.
+template <std::floating_point T>
+std::vector<std::int32_t>
+mark_equidistribution_squared(MPI_Comm comm,
+                              std::span<const T> squared_indicators, T theta)
+{
+  if ((theta <= 0) || (theta >= 1))
+    throw std::invalid_argument("theta must fulfill 0 < theta < 1.");
+
+  T norm = std::accumulate(squared_indicators.begin(), squared_indicators.end(),
+                           T{0});
+
+  MPI_Allreduce(MPI_IN_PLACE, &norm, 1, dolfinx::MPI::mpi_t<T>, MPI_SUM, comm);
+
+  // int64_t gives headroom for global sum across ranks.
+  std::int64_t count = squared_indicators.size();
+  MPI_Allreduce(MPI_IN_PLACE, &count, 1, dolfinx::MPI::mpi_t<std::int64_t>,
+                MPI_SUM, comm);
+
+  std::vector<std::int32_t> indices = impl::mark_threshold<T>(
+      squared_indicators, theta * theta * norm / static_cast<T>(count));
+
+  spdlog::info("Marking (equidistribution) {} of {} local entities.",
+               indices.size(), squared_indicators.size());
 
   return indices;
 }

cpp/test/mesh/refinement/mark.cpp

@@ -17,37 +17,95 @@ using namespace dolfinx::refinement;
 TEMPLATE_TEST_CASE("Mark maximum empty", "[refinement][mark][maximum]", double,
                    float)
 {
-  std::vector<TestType> marker;
-  auto indices = mark_maximum<TestType>(marker, .5, MPI_COMM_WORLD);
+  std::vector<TestType> indicators;
+  auto indices = mark_maximum<TestType>(MPI_COMM_WORLD, indicators, .5);
   CHECK(indices.size() == 0);
 }
 
 TEMPLATE_TEST_CASE("Mark maximum", "[refinement][mark][maximum]", double, float)
 {
   MPI_Comm comm = MPI_COMM_WORLD;
 
-  std::vector<TestType> marker;
-  marker.reserve(10);
+  std::vector<TestType> indicators;
+  indicators.reserve(10);
   for (std::size_t i = 0; i < 10; i++)
-    marker.push_back(10 * dolfinx::MPI::rank(comm) + i);
+    indicators.push_back(10 * dolfinx::MPI::rank(comm) + i);
 
   TestType theta = 0.5;
-  auto indices = mark_maximum<TestType>(marker, theta, comm);
+  auto indices = mark_maximum<TestType>(comm, indicators, theta);
 
   CHECK(std::ranges::all_of(
-      indices, [&](auto e)
-      { return (0 <= e) && (e <= static_cast<std::int32_t>(marker.size())); }));
+      indices,
+      [&](auto e)
+      {
+        return (0 <= e) && (e <= static_cast<std::int32_t>(indicators.size()));
+      }));
 
   TestType max = dolfinx::MPI::size(comm) * 10 - 1;
   auto mark = [=](auto e) { return e > theta * max; };
 
-  CHECK(std::ranges::count_if(marker, mark)
+  CHECK(std::ranges::count_if(indicators, mark)
         == static_cast<std::int32_t>(indices.size()));
 
-  for (std::int32_t i = 0; i < static_cast<std::int32_t>(marker.size()); ++i)
+  for (std::int32_t i = 0; i < static_cast<std::int32_t>(indicators.size());
+       ++i)
   {
-    bool expect_marked = mark(marker[i]);
+    bool expect_marked = mark(indicators[i]);
     bool marked = std::ranges::find(indices, i) != indices.end();
     CHECK(expect_marked == marked);
   }
 }
+
+TEMPLATE_TEST_CASE("Mark equidistribution empty",
+                   "[refinement][mark][equidistribution]", double, float)
+{
+  std::vector<TestType> indicators;
+  auto indices
+      = mark_equidistribution<TestType>(MPI_COMM_WORLD, indicators, .5);
+  CHECK(indices.size() == 0);
+}
+
+TEMPLATE_TEST_CASE("Mark equidistribution",
+                   "[refinement][mark][equidistribution]", double, float)
+{
+  MPI_Comm comm = MPI_COMM_WORLD;
+
+  std::vector<TestType> indicators;
+  indicators.reserve(10);
+  for (std::size_t i = 0; i < 10; i++)
+    indicators.push_back(10 * dolfinx::MPI::rank(comm) + i);
+
+  TestType theta = 0.5;
+  auto indices = mark_equidistribution<TestType>(comm, indicators, theta);
+
+  CHECK(std::ranges::all_of(
+      indices,
+      [&](auto e)
+      {
+        return (0 <= e) && (e <= static_cast<std::int32_t>(indicators.size()));
+      }));
+
+  std::int32_t n = dolfinx::MPI::size(comm) * 10 - 1;
+  std::int32_t count = dolfinx::MPI::size(comm) * 10;
+  TestType norm = std::sqrt(n * (n + 1) * (2 * n + 1) / 6);
+
+  auto mark = [=](auto e) { return e > theta * norm / std::sqrt(count); };
+
+  CHECK(std::ranges::count_if(indicators, mark)
+        == static_cast<std::int32_t>(indices.size()));
+
+  for (std::int32_t i = 0; i < static_cast<std::int32_t>(indicators.size());
+       ++i)
+  {
+    bool expect_marked = mark(indicators[i]);
+    bool marked = std::ranges::find(indices, i) != indices.end();
+    CHECK(expect_marked == marked);
+  }
+
+  // squared input
+  auto indicators_sq = indicators;
+  std::ranges::for_each(indicators_sq, [](auto& e) { e = std::pow(e, 2); });
+
+  CHECK(std::ranges::equal(indices, mark_equidistribution_squared<TestType>(
+                                        comm, indicators_sq, theta)));
+}