Add MaskRandom

docs/source/api/processing.rst

@@ -51,6 +51,18 @@ Categorical Processors
    ~sdm.processing.categorical.{{ name }}
 {% endfor %}
 
+Missingness Processors
+----------------------
+
+.. autosummary::
+   :toctree: generated
+   :template: processor_class
+   :nosignatures:
+
+{% for name in api_names("sdm.processing.missing") %}
+   ~sdm.processing.missing.{{ name }}
+{% endfor %}
+
 Text Processors
 ---------------
 

sdm/processing/__init__.py

@@ -49,6 +49,7 @@
     ImputeMode,
     AddCategoryCounts,
 )
+from sdm.processing.missing import MaskRandom
 from sdm.processing.datetime import AddCalendarFields
 from sdm.processing.output import AverageEstimators, Softmax, SortQuantiles
 from sdm.processing.recipe import Recipe
@@ -94,6 +95,7 @@
     "ShuffleCategories",
     "ImputeMode",
     "AddCategoryCounts",
+    "MaskRandom",
     "AddCalendarFields",
     "AverageEstimators",
     "Softmax",

sdm/processing/_random.py

@@ -0,0 +1,46 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Sequence
+from typing import cast
+
+import torch
+from torch import Tensor
+
+
+class SeedMixin:
+    r"""Store a seed at fit, from which transform draws its row values."""
+
+    _seed = 0
+
+    def get_extra_state(self) -> object:
+        r""":meta private:"""  # noqa: D415
+        return self._seed
+
+    def set_extra_state(self, state: object) -> None:
+        r""":meta private:"""  # noqa: D415
+        self._seed = cast(int, state)
+
+    def _draw_seed(self, generator: torch.Generator | None) -> None:
+        device = None if generator is None else generator.device
+        self._seed = int(
+            torch.randint(2**62, (), device=device, generator=generator)
+        )
+
+    def _transform_generator(self) -> torch.Generator:
+        # A CPU generator makes the output independent of the device.
+        return torch.Generator().manual_seed(self._seed)
+
+
+def uniform(
+    low: float,
+    high: float,
+    size: Sequence[int],
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+    generator: torch.Generator | None,
+) -> Tensor:
+    return torch.empty(size, device=device, dtype=dtype).uniform_(
+        low, high, generator=generator
+    )

