diff --git a/cpp/dolfinx/refinement/CMakeLists.txt b/cpp/dolfinx/refinement/CMakeLists.txt index 9560e45c82..b11546361d 100644 --- a/cpp/dolfinx/refinement/CMakeLists.txt +++ b/cpp/dolfinx/refinement/CMakeLists.txt @@ -11,4 +11,5 @@ set(HEADERS_refinement target_sources( dolfinx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/plaza.cpp ${CMAKE_CURRENT_SOURCE_DIR}/utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/refine.cpp ) diff --git a/cpp/dolfinx/refinement/refine.cpp b/cpp/dolfinx/refinement/refine.cpp new file mode 100644 index 0000000000..3fba0e2410 --- /dev/null +++ b/cpp/dolfinx/refinement/refine.cpp @@ -0,0 +1,22 @@ +// Copyright (C) 2024 Paul T. Kühner +// +// This file is part of DOLFINx (https://www.fenicsproject.org) +// +// SPDX-License-Identifier: LGPL-3.0-or-later + +#include "refine.h" +#include + +using namespace dolfinx; + +graph::AdjacencyList refinement::maintain_coarse_partitioner( + MPI_Comm comm, int /*nparts*/, std::vector cell_types, + std::vector> cells) +{ + int num_cell_vertices = mesh::num_cell_vertices(cell_types.front()); + std::int32_t num_cells = cells.front().size() / num_cell_vertices; + std::vector destinations(num_cells, dolfinx::MPI::rank(comm)); + std::vector dest_offsets(num_cells + 1); + std::iota(dest_offsets.begin(), dest_offsets.end(), 0); + return graph::AdjacencyList(std::move(destinations), std::move(dest_offsets)); +} diff --git a/cpp/dolfinx/refinement/refine.h b/cpp/dolfinx/refinement/refine.h index 32f18d8982..07a5b411a8 100644 --- a/cpp/dolfinx/refinement/refine.h +++ b/cpp/dolfinx/refinement/refine.h @@ -1,4 +1,4 @@ -// 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) // @@ -6,6 +6,7 @@ #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,28 @@ namespace dolfinx::refinement { -namespace impl -{ -/// @brief create the refined mesh by optionally redistributing it -template -mesh::Mesh -create_refined_mesh(const mesh::Mesh& mesh, - const graph::AdjacencyList& cell_adj, - std::span new_vertex_coords, - std::array 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(mesh, cell_adj, new_vertex_coords, xshape, redistribute, - ghost_mode); - } -} -} // namespace impl +/// @brief Partitioner that maintains the regional distribution for a +/// refined mesh. Meaning, process local data is just maintained and no +/// redistribution happens. +/// +/// @param[in] comm MPI Communicator +/// @param[in] nparts Number of partitions (input has no effect) +/// @param[in] cell_types Cell types in the mesh +/// @param[in] cells Lists of cells of each cell type. +/// @return Destination ranks for each cell on this process +graph::AdjacencyList +maintain_coarse_partitioner(MPI_Comm comm, int nparts, + std::vector cell_types, + std::vector> cells); /// @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 +50,26 @@ template std::tuple, std::optional>, std::optional>> refine(const mesh::Mesh& mesh, - std::optional> edges, bool redistribute, - mesh::GhostMode ghost_mode = mesh::GhostMode::shared_facet, + std::optional> 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 refined_mesh = impl::create_refined_mesh( - mesh, cell_adj, std::span(new_vertex_coords), xshape, - redistribute, ghost_mode); + mesh::Mesh 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(); diff --git a/cpp/dolfinx/refinement/utils.h b/cpp/dolfinx/refinement/utils.h index 275b40d7cc..1c6af07755 100644 --- a/cpp/dolfinx/refinement/utils.h +++ b/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 -mesh::Mesh partition(const mesh::Mesh& old_mesh, - const graph::AdjacencyList& cell_topology, - std::span new_coords, - std::array 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& cell_types, - const std::vector>& 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 destinations(num_cells, mpi_rank); - std::vector 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 diff --git a/cpp/test/mesh/refinement/interval.cpp b/cpp/test/mesh/refinement/interval.cpp index 763418ac4b..74b3f43462 100644 --- a/cpp/test/mesh/refinement/interval.cpp +++ b/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 expected_x = { @@ -114,7 +114,8 @@ TEMPLATE_TEST_CASE("Interval adaptive refinement", std::vector 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 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(rank); diff --git a/cpp/test/mesh/refinement/rectangle.cpp b/cpp/test/mesh/refinement/rectangle.cpp index 954879abb1..ff2a67979d 100644 --- a/cpp/test/mesh/refinement/rectangle.cpp +++ b/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 diff --git a/python/dolfinx/mesh.py b/python/dolfinx/mesh.py index 28b9eb759e..bac5232cbb 100644 --- a/python/dolfinx/mesh.py +++ b/python/dolfinx/mesh.py @@ -535,8 +535,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. @@ -545,17 +544,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 ) # Create new ufl domain as it will carry a reference to the C++ mesh in the ufl_cargo ufl_domain = ufl.Mesh(mesh._ufl_domain.ufl_coordinate_element()) # type: ignore diff --git a/python/dolfinx/wrappers/refinement.cpp b/python/dolfinx/wrappers/refinement.cpp index 37ca393799..8bb89b7260 100644 --- a/python/dolfinx/wrappers/refinement.cpp +++ b/python/dolfinx/wrappers/refinement.cpp @@ -4,6 +4,7 @@ // // SPDX-License-Identifier: LGPL-3.0-or-later +#include "MPICommWrapper.h" #include "array.h" #include #include @@ -14,8 +15,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -27,17 +30,57 @@ namespace nb = nanobind; namespace dolfinx_wrappers { +namespace +{ +using PythonCellPartitionFunction + = std::function( + dolfinx_wrappers::MPICommWrapper, int, + const std::vector&, + std::vector>)>; + +using CppCellPartitionFunction + = std::function( + MPI_Comm, int, const std::vector& q, + const std::vector>&)>; + +/// 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& cell_types, + const std::vector>& cells) + { + std::vector> cells_nb; + std::ranges::transform( + cells, std::back_inserter(cells_nb), + [](auto c) + { + return nb::ndarray( + c.data(), {c.size()}, nb::handle()); + }); + + return p(dolfinx_wrappers::MPICommWrapper(comm), n, cell_types, cells_nb); + }; + } + else + return nullptr; +} +} // namespace template void export_refinement_with_variable_mesh_type(nb::module_& m) { + m.def( "refine", [](const dolfinx::mesh::Mesh& mesh, std::optional< nb::ndarray, nb::c_contig>> edges, - bool redistribute, dolfinx::mesh::GhostMode ghost_mode, + std::optional partitioner, dolfinx::refinement::Option option) { std::optional> cpp_edges(std::nullopt); @@ -47,8 +90,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> python_cell( std::nullopt); @@ -63,8 +110,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) diff --git a/python/test/unit/refinement/test_interval.py b/python/test/unit/refinement/test_interval.py index 2877964484..27040fb0b9 100644 --- a/python/test/unit/refinement/test_interval.py +++ b/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 diff --git a/python/test/unit/refinement/test_refinement.py b/python/test/unit/refinement/test_refinement.py index 6aaa4090df..202fed2339 100644 --- a/python/test/unit/refinement/test_refinement.py +++ b/python/test/unit/refinement/test_refinement.py @@ -32,7 +32,7 @@ def test_refine_create_unit_square(): """Refine mesh of unit square.""" mesh = create_unit_square(MPI.COMM_WORLD, 5, 7, ghost_mode=GhostMode.none) mesh.topology.create_entities(1) - mesh_refined, _, _ = refine(mesh, redistribute=False) + mesh_refined, _, _ = refine(mesh) assert mesh_refined.topology.index_map(0).size_global == 165 assert mesh_refined.topology.index_map(2).size_global == 280 @@ -46,7 +46,7 @@ def test_refine_create_unit_cube(ghost_mode, redistribute): """Refine mesh of unit cube.""" mesh = create_unit_cube(MPI.COMM_WORLD, 5, 7, 9, ghost_mode=ghost_mode) mesh.topology.create_entities(1) - mesh, _, _ = refine(mesh, redistribute=redistribute) + mesh, _, _ = refine(mesh) # , partitioner=create_cell_partitioner(ghost_mode)) assert mesh.topology.index_map(0).size_global == 3135 assert mesh.topology.index_map(3).size_global == 15120 @@ -58,7 +58,7 @@ def test_refine_create_form(): """Check that forms can be assembled on refined mesh""" mesh = create_unit_cube(MPI.COMM_WORLD, 3, 3, 3) mesh.topology.create_entities(1) - mesh, _, _ = refine(mesh, redistribute=True) + mesh, _, _ = refine(mesh) # TODO: recover redistribute=True behavior V = functionspace(mesh, ("Lagrange", 1)) @@ -83,7 +83,7 @@ def left_corner_edge(x, tol=1e-7): if MPI.COMM_WORLD.size == 0: assert edges == 1 - mesh2, _, _ = refine(mesh, edges, redistribute=False) + mesh2, _, _ = refine(mesh, edges) assert mesh.topology.index_map(2).size_global + 3 == mesh2.topology.index_map(2).size_global @@ -105,7 +105,7 @@ def left_side(x, tol=1e-16): edges = compute_incident_entities(msh.topology, cells, 2, 1) if MPI.COMM_WORLD.size == 0: assert edges.__len__() == Nx // 2 * (2 * Ny + 1) + (Nx // 2 + 1) * Ny - mesh2, _, _ = refine(msh, edges, redistribute=True) + mesh2, _, _ = refine(msh, edges) # TODO: redistribute=True num_cells_global = mesh2.topology.index_map(2).size_global actual_cells = 3 * (Nx * Ny) + 3 * Ny + 2 * Nx * Ny @@ -116,13 +116,10 @@ def left_side(x, tol=1e-16): @pytest.mark.parametrize( "refine_plaza_wrapper", [ - lambda mesh: refine( - mesh, redistribute=False, option=RefinementOption.parent_cell_and_facet - ), + lambda mesh: refine(mesh, option=RefinementOption.parent_cell_and_facet), lambda mesh: refine( mesh, np.arange(mesh.topology.index_map(1).size_local), - redistribute=False, option=RefinementOption.parent_cell_and_facet, ), ], @@ -178,13 +175,10 @@ def test_refine_facet_meshtag(tdim, refine_plaza_wrapper): @pytest.mark.parametrize( "refine_plaza_wrapper", [ - lambda mesh: refine( - mesh, redistribute=False, option=RefinementOption.parent_cell_and_facet - ), + lambda mesh: refine(mesh, option=RefinementOption.parent_cell_and_facet), lambda mesh: refine( mesh, np.arange(mesh.topology.index_map(1).size_local), - redistribute=False, option=RefinementOption.parent_cell_and_facet, ), ], @@ -217,5 +211,5 @@ def test_refine_cell_meshtag(tdim, refine_plaza_wrapper): def test_refine_ufl_cargo(): mesh = create_unit_cube(MPI.COMM_WORLD, 4, 3, 3) mesh.topology.create_entities(1) - refined_mesh, _, _ = refine(mesh, redistribute=False) + refined_mesh, _, _ = refine(mesh) assert refined_mesh.ufl_domain().ufl_cargo() != mesh.ufl_domain().ufl_cargo()