diff --git a/sdm/tensor/categorical.py b/sdm/tensor/categorical.py index 2d21727627..2f7e56d249 100644 --- a/sdm/tensor/categorical.py +++ b/sdm/tensor/categorical.py @@ -14,7 +14,7 @@ from torch.utils import _pytree as pytree from typing_extensions import override -from sdm.tensor import StringTensor, VarLenTensor +from sdm.tensor import NullableTensor, StringTensor, VarLenTensor from sdm.tensor.io import ( arrow_as_tensor, combine_arrow_chunks, @@ -350,6 +350,43 @@ def categories(self) -> tuple[Tensor, ...]: r"""Return category vector for each categorical column.""" return self._categories + # Methods ################################################################# + + def decode(self) -> Tensor: + r"""Decode this tensor to its category values.""" + from sdm.tensor import ColumnarTensor # noqa: PLC0415 + + if self.size(-1) == 0: + raise ValueError("Decoding requires at least one column") + + tensors = [] + for tensor in self.split(1, dim=-1): + code = tensor.code.squeeze(-1) + valid = code >= 0 + if not valid.all(): + tensor = tensor.categories[0][code.clamp(min=0)] + if isinstance(tensor, VarLenTensor): + data, offset = tensor.contiguous().data_offset + tensor = tensor.__class__( + data=data, + offset=offset, + valid=valid.reshape(-1), + size=tensor.size(), + ) + elif tensor.is_floating_point(): + tensor[~valid] = torch.nan + else: + tensor = NullableTensor(tensor, valid) + else: + tensor = tensor.categories[0][code] + + tensors.append(tensor) + + if len(tensors) == 1: + return tensors[0].unsqueeze(-1) + + return ColumnarTensor(tensors) + # Decorators ############################################################## @classmethod diff --git a/test/tensor/test_categorical.py b/test/tensor/test_categorical.py index 8596301dab..7709fec21c 100644 --- a/test/tensor/test_categorical.py +++ b/test/tensor/test_categorical.py @@ -8,7 +8,7 @@ import pytest import torch -from sdm import CategoricalTensor, StringTensor +from sdm import CategoricalTensor, ColumnarTensor, NullableTensor, StringTensor from sdm.testing import onlyCUDA, withCUDA @@ -74,6 +74,34 @@ def test_to_copy() -> None: assert out.dtype == torch.float32 +@withCUDA +def test_decode(device: torch.device) -> None: + tensor = CategoricalTensor( + code=torch.tensor([[0, 1, 0], [-1, -1, -1], [1, 0, 1]], device=device), + categories=( + StringTensor.from_list(["a", "b"], device=device), + torch.tensor([10, 20], dtype=torch.int64, device=device), + torch.tensor([1.5, 2.5], device=device), + ), + ) + + out = tensor.decode() + + assert isinstance(out, ColumnarTensor) + assert out.shape == (3, 3) + assert out.device == device + assert isinstance(out[:, 0], StringTensor) + assert out[:, 0].tolist() == ["a", None, "b"] + assert isinstance(out[:, 1], NullableTensor) + assert out[:, 1].dtype == torch.int64 + assert out[:, 1].tolist() == [20, None, 10] + torch.testing.assert_close( + out[:, 2], + torch.tensor([1.5, torch.nan, 2.5], device=device), + equal_nan=True, + ) + + def test_from_arrow_string_values() -> None: tensor = CategoricalTensor.from_arrow( pa.array(["b", "a", None, "b"]),