Increase ensemble member diversity for Kumo Tabular

sdm/models/kumo/tabular/recipe.py

@@ -22,6 +22,7 @@ def numerical_processor() -> sp.Sequential:
                     sp.RobustScale(),
                     sp.ClipSoft(3.0),
                 ],
+                sp.RankGaussian(),
                 method="round_robin",
             ),
             sp.ClipSigma(threshold=4.0),
@@ -49,7 +50,7 @@ def numerical_processor() -> sp.Sequential:
             sp.StypeDispatch(
                 categorical=[
                     sp.AlignCategories(),
-                    sp.ShuffleCategories(method="shift"),
+                    sp.ShuffleCategories(method="balanced_shift"),
                 ],
                 numerical=[
                     sp.Standardize(),

sdm/processing/__init__.py

@@ -32,6 +32,7 @@
     ImputeMean,
     PowerTransform,
     QuantileTransform,
+    RankGaussian,
     Standardize,
     RobustScale,
     FlipSign,
@@ -76,6 +77,7 @@
     "ImputeMean",
     "PowerTransform",
     "QuantileTransform",
+    "RankGaussian",
     "Standardize",
     "RobustScale",
     "FlipSign",

sdm/processing/categorical/shuffle.py

@@ -12,7 +12,7 @@
 
 
 class ShuffleCategories(EnsembleProcessor):
-    """Independently permute the integer codes of categorical columns.
+    """Permute the integer codes of categorical columns.
 
     One permutation per categorical column is drawn when the processor is
     fitted. Codes and their corresponding category vectors are permuted
@@ -24,15 +24,18 @@ class ShuffleCategories(EnsembleProcessor):
     Args:
         method: Permutation strategy. ``"shift"`` cyclically shifts the
             codes by a drawn offset, and ``"random"`` remaps the codes with
-            a drawn permutation.
+            a drawn permutation. ``"balanced_shift"`` draws cyclic offsets
+            without replacement before starting another cycle, balancing
+            their counts across ensemble members for each column and class
+            count.
     """
 
     handles_stypes = frozenset({Stype.categorical})
     requires_fit = True
 
     def __init__(
         self,
-        method: Literal["shift", "random"] = "random",
+        method: Literal["shift", "random", "balanced_shift"] = "random",
     ) -> None:
         super().__init__()
         self.method = method
@@ -51,11 +54,12 @@ def _draw_permutations(
         self,
         table: TableTensor,
         *,
+        shifts: dict[tuple[torch.device, int, int], Tensor],
         generator: torch.Generator | None = None,
     ) -> list[Tensor]:
         device = table.categorical.device
         permutations: list[Tensor] = []
-        for category in table.categorical.categories:
+        for column, category in enumerate(table.categorical.categories):
             n_classes = category.numel()
             if n_classes <= 1:
                 permutation = torch.arange(n_classes, device=device)
@@ -69,6 +73,20 @@ def _draw_permutations(
                 permutation = (
                     torch.arange(n_classes, device=device) - offset
                 ) % n_classes
+            elif self.method == "balanced_shift":
+                key = (device, column, n_classes)
+                remaining = shifts.get(key)
+                if remaining is None or remaining.numel() == 0:
+                    remaining = torch.randperm(
+                        n_classes,
+                        generator=generator,
+                        device=device,
+                    )
+                offset = remaining[:1]
+                shifts[key] = remaining[1:]
+                permutation = (
+                    torch.arange(n_classes, device=device) - offset
+                ) % n_classes
             else:
                 assert self.method == "random"
                 permutation = torch.randperm(
@@ -92,9 +110,11 @@ def _fit_ensemble(
             tuple[torch.device, tuple[tuple[int, ...], ...]], int
         ] = {}
 
+        shifts: dict[tuple[torch.device, int, int], Tensor] = {}
         for member_id in range(len(ensemble_table)):
             permutations = self._draw_permutations(
                 ensemble_table[member_id],
+                shifts=shifts,
                 generator=generator,
             )
             key = (

sdm/processing/numerical/__init__.py

@@ -11,6 +11,7 @@
 from sdm.processing.numerical.impute import ImputeMean
 from sdm.processing.numerical.power import PowerTransform
 from sdm.processing.numerical.quantile import QuantileTransform
+from sdm.processing.numerical.rank_gaussian import RankGaussian
 from sdm.processing.numerical.standardize import Standardize
 from sdm.processing.numerical.robust_scale import RobustScale
 from sdm.processing.numerical.flip_sign import FlipSign
@@ -27,6 +28,7 @@
     "ImputeMean",
     "PowerTransform",
     "QuantileTransform",
+    "RankGaussian",
     "Standardize",
     "RobustScale",
     "FlipSign",

sdm/processing/numerical/rank_gaussian.py

@@ -0,0 +1,72 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from sdm import Stype, TableTensor
+from sdm.processing import Processor
+from sdm.processing.numerical.quantile import _batched_interp
+
+
+class RankGaussian(Processor):
+    """Map interpolated empirical mid-ranks to standard normal quantiles.
+
+    Each fitted value receives probability ``(L + R) / (2 * N)``, where
+    ``L`` and ``R`` count finite fitted values strictly below and at or below
+    it, and ``N`` is the number of finite fitted values. Ties share a rank.
+    Query probabilities interpolate between fitted values and clamp to the
+    endpoint probabilities, keeping finite outputs even outside the range.
+
+    NaN and infinite values are ignored during fitting. NaNs are preserved
+    during transformation. Constant columns map to zero; columns without
+    finite fitted values produce NaNs. All fitted values are retained.
+    """
+
+    handles_stypes = frozenset({Stype.numerical})
+    requires_fit = True
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.register_buffer("_values", torch.empty(0))
+        self.register_buffer("_probabilities", torch.empty(0))
+
+    def _fit(
+        self,
+        table: TableTensor,
+        *,
+        generator: torch.Generator | None = None,
+    ) -> None:
+        # [..., N, C] -> [..., C, N]; double precision keeps tail ranks open.
+        numerical = table.numerical.double().movedim(-1, -2)
+        values = numerical.masked_fill(~numerical.isfinite(), torch.inf)
+        values = values.sort(dim=-1).values.contiguous()
+        finite = values.isfinite()
+        count = finite.sum(dim=-1, keepdim=True)
+        left = torch.searchsorted(values, values, right=False)
+        right = torch.searchsorted(values, values, right=True)
+        probabilities = (left + right).double() / (2 * count.clamp_min(1))
+
+        # Pad missing observations with the last finite knot and its rank.
+        last = (count - 1).clamp_min(0)
+        self._values = torch.where(finite, values, values.gather(-1, last))
+        self._probabilities = torch.where(
+            finite, probabilities, probabilities.gather(-1, last)
+        )
+        self._values.masked_fill_(count == 0, torch.nan)
+        self._probabilities.masked_fill_(count == 0, torch.nan)
+
+    def _transform(self, table: TableTensor) -> TableTensor:
+        numerical = table.numerical
+        columns = numerical.movedim(-1, -2)
+        n_rows = columns.size(-1)
+        n_fitted = self._values.size(-1)
+        probabilities = _batched_interp(
+            columns.reshape(-1, n_rows).to(self._values.dtype).contiguous(),
+            self._values.reshape(-1, n_fitted),
+            self._probabilities.reshape(-1, n_fitted),
+        )
+        output = torch.special.ndtri(probabilities).reshape(columns.shape)
+        output = output.masked_fill(columns.isnan(), torch.nan)
+        return table.replace_blocks(
+            numerical=output.movedim(-1, -2).to(numerical.dtype)
+        )

test/processing/categorical/test_shuffle.py

@@ -84,9 +84,9 @@ def test_shuffle_categories_random_permutes_each_categorical_column(
     assert transformed.categorical.tolist() == features.categorical.tolist()
 
 
-@pytest.mark.parametrize("method", ["shift", "random"])
+@pytest.mark.parametrize("method", ["shift", "random", "balanced_shift"])
 def test_shuffle_categories_is_reproducible_with_generator(
-    method: Literal["shift", "random"],
+    method: Literal["shift", "random", "balanced_shift"],
 ) -> None:
     features = _table(
         [[0, 0], [1, 1], [2, -1], [1, 0]],
@@ -201,3 +201,49 @@ def test_shuffle_categories_refit_replaces_ensemble_state() -> None:
     )
 
     assert output.equal(expected)
+
+
+@withCUDA
+@pytest.mark.parametrize("num_classes", [1, 2, 3, 10])
+@pytest.mark.parametrize("num_members", [1, 8, 21])
+def test_balanced_shifts_cover_classes(
+    device: torch.device,
+    num_classes: int,
+    num_members: int,
+) -> None:
+    target = _table(
+        [[i] for i in range(num_classes)] + [[-1]],
+        (tuple(str(i) for i in range(num_classes)),),
+        device=device,
+    )
+    ensemble = EnsembleTable.from_table(target, num_members=num_members)
+    processor = ShuffleCategories(method="balanced_shift")
+    output = processor.fit_transform_ensemble(ensemble)
+    shifts = torch.stack(
+        [output[i].categorical.code[0, 0] for i in range(num_members)]
+    ).long()
+    counts = shifts.bincount(minlength=num_classes)
+    assert int(counts.max() - counts.min()) <= 1
+    for start in range(0, num_members, num_classes):
+        cycle = shifts[start : start + num_classes]
+        assert cycle.unique().numel() == cycle.numel()
+    query_output = processor.transform_ensemble(ensemble)
+    for i in range(num_members):
+        assert output[i].categorical.tolist() == target.categorical.tolist()
+        assert output[i].equal(query_output[i])
+
+
+def test_balanced_shifts_reproducible_ensemble() -> None:
+    target = _table([[0], [1], [2], [-1]], (("a", "b", "c"),))
+    ensemble = EnsembleTable.from_table(target, num_members=8)
+    processor = ShuffleCategories(method="balanced_shift")
+    first = processor.fit_transform_ensemble(
+        ensemble,
+        generator=torch.Generator().manual_seed(7),
+    )
+    second = processor.fit_transform_ensemble(
+        ensemble,
+        generator=torch.Generator().manual_seed(7),
+    )
+    for i in range(8):
+        assert first[i].equal(second[i])

test/processing/numerical/test_rank_gaussian.py

@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+import torch
+
+from sdm import TableTensor
+from sdm.processing import RankGaussian
+from sdm.testing import withCUDA
+
+
+@withCUDA
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_mid_ranks_and_interpolated_query(
+    dtype: torch.dtype, device: torch.device
+) -> None:
+    context = torch.tensor(
+        [[0.0], [1.0], [1.0], [2.0]], dtype=dtype, device=device
+    )
+    processor = RankGaussian().fit(TableTensor.from_tensor(context))
+    probabilities = context.new_tensor([[0.125], [0.5], [0.5], [0.875]])
+    torch.testing.assert_close(
+        processor.transform(TableTensor.from_tensor(context)).numerical,
+        torch.special.ndtri(probabilities),
+    )
+    query = context.new_tensor([[-100.0], [0.5], [1.5], [100.0], [torch.nan]])
+    expected = context.new_tensor(
+        [[0.125], [0.3125], [0.6875], [0.875], [torch.nan]]
+    )
+    torch.testing.assert_close(
+        processor.transform(TableTensor.from_tensor(query)).numerical,
+        torch.special.ndtri(expected),
+        equal_nan=True,
+    )
+
+
+@withCUDA
+def test_ties_at_endpoints_use_mid_ranks(device: torch.device) -> None:
+    context = torch.tensor([[0.0], [0.0], [2.0], [2.0]], device=device)
+    output = RankGaussian().fit_transform(TableTensor.from_tensor(context))
+    probabilities = context.new_tensor([[0.25], [0.25], [0.75], [0.75]])
+    torch.testing.assert_close(
+        output.numerical, torch.special.ndtri(probabilities)
+    )
+
+
+@withCUDA
+@pytest.mark.parametrize("n_rows", [1, 4])
+def test_constants_and_missing_columns(
+    n_rows: int, device: torch.device
+) -> None:
+    context = torch.tensor([[7.0, torch.nan]], device=device).expand(
+        n_rows, -1
+    )
+    processor = RankGaussian().fit(TableTensor.from_tensor(context))
+    query = context.new_tensor(
+        [[7.0, 8.0], [torch.nan, torch.nan], [100.0, 3.0]]
+    )
+    expected = context.new_tensor(
+        [[0.0, torch.nan], [torch.nan, torch.nan], [0.0, torch.nan]]
+    )
+    torch.testing.assert_close(
+        processor.transform(TableTensor.from_tensor(query)).numerical,
+        expected,
+        equal_nan=True,
+    )
+
+
+@withCUDA
+def test_nonfinite_context_does_not_change_observed_ranks(
+    device: torch.device,
+) -> None:
+    context = torch.tensor(
+        [[0.0], [1.0], [torch.nan], [1.0], [torch.inf], [-torch.inf], [2.0]],
+        device=device,
+    )
+    processor = RankGaussian().fit(TableTensor.from_tensor(context))
+    reference = RankGaussian().fit(
+        TableTensor.from_tensor(context[[0, 1, 3, 6]])
+    )
+    query = TableTensor.from_tensor(context)
+    output = processor.transform(query).numerical
+    torch.testing.assert_close(
+        output, reference.transform(query).numerical, equal_nan=True
+    )
+    assert torch.equal(output.isnan(), context.isnan())
+    assert output[~context.isnan()].isfinite().all()
+
+
+@withCUDA
+def test_query_batch_does_not_change_fitted_ranks(
+    device: torch.device,
+) -> None:
+    context = torch.arange(10.0, device=device).unsqueeze(-1)
+    processor = RankGaussian().fit(TableTensor.from_tensor(context))
+    query = context.new_tensor([[1.5], [torch.nan], [5.5]])
+    expected = processor.transform(TableTensor.from_tensor(query)).numerical
+    extended = torch.cat([query, context.new_tensor([[-1e12], [1e12]])])
+    actual = processor.transform(TableTensor.from_tensor(extended)).numerical
+    torch.testing.assert_close(actual[:3], expected, equal_nan=True)
+
+
+@withCUDA
+@pytest.mark.parametrize("batch_shape", [(), (2,), (2, 3)])
+def test_batched_missing_values_match_independent_columns(
+    batch_shape: tuple[int, ...], device: torch.device
+) -> None:
+    context = torch.randn(*batch_shape, 17, 4, device=device)
+    context[..., 1, 0] = torch.nan
+    context[..., :3, 1] = torch.nan
+    context[..., :, 2] = 2.0
+    context[..., :, 3] = torch.nan
+    query = torch.randn(*batch_shape, 7, 4, device=device)
+    query[..., 0, 0] = torch.nan
+    processor = RankGaussian().fit(TableTensor.from_tensor(context))
+    output = processor.transform(TableTensor.from_tensor(query)).numerical
+    for train, test, actual in zip(
+        context.reshape(-1, 17, 4),
+        query.reshape(-1, 7, 4),
+        output.reshape(-1, 7, 4),
+        strict=True,
+    ):
+        for column in range(4):
+            reference = RankGaussian().fit(
+                TableTensor.from_tensor(train[:, column : column + 1])
+            )
+            expected = reference.transform(
+                TableTensor.from_tensor(test[:, column : column + 1])
+            ).numerical
+            torch.testing.assert_close(
+                actual[:, column : column + 1], expected, equal_nan=True
+            )