def quantize(coords): r"""Returns a unique index map and an inverse index map. Args: :attr:`coords` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a matrix of size :math:`N \times D` where :math:`N` is the number of points in the :math:`D` dimensional space. Returns: :attr:`unique_map` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a list of indices that defines unique coordinates. :attr:`coords[unique_map]` is the unique coordinates. :attr:`inverse_map` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a list of indices that defines the inverse map that recovers the original coordinates. :attr:`coords[unique_map[inverse_map]] == coords` Example:: >>> unique_map, inverse_map = quantize(coords) >>> unique_coords = coords[unique_map] >>> print(unique_coords[inverse_map] == coords) # True, ..., True >>> print(coords[unique_map[inverse_map]] == coords) # True, ..., True """ assert isinstance(coords, np.ndarray) or isinstance( coords, torch.Tensor), "Invalid coords type" if isinstance(coords, np.ndarray): assert (coords.dtype == np.int32 ), f"Invalid coords type {coords.dtype} != np.int32" return MEB.quantize_np(coords.astype(np.int32)) else: # Type check done inside return MEB.quantize_th(coords.int())
def test_mapping(self): N = 16575 coords = (np.random.rand(N, 3) * 100).astype(np.int32) mapping, inverse_mapping = MEB.quantize_np(coords) print("N unique:", len(mapping), "N:", N) self.assertTrue((coords == coords[mapping][inverse_mapping]).all()) self.assertTrue((coords == coords[mapping[inverse_mapping]]).all()) coords = torch.from_numpy(coords) mapping, inverse_mapping = MEB.quantize_th(coords) print("N unique:", len(mapping), "N:", N) self.assertTrue((coords == coords[mapping[inverse_mapping]]).all()) unique_coords, index, reverse_index = sparse_quantize( coords, return_index=True, return_inverse=True) self.assertTrue((coords == coords[index[reverse_index]]).all())