Update test_quality.py

Author: yoid2000

tests/test_quality.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 from pytest import approx
-from scipy.stats import pearsonr, chi2_contingency, ks_2samp
+from scipy.stats import chi2_contingency, ks_2samp, pearsonr
 from sklearn.metrics import mutual_info_score
 
 from syndiffix import Synthesizer
@@ -12,23 +12,20 @@
 def test_quality_float_str() -> None:
     """Test quality of synthetic data for float-string columns with dependency."""
     np.random.seed(42)
-    
+
     num_rows = 1000
     # Create float column with normal distribution
     float_col = np.random.normal(50, 15, num_rows)
 
     # Create categorical string column with dependency on float values
     # Higher float values more likely to be in categories 6-9
     str_categories = [f"cat_{i}" for i in range(10)]
-    
+
     # Create dependency: map float values to categories with some noise
     normalized_float = (float_col - float_col.min()) / (float_col.max() - float_col.min())
-    category_indices = np.clip(
-        (normalized_float * 8 + np.random.normal(0, 1, num_rows)).astype(int), 
-        0, 9
-    )
+    category_indices = np.clip((normalized_float * 8 + np.random.normal(0, 1, num_rows)).astype(int), 0, 9)
     str_col = [str_categories[i] for i in category_indices]
-    
+
     df = pd.DataFrame({"float_col": float_col, "str_col": str_col})
     # Ensure at least 20 instances of every category in str_col
     min_count = 20
@@ -44,103 +41,102 @@ def test_quality_float_str() -> None:
                 df = pd.concat([df, replicated], ignore_index=True)
             else:
                 # If no rows exist, create new rows with random float_col
-                new_rows = pd.DataFrame({
-                    "float_col": np.random.normal(50, 15, min_count),
-                    "str_col": [cat]*min_count
-                })
+                new_rows = pd.DataFrame(
+                    {"float_col": np.random.normal(50, 15, min_count), "str_col": [cat] * min_count}
+                )
                 df = pd.concat([df, new_rows], ignore_index=True)
-    
+
     # Generate synthetic data
     df_syn = Synthesizer(df, anonymization_params=NOISELESS_PARAMS).sample()
-    
+
     # 1. Single-column similarity tests
-    
+
     # Float column: test distribution similarity using KS test
     ks_stat, ks_pvalue = ks_2samp(df["float_col"], df_syn["float_col"])
     assert ks_stat < 0.15, f"Float column distributions too different (KS statistic: {ks_stat})"
-    
+
     # Float column: test mean and std similarity
     assert df_syn["float_col"].mean() == approx(df["float_col"].mean(), abs=1.0)
     assert df_syn["float_col"].std() == approx(df["float_col"].std(), abs=0.5)
-    
+
     # String column: test category distribution similarity
     orig_counts = df["str_col"].value_counts().sort_index()
     syn_counts = df_syn["str_col"].value_counts().sort_index()
-    
+
     # Ensure all original categories are present in synthetic data
     for cat in orig_counts.index:
         assert cat in syn_counts.index, f"Category {cat} missing from synthetic data"
-    
+
     # Test count similarity (allowing for some variation)
     for cat in orig_counts.index:
         orig_pct = orig_counts[cat] / len(df)
         syn_pct = syn_counts[cat] / len(df_syn)
         assert syn_pct == approx(orig_pct, abs=0.01), f"Category {cat} proportion differs too much"
-    
+
     # 2. Dependency/correlation tests
-    
+
     # For float-string dependency, use mutual information
     # Discretize float column for mutual information calculation
     orig_float_binned = pd.cut(df["float_col"], bins=10, labels=False)
     syn_float_binned = pd.cut(df_syn["float_col"], bins=10, labels=False)
-    
+
     orig_str_encoded = pd.Categorical(df["str_col"]).codes
     syn_str_encoded = pd.Categorical(df_syn["str_col"]).codes
