DataTransformer init parameters

ctgan/data_transformer.py

@@ -19,20 +19,27 @@ class DataTransformer(object):
     Discrete columns are encoded using a scikit-learn OneHotEncoder.
     """
 
-    def __init__(self, max_clusters=10, weight_threshold=0.005):
+    def __init__(self, max_clusters=10, weight_threshold=0.005, max_gm_samples=None):
         """Create a data transformer.
 
         Args:
             max_clusters (int):
                 Maximum number of Gaussian distributions in Bayesian GMM.
             weight_threshold (float):
                 Weight threshold for a Gaussian distribution to be kept.
+            max_gm_samples (int):
+                Maximum number of samples to use during GMM fit.
         """
         self._max_clusters = max_clusters
         self._weight_threshold = weight_threshold
+        self._max_gm_samples = np.inf if max_gm_samples is None else max_gm_samples
 
     def _fit_continuous(self, column_name, raw_column_data):
         """Train Bayesian GMM for continuous column."""
+        if self._max_gm_samples <= raw_column_data.shape[0]:
+            raw_column_data = np.random.choice(raw_column_data,
+                                               size=self._max_gm_samples,
+                                               replace=False)
         gm = BayesianGaussianMixture(
             self._max_clusters,
             weight_concentration_prior_type='dirichlet_process',

ctgan/synthesizers/ctgan.py

@@ -267,7 +267,8 @@ def _validate_discrete_columns(self, train_data, discrete_columns):
         if invalid_columns:
             raise ValueError('Invalid columns found: {}'.format(invalid_columns))
 
-    def fit(self, train_data, discrete_columns=tuple(), epochs=None):
+    def fit(self, train_data, discrete_columns=tuple(), epochs=None,
+            data_transformer_params={}):
         """Fit the CTGAN Synthesizer models to the training data.
 
         Args:
@@ -278,6 +279,8 @@ def fit(self, train_data, discrete_columns=tuple(), epochs=None):
                 Vector. If ``train_data`` is a Numpy array, this list should
                 contain the integer indices of the columns. Otherwise, if it is
                 a ``pandas.DataFrame``, this list should contain the column names.
+            data_transformer_params (dict):
+                Dictionary of parameters for ``DataTransformer`` initialization.
         """
         self._validate_discrete_columns(train_data, discrete_columns)
 
@@ -290,7 +293,7 @@ def fit(self, train_data, discrete_columns=tuple(), epochs=None):
                 DeprecationWarning
             )
 
-        self._transformer = DataTransformer()
+        self._transformer = DataTransformer(**data_transformer_params)
         self._transformer.fit(train_data, discrete_columns)
 
         train_data = self._transformer.transform(train_data)

tests/integration/test_ctgan.py

@@ -184,3 +184,14 @@ def test_wrong_sampling_conditions():
 
     with pytest.raises(ValueError):
         ctgan.sample(1, 'discrete', "d")
+
+
+def test_ctgan_data_transformer_params():
+    data = pd.DataFrame({
+        'continuous': np.random.random(1000)
+    })
+
+    ctgan = CTGANSynthesizer(epochs=1)
+    ctgan.fit(data, [], data_transformer_params={'max_gm_samples': 100})
+
+    assert ctgan._transformer._max_gm_samples == 100