Treat low-cardinality numeric features as categorical

Author: bogdan-kulynych

risksyn/generator.py

@@ -12,6 +12,35 @@
 # Empirically tested default from dpmm library
 _DEFAULT_PROC_EPSILON = 0.1
 
+# Numeric columns with at most this many unique values are auto-treated as categorical
+_AUTO_CATEGORICAL_THRESHOLD = 10
+
+
+def _auto_categorize(
+    data: pd.DataFrame, domain: Optional[dict]
+) -> tuple[pd.DataFrame, dict]:
+    """Auto-detect low-cardinality numeric columns and treat as categorical.
+
+    Numeric columns with <= _AUTO_CATEGORICAL_THRESHOLD unique values are
+    cast to string dtype and given categorical domains to avoid private
+    bounds estimation.
+
+    Returns a (possibly modified) copy of data and the augmented domain.
+    """
+    domain = dict(domain) if domain else {}
+    cols_to_cast = []
+    for col, series in data.items():
+        if col in domain:
+            continue
+        if series.dtype.kind in "ui":  # uint, int only
+            if series.nunique() <= _AUTO_CATEGORICAL_THRESHOLD:
+                domain[col] = sorted(str(v) for v in series.unique())
+                cols_to_cast.append(col)
+    if cols_to_cast:
+        data = data.copy()
+        data[cols_to_cast] = data[cols_to_cast].astype(str)
+    return data, domain
+
 
 def _requires_private_preprocessing(data: pd.DataFrame, domain: Optional[dict]) -> bool:
     """Check if any numeric column lacks bounds in domain.
@@ -24,8 +53,10 @@ def _requires_private_preprocessing(data: pd.DataFrame, domain: Optional[dict])
             if domain is None:
                 return True
             col_domain = domain.get(col, {})
+            if isinstance(col_domain, list):
+                continue  # categorical domain provided, no bounds needed
             if not isinstance(col_domain, dict):
-                return True  # categorical-style domain for numeric column
+                return True
             if col_domain.get("lower") is None or col_domain.get("upper") is None:
                 return True
     return False
@@ -103,6 +134,8 @@ def fit(self, data: pd.DataFrame, domain: Optional[dict] = None) -> "AIMGenerato
         UserWarning
             If privacy budget for generation is smaller than for processing.
         """
+        data, domain = _auto_categorize(data, domain)
+
         needs_preprocessing = _requires_private_preprocessing(data, domain)
 
         if needs_preprocessing:
@@ -135,7 +168,22 @@ def fit(self, data: pd.DataFrame, domain: Optional[dict] = None) -> "AIMGenerato
             proc_epsilon=params.get("proc_epsilon"),
             gen_kwargs={"degree": self._degree},
         )
-        self._pipeline.fit(data, domain)
+        _BOUNDS_ERROR_MSG = (
+            "Private bounds estimation failed for one or more numeric columns. "
+            "This typically happens when the privacy budget is too small to detect "
+            "data bounds. Remedies: (1) provide explicit domain bounds for numeric "
+            "columns via the domain parameter, e.g. domain={'col': {'lower': 0, "
+            "'upper': 100}}, (2) increase proc_epsilon, or (3) relax the risk "
+            "requirement."
+        )
+        try:
+            self._pipeline.fit(data, domain)
+        except (TypeError, KeyError) as e:
+            raise ValueError(_BOUNDS_ERROR_MSG) from e
+        except ValueError as e:
+            if "Private bounds estimation failed" not in str(e):
+                raise ValueError(_BOUNDS_ERROR_MSG) from e
+            raise
         return self
 
     def generate(self, count: int) -> pd.DataFrame:

tests/test_generator.py

@@ -4,6 +4,7 @@
 from sklearn.datasets import load_wine
 
 from risksyn import AIMGenerator, Risk
+from risksyn.generator import _auto_categorize, _requires_private_preprocessing
 
 
 # Simple dataset for fast tests
@@ -112,6 +113,71 @@ def test_generate_before_fit_raises():
         gen.generate(count=10)
 
 
+def test_auto_categorize_binary_int_columns():
+    """Low-cardinality int columns should be auto-categorized and cast to str."""
+    df = pd.DataFrame({
+        "binary": [0, 1, 0, 1, 1],
+        "ternary": [0, 1, 2, 0, 1],
+        "continuous": np.random.uniform(0, 100, 5),
+        "cat": ["a", "b", "c", "a", "b"],
+    })
+    out_data, domain = _auto_categorize(df, None)
+    assert domain["binary"] == ["0", "1"]
+    assert domain["ternary"] == ["0", "1", "2"]
+    assert out_data["binary"].dtype == object
+    assert out_data["ternary"].dtype == object
+    assert "continuous" not in domain  # float, not int
+    assert "cat" not in domain  # not numeric
+
+
+def test_auto_categorize_respects_existing_domain():
+    """Auto-categorization should not override user-provided domain."""
+    df = pd.DataFrame({"x": [0, 1, 0, 1]})
+    user_domain = {"x": {"lower": 0, "upper": 1}}
+    _, domain = _auto_categorize(df, user_domain)
+    assert domain["x"] == {"lower": 0, "upper": 1}
+
+
+def test_auto_categorize_skips_high_cardinality():
+    """Int columns with >10 unique values should not be auto-categorized."""
+    df = pd.DataFrame({"x": list(range(11))})
+    _, domain = _auto_categorize(df, None)
+    assert "x" not in domain
+
+
+def test_requires_private_preprocessing_false_for_list_domain():
+    """List domain on a numeric column means no preprocessing needed."""
+    df = pd.DataFrame({"x": [0, 1, 0, 1]})
+    assert not _requires_private_preprocessing(df, {"x": [0, 1]})
+
+
+def test_binary_int_columns_fit_without_domain():
+    """Binary int columns should fit without explicit domain via auto-categorization."""
+    np.random.seed(42)
+    df = pd.DataFrame({
+        "a": np.random.choice([0, 1], 100),
+        "b": np.random.choice([0, 1], 100),
+        "cat": np.random.choice(["x", "y"], 100),
+    })
+    risk = Risk.from_advantage(0.25)
+    gen = AIMGenerator(risk=risk)
+    gen.fit(df)
+    synth = gen.generate(count=10)
+    assert len(synth) == 10
+    assert list(synth.columns) == list(df.columns)
+
+
+def test_bounds_estimation_failure_raises_value_error():
+    """Should raise ValueError (not TypeError) when private bounds estimation fails."""
+    np.random.seed(42)
+    # High-cardinality float column with tiny budget -> approx_bounds will fail
+    df = pd.DataFrame({"x": np.random.uniform(0, 1, 50)})
+    risk = Risk.from_zcdp(0.001)
+    gen = AIMGenerator(risk=risk, proc_epsilon=0.001)
+    with pytest.raises(ValueError, match="Private bounds estimation failed"):
+        gen.fit(df)
+
+
 def test_categorical_only_no_preprocessing():
     """Categorical-only data should not require preprocessing."""
     df = pd.DataFrame({