diff --git a/cpp/dolfinx/mesh/utils.h b/cpp/dolfinx/mesh/utils.h index 12d334fdf75..a77112a043d 100644 --- a/cpp/dolfinx/mesh/utils.h +++ b/cpp/dolfinx/mesh/utils.h @@ -747,7 +747,7 @@ compute_incident_entities(const Topology& topology, /// 'nodes' will be included. See dolfinx::io::cells for examples of the /// Basix ordering. /// @param[in] element Coordinate element for the cells. -/// @param[in] commg +/// @param[in] commg Communicator for geometry /// @param[in] x Geometry data ('node' coordinates). Row-major storage. /// The global index of the `i`th node (row) in `x` is taken as `i` plus /// the process offset on`comm`, The offset is the sum of `x` rows on @@ -888,6 +888,237 @@ Mesh> create_mesh( std::move(geometry)); } +/// @brief Create a distributed mixed-topology mesh from mesh data using a +/// provided graph partitioning function for determining the parallel +/// distribution of the mesh. +/// +/// From mesh input data that is distributed across processes, a +/// distributed mesh::Mesh is created. If the partitioning function is +/// not callable, i.e. it does not store a callable function, no +/// re-distribution of cells is done. +/// +/// @note This is an experimental specialised version of `create_mesh` for mixed +/// topology meshes, and does not include cell reordering. +/// +/// @param[in] comm Communicator to build the mesh on. +/// @param[in] commt Communicator that the topology data (`cells`) is +/// distributed on. This should be `MPI_COMM_NULL` for ranks that should +/// not participate in computing the topology partitioning. +/// @param[in] cells Cells on the calling process, as a list of lists, +/// one list for each cell type (or an empty list if there are no cells of that +/// type on this process). The cells are defined by their 'nodes' (using global +/// indices) following the Basix ordering, and concatenated to form a flattened +/// list. For lowest order cells this will be just the cell vertices. For +/// higher-order cells, other cells 'nodes' will be included. See +/// dolfinx::io::cells for examples of the Basix ordering. +/// @param[in] elements Coordinate elements for the cells, one for each cell +/// type in the mesh. In parallel, these must be the same on all processes. +/// @param[in] commg Communicator for geometry +/// @param[in] x Geometry data ('node' coordinates). Row-major storage. +/// The global index of the `i`th node (row) in `x` is taken as `i` plus +/// the process offset on`comm`, The offset is the sum of `x` rows on +/// all processed with a lower rank than the caller. +/// @param[in] xshape Shape of the `x` data. +/// @param[in] partitioner Graph partitioner that computes the owning +/// rank for each cell. If not callable, cells are not redistributed. +/// @return A mesh distributed on the communicator `comm`. +template +Mesh> create_mesh( + MPI_Comm comm, MPI_Comm commt, + const std::vector>& cells, + const std::vector>>& elements, + MPI_Comm commg, const U& x, std::array xshape, + const CellPartitionFunction& partitioner) +{ + assert(cells.size() == elements.size()); + std::int32_t num_cell_types = cells.size(); + std::vector celltypes; + std::transform(elements.cbegin(), elements.cend(), + std::back_inserter(celltypes), + [](auto e) { return e.cell_shape(); }); + std::vector doflayouts; + std::transform(elements.cbegin(), elements.cend(), + std::back_inserter(doflayouts), + [](auto e) { return e.create_dof_layout(); }); + + // Note: `extract_topology` extracts topology data, i.e. just the + // vertices. For P1 geometry this should just be the identity + // operator. For other elements the filtered lists may have 'gaps', + // i.e. the indices might not be contiguous. + // + // `extract_topology` could be skipped for 'P1 geometry' elements + + // -- Partition topology across ranks of comm + std::vector> cells1(num_cell_types); + std::vector> original_idx1(num_cell_types); + std::vector> ghost_owners(num_cell_types); + if (partitioner) + { + spdlog::info("Using partitioner with cell data ({} cell types)", + num_cell_types); + graph::AdjacencyList dest(0); + if (commt != MPI_COMM_NULL) + { + int size = dolfinx::MPI::size(comm); + std::vector> t(num_cell_types); + std::vector> tspan(num_cell_types); + for (std::int32_t i = 0; i < num_cell_types; ++i) + { + t[i] = extract_topology(celltypes[i], doflayouts[i], cells[i]); + tspan[i] = std::span(t[i]); + } + dest = partitioner(commt, size, celltypes, tspan); + } + + std::int32_t cell_offset = 0; + for (std::int32_t i = 0; i < num_cell_types; ++i) + { + std::size_t num_cell_nodes = doflayouts[i].num_dofs(); + assert(cells[i].size() % num_cell_nodes == 0); + std::size_t num_cells = cells[i].size() / num_cell_nodes; + + // Extract destination AdjacencyList for this cell type + std::vector offsets_i( + std::next(dest.offsets().begin(), cell_offset), + std::next(dest.offsets().begin(), cell_offset + num_cells + 1)); + std::vector data_i( + std::next(dest.array().begin(), offsets_i.front()), + std::next(dest.array().begin(), offsets_i.back())); + std::int32_t offset_0 = offsets_i.front(); + std::for_each(offsets_i.begin(), offsets_i.end(), + [&offset_0](std::int32_t& j) { j -= offset_0; }); + graph::AdjacencyList dest_i(data_i, offsets_i); + cell_offset += num_cells; + + // Distribute cells (topology, includes higher-order 'nodes') to + // destination rank + std::tie(cells1[i], original_idx1[i], ghost_owners[i]) + = graph::build::distribute(comm, cells[i], + {num_cells, num_cell_nodes}, dest_i); + spdlog::debug("Got {} cells from distribution", cells1[i].size()); + } + } + else + { + // No partitioning, construct a global index + std::int64_t num_owned = 0; + for (std::int32_t i = 0; i < num_cell_types; ++i) + { + cells1[i] = std::vector(cells[i].begin(), cells[i].end()); + std::int32_t num_cell_nodes = doflayouts[i].num_dofs(); + assert(cells1[i].size() % num_cell_nodes == 0); + original_idx1[i].resize(cells1[i].size() / num_cell_nodes); + num_owned += original_idx1[i].size(); + } + // Add on global offset + std::int64_t global_offset = 0; + MPI_Exscan(&num_owned, &global_offset, 1, MPI_INT64_T, MPI_SUM, comm); + for (std::int32_t i = 0; i < num_cell_types; ++i) + { + std::iota(original_idx1[i].begin(), original_idx1[i].end(), + global_offset); + global_offset += original_idx1[i].size(); + } + } + + // Extract cell 'topology', i.e. extract the vertices for each cell + // and discard any 'higher-order' nodes + std::vector> cells1_v(num_cell_types); + for (std::int32_t i = 0; i < num_cell_types; ++i) + { + cells1_v[i] = extract_topology(celltypes[i], doflayouts[i], cells1[i]); + spdlog::info("Extract basic topology: {}->{}", cells1[i].size(), + cells1_v[i].size()); + } + + // Build local dual graph for owned cells to (i) get list of vertices + // on the process boundary and (ii) apply re-ordering to cells for + // locality + std::vector boundary_v; + { + std::vector> cells1_v_local_cells; + + for (std::int32_t i = 0; i < num_cell_types; ++i) + { + std::int32_t num_cell_vertices = mesh::num_cell_vertices(celltypes[i]); + std::int32_t num_owned_cells + = cells1_v[i].size() / num_cell_vertices - ghost_owners[i].size(); + cells1_v_local_cells.push_back( + std::span(cells1_v[i].data(), num_owned_cells * num_cell_vertices)); + } + spdlog::info("Build local dual graph"); + auto [graph, unmatched_facets, max_v, facet_attached_cells] + = build_local_dual_graph(celltypes, cells1_v_local_cells); + + // TODO: in the original create_mesh(), cell reordering is done here. + // This needs reworking for mixed cell topology + + // Boundary vertices are marked as 'unknown' + boundary_v = unmatched_facets; + std::sort(boundary_v.begin(), boundary_v.end()); + boundary_v.erase(std::unique(boundary_v.begin(), boundary_v.end()), + boundary_v.end()); + + // Remove -1 if it occurs in boundary vertices (may occur in mixed + // topology) + if (!boundary_v.empty() > 0 and boundary_v[0] == -1) + boundary_v.erase(boundary_v.begin()); + } + + spdlog::debug("Got {} boundary vertices", boundary_v.size()); + + // Create Topology + + std::vector> cells1_v_span; + std::transform(cells1_v.cbegin(), cells1_v.cend(), + std::back_inserter(cells1_v_span), + [](auto& c) { return std::span(c); }); + std::vector> original_idx1_span; + std::transform(original_idx1.cbegin(), original_idx1.cend(), + std::back_inserter(original_idx1_span), + [](auto& c) { return std::span(c); }); + std::vector> ghost_owners_span; + std::transform(ghost_owners.cbegin(), ghost_owners.cend(), + std::back_inserter(ghost_owners_span), + [](auto& c) { return std::span(c); }); + + Topology topology + = create_topology(comm, celltypes, cells1_v_span, original_idx1_span, + ghost_owners_span, boundary_v); + + // Create connectivities required higher-order geometries for creating + // a Geometry object + for (int i = 0; i < num_cell_types; ++i) + { + for (int e = 1; e < topology.dim(); ++e) + if (doflayouts[i].num_entity_dofs(e) > 0) + topology.create_entities(e); + if (elements[i].needs_dof_permutations()) + topology.create_entity_permutations(); + } + + // Build list of unique (global) node indices from cells1 and + // distribute coordinate data + std::vector nodes1, nodes2; + for (std::vector& c : cells1) + nodes1.insert(nodes1.end(), c.begin(), c.end()); + for (std::vector& c : cells1) + nodes2.insert(nodes2.end(), c.begin(), c.end()); + + dolfinx::radix_sort(std::span(nodes1)); + nodes1.erase(std::unique(nodes1.begin(), nodes1.end()), nodes1.end()); + std::vector coords + = dolfinx::MPI::distribute_data(comm, nodes1, commg, x, xshape[1]); + + // Create geometry object + Geometry geometry + = create_geometry(topology, elements, nodes1, nodes2, coords, xshape[1]); + + return Mesh(comm, std::make_shared(std::move(topology)), + std::move(geometry)); +} + /// @brief Create a distributed mesh from mesh data using the default /// graph partitioner to determine the parallel distribution of the /// mesh. @@ -922,10 +1153,10 @@ create_mesh(MPI_Comm comm, std::span cells, } } -/// @brief Create a sub-geometry from a mesh and a subset of mesh entities to be -/// included. A sub-geometry is simply a `Geometry` object containing only the -/// geometric information for the subset of entities. The entities may differ in -/// topological dimension from the original mesh. +/// @brief Create a sub-geometry from a mesh and a subset of mesh entities to +/// be included. A sub-geometry is simply a `Geometry` object containing only +/// the geometric information for the subset of entities. The entities may +/// differ in topological dimension from the original mesh. /// /// @param[in] mesh The full mesh. /// @param[in] dim Topological dimension of the sub-topology. diff --git a/python/dolfinx/wrappers/mesh.cpp b/python/dolfinx/wrappers/mesh.cpp index 5e8b72bf77c..a68cdc5f7a6 100644 --- a/python/dolfinx/wrappers/mesh.cpp +++ b/python/dolfinx/wrappers/mesh.cpp @@ -279,6 +279,50 @@ void declare_mesh(nb::module_& m, std::string type) }, nb::arg("comm"), nb::arg("p"), nb::arg("n"), nb::arg("celltype"), nb::arg("partitioner").none()); + + m.def("create_mesh", + [](MPICommWrapper comm, + const std::vector, + nb::c_contig>>& cells_nb, + const std::vector>& elements, + nb::ndarray x, + const PythonCellPartitionFunction& p) + { + std::size_t shape1 = x.ndim() == 1 ? 1 : x.shape(1); + + std::vector> cells; + std::transform( + cells_nb.begin(), cells_nb.end(), std::back_inserter(cells), + [](auto c) + { return std::span(c.data(), c.size()); }); + + if (p) + { + auto p_wrap + = [p](MPI_Comm comm, int n, + const std::vector& cell_types, + const std::vector>& cells) + { + std::vector> cells_nb; + std::transform( + cells.begin(), cells.end(), std::back_inserter(cells_nb), + [](auto c) + { + return nb::ndarray( + c.data(), {c.size()}, nb::handle()); + }); + return p(MPICommWrapper(comm), n, cell_types, cells_nb); + }; + return dolfinx::mesh::create_mesh( + comm.get(), comm.get(), cells, elements, comm.get(), + std::span(x.data(), x.size()), {x.shape(0), shape1}, p_wrap); + } + else + return dolfinx::mesh::create_mesh( + comm.get(), comm.get(), cells, elements, comm.get(), + std::span(x.data(), x.size()), {x.shape(0), shape1}, nullptr); + }); + m.def( "create_mesh", [](MPICommWrapper comm, diff --git a/python/test/unit/mesh/test_create_mixed_mesh.py b/python/test/unit/mesh/test_create_mixed_mesh.py new file mode 100644 index 00000000000..83be54ae3bc --- /dev/null +++ b/python/test/unit/mesh/test_create_mixed_mesh.py @@ -0,0 +1,107 @@ +from mpi4py import MPI + +import numpy as np + +import dolfinx.cpp as _cpp +from dolfinx.cpp.mesh import ( + GhostMode, + create_mesh, +) +from dolfinx.fem import coordinate_element +from dolfinx.mesh import CellType + + +def test_create_mixed_mesh(): + nx = 7 + ny = 11 + nz = 8 + n_cells = nx * ny * nz + cells: list = [[], [], []] + orig_idx: list = [[], [], []] + idx = 0 + for i in range(n_cells): + iz = i // (nx * ny) + j = i % (nx * ny) + iy = j // nx + ix = j % nx + + v0 = (iz * (ny + 1) + iy) * (nx + 1) + ix + v1 = v0 + 1 + v2 = v0 + (nx + 1) + v3 = v1 + (nx + 1) + v4 = v0 + (nx + 1) * (ny + 1) + v5 = v1 + (nx + 1) * (ny + 1) + v6 = v2 + (nx + 1) * (ny + 1) + v7 = v3 + (nx + 1) * (ny + 1) + if iz < nz // 2: + cells[0] += [v0, v1, v2, v3, v4, v5, v6, v7] + orig_idx[0] += [idx] + idx += 1 + elif iz == nz // 2: + # pyramid + cells[1] += [v0, v1, v2, v3, v6] + orig_idx[1] += [idx] + idx += 1 + # tet + cells[2] += [v0, v1, v4, v6] + cells[2] += [v4, v6, v5, v1] + cells[2] += [v5, v6, v7, v1] + cells[2] += [v6, v7, v3, v1] + orig_idx[2] += [idx, idx + 1, idx + 2, idx + 3] + idx += 4 + else: + # tet + cells[2] += [v0, v1, v2, v6] + cells[2] += [v1, v2, v3, v6] + cells[2] += [v0, v1, v4, v6] + cells[2] += [v4, v6, v5, v1] + cells[2] += [v5, v6, v7, v1] + cells[2] += [v6, v7, v3, v1] + orig_idx[2] += [idx, idx + 1, idx + 2, idx + 3, idx + 4, idx + 5] + idx += 6 + + n_points = (nx + 1) * (ny + 1) * (nz + 1) + sqxy = (nx + 1) * (ny + 1) + geom = [] + for v in range(n_points): + iz = v // sqxy + p = v % sqxy + iy = p // (nx + 1) + ix = p % (nx + 1) + geom += [[ix / nx, iy / ny, iz / nz]] + + if MPI.COMM_WORLD.rank == 0: + cells_np = [np.array(c) for c in cells] + geomx = np.array(geom, dtype=np.float64) + else: + cells_np = [np.zeros(0) for c in cells] + geomx = np.zeros((0, 3), dtype=np.float64) + + hexahedron = coordinate_element(CellType.hexahedron, 1) + pyramid = coordinate_element(CellType.pyramid, 1, variant=1) + tetrahedron = coordinate_element(CellType.tetrahedron, 1) + + part = _cpp.mesh.create_cell_partitioner(GhostMode.none) + mesh = create_mesh( + MPI.COMM_WORLD, + cells_np, + [hexahedron._cpp_object, pyramid._cpp_object, tetrahedron._cpp_object], + geomx, + part, + ) + + entity_types = mesh.topology.entity_types[3] + assert entity_types[0] == CellType.hexahedron + assert entity_types[1] == CellType.pyramid + assert entity_types[2] == CellType.tetrahedron + + for i in range(3): + assert ( + mesh.topology.connectivity((3, i), (0, 0)).num_nodes + == mesh.topology.index_maps(3)[i].size_local + ) + + num_cells = sum(mesh.topology.index_maps(3)[i].size_local for i in range(3)) + all_cells = mesh.comm.allreduce(num_cells) + + assert all_cells == 2079 diff --git a/python/test/unit/mesh/test_mesh_partitioners.py b/python/test/unit/mesh/test_mesh_partitioners.py index 18947d84bd9..39925b9b28e 100644 --- a/python/test/unit/mesh/test_mesh_partitioners.py +++ b/python/test/unit/mesh/test_mesh_partitioners.py @@ -195,11 +195,11 @@ def test_mixed_topology_partitioning(): v5 = v1 + (nx + 1) * (ny + 1) v6 = v2 + (nx + 1) * (ny + 1) v7 = v3 + (nx + 1) * (ny + 1) - if ix < nx // 2: + if iz < nz // 2: cells[0] += [v0, v1, v2, v3, v4, v5, v6, v7] orig_idx[0] += [idx] idx += 1 - elif ix == nx // 2: + elif iz == nz // 2: # pyramid cells[1] += [v0, v1, v2, v3, v6] orig_idx[1] += [idx]