| 212 | class TestGridRayIntersection(unittest.TestCase): |
| 213 | |
| 214 | def test(self, b=10, m=1000, max_size=40, seed=0): |
| 215 | |
| 216 | if seed is not None: |
| 217 | torch.manual_seed(seed) |
| 218 | random.seed(seed) |
| 219 | np.random.seed(seed) |
| 220 | |
| 221 | ray_origins = torch.randn(b, m, 3) # (b, m, 3) |
| 222 | ray_directions = torch.nn.functional.normalize(torch.randn(b, m, 3), dim=-1) # (b, m, 3) |
| 223 | ray_radius = torch.rand(b,)*0.8 + 1e-3 # (b,) |
| 224 | grid_size = (torch.rand(b, 3) * (max_size-1) + 1).floor().long() # (b, 3) |
| 225 | total_cells = torch.prod(grid_size, dim=-1) |
| 226 | grid_center = torch.randn(b, 3) |
| 227 | grid_width = torch.rand(b, 3) * 20 + 10 |
| 228 | |
| 229 | stime = timer() |
| 230 | all_grid_idxs_gt = naive.grid_ray_intersection( |
| 231 | ray_origins=ray_origins, |
| 232 | ray_directions=ray_directions, |
| 233 | ray_radius=ray_radius, |
| 234 | grid_size=grid_size, |
| 235 | grid_center=grid_center, |
| 236 | grid_width=grid_width, |
| 237 | ) |
| 238 | time_python = timer() - stime |
| 239 | |
| 240 | stime = timer() |
| 241 | _all_grid_idxs_cpp = pr_cpp.grid_ray_intersection_v2( |
| 242 | ray_origins, |
| 243 | ray_directions, |
| 244 | ray_radius, |
| 245 | grid_size, |
| 246 | grid_center, |
| 247 | grid_width, |
| 248 | ) |
| 249 | time_cpp = timer() - stime |
| 250 | |
| 251 | # convert list of idxmap to list of list of list |
| 252 | _all_grid_idxs_cpp = [c.get_table() for c in _all_grid_idxs_cpp] |
| 253 | all_grid_idxs_cpp = [] |
| 254 | for b in range(len(_all_grid_idxs_cpp)): |
| 255 | grid_idxs = [[] for i in range(m)] |
| 256 | for midx in _all_grid_idxs_cpp[b]: |
| 257 | grid_idxs[midx] = list(_all_grid_idxs_cpp[b][midx]) |
| 258 | all_grid_idxs_cpp.append(grid_idxs) |
| 259 | |
| 260 | ray_origins = ray_origins.cuda() |
| 261 | ray_directions = ray_directions.cuda() |
| 262 | ray_radius = ray_radius.cuda() |
| 263 | grid_size = grid_size.cuda() |
| 264 | grid_center = grid_center.cuda() |
| 265 | grid_width = grid_width.cuda() |
| 266 | |
| 267 | stime = timer() |
| 268 | _all_grid_idxs, n_all_grid_idx = pr_cuda.grid_ray_intersection( |
| 269 | ray_origins, |
| 270 | ray_directions, |
| 271 | ray_radius, |