Introduce create_identity_partitioner for refined mesh partitioning

cpp/dolfinx/graph/partitioners.cpp

@@ -35,21 +35,8 @@ extern "C"
 
 using namespace dolfinx;
 
-namespace
-{
-/// @todo Is it un-documented that the owning rank must come first in
-/// reach list of edges?
-///
-/// @param[in] comm The communicator
-/// @param[in] graph Graph, using global indices for graph edges
-/// @param[in] node_disp The distribution of graph nodes across MPI
-/// ranks. The global index `gidx` of local index `lidx` is `lidx +
-/// node_disp[my_rank]`.
-/// @param[in] part The destination rank for owned nodes, i.e. `dest[i]`
-/// is the destination of the node with local index `i`.
-/// @return Destination ranks for each local node.
 template <typename T>
-graph::AdjacencyList<int> compute_destination_ranks(
+graph::AdjacencyList<int> dolfinx::graph::compute_destination_ranks(
     MPI_Comm comm, const graph::AdjacencyList<std::int64_t>& graph,
     const std::vector<T>& node_disp, const std::vector<T>& part)
 {
@@ -202,6 +189,18 @@ graph::AdjacencyList<int> compute_destination_ranks(
 
   return g;
 }
+
+/// @cond
+template graph::AdjacencyList<int> dolfinx::graph::compute_destination_ranks(
+    MPI_Comm comm, const graph::AdjacencyList<std::int64_t>& graph,
+    const std::vector<int>& node_disp, const std::vector<int>& part);
+
+template graph::AdjacencyList<int> dolfinx::graph::compute_destination_ranks(
+    MPI_Comm comm, const graph::AdjacencyList<std::int64_t>& graph,
+    const std::vector<unsigned long long>& node_disp,
+    const std::vector<unsigned long long>& part);
+/// @endcond
+
 //-----------------------------------------------------------------------------
 #ifdef HAS_PARMETIS
 template <typename T>
@@ -304,7 +303,6 @@ std::vector<int> refine(MPI_Comm comm, const graph::AdjacencyList<T>& adj_graph)
   //-----------------------------------------------------------------------------
 }
 #endif
-} // namespace
 
 //-----------------------------------------------------------------------------
 #ifdef HAS_PTSCOTCH
@@ -587,7 +585,7 @@ graph::partition_fn graph::parmetis::partitioner(double imbalance,
     {
       // FIXME: Is it implicit that the first entry is the owner?
       graph::AdjacencyList<int> dest
-          = compute_destination_ranks(pcomm, graph, node_disp, part);
+          = graph::compute_destination_ranks(pcomm, graph, node_disp, part);
       if (split_comm)
         MPI_Comm_free(&pcomm);
       return dest;
@@ -653,7 +651,7 @@ graph::partition_fn graph::kahip::partitioner(int mode, int seed,
     timer2.stop();
 
     if (ghosting)
-      return compute_destination_ranks(comm, graph, node_disp, part);
+      return graph::compute_destination_ranks(comm, graph, node_disp, part);
     else
     {
       return regular_adjacency_list(std::vector<int>(part.begin(), part.end()),

cpp/dolfinx/graph/partitioners.h

@@ -11,6 +11,23 @@
 
 namespace dolfinx::graph
 {
+
+/// @todo Is it un-documented that the owning rank must come first in
+/// reach list of edges?
+///
+/// @param[in] comm The communicator
+/// @param[in] graph Graph, using global indices for graph edges
+/// @param[in] node_disp The distribution of graph nodes across MPI
+/// ranks. The global index `gidx` of local index `lidx` is `lidx +
+/// node_disp[my_rank]`.
+/// @param[in] part The destination rank for owned nodes, i.e. `dest[i]`
+/// is the destination of the node with local index `i`.
+/// @return Destination ranks for each local node.
+template <typename T>
+graph::AdjacencyList<int> compute_destination_ranks(
+    MPI_Comm comm, const graph::AdjacencyList<std::int64_t>& graph,
+    const std::vector<T>& node_disp, const std::vector<T>& part);
+
 namespace scotch
 {
 #ifdef HAS_PTSCOTCH

cpp/dolfinx/refinement/refine.h

@@ -6,21 +6,90 @@
 
 #pragma once
 
+#include "dolfinx/common/MPI.h"
 #include "dolfinx/graph/AdjacencyList.h"
+#include "dolfinx/graph/partitioners.h"
 #include "dolfinx/mesh/Mesh.h"
 #include "dolfinx/mesh/Topology.h"
 #include "dolfinx/mesh/cell_types.h"
+#include "dolfinx/mesh/graphbuild.h"
 #include "dolfinx/mesh/utils.h"
 #include "interval.h"
 #include "plaza.h"
 #include <algorithm>
 #include <concepts>
+#include <mpi.h>
 #include <optional>
 #include <spdlog/spdlog.h>
+#include <stdexcept>
 #include <utility>
+#include <variant>
 
 namespace dolfinx::refinement
 {
+
+/**
+ * @brief Create a cell partitioner which maintains the partition of a coarse
+ * mesh.
+ *
+ * @tparam T floating point type of mesh geometry.
+ * @param parent_mesh mesh indicating the partition scheme to use, i.e. the
+ * coarse mesh.
+ * @param parent_cell list of cell indices mapping cells of the new refined mesh
+ * into the coarse mesh, usually output of `refinement::refine`.
+ * @return The created cell partition function.
+ */
+template <std::floating_point T>
+mesh::CellPartitionFunction
+create_identity_partitioner(const mesh::Mesh<T>& parent_mesh,
+                            std::span<std::int32_t> parent_cell)
+{
+  // TODO: optimize for non ghosted mesh?
+
+  return [&parent_mesh,
+          parent_cell](MPI_Comm comm, int /*nparts*/,
+                       std::vector<mesh::CellType> cell_types,
+                       std::vector<std::span<const std::int64_t>> cells)
+             -> graph::AdjacencyList<std::int32_t>
+  {
+    const auto parent_cell_im
+        = parent_mesh.topology()->index_map(parent_mesh.topology()->dim());
+    std::int32_t parent_num_cells = parent_cell_im->size_local();
+    std::span parent_cell_owners = parent_cell_im->owners();
+
+    std::int32_t num_cells
+        = cells.front().size() / mesh::num_cell_vertices(cell_types.front());
+    std::vector<std::int32_t> destinations(num_cells);
+
+    int rank = dolfinx::MPI::rank(comm);
+    for (std::size_t i = 0; i < destinations.size(); i++)
+    {
+      bool parent_is_ghost_cell = parent_cell[i] > parent_num_cells;
+      if (parent_is_ghost_cell)
+        destinations[i] = parent_cell_owners[parent_cell[i] - parent_num_cells];
+      else
+        destinations[i] = rank;
+    }
+
+    if (comm == MPI_COMM_NULL)
+      return graph::regular_adjacency_list(std::move(destinations), 1);
+
+    auto dual_graph = mesh::build_dual_graph(comm, cell_types, cells);
+    std::vector<std::int32_t> node_disp(MPI::size(comm) + 1, 0);
+    std::int32_t local_size = dual_graph.num_nodes();
+    MPI_Allgather(&local_size, 1, dolfinx::MPI::mpi_t<std::int32_t>,
+                  node_disp.data() + 1, 1, dolfinx::MPI::mpi_t<std::int32_t>,
+                  comm);
+    std::partial_sum(node_disp.begin(), node_disp.end(), node_disp.begin());
+    return compute_destination_ranks(comm, dual_graph, node_disp, destinations);
+  };
+}
+
+/// @brief Placeholder for the creation of an identity partitioner in refine.
+struct IdentityPartitionerPlaceholder
+{
+};
+
 /// @brief Refine a mesh with markers.
 ///
 /// The refined mesh can be optionally re-partitioned across processes.
@@ -35,30 +104,28 @@ namespace dolfinx::refinement
 /// refined mesh will **not** include ghosts cells (cells connected by
 /// facet to an owned cell) even if the parent mesh is ghosted.
 ///
-/// @warning Passing `nullptr` for `partitioner`, the refined mesh will
-/// **not** have ghosts cells even if the parent mesh is ghosted. The
-/// possibility to not re-partition the refined mesh and include ghost
-/// cells in the refined mesh will be added in a future release.
 ///
 /// @param[in] mesh Input mesh to be refined.
 /// @param[in] edges Indices of the edges that should be split in the
 /// refinement. If not provided (`std::nullopt`), uniform refinement is
 /// performed.
-/// @param[in] partitioner Partitioner to be used to distribute the
-/// refined mesh. If not callable, refined cells will be on the same
+/// @param[in] partitioner (Optional) partitioner to be used to distribute the
+/// refined mesh. If not provided (`std::nullopt`) the partition of the coarse
+/// mesh will be maintained. If not callable, refined cells will be on the same
 /// process as the parent cell.
 /// @param[in] option Control the computation of parent facets, parent
-/// cells. If an option is not selected, an empty list is returned.
+/// cells.
 /// @return New mesh, and optional parent cell indices and parent facet
 /// indices.
 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,
-       const mesh::CellPartitionFunction& partitioner
-       = mesh::create_cell_partitioner(mesh::GhostMode::none),
-       Option option = Option::none)
+       std::variant<IdentityPartitionerPlaceholder, mesh::CellPartitionFunction>
+           partitioner
+       = IdentityPartitionerPlaceholder(),
+       Option option = Option::parent_cell)
 {
   auto topology = mesh.topology();
   assert(topology);
@@ -70,9 +137,23 @@ refine(const mesh::Mesh<T>& mesh,
             ? interval::compute_refinement_data(mesh, edges, option)
             : plaza::compute_refinement_data(mesh, edges, option);
 
+  if (std::holds_alternative<IdentityPartitionerPlaceholder>(partitioner))
+  {
+    if (!parent_cell)
+    {
+      throw std::runtime_error(
+          "Identity partitioner relies on parent cell computation");
+    }
+    assert(parent_cell);
+    partitioner = create_identity_partitioner(mesh, parent_cell.value());
+  }
+
+  assert(std::holds_alternative<mesh::CellPartitionFunction>(partitioner));
+
   mesh::Mesh<T> mesh1 = mesh::create_mesh(
       mesh.comm(), mesh.comm(), cell_adj.array(), mesh.geometry().cmap(),
-      mesh.comm(), new_vertex_coords, xshape, partitioner);
+      mesh.comm(), new_vertex_coords, xshape,
+      std::get<mesh::CellPartitionFunction>(partitioner));
 
   // Report the number of refined cells
   const int D = topology->dim();
@@ -84,4 +165,5 @@ refine(const mesh::Mesh<T>& mesh,
 
   return {std::move(mesh1), std::move(parent_cell), std::move(parent_facet)};
 }
+
 } // namespace dolfinx::refinement

python/dolfinx/mesh.py

@@ -31,7 +31,7 @@
     to_string,
     to_type,
 )
-from dolfinx.cpp.refinement import RefinementOption
+from dolfinx.cpp.refinement import IdentityPartitionerPlaceholder, RefinementOption
 from dolfinx.fem import CoordinateElement as _CoordinateElement
 from dolfinx.fem import coordinate_element as _coordinate_element
 from dolfinx.graph import AdjacencyList
@@ -550,8 +550,10 @@ def transfer_meshtag(
 def refine(
     msh: Mesh,
     edges: typing.Optional[np.ndarray] = None,
-    partitioner: typing.Optional[typing.Callable] = create_cell_partitioner(GhostMode.none),
-    option: RefinementOption = RefinementOption.none,
+    partitioner: typing.Union[
+        typing.Callable, IdentityPartitionerPlaceholder
+    ] = IdentityPartitionerPlaceholder(),
+    option: RefinementOption = RefinementOption.parent_cell,
 ) -> tuple[Mesh, npt.NDArray[np.int32], npt.NDArray[np.int8]]:
     """Refine a mesh.
 
@@ -563,20 +565,17 @@ def refine(
         mesh will **not** include ghosts cells (cells connected by facet
         to an owned cells) even if the parent mesh is ghosted.
 
-    Warning:
-        Passing ``None`` for ``partitioner``, the refined mesh will
-        **not** have ghosts cells even if the parent mesh has ghost
-        cells. The possibility to not re-partition the refined mesh and
-        include ghost cells in the refined mesh will be added in a
-        future release.
-
     Args:
         msh: Mesh from which to create the refined mesh.
         edges: Indices of edges to split during refinement. If ``None``,
             mesh refinement is uniform.
-        partitioner: Partitioner to distribute the refined mesh. If
-            ``None`` no redistribution is performed, i.e. refined cells
-            remain on the same process as the parent cell.
+        partitioner: Partitioner to distribute the refined mesh. If a
+            ``IdentityPartitionerPlaceholder`` is passed (default) no
+            redistribution is performed, i.e. refined cells remain on the
+            same process as the parent cell, but the ghost layer is
+            updated. If a custom partitioner is passed, it will be used for
+            distributing the refined mesh. If ``None`` is passed no
+            redistribution will happen.
         option: Controls whether parent cells and/or parent facets are
             computed.
 

python/dolfinx/wrappers/refinement.cpp

@@ -1,4 +1,5 @@
-// Copyright (C) 2018-2024 Chris N. Richardson and Garth N. Wells
+// Copyright (C) 2018-2024 Chris N. Richardson, Garth N. Wells and Paul T.
+// Kühner
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -23,9 +24,12 @@
 #include <nanobind/stl/optional.h>
 #include <nanobind/stl/shared_ptr.h>
 #include <nanobind/stl/tuple.h>
+#include <nanobind/stl/variant.h>
 #include <nanobind/stl/vector.h>
 #include <optional>
 #include <span>
+#include <stdexcept>
+#include <variant>
 
 namespace nb = nanobind;
 
@@ -45,8 +49,9 @@ void export_refinement(nb::module_& m)
          std::optional<
              nb::ndarray<const std::int32_t, nb::ndim<1>, nb::c_contig>>
              edges,
-         std::optional<
-             dolfinx_wrappers::part::impl::PythonCellPartitionFunction>
+         std::variant<dolfinx::refinement::IdentityPartitionerPlaceholder,
+                      std::optional<dolfinx_wrappers::part::impl::
+                                        PythonCellPartitionFunction>>
              partitioner,
          dolfinx::refinement::Option option)
       {
@@ -57,11 +62,29 @@ void export_refinement(nb::module_& m)
               std::span(edges.value().data(), edges.value().size()));
         }
 
-        dolfinx_wrappers::part::impl::CppCellPartitionFunction cpp_partitioner
-            = partitioner.has_value()
-                  ? dolfinx_wrappers::part::impl::create_cell_partitioner_cpp(
-                        partitioner.value())
-                  : nullptr;
+        std::variant<dolfinx::refinement::IdentityPartitionerPlaceholder,
+                     dolfinx_wrappers::part::impl::CppCellPartitionFunction>
+            cpp_partitioner
+            = dolfinx::refinement::IdentityPartitionerPlaceholder();
+        if (std::holds_alternative<std::optional<
+                dolfinx_wrappers::part::impl::PythonCellPartitionFunction>>(
+                partitioner))
+        {
+          auto optional = std::get<std::optional<
+              dolfinx_wrappers::part::impl::PythonCellPartitionFunction>>(
+              partitioner);
+          if (!optional.has_value())
+            cpp_partitioner
+                = dolfinx_wrappers::part::impl::CppCellPartitionFunction(
+                    nullptr);
+          else
+          {
+            cpp_partitioner
+                = dolfinx_wrappers::part::impl::create_cell_partitioner_cpp(
+                    optional.value());
+          }
+        }
+
         auto [mesh1, cell, facet] = dolfinx::refinement::refine(
             mesh, cpp_edges, cpp_partitioner, option);
 
@@ -97,6 +120,10 @@ void refinement(nb::module_& m)
   export_refinement<float>(m);
   export_refinement<double>(m);
 
+  nb::class_<dolfinx::refinement::IdentityPartitionerPlaceholder>(
+      m, "IdentityPartitionerPlaceholder")
+      .def(nb::init<>());
+
   nb::enum_<dolfinx::refinement::Option>(m, "RefinementOption")
       .value("none", dolfinx::refinement::Option::none)
       .value("parent_facet", dolfinx::refinement::Option::parent_facet)

python/test/unit/refinement/test_interval.py

@@ -22,8 +22,7 @@
 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,
-        option=option,
+        msh, option=option, partitioner=mesh.create_cell_partitioner(ghost_mode_refined)
     )
     # TODO: add create_cell_partitioner(ghost_mode) when works