Improve python exposure of determine_point_ownership

cpp/dolfinx/geometry/utils.h

@@ -663,40 +663,40 @@ 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
 /// 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)`,
-/// 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.
+/// @return Point ownership data.
 ///
 /// @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 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
 /// intersecting bounding box, which is an expensive operation to perform.
 template <std::floating_point T>
 PointOwnershipData<T> determine_point_ownership(const mesh::Mesh<T>& mesh,
                                                 std::span<const T> points,
-                                                T padding)
+                                                T padding,
+                                                std::span<const std::int32_t> cells = {})
 {
   MPI_Comm comm = mesh.comm();
 
+  const int tdim = mesh.topology()->dim();
+
+  std::vector<std::int32_t> local_cells;
+  if (cells.empty()) {
+    auto cell_map = mesh.topology()->index_map(tdim);
+    local_cells.resize(cell_map->size_local());
+    std::iota(local_cells.begin(), local_cells.end(), 0);
+    cells = std::span<const std::int32_t>(local_cells.data(), local_cells.size());
+  }
   // 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/fem/bcs.py

@@ -17,6 +17,7 @@
     from dolfinx.fem.function import Constant, Function
 
 import numpy as np
+import numpy.typing as npt
 
 import dolfinx
 from dolfinx import cpp as _cpp
@@ -57,7 +58,7 @@ def locate_dofs_geometrical(
 def locate_dofs_topological(
     V: typing.Union[dolfinx.fem.FunctionSpace, typing.Iterable[dolfinx.fem.FunctionSpace]],
     entity_dim: int,
-    entities: numpy.typing.NDArray[np.int32],
+    entities: npt.NDArray[np.int32],
     remote: bool = True,
 ) -> np.ndarray:
     """Locate degrees-of-freedom belonging to mesh entities topologically.
@@ -150,7 +151,7 @@ def dof_indices(self) -> tuple[np.ndarray, int]:
 
 def dirichletbc(
     value: typing.Union[Function, Constant, np.ndarray],
-    dofs: numpy.typing.NDArray[np.int32],
+    dofs: npt.NDArray[np.int32],
     V: typing.Optional[dolfinx.fem.FunctionSpace] = None,
 ) -> DirichletBC:
     """Create a representation of Dirichlet boundary condition which

python/dolfinx/fem/forms.py

@@ -111,7 +111,7 @@ def integral_types(self):
 
 def get_integration_domains(
     integral_type: IntegralType,
-    subdomain: typing.Optional[typing.Union[MeshTags, list[tuple[int, np.ndarray]]]],
+    subdomain: typing.Optional[typing.Union[MeshTags, list[tuple[int, npt.NDArray[np.int32]]]]],
     subdomain_ids: list[int],
 ) -> list[tuple[int, np.ndarray]]:
     """Get integration domains from subdomain data.

python/dolfinx/fem/function.py

@@ -90,7 +90,7 @@ class Expression:
     def __init__(
         self,
         e: ufl.core.expr.Expr,
-        X: np.ndarray,
+        X: typing.Union[npt.NDArray[np.float32], npt.NDArray[np.float64]],
         comm: typing.Optional[_MPI.Comm] = None,
         form_compiler_options: typing.Optional[dict] = None,
         jit_options: typing.Optional[dict] = None,
@@ -196,7 +196,7 @@ def _create_expression(dtype):
     def eval(
         self,
         mesh: Mesh,
-        entities: np.ndarray,
+        entities: npt.NDArray[np.int32],
         values: typing.Optional[np.ndarray] = None,
     ) -> np.ndarray:
         """Evaluate Expression on entities.
@@ -412,8 +412,8 @@ def interpolate_nonmatching(
     def interpolate(
         self,
         u0: typing.Union[typing.Callable, Expression, Function],
-        cells0: typing.Optional[np.ndarray] = None,
-        cells1: typing.Optional[np.ndarray] = None,
+        cells0: typing.Optional[npt.NDArray[np.int32]] = None,
+        cells1: typing.Optional[npt.NDArray[np.int32]] = None,
     ) -> None:
         """Interpolate an expression.
 
@@ -584,7 +584,7 @@ def _create_dolfinx_element(
     comm: _MPI.Intracomm,
     cell_type: _cpp.mesh.CellType,
     ufl_e: ufl.FiniteElementBase,
-    dtype: np.dtype,
+    dtype: npt.DTypeLike,
 ) -> typing.Union[_cpp.fem.FiniteElement_float32, _cpp.fem.FiniteElement_float64]:
     """Create a DOLFINx element from a basix.ufl element."""
     if np.issubdtype(dtype, np.float32):

python/dolfinx/geometry.py

@@ -270,3 +270,45 @@ def compute_distance_gjk(
 
     """
     return _cpp.geometry.compute_distance_gjk(p, q)
+
+
+def determine_point_ownership(
+    mesh: Mesh,
+    points: npt.NDArray[np.floating],
+    padding: float,
+    cells: typing.Optional[npt.NDArray[np.int32]] = None,
+) -> 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)``)
+        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.
+        cells: Cells to check for ownership
+            If ``None`` then all cells are considered.
+
+    Returns:
+        Point ownership data
+
+    Note:
+        ``dest_owner`` is sorted
+
+        ``src_owner`` is -1 if no colliding process is found
+
+        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)
+    )

