Pass graph partitioner to mesh::refine (#3444)

* Remove refine option `distribute` and `ghost_mode` in favor of an explicit partitioner

* Remove wrongly added transfer_matrix -> later PR

* Add author name

* Doc fix

* Fix typi

* Move maintain_coarse_partitioner to refine.cpp

* Works better with the actual file

* Add doc string to partitioner

* Make doxygen happy

* Switch to general partitioner naming

* Small updates

* Fix test failure

* Simplifications

* Small fix

* Small fix

* Work on docs

* Simplify

* Improve function name

* Namespace fix

* Compile fix

* Add nanobind include

* Small re-factor

* Fixes

* More updates

* Fix broken tests

* Doc update

* Update test

* Test updates

* Work on docs

* Lint

* Minor doc updates

---------

Co-authored-by: schnellerhase <[email protected]>
Co-authored-by: Jørgen Schartum Dokken <[email protected]>

cpp/dolfinx/mesh/MeshTags.h

@@ -127,8 +127,8 @@ class MeshTags
 };
 
 /// @brief Create MeshTags from arrays
-/// @param[in] topology Mesh topology that the tags are associated with
-/// @param[in] dim Topological dimension of tagged entities
+/// @param[in] topology Mesh topology that the tags are associated with.
+/// @param[in] dim Topological dimension of tagged entities.
 /// @param[in] entities Local vertex indices for tagged entities.
 /// @param[in] values Tag values for each entity in `entities`. The
 /// length of `values` must be equal to number of rows in `entities`.

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"
@@ -15,85 +16,72 @@
 #include <algorithm>
 #include <concepts>
 #include <optional>
+#include <spdlog/spdlog.h>
 #include <utility>
 
 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)
-{
- 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);
- }
-}
-} // namespace impl
-
-/// @brief Refine with markers, optionally redistributing, and
-/// optionally calculating the parent-child relationships.
+/// @brief Refine a mesh with markers.
+///
+/// The refined mesh can be optionally re-partitioned across processes.
+/// Passing `nullptr` for `partitioner`, refined cells will be on the
+/// same process as the parent cell.
 ///
-/// @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.
+/// Parent-child relationships can be optionally computed. Parent-child
+/// relationships can be used to create MeshTags on the refined mesh
+/// from MeshTags on the parent mesh.
+///
+/// @warning Using the default partitioner for a refined mesh, the
+/// 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
+/// process as the parent cell.
 /// @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
+/// cells. If an option is not selected, an empty list is returned.
+/// @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, bool redistribute,
- mesh::GhostMode ghost_mode = mesh::GhostMode::shared_facet,
+ std::optional<std::span<const std::int32_t>> edges,
+ const mesh::CellPartitionFunction& partitioner
+ = mesh::create_cell_partitioner(mesh::GhostMode::none),
  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);
+ = (topology->cell_type() == mesh::CellType::interval)
+ ? 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> mesh1 = mesh::create_mesh(
+ mesh.comm(), mesh.comm(), cell_adj.array(), mesh.geometry().cmap(),
+ mesh.comm(), 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();
+ const std::int64_t n1 = mesh1.topology()->index_map(D)->size_global();
  spdlog::info(
  "Number of cells increased from {} to {} ({}% increase).", n0, n1,
  100.0 * (static_cast<double>(n1) / static_cast<double>(n0) - 1.0));
 
- return {std::move(refined_mesh), std::move(parent_cell),
- std::move(parent_facet)};
+ return {std::move(mesh1), std::move(parent_cell), std::move(parent_facet)};
 }
-
 } // namespace dolfinx::refinement

cpp/dolfinx/refinement/utils.h

