point_quadtree.cuh

/*
 * 版权所有 (c) 2022-2024, NVIDIA CORPORATION.
 *
 * 获得 Apache 许可证 2.0 版 ("许可证") 的许可；
 * 除非遵守许可证，否则您不得使用此文件。
 * 您可以在以下地址获取许可证的副本
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，软件是
 * 根据许可证分发是基于“原样”分发的，
 * 不附带任何形式的明示或暗示的担保或条件。
 * 有关管理许可和
 * 限制的特定语言，请参阅许可证。
 */
 
#pragma once
 
#include <cuspatial/geometry/vec_2d.hpp>
#include <cuspatial/traits.hpp>
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <thrust/iterator/zip_iterator.h>
#include <thrust/tuple.h>
 
#include <tuple>
 
namespace cuspatial {

 
struct point_quadtree {
 // quadtree 节点键的 uint32_t 向量
 rmm::device_uvector<uint32_t> key;
 // quadtree 层级的 uint8_t 向量
 rmm::device_uvector<uint8_t> level;
 // bool 向量，指示节点是父节点 (true) 还是叶节点 (false)
 rmm::device_uvector<bool> internal_node_flag;
 // 子节点数（如果是内部节点）或点数的 uint32_t 向量
 rmm::device_uvector<uint32_t> length;
 // 第一个子节点位置（如果是内部节点）或第一个点位置的 uint32_t 向量
 rmm::device_uvector<uint32_t> offset;
};
 
struct point_quadtree_ref {
 using key_iterator = decltype(point_quadtree::key)::const_iterator;
 using level_iterator = decltype(point_quadtree::level)::const_iterator;
 using internal_node_flag_iterator = decltype(point_quadtree::internal_node_flag)::const_iterator;
 using length_iterator = decltype(point_quadtree::length)::const_iterator;
 using offset_iterator = decltype(point_quadtree::offset)::const_iterator;

 point_quadtree_ref(point_quadtree const& quadtree)
 : _key_begin(quadtree.key.begin()),
 _key_end(quadtree.key.end()),
 _level_begin(quadtree.level.begin()),
 _level_end(quadtree.level.end()),
 _internal_node_flag_begin(quadtree.internal_node_flag.begin()),
 _internal_node_flag_end(quadtree.internal_node_flag.end()),
 _length_begin(quadtree.length.begin()),
 _length_end(quadtree.length.end()),
 _offset_begin(quadtree.offset.begin()),
 _offset_end(quadtree.offset.end())
  {
  }

 point_quadtree_ref(key_iterator key_begin,
 key_iterator key_end,
 level_iterator level_begin,
 internal_node_flag_iterator internal_node_flag_begin,
 length_iterator length_begin,
 offset_iterator offset_begin)
 : _key_begin(key_begin),
 _key_end(key_end),
 _level_begin(level_begin),
 _level_end(level_begin + std::distance(key_begin, key_end)),
 _internal_node_flag_begin(internal_node_flag_begin),
 _internal_node_flag_end(internal_node_flag_begin + std::distance(key_begin, key_end)),
 _length_begin(length_begin),
 _length_end(length_begin + std::distance(key_begin, key_end)),
 _offset_begin(offset_begin),
 _offset_end(offset_begin + std::distance(key_begin, key_end))
  {
  }

 CUSPATIAL_HOST_DEVICE auto num_nodes() const { return std::distance(_key_begin, _key_end); }

 CUSPATIAL_HOST_DEVICE auto key_begin() const { return _key_begin; }
 CUSPATIAL_HOST_DEVICE auto key_end() const { return _key_end; }

 CUSPATIAL_HOST_DEVICE auto level_begin() const { return _level_begin; }
 CUSPATIAL_HOST_DEVICE auto level_end() const { return _level_end; }

 CUSPATIAL_HOST_DEVICE auto internal_node_flag_begin() const { return _internal_node_flag_begin; }
 CUSPATIAL_HOST_DEVICE auto internal_node_flag_end() const { return _internal_node_flag_end; }

 CUSPATIAL_HOST_DEVICE auto length_begin() const { return _length_begin; }
 CUSPATIAL_HOST_DEVICE auto length_end() const { return _length_end; }

 CUSPATIAL_HOST_DEVICE auto offset_begin() const { return _offset_begin; }
 CUSPATIAL_HOST_DEVICE auto offset_end() const { return _offset_end; }
 
 protected
 key_iterator _key_begin;
 key_iterator _key_end;
 level_iterator _level_begin;
 level_iterator _level_end;
 internal_node_flag_iterator _internal_node_flag_begin;
 internal_node_flag_iterator _internal_node_flag_end;
 length_iterator _length_begin;
 length_iterator _length_end;
 offset_iterator _offset_begin;
 offset_iterator _offset_end;
};

template <class PointIterator, class T = typename cuspatial::iterator_value_type<PointIterator>>
std::pair<rmm::device_uvector<uint32_t>, point_quadtree> quadtree_on_points(
 PointIterator points_first,
 PointIterator points_last,
 vec_2d<T> vertex_1,
 vec_2d<T> vertex_2,
 T scale,
 int8_t max_depth,
 int32_t max_size,
 rmm::cuda_stream_view stream = rmm::cuda_stream_default,
 rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());

 
} // 命名空间 cuspatial
 
#include <cuspatial/detail/point_quadtree.cuh>