Header-only quadtree_point_in_polygon

cpp/benchmarks/floating_point_equality.cu

@@ -45,14 +45,14 @@ template <class FloatsIter>
 void generate_floats(FloatsIter begin, FloatsIter end)
 {
   using T       = typename std::iterator_traits<FloatsIter>::value_type;
-  auto engine_x = deterministic_engine(std::distance(begin, end));
+  auto engine_x = cuspatial::test::deterministic_engine(std::distance(begin, end));
 
   auto lo = std::numeric_limits<T>::min();
   auto hi = std::numeric_limits<T>::max();
 
-  auto x_dist = make_uniform_dist(lo, hi);
+  auto x_dist = cuspatial::test::make_uniform_dist(lo, hi);
 
-  auto x_gen = value_generator{lo, hi, engine_x, x_dist};
+  auto x_gen = cuspatial::test::value_generator{lo, hi, engine_x, x_dist};
 
   thrust::tabulate(
     rmm::exec_policy(), begin, end, [x_gen] __device__(size_t n) mutable { return x_gen(n); });

cpp/benchmarks/points_in_range.cu

@@ -15,6 +15,7 @@
  */
 
 #include <benchmarks/fixture/rmm_pool_raii.hpp>
+
 #include <cuspatial_test/random.cuh>
 
 #include <cuspatial/detail/iterator.hpp>
@@ -54,14 +55,14 @@ using namespace cuspatial;
 template <class PointsIter, typename T>
 void generate_points(PointsIter begin, PointsIter end, vec_2d<T> range_min, vec_2d<T> range_max)
 {
-  auto engine_x = deterministic_engine(std::distance(begin, end));
-  auto engine_y = deterministic_engine(2 * std::distance(begin, end));
+  auto engine_x = cuspatial::test::deterministic_engine(std::distance(begin, end));
+  auto engine_y = cuspatial::test::deterministic_engine(2 * std::distance(begin, end));
 
-  auto x_dist = make_uniform_dist(range_min.x, range_max.x);
-  auto y_dist = make_uniform_dist(range_min.y, range_max.y);
+  auto x_dist = cuspatial::test::make_uniform_dist(range_min.x, range_max.x);
+  auto y_dist = cuspatial::test::make_uniform_dist(range_min.y, range_max.y);
 
-  auto x_gen = value_generator{range_min.x, range_max.x, engine_x, x_dist};
-  auto y_gen = value_generator{range_min.y, range_max.y, engine_y, y_dist};
+  auto x_gen = cuspatial::test::value_generator{range_min.x, range_max.x, engine_x, x_dist};
+  auto y_gen = cuspatial::test::value_generator{range_min.y, range_max.y, engine_y, y_dist};
 
   thrust::tabulate(rmm::exec_policy(), begin, end, [x_gen, y_gen] __device__(size_t n) mutable {
     return vec_2d<T>{x_gen(n), y_gen(n)};

cpp/include/cuspatial/experimental/detail/algorithm/is_point_in_polygon.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

cpp/include/cuspatial/experimental/detail/join/get_quad_and_local_point_indices.cuh

@@ -23,8 +23,11 @@
 namespace cuspatial {
 namespace detail {
 
-inline __device__ std::pair<uint32_t, uint32_t> get_quad_and_local_point_indices(
-  uint32_t const global_index, uint32_t const* point_offsets, uint32_t const* point_offsets_end)
+template <typename IndexType, typename PointOffsetIterator>
+inline __device__ std::pair<IndexType, IndexType> get_quad_and_local_point_indices(
+  IndexType const global_index,
+  PointOffsetIterator point_offsets_begin,
+  PointOffsetIterator point_offsets_end)
 {
   // Calculate the position in "point_offsets" that `global_index` falls between.
   // This position is the index of the poly/quad pair for this `global_index`.
@@ -33,10 +36,10 @@ inline __device__ std::pair<uint32_t, uint32_t> get_quad_and_local_point_indices
   // quadrant. Adding this zero-based position to the quadrant's first point position in the
   // quadtree yields the "global" position in the `point_indices` map.
   auto const local_point_offset =
-    thrust::upper_bound(thrust::seq, point_offsets, point_offsets_end, global_index) - 1;
+    thrust::upper_bound(thrust::seq, point_offsets_begin, point_offsets_end, global_index) - 1;
   return std::make_pair(
     // quad_poly_index
-    thrust::distance(point_offsets, local_point_offset),
+    thrust::distance(point_offsets_begin, local_point_offset),
     // local_point_index
     global_index - *local_point_offset);
 }

cpp/include/cuspatial/experimental/detail/join/intersection.cuh

@@ -17,6 +17,8 @@
 #pragma once
 
 #include <cuspatial/detail/utility/z_order.cuh>
+#include <cuspatial/experimental/geometry/box.hpp>
+#include <cuspatial/experimental/point_quadtree.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
@@ -69,37 +71,32 @@ inline int32_t remove_non_quad_intersections(InputIterator input_begin,
 }
 
 template <class T,
-          class KeyIterator,
-          class LevelIterator,
-          class IsInternalIterator,
           class BoundingBoxIterator,
           class NodeIndicesIterator,
           class BBoxIndicesIterator,
           class NodePairsIterator,
           class LeafPairsIterator>
-inline std::pair<int32_t, int32_t> find_intersections(KeyIterator keys_first,
-                                                      LevelIterator levels_first,
-                                                      IsInternalIterator is_internal_node_first,
+inline std::pair<int32_t, int32_t> find_intersections(point_quadtree_ref quadtree,
                                                       BoundingBoxIterator bounding_box_first,
                                                       NodeIndicesIterator node_indices,
                                                       BBoxIndicesIterator bbox_indices,
                                                       NodePairsIterator node_pairs,
                                                       LeafPairsIterator leaf_pairs,
                                                       int32_t num_pairs,
-                                                      T x_min,
-                                                      T y_min,
+                                                      vec_2d<T> const& v_min,
                                                       T scale,
                                                       int8_t max_depth,
                                                       rmm::cuda_stream_view stream)
 {
-  auto nodes_first = thrust::make_zip_iterator(keys_first, levels_first, is_internal_node_first);
+  auto nodes_first = thrust::make_zip_iterator(
+    quadtree.key_begin(), quadtree.level_begin(), quadtree.internal_node_flag_begin());
 
   thrust::transform(
     rmm::exec_policy(stream),
     thrust::make_zip_iterator(node_indices, bbox_indices),
     thrust::make_zip_iterator(node_indices, bbox_indices) + num_pairs,
     node_pairs,
-    [x_min, y_min, scale, max_depth, nodes = nodes_first, bboxes = bounding_box_first] __device__(
+    [v_min, scale, max_depth, nodes = nodes_first, bboxes = bounding_box_first] __device__(
       auto const& node_and_bbox) {
       auto const& node_idx = thrust::get<0>(node_and_bbox);
       auto const& bbox_idx = thrust::get<1>(node_and_bbox);
@@ -109,24 +106,20 @@ inline std::pair<int32_t, int32_t> find_intersections(KeyIterator keys_first,
       uint8_t const& level            = thrust::get<1>(node);
       uint8_t const& is_internal_node = thrust::get<2>(node);
 
-      auto const& bbox       = bboxes[bbox_idx];
-      auto const& bbox_min   = bbox.v1;
-      auto const& bbox_max   = bbox.v2;
-      auto const& poly_x_min = bbox_min.x;
-      auto const& poly_y_min = bbox_min.y;
-      auto const& poly_x_max = bbox_max.x;
-      auto const& poly_y_max = bbox_max.y;
+      box<T> const bbox        = bboxes[bbox_idx];
+      vec_2d<T> const bbox_min = bbox.v1;
+      vec_2d<T> const bbox_max = bbox.v2;
 
       T const key_x       = utility::z_order_x(key);
       T const key_y       = utility::z_order_y(key);
       T const level_scale = scale * (1 << (max_depth - 1 - level));
-      T const node_x_min  = x_min + (key_x + 0) * level_scale;
-      T const node_y_min  = y_min + (key_y + 0) * level_scale;
-      T const node_x_max  = x_min + (key_x + 1) * level_scale;
-      T const node_y_max  = y_min + (key_y + 1) * level_scale;
+      T const node_x_min  = v_min.x + (key_x + 0) * level_scale;
+      T const node_y_min  = v_min.y + (key_y + 0) * level_scale;
+      T const node_x_max  = v_min.x + (key_x + 1) * level_scale;
+      T const node_y_max  = v_min.y + (key_y + 1) * level_scale;
 
-      if ((node_x_min > poly_x_max) || (node_x_max < poly_x_min) || (node_y_min > poly_y_max) ||
-          (node_y_max < poly_y_min)) {
+      if ((node_x_min > bbox_max.x) || (node_x_max < bbox_min.x) || (node_y_min > bbox_max.y) ||
+          (node_y_max < bbox_min.y)) {
         // if no overlap, return type = none_indicator
         return thrust::make_tuple(none_indicator, level, node_idx, bbox_idx);
       }

cpp/include/cuspatial/experimental/detail/point_quadtree.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -241,8 +241,8 @@ std::pair<rmm::device_uvector<uint32_t>, point_quadtree> quadtree_on_points(
   T scale,
   int8_t max_depth,
   int32_t max_size,
-  rmm::mr::device_memory_resource* mr,
-  rmm::cuda_stream_view stream)
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr)
 {
   auto num_points = thrust::distance(points_first, points_last);
   if (num_points <= 0) {

cpp/include/cuspatial/experimental/detail/quadtree_bbox_filtering.cuh

@@ -29,38 +29,27 @@
 
 namespace cuspatial {
 
-template <class KeyIterator,
-          class LevelIterator,
-          class IsInternalIterator,
-          class BoundingBoxIterator,
-          class T>
+template <class BoundingBoxIterator, class T>
 std::pair<rmm::device_uvector<uint32_t>, rmm::device_uvector<uint32_t>>
-join_quadtree_and_bounding_boxes(KeyIterator keys_first,
-                                 KeyIterator keys_last,
-                                 LevelIterator levels_first,
-                                 IsInternalIterator is_internal_nodes_first,
-                                 KeyIterator lengths_first,
-                                 KeyIterator offsets_first,
+join_quadtree_and_bounding_boxes(point_quadtree_ref quadtree,
                                  BoundingBoxIterator bounding_boxes_first,
                                  BoundingBoxIterator bounding_boxes_last,
-                                 T x_min,
-                                 T y_min,
+                                 vec_2d<T> const& v_min,
                                  T scale,
                                  int8_t max_depth,
-                                 rmm::mr::device_memory_resource* mr,
-                                 rmm::cuda_stream_view stream)
+                                 rmm::cuda_stream_view stream,
+                                 rmm::mr::device_memory_resource* mr)
 {
   static_assert(is_same<T, cuspatial::iterator_vec_base_type<BoundingBoxIterator>>(),
                 "Iterator value_type mismatch");
 
-  auto const num_nodes = std::distance(keys_first, keys_last);
   auto const num_boxes = std::distance(bounding_boxes_first, bounding_boxes_last);
 
   // Count the number of top-level nodes to start.
   // This could be provided explicitly, but count_if should be fast enough.
   auto num_top_level_leaves = thrust::count_if(rmm::exec_policy(stream),
-                                               levels_first,
-                                               levels_first + num_nodes,
+                                               quadtree.level_begin(),
+                                               quadtree.level_end(),
                                                thrust::placeholders::_1 == 0);
 
   auto num_pairs = num_top_level_leaves * num_boxes;
@@ -96,9 +85,7 @@ join_quadtree_and_bounding_boxes(KeyIterator keys_first,
 
   // Find intersections for all the top level quadrants and bounding boxes
   std::tie(num_parents, num_leaves) =
-    detail::find_intersections(keys_first,
-                               levels_first,
-                               is_internal_nodes_first,
+    detail::find_intersections(quadtree,
                                bounding_boxes_first,
                                // The top-level node indices
                                detail::make_counting_transform_iterator(
@@ -111,8 +98,7 @@ join_quadtree_and_bounding_boxes(KeyIterator keys_first,
                                // found intersecting quadrant and bbox indices for output
                                make_output_values_iter(),
                                num_pairs,
-                               x_min,
-                               y_min,
+                               v_min,
                                scale,
                                max_depth,
                                stream);
@@ -141,8 +127,8 @@ join_quadtree_and_bounding_boxes(KeyIterator keys_first,
     // Walk one level down and fill the current level's vectors with
     // the next level's quadrant info and bbox indices.
     std::tie(num_pairs, cur_types, cur_levels, cur_node_idxs, cur_bbox_idxs) =
-      detail::descend_quadtree(lengths_first,
-                               offsets_first,
+      detail::descend_quadtree(quadtree.length_begin(),
+                               quadtree.offset_begin(),
                                num_parents,
                                cur_types,
                                cur_levels,
@@ -152,9 +138,7 @@ join_quadtree_and_bounding_boxes(KeyIterator keys_first,
 
     // Find intersections for the the next level's quadrants and bboxes
     std::tie(num_parents, num_leaves) =
-      detail::find_intersections(keys_first,
-                                 levels_first,
-                                 is_internal_nodes_first,
+      detail::find_intersections(quadtree,
                                  bounding_boxes_first,
                                  cur_node_idxs.begin(),
                                  cur_bbox_idxs.begin(),
@@ -163,8 +147,7 @@ join_quadtree_and_bounding_boxes(KeyIterator keys_first,
                                  // found intersecting quadrant and bbox indices for output
                                  make_output_values_iter(),
                                  num_pairs,
-                                 x_min,
-                                 y_min,
+                                 v_min,
                                  scale,
                                  max_depth,
                                  stream);
@@ -177,7 +160,8 @@ join_quadtree_and_bounding_boxes(KeyIterator keys_first,
     // Copy the relevant node offsets into a temporary vector so we don't modify the quadtree column
     rmm::device_uvector<uint32_t> tmp_node_offsets(num_results, stream);
 
-    auto const iter = thrust::make_permutation_iterator(offsets_first, out_node_idxs.begin());
+    auto const iter =
+      thrust::make_permutation_iterator(quadtree.offset_begin(), out_node_idxs.begin());
 
     thrust::copy(rmm::exec_policy(stream), iter, iter + num_results, tmp_node_offsets.begin());