Function find_k_neighbor_points_of_rays

pointersect/pr/cuda/pr_cuda.cpp:514–547

Find the k nearest points within `radius` of a ray, ie, the vertical distance from the point to ray <= radius. If not enough points found, return dummy index Returns: ray2pidx: (b, M), all the points. The pidxs of a ray is found via ray2pidx[b][ray_start_idx[m]:ray_end_idx[m]] ray_start_idx: (b, m), ray_end_idx: (b, m) Note: Our algorithm is very simple. In order to be parallelized on gpu easily

//     5. given grid idxs for each ray, gather point idxs
// 
std::tuple<torch::Tensor, torch::Tensor> find_k_neighbor_points_of_rays(
        torch::Tensor points,  // (b, n, 3), float
        torch::Tensor ray_origins, // (b, m, 3), float
        torch::Tensor ray_directions,  // (b, m, 3), float
        torch::Tensor ray_radius,  // (b,), float
        torch::Tensor grid_size,  // (b, 3), long
        torch::Tensor grid_center,  // (b, 3), float
        torch::Tensor grid_width,  // (b, 3), float
        int64_t k,
        float t_min = 0.,
        float t_max = 1.e12
) {
    
    CHECK_INPUT(points);
    CHECK_INPUT(ray_origins);
    CHECK_INPUT(ray_directions);
    CHECK_INPUT(ray_radius);
    CHECK_INPUT(grid_size);
    CHECK_INPUT(grid_center);
    CHECK_INPUT(grid_width);

    return find_k_neighbor_points_of_rays_cuda(
        points,  // (b, n, 3), 
        ray_origins, // (b, m, 3), float
        ray_directions,  // (b, m, 3), float
        ray_radius,  // (b,), float
        grid_size,  // (b, 3), long
        grid_center,  // (b, 3), float
        grid_width,  // (b, 3), float
        k,
        t_min,
        t_max
    );
}


torch::Tensor keep_min_k_values(