def main():
args = _parse_args()
if args.tsv:
data, discrete_columns = read_tsv(args.data, args.metadata)
else:
data, discrete_columns = read_csv(
args.data, args.metadata, args.header, args.discrete
)
if args.load:
model = CTGANSynthesizer.load(args.load)
else:
model = CTGANSynthesizer()
model.fit(data, discrete_columns, args.epochs)
if args.save is not None:
model.save(args.save)
num_samples = args.num_samples or len(data)
if args.sample_condition_column is not None:
assert args.sample_condition_column_value is not None
sampled = model.sample(
num_samples, args.sample_condition_column, args.sample_condition_column_value
)
if args.tsv:
write_tsv(sampled, args.metadata, args.output)
else:
sampled.to_csv(args.output, index=False)
def test_synthesizer_sample():
data = pd.DataFrame({"discrete": np.random.choice(["a", "b", "c"], 100)})
discrete_columns = ["discrete"]
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=1)
samples = ctgan.sample(1000, "discrete", "a")
assert isinstance(samples, pd.DataFrame)
def test_synthesizer_sample():
data = pd.DataFrame({'discrete': np.random.choice(['a', 'b', 'c'], 100)})
discrete_columns = ['discrete']
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=1)
samples = ctgan.sample(1000, 'discrete', 'a')
assert isinstance(samples, pd.DataFrame)
def test_log_frequency():
data = pd.DataFrame({
"continuous": np.random.random(1000),
"discrete": np.repeat(["a", "b", "c"], [950, 25, 25]),
})
discrete_columns = ["discrete"]
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=100)
sampled = ctgan.sample(10000)
counts = sampled["discrete"].value_counts()
assert counts["a"] < 6500
ctgan = CTGANSynthesizer(log_frequency=False)
ctgan.fit(data, discrete_columns, epochs=100)
sampled = ctgan.sample(10000)
counts = sampled["discrete"].value_counts()
assert counts["a"] > 9000
def test_log_frequency():
data = pd.DataFrame({
'continuous': np.random.random(1000),
'discrete': np.repeat(['a', 'b', 'c'], [950, 25, 25])
})
discrete_columns = ['discrete']
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=100)
sampled = ctgan.sample(10000)
counts = sampled['discrete'].value_counts()
assert counts['a'] < 6500
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=100, log_frequency=False)
sampled = ctgan.sample(10000)
counts = sampled['discrete'].value_counts()
assert counts['a'] > 9000
def test_sample(self):
np.random.seed(0)
tf.random.set_seed(0)
data, discrete = generate_data(self._vars['batch_size'])
model = CTGANSynthesizer(
batch_size=self._vars['batch_size'], pac=self._vars['pac'])
self.assertIsNotNone(model)
model.train(data, discrete, epochs=1)
output = model.sample(self._n_samples).values
expected_output = np.array([[0.4139329, 3.0]])
np.testing.assert_almost_equal(
output, expected_output, decimal=self._vars['decimal'])
def test_ctgan_numpy():
data = pd.DataFrame({
"continuous": np.random.random(100),
"discrete": np.random.choice(["a", "b", "c"], 100),
})
discrete_columns = [1]
ctgan = CTGANSynthesizer()
ctgan.fit(data.values, discrete_columns, epochs=1)
sampled = ctgan.sample(100)
assert sampled.shape == (100, 2)
assert isinstance(sampled, np.ndarray)
assert set(np.unique(sampled[:, 1])) == {"a", "b", "c"}
def test_ctgan_dataframe():
data = pd.DataFrame({
"continuous": np.random.random(100),
"discrete": np.random.choice(["a", "b", "c"], 100),
})
discrete_columns = ["discrete"]
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=1)
sampled = ctgan.sample(100)
assert sampled.shape == (100, 2)
assert isinstance(sampled, pd.DataFrame)
assert set(sampled.columns) == {"continuous", "discrete"}
assert set(sampled["discrete"].unique()) == {"a", "b", "c"}
def test_ctgan_dataframe():
data = pd.DataFrame({
'continuous': np.random.random(100),
'discrete': np.random.choice(['a', 'b', 'c'], 100)
})
discrete_columns = ['discrete']
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=1)
sampled = ctgan.sample(100)
assert sampled.shape == (100, 2)
assert isinstance(sampled, pd.DataFrame)
assert set(sampled.columns) == {'continuous', 'discrete'}
assert set(sampled['discrete'].unique()) == {'a', 'b', 'c'}
def main():
args = _parse_args()
if args.tsv:
data, discrete_columns = read_tsv(args.data, args.metadata)
else:
data, discrete_columns = read_csv(args.data, args.metadata, args.header, args.discrete)
model = CTGANSynthesizer()
model.fit(data, discrete_columns, args.epochs)
num_samples = args.num_samples or len(data)
sampled = model.sample(num_samples)
if args.tsv:
write_tsv(sampled, args.metadata, args.output)
else:
sampled.to_csv(args.output, index=False)
def test_categorical_nan():
data = pd.DataFrame({
"continuous": np.random.random(30),
# This must be a list (not a np.array) or NaN will be cast to a string.
"discrete": [np.nan, "b", "c"] * 10,
})
discrete_columns = ["discrete"]
ctgan = CTGANSynthesizer()
ctgan.fit(data, discrete_columns, epochs=1)
sampled = ctgan.sample(100)
assert sampled.shape == (100, 2)
assert isinstance(sampled, pd.DataFrame)
assert set(sampled.columns) == {"continuous", "discrete"}
# since np.nan != np.nan, we need to be careful here
values = set(sampled["discrete"].unique())
assert len(values) == 3
assert any(pd.isnull(x) for x in values)
assert {"b", "c"}.issubset(values)
