Introduce threading for collision detection (#4123)

* compute distances for multiple bodies towards single convex hull

* Fix tolerance in GJK algorithm and remove looping in tests

* Try to easen tolerance a bit

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Fix tolerance and simplify chunking

* Stricter tolerance again

* Add new termination criterion

* Use <=

Author: jorgensd

cpp/dolfinx/geometry/gjk.h

@@ -318,7 +318,7 @@ std::array<T, 3> compute_distance_gjk(std::span<const T> p0,
   constexpr int maxk = 15; // Maximum number of iterations of the GJK algorithm
 
   // Tolerance
-  const U eps = 1.0e4 * std::numeric_limits<T>::epsilon();
+  const U eps = 1000 * std::numeric_limits<U>::epsilon();
 
   // Initialise vector and simplex
   std::array<U, 3> v = {p[0] - q[0], p[1] - q[1], p[2] - q[2]};
@@ -382,16 +382,75 @@ std::array<T, 3> compute_distance_gjk(std::span<const T> p0,
     SPDLOG_DEBUG("new s size={}", 3 * j);
     s.resize(3 * j);
 
+    // 2nd exit condition - strict monotonicity
+    // Floating point can cause the algorithm to stagnate. Then we terminate.
     const U vn = impl_gjk::dot3(v, v);
-    // 2nd exit condition - intersecting or touching
+    if (vnorm2 <= vn)
+      break;
+
+    // 3nd exit condition - intersecting or touching
     if (vn < eps * eps)
       break;
   }
 
   if (k == maxk)
     throw std::runtime_error("GJK error - max iteration limit reached");
-
   return {static_cast<T>(v[0]), static_cast<T>(v[1]), static_cast<T>(v[2])};
 }
 
+/// @brief Compute the distance between a sequence of convex bodies `p0, ...,
+/// pN` and `q`, each defined by a set of points.
+///
+/// Uses the Gilbert–Johnson–Keerthi (GJK) distance algorithm.
+///
+/// @param[in] bodies List of the list of points that make up each of  N bodies
+/// considered as body 1. `shape=(num_bodies, (num_points_body_j, 3)`. Row-major
+/// storage.
+/// @param[in] q Body 2 list of points, `shape=(num_points, 3)`. Row-major
+/// storage.
+/// @tparam T Floating point type
+/// @tparam U Floating point type used for geometry computations internally,
+/// which should be higher precision than T, to maintain accuracy.
+/// @return For each body in `p_j`, return the shortest distance vector to
+/// body 2. Shape (num_points, 3).
+template <std::floating_point T,
+          typename U = boost::multiprecision::cpp_bin_float_double_extended>
+std::vector<T>
+compute_distances_gjk(const std::vector<std::span<const T>>& bodies,
+                      std::span<const T> q, size_t num_threads)
+{
+  size_t total_size = bodies.size();
+  num_threads = std::min(num_threads, total_size);
+
+  std::vector<T> results(total_size * 3);
+  auto compute_chunk
+      = [&results, &bodies](size_t c0, size_t c1, std::span<const T> q_ref)
+  {
+    for (size_t i = c0; i < c1; ++i)
+    {
+      // Using U explicitly as the internal precision type
+      std::array<T, 3> dist = compute_distance_gjk<T, U>(bodies[i], q_ref);
+      results[3 * i + 0] = dist[0];
+      results[3 * i + 1] = dist[1];
+      results[3 * i + 2] = dist[2];
+    }
+  };
+
+  if (num_threads <= 1)
+  {
+    compute_chunk(0, total_size, q);
+  }
+  else
+  {
+    std::vector<std::jthread> threads(num_threads);
+    for (size_t i = 0; i < num_threads; ++i)
+    {
+      auto [c0, c1] = dolfinx::MPI::local_range(i, total_size, num_threads);
+      threads[i] = std::jthread(compute_chunk, c0, c1, q);
+    }
+  }
+
+  return results;
+}
+
 } // namespace dolfinx::geometry

cpp/dolfinx/geometry/utils.h

