| 7 | from timeit import default_timer as timer |
| 8 | |
| 9 | class TestGetGridIdx(unittest.TestCase): |
| 10 | def _test( |
| 11 | self, |
| 12 | b: int, |
| 13 | n: int, |
| 14 | grid_size, |
| 15 | center, |
| 16 | grid_width, |
| 17 | ): |
| 18 | points = torch.rand(b, n, 3) |
| 19 | gidx, valid_mask = naive.get_grid_idx( |
| 20 | points=points, |
| 21 | grid_size=grid_size, |
| 22 | center=center, |
| 23 | grid_width=grid_width, |
| 24 | ) |
| 25 | if isinstance(grid_size, int): |
| 26 | total_cells = grid_size * grid_size * grid_size |
| 27 | else: |
| 28 | total_cells = torch.prod(grid_size, dim=-1, keepdim=True) # (*, 1) |
| 29 | torch.prod(gidx) |
| 30 | assert torch.logical_or(gidx >= 0, torch.logical_not(valid_mask)).all() |
| 31 | assert torch.logical_or(gidx < total_cells, torch.logical_not(valid_mask)).all() |
| 32 | |
| 33 | |
| 34 | def test1(self): |
| 35 | self._test( |
| 36 | b=3, |
| 37 | n=10, |
| 38 | grid_size=5, |
| 39 | center=0., |
| 40 | grid_width=1., |
| 41 | ) |
| 42 | |
| 43 | def test2(self): |
| 44 | b = 3 |
| 45 | self._test( |
| 46 | b=b, |
| 47 | n=10, |
| 48 | grid_size=(torch.rand(b, 3) * 10).long() + 1, |
| 49 | center=torch.randn(b, 3), |
| 50 | grid_width=torch.rand(b, 3), |
| 51 | ) |
| 52 | |
| 53 | |
| 54 |
nothing calls this directly
no outgoing calls
no test coverage detected