expose log_frequency param

Author: kevinykuo

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,