Function grid_ray_intersection

pointersect/pr/cpp/pr.cpp:313–512

Compute the intersection between grid cells and the ray. Args: ray_origins: (b, m, 3) ray_directions: (b, m, 3) ray_radius: (b,) grid_size: (b, 3) long, xyz grid_center: (b, 3), xyz grid_width: (b, 3), xyz Returns: list of list of list: b -> ray_idx -> grid_idx (including -1, outside the grid) Algorithm: for each ray (px, py, pz, dx, dy, dz) if dx.abs() < 1e-8: # use y=c plane else: # use x=c

//             surrounding the point.
// 
std::vector<idxmap> grid_ray_intersection(
        const torch::Tensor & ray_origins,  // (b, m, 3), float
        const torch::Tensor & ray_directions,  // (b, m, 3), float
        const torch::Tensor & ray_radius,   //  (b, )m float
        const torch::Tensor & grid_size,  // (b, 3), long
        const torch::Tensor & grid_center,  // (b, 3) float
        const torch::Tensor & grid_width  // (b, 3) float
) {

    assert(grid_size.dtype() == torch::kLong);

    auto batch_size = ray_origins.size(0);
    auto n_rays = ray_origins.size(1);
    auto device = ray_origins.device();

    auto grid_from = grid_center - grid_width * 0.5;  // (b, 3)
    auto cell_width = grid_width / grid_size;  // (b, 3)
    auto total_cells = torch::prod(grid_size, -1);  // (b,)

    // determine whether to intersect with x=c plane, y=c plane, or z=c plane.
    // It is important to select a direction so that the intersection point on two nearby planes
    // do not exceed one grid in one of the rest of the two directions.
    // Fortunately, it can be shown that it will always happen when we select wisely.

    // Let inv_tx = |dx / wx|  (1 / time to travel one x grid)
    //     inv_ty = |dy / wy|  (1 / time to travel one y grid)
    //     inv_tz = |dz / wz|  (1 / time to travel one z grid)
    // If inv_tx is the largest of (inv_tx, inv_ty, inv_tz) -> we will intersect with x=c planes.
    // vise versa for inv_ty and inv_tz.
    // When we select inv_tx, it ensures we move in the x direction in 1 grid, y and z in <= 1 grid.

    auto inv_t = (ray_directions / cell_width.unsqueeze(-2)).abs();  // (b, m, 3)
    auto idx_to_use = torch::argmax(inv_t, -1);  // (b, m)  (0, 1, 2)

    // create accessors to each tensor
    // auto start = high_resolution_clock::now();
    auto a_grid_from = grid_from.accessor<float, 2>();  // (b, 3)
    auto a_cell_width = cell_width.accessor<float, 2>();  // (b, 3)
    auto a_grid_size = grid_size.accessor<int64_t, 2>();  // (b, 3)
    auto a_idx_to_use = idx_to_use.accessor<int64_t, 2>();  // (b, m)  0, 1, 2
    // auto stop = high_resolution_clock::now();
    // auto duration = duration_cast<microseconds>(stop - start);
    // std::cout<<"accessor_time: "<<duration.count()<<std::endl;

    auto l_options = torch::TensorOptions().dtype(torch::kLong).device(device);
    auto f_options = torch::TensorOptions().dtype(torch::kFloat32).device(device);


    std::vector<idxmap> all_grid_idxs;
    all_grid_idxs.reserve(batch_size);
    for (auto b = 0; b < batch_size; ++b) {

        auto a_grid_size_b = a_grid_size[b];

        // build a meshgrid for local grid idx.  This will determine the neighboring grids to gather.
        // the local grid width = 2 * local_grid_radius
        auto local_grid_radius = (torch::ceil(ray_radius.index({b, None}) / cell_width[b] + 1)).to(torch::kLong);  // (3,) long,  xyz
        auto grid_idx_from = -1 * local_grid_radius;  // (3,) long, xyz