add test

Author: kevinykuo

tests/integration/test_ctgan.py

@@ -48,3 +48,26 @@ def test_ctgan_numpy():
     assert sampled.shape == (100, 2)
     assert isinstance(sampled, np.ndarray)
     assert set(np.unique(sampled[:, 1])) == {'a', 'b', 'c'}
+
+
+def test_log_frequency():
+    data = pd.DataFrame({
+        'continuous': np.random.random(1000),
+        'discrete': np.random.choice(['a', 'b', 'c'], 1000, p=[0.95, 0.025, 0.025])
+    })
+
+    discrete_columns = ['discrete']
+
+    ctgan = CTGANSynthesizer()
+    ctgan.fit(data, discrete_columns, epochs=100)
+
+    sampled = ctgan.sample(1000)
+    counts = sampled['discrete'].value_counts()
+    assert counts['a'] < 650
+
+    ctgan = CTGANSynthesizer()
+    ctgan.fit(data, discrete_columns, epochs=100, log_frequency=False)
+
+    sampled = ctgan.sample(1000)
+    counts = sampled['discrete'].value_counts()
+    assert counts['a'] > 900