Method test_chunk

tests/plib/test_utils.py:179–208

(self)

class Test_KNN(unittest.TestCase):

    def test_chunk(self):
        b_shape = [3, 5]
        n = 10
        m = 7
        k=6
        points = torch.randn(*b_shape, n, 3)
        ray_origins = torch.randn(*b_shape, m, 3)
        ray_directions = torch.zeros(*b_shape, m, 3)
        ray_directions[..., 2] = 1.

        # standard (no chunking)
        out_dict_gt = utils.get_k_neighbor_points(
            points=points,
            ray_origins=ray_origins,
            ray_directions=ray_directions,
            k=k,
        )

        # with chunking
        mn = m * n
        for max_chunk_size in [int(1e9), mn//2, mn+1]:
            out_dict = utils.get_k_neighbor_points_in_chunks(
                points=points,
                ray_origins=ray_origins,
                ray_directions=ray_directions,
                k=k,
                max_chunk_size=max_chunk_size,
            )
            for key in out_dict_gt:
                assert torch.allclose(out_dict_gt[key], out_dict[key]), f'{max_chunk_size}'