{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from iwae import IWAE\n",
    "from data import list_datasets, collate, collate_sort, load_dataset\n",
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.stats import wasserstein_distance\n",
    "from sklearn.neighbors import KDTree\n",
    "import torch\n",
    "torch.set_default_tensor_type('torch.cuda.FloatTensor') # Comment this out if you want to train without GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = 'mixture'\n",
    "batch_size = 64\n",
    "learning_rate = 1e-3\n",
    "weight_decay = 0\n",
    "epochs = 50\n",
    "display_step = 5\n",
    "patience = 10\n",
    "# params\n",
    "hidden_dim = 64\n",
    "num_layers = 2\n",
    "n_bins = 5\n",
    "att_layers = 5\n",
    "n_heads = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch    5, loss_train = -1.6041, loss_val = -1.6286\n",
      "Epoch   10, loss_train = -1.8546, loss_val = -1.8500\n",
      "Epoch   15, loss_train = -1.9103, loss_val = -1.9036\n",
      "Epoch   20, loss_train = -1.9112, loss_val = -1.9214\n",
      "Epoch   25, loss_train = -1.9286, loss_val = -1.9334\n",
      "Epoch   30, loss_train = -1.8373, loss_val = -1.8864\n",
      "Breaking due to early stopping at epoch 33\n"
     ]
    }
   ],
   "source": [
    "## Load data\n",
    "dset = load_dataset(dataset)\n",
    "trainset, valset, testset = dset.split_train_val_test()\n",
    "\n",
    "collate = collate\n",
    "dl_train = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, collate_fn=collate)\n",
    "dl_val = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False, collate_fn=collate)\n",
    "dl_test = torch.utils.data.DataLoader(testset, batch_size=1, shuffle=False, collate_fn=collate)\n",
    "\n",
    "## Train model\n",
    "model = IWAE(dim=trainset.dim,\n",
    "             hidden_dim=hidden_dim,\n",
    "             num_layers=num_layers,\n",
    "             n_bins=n_bins,\n",
    "             att_layers=att_layers,\n",
    "             n_heads=n_heads)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)\n",
    "\n",
    "impatient = 0\n",
    "best_loss = np.inf\n",
    "best_model = deepcopy(model.state_dict())\n",
    "training_val_losses = []\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    # Optimization\n",
    "    model.train()\n",
    "    for batch in dl_train:\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        x, m, _ = batch\n",
    "        loss = model(x, m)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    # Validation\n",
    "    model.eval()\n",
    "    loss_val = 0\n",
    "    for i, batch in enumerate(dl_val):\n",
    "        x, m, _ = batch\n",
    "        loss_val += model(x, m).item() / len(dl_val)\n",
    "\n",
    "    training_val_losses.append(loss_val)\n",
    "\n",
    "    # Early stopping\n",
    "    if (best_loss - loss_val) < 1e-4:\n",
    "        impatient += 1\n",
    "        if loss_val < best_loss:\n",
    "            best_loss = loss_val\n",
    "            best_model = deepcopy(model.state_dict())\n",
    "    else:\n",
    "        best_loss = loss_val\n",
    "        best_model = deepcopy(model.state_dict())\n",
    "        impatient = 0\n",
    "\n",
    "    if impatient >= patience:\n",
    "        print(f'Breaking due to early stopping at epoch {epoch}')\n",
    "        break\n",
    "\n",
    "    if (epoch + 1) % display_step == 0:\n",
    "        print(f\"Epoch {epoch+1:4d}, loss_train = {loss:.4f}, loss_val = {loss_val:.4f}\")\n",
    "\n",
    "\n",
    "## Test model\n",
    "model.load_state_dict(best_model)\n",
    "model.eval()\n",
    "\n",
    "test_loss = 0\n",
    "for i, batch in enumerate(dl_test):\n",
    "    x, m, _ = batch\n",
    "    loss = model(x, m, num_samples=5000)\n",
    "    test_loss += loss.item() / len(dl_test)\n",
    "\n",
    "\n",
    "## Sampling quality -- Wasserstein score\n",
    "dist_test, dist_samples = [], []\n",
    "for x in testset:\n",
    "    if len(x[0]) > 2:\n",
    "        dist_test.append(KDTree(x[0]).query(x[0], k=2)[0][:,1])\n",
    "        samples = model.sample(len(x[0])).detach().cpu().numpy()\n",
    "        dist_samples.append(KDTree(samples).query(samples, k=2)[0][:,1])\n",
    "\n",
    "dist_test = np.concatenate(dist_test, 0)\n",
    "dist_samples = np.concatenate(dist_samples, 0)\n",
    "\n",
    "wasserstein = float(wasserstein_distance(dist_test, dist_samples))\n",
    "\n",
    "\n",
    "## Save results\n",
    "results = { 'test_loss': test_loss, 'training_val_losses': training_val_losses,\n",
    "            'final_epoch': epoch, 'wasserstein': wasserstein }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss -1.9830\n"
     ]
    }
   ],
   "source": [
    "print(f'Test loss {test_loss:.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 9512 characters]\n",
      "text/plain": [
       "<Figure size 288x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x = model.sample(100).detach().cpu().numpy()\n",
    "plt.figure(figsize=(4, 4))\n",
    "plt.scatter(x[:,0], x[:,1])\n",
    "plt.tight_layout()\n",
    "plt.xlim([0, 1])\n",
    "plt.ylim([0, 1])\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