sdm/processing/missing/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Processors that mask numerical and categorical cells."""
+
+from sdm.processing.missing.random import MaskRandom
+
+__all__ = [
+    "MaskRandom",
+]

sdm/processing/missing/_cells.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch import Tensor
+
+from sdm import CategoricalTensor, TableTensor
+
+
+def mask_cells(table: TableTensor, missing: Tensor) -> TableTensor:
+    r"""Mask numerical and categorical cells where ``missing`` is true."""
+    num_numerical = table.numerical.size(-1)
+    code = table.categorical.code.masked_fill(missing[..., num_numerical:], -1)
+    return table.replace_blocks(
+        numerical=table.numerical.masked_fill(
+            missing[..., :num_numerical], torch.nan
+        ),
+        categorical=CategoricalTensor(
+            code=code,
+            categories=table.categorical.categories,
+        ),
+    )

sdm/processing/missing/random.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch import Tensor
+
+from sdm import Stype, TableTensor
+from sdm.processing import Processor
+from sdm.processing._random import SeedMixin, uniform
+from sdm.processing.missing._cells import mask_cells
+
+
+class MaskRandom(SeedMixin, Processor):
+    """Mask numerical and categorical cells completely at random.
+
+    Each table draws a missing-cell probability when the processor is fitted.
+    Transform masks each cell independently with this probability. Apply the
+    processor to all rows in one call.
+
+    Args:
+        rate: Range from which to draw each table's missing-cell probability.
+    """
+
+    handles_stypes = (Stype.numerical, Stype.categorical)
+    requires_fit = True
+
+    _missing_rate: Tensor
+
+    def __init__(self, *, rate: tuple[float, float] = (0.01, 0.08)) -> None:
+        super().__init__()
+        self.rate = rate
+        self.register_buffer("_missing_rate", torch.empty(0))
+
+    def _fit(
+        self,
+        table: TableTensor,
+        *,
+        generator: torch.Generator | None = None,
+    ) -> None:
+        self._draw_seed(generator)
+        self._set_buffer(
+            "_missing_rate",
+            uniform(
+                low=self.rate[0],
+                high=self.rate[1],
+                size=table.size()[:-2],
+                device=table.device,
+                dtype=table.numerical.dtype,
+                generator=generator,
+            ),
+        )
+
+    def _transform(self, table: TableTensor) -> TableTensor:
+        num_columns = table.numerical.size(-1) + table.categorical.size(-1)
+        u = torch.rand(
+            (*table.size()[:-1], num_columns),
+            generator=self._transform_generator(),
+        )
+        missing = u.to(table.device) < self._missing_rate[..., None, None]
+        return mask_cells(table, missing)
+
+    def __repr__(self, *, indent: int = 0) -> str:
+        return f"{' ' * indent}{self.__class__.__name__}(rate={self.rate})"

test/processing/missing/test_random.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+
+from sdm import CategoricalTensor, TableTensor
+from sdm.processing import MaskRandom
+
+
+def test_masks_cells_at_the_configured_rate() -> None:
+    rate = 0.05
+    table = TableTensor(
+        numerical=torch.randn(4, 4000, 3),
+        categorical=CategoricalTensor(
+            code=torch.randint(3, (4, 4000, 2)),
+            categories=(torch.arange(3),) * 2,
+        ),
+    )
+
+    output = MaskRandom(rate=(rate, rate)).fit_transform(table)
+
+    numerical_missing = output.numerical.isnan()
+    categorical_missing = output.categorical.code == -1
+    missing = torch.cat((numerical_missing, categorical_missing), dim=-1)
+    actual_rate = missing.float().mean(dim=(-2, -1))
+    torch.testing.assert_close(
+        actual_rate,
+        torch.full_like(actual_rate, rate),
+        rtol=0.0,
+        atol=0.01,
+    )
+    assert output.numerical[~numerical_missing].equal(
+        table.numerical[~numerical_missing]
+    )
+    assert output.categorical.code[~categorical_missing].equal(
+        table.categorical.code[~categorical_missing]
+    )
+    assert all(
+        actual.equal(expected)
+        for actual, expected in zip(
+            output.categorical.categories,
+            table.categorical.categories,
+            strict=True,
+        )
+    )

test/processing/test_contract.py

@@ -138,6 +138,7 @@ def _make_processor_pair(
     ProcessorCase(sp.ShuffleCategories()),
     ProcessorCase(sp.ImputeMode()),
     ProcessorCase(sp.AddCategoryCounts()),
+    ProcessorCase(sp.MaskRandom(rate=(0.5, 0.5))),
     ProcessorCase(sp.AddCalendarFields(["month"])),
     ProcessorCase(sp.Softmax()),
     ProcessorCase(sp.SortQuantiles()),
@@ -186,6 +187,11 @@ def _make_processor_pair(
 )
 
 
+def _assert_equal(actual: TableTensor, expected: TableTensor) -> None:
+    # Processors can emit NaN for missing values.
+    assert actual.allclose(expected, rtol=0.0, atol=0.0, equal_nan=True)
+
+
 def test_all_public_processors_have_contract_cases() -> None:
     public_processors = set()
     for name in sp.__all__:
@@ -230,7 +236,7 @@ def test_fit_transform_matches_fit_then_transform(
     )
     assert len(actual) == len(expected)
     for member_id in range(len(actual)):
-        assert actual[member_id].equal(expected[member_id])
+        _assert_equal(actual[member_id], expected[member_id])
 
 
 @pytest.mark.parametrize(
@@ -249,7 +255,7 @@ def test_save_and_load_preserves_fitted_processor_behavior(
     actual = processor_b.transform_ensemble(table)
     assert len(actual) == len(expected)
     for member_id in range(len(actual)):
-        assert actual[member_id].equal(expected[member_id])
+        _assert_equal(actual[member_id], expected[member_id])
 
 
 @pytest.mark.parametrize(