POTNet is a deep generative model for generating mixed-type tabular data (continuous/discrete/categorical) based on the marginally-penalized Wasserstein loss.
This repository contains an implementation of the POTNet model as described in the paper: Efficient Generative Modeling via Penalized Optimal Transport Network.
Python version: 3.9.13
Intall dependences:
pip install -r requirements.txtTo run the following package from command line, execute the following lines:
conda create -n env_name python=3.9 -n
conda activate env_name
pip install -r requirements.txt
pip install .
python3 demo.py-
Input: When using POTNet, continuous and discrete features should be represented as
float. Categorical columns can be represented as eitherstr,int, orfloat. -
Categorical features: You will need to specify which columns correspond to categorical features via the
categorical_colsargument using either column names (e.g.['SEX', 'EDUCATION']) or column indices ([1, 2]). -
Numeric output:
- If all numeric features are discrete, then specify
numeric_output_data_type = "integer". - If numeric features consist of both discrete and continuous types, then you need to specify each numeric feature type using a
dict:
{'integer': [col1, col2], 'continuous': [col3, col4]}- By default, all numeric features are assumed to be continuous. You can also manually specify this via
numeric_output_data_type = "continuous"
- If all numeric features are discrete, then specify
-
Conditional POTNet: You can condition on a selection of numeric features by setting
conditional = Truein model initialization, and subsequently passing the features you wish to condition on viaconditioning_datato POTNet during the fitting stage. For a concrete example, please see./example/demo_lfi.ipynbunder section Conditional POTNet. -
Missing values: The data should not contain any missing values.
We provide two examples for using POTNet, both deposited in the ./examples folder:
-
Likelihood-free inference (
./examples/demo_lfi.ipynb): We generate 3,000 samples consisting of all continous features from the model-
$\theta_i \sim \mathrm{Unif}[-3, 3]$ for$i = 1, \dots, 5$ $\mu = (\theta_1, \theta_2)$ $\sigma_1 = \theta_3^2$ $\sigma_2 = \theta_4^2$ $\rho = \tanh(\theta_5)$ $\Sigma = ((\sigma_1^2, ~\rho \sigma_1 \sigma_2 ), (\rho \sigma_1 \sigma_2, ~\sigma_2^2))$ -
$X_j \sim \mathcal{N}(\mu, \Sigma)$ for$j = 1, \dots, 4$
-
-
Default of credit card clients (
./examples/demo_credit.ipynb): For an example illustrating usage of POTNet for mixed-data types, We subsampled 1,000 samples from the credit default dataset, deposited in./data/credit_card.csv.
Below, we provide a simple template for using POTNet:
from potnet import *
# load data
data = ...
# specify categorical columns
cat_cols = [col1, col2, col3]
# initialize POTNet
potnet_model= POTNet(embedding_dim=data.shape[1],
categorical_cols=cat_cols,
numeric_output_data_type = 'continuous', # continuous data
epochs=500,
batch_size=256,
save_checkpoint=True, # save checkpoint periodically
checkpoint_epoch = 50 # save checkpoint every 50 epochs
)
# fit POTNet
potnet_model.fit(data)
# generate 1000 synthetic samples
gen_data = potnet_model.generate(1000)
# save model
potnet_model.save('potnet_model.pt')If you use POTNet in your research or project, please cite the following paper:
@article{lu2024generative,
title={Efficient Generative Modeling via Penalized Optimal Transport Network},
author={Lu, Wenhui Sophia and Zhong, Chenyang and Wong, Wing Hung},
journal={arXiv preprint arXiv:2402.10456},
year={2025}
}