python/dolfinx/graph.py

@@ -7,7 +7,10 @@
 
 from __future__ import annotations
 
+import typing
+
 import numpy as np
+import numpy.typing as npt
 
 from dolfinx import cpp as _cpp
 from dolfinx.cpp.graph import partitioner
@@ -31,7 +34,10 @@
 __all__ = ["adjacencylist", "partitioner"]
 
 
-def adjacencylist(data: np.ndarray, offsets=None):
+def adjacencylist(
+    data: typing.Union[npt.NDArray[np.int32], npt.NDArray[np.int64]],
+    offsets: typing.Optional[npt.NDArray[np.int32]] = None,
+):
     """Create an AdjacencyList for int32 or int64 datasets.
 
     Args:

python/dolfinx/mesh.py

@@ -453,7 +453,7 @@ def compute_midpoints(mesh: Mesh, dim: int, entities: npt.NDArray[np.int32]):
     return _cpp.mesh.compute_midpoints(mesh._cpp_object, dim, entities)
 
 
-def locate_entities(mesh: Mesh, dim: int, marker: typing.Callable) -> np.ndarray:
+def locate_entities(mesh: Mesh, dim: int, marker: typing.Callable) -> npt.NDArray[np.int32]:
     """Compute mesh entities satisfying a geometric marking function.
 
     Args:
@@ -470,7 +470,9 @@ def locate_entities(mesh: Mesh, dim: int, marker: typing.Callable) -> np.ndarray
     return _cpp.mesh.locate_entities(mesh._cpp_object, dim, marker)
 
 
-def locate_entities_boundary(mesh: Mesh, dim: int, marker: typing.Callable) -> np.ndarray:
+def locate_entities_boundary(
+    mesh: Mesh, dim: int, marker: typing.Callable
+) -> npt.NDArray[np.int32]:
     """Compute mesh entities that are connected to an owned boundary
     facet and satisfy a geometric marking function.
 
@@ -534,7 +536,7 @@ def transfer_meshtag(
 
 def refine(
     mesh: Mesh,
-    edges: typing.Optional[np.ndarray] = None,
+    edges: typing.Optional[npt.NDArray[np.int32]] = None,
     redistribute: bool = True,
     ghost_mode: GhostMode = GhostMode.shared_facet,
     option: RefinementOption = RefinementOption.none,

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,
-                                                                 padding);
-        });
+                                                                 padding,
+                                                                 std::span(cells.data(), cells.size()));
+        },
+      nb::arg("mesh"), nb::arg("points"), nb::arg("padding"), nb::arg("cells"),
+      "Compute point ownership data for mesh-points pair.");
 
   std::string pod_pyclass_name = "PointOwnershipData_" + type;
   nb::class_<dolfinx::geometry::PointOwnershipData<T>>(m,