geometry_generator.cuh

/*
 * 版权所有 (c) 2023-2024, NVIDIA CORPORATION。
 *
 * 获得 Apache License, Version 2.0 (以下简称“License”) 的许可；
 * 除非遵守 License，否则您不得使用此文件。
 * 您可在以下位置获得 License 的副本
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，否则软件
 * 根据 License 分发时按“原样”提供，
 * 不提供任何明示或默示担保。
 * 有关管理权限和
 * 限制条件的特定语言，请参见 License。
 */
 
#pragma once
 
#include <cuspatial_test/random.cuh>
#include <cuspatial_test/vector_factories.cuh>
 
#include <cuspatial/cuda_utils.hpp>
#include <cuspatial/error.hpp>
#include <cuspatial/geometry/vec_2d.hpp>
#include <cuspatial/range/multipolygon_range.cuh>
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/exec_policy.hpp>
 
#include <thrust/sequence.h>
#include <thrust/tabulate.h>
 
#include <ranger/ranger.hpp>
 
#include <cmath>
 
namespace cuspatial {
namespace test {
 
namespace detail {
 
template <typename T, typename index_t>
struct tabulate_direction_functor {
 vec_2d<T> __device__ operator()(index_t i)
  {
 return vec_2d<T>{cos(static_cast<T>(i)), sin(static_cast<T>(i))};
  }
};
 
template <typename T>
struct random_walk_functor {
 T segment_length;
 
 vec_2d<T> __device__ operator()(vec_2d<T> prev, vec_2d<T> rad)
  {
 return prev + segment_length * rad;
  }
};
 
} // detail 命名空间

template <typename T>
struct multipolygon_generator_parameter {
 using element_t = T;
 
 std::size_t num_multipolygons;
 std::size_t num_polygons_per_multipolygon;
 std::size_t num_holes_per_polygon;
 std::size_t num_edges_per_ring;
 vec_2d<T> centroid;
 T radius;
 
 CUSPATIAL_HOST_DEVICE std::size_t num_polygons()
  {
 return num_multipolygons * num_polygons_per_multipolygon;
  }
 CUSPATIAL_HOST_DEVICE std::size_t num_rings() { return num_polygons() * num_rings_per_polygon(); }
 CUSPATIAL_HOST_DEVICE std::size_t num_coords() { return num_rings() * num_vertices_per_ring(); }
 CUSPATIAL_HOST_DEVICE std::size_t num_vertices_per_ring() { return num_edges_per_ring + 1; }
 CUSPATIAL_HOST_DEVICE std::size_t num_rings_per_polygon() { return num_holes_per_polygon + 1; }
 CUSPATIAL_HOST_DEVICE T hole_radius() { return radius / (num_holes_per_polygon + 1); }
};

template <typename T>
vec_2d<T> __device__ generate_ring_coordinate(std::size_t point_local_idx,
 std::size_t num_edges,
 vec_2d<T> centroid,
 T radius)
{
 // 覆盖最后一个坐标，确保环是闭合的。
 if (point_local_idx == num_edges) return vec_2d<T>{centroid.x + radius, centroid.y};
 
 T angle = (2.0 * M_PI * point_local_idx) / num_edges;
 
 return vec_2d<T>{centroid.x + radius * cos(angle), centroid.y + radius * sin(angle)};
}

template <typename T>
vec_2d<T> __device__ polygon_centroid_displacement(vec_2d<T> centroid,
 std::size_t part_local_idx,
 T radius)
{
 return centroid + vec_2d<T>{part_local_idx * radius * T{3.0}, T{0.0}};
}

template <typename T>
vec_2d<T> __device__
ring_centroid_displacement(vec_2d<T> centroid, std::size_t ring_local_idx, T radius, T hole_radius)
{
 // 这是外壳
 if (ring_local_idx == 0) { return centroid; }
 
 // 这是孔洞
 ring_local_idx -= 1; // 偏移孔洞索引，使其从0开始
 T max_hole_displacement = radius - hole_radius;
 T displacement_x = -max_hole_displacement + ring_local_idx * hole_radius * 2;
 T displacement_y = 0.0;
 return centroid + vec_2d<T>{displacement_x, displacement_y};
}

template <typename T, typename MultipolygonRange>
CUSPATIAL_KERNEL void generate_multipolygon_array_coordinates(
 MultipolygonRange multipolygons, multipolygon_generator_parameter<T> params)
{
 for (auto idx : ranger::grid_stride_range(multipolygons.num_points())) {
 auto ring_idx = multipolygons.ring_idx_from_point_idx(idx);
 auto part_idx = multipolygons.part_idx_from_ring_idx(ring_idx);
 auto geometry_idx = multipolygons.geometry_idx_from_part_idx(part_idx);
 
 auto point_local_idx = idx - params.num_vertices_per_ring() * ring_idx;
 auto ring_local_idx = ring_idx - params.num_rings_per_polygon() * part_idx;
 auto part_local_idx = part_idx - params.num_polygons_per_multipolygon * geometry_idx;
 
 auto centroid = ring_centroid_displacement(
 polygon_centroid_displacement(params.centroid, part_local_idx, params.radius),
 ring_local_idx,
 params.radius,
 params.hole_radius());
 
 if (ring_local_idx == 0) // Generate coordinate for shell
 multipolygons.point_begin()[idx] = generate_ring_coordinate(
 point_local_idx, params.num_edges_per_ring, centroid, params.radius);
 else // Generate coordinate for holes
 multipolygons.point_begin()[idx] = generate_ring_coordinate(
 point_local_idx, params.num_edges_per_ring, centroid, params.hole_radius());
  }
}

template <typename T>
auto generate_multipolygon_array(multipolygon_generator_parameter<T> params,
 rmm::cuda_stream_view stream)
{
 rmm::device_uvector<std::size_t> geometry_offsets(params.num_multipolygons + 1, stream);
 rmm::device_uvector<std::size_t> part_offsets(params.num_polygons() + 1, stream);
 rmm::device_uvector<std::size_t> ring_offsets(params.num_rings() + 1, stream);
 rmm::device_uvector<vec_2d<T>> coordinates(params.num_coords(), stream);
 
 thrust::sequence(rmm::exec_policy(stream),
 ring_offsets.begin(),
 ring_offsets.end(),
 std::size_t{0},
 params.num_vertices_per_ring());
 
 thrust::sequence(rmm::exec_policy(stream),
 part_offsets.begin(),
 part_offsets.end(),
 std::size_t{0},
 params.num_rings_per_polygon());
 
 thrust::sequence(rmm::exec_policy(stream),
 geometry_offsets.begin(),
 geometry_offsets.end(),
 std::size_t{0},
 params.num_polygons_per_multipolygon);
 
 auto multipolygons = multipolygon_range(geometry_offsets.begin(),
 geometry_offsets.end(),
 part_offsets.begin(),
 part_offsets.end(),
 ring_offsets.begin(),
 ring_offsets.end(),
 coordinates.begin(),
 coordinates.end());
 
 auto [tpb, nblocks] = grid_1d(multipolygons.num_points());
 
 generate_multipolygon_array_coordinates<T><<<nblocks, tpb, 0, stream>>>(multipolygons, params);
 
 CUSPATIAL_CHECK_CUDA(stream.value());
 
 return make_multipolygon_array<std::size_t, vec_2d<T>>(std::move(geometry_offsets),
 std::move(part_offsets),
 std::move(ring_offsets),
 std::move(coordinates));
}

template <typename T>
struct multilinestring_generator_parameter {
 std::size_t num_multilinestrings;
 std::size_t num_linestrings_per_multilinestring;
 std::size_t num_segments_per_linestring;
 T segment_length;
 vec_2d<T> origin;
 
