Class TestIndexing

tests/pointersect/pr/cpp/test_cpp.py:18–60

class TestIndexing(unittest.TestCase):
    def test_sub2ind(self, b=10, n=20, max_size=20):
        size = torch.randint(max_size-1, size=[b, 3]) + 1 # (b, 3)
        total_cells = torch.prod(size, dim=-1)  # (b,)
        subidx = torch.rand(b, n, 3).floor().long()

        # python
        stime = timer()
        ind_gt = naive.sub2ind(subidx, size)
        time_python = timer() - stime

        # cpp
        stime = timer()
        ind_cpp = pr_cpp.sub2ind(subidx, size)
        time_cpp = timer() - stime

        print('sub2ind')
        print(f'python: {time_python * 1000.:.3f} ms')
        print(f'cpp: {time_cpp * 1000.:.3f} ms ({time_python/time_cpp:.2f} speed up)')

        assert torch.allclose(ind_gt, ind_cpp)


    def test_ind2sub(self, b=10, n=20, max_size=20):
        size = torch.randint(max_size-1, size=[b, 3]) + 1  # (b, 3)
        total_cells = torch.prod(size, dim=-1)  # (b,)
        ind = (torch.rand(b, n) * total_cells.unsqueeze(-1)).floor().long()

        # python
        stime = timer()
        out_gt = naive.ind2sub(ind, size)
        time_python = timer() - stime

        # cpp
        stime = timer()
        out_cpp = pr_cpp.ind2sub(ind, size)
        time_cpp = timer() - stime

        print('ind2sub')
        print(f'python: {time_python * 1000.:.3f} ms')
        print(f'cpp: {time_cpp * 1000.:.3f} ms ({time_python/time_cpp:.2f} speed up)')

        assert torch.allclose(out_gt, out_cpp)


class TestGatherPoints(unittest.TestCase):