Add CAG support to single table synthesizers (#2419)

Author: fealho

sdv/evaluation/single_table.py

@@ -31,9 +31,6 @@ def evaluate_quality(real_data, synthetic_data, metadata, verbose=True):
         metadata = metadata._convert_to_single_table()
 
     quality_report = QualityReport()
-    if isinstance(metadata, Metadata):
-        metadata = metadata._convert_to_single_table()
-
     quality_report.generate(real_data, synthetic_data, metadata.to_dict(), verbose)
     return quality_report
 

sdv/single_table/base.py

@@ -10,6 +10,7 @@
 import uuid
 import warnings
 from collections import defaultdict
+from copy import deepcopy
 
 import cloudpickle
 import copulas
@@ -26,6 +27,7 @@
     generate_synthesizer_id,
     get_possible_chars,
 )
+from sdv.cag._errors import PatternNotMetError
 from sdv.constraints.errors import AggregateConstraintsError
 from sdv.data_processing.data_processor import DataProcessor
 from sdv.errors import (
@@ -610,6 +612,118 @@ class BaseSingleTableSynthesizer(BaseSynthesizer):
     for all single-table synthesizers.
     """
 
+    def __init__(
+        self,
+        metadata,
+        enforce_min_max_values=True,
+        enforce_rounding=True,
+        locales=['en_US'],
+    ):
+        super().__init__(metadata, enforce_min_max_values, enforce_rounding, locales)
+        self._chained_patterns = []  # chain of patterns used to preprocess the data
+        self._reject_sampling_patterns = []  # patterns used only for reject sampling
+        self._original_metadata = self.metadata
+
+    def add_cag(self, patterns):
+        """Add the list of constraint-augmented generation patterns to the synthesizer.
+
+        Args:
+            patterns (list):
+                A list of CAG patterns to apply to the synthesizer.
+        """
+        for pattern in patterns:
+            try:
+                self.metadata = pattern.get_updated_metadata(self.metadata)
+                self._chained_patterns.append(pattern)
+            except PatternNotMetError as e:
+                LOGGER.info(
+                    'Enforcing pattern %s using reject sampling.', pattern.__class__.__name__
+                )
+
+                try:
+                    pattern.get_updated_metadata(self._original_metadata)
+                    self._reject_sampling_patterns.append(pattern)
+                except PatternNotMetError:
+                    raise e
+
+        self._data_processor = DataProcessor(
+            metadata=self.metadata._convert_to_single_table(),
+            enforce_rounding=self.enforce_rounding,
+            enforce_min_max_values=self.enforce_min_max_values,
+            locales=self.locales,
+        )
+
+    def get_cag(self):
+        """Get a list of constraint-augmented generation patterns applied to the synthesizer."""
+        return deepcopy(self._chained_patterns + self._reject_sampling_patterns)
+
+    def get_metadata(self, version='original'):
+        """Get the metadata, either original or modified after applying CAG patterns.
+
+        Args:
+            version (str, optional):
+                The version of metadata to return, must be one of 'original' or 'modified'. If
+                'original', will return the original metadata used to instantiate the
+                synthesizer. If 'modified', will return the modified metadata after applying this
+                synthesizer's CAG patterns. Defaults to 'original'.
+        """
+        if version not in ('original', 'modified'):
+            raise ValueError(
+                f"Unrecognized version '{version}', please use 'original' or 'modified'."
+            )
+
+        return self._original_metadata if version == 'original' else self.metadata
+
+    def _transform_helper(self, data):
+        """Validate and transform all CAG patterns during preprocessing.
+
+        Args:
+            data (dict[str, pd.DataFrame]):
+                The data dictionary.
+        """
+        if self._fitted:
+            for pattern in self._chained_patterns:
+                data = pattern.transform(data)
+            return data
+
+        metadata = self._original_metadata
+        original_data = data
+        for pattern in self._chained_patterns:
+            pattern.fit(data, metadata)
+            metadata = pattern.get_updated_metadata(metadata)
+            data = pattern.transform(data)
+
+        for pattern in self._reject_sampling_patterns:
+            pattern.fit(original_data, self._original_metadata)
+
+        return data
+
+    def preprocess(self, data):
+        """Transform the raw data to numerical space.
+
+        Args:
+            data (pandas.DataFrame):
+                The raw data to be transformed.
+
+        Returns:
+            pandas.DataFrame:
+                The preprocessed data.
+        """
+        if self._fitted:
+            warnings.warn(
+                'This model has already been fitted. To use the new preprocessed data, '
+                "please refit the model using 'fit' or 'fit_processed_data'."
+            )
+
+        is_converted = self._store_and_convert_original_cols(data)
+        data = self._transform_helper(data)
+        preprocess_data = self._preprocess(data)
+
+        if is_converted:
+            data.columns = self._original_columns
+
+        return preprocess_data
+
     def _set_random_state(self, random_state):
         """Set the random state of the model's random number generator.
 
@@ -720,6 +834,15 @@ def _sample_rows(
                 )
                 sampled = pd.concat([sampled, raw_sampled[missing_cols]], axis=1)
 
+            for pattern in reversed(self._chained_patterns):
+                sampled = pattern.reverse_transform(sampled)
+                valid_rows = pattern.is_valid(sampled)
+                sampled = sampled[valid_rows]
+
+            for pattern in reversed(self._reject_sampling_patterns):
+                valid_rows = pattern.is_valid(sampled)
+                sampled = sampled[valid_rows]
+
             if previous_rows is not None:
                 sampled = pd.concat([previous_rows, sampled], ignore_index=True)
             sampled = self._data_processor.filter_valid(sampled)

tests/integration/cag/test_fixed_combinations.py

@@ -5,6 +5,7 @@
 
 from sdv.cag import FixedCombinations
 from sdv.metadata import Metadata
+from sdv.single_table.copulas import GaussianCopulaSynthesizer
 from tests.utils import run_pattern
 
 
@@ -117,3 +118,221 @@ def test_fixed_null_combinations_with_multi_table():
     assert set(data.keys()) == set(reverse_transformed.keys())
     for table_name, table in data.items():
         pd.testing.assert_frame_equal(table, reverse_transformed[table_name])
+
+
+def test_fixed_combinations_multiple_patterns():
+    """Test that FixedCombinations pattern works with multiple patterns."""
+    # Setup
+    data = pd.DataFrame({
+        'A': [1, 2, 3, 1, 2, 1],
+        'B': [10, 20, 30, 10, 20, 10],
+        'C': [100, 200, 300, 100, 200, 100],
+        'D': [1000, 2000, 3000, 1000, 2000, 1000],
+    })
+    metadata = Metadata.load_from_dict({
+        'columns': {
+            'A': {'sdtype': 'categorical'},
+            'B': {'sdtype': 'categorical'},
+            'C': {'sdtype': 'categorical'},
+            'D': {'sdtype': 'categorical'},
+        }
+    })
+    pattern1 = FixedCombinations(['A', 'B'])
+    pattern2 = FixedCombinations(['C', 'D'])
+
+    # Run
+    synthesizer = GaussianCopulaSynthesizer(metadata)
+    synthesizer.add_cag(patterns=[pattern1, pattern2])
+    synthesizer.fit(data)
+    samples = synthesizer.sample(100)
+    updated_metadata = synthesizer.get_metadata('modified')
+    original_metadata = synthesizer.get_metadata('original')
+
+    # Assert
+    expected_updated_metadata = Metadata.load_from_dict({
+        'columns': {
+            'A#B': {'sdtype': 'categorical'},
+            'C#D': {'sdtype': 'categorical'},
+        }
+    }).to_dict()
+    assert expected_updated_metadata == updated_metadata.to_dict()
+
+    assert original_metadata.to_dict() == metadata.to_dict()
+
+    # Get unique combinations from original data
+    original_ab_combos = set(zip(data['A'], data['B']))
+    original_cd_combos = set(zip(data['C'], data['D']))
+
+    # Get unique combinations from synthetic data
+    synthetic_ab_combos = set(zip(samples['A'], samples['B']))
+    synthetic_cd_combos = set(zip(samples['C'], samples['D']))
+
+    # Assert combinations match
+    assert original_ab_combos == synthetic_ab_combos
+    assert original_cd_combos == synthetic_cd_combos
+
+
+def test_fixed_combinations_multiple_patterns_reject_sampling():
+    """Test that FixedCombinations pattern works with multiple patterns and reject sampling."""
+    # Setup
+    data = pd.DataFrame({
+        'A': [1, 2, 3, 1, 2, 1],
+        'B': [10, 20, 30, 10, 20, 10],
+        'C': [100, 200, 300, 100, 200, 100],
+    })
+    metadata = Metadata.load_from_dict({
+        'columns': {
+            'A': {'sdtype': 'categorical'},
+            'B': {'sdtype': 'categorical'},
+            'C': {'sdtype': 'categorical'},
+        }
+    })
+    pattern1 = FixedCombinations(['A', 'B'])
+    pattern2 = FixedCombinations(['A', 'C'])
+
+    # Run
+    synthesizer = GaussianCopulaSynthesizer(metadata)
+    synthesizer.add_cag(patterns=[pattern1, pattern2])
+    synthesizer.fit(data)
+    samples = synthesizer.sample(100)
+    updated_metadata = synthesizer.get_metadata('modified')
+    original_metadata = synthesizer.get_metadata('original')
+
+    # Assert
+    expected_updated_metadata = Metadata.load_from_dict({
+        'columns': {
+            'A#B': {'sdtype': 'categorical'},
+            'C': {'sdtype': 'categorical'},
+        }
+    }).to_dict()
+    assert expected_updated_metadata == updated_metadata.to_dict()
+
+    assert original_metadata.to_dict() == metadata.to_dict()
+
+    # Get unique combinations from original data
+    original_ab_combos = set(zip(data['A'], data['B']))
+    original_ac_combos = set(zip(data['A'], data['C']))
+
+    # Get unique combinations from synthetic data
+    synthetic_ab_combos = set(zip(samples['A'], samples['B']))
+    synthetic_ac_combos = set(zip(samples['A'], samples['C']))
+
+    # Assert combinations match
+    assert original_ab_combos == synthetic_ab_combos
+    assert original_ac_combos == synthetic_ac_combos
+
+
+def test_fixed_combinations_multiple_patterns_three_patterns():
+    """Test that FixedCombinations pattern works with multiple patterns."""
+    # Setup
+    data = pd.DataFrame({
+        'A': [1, 2, 3, 1, 2, 1],
+        'B': [10, 20, 30, 10, 20, 10],
+        'C': [100, 200, 300, 100, 200, 100],
+        'D': [1000, 2000, 3000, 1000, 2000, 1000],
+    })
+    metadata = Metadata.load_from_dict({
+        'columns': {
+            'A': {'sdtype': 'categorical'},
+            'B': {'sdtype': 'categorical'},
+            'C': {'sdtype': 'categorical'},
+            'D': {'sdtype': 'categorical'},
+        }
+    })
+    pattern1 = FixedCombinations(['A', 'B'])
+    pattern2 = FixedCombinations(['C', 'D'])
+    pattern3 = FixedCombinations(['A', 'C'])
+
+    # Run
+    synthesizer = GaussianCopulaSynthesizer(metadata)
+    synthesizer.add_cag(patterns=[pattern1, pattern2, pattern3])
+    synthesizer.fit(data)
+    samples = synthesizer.sample(100)
+    updated_metadata = synthesizer.get_metadata('modified')
+    original_metadata = synthesizer.get_metadata('original')
+
+    # Assert
+    expected_updated_metadata = Metadata.load_from_dict({
+        'columns': {
+            'A#B': {'sdtype': 'categorical'},
+            'C#D': {'sdtype': 'categorical'},
+        }
+    }).to_dict()
+    assert expected_updated_metadata == updated_metadata.to_dict()
+
+    assert original_metadata.to_dict() == metadata.to_dict()
+
+    # Get unique combinations from original data
+    original_ab_combos = set(zip(data['A'], data['B']))
+    original_cd_combos = set(zip(data['C'], data['D']))
+    original_ac_combos = set(zip(data['A'], data['C']))
+
+    # Get unique combinations from synthetic data
+    synthetic_ab_combos = set(zip(samples['A'], samples['B']))
+    synthetic_cd_combos = set(zip(samples['C'], samples['D']))
+    synthetic_ac_combos = set(zip(samples['A'], samples['C']))
+
+    # Assert combinations match
+    assert original_ab_combos == synthetic_ab_combos
+    assert original_cd_combos == synthetic_cd_combos
+    assert original_ac_combos == synthetic_ac_combos
+
+
+def test_fixed_combinations_multiple_patterns_three_patterns_reject_sampling():
+    """Test that FixedCombinations pattern works with multiple patterns.
+
+    Test that when the second pattern in the chain fails, the third pattern still works.
+    """
+    # Setup
+    data = pd.DataFrame({
+        'A': [1, 2, 3, 1, 2, 1],
+        'B': [10, 20, 30, 10, 20, 10],
+        'C': [100, 200, 300, 100, 200, 100],
+        'D': [1000, 2000, 3000, 1000, 2000, 1000],
+    })
+    metadata = Metadata.load_from_dict({
+        'columns': {
+            'A': {'sdtype': 'categorical'},
+            'B': {'sdtype': 'categorical'},
+            'C': {'sdtype': 'categorical'},
+            'D': {'sdtype': 'categorical'},
+        }
+    })
+    pattern1 = FixedCombinations(['A', 'B'])
+    pattern2 = FixedCombinations(['C', 'D'])
+    pattern3 = FixedCombinations(['A', 'C'])
+
+    # Run
+    synthesizer = GaussianCopulaSynthesizer(metadata)
+    synthesizer.add_cag(patterns=[pattern1, pattern3, pattern2])
+    synthesizer.fit(data)
+    samples = synthesizer.sample(100)
+    updated_metadata = synthesizer.get_metadata('modified')
+    original_metadata = synthesizer.get_metadata('original')
+
+    # Assert
+    expected_updated_metadata = Metadata.load_from_dict({
+        'columns': {
+            'A#B': {'sdtype': 'categorical'},
+            'C#D': {'sdtype': 'categorical'},
+        }
+    }).to_dict()
+
+    assert expected_updated_metadata == updated_metadata.to_dict()
+
+    assert original_metadata.to_dict() == metadata.to_dict()
+
+    # Get unique combinations from original data
+    original_ab_combos = set(zip(data['A'], data['B']))
+    original_cd_combos = set(zip(data['C'], data['D']))
+    original_ac_combos = set(zip(data['A'], data['C']))
+
+    # Get unique combinations from synthetic data
+    synthetic_ab_combos = set(zip(samples['A'], samples['B']))
+    synthetic_cd_combos = set(zip(samples['C'], samples['D']))
+    synthetic_ac_combos = set(zip(samples['A'], samples['C']))
+
+    # Assert combinations match
+    assert original_ab_combos == synthetic_ab_combos
+    assert original_cd_combos == synthetic_cd_combos
+    assert original_ac_combos == synthetic_ac_combos