VAE.py

# -*- coding: utf-8 -*-
"""VAE PyTorch.ipynb
Automatically generated by Colaboratory.
Original file is located at
    https://colab.research.google.com/drive/1484hwapatM9DqxsXJA-bWSPn_KJX49HW
"""
import matplotlib.pyplot as plt
import seaborn as sns
import torch
import torch.nn as nn
from scipy.stats import entropy
from torch.utils.data import Dataset, DataLoader
from sklearn.manifold import TSNE
from scipy.stats import wasserstein_distance
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import pandas as pd
import os
from sklearn.decomposition import PCA
import random
import sys
import glob

torch.manual_seed(0)
random.seed(0)
np.random.seed(0)
original = sys.stdout


"""# Dataset"""
class CustomDataset(Dataset):
    def __init__(self, actual, shuffle=True, scale=False):
        if isinstance(actual, str):
            df = pd.read_csv(actual)
            self.df = df
        elif isinstance(actual, np.ndarray):
            self.df = pd.DataFrame(actual)
        else:
            self.df = actual

        self.len = self.df.shape[0]
        self.df = self.df.values

    def __len__(self):
        return self.len

    def __getitem__(self, idx):
        return self.df[idx, :]


# Does mean and variance need to be shared in the same layer ?

""" Model"""
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
""" Encoder"""