@@ -872,7 +872,7 @@ determine_point_ownership(const mesh::Mesh<T>& mesh, std::span<const T> points,
       std::array<T, 3> point;
       std::copy_n(std::next(received_points.begin(), 3 * i), 3, point.begin());
 
-      // Find shortest distance among cells with colldiing bounding box
+      // Find shortest distance among cells with colliding bounding box
       T shortest_distance = std::numeric_limits<T>::max();
       std::int32_t closest_cell = -1;
       for (auto cell : candidate_collisions.links(i))

python/dolfinx/geometry.py

@@ -28,6 +28,7 @@
     "compute_collisions_points",
     "compute_collisions_trees",
     "compute_distance_gjk",
+    "compute_distances_gjk",
     "create_midpoint_tree",
     "determine_point_ownership",
     "squared_distance",
@@ -290,6 +291,34 @@ def compute_distance_gjk(
     raise RuntimeError("Invalid dtype in compute_distance_gjk")
 
 
+def compute_distances_gjk(
+    bodies: list[npt.NDArray[np.floating]], q: npt.NDArray[np.floating], num_threads: int
+) -> npt.NDArray[np.floating]:
+    """Compute the distance between a set of convex bodies.
+
+    For each convex body defined in `bodies`;
+    (a set of 3D points for each body) find the shortest distance vector
+    to to the body `q` defined by another set of 3D points.
+    The method uses the
+    Gilbert-Johnson-Keerthi (GJK) distance algorithm.
+
+    Args:
+        bodies: List of bodies, where each body is an array of
+            (``shape=(num_points_i, 3, gdim)``).
+        q: Body 2 list of points (``shape=(num_points_2, 3)``).
+        num_threads: Number of threads to use for GJK computation.
+
+    Returns:
+        Shortest vector between the two bodies.
+    """
+    assert all([p.dtype == q.dtype for p in bodies])
+    if np.issubdtype(q.dtype, np.float32):
+        return _cpp.geometry.compute_distances_gjk_float32(bodies, q, num_threads)
+    elif np.issubdtype(q.dtype, np.float64):
+        return _cpp.geometry.compute_distances_gjk_float64(bodies, q, num_threads)
+    raise RuntimeError("Invalid dtype in compute_distances_gjk")
+
+
 def determine_point_ownership(
     mesh: Mesh,
     points: npt.NDArray[np.floating],

python/dolfinx/wrappers/dolfinx_wrappers/geometry.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2017-2025 Chris N. Richardson and Garth N. Wells
+// Copyright (C) 2017-2026 Chris N. Richardson and Garth N. Wells
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -18,6 +18,7 @@
 #include <nanobind/nanobind.h>
 #include <nanobind/ndarray.h>
 #include <nanobind/stl/optional.h>
+#include <nanobind/stl/vector.h>
 #include <optional>
 #include <span>
 #include <string>
@@ -196,6 +197,40 @@ void declare_bbtree(nb::module_& m, const std::string& type)
       },
       nb::arg("p"), nb::arg("q"));
 
+  std::string gjks_name = "compute_distances_gjk_" + type;
+  m.def(
+      gjks_name.c_str(),
+      [](const std::vector<nb::ndarray<const T, nb::c_contig>>& bodies,
+         nb::ndarray<const T, nb::c_contig> q, size_t num_threads)
+      {
+        // If array is 1D assume single point
+        std::size_t q_s0 = q.ndim() == 1 ? 1 : q.shape(0);
+        std::span<const T> _q(q.data(), 3 * q_s0);
+
+        std::vector<std::span<const T>> _bodies;
+        _bodies.reserve(bodies.size());
+
+        std::ranges::transform(
+            bodies, std::back_inserter(_bodies),
+            [](auto& body)
+            {
+              // If sub-array in 1D assume single point
+              std::size_t body_s0 = body.ndim() == 1 ? 1 : body.shape(0);
+              return std::span<const T>(body.data(), 3 * body_s0);
+            });
+
+        using U = typename std::conditional<
+            std::is_same_v<T, float>, double,
+            boost::multiprecision::cpp_bin_float_double_extended>::type;
+
+        std::vector<T> distances
+            = dolfinx::geometry::compute_distances_gjk<T, U>(_bodies, _q,
+                                                             num_threads);
+        return dolfinx_wrappers::as_nbarray(std::move(distances),
+                                            {distances.size() / 3, 3});
+      },
+      nb::arg("bodies"), nb::arg("q"), nb::arg("num_threads"));
+
   m.def(
       "squared_distance",
       [](const dolfinx::mesh::Mesh<T>& mesh, int dim,

python/test/unit/geometry/test_bounding_box_tree.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2013-2025 Anders Logg, Jørgen S. Dokken, Chris Richardson
+# Copyright (C) 2013-2026 Anders Logg, Jørgen S. Dokken, Chris Richardson
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -17,6 +17,7 @@
     compute_collisions_points,
     compute_collisions_trees,
     compute_distance_gjk,
+    compute_distances_gjk,
     create_midpoint_tree,
     determine_point_ownership,
 )
@@ -50,47 +51,47 @@ def extract_geometricial_data(mesh, dim, entities):
     return mesh_nodes
 
 
-def expand_bbox(bbox, dtype):
-    """Expand min max bbox to convex hull."""
-    return np.array(
-        [
-            [bbox[0][0], bbox[0][1], bbox[0][2]],
-            [bbox[0][0], bbox[0][1], bbox[1][2]],
-            [bbox[0][0], bbox[1][1], bbox[0][2]],
-            [bbox[1][0], bbox[0][1], bbox[0][2]],
-            [bbox[1][0], bbox[0][1], bbox[1][2]],
-            [bbox[1][0], bbox[1][1], bbox[0][2]],
-            [bbox[0][0], bbox[1][1], bbox[1][2]],
-            [bbox[1][0], bbox[1][1], bbox[1][2]],
-        ],
-        dtype=dtype,
-    )
+def expand_bboxes(bboxes, dtype):
+    """Expand an array of min/max bboxes to convex hulls."""
+    if len(bboxes.shape) == 2:
+        bboxes = bboxes.reshape(1, *bboxes.shape)
+    idx_x = [0, 0, 0, 1, 1, 1, 0, 1]
+    idx_y = [0, 0, 1, 0, 0, 1, 1, 1]
+    idx_z = [0, 1, 0, 0, 1, 0, 1, 1]
+
+    x = bboxes[:, idx_x, 0]
+    y = bboxes[:, idx_y, 1]
+    z = bboxes[:, idx_z, 2]
+    return np.stack((x, y, z), axis=-1).astype(dtype)
 
 
-def find_colliding_cells(mesh, bbox, dtype):
+def find_colliding_cells(mesh, bbox, dtype, num_threads):
     """Given a mesh and a bounding box((xmin, ymin, zmin), (xmax, ymax,
     zmax)) find all colliding cells.
     """
     # Find actual cells using known bounding box tree
-    colliding_cells = []
     num_cells = mesh.topology.index_map(mesh.topology.dim).size_local
     x_indices = entities_to_geometry(
         mesh, mesh.topology.dim, np.arange(num_cells, dtype=np.int32), False
     )
     points = mesh.geometry.x
-    bounding_box = expand_bbox(bbox, dtype)
-    for cell in range(num_cells):
-        vertex_coords = points[x_indices[cell]]
-        bbox_cell = np.array([vertex_coords[0], vertex_coords[0]])
-        # Create bounding box for cell
-        for i in range(1, vertex_coords.shape[0]):
-            for j in range(3):
-                bbox_cell[0, j] = min(bbox_cell[0, j], vertex_coords[i, j])
-                bbox_cell[1, j] = max(bbox_cell[1, j], vertex_coords[i, j])
-        distance = compute_distance_gjk(expand_bbox(bbox_cell, dtype), bounding_box)
-        if np.dot(distance, distance) < 1e-16:
-            colliding_cells.append(cell)
-
+    bounding_box = expand_bboxes(bbox, dtype)[0]
+
+    # Pack the data for each of the cell bounding boxes
+    # First pack min and max values for each cell
+    min_in_cell = np.min(points[x_indices], axis=1)
+    max_in_cell = np.max(points[x_indices], axis=1)
+    bboxes = np.zeros((num_cells, 2, 3))
+    bboxes[:, 0, :] = min_in_cell
+    bboxes[:, 1, :] = max_in_cell
+    # Expand min and max values to bounding box
+    body_1 = list(expand_bboxes(bboxes, dtype))
+
+    # Compute distances and check for collision
+    distances = compute_distances_gjk(body_1, bounding_box, num_threads)
+    norm_dist = np.linalg.norm(distances, axis=1) ** 2
+    tol = 10 * np.finfo(dtype).eps
+    colliding_cells = np.flatnonzero(norm_dist < tol).astype(np.int32)
     return colliding_cells
 
 
@@ -224,9 +225,10 @@ def locator_B(x):
 
 
 @pytest.mark.skip_in_parallel
+@pytest.mark.parametrize("num_threads", [1, 2])
 @pytest.mark.parametrize("point", [np.array([0.52, 0.51, 0.0]), np.array([0.9, -0.9, 0.0])])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-def test_compute_collisions_tree_2d(point, dtype):
+def test_compute_collisions_tree_2d(point, dtype, num_threads):
     mesh_A = create_unit_square(MPI.COMM_WORLD, 3, 3, dtype=dtype)
     mesh_B = create_unit_square(MPI.COMM_WORLD, 5, 5, dtype=dtype)
     bgeom = mesh_B.geometry.x
@@ -237,18 +239,24 @@ def test_compute_collisions_tree_2d(point, dtype):
 
     entities_A = np.sort(np.unique([q[0] for q in entities]))
     entities_B = np.sort(np.unique([q[1] for q in entities]))
-    cells_A = find_colliding_cells(mesh_A, tree_B.get_bbox(tree_B.num_bboxes - 1), dtype)
-    cells_B = find_colliding_cells(mesh_B, tree_A.get_bbox(tree_A.num_bboxes - 1), dtype)
+    cells_A = find_colliding_cells(
+        mesh_A, tree_B.get_bbox(tree_B.num_bboxes - 1), dtype, num_threads
+    )
+    cells_B = find_colliding_cells(
+        mesh_B, tree_A.get_bbox(tree_A.num_bboxes - 1), dtype, num_threads
+    )
     assert np.allclose(entities_A, cells_A)
     assert np.allclose(entities_B, cells_B)
 
 
 @pytest.mark.skip_in_parallel
+@pytest.mark.parametrize("num_threads", [1, 2])
 @pytest.mark.parametrize("point", [np.array([0.52, 0.51, 0.3]), np.array([0.9, -0.9, 0.3])])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-def test_compute_collisions_tree_3d(point, dtype):
-    mesh_A = create_unit_cube(MPI.COMM_WORLD, 2, 2, 2, dtype=dtype)
-    mesh_B = create_unit_cube(MPI.COMM_WORLD, 2, 2, 2, dtype=dtype)
+def test_compute_collisions_tree_3d(point, dtype, num_threads):
+    M = 10
+    mesh_A = create_unit_cube(MPI.COMM_WORLD, M, M, M, dtype=dtype)
+    mesh_B = create_unit_cube(MPI.COMM_WORLD, M, M, M, dtype=dtype)
 
     bgeom = mesh_B.geometry.x
     bgeom += point
@@ -258,8 +266,12 @@ def test_compute_collisions_tree_3d(point, dtype):
     entities = compute_collisions_trees(tree_A, tree_B)
     entities_A = np.sort(np.unique([q[0] for q in entities]))
     entities_B = np.sort(np.unique([q[1] for q in entities]))
-    cells_A = find_colliding_cells(mesh_A, tree_B.get_bbox(tree_B.num_bboxes - 1), dtype)
-    cells_B = find_colliding_cells(mesh_B, tree_A.get_bbox(tree_A.num_bboxes - 1), dtype)
+    cells_A = find_colliding_cells(
+        mesh_A, tree_B.get_bbox(tree_B.num_bboxes - 1), dtype, num_threads
+    )
+    cells_B = find_colliding_cells(
+        mesh_B, tree_A.get_bbox(tree_A.num_bboxes - 1), dtype, num_threads
+    )
     assert np.allclose(entities_A, cells_A)
     assert np.allclose(entities_B, cells_B)