-    
+
     orig_mi = mutual_info_score(orig_float_binned, orig_str_encoded)
     syn_mi = mutual_info_score(syn_float_binned, syn_str_encoded)
-    
+
     assert syn_mi == approx(orig_mi, rel=0.4), f"Mutual information differs too much: orig={orig_mi}, syn={syn_mi}"
 
 
 def test_quality_float_float() -> None:
     """Test quality of synthetic data for float-float columns with correlation."""
     np.random.seed(42)
-    
+
     # Create correlated float columns
     x = np.random.normal(0, 1, 1000)
     y = 0.7 * x + np.random.normal(0, 0.5, 1000)  # Correlation ~0.8
-    
+
     df = pd.DataFrame({"x": x, "y": y})
-    
+
     # Generate synthetic data
     df_syn = Synthesizer(df, anonymization_params=NOISELESS_PARAMS).sample()
-    
+
     # 1. Single-column similarity tests
-    
+
     # X column
     ks_stat_x, _ = ks_2samp(df["x"], df_syn["x"])
     assert ks_stat_x < 0.15, f"X column distributions too different (KS statistic: {ks_stat_x})"
     assert df_syn["x"].mean() == approx(df["x"].mean(), abs=0.02)
     assert df_syn["x"].std() == approx(df["x"].std(), abs=0.04)
-    
+
     # Y column
     ks_stat_y, _ = ks_2samp(df["y"], df_syn["y"])
     assert ks_stat_y < 0.15, f"Y column distributions too different (KS statistic: {ks_stat_y})"
     assert df_syn["y"].mean() == approx(df["y"].mean(), abs=0.02)
     assert df_syn["y"].std() == approx(df["y"].std(), abs=0.04)
-    
+
     # 2. Correlation tests
-    
+
     orig_corr, _ = pearsonr(df["x"], df["y"])
     syn_corr, _ = pearsonr(df_syn["x"], df_syn["y"])
-    
+
     assert syn_corr == approx(orig_corr, abs=0.01), f"Correlation differs too much: orig={orig_corr}, syn={syn_corr}"
 
 
 def test_quality_str_str() -> None:
     """Test quality of synthetic data for string-string columns with dependency."""
     np.random.seed(42)
-    
+
     # Create two categorical columns with dependency
     categories_a = [f"group_{i}" for i in range(10)]
     categories_b = [f"type_{i}" for i in range(10)]
-    
+
     # Create dependency: certain groups prefer certain types
     group_prefs = {
-        0: [0, 1, 2],      # group_0 prefers type_0, type_1, type_2
-        1: [1, 2, 3],      # group_1 prefers type_1, type_2, type_3
-        2: [2, 3, 4],      # etc.
+        0: [0, 1, 2],  # group_0 prefers type_0, type_1, type_2
+        1: [1, 2, 3],  # group_1 prefers type_1, type_2, type_3
+        2: [2, 3, 4],  # etc.
         3: [3, 4, 5],
         4: [4, 5, 6],
         5: [5, 6, 7],
@@ -149,64 +145,64 @@ def test_quality_str_str() -> None:
         8: [8, 9, 0],
         9: [9, 0, 1],
     }
-    
+
     # Generate data with dependency
     col_a = []
     col_b = []
-    
+
     for _ in range(1000):
         # Choose group randomly
         group_idx = np.random.randint(0, 10)
         group = categories_a[group_idx]
-        
+
         # Choose type based on group preference (80% of time) or random (20% of time)
         if np.random.random() < 0.8:
             type_idx = np.random.choice(group_prefs[group_idx])
         else:
             type_idx = np.random.randint(0, 10)
-        
+
         type_val = categories_b[type_idx]
-        
+
         col_a.append(group)
         col_b.append(type_val)
-    
+
     df = pd.DataFrame({"group": col_a, "type": col_b})
-    
+
     # Generate synthetic data
     df_syn = Synthesizer(df, anonymization_params=NOISELESS_PARAMS).sample()
-    
+
     # 1. Single-column similarity tests
