Improve python exposure determine_point_ownership

* Python wrapper around nanobind exposed function
* Extra optional arguments `cells`

Author: ordinary-slim

cpp/dolfinx/fem/interpolate.h

@@ -1070,7 +1070,12 @@ geometry::PointOwnershipData<T> create_interpolation_data(
       x[3 * i + j] = coords[i + j * num_points];
 
   // Determine ownership of each point
-  return geometry::determine_point_ownership<T>(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<std::int32_t> cells1(num_cells1, 0);
+  std::iota(cells1.begin(), cells1.end(), 0);
+  return geometry::determine_point_ownership<T>(mesh1, x, cells1, padding);
 }
 
 /// @brief Interpolate a finite element Function defined on a mesh to a

cpp/dolfinx/geometry/utils.h

@@ -663,6 +663,7 @@ graph::AdjacencyList<std::int32_t> 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
@@ -685,18 +686,14 @@ graph::AdjacencyList<std::int32_t> compute_colliding_cells(
 template <std::floating_point T>
 PointOwnershipData<T> determine_point_ownership(const mesh::Mesh<T>& mesh,
                                                 std::span<const T> points,
+                                                std::span<const std::int32_t> cells,
                                                 T padding)
 {
   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<std::int32_t> 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);
 

python/dolfinx/geometry.py

@@ -270,3 +270,32 @@ 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[np.ndarray] = None,
+    padding: float = 0.0
+) -> PointOwnershipData:
+    """Build PointOwnershipData for a mesh-points pair.
+    First, potential collisions are found
+    by building BoundingBoxTrees for the points
+    and the mesh cells and finding which leaves collide.
+    Then, actual cell-point containment pairs are determined
+    using the GJK algorithm.
+
+    Args:
+        mesh: The mesh
+        points: The points
+        cells: Cells in mesh who potentially contain the points.
+            If ``None`` then all cells are considered.
+        padding: Padding added to the bounding boxes.
+
+    Returns:
+        PointOwnershipData
+
+    """
+    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))

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<T>& mesh,
-           nb::ndarray<const T, nb::c_contig> points, const T padding)
+           nb::ndarray<const T, nb::c_contig> points,
+           nb::ndarray<const std::int32_t, nb::ndim<1>, nb::c_contig> cells,
+           const T padding)
         {
           const std::size_t p_s0 = points.ndim() == 1 ? 1 : points.shape(0);
           std::span<const T> _p(points.data(), 3 * p_s0);
           return dolfinx::geometry::determine_point_ownership<T>(mesh, _p,
+                                                                 std::span(cells.data(), cells.size()),
                                                                  padding);
-        });
+        },
+      nb::arg("mesh"), nb::arg("points"), nb::arg("cells"), nb::arg("padding"),
+      "Compute PointOwnershipData for mesh-points pair.");
 
   std::string pod_pyclass_name = "PointOwnershipData_" + type;
   nb::class_<dolfinx::geometry::PointOwnershipData<T>>(m,