Function gather_points

pointersect/pr/cpp/pr.cpp:160–197

Gather the points belonging to each grid cell. Args: grid_idxs: (b, n), grid index of each point total_cells: (b,) total grid cells valid_mask: (b, n), whether to record the point Returns: vector of idxmap, b -> (total_cells[bidx] -> n_idx of points in the cell)

//     vector of idxmap, b -> (total_cells[bidx] -> n_idx of points in the cell)
// 
std::vector<idxmap> gather_points(
    // input
    const torch::Tensor & grid_idxs,   // (b, n), int64_t
    const torch::Tensor & total_cells, // (b,) int64_t
    const torch::Tensor & valid_mask   // (b, n), bool
    // // output
    // std::vector<idxmap> & all_cell2pidx  // b -> gidx -> pidx
) {

    assert(grid_idxs.dtype() == torch::kLong);
    assert(total_cells.dtype() == torch::kLong);
    assert(valid_mask.dtype() == torch::kBool);

    auto batch_size = grid_idxs.size(0);
    auto n_points = grid_idxs.size(1);

    // accessor 
    auto a_grid_idxs = grid_idxs.accessor<int64_t, 2>();
    auto a_total_cells = total_cells.accessor<int64_t, 1>();
    auto a_valid_mask = valid_mask.accessor<bool, 2>();

    std::vector<idxmap> all_cell2pidx;
    all_cell2pidx.reserve(batch_size);

    // loop through each batch element
    for (auto b = 0; b < batch_size; ++b) {    
        all_cell2pidx.push_back(idxmap(a_total_cells[b]));
        idxmap& cell2pidx = all_cell2pidx[b];

        // loop through each point
        for (auto i = 0; i < n_points; ++i) {
            if (a_valid_mask[b][i]) {
                cell2pidx.insert(a_grid_idxs[b][i], i);
            }
        }
    }
    return all_cell2pidx;
}