-    
+
     # Group column
     orig_group_counts = df["group"].value_counts().sort_index()
     syn_group_counts = df_syn["group"].value_counts().sort_index()
-    
+
     for group in orig_group_counts.index:
         assert group in syn_group_counts.index, f"Group {group} missing from synthetic data"
         orig_pct = orig_group_counts[group] / len(df)
         syn_pct = syn_group_counts[group] / len(df_syn)
         assert syn_pct == approx(orig_pct, abs=0.01), f"Group {group} proportion differs too much"
-    
+
     # Type column
     orig_type_counts = df["type"].value_counts().sort_index()
     syn_type_counts = df_syn["type"].value_counts().sort_index()
-    
+
     for type_val in orig_type_counts.index:
         assert type_val in syn_type_counts.index, f"Type {type_val} missing from synthetic data"
         orig_pct = orig_type_counts[type_val] / len(df)
         syn_pct = syn_type_counts[type_val] / len(df_syn)
         assert syn_pct == approx(orig_pct, abs=0.005), f"Type {type_val} proportion differs too much"
-    
+
     # 2. Dependency tests using contingency table analysis
-    
+
     # Create contingency tables
     orig_contingency = pd.crosstab(df["group"], df["type"])
     syn_contingency = pd.crosstab(df_syn["group"], df_syn["type"])
-    
+
     # Ensure synthetic contingency table has same shape as original
     assert orig_contingency.shape[0] <= syn_contingency.shape[0], "Missing groups in synthetic data"
     assert orig_contingency.shape[1] <= syn_contingency.shape[1], "Missing types in synthetic data"
-    
+
     # Calculate Cramér's V (measure of association between categorical variables)
     def cramers_v(contingency_table):
         chi2, _, _, _ = chi2_contingency(contingency_table)
@@ -215,38 +211,37 @@ def cramers_v(contingency_table):
         if min_dim == 0:
             return 0
         return np.sqrt(chi2 / (n * min_dim))
-    
+
     # Align contingency tables for comparison
     common_groups = set(orig_contingency.index) & set(syn_contingency.index)
     common_types = set(orig_contingency.columns) & set(syn_contingency.columns)
-    
+
     orig_aligned = orig_contingency.loc[list(common_groups), list(common_types)]
     syn_aligned = syn_contingency.loc[list(common_groups), list(common_types)]
-    
+
     orig_cramers_v = cramers_v(orig_aligned)
     syn_cramers_v = cramers_v(syn_aligned)
-    
-    assert syn_cramers_v == approx(orig_cramers_v, rel=0.4), f"Cramér's V differs too much: orig={orig_cramers_v}, syn={syn_cramers_v}"
+
+    assert syn_cramers_v == approx(
+        orig_cramers_v, rel=0.4
+    ), f"Cramér's V differs too much: orig={orig_cramers_v}, syn={syn_cramers_v}"
 
 
 def test_synthetic_data_size_consistency() -> None:
     """Test that synthetic data has reasonable size compared to original."""
     np.random.seed(42)
-    
+
     # Create simple test data
-    df = pd.DataFrame({
-        "col1": np.random.normal(0, 1, 1000),
-        "col2": np.random.choice(["A", "B", "C"], 1000)
-    })
-    
+    df = pd.DataFrame({"col1": np.random.normal(0, 1, 1000), "col2": np.random.choice(["A", "B", "C"], 1000)})
+
     df_syn = Synthesizer(df, anonymization_params=NOISELESS_PARAMS).sample()
-    
+
     # Synthetic data should have reasonable size (within 50% of original)
     assert len(df_syn) > 0.5 * len(df), "Synthetic data too small"
     assert len(df_syn) < 2.0 * len(df), "Synthetic data too large"
-    
+
     # Should have same number of columns
     assert len(df_syn.columns) == len(df.columns), "Different number of columns"
-    
+
     # Should have same column names
     assert list(df_syn.columns) == list(df.columns), "Different column names"