Function find_neighbor_points_of_rays

pointersect/pr/cpp/pr.cpp:803–952

Find the points within `radius` of a ray, ie, the vertical distance from the point to ray <= radius. Returns: list of list of list: b -> m -> n_idx Note: Our algorithm is very simple. In order to be parallelized on gpu easily, we want to every thread to have as few branching conditions (if/else) as possible. We will: - parallelize on ray - do the same thing for all rays. - ignore all points ou

//     5. given grid idxs for each ray, gather point idxs
// 
std::vector<llist> find_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
        float t_min = 0.,
        float t_max = 1.e12,
        std::string version = "v2"  // v2 is faster
) {
    
    assert(grid_size.dtype() == torch::kLong);

    auto batch_size = ray_origins.size(0);
    auto n_points = points.size(1);
    auto n_rays = ray_origins.size(1);
    auto device = ray_origins.device();
    auto total_cells = torch::prod(grid_size, -1);  // (b,) long

    // get grid idx of each point
    auto start = high_resolution_clock::now();
    auto tmp_grid_out = get_grid_idx(
        points,  
        grid_size,
        grid_center,
        grid_width,
        "ind"
    );  
    torch::Tensor &  grid_idxs = tmp_grid_out[0];  // (b, n), long  [0, total_cells-1]
    torch::Tensor & valid_mask = tmp_grid_out[1];  // (b, n), bool 
    auto stop = high_resolution_clock::now();
    auto duration = duration_cast<microseconds>(stop - start);
    // std::cout<<"get_grid_idx: "<<duration.count()<<std::endl;
    

    // gather points belonging to each cell
    start = high_resolution_clock::now();
    std::vector<idxmap> all_cell2pidx = gather_points(
        grid_idxs,  // (b, n)
        total_cells,  // (b,)
        valid_mask  // (b, n)
    );  // b -> gidx -> point_idxs
    stop = high_resolution_clock::now();
    duration = duration_cast<microseconds>(stop - start);
    // std::cout<<"gather points: "<<duration.count()<<std::endl;


    // gather the cells intersected by the ray
    start = high_resolution_clock::now();
    std::vector<idxmap> all_ray2gidxs;
    if (version.compare("v1") == 0) {
        all_ray2gidxs = std::move(grid_ray_intersection(
            ray_origins,
            ray_directions,
            ray_radius,
            grid_size,