Function get_grid_idx

pointersect/pr/cuda/pr_cuda.cpp:189–203

Compute the grid index given xyz_w. Args: points: (b, n, 3) grid_size: (b, 3) long. number of grid cells in x y z. center: (b, 3) center of the grid grid_width: (b, 3) length (full width) of the grid in xyz mode: 'subidx': return sub idx 'ind': return linear index Returns: grid_idx: if mode == 'subidx': (*, n, 3) long elif mode == 'ind': (*, n) long valid_mask: (b, n) bool cell_counts: (b,

//         grid_idx = z_idx + y_dix * grid_size_z + x_idx * (grid_size_y * grid_size_z)
// 
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> get_grid_idx(
        const torch::Tensor & points,  // (*, n, 3), float
        const torch::Tensor & grid_size,  // (*, 3), long
        const torch::Tensor & grid_center,  // (*, 3), float
        const torch::Tensor & grid_width,  // (*, 3), float
        const std::string & mode = "ind"  
) { 

    CHECK_INPUT(points);
    CHECK_INPUT(grid_size);
    CHECK_INPUT(grid_center);
    CHECK_INPUT(grid_width);

    return get_grid_idx_cuda(points, grid_size, grid_center, grid_width, mode);
}