Function get_grid_idx

pointersect/pr/cpp/pr.cpp:239–274

Compute the grid index given xyz_w. Args: points: (*, n, 3) grid_size: (*, 3) long. number of grid cells in x y z. center: (*, 3) center of the grid grid_width: (*, 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: (*, n) bool Algorithm: Let x_

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

    assert(grid_size.dtype() == torch::kLong);
    
    auto grid_from = (center - grid_width / 2).unsqueeze(-2);  // (*, 1, 3)
    auto grid_to = (center + grid_width / 2).unsqueeze(-2);  // (*, 1, 3)
    auto cell_width = (grid_width / grid_size).unsqueeze(-2);  // (*, 1, 3)
    auto p_idx = ((points - grid_from) / cell_width).floor().to(torch::kLong);  // (*, n, 3), sub_idx on the grid

    // we will mark any out-of-bound points invalid
    auto valid_mask = torch::logical_and(
        points >= grid_from,
        points <= grid_to
    ).all(-1);  // (*, n),  bool

    if (mode.compare("ind") == 0) {
        auto grid_idx = sub2ind(
            p_idx,
            grid_size
        );  // (*, n) long
        return {grid_idx, valid_mask};
    }
    else if (mode.compare("subidx") == 0) {
        return {p_idx, valid_mask};
    }
    else {
        throw;
    }

}