Add learning_strategy parameter to OneHotEncoding constraint (#2658)

Author: fealho

sdv/cag/one_hot_encoding.py

@@ -4,10 +4,12 @@
 
 import numpy as np
 
+from sdv._utils import _create_unique_name
 from sdv.cag._errors import ConstraintNotMetError
 from sdv.cag._utils import (
     _get_is_valid_dict,
     _is_list_of_type,
+    _remove_columns_from_metadata,
     _validate_table_and_column_names,
     _validate_table_name_if_defined,
 )
@@ -30,20 +32,30 @@ class OneHotEncoding(BaseConstraint):
         table_name (str, optional):
             The name of the table that contains the columns. Optional if the
             data is only a single table. Defaults to None.
+        learning_strategy (str, optional):
+            Strategy for how the model should learn the one-hot fields. Supported values:
+            - 'one_hot' (default): Learn each one-hot column separately.
+            - 'categorical': Internally collapse the one-hot columns into a single categorical
+              column for the model to learn, then expand back to one-hot at sampling time.
     """
 
     @staticmethod
-    def _validate_init_inputs(column_names, table_name):
+    def _validate_init_inputs(column_names, table_name, learning_strategy):
         if not _is_list_of_type(column_names):
             raise ValueError('`column_names` must be a list of strings.')
 
         _validate_table_name_if_defined(table_name)
 
-    def __init__(self, column_names, table_name=None):
+        if learning_strategy not in ['one_hot', 'categorical']:
+            raise ValueError("`learning_strategy` must be either 'one_hot' or 'categorical'.")
+
+    def __init__(self, column_names, table_name=None, learning_strategy='one_hot'):
         super().__init__()
-        self._validate_init_inputs(column_names, table_name)
+        self._validate_init_inputs(column_names, table_name, learning_strategy)
         self._column_names = column_names
         self.table_name = table_name
+        self.learning_strategy = learning_strategy
+        self._categorical_column = '#'.join(self._column_names)
 
     def _validate_constraint_with_metadata(self, metadata):
         """Validate the constraint is compatible with the provided metadata.
@@ -94,12 +106,25 @@ def _fit(self, data, metadata):
 
     def _get_updated_metadata(self, metadata):
         table_name = self._get_single_table_name(metadata)
-        metadata = deepcopy(metadata)
-        for column in self._column_names:
-            if metadata.tables[table_name].columns[column]['sdtype'] in ['categorical', 'boolean']:
-                metadata.tables[table_name].columns[column]['sdtype'] = 'numerical'
-
-        return metadata
+        if self.learning_strategy == 'categorical':
+            self._categorical_column = _create_unique_name(
+                self._categorical_column, metadata.tables[table_name].columns
+            )
+            md = metadata.to_dict()
+            md['tables'][table_name]['columns'][self._categorical_column] = {
+                'sdtype': 'categorical'
+            }
+            return _remove_columns_from_metadata(md, table_name, columns_to_drop=self._column_names)
+
+        else:
+            metadata = deepcopy(metadata)
+            for column in self._column_names:
+                if metadata.tables[table_name].columns[column]['sdtype'] in [
+                    'categorical',
+                    'boolean',
+                ]:
+                    metadata.tables[table_name].columns[column]['sdtype'] = 'numerical'
+            return metadata
 
     def _transform(self, data):
         """Transform the data.
@@ -113,9 +138,15 @@ def _transform(self, data):
                 Transformed data.
         """
         table_name = self._get_single_table_name(self.metadata)
-        one_hot_data = data[table_name][self._column_names]
-        one_hot_data = np.where(one_hot_data == 0, EPSILON, 1 - EPSILON)
-        data[table_name][self._column_names] = one_hot_data
+        if self.learning_strategy == 'categorical':
+            table_data = data[table_name]
+            categories = table_data[self._column_names].idxmax(axis=1)
+            table_data[self._categorical_column] = categories
+            data[table_name] = table_data.drop(self._column_names, axis=1)
+        else:
+            one_hot_data = data[table_name][self._column_names]
+            one_hot_data = np.where(one_hot_data == 0, EPSILON, 1 - EPSILON)
+            data[table_name][self._column_names] = one_hot_data
 
         return data
 
@@ -134,13 +165,28 @@ def _reverse_transform(self, data):
         """
         table_name = self._get_single_table_name(self.metadata)
         table_data = data[table_name]
-        one_hot_data = table_data[self._column_names]
-        transformed_data = np.zeros_like(one_hot_data.to_numpy())
-        max_category_indices = np.argmax(one_hot_data.to_numpy(), axis=1)
-        transformed_data[np.arange(len(one_hot_data)), max_category_indices] = 1
-        table_data[self._column_names] = transformed_data
-        data[table_name] = table_data
 
+        if self.learning_strategy == 'categorical':
+            categories = table_data.pop(self._categorical_column)
+            num_rows = len(table_data)
+            num_cols = len(self._column_names)
+            transformed = np.zeros((num_rows, num_cols), dtype=float)
+
+            column_to_index = {name: idx for idx, name in enumerate(self._column_names)}
+            indices = categories.map(lambda x: column_to_index[x]).to_numpy()
+            transformed[np.arange(num_rows), indices] = 1
+
+            for idx, col in enumerate(self._column_names):
+                table_data[col] = transformed[:, idx]
+
+        else:
+            one_hot_data = table_data[self._column_names]
+            transformed_data = np.zeros_like(one_hot_data.to_numpy())
+            max_category_indices = np.argmax(one_hot_data.to_numpy(), axis=1)
+            transformed_data[np.arange(len(one_hot_data)), max_category_indices] = 1
+            table_data[self._column_names] = transformed_data
+
+        data[table_name] = table_data
         return data
 
     def _is_valid(self, data, metadata):

tests/integration/cag/test_one_hot_encoding.py

@@ -8,7 +8,7 @@
 from sdv.cag._errors import ConstraintNotMetError
 from sdv.metadata import Metadata
 from sdv.single_table import GaussianCopulaSynthesizer
-from tests.utils import run_copula, run_hma
+from tests.utils import run_constraint, run_copula, run_hma
 
 
 @pytest.fixture()
@@ -201,3 +201,144 @@ def test_end_to_end_boolean():
     assert (samples.sum(axis=1) == 1).all()
     for col in columns:
         assert sorted(samples[col].unique().tolist()) == [0, 1]
+
+
+def test_end_to_end_categorical_single(data, metadata):
+    """End-to-end with learning_strategy='categorical' for single-table data."""
+    # Setup
+    constraint = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+    # Run
+    synthesizer = run_copula(data, metadata, [constraint])
+    synthetic_data = synthesizer.sample(200)
+    synthesizer.validate_constraints(synthetic_data=synthetic_data)
+
+    # Assert
+    assert set(synthetic_data.columns) == {'a', 'b', 'c'}
+    for col in ['a', 'b', 'c']:
+        assert set(synthetic_data[col]) == {0, 1}
+    assert (synthetic_data[['a', 'b', 'c']].sum(axis=1) == 1).all()
+
+
+def test_end_to_end_categorical_single_raises(data, metadata):
+    """Invalid synthetic data should raise with learning_strategy='categorical'."""
+    # Setup
+    invalid_data = pd.DataFrame({
+        'a': [1, 2, 0],
+        'b': [0, 1, np.nan],
+        'c': [0, 0, 3],
+    })
+    constraint = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+    # Run and Assert
+    msg = re.escape(
+        "Data is not valid for the 'OneHotEncoding' constraint in table 'table':\n"
+        '   a    b  c\n'
+        '1  2  1.0  0\n'
+        '2  0  NaN  3'
+    )
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        run_copula(invalid_data, metadata, [constraint])
+
+    # Run and Assert
+    msg = re.escape('The one hot encoding requirement is not met for row indices: 1, 2')
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        synthesizer = run_copula(data, metadata, [constraint])
+        synthesizer.validate_constraints(synthetic_data=invalid_data)
+
+
+def test_end_to_end_categorical_multi(data_multi, metadata_multi):
+    """End-to-end with learning_strategy='categorical' for multi-table data."""
+    # Setup
+    constraint = OneHotEncoding(
+        column_names=['a', 'b', 'c'], table_name='table1', learning_strategy='categorical'
+    )
+
+    # Run
+    synthesizer = run_hma(data_multi, metadata_multi, [constraint])
+    synthetic = synthesizer.sample(200)
+    synthesizer.validate_constraints(synthetic_data=synthetic)
+
+    # Assert
+    assert set(synthetic['table1'].columns) == {'a', 'b', 'c'}
+    for col in ['a', 'b', 'c']:
+        assert set(synthetic['table1'][col]) == {0, 1}
+    assert (synthetic['table1'][['a', 'b', 'c']].sum(axis=1) == 1).all()
+
+
+def test_end_to_end_categorical_multi_raises(data_multi, metadata_multi):
+    """Invalid multi-table synthetic data should raise with learning_strategy='categorical'."""
+    # Setup
+    constraint = OneHotEncoding(
+        column_names=['a', 'b', 'c'], table_name='table1', learning_strategy='categorical'
+    )
+    invalid = {
+        'table1': pd.DataFrame({
+            'a': [1, 2, 0],
+            'b': [0, 1, np.nan],
+            'c': [0, 0, 3],
+        }),
+        'table2': pd.DataFrame({'id': range(5)}),
+    }
+
+    # Run and Assert
+    msg = re.escape(
+        "Data is not valid for the 'OneHotEncoding' constraint in table 'table1':\n   "
+        'a    b  c\n1  2  1.0  0\n2  0  NaN  3'
+    )
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        run_hma(invalid, metadata_multi, [constraint])
+
+    # Run and Assert
+    msg = "Table 'table1': The one hot encoding requirement is not met for row indices: 1, 2."
+    with pytest.raises(ConstraintNotMetError, match=msg):
+        synthesizer = run_hma(data_multi, metadata_multi, [constraint])
+        synthesizer.validate_constraints(synthetic_data=invalid)
+
+
+def test_constraint_pipeline_categorical_single(data, metadata):
+    """Constraint pipeline behavior for categorical strategy (single table)."""
+    # Setup
+    constraint = OneHotEncoding(column_names=['a', 'b', 'c'], learning_strategy='categorical')
+
+    # Run
+    updated_metadata, transformed, reverse_transformed = run_constraint(constraint, data, metadata)
+
+    # Assert metadata
+    assert updated_metadata.get_column_names() == ['a#b#c']
+
+    # Assert transform
+    assert transformed.shape[1] == 1
+    assert not any(col in transformed.columns for col in ['a', 'b', 'c'])
+    assert set(transformed.columns) == {'a#b#c'}
+
+    # Assert reverse_transform
+    assert set(reverse_transformed.columns) == {'a', 'b', 'c'}
+    assert (reverse_transformed[['a', 'b', 'c']].sum(axis=1) == 1).all()
+    assert set(reverse_transformed.columns) == {'a', 'b', 'c'}
+
+
+def test_constraint_pipeline_categorical_multi(data_multi, metadata_multi):
+    """Constraint pipeline behavior for categorical strategy (multi table)."""
+    # Setup
+    orig_cols = ['a', 'b', 'c']
+    constraint = OneHotEncoding(
+        column_names=orig_cols, table_name='table1', learning_strategy='categorical'
+    )
+
+    # Run
+    updated_metadata, transformed, reverse_transformed = run_constraint(
+        constraint, data_multi, metadata_multi
+    )
+
+    # Assert metadata
+    assert updated_metadata.tables['table1'].get_column_names() == ['a#b#c']
+
+    # Assert transform
+    assert list(transformed['table1'].columns) != orig_cols
+    assert transformed['table1'].shape[1] == 1
+    assert list(transformed['table2'].columns) == list(data_multi['table2'].columns)
+
+    # Assert reverse_transform
+    assert set(reverse_transformed['table1'].columns) == set(orig_cols)
+    assert (reverse_transformed['table1'][orig_cols].sum(axis=1) == 1).all()