def test_save_load():
data = pd.DataFrame({
'continuous': np.random.random(100),
'discrete': np.random.choice(['a', 'b'], 100)
})
discrete_columns = ['discrete']
tvae = TVAESynthesizer(epochs=10)
tvae.fit(data, discrete_columns)
with tf.TemporaryDirectory() as temporary_directory:
tvae.save(temporary_directory + "test_tvae.pkl")
tvae = TVAESynthesizer.load(temporary_directory + "test_tvae.pkl")
sampled = tvae.sample(1000)
assert set(sampled.columns) == {'continuous', 'discrete'}
assert set(sampled['discrete'].unique()) == {'a', 'b'}
def test_tvae(tmpdir):
iris = datasets.load_iris()
data = pd.DataFrame(iris.data, columns=iris.feature_names)
data['class'] = pd.Series(iris.target).map(iris.target_names.__getitem__)
tvae = TVAESynthesizer(epochs=10)
tvae.fit(data, ['class'])
path = str(tmpdir / 'test_tvae.pkl')
tvae.save(path)
tvae = TVAESynthesizer.load(path)
sampled = tvae.sample(100)
assert sampled.shape == (100, 5)
assert isinstance(sampled, pd.DataFrame)
assert set(sampled.columns) == set(data.columns)
assert set(sampled.dtypes) == set(data.dtypes)