 std::size_t num_linestrings()
  {
 return num_multilinestrings * num_linestrings_per_multilinestring;
  }
 
 std::size_t num_points_per_linestring() { return num_segments_per_linestring + 1; }
 
 std::size_t num_segments() { return num_linestrings() * num_segments_per_linestring; }
 std::size_t num_points() { return num_linestrings() * num_points_per_linestring(); }
};

template <typename T>
auto generate_multilinestring_array(multilinestring_generator_parameter<T> params,
 rmm::cuda_stream_view stream)
{
 rmm::device_uvector<std::size_t> geometry_offset(params.num_multilinestrings + 1, stream);
 rmm::device_uvector<std::size_t> part_offset(params.num_linestrings() + 1, stream);
 rmm::device_uvector<vec_2d<T>> points(params.num_points(), stream);
 
 thrust::sequence(rmm::exec_policy(stream),
 geometry_offset.begin(),
 geometry_offset.end(),
 static_cast<std::size_t>(0),
 params.num_linestrings_per_multilinestring);
 
 thrust::sequence(rmm::exec_policy(stream),
 part_offset.begin(),
 part_offset.end(),
 static_cast<std::size_t>(0),
 params.num_segments_per_linestring + 1);
 
 thrust::tabulate(rmm::exec_policy(stream),
 points.begin(),
 points.end(),
 detail::tabulate_direction_functor<T, std::size_t>{});
 
 thrust::exclusive_scan(rmm::exec_policy(stream),
 points.begin(),
 points.end(),
 points.begin(),
 params.origin,
 detail::random_walk_functor<T>{params.segment_length});
 
 return make_multilinestring_array<std::size_t, T>(
 std::move(geometry_offset), std::move(part_offset), std::move(points));
}

template <typename T>
struct multipoint_generator_parameter {
 using element_t = T;
 
 std::size_t num_multipoints;
 std::size_t num_points_per_multipoints;
 vec_2d<T> lower_left;
 vec_2d<T> upper_right;
 
 CUSPATIAL_HOST_DEVICE std::size_t num_points()
  {
 return num_multipoints * num_points_per_multipoints;
  }
};

template <typename T>
auto generate_multipoint_array(multipoint_generator_parameter<T> params,
 rmm::cuda_stream_view stream)
{
 rmm::device_uvector<vec_2d<T>> coordinates(params.num_points(), stream);
 rmm::device_uvector<std::size_t> offsets(params.num_multipoints + 1, stream);
 
 thrust::sequence(rmm::exec_policy(stream),
 offsets.begin(),
 offsets.end(),
 std::size_t{0},
 params.num_points_per_multipoints);
 
 auto golden_ratio = (1 + std::sqrt(T{5})) / 2;
 auto engine_x = deterministic_engine(golden_ratio * params.num_points());
 auto engine_y = deterministic_engine((1 / golden_ratio) * params.num_points());
 
 auto x_dist = make_uniform_dist(params.lower_left.x, params.upper_right.x);
 auto y_dist = make_uniform_dist(params.lower_left.y, params.upper_right.y);
 
 auto point_gen =
 point_generator(params.lower_left, params.upper_right, engine_x, engine_y, x_dist, y_dist);
 
 thrust::tabulate(rmm::exec_policy(stream), coordinates.begin(), coordinates.end(), point_gen);
 
 return make_multipoint_array(std::move(offsets), std::move(coordinates));
}
 
} // test 命名空间
} // cuspatial 命名空间