class TestPytorchDPSynthesizer_DPGAN:
def setup(self):
self.dpgan = PytorchDPSynthesizer(DPGAN(), GeneralTransformer())
def test_fit(self):
self.dpgan.fit(df)
assert self.dpgan.gan.generator
def test_sample(self):
self.dpgan.fit(df)
sample_size = len(df)
synth_data = self.dpgan.sample(sample_size)
assert synth_data.shape == df.shape
class TestPytorchDPSynthesizer_PATECTGAN:
def setup(self):
self.patectgan = PytorchDPSynthesizer(PATECTGAN(), None)
def test_fit(self):
self.patectgan.fit(
df, categorical_columns=['sex', 'educ', 'race', 'married'])
assert self.patectgan.gan.generator
def test_sample(self):
self.patectgan.fit(
df, categorical_columns=['sex', 'educ', 'race', 'married'])
sample_size = len(df)
synth_data = self.patectgan.sample(sample_size)
assert synth_data.shape == df.shape