class Encoder(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(Encoder, self).__init__()
        self.mean = nn.Sequential(
            nn.Linear(input_dim, latent_dim),
            nn.BatchNorm1d(latent_dim),
            nn.ReLU(inplace=True),
            nn.Linear(latent_dim, latent_dim),
            nn.BatchNorm1d(latent_dim),
            nn.ReLU(inplace=True),
            nn.Linear(latent_dim, latent_dim),
            nn.BatchNorm1d(latent_dim),
            nn.ReLU(inplace=True),
        )
        self.var = nn.Sequential(
            nn.Linear(input_dim, latent_dim),
            nn.BatchNorm1d(latent_dim),
            nn.ReLU(inplace=True),
            nn.Linear(latent_dim, latent_dim),
            nn.BatchNorm1d(latent_dim),
            nn.ReLU(inplace=True),
            nn.Linear(latent_dim, latent_dim),
            nn.BatchNorm1d(latent_dim),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        mean = self.mean(x)
        std = self.var(x)
        return mean, std


"""## Decoder"""


class Decoder(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(Decoder, self).__init__()
        self.layer = nn.Sequential(
            nn.Linear(latent_dim, latent_dim),
            nn.ReLU(inplace=True),
            nn.Linear(latent_dim, latent_dim),
            nn.ReLU(inplace=True),
            nn.Linear(latent_dim, input_dim),
            nn.Sigmoid(),
        )

    def forward(self, x):
        return self.layer(x)


"""## Connect Encoder and Decoder"""


class VAE(nn.Module):
    def __init__(self, input_dim, latent_dim):
        super(VAE, self).__init__()
        global device
        self.enc = Encoder(input_dim, latent_dim).to(device)
        self.dec = Decoder(input_dim, latent_dim).to(device)

    def forward(self, x):
        z_mean, z_var = self.enc(x)
        std = torch.exp(z_var / 2)
        eps = torch.randn_like(std)
        sample = eps.mul(std).add(z_mean)
        gen = self.dec(sample)
        return gen, z_mean, z_var


"""# Training Loop"""


def prep(data):
    if "Unnamed: 0" in data.columns.values:
        data.drop(["Unnamed: 0"], axis=1, inplace=True)
    if "Unnamed: 0.1" in data.columns.values:
        data.drop(["Unnamed: 0.1"], axis=1, inplace=True)
    if "Unnamed: 0.1.1" in data.columns.values:
        data.drop(["Unnamed: 0.1.1"], axis=1, inplace=True)
    return data


def js_divergence(p, q):
    m = 0.5 * (p + q)
    return 0.5 * (entropy(p, m) + entropy(q, m))


def wstd(org, gen):
    org = prep(org)
    gen = prep(gen)
    wsd = 0
    orig_data = org.values
    gen_data = gen.values
    # if orig_data.shape[1] != gen_data.shape[1]:
    #    pca = PCA(n_components=gen_data.shape[1])
    #    pca = pca.fit(orig_data)
    #    gen_data = pca.inverse_transform(gen_data)
    columns = orig_data.shape[1]
    w = 0
    print(orig_data.shape, gen_data.shape, " Test WSTD Function")
    for i in range(columns):
        orig_data_ = orig_data[:, i].reshape(-1, 1)
        gen_data_ = gen_data[:, i].reshape(-1, 1)
        orig_data_ = orig_data_.reshape(
            len(orig_data),
        )
        gen_data_ = gen_data_.reshape(
            len(gen_data),
        )
        # orig_data_ += 1
        # gen_data_ += 1
        w += wasserstein_distance(orig_data_, gen_data_)
        # break
    w /= columns
    wsd += w
    return wsd


def calculate_jsd(org, gen):
    org = prep(org)
    gen = prep(gen)
    KLD = 0
    orig_data = org.values
    gen_data = gen.values
    # if orig_data.shape[1] != gen_data.shape[1]:
    #    pca = PCA(n_components=gen_data.shape[1])
    #    pca = pca.fit(orig_data)
    #    gen_data = pca.inverse_transform(gen_data)
    columns = orig_data.shape[1]
    scaler = MinMaxScaler()
    softmax = torch.nn.Softmax(0)
    kld = 0
    print(orig_data.shape, gen_data.shape, " Test JSD Function")
    for i in range(columns):
        orig_data_ = scaler.fit_transform(orig_data[:, i].reshape(-1, 1))
        gen_data_ = scaler.fit_transform(gen_data[:, i].reshape(-1, 1))
        orig_data_ = orig_data_.reshape(
            len(orig_data),
        )
        gen_data_ = gen_data_.reshape(
            len(gen_data),
        )
        orig_data_ = softmax(torch.from_numpy(orig_data_)).numpy()
        gen_data_ = softmax(torch.from_numpy(gen_data_)).numpy()
        # orig_data_ += 1
        # gen_data_ += 1
        kld += js_divergence(orig_data_, gen_data_)
        # break
    kld /= columns
    KLD += kld
    return KLD


def train(original_dim, latent_dim, learning_rate, epochs, dataloader, rawdataset):
    global device
    vae = VAE(original_dim, latent_dim).to(device)
    optimizer = torch.optim.Adam(vae.parameters(), lr=learning_rate)
    losslist = []
    allowed_epochs = [1, 10, 100, 500, epochs]
    for epoch in range(epochs):
        running_loss = 0
        for _, data in enumerate(dataloader):
            data = data.float().to(device)
            optimizer.zero_grad()
            sample, z_mean, z_var = vae(data)
            reconstruction_loss = nn.functional.binary_cross_entropy(sample, data)
            kl_loss = 0.0001 * (torch.sum(torch.exp(z_var) + z_mean ** 2 - 1.0 - z_var))
            loss = reconstruction_loss + kl_loss
            loss.backward()
            running_loss += loss.item()
            optimizer.step()
        losslist.append(running_loss / len(rawdataset))
        if epoch % 500 == 499:
            print(
                f"[{epoch+1}/{epochs}] Done. Average loss = {sum(losslist)/len(losslist)}"
            )

    return vae, losslist


def generate(decoder, numsamples, latent_dim, doMinMax, doPca, pca, scaler, cols):
    global device
    noise = np.random.randn(numsamples, latent_dim)
    noise = torch.from_numpy(noise).float().to(device)
    decoder.eval()
    with torch.no_grad():
        gen = decoder(noise).cpu().numpy()
    if doPca:
        gen = pca.inverse_transform(gen)

    if doMinMax:
        gen = scaler.inverse_transform(gen)
    gen = pd.DataFrame(gen, columns=cols)
    return gen


def plot(plotpath, losslist, n):
    x = list(range(1, len(losslist) + 1))
    plt.xlabel("Epochs")
    plt.ylabel("Loss")
    plt.title("Loss vs Epoch")
    plt.plot(x, losslist)
    plt.savefig(os.path.join(plotpath, "LossVsEpoch_" + str(n) + ".png"))
    plt.close()


def save(encoder, decoder, losslist, modelpath, n):
    encoder.cpu()
    decoder.cpu()
    torch.save(
        {
            "Enc": encoder,
            "Dec": decoder,
            "state_dict": {"enc": encoder.state_dict(), "dec": decoder.state_dict()},
        },
        os.path.join(modelpath, "models.pth"),
    )
    f = open("kl_losslist_" + str(n) + ".txt", "w+")
    for i, val in enumerate(losslist):
        f.write(f"Epoch: {i+1}, Loss: {losslist}\n")
    f.close()
    f = open("Info_" + str(n) + ".txt", "w+")
    f.write("Models saved in form of dictionary of the following structure\n")
    f.write(
        """
            'Enc'+:encoder,
            'Dec':decoder,
            'state_dict':{
              'enc':encoder state dict,
              'dec':decoder state dict
            } 
            \n"""
    )
    f.close()


def main(root, gen_counter, doPca, doMinMax):
    # origpath = os.path.join(root, "origdata", "data.csv")
    origpath = root
    print(origpath)

    genpath = os.path.join("../OutputFiles/VAE/", "gen_final_org" + str(gen_counter))
    modelpath = os.path.join("../OutputFiles/VAE/", "model_final_org" + str(gen_counter))
    plotpath = os.path.join("../OutputFiles/VAE/", "plot_final_org" + str(gen_counter))
    os.makedirs(genpath, exist_ok=True)
    os.makedirs(modelpath, exist_ok=True)
    os.makedirs(plotpath, exist_ok=True)
    print(genpath)
    print(modelpath)
    print(plotpath)
    # Hyperparameters
    latent_dim = 100
    batch_size = 10
    epochs = 4000
    learning_rate = 0.0005
    print(f"[INFO] Loading Dataset {origpath}...")

    print("[INFO] Load Data")
    sample_sizes = [83]
    for sample_size in sample_sizes:
        data = pd.read_csv(origpath)
        if "Unnamed: 0" in data.columns.values:
            data.drop(["Unnamed: 0"], axis=1, inplace=True)
        if "Unnamed: 0.1" in data.columns.values:
            data.drop(["Unnamed: 0.1"], axis=1, inplace=True)
        print(data.head(1))
        try:
            data.drop(["host_name", "tuberculosis", "hiv"], axis=1, inplace=True)
        except:
            pass
        columns = data.columns
        orig_data = data
        data = data.sample(sample_size)
        print(len(columns), "Num Columns")
        scaler = MinMaxScaler()
        pca = PCA(0.99)

        if doMinMax:
            print("In MinMax")
            data = pd.DataFrame(scaler.fit_transform(data), columns=columns)
        if doPca:
            data = pd.DataFrame(pca.fit_transform(data))

        print(data.columns, "FEATURE Space")

        rawdataset = CustomDataset(data)
        dataloader = DataLoader(rawdataset, batch_size=batch_size, shuffle=True)

        # rawdataset = CustomDataset(origpath, scale=True)
        # #print(type(rawdataset), " <- Raw Dataset")
        # dataloader = DataLoader(rawdataset, batch_size=batch_size, shuffle=True)
        original_dim = data.shape[1]
        # print(type(dataloader))
        print("[INFO] Training Begin")
        vae, losslist = train(
            original_dim,
            latent_dim,
            learning_rate,
            epochs,
            dataloader,
            rawdataset,
        )
        encoder = vae.enc
        decoder = vae.dec
        print("[INFO] Generate Samples")
        gen_sample = generate(
            decoder, len(orig_data), latent_dim, doMinMax, doPca, pca, scaler, columns
        )
        gen_sample.to_csv(os.path.join(genpath, f"VAE_{sample_size}.csv"), index=False)
        jsd = calculate_jsd(orig_data, gen_sample)
        wass_dist = wstd(orig_data, gen_sample)
        sys.stdout = open(
            f"{genpath}/model_sample_Axes_{sample_size}_.txt",
            "w+",
        )
        print("Dataset Sizes", orig_data.shape, gen_sample.shape)
        print(
            "JSD WSD",
            jsd,
            wass_dist,
        )
        sys.stdout = original
        # if size == sample_size[-1]:
        #     gen_sample.to_csv(os.path.join(genpath, "orggendata_" + str(size) + ".csv"))
        # gen_sample = generate(decoder, 250, latent_dim, doMinMax, scaler, columns)
        # gen_sample.to_csv(os.path.join(genpath, "250gendata.csv"))
        # gen_sample = generate(decoder, 550, latent_dim, doMinMax, scaler, columns)
        # gen_sample.to_csv(os.path.join(genpath, "550gendata.csv"))
        # save(encoder, decoder, modelpath, size)
        encoder.cpu()
        decoder.cpu()
        vae.cpu()
        del encoder
        del decoder
        del vae
        del rawdataset
        del dataloader
        del losslist


if __name__ == "__main__":
    import time

    roots = sorted(glob.glob("../InputFiles/Group*Data.csv"))
    gen_counter = 0
    roots = roots[:2]
    print(roots)
    for root in roots:
        gen_counter += 1
        print(f"{'-'*5} {root} {'-'*5}")
        start = time.time()
        main(root, gen_counter, False, False)
        print(f"Time required = {time.time()-start}")
        print(f"-" * 20)

    roots = sorted(glob.glob("../InputFiles/Group*Axes.csv"))
    gen_counter = 0
    roots = roots[:2]
    print(roots)
    for root in roots:
        gen_counter += 1
        print(f"{'-'*5} {root} {'-'*5}")
        start = time.time()
        main(root, gen_counter, False, False)
        print(f"Time required = {time.time()-start}")
        print(f"-" * 20)