From 69123219aae45316ecf423d20c187be2a84681aa Mon Sep 17 00:00:00 2001 From: ordinary-slim Date: Mon, 26 Aug 2024 10:55:43 +0000 Subject: [PATCH] Improve python exposure determine_point_ownership * Python wrapper around nanobind exposed function * Extra optional arguments `cells` --- cpp/dolfinx/fem/interpolate.h | 7 +- cpp/dolfinx/geometry/utils.h | 15 +- python/dolfinx/geometry.py | 45 +++++ python/dolfinx/wrappers/geometry.cpp | 9 +- .../unit/geometry/test_bounding_box_tree.py | 158 ++++++++++++++++++ 5 files changed, 222 insertions(+), 12 deletions(-) diff --git a/cpp/dolfinx/fem/interpolate.h b/cpp/dolfinx/fem/interpolate.h index 12956057cb7..55a4cee93bd 100644 --- a/cpp/dolfinx/fem/interpolate.h +++ b/cpp/dolfinx/fem/interpolate.h @@ -1070,7 +1070,12 @@ geometry::PointOwnershipData create_interpolation_data( x[3 * i + j] = coords[i + j * num_points]; // Determine ownership of each point - return geometry::determine_point_ownership(mesh1, x, padding); + const int tdim = mesh1.topology()->dim(); + auto cell_map1 = mesh1.topology()->index_map(tdim); + const std::int32_t num_cells1 = cell_map1->size_local(); + std::vector cells1(num_cells1, 0); + std::iota(cells1.begin(), cells1.end(), 0); + return geometry::determine_point_ownership(mesh1, x, cells1, padding); } /// @brief Interpolate a finite element Function defined on a mesh to a diff --git a/cpp/dolfinx/geometry/utils.h b/cpp/dolfinx/geometry/utils.h index 0f643e62643..f38a8fbbd18 100644 --- a/cpp/dolfinx/geometry/utils.h +++ b/cpp/dolfinx/geometry/utils.h @@ -663,21 +663,22 @@ graph::AdjacencyList compute_colliding_cells( /// @param[in] mesh The mesh /// @param[in] points Points to check for collision (`shape=(num_points, /// 3)`). Storage is row-major. +/// @param[in] cells Cells to check for ownership /// @param[in] padding Amount of absolute padding of bounding boxes of the mesh. /// Each bounding box of the mesh is padded with this amount, to increase /// the number of candidates, avoiding rounding errors in determining the owner /// of a point if the point is on the surface of a cell in the mesh. -/// @return Tuple `(src_owner, dest_owner, dest_points, dest_cells)`, +/// @return Point ownership data with fields src_owner, dest_owner, dest_points, dest_cells, /// where src_owner is a list of ranks corresponding to the input /// points. dest_owner is a list of ranks corresponding to dest_points, /// the points that this process owns. dest_cells contains the /// corresponding cell for each entry in dest_points. /// /// @note `dest_owner` is sorted -/// @note Returns -1 if no colliding process is found +/// @note `src_owner` is -1 if no colliding process is found /// @note dest_points is flattened row-major, shape `(dest_owner.size(), /// 3)` -/// @note Only looks through cells owned by the process +/// @note Only returns cells owned by the process /// @note A large padding value can increase the runtime of the function by /// orders of magnitude, because for non-colliding cells /// one has to determine the closest cell among all processes with an @@ -685,18 +686,14 @@ graph::AdjacencyList compute_colliding_cells( template PointOwnershipData determine_point_ownership(const mesh::Mesh& mesh, std::span points, - T padding) + std::span cells, + T padding = 0.0) { MPI_Comm comm = mesh.comm(); // Create a global bounding-box tree to find candidate processes with // cells that could collide with the points const int tdim = mesh.topology()->dim(); - auto cell_map = mesh.topology()->index_map(tdim); - const std::int32_t num_cells = cell_map->size_local(); - // NOTE: Should we send the cells in as input? - std::vector cells(num_cells, 0); - std::iota(cells.begin(), cells.end(), 0); BoundingBoxTree bb(mesh, tdim, cells, padding); BoundingBoxTree global_bbtree = bb.create_global_tree(comm); diff --git a/python/dolfinx/geometry.py b/python/dolfinx/geometry.py index 03cd0870b24..1d5db202ca9 100644 --- a/python/dolfinx/geometry.py +++ b/python/dolfinx/geometry.py @@ -270,3 +270,48 @@ def compute_distance_gjk( """ return _cpp.geometry.compute_distance_gjk(p, q) + + +def determine_point_ownership( + mesh: Mesh, + points: npt.NDArray[np.floating], + cells: typing.Optional[npt.NDArray[np.int32]] = None, + padding: float = 0.0, +) -> PointOwnershipData: + """Build point ownership data for a mesh-points pair. + + First, potential collisions are found by computing intersections + between the bounding boxes of the cells and the set of points. + Then, actual containment pairs are determined using the GJK algorithm. + + Args: + mesh: The mesh + points: Points to check for collision (``shape=(num_points, gdim)``) + cells: Cells to check for ownership + If ``None`` then all cells are considered. + padding: Amount of absolute padding of bounding boxes of the mesh. + Each bounding box of the mesh is padded with this amount, to increase + the number of candidates, avoiding rounding errors in determining the owner + of a point if the point is on the surface of a cell in the mesh. + + Returns: + Point ownership data + + Note: + `dest_owner` is sorted + + `src_owner` is -1 if no colliding process is found + + Only returns cells owned by the process + + A large padding value will increase the run-time of the code by orders + of magnitude. General advice is to use a padding on the scale of the + cell size. + + """ + if cells is None: + map = mesh.topology.index_map(mesh.topology.dim) + cells = np.arange(map.size_local, dtype=np.int32) + return PointOwnershipData( + _cpp.geometry.determine_point_ownership(mesh._cpp_object, points, cells, padding) + ) diff --git a/python/dolfinx/wrappers/geometry.cpp b/python/dolfinx/wrappers/geometry.cpp index d64e70cb584..ce0931134af 100644 --- a/python/dolfinx/wrappers/geometry.cpp +++ b/python/dolfinx/wrappers/geometry.cpp @@ -180,13 +180,18 @@ void declare_bbtree(nb::module_& m, std::string type) nb::arg("mesh"), nb::arg("dim"), nb::arg("indices"), nb::arg("points")); m.def("determine_point_ownership", [](const dolfinx::mesh::Mesh& mesh, - nb::ndarray points, const T padding) + nb::ndarray points, + nb::ndarray, nb::c_contig> cells, + const T padding) { const std::size_t p_s0 = points.ndim() == 1 ? 1 : points.shape(0); std::span _p(points.data(), 3 * p_s0); return dolfinx::geometry::determine_point_ownership(mesh, _p, + std::span(cells.data(), cells.size()), padding); - }); + }, + nb::arg("mesh"), nb::arg("points"), nb::arg("cells"), nb::arg("padding") = 0.0, + "Compute point ownership data for mesh-points pair."); std::string pod_pyclass_name = "PointOwnershipData_" + type; nb::class_>(m, diff --git a/python/test/unit/geometry/test_bounding_box_tree.py b/python/test/unit/geometry/test_bounding_box_tree.py index 2cc328300b6..f76d1ccd05f 100644 --- a/python/test/unit/geometry/test_bounding_box_tree.py +++ b/python/test/unit/geometry/test_bounding_box_tree.py @@ -12,6 +12,7 @@ from dolfinx import cpp as _cpp from dolfinx.geometry import ( + PointOwnershipData, bb_tree, compute_closest_entity, compute_colliding_cells, @@ -19,6 +20,7 @@ compute_collisions_trees, compute_distance_gjk, create_midpoint_tree, + determine_point_ownership, ) from dolfinx.mesh import ( CellType, @@ -496,3 +498,159 @@ def test_surface_bbtree_collision(dtype): collisions = compute_collisions_trees(bbtree1, bbtree2) assert len(collisions) == 1 + + +@pytest.mark.parametrize("dim", [2, 3]) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_determine_point_ownership(dim, dtype): + """Find point owners (ranks and cells) using bounding box trees + global communication + and compare to point ownership data results.""" + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + tdim = dim + num_cells_side = 4 + h = dtype(1.0 / num_cells_side) + if tdim == 2: + mesh = create_unit_square( + MPI.COMM_WORLD, num_cells_side, num_cells_side, CellType.quadrilateral, dtype=dtype + ) + else: + mesh = create_unit_cube( + MPI.COMM_WORLD, + num_cells_side, + num_cells_side, + num_cells_side, + CellType.hexahedron, + dtype=dtype, + ) + cell_map = mesh.topology.index_map(tdim) + + tree = bb_tree(mesh, mesh.topology.dim, np.arange(cell_map.size_local)) + num_global_cells = num_cells_side**dim + all_midpoints = np.zeros((num_global_cells, 3), dtype=dtype) + counter = 0 + Xs = [(i + 0.5) * h for i in range(num_cells_side)] + Ys = Xs + Zs = [0.0] if dim == 2 else Xs + for x in Xs: + for y in Ys: + for z in Zs: + all_midpoints[counter, :] = np.array([x, y, z], dtype=dtype) + counter += 1 + # Since ghost cells are left out and the points considered are midpoints + # of cells, they are only contained in a single process / cell + # Moreover, the bounding boxes of the elements correspond to the elements, + # hence the tree collisions are the exact collisions + tree_col = compute_collisions_points(tree, all_midpoints) + processwise_owners = np.zeros(2 * num_global_cells, dtype=np.int32) + owners = np.empty_like(processwise_owners) + for ipoint in range(num_global_cells): + cell = tree_col.links(ipoint) + if len(cell) > 0: + processwise_owners[2 * ipoint] = rank + processwise_owners[2 * ipoint + 1] = cell[0] + + # The value at a given index is null if it doesn't correspond + # to the current process. + # We can sum the processwise arrays to obtain a global array + comm.Allreduce(processwise_owners, owners, op=MPI.SUM) + owner_ranks = owners[[2 * i for i in range(num_global_cells)]] + owner_cells = owners[[2 * i + 1 for i in range(num_global_cells)]] + + # Reorganize ownership data (point, index, rank, cell) into dictionary + ownership_data = {} + for ipoint in range(num_global_cells): + ownership_data[tuple(all_midpoints[ipoint])] = ( + ipoint, + owner_ranks[ipoint], + owner_cells[ipoint], + ) + + def check_po(po: PointOwnershipData, src_points, ownership_data, global_dest_owners): + """ + Check point ownership data + + po: PointOwnershipData object to check + src_points: Points sent by process + ownership_data: {point:(global_index,rank,cell} + global_dest_owners: Rank who sent each point + """ + # Check src_owner: Check owner ranks of sent points + src_owner = po.src_owner() + for ipoint in range(src_points.shape[0]): + assert ownership_data[tuple(src_points[ipoint])][1] == src_owner[ipoint] + + dest_points = po.dest_points() + dest_owners = po.dest_owner() + dest_cells = po.dest_cells() + + # Check dest_points: All points that should have been found have been found + dest_points_indices = list(range(dest_points.shape[0])) + for point, data in ownership_data.items(): + (iglobal, processor, _) = data + if processor == rank: + found = False + point = np.array(point, dtype) + for jpoint in dest_points_indices: + found = np.allclose(point, dest_points[jpoint]) + if found: + break + assert found + dest_points_indices.remove(jpoint) + + # Check dest_owners and dest_cells + # dest_owners: Ranks that asked about the points we own + # dest_cells: Local index of cell that contains the points we own + for ipoint in range(dest_points.shape[0]): + iglobal = ownership_data[tuple(dest_points[ipoint])][0] + c = ownership_data[tuple(dest_points[ipoint])][2] + assert dest_owners[ipoint] == global_dest_owners[iglobal] + assert dest_cells[ipoint] == c + + def set_local_range(array): + N = array.shape[0] + n = N // comm.size + r = N % comm.size + # First r processes has one extra value + if rank < r: + (start, stop) = [rank * (n + 1), (rank + 1) * (n + 1)] + else: + (start, stop) = [rank * n + r, (rank + 1) * n + r] + return array[start:stop], start, stop + + def compute_global_owners(N, start, stop): + """Compute array of ranks who own each point""" + mask_points_owned = np.zeros(N, np.int32) + global_owners = np.empty_like(mask_points_owned) + mask_points_owned[start:stop] = rank + comm.Allreduce(mask_points_owned, global_owners, op=MPI.SUM) + return global_owners + + # All cells + points, start, stop = set_local_range(all_midpoints) + owners = compute_global_owners(np.int64(all_midpoints.shape[0]), start, stop) + all_cells = np.arange(cell_map.size_local, dtype=np.int32) + po = determine_point_ownership(mesh, points, all_cells) + + check_po(po, points, ownership_data, owners) + + # Left half + num_left_cells = np.rint((num_cells_side**dim) / 2).astype(np.int32) + left_midpoints = all_midpoints[np.arange(num_left_cells), :] + points, start, stop = set_local_range(left_midpoints) + owners = compute_global_owners(np.int64(all_midpoints.shape[0]), start, stop) + left_cells = locate_entities(mesh, tdim, lambda x: x[0] <= 0.5) + left_cells = np.array( + [cell for cell in left_cells if cell < cell_map.size_local], dtype=np.int32 + ) # Filter out ghost cells + lpo = determine_point_ownership(mesh, points, left_cells) + + left_ownership_data = {} + for ipoint in range(num_left_cells): + left_ownership_data[tuple(all_midpoints[ipoint])] = ( + ipoint, + owner_ranks[ipoint], + owner_cells[ipoint], + ) + check_po(lpo, points, left_ownership_data, owners)