Merge pull request #20 from kevinykuo/feature/sampling-freq

Expose `log_frequency` parameter for conditional sampling

Author: leix28

ctgan/conditional.py

@@ -2,7 +2,7 @@
 
 
 class ConditionalGenerator(object):
-    def __init__(self, data, output_info):
+    def __init__(self, data, output_info, log_frequency):
         self.model = []
 
         start = 0
@@ -50,7 +50,8 @@ def __init__(self, data, output_info):
                     continue
                 end = start + item[0]
                 tmp = np.sum(data[:, start:end], axis=0)
-                tmp = np.log(tmp + 1)
+                if log_frequency:
+                    tmp = np.log(tmp + 1)
                 tmp = tmp / np.sum(tmp)
                 self.p[self.n_col, :item[0]] = tmp
                 self.interval.append((self.n_opt, item[0]))

ctgan/synthesizer.py

@@ -95,7 +95,7 @@ def _cond_loss(self, data, c, m):
 
         return (loss * m).sum() / data.size()[0]
 
-    def fit(self, train_data, discrete_columns=tuple(), epochs=300):
+    def fit(self, train_data, discrete_columns=tuple(), epochs=300, log_frequency=True):
         """Fit the CTGAN Synthesizer models to the training data.
 
         Args:
@@ -109,6 +109,9 @@ def fit(self, train_data, discrete_columns=tuple(), epochs=300):
                 a ``pandas.DataFrame``, this list should contain the column names.
             epochs (int):
                 Number of training epochs. Defaults to 300.
+            log_frequency (boolean):
+                Whether to use log frequency of categorical levels in conditional
+                sampling. Defaults to ``True``.
         """
 
         self.transformer = DataTransformer()
@@ -118,7 +121,11 @@ def fit(self, train_data, discrete_columns=tuple(), epochs=300):
         data_sampler = Sampler(train_data, self.transformer.output_info)
 
         data_dim = self.transformer.output_dimensions
-        self.cond_generator = ConditionalGenerator(train_data, self.transformer.output_info)
+        self.cond_generator = ConditionalGenerator(
+            train_data,
+            self.transformer.output_info,
+            log_frequency
+        )
 
         self.generator = Generator(
             self.embedding_dim + self.cond_generator.n_opt,

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