@@ -264,85 +264,48 @@ 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
+/// @todo Improve docstring.
+///
+/// @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
-/// addition of new local indices.
-/// @param[in] map Index map
-/// @param[in] n Number of new local indices
-/// @return New global indices for both owned and ghosted indices in input
-/// index map.
+/// @note The returned global indices (local and ghosts) are adjust for
+/// the addition of new local indices.
+///
+/// @param[in] map Index map.
+/// @param[in] n Number of new local indices.
+/// @return New global indices for both owned and ghosted indices in
+/// input index map.
 std::vector<std::int64_t> adjust_indices(const common::IndexMap& map,
  std::int32_t n);
 
-/// @brief Transfer facet MeshTags from coarse mesh to refined mesh
-/// @note The refined mesh must not have been redistributed during
-/// refinement
-/// @note GhostMode must be GhostMode.none
-/// @param[in] tags0 Tags on the parent mesh
-/// @param[in] topology1 Refined mesh topology
+/// @brief Transfer facet MeshTags from coarse mesh to refined mesh.
+///
+/// @warning The refined mesh must not have been redistributed during
+/// refinement.
+///
+/// @warning Mesh/topology `GhostMode` must be mesh::GhostMode::none.
+///
+/// @param[in] tags0 Tags on the parent mesh.
+/// @param[in] topology1 Refined mesh topology.
 /// @param[in] cell Parent cell of each cell in refined mesh
-/// @param[in] facet Local facets of parent in each cell in refined mesh
-/// @return (0) entities and (1) values on the refined topology
+/// @param[in] facet Local facets of parent in each cell in refined mesh.
+/// @return (0) entities and (1) values on the refined topology.
 std::array<std::vector<std::int32_t>, 2> transfer_facet_meshtag(
  const mesh::MeshTags<std::int32_t>& tags0, const mesh::Topology& topology1,
  std::span<const std::int32_t> cell, std::span<const std::int8_t> facet);
 
 /// @brief Transfer cell MeshTags from coarse mesh to refined mesh.
 ///
-/// @note The refined mesh must not have been redistributed during
+/// @warning The refined mesh must not have been redistributed during
 /// refinement.
-/// @note GhostMode must be GhostMode.none
 ///
-/// @param[in] tags0 Tags on the parent mesh
-/// @param[in] topology1 Refined mesh topology
-/// @param[in] parent_cell Parent cell of each cell in refined mesh
-/// @return (0) entities and (1) values on the refined topology
+/// @warning Mesh/topology `GhostMode` must be mesh::GhostMode::none.
+///
+/// @param[in] tags0 Tags on the parent mesh.
+/// @param[in] topology1 Refined mesh topology.
+/// @param[in] parent_cell Parent cell of each cell in refined mesh.
+/// @return (0) entities and (1) values on the refined topology.
 std::array<std::vector<std::int32_t>, 2>
 transfer_cell_meshtag(const mesh::MeshTags<std::int32_t>& tags0,
  const mesh::Topology& topology1,

cpp/test/mesh/refinement/interval.cpp

@@ -73,8 +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,
- refinement::Option::parent_cell);
+ mesh, std::nullopt, nullptr, refinement::Option::parent_cell);
 
  std::vector<T> expected_x = {
  /* v_0 */ 0.0, 0.0, 0.0,
@@ -114,7 +113,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 = {
@@ -190,15 +190,13 @@ TEMPLATE_TEST_CASE("Interval Refinement (parallel)",
  mesh::Mesh<T> mesh = create_mesh();
  mesh.topology()->create_connectivity(1, 0);
 
- // TODO: parent_facet
  auto [refined_mesh, parent_edges, parent_facet] = refinement::refine(
- mesh, std::nullopt, false, mesh::GhostMode::shared_facet,
- refinement::Option::parent_cell);
+ mesh, std::nullopt, nullptr, refinement::Option::parent_cell);
 
  T rank_d = static_cast<T>(rank);
  T comm_size_d = static_cast<T>(comm_size);
 
- auto x = refined_mesh.geometry().x();
+ std::span x = refined_mesh.geometry().x();
  std::ranges::sort(x);
  std::vector<T> expected_x
  = {rank_d / comm_size_d,

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