{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "41507fab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' \n",
    "import numpy as np\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('seaborn-paper')\n",
    "from data_loader import load_data\n",
    "from data_preprocesser import preprocess_data\n",
    "from radial import RadialFlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "48b7ee2d",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data successfully loaded...\n",
      "\n",
      "Data successfully preprocessed...\n",
      "\n",
      "x0 shape: (1, 2)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 19116 characters]\n",
      "text/plain": [
       "<Figure size 460.8x316.8 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 21412 characters]\n",
      "text/plain": [
       "<Figure size 460.8x316.8 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 23808 characters]\n",
      "text/plain": [
       "<Figure size 460.8x316.8 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 24708 characters]\n",
      "text/plain": [
       "<Figure size 460.8x316.8 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 28420 characters]\n",
      "text/plain": [
       "<Figure size 864x360 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TensorSpec(shape=(None, 2), dtype=tf.float32, name=None)\n",
      "TensorSpec(shape=(None, 2), dtype=tf.float32, name=None)\n",
      "\n",
      "Epoch 0...\n",
      "Epoch 1...\n",
      "Epoch 2...\n",
      "Epoch 3...\n",
      "Epoch 4...\n",
      "Epoch 5...\n",
      "Epoch 6...\n",
      "Epoch 7...\n",
      "Epoch 8...\n",
      "Epoch 9...\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 22276 characters]\n",
      "text/plain": [
       "<Figure size 460.8x316.8 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 30340 characters]\n",
      "text/plain": [
       "<Figure size 864x360 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def main():\n",
    "    \n",
    "    tfd = tfp.distributions\n",
    "    tfb = tfp.bijectors\n",
    "    \n",
    "    \"\"\" load data \"\"\"\n",
    "\n",
    "    filename = 'prostate.xls'\n",
    "    directory = '/Users/kaanguney.keklikci/Data/'\n",
    "\n",
    "    loader = load_data(filename, directory)\n",
    "    loader.create_directory(directory)\n",
    "    data = loader.read_data(directory, filename)\n",
    "    print('Data successfully loaded...\\n')\n",
    "\n",
    "    \"\"\" preprocess data \"\"\"\n",
    "\n",
    "    fillna_vals = ['sz', 'sg', 'wt']\n",
    "    dropna_vals = ['ekg', 'age']\n",
    "    drop_vals = ['patno', 'sdate']\n",
    "\n",
    "    preprocesser = preprocess_data(StandardScaler(), fillna_vals, dropna_vals, drop_vals)\n",
    "    data = preprocesser.dropna_features(data)\n",
    "    data = preprocesser.impute(data)\n",
    "    data = preprocesser.drop_features(data)\n",
    "    data = preprocesser.encode_categorical(data)\n",
    "    data = preprocesser.scale(data)\n",
    "    print('Data successfully preprocessed...\\n')\n",
    "    \n",
    "    \"\"\" define the base distributon as bivariate gaussian \"\"\"\n",
    "\n",
    "    base_dist = tfd.Independent(tfd.Normal(loc=[2., -0.5], scale=[1.,1.]),\n",
    "                              reinterpreted_batch_ndims=1)\n",
    "    \n",
    "    \"\"\" instantiate the bijector (a,b,x0) \"\"\"\n",
    "\n",
    "    n = 1000\n",
    "    a = 10.\n",
    "    b = -10.\n",
    "    x0 = np.array([-0.5,1.]).astype(np.float32).reshape(-1, 2)\n",
    "    bijector = RadialFlow(a, b, x0)\n",
    "    print(f'x0 shape: {x0.shape}')\n",
    "\n",
    "    x = np.linspace(start=-4, stop=4, num=n).astype(np.float32).reshape(-1, 2)\n",
    "    plt.title('Forward transformation')\n",
    "    plt.plot(x, bijector.forward(x))\n",
    "    plt.show()\n",
    "  \n",
    "    plt.plot(x, bijector.inverse(x))\n",
    "    plt.title('Inverse transformation')\n",
    "    plt.show()\n",
    " \n",
    "    plt.plot(x, bijector.forward_log_det_jacobian(x))\n",
    "    plt.title('Jacobian determinant')\n",
    "    plt.show()\n",
    "    \n",
    "    plt.plot(x, bijector.inverse_log_det_jacobian(x))\n",
    "    plt.title('Inverse Jacobian determinant')\n",
    "    plt.show()\n",
    "    \n",
    "    \"\"\" create transformed distribution \"\"\"\n",
    "    tfd_dist = tfd.TransformedDistribution(distribution=base_dist, \n",
    "                                           bijector=bijector\n",
    "                                         )\n",
    "    # prior training\n",
    "\n",
    "    plt.figure(figsize=(12,5))\n",
    "    plt.plot(tfd_dist.prob(x), label='Trainable')\n",
    "    plt.plot(base_dist.prob(x), label='Base')\n",
    "    plt.title('Target and Trainable distribution')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "    # sample, batch -- train, validation \n",
    "    \n",
    "    x_train = base_dist.sample(10000)\n",
    "    x_train = tf.data.Dataset.from_tensor_slices(x_train)\n",
    "    x_train = x_train.batch(int(n/2))\n",
    "\n",
    "    x_valid = base_dist.sample(1000)\n",
    "    x_valid = tf.data.Dataset.from_tensor_slices(x_valid)\n",
    "    x_valid = x_valid.batch(int(n/2))\n",
    "\n",
    "    print(x_train.element_spec)\n",
    "    print(x_valid.element_spec)\n",
    "    print()\n",
    "    \n",
    "    # instantiate trainable bijector \n",
    "\n",
    "    trainable_bijector = RadialFlow(tf.Variable(a, \n",
    "                                            name='alpha'), \n",
    "                                    tf.Variable(b, \n",
    "                                            name='beta'), \n",
    "                                    tf.Variable(x0, \n",
    "                                            name='ref'))\n",
    "                                   \n",
    "\n",
    "    # instantiate trainable distribution \n",
    "\n",
    "    trainable_dist = tfd.TransformedDistribution(tfd_dist, \n",
    "                                                 trainable_bijector\n",
    "                                                )\n",
    "    \n",
    "    # Train the bijector\n",
    "    \n",
    "    num_epochs = 10\n",
    "    opt = tf.keras.optimizers.Adam()\n",
    "    train_losses = []\n",
    "    valid_losses = []\n",
    "    norm = 1e3\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        print(\"Epoch {}...\".format(epoch))\n",
    "        train_loss = tf.keras.metrics.Mean()\n",
    "        val_loss = tf.keras.metrics.Mean()\n",
    "        for train_batch in x_train:\n",
    "            with tf.GradientTape() as tape:\n",
    "                tape.watch(trainable_bijector.trainable_variables)\n",
    "                loss = -trainable_dist.log_prob(train_batch)\n",
    "            train_loss(loss)\n",
    "            grads = tape.gradient(loss, trainable_bijector.trainable_variables)\n",
    "            grads, _ = tf.clip_by_global_norm(grads, norm) \n",
    "            # grads = tf.reshape(tf.nn.softmax(grads[-1], axis=1), [2])\n",
    "            # note that both alternatives work almost identically\n",
    "            opt.apply_gradients(zip(grads, trainable_bijector.trainable_variables))\n",
    "        train_losses.append(train_loss.result().numpy())\n",
    "        # validation\n",
    "        for valid_batch in x_valid:\n",
    "            loss = -trainable_dist.log_prob(valid_batch)\n",
    "            val_loss(loss)\n",
    "        valid_losses.append(val_loss.result().numpy())\n",
    "        \n",
    "    # Plot the learning curves\n",
    "\n",
    "    plt.plot(train_losses, label='train')\n",
    "    plt.plot(valid_losses, label='valid')\n",
    "    plt.legend()\n",
    "    plt.xlabel(\"Epochs\")\n",
    "    plt.ylabel(\"Negative log likelihood\")\n",
    "    plt.title(\"Training and validation loss curves\")\n",
    "    plt.show()\n",
    "    \n",
    "    # Plot the data and learned distributions\n",
    "\n",
    "    plt.figure(figsize=(12,5))\n",
    "    plt.plot(trainable_dist.prob(x), label='Learned')\n",
    "    plt.plot(base_dist.prob(x), label='Data')\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "    \n",
    "    ### best result obtained with tuning displayed as above\n",
    "    ### radial flows converge fast, more epochs overfit\n",
    "    ### DO NOT CHANGE validate_args=True\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "a2fb12cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "### DOES NOT USE DATASET YET\n",
    "### FOR VISUALIZATION PURPOSES IN 2D\n",
    "### WILL INTEGRATE DATASET AFTER LEARNING THE DISTRIBUTION"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
