diff --git a/sdm/models/kumo/tabular/recipe.py b/sdm/models/kumo/tabular/recipe.py index 38173f3b4..9a22d4c59 100644 --- a/sdm/models/kumo/tabular/recipe.py +++ b/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(), diff --git a/sdm/processing/__init__.py b/sdm/processing/__init__.py index 400a11d9a..ad88657d6 100644 --- a/sdm/processing/__init__.py +++ b/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", diff --git a/sdm/processing/categorical/shuffle.py b/sdm/processing/categorical/shuffle.py index 7b160c46e..cc228abf1 100644 --- a/sdm/processing/categorical/shuffle.py +++ b/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,7 +24,10 @@ 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}) @@ -32,7 +35,7 @@ class ShuffleCategories(EnsembleProcessor): 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 = ( diff --git a/sdm/processing/numerical/__init__.py b/sdm/processing/numerical/__init__.py index 6f670deb4..fbbba0ec8 100644 --- a/sdm/processing/numerical/__init__.py +++ b/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", diff --git a/sdm/processing/numerical/rank_gaussian.py b/sdm/processing/numerical/rank_gaussian.py new file mode 100644 index 000000000..34bbb4ce4 --- /dev/null +++ b/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) + ) diff --git a/test/processing/categorical/test_shuffle.py b/test/processing/categorical/test_shuffle.py index 392e428bc..6757830f5 100644 --- a/test/processing/categorical/test_shuffle.py +++ b/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]) diff --git a/test/processing/numerical/test_rank_gaussian.py b/test/processing/numerical/test_rank_gaussian.py new file mode 100644 index 000000000..f71fafcb3 --- /dev/null +++ b/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 + ) diff --git a/test/processing/test_contract.py b/test/processing/test_contract.py index 43692e47b..3646dbcf3 100644 --- a/test/processing/test_contract.py +++ b/test/processing/test_contract.py @@ -126,6 +126,7 @@ def _make_processor_pair( sp.QuantileTransform(n_quantiles=4, subsample=None), ), ProcessorCase(sp.Standardize()), + ProcessorCase(sp.RankGaussian()), ProcessorCase(sp.RobustScale()), ProcessorCase(sp.FlipSign()), ProcessorCase(sp.DropConstantColumns()), @@ -133,6 +134,7 @@ def _make_processor_pair( ProcessorCase(sp.RandomProjection(2)), ProcessorCase(sp.AlignCategories(), _make_align_categories_table), ProcessorCase(sp.ShuffleCategories()), + ProcessorCase(sp.ShuffleCategories(method="balanced_shift")), ProcessorCase(sp.ImputeMode()), ProcessorCase(sp.AddCategoryCounts()), ProcessorCase(sp.AddCalendarFields(["month"])),