Add categorical value decoding

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

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"]),