Remove refine option redistribute and ghost_mode in favor of an explicit partitioner

cpp/dolfinx/refinement/refine.h

@@ -1,11 +1,12 @@
-// Copyright (C) 2010-2023 Garth N. Wells
+// Copyright (C) 2010-2024 Garth N. Wells and Paul T. Kühner
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
 // SPDX-License-Identifier:    LGPL-3.0-or-later
 
 #pragma once
 
+#include "dolfinx/graph/AdjacencyList.h"
 #include "dolfinx/mesh/Mesh.h"
 #include "dolfinx/mesh/Topology.h"
 #include "dolfinx/mesh/cell_types.h"
@@ -19,44 +20,29 @@
 
 namespace dolfinx::refinement
 {
-namespace impl
-{
-/// @brief create the refined mesh by optionally redistributing it
-template <std::floating_point T>
-mesh::Mesh<T>
-create_refined_mesh(const mesh::Mesh<T>& mesh,
-                    const graph::AdjacencyList<std::int64_t>& cell_adj,
-                    std::span<const T> new_vertex_coords,
-                    std::array<std::size_t, 2> xshape, bool redistribute,
-                    mesh::GhostMode ghost_mode)
+
+/// TODO: Document function
+inline graph::AdjacencyList<std::int32_t> maintain_coarse_partitioner(
+    MPI_Comm comm, int, const std::vector<mesh::CellType>& cell_types,
+    const std::vector<std::span<const std::int64_t>>& cell_topology)
 {
-  if (dolfinx::MPI::size(mesh.comm()) == 1)
-  {
-    // No parallel construction necessary, i.e. redistribute also has no
-    // effect
-    return mesh::create_mesh(mesh.comm(), cell_adj.array(),
-                             mesh.geometry().cmap(), new_vertex_coords, xshape,
-                             ghost_mode);
-  }
-  else
-  {
-    // Let partition handle the parallel construction of the mesh
-    return partition<T>(mesh, cell_adj, new_vertex_coords, xshape, redistribute,
-                        ghost_mode);
-  }
+  int mpi_rank = MPI::rank(comm);
+  int num_cell_vertices = mesh::num_cell_vertices(cell_types.front());
+  std::int32_t num_cells = cell_topology.front().size() / num_cell_vertices;
+  std::vector<std::int32_t> destinations(num_cells, mpi_rank);
+  std::vector<std::int32_t> dest_offsets(num_cells + 1);
+  std::iota(dest_offsets.begin(), dest_offsets.end(), 0);
+  return graph::AdjacencyList(std::move(destinations), std::move(dest_offsets));
 }
-} // namespace impl
 
 /// @brief Refine with markers, optionally redistributing, and
 /// optionally calculating the parent-child relationships.
 ///
-/// @param[in] mesh Input mesh to be refined
+/// @param[in] mesh Input mesh to be refined.
 /// @param[in] edges Optional indices of the edges that should be split
 /// by this refinement, if optional is not set, a uniform refinement is
 /// performed, same behavior as passing a list of all indices.
-/// @param[in] redistribute Flag to call the Mesh Partitioner to
-/// redistribute after refinement.
-/// @param[in] ghost_mode Ghost mode of the refined mesh.
+/// @param[in] partitioner partitioner to be used for the refined mesh
 /// @param[in] option Control the computation of parent facets, parent
 /// cells. If an option is unselected, an empty list is returned.
 /// @return New Mesh and optional parent cell index, parent facet
@@ -65,26 +51,26 @@ template <std::floating_point T>
 std::tuple<mesh::Mesh<T>, std::optional<std::vector<std::int32_t>>,
            std::optional<std::vector<std::int8_t>>>
 refine(const mesh::Mesh<T>& mesh,
-       std::optional<std::span<const std::int32_t>> edges, bool redistribute,
-       mesh::GhostMode ghost_mode = mesh::GhostMode::shared_facet,
+       std::optional<std::span<const std::int32_t>> edges,
+       mesh::CellPartitionFunction partitioner = maintain_coarse_partitioner,
        Option option = Option::none)
 {
   auto topology = mesh.topology();
   assert(topology);
-
-  mesh::CellType cell_t = topology->cell_type();
-  if (!mesh::is_simplex(cell_t))
+  if (!mesh::is_simplex(topology->cell_type()))
     throw std::runtime_error("Refinement only defined for simplices");
+
   bool oned = topology->cell_type() == mesh::CellType::interval;
   auto [cell_adj, new_vertex_coords, xshape, parent_cell, parent_facet]
       = oned ? interval::compute_refinement_data(mesh, edges, option)
              : plaza::compute_refinement_data(mesh, edges, option);
 
-  mesh::Mesh<T> refined_mesh = impl::create_refined_mesh(
-      mesh, cell_adj, std::span<const T>(new_vertex_coords), xshape,
-      redistribute, ghost_mode);
+  mesh::Mesh<T> refined_mesh
+      = mesh::create_mesh(mesh.comm(), mesh.comm(), std::move(cell_adj).array(),
+                          mesh.geometry().cmap(), mesh.comm(),
+                          std::move(new_vertex_coords), xshape, partitioner);
 
-  // Report the number of refined cellse
+  // Report the number of refined cells
   const int D = topology->dim();
   const std::int64_t n0 = topology->index_map(D)->size_global();
   const std::int64_t n1 = refined_mesh.topology()->index_map(D)->size_global();

cpp/dolfinx/refinement/utils.h

@@ -264,51 +264,6 @@ create_new_vertices(MPI_Comm comm,
           xshape};
 }
 
-/// Use vertex and topology data to partition new mesh across
-/// processes
-/// @param[in] old_mesh
-/// @param[in] cell_topology Topology of cells, (vertex indices)
-/// @param[in] new_coords New coordinates, row-major storage
-/// @param[in] xshape The shape of `new_coords`
-/// @param[in] redistribute Call graph partitioner if true
-/// @param[in] ghost_mode None or shared_facet
-/// @return New mesh
-template <std::floating_point T>
-mesh::Mesh<T> partition(const mesh::Mesh<T>& old_mesh,
-                        const graph::AdjacencyList<std::int64_t>& cell_topology,
-                        std::span<const T> new_coords,
-                        std::array<std::size_t, 2> xshape, bool redistribute,
-                        mesh::GhostMode ghost_mode)
-{
-  if (redistribute)
-  {
-    return mesh::create_mesh(old_mesh.comm(), cell_topology.array(),
-                             old_mesh.geometry().cmap(), new_coords, xshape,
-                             ghost_mode);
-  }
-  else
-  {
-    auto partitioner
-        = [](MPI_Comm comm, int, const std::vector<mesh::CellType>& cell_types,
-             const std::vector<std::span<const std::int64_t>>& cell_topology)
-    {
-      std::int32_t mpi_rank = MPI::rank(comm);
-      std::int32_t num_cell_vertices
-          = mesh::num_cell_vertices(cell_types.front());
-      std::int32_t num_cells = cell_topology.front().size() / num_cell_vertices;
-      std::vector<std::int32_t> destinations(num_cells, mpi_rank);
-      std::vector<std::int32_t> dest_offsets(num_cells + 1);
-      std::iota(dest_offsets.begin(), dest_offsets.end(), 0);
-      return graph::AdjacencyList(std::move(destinations),
-                                  std::move(dest_offsets));
-    };
-
-    return mesh::create_mesh(old_mesh.comm(), old_mesh.comm(),
-                             cell_topology.array(), old_mesh.geometry().cmap(),
-                             old_mesh.comm(), new_coords, xshape, partitioner);
-  }
-}
-
 /// @brief Given an index map, add "n" extra indices at the end of local range
 ///
 /// @note The returned global indices (local and ghosts) are adjust for the

cpp/test/mesh/refinement/interval.cpp

@@ -73,7 +73,7 @@ TEMPLATE_TEST_CASE("Interval uniform refinement",
 
   // TODO: parent_facet
   auto [refined_mesh, parent_edge, parent_facet] = refinement::refine(
-      mesh, std::nullopt, false, mesh::GhostMode::shared_facet,
+      mesh, std::nullopt, &refinement::maintain_coarse_partitioner,
       refinement::Option::parent_cell);
 
   std::vector<T> expected_x = {
@@ -114,7 +114,8 @@ TEMPLATE_TEST_CASE("Interval adaptive refinement",
   std::vector<std::int32_t> edges{1};
   // TODO: parent_facet
   auto [refined_mesh, parent_edge, parent_facet] = refinement::refine(
-      mesh, std::span(edges), false, mesh::GhostMode::shared_facet,
+      mesh, std::span(edges),
+      mesh::create_cell_partitioner(mesh::GhostMode::shared_facet),
       refinement::Option::parent_cell);
 
   std::vector<T> expected_x = {
@@ -192,7 +193,7 @@ TEMPLATE_TEST_CASE("Interval Refinement (parallel)",
 
   // TODO: parent_facet
   auto [refined_mesh, parent_edges, parent_facet] = refinement::refine(
-      mesh, std::nullopt, false, mesh::GhostMode::shared_facet,
+      mesh, std::nullopt, &refinement::maintain_coarse_partitioner,
       refinement::Option::parent_cell);
 
   T rank_d = static_cast<T>(rank);

cpp/test/mesh/refinement/rectangle.cpp

@@ -99,9 +99,9 @@ plotter.show()
   // plaza requires the edges to be pre initialized!
   mesh.topology()->create_entities(1);
 
-  auto [mesh_fine, parent_cell, parent_facet]
-      = refinement::refine(mesh, std::nullopt, false, mesh::GhostMode::none,
-                           refinement::Option::parent_cell_and_facet);
+  auto [mesh_fine, parent_cell, parent_facet] = refinement::refine(
+      mesh, std::nullopt, mesh::create_cell_partitioner(mesh::GhostMode::none),
+      refinement::Option::parent_cell_and_facet);
 
   // vertex layout:
   // 8---7---5

python/dolfinx/mesh.py

@@ -370,8 +370,7 @@ def transfer_meshtag(
 def refine(
     mesh: Mesh,
     edges: typing.Optional[np.ndarray] = None,
-    redistribute: bool = True,
-    ghost_mode: GhostMode = GhostMode.shared_facet,
+    partitioner: typing.Optional[typing.Callable] = None,
     option: RefinementOption = RefinementOption.none,
 ) -> tuple[Mesh, npt.NDArray[np.int32], npt.NDArray[np.int8]]:
     """Refine a mesh.
@@ -380,17 +379,16 @@ def refine(
         mesh: Mesh from which to create the refined mesh.
         edges: Indices of edges to split during refinement. If ``None``,
             mesh refinement is uniform.
-        redistribute:
-            Refined mesh is re-partitioned if ``True``.
-        ghost_mode: Ghost mode to use for the refined mesh.
+        partitioner:
+            partitioner to use for the refined mesh, If ``None`` no redistribution is performed,
+            i.e. previous local mesh is equally parallelized now with new vertices.s
         option: Controls whether parent cells and/or parent facets are computed.
 
-
     Returns:
        Refined mesh, (optional) parent cells, (optional) parent facets
     """
     mesh1, parent_cell, parent_facet = _cpp.refinement.refine(
-        mesh._cpp_object, edges, redistribute, ghost_mode, option
+        mesh._cpp_object, edges, partitioner, option
     )
     return Mesh(mesh1, mesh._ufl_domain), parent_cell, parent_facet
 

python/dolfinx/wrappers/refinement.cpp

@@ -4,40 +4,88 @@
 //
 // SPDX-License-Identifier:    LGPL-3.0-or-later
 
-#include "array.h"
 #include <concepts>
 #include <cstdint>
+#include <functional>
+#include <optional>
+#include <span>
+
+#include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/array.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/shared_ptr.h>
+#include <nanobind/stl/tuple.h>
+#include <nanobind/stl/vector.h>
+
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/MeshTags.h>
 #include <dolfinx/refinement/interval.h>
 #include <dolfinx/refinement/option.h>
 #include <dolfinx/refinement/plaza.h>
 #include <dolfinx/refinement/refine.h>
 #include <dolfinx/refinement/utils.h>
-#include <nanobind/nanobind.h>
-#include <nanobind/ndarray.h>
-#include <nanobind/stl/optional.h>
-#include <nanobind/stl/shared_ptr.h>
-#include <nanobind/stl/tuple.h>
-#include <nanobind/stl/vector.h>
-#include <optional>
-#include <span>
+
+#include "MPICommWrapper.h"
+#include "array.h"
 
 namespace nb = nanobind;
 
 namespace dolfinx_wrappers
 {
 
+namespace
+{
+
+using PythonCellPartitionFunction
+    = std::function<dolfinx::graph::AdjacencyList<std::int32_t>(
+        dolfinx_wrappers::MPICommWrapper, int,
+        const std::vector<dolfinx::mesh::CellType>&,
+        std::vector<nb::ndarray<const std::int64_t, nb::numpy>>)>;
+
+using CppCellPartitionFunction
+    = std::function<dolfinx::graph::AdjacencyList<std::int32_t>(
+        MPI_Comm, int, const std::vector<dolfinx::mesh::CellType>& q,
+        const std::vector<std::span<const std::int64_t>>&)>;
+
+/// Wrap a Python graph partitioning function as a C++ function
+CppCellPartitionFunction
+create_cell_partitioner_cpp(const PythonCellPartitionFunction& p)
+{
+  if (p)
+  {
+    return [p](MPI_Comm comm, int n,
+               const std::vector<dolfinx::mesh::CellType>& cell_types,
+               const std::vector<std::span<const std::int64_t>>& cells)
+    {
+      std::vector<nb::ndarray<const std::int64_t, nb::numpy>> cells_nb;
+      std::ranges::transform(
+          cells, std::back_inserter(cells_nb),
+          [](auto c)
+          {
+            return nb::ndarray<const std::int64_t, nb::numpy>(
+                c.data(), {c.size()}, nb::handle());
+          });
+
+      return p(dolfinx_wrappers::MPICommWrapper(comm), n, cell_types, cells_nb);
+    };
+  }
+  else
+    return nullptr;
+}
+} // namespace
+
 template <std::floating_point T>
 void export_refinement_with_variable_mesh_type(nb::module_& m)
 {
+
   m.def(
       "refine",
       [](const dolfinx::mesh::Mesh<T>& mesh,
          std::optional<
              nb::ndarray<const std::int32_t, nb::ndim<1>, nb::c_contig>>
              edges,
-         bool redistribute, dolfinx::mesh::GhostMode ghost_mode,
+         std::optional<PythonCellPartitionFunction> partitioner,
          dolfinx::refinement::Option option)
       {
         std::optional<std::span<const std::int32_t>> cpp_edges(std::nullopt);
@@ -47,8 +95,12 @@ void export_refinement_with_variable_mesh_type(nb::module_& m)
               std::span(edges.value().data(), edges.value().size()));
         }
 
+        CppCellPartitionFunction cpp_partitioner
+            = partitioner.has_value()
+                  ? create_cell_partitioner_cpp(partitioner.value())
+                  : dolfinx::refinement::maintain_coarse_partitioner;
         auto [mesh1, cell, facet] = dolfinx::refinement::refine(
-            mesh, cpp_edges, redistribute, ghost_mode, option);
+            mesh, cpp_edges, cpp_partitioner, option);
 
         std::optional<nb::ndarray<std::int32_t, nb::numpy>> python_cell(
             std::nullopt);
@@ -63,8 +115,8 @@ void export_refinement_with_variable_mesh_type(nb::module_& m)
         return std::tuple{std::move(mesh1), std::move(python_cell),
                           std::move(python_facet)};
       },
-      nb::arg("mesh"), nb::arg("edges") = nb::none(), nb::arg("redistribute"),
-      nb::arg("ghost_mode"), nb::arg("option"));
+      nb::arg("mesh"), nb::arg("edges") = nb::none(),
+      nb::arg("partitioner") = nb::none(), nb::arg("option"));
 }
 
 void refinement(nb::module_& m)

python/test/unit/refinement/test_interval.py

@@ -20,16 +20,14 @@
     "ghost_mode_refined",
     [mesh.GhostMode.none, mesh.GhostMode.shared_vertex, mesh.GhostMode.shared_facet],
 )
-@pytest.mark.parametrize("redistribute", [True, False])
 @pytest.mark.parametrize("option", [mesh.RefinementOption.none, mesh.RefinementOption.parent_cell])
-def test_refine_interval(n, ghost_mode, redistribute, ghost_mode_refined, option):
+def test_refine_interval(n, ghost_mode, ghost_mode_refined, option):
     msh = mesh.create_interval(MPI.COMM_WORLD, n, [0, 1], ghost_mode=ghost_mode)
     msh_refined, edges, vertices = mesh.refine(
         msh,
-        redistribute=redistribute,
-        ghost_mode=ghost_mode_refined,
         option=option,
     )
+    # TODO: add create_cell_partitioner(ghost_mode) when works
 
     # vertex count
     assert msh_refined.topology.index_map(0).size_global == 2 * n + 1
@@ -52,16 +50,14 @@ def test_refine_interval(n, ghost_mode, redistribute, ghost_mode_refined, option
     "ghost_mode_refined",
     [mesh.GhostMode.none, mesh.GhostMode.shared_vertex, mesh.GhostMode.shared_facet],
 )
-@pytest.mark.parametrize("redistribute", [True, False])
-def test_refine_interval_adaptive(n, ghost_mode, redistribute, ghost_mode_refined):
+def test_refine_interval_adaptive(n, ghost_mode, ghost_mode_refined):
     msh = mesh.create_interval(MPI.COMM_WORLD, n, [0, 1], ghost_mode=ghost_mode)
     msh_refined, edges, vertices = mesh.refine(
         msh,
         np.arange(10, dtype=np.int32),
-        redistribute=redistribute,
-        ghost_mode=ghost_mode_refined,
         option=mesh.RefinementOption.parent_cell,
     )
+    # TODO: add create_cell_partitioner(ghost_mode) when works
 
     # vertex count
     assert msh_refined.topology.index_map(0).size_global == n + 1 + 10 * MPI.COMM_WORLD.size