Add OneHotEncoding CAG (#2414)

Author: fealho

sdv/cag/__init__.py

@@ -4,5 +4,12 @@
 from sdv.cag.fixed_increments import FixedIncrements
 from sdv.cag.inequality import Inequality
 from sdv.cag.range import Range
+from sdv.cag.one_hot_encoding import OneHotEncoding
 
-__all__ = ('FixedCombinations', 'FixedIncrements', 'Inequality', 'Range')
+__all__ = (
+    'FixedCombinations',
+    'FixedIncrements',
+    'Inequality',
+    'Range',
+    'OneHotEncoding',
+)

sdv/cag/_utils.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pandas as pd
 
 from sdv.cag._errors import PatternNotMetError
 from sdv.metadata import Metadata
@@ -93,9 +94,9 @@ def _remove_columns_from_metadata(metadata, table_name, columns_to_drop):
     return Metadata.load_from_dict(metadata)
 
 
-def _is_list_of_strings(values):
-    """Checks that a list contains all strings."""
-    return isinstance(values, list) and all(isinstance(value, str) for value in values)
+def _is_list_of_type(values, type_to_check=str):
+    """Checks that 'values' is a list and all elements are of type 'type_to_check'."""
+    return isinstance(values, list) and all(isinstance(value, type_to_check) for value in values)
 
 
 def _get_invalid_rows(valid):
@@ -117,3 +118,26 @@ def _get_invalid_rows(valid):
         remaining = len(invalid_rows) - 5
         invalid_rows_str = f'{first_five}, +{remaining} more'
     return invalid_rows_str
+
+
+def _get_is_valid_dict(data, table_name):
+    """Create a dictionary of True values for each table besides table_name.
+
+    Besides table_name, all rows of every other table are considered valid,
+    so the boolean Series will be True for all rows of every other table.
+
+    Args:
+        data (dict):
+            The data.
+        table_name (str):
+            The name of the table to exclude from the dictionary.
+
+    Returns:
+        dict:
+            Dictionary of table names to boolean Series of True values.
+    """
+    return {
+        table: pd.Series(True, index=table_data.index)
+        for table, table_data in data.items()
+        if table != table_name
+    }

sdv/cag/fixed_combinations.py

@@ -8,7 +8,8 @@
 from sdv._utils import _create_unique_name
 from sdv.cag._errors import PatternNotMetError
 from sdv.cag._utils import (
-    _is_list_of_strings,
+    _get_is_valid_dict,
+    _is_list_of_type,
     _validate_table_and_column_names,
     _validate_table_name_if_defined,
 )
@@ -41,7 +42,7 @@ class FixedCombinations(BasePattern):
 
     def __init__(self, column_names, table_name=None):
         super().__init__()
-        if not _is_list_of_strings(column_names):
+        if not _is_list_of_type(column_names):
             raise ValueError('`column_names` must be a list of strings.')
 
         if len(column_names) < 2:
@@ -192,11 +193,7 @@ def _reverse_transform(self, data):
     def _is_valid(self, data):
         """Determine whether the data matches the pattern."""
         table_name = self._get_single_table_name(self.metadata)
-        is_valid = {
-            table: pd.Series(True, index=table_data.index)
-            for table, table_data in data.items()
-            if table != table_name
-        }
+        is_valid = _get_is_valid_dict(data, table_name)
         merged = data[table_name].merge(
             self._combinations, how='left', on=self.column_names, indicator=self._joint_column
         )

sdv/cag/fixed_increments.py

@@ -6,6 +6,7 @@
 from sdv.cag._errors import PatternNotMetError
 from sdv.cag._utils import (
     _get_invalid_rows,
+    _get_is_valid_dict,
     _remove_columns_from_metadata,
     _validate_table_and_column_names,
     _validate_table_name_if_defined,
@@ -173,11 +174,7 @@ def _is_valid(self, data):
                 table names.
         """
         table_name = self._get_single_table_name(self.metadata)
-        is_valid = {
-            table: pd.Series(True, index=table_data.index)
-            for table, table_data in data.items()
-            if table != table_name
-        }
+        is_valid = _get_is_valid_dict(data, table_name)
         valid = self._check_if_divisible(data, table_name, self.column_name, self.increment_value)
         is_valid[table_name] = valid
         return is_valid

sdv/cag/inequality.py

@@ -5,7 +5,10 @@
 
 from sdv._utils import _convert_to_timedelta, _create_unique_name
 from sdv.cag._errors import PatternNotMetError
-from sdv.cag._utils import _validate_table_and_column_names
+from sdv.cag._utils import (
+    _get_is_valid_dict,
+    _validate_table_and_column_names,
+)
 from sdv.cag.base import BasePattern
 from sdv.constraints.utils import (
     cast_to_datetime64,
@@ -293,12 +296,7 @@ def _is_valid(self, data):
                 Whether each row is valid.
         """
         table_name = self._get_single_table_name(self.metadata)
-        is_valid = {
-            table: pd.Series(True, index=table_data.index)
-            for table, table_data in data.items()
-            if table != table_name
-        }
-
+        is_valid = _get_is_valid_dict(data, table_name)
         table_data = data[table_name]
         low, high = self._get_data(table_data)
         if self._is_datetime and self._dtype == 'O':

sdv/cag/one_hot_encoding.py

@@ -0,0 +1,145 @@
+"""One Hot Encoding CAG pattern."""
+
+import numpy as np
+
+from sdv.cag._errors import PatternNotMetError
+from sdv.cag._utils import (
+    _get_invalid_rows,
+    _get_is_valid_dict,
+    _is_list_of_type,
+    _validate_table_and_column_names,
+)
+from sdv.cag.base import BasePattern
+
+
+class OneHotEncoding(BasePattern):
+    """Ensure the appropriate columns are one hot encoded.
+
+    This constraint allows the user to specify a list of columns where each row
+    is a one hot vector. During the reverse transform, the output of the model
+    is transformed so that the column with the largest value is set to 1 while
+    all other columns are set to 0.
+
+    Args:
+        column_names (list[str]):
+            Names of the columns containing one hot rows.
+        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.
+    """
+
+    @staticmethod
+    def _validate_init_inputs(column_names, table_name):
+        if not _is_list_of_type(column_names):
+            raise ValueError('`column_names` must be a list of strings.')
+
+        if table_name and not isinstance(table_name, str):
+            raise ValueError('`table_name` must be a string or None.')
+
+    def __init__(self, column_names, table_name=None):
+        super().__init__()
+        self._validate_init_inputs(column_names, table_name)
+        self._column_names = column_names
+        self.table_name = table_name
+
+    def _validate_pattern_with_metadata(self, metadata):
+        """Validate the pattern is compatible with the provided metadata.
+
+        Validates that:
+        - If no table_name is provided the metadata contains a single table
+        - All input columns exist in the table in the metadata.
+
+        Args:
+            metadata (sdv.metadata.Metadata):
+                The metadata to validate against.
+
+        Raises:
+            PatternNotMetError:
+                If any of the validations fail.
+        """
+        _validate_table_and_column_names(self.table_name, self._column_names, metadata)
+
+    def _get_valid_table_data(self, table_data):
+        one_hot_data = table_data[self._column_names]
+
+        sum_one = one_hot_data.sum(axis=1) == 1.0
+        max_one = one_hot_data.max(axis=1) == 1.0
+        min_zero = one_hot_data.min(axis=1) == 0.0
+        no_nans = ~one_hot_data.isna().any(axis=1)
+
+        return sum_one & max_one & min_zero & no_nans
+
+    def _validate_pattern_with_data(self, data, metadata):
+        """Validate the data is compatible with the pattern."""
+        table_name = self._get_single_table_name(metadata)
+        valid = self._get_valid_table_data(data[table_name])
+        if not valid.all():
+            invalid_rows_str = _get_invalid_rows(valid)
+            raise PatternNotMetError(
+                f'The one hot encoding requirement is not met for row indices: [{invalid_rows_str}]'
+            )
+
+    def _fit(self, data, metadata):
+        """Fit the pattern.
+
+        Args:
+            data (dict[str, pd.DataFrame]):
+                Table data.
+            metadata (sdv.metadata.Metadata):
+                Metadata.
+        """
+        pass
+
+    def _transform(self, data):
+        """Transform the data.
+
+        Args:
+            data (dict[str, pd.DataFrame]):
+                Table data.
+
+        Returns:
+            dict[str, pd.DataFrame]:
+                Transformed data.
+        """
+        return data
+
+    def _reverse_transform(self, data):
+        """Reverse transform the table data.
+
+        Set the column with the largest value to one, set all other columns to zero.
+
+        Args:
+            data (dict[str, pd.DataFrame]):
+                Table data.
+
+        Returns:
+            dict[str, pd.DataFrame]:
+                Transformed 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
+
+        return data
+
+    def _is_valid(self, data):
+        """Check whether the data satisfies the one-hot constraint.
+
+        Args:
+            data (dict[str, pd.DataFrame]):
+                Table data.
+
+        Returns:
+            dict[str, pd.Series]:
+                Whether each row is valid.
+        """
+        table_name = self._get_single_table_name(self.metadata)
+        is_valid = _get_is_valid_dict(data, table_name)
+        is_valid[table_name] = self._get_valid_table_data(data[table_name])
+
+        return is_valid

sdv/cag/range.py

@@ -7,7 +7,10 @@
 
 from sdv._utils import _convert_to_timedelta, _create_unique_name
 from sdv.cag._errors import PatternNotMetError
-from sdv.cag._utils import _validate_table_and_column_names
+from sdv.cag._utils import (
+    _get_is_valid_dict,
+    _validate_table_and_column_names,
+)
 from sdv.cag.base import BasePattern
 from sdv.constraints.utils import (
     cast_to_datetime64,
@@ -355,11 +358,7 @@ def _is_valid(self, data):
                 Whether each row is valid.
         """
         table_name = self._get_single_table_name(self.metadata)
-        is_valid = {
-            table: pd.Series(True, index=table_data.index)
-            for table, table_data in data.items()
-            if table != table_name
-        }
+        is_valid = _get_is_valid_dict(data, table_name)
         is_valid[table_name] = self._get_valid_table_data(data[table_name])
 
         return is_valid

tests/unit/cag/test__utils.py

@@ -7,7 +7,7 @@
 
 from sdv.cag._errors import PatternNotMetError
 from sdv.cag._utils import (
-    _is_list_of_strings,
+    _is_list_of_type,
     _remove_columns_from_metadata,
     _validate_table_and_column_names,
     _validate_table_name_if_defined,
@@ -192,10 +192,10 @@ def test__remove_columns_from_metadata_raises_pk():
         )
 
 
-def test__is_list_of_strings():
-    """Test `_is_list_of_strings` method"""
-    assert _is_list_of_strings(['a', 'b'])
-    assert not _is_list_of_strings(['a', 1])
-    assert not _is_list_of_strings([1, 2])
-    assert not _is_list_of_strings(1)
-    assert not _is_list_of_strings('a')
+def test__is_list_of_type():
+    """Test `_is_list_of_type` method"""
+    assert _is_list_of_type(['a', 'b'])
+    assert not _is_list_of_type(['a', 1])
+    assert not _is_list_of_type([1, 2])
+    assert not _is_list_of_type(1)
+    assert not _is_list_of_type('a')