From 24ae670440116b7dbae83aa3805cc82f6750d58a Mon Sep 17 00:00:00 2001
From: Marco Feder <marco.feder@sissa.it>
Date: Fri, 30 Jun 2023 01:33:55 +0200
Subject: [PATCH] Use arborX for spatial search in
 distributed_compute_point_locations

---
 source/grid/grid_tools.cc | 109 ++++++++++++++++++++++++--------------
 1 file changed, 70 insertions(+), 39 deletions(-)

diff --git a/source/grid/grid_tools.cc b/source/grid/grid_tools.cc
index e2920086f7..d40be81d75 100644
--- a/source/grid/grid_tools.cc
+++ b/source/grid/grid_tools.cc
@@ -20,6 +20,11 @@
 #include <deal.II/base/quadrature_lib.h>
 #include <deal.II/base/thread_management.h>
 
+#ifdef DEAL_II_WITH_ARBORX
+#  include <deal.II/arborx/access_traits.h>
+#  include <deal.II/arborx/distributed_tree.h>
+#endif
+
 #ifdef DEAL_II_WITH_CGAL
 #  include <deal.II/cgal/intersections.h>
 #  include <deal.II/cgal/utilities.h>
@@ -5880,67 +5885,93 @@ namespace GridTools
       const std::vector<T> &                                 entities,
       const double                                           tolerance)
     {
+      std::vector<std::pair<unsigned int, unsigned int>> ranks_and_indices;
+      ranks_and_indices.reserve(entities.size());
+
+#if !defined(DEAL_II_WITH_ARBORX)
+
       std::vector<std::vector<BoundingBox<spacedim>>> global_bboxes_temp;
       auto *global_bboxes_to_be_used = &global_bboxes;
-
-      if (global_bboxes.size() == 1) // TODO: and not ArborX installed
+      if (global_bboxes.size() == 1)
         {
           global_bboxes_temp =
             Utilities::MPI::all_gather(comm, global_bboxes[0]);
           global_bboxes_to_be_used = &global_bboxes_temp;
         }
 
-      std::vector<std::pair<unsigned int, unsigned int>> ranks_and_indices;
-      ranks_and_indices.reserve(entities.size());
 
-      if (true)
+      // helper function to determine if a bounding box is valid
+      const auto is_valid = [](const auto &bb) {
+        for (unsigned int i = 0; i < spacedim; ++i)
+          if (bb.get_boundary_points().first[i] >
+              bb.get_boundary_points().second[i])
+            return false;
+
+        return true;
+      };
+
+      // linearize vector of vectors
+      std::vector<std::pair<BoundingBox<spacedim>, unsigned int>>
+        boxes_and_ranks;
+
+      for (unsigned rank = 0; rank < global_bboxes_to_be_used->size(); ++rank)
+        for (const auto &box : (*global_bboxes_to_be_used)[rank])
+          if (is_valid(box))
+            boxes_and_ranks.emplace_back(box, rank);
+
+      // pack boxes into r-tree
+      const auto tree = pack_rtree(boxes_and_ranks);
+
+      // loop over all entities
+      for (unsigned int i = 0; i < entities.size(); ++i)
         {
-          // helper function to determine if a bounding box is valid
-          const auto is_valid = [](const auto &bb) {
-            for (unsigned int i = 0; i < spacedim; ++i)
-              if (bb.get_boundary_points().first[i] >
-                  bb.get_boundary_points().second[i])
-                return false;
+          // create a bounding box with tolerance
+          const auto bb =
+            BoundingBox<spacedim>(entities[i]).create_extended(tolerance);
 
-            return true;
-          };
+          // determine ranks potentially owning point/bounding box
+          std::set<unsigned int> my_ranks;
 
-          // linearize vector of vectors
-          std::vector<std::pair<BoundingBox<spacedim>, unsigned int>>
-            boxes_and_ranks;
+          for (const auto &box_and_rank :
+               tree | boost::geometry::index::adaptors::queried(
+                        boost::geometry::index::intersects(bb)))
+            my_ranks.insert(box_and_rank.second);
 
-          for (unsigned rank = 0; rank < global_bboxes_to_be_used->size();
-               ++rank)
-            for (const auto &box : (*global_bboxes_to_be_used)[rank])
-              if (is_valid(box))
-                boxes_and_ranks.emplace_back(box, rank);
+          for (const auto rank : my_ranks)
+            ranks_and_indices.emplace_back(rank, i);
+        }
 
-          // pack boxes into r-tree
-          const auto tree = pack_rtree(boxes_and_ranks);
+#else
+      if (global_bboxes.size() == 1)
+        {
+          // Create query bounding boxes out of entities
+          std::vector<BoundingBox<spacedim>> query_bounding_boxes;
+          for (const auto &entity : entities)
+            query_bounding_boxes.push_back(
+              BoundingBox<spacedim>(entity).create_extended(tolerance));
+          ArborXWrappers::BoundingBoxIntersectPredicate bb_intersect(
+            query_bounding_boxes);
 
-          // loop over all entities
-          for (unsigned int i = 0; i < entities.size(); ++i)
-            {
-              // create a bounding box with tolerance
-              const auto bb =
-                BoundingBox<spacedim>(entities[i]).create_extended(tolerance);
+          // Create bounding volume hierarchy out of input vector of bounding
+          ArborXWrappers::DistributedTree distributed_tree(comm,
+                                                           global_bboxes[0]);
 
-              // determine ranks potentially owning point/bounding box
-              std::set<unsigned int> my_ranks;
+          auto indices_ranks_offset = distributed_tree.query(bb_intersect);
 
-              for (const auto &box_and_rank :
-                   tree | boost::geometry::index::adaptors::queried(
-                            boost::geometry::index::intersects(bb)))
-                my_ranks.insert(box_and_rank.second);
+          auto indices_ranks = indices_ranks_offset.first;
+          auto offsets       = indices_ranks_offset.second;
 
+          std::set<unsigned int> my_ranks;
+          for (unsigned int i = 0; i < offsets.size() - 1; ++i)
+            {
+              for (int j = offsets[i]; j < offsets[i + 1]; ++j)
+                my_ranks.insert(indices_ranks[j].second);
               for (const auto rank : my_ranks)
                 ranks_and_indices.emplace_back(rank, i);
+              my_ranks.clear();
             }
         }
-      else
-        {
-          // TODO: use ArborX
-        }
+#endif
 
       // convert to CRS
       std::sort(ranks_and_indices.begin(), ranks_and_indices.end());
-- 
2.39.5