Simplify Kumo transform schedule and shift dispatch

Author: ValterH

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),
@@ -31,32 +32,7 @@ def numerical_processor() -> sp.Sequential:
         features=[
             sp.StypeDispatch(
                 numerical=[
-                    sp.Cast(torch.float64),
-                    sp.DropConstantColumns(),
-                    # Period 12 preserves the original three-way schedule
-                    # while replacing every fourth view with Gaussian ranks.
-                    sp.Choice(
-                        *[
-                            [
-                                sp.RankGaussian(),
-                                sp.Standardize(),
-                                sp.ClipSigma(threshold=4.0),
-                            ]
-                            if i % 4 == 3
-                            else [
-                                sp.Standardize(eps=1e-6),
-                                sp.Clip(-100.0, 100.0),
-                                (
-                                    sp.Identity(),
-                                    sp.PowerTransform(),
-                                    [sp.RobustScale(), sp.ClipSoft(3.0)],
-                                )[i % 3],
-                                sp.ClipSigma(threshold=4.0),
-                            ]
-                            for i in range(12)
-                        ],
-                        method="round_robin",
-                    ),
+                    numerical_processor(),
                     sp.FlipSign(),
                 ],
                 categorical=[

sdm/processing/categorical/shuffle.py

@@ -63,25 +63,27 @@ def _draw_permutations(
             n_classes = category.numel()
             if n_classes <= 1:
                 permutation = torch.arange(n_classes, device=device)
-            elif self.method in ("shift", "balanced_shift"):
-                if 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:]
-                else:
-                    offset = torch.randint(
+            elif self.method == "shift":
+                offset = torch.randint(
+                    n_classes,
+                    (1,),
+                    generator=generator,
+                    device=device,
+                )
+                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,
-                        (1,),
                         generator=generator,
                         device=device,
                     )
+                offset = remaining[:1]
+                shifts[key] = remaining[1:]
                 permutation = (
                     torch.arange(n_classes, device=device) - offset
                 ) % n_classes

test/models/kumo/tabular/test_recipe.py

@@ -169,23 +169,17 @@ def test_default_recipe_numerical_schedule(
         EnsembleTable.from_table(query, num_members=num_members)
     )
     for i in range(num_members):
-        if i % 4 == 3:
-            processor = sp.Sequential(
+        processor = sp.Sequential(
+            sp.Standardize(eps=1e-6),
+            sp.Clip(-100.0, 100.0),
+            (
+                sp.Identity(),
+                sp.PowerTransform(),
+                [sp.RobustScale(), sp.ClipSoft(3.0)],
                 sp.RankGaussian(),
-                sp.Standardize(),
-                sp.ClipSigma(threshold=4.0),
-            )
-        else:
-            processor = sp.Sequential(
-                sp.Standardize(eps=1e-6),
-                sp.Clip(-100.0, 100.0),
-                (
-                    sp.Identity(),
-                    sp.PowerTransform(),
-                    [sp.RobustScale(), sp.ClipSoft(3.0)],
-                )[i % 3],
-                sp.ClipSigma(threshold=4.0),
-            )
+            )[i % 4],
+            sp.ClipSigma(threshold=4.0),
+        )
         expected = processor.fit_transform(features).numerical.float()
         # Independent sign flips are allowed; magnitudes identify each view.
         torch.testing.assert_close(