From 6fcc15cd2935b4ac07c14c7e67a636772ad74916 Mon Sep 17 00:00:00 2001 From: clorenz Date: Fri, 16 Oct 2026 04:56:24 -0700 Subject: [PATCH] Add MaskRandom --- docs/source/api/processing.rst | 12 +++++ sdm/processing/__init__.py | 2 + sdm/processing/_random.py | 46 +++++++++++++++++++ sdm/processing/missing/__init__.py | 10 ++++ sdm/processing/missing/_cells.py | 22 +++++++++ sdm/processing/missing/random.py | 63 ++++++++++++++++++++++++++ test/processing/missing/test_random.py | 45 ++++++++++++++++++ test/processing/test_contract.py | 10 +++- 8 files changed, 208 insertions(+), 2 deletions(-) create mode 100644 sdm/processing/_random.py create mode 100644 sdm/processing/missing/__init__.py create mode 100644 sdm/processing/missing/_cells.py create mode 100644 sdm/processing/missing/random.py create mode 100644 test/processing/missing/test_random.py diff --git a/docs/source/api/processing.rst b/docs/source/api/processing.rst index 052634e7a0..880a960ba4 100644 --- a/docs/source/api/processing.rst +++ b/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 --------------- diff --git a/sdm/processing/__init__.py b/sdm/processing/__init__.py index 45a6d4a48f..1d20bcea77 100644 --- a/sdm/processing/__init__.py +++ b/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", diff --git a/sdm/processing/_random.py b/sdm/processing/_random.py new file mode 100644 index 0000000000..c6d73c08a2 --- /dev/null +++ b/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 + ) diff --git a/sdm/processing/missing/__init__.py b/sdm/processing/missing/__init__.py new file mode 100644 index 0000000000..29fdad201b --- /dev/null +++ b/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", +] diff --git a/sdm/processing/missing/_cells.py b/sdm/processing/missing/_cells.py new file mode 100644 index 0000000000..4a66c37b57 --- /dev/null +++ b/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, + ), + ) diff --git a/sdm/processing/missing/random.py b/sdm/processing/missing/random.py new file mode 100644 index 0000000000..dc1eab6aee --- /dev/null +++ b/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})" diff --git a/test/processing/missing/test_random.py b/test/processing/missing/test_random.py new file mode 100644 index 0000000000..b0e58158a2 --- /dev/null +++ b/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, + ) + ) diff --git a/test/processing/test_contract.py b/test/processing/test_contract.py index 7525b3b8c7..a7affecec9 100644 --- a/test/processing/test_contract.py +++ b/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(