| 62 | |
| 63 | class TestGatherPoints(unittest.TestCase): |
| 64 | def test(self, b=10, n=20, max_size=20): |
| 65 | size = torch.randint(max_size-1, size=[b, 3]) + 1 # (b, 3) |
| 66 | total_cells = torch.prod(size, dim=-1) # (b,) |
| 67 | grid_idxs = (torch.rand(b, n) * total_cells.unsqueeze(1)).floor().long() |
| 68 | valid_mask = torch.randn(b, n) > 0.5 |
| 69 | |
| 70 | stime = timer() |
| 71 | all_cell2pidx_gt = naive.gather_points( |
| 72 | grid_idxs=grid_idxs, |
| 73 | total_cells=total_cells, |
| 74 | valid_mask=valid_mask, |
| 75 | ) |
| 76 | time_python = timer() - stime |
| 77 | |
| 78 | stime = timer() |
| 79 | _all_cell2pidx = pr_cpp.gather_points( |
| 80 | grid_idxs, |
| 81 | total_cells, |
| 82 | valid_mask, |
| 83 | ) |
| 84 | time_cpp = timer() - stime |
| 85 | |
| 86 | # convert list of idxmap to list of list of list |
| 87 | _all_cell2pidx = [c.get_table() for c in _all_cell2pidx] |
| 88 | all_cell2pidx = [] |
| 89 | for b in range(len(_all_cell2pidx)): |
| 90 | cell2pidx = [[] for i in range(total_cells[b])] |
| 91 | for gidx in _all_cell2pidx[b]: |
| 92 | cell2pidx[gidx] = list(_all_cell2pidx[b][gidx]) |
| 93 | all_cell2pidx.append(cell2pidx) |
| 94 | |
| 95 | print('gather points:') |
| 96 | print(f'python: {time_python * 1000.:.3f} ms') |
| 97 | print(f'cpp: {time_cpp * 1000.:.3f} ms ({time_python / time_cpp:.2f} speed up)') |
| 98 | |
| 99 | assert len(all_cell2pidx_gt) == len(all_cell2pidx) |
| 100 | for b in range(len(all_cell2pidx_gt)): |
| 101 | cell2pidx_gt = all_cell2pidx_gt[b] |
| 102 | cell2pidx = all_cell2pidx[b] |
| 103 | assert len(cell2pidx_gt) == len(cell2pidx) |
| 104 | for gidx in range(len(cell2pidx_gt)): |
| 105 | pidxs_gt = np.array(sorted(list(cell2pidx_gt[gidx]))) |
| 106 | pidxs = np.array(sorted(list(cell2pidx[gidx]))) |
| 107 | assert np.allclose(pidxs_gt, pidxs) |
| 108 | |
| 109 | |
| 110 | |