{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fac0167e",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7008 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 6376 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "%run 'preprocessing/preprocessing_v1.1.ipynb'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9f0b013a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os \n",
    "import time \n",
    "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
    "import numpy as np \n",
    "import pandas as pd\n",
    "import seaborn as sns \n",
    "sns.set_style('white')\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline \n",
    "from sklearn.datasets import make_moons\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "from tensorflow.keras.callbacks import LambdaCallback\n",
    "import tensorflow as tf \n",
    "tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR) # for gradient warning (variables are still updated, therefore supressing is fine)\n",
    "import tensorflow_probability as tfp \n",
    "from tensorflow.keras import Model, Sequential\n",
    "from tensorflow.keras.layers import Dense, Flatten, Reshape, Input\n",
    "from tensorflow.keras.losses import Loss\n",
    "import scipy\n",
    "\n",
    "tfd = tfp.distributions\n",
    "tfb = tfp.bijectors\n",
    "tfpl = tfp.layers\n",
    "\n",
    "color_list = ['#bcad', '#dacb']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17da99b5",
   "metadata": {},
   "source": [
    "### Define prior distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9ed6efd6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_prior(num_modes, latent_dim):\n",
    "    \n",
    "    mixture_distribution = tfd.Categorical(probs=[1./num_modes] * num_modes)\n",
    "    components_distribution = tfd.MultivariateNormalDiag(loc=tf.Variable(tf.random.normal((num_modes, latent_dim))),\n",
    "                                                        scale_diag=tfp.util.TransformedVariable(tf.ones((num_modes,latent_dim)),\n",
    "                                                                                               bijector=tfb.Softplus())\n",
    "                                                        )\n",
    "    prior = tfd.MixtureSameFamily(mixture_distribution,\n",
    "                                  components_distribution\n",
    "                                 )\n",
    "    return prior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "956433fd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prior event shape: 2\n",
      "# of Gaussions: 2\n",
      "Covariance matrix: MultivariateNormalDiag\n"
     ]
    }
   ],
   "source": [
    "latent_dim = 2\n",
    "input_shape = 2\n",
    "\n",
    "prior = get_prior(num_modes=latent_dim, latent_dim=input_shape) \n",
    "print(f'Prior event shape: {prior.event_shape[0]}')\n",
    "print(f'# of Gaussions: {prior.components_distribution.batch_shape[0]}') \n",
    "print(f'Covariance matrix: {prior.components_distribution.name}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf7ab036",
   "metadata": {},
   "source": [
    "### Define KL divergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d33eb9a1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# set weight for more emphasis on KLDivergence term rather than reconstruction loss\n",
    "# average over both samples and batches\n",
    "\n",
    "def get_KL_regularizer(prior, weight=4.):\n",
    "    \n",
    "    regularizer = tfpl.KLDivergenceRegularizer(prior, \n",
    "                                        use_exact_kl=False,\n",
    "                                        test_points_reduce_axis=(),\n",
    "                                        test_points_fn=lambda q: q.sample(10),\n",
    "                                        weight=weight\n",
    "                                        )\n",
    "    return regularizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "859c9d44",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "KLDivergence_regularizer = get_KL_regularizer(prior)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57682c39",
   "metadata": {},
   "source": [
    "### Define the encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c011c224",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_encoder(input_shape, latent_dim, KL_regularizer):\n",
    "    \n",
    "    encoder = Sequential([\n",
    "        Dense(input_shape=input_shape, units=256, activation='relu'),\n",
    "        Dense(units=128, activation='relu'),\n",
    "        Dense(units=64, activation='relu'),\n",
    "        Dense(units=32, activation='relu'),\n",
    "        Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim)),\n",
    "        tfpl.MultivariateNormalTriL(latent_dim, \n",
    "                                   activity_regularizer=KL_regularizer),\n",
    "    ])\n",
    "    \n",
    "    return encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ffea5058",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense (Dense)                (None, 256)               768       \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 128)               32896     \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 64)                8256      \n",
      "_________________________________________________________________\n",
      "dense_3 (Dense)              (None, 32)                2080      \n",
      "_________________________________________________________________\n",
      "dense_4 (Dense)              (None, 5)                 165       \n",
      "_________________________________________________________________\n",
      "multivariate_normal_tri_l (M multiple                  8         \n",
      "=================================================================\n",
      "Total params: 44,173\n",
      "Trainable params: 44,173\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "encoder = get_encoder(input_shape=(input_shape,), latent_dim=latent_dim, KL_regularizer=KLDivergence_regularizer)\n",
    "encoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad3eb53a",
   "metadata": {},
   "source": [
    "### Define the decoder "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2fcb68ca",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_decoder(latent_dim):\n",
    "    \n",
    "    decoder = Sequential([\n",
    "        Dense(input_shape=(latent_dim,), units=5, activation='relu'),\n",
    "        Dense(units=64, activation='relu'),\n",
    "        Dense(units=128, activation='relu'),\n",
    "        Dense(units=256, activation='relu'),\n",
    "        Dense(tfpl.MultivariateNormalTriL.params_size(latent_dim)),\n",
    "        tfpl.MultivariateNormalTriL(latent_dim)\n",
    "    ])\n",
    "    \n",
    "    return decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f7abd481",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_1\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense_5 (Dense)              (None, 5)                 15        \n",
      "_________________________________________________________________\n",
      "dense_6 (Dense)              (None, 64)                384       \n",
      "_________________________________________________________________\n",
      "dense_7 (Dense)              (None, 128)               8320      \n",
      "_________________________________________________________________\n",
      "dense_8 (Dense)              (None, 256)               33024     \n",
      "_________________________________________________________________\n",
      "dense_9 (Dense)              (None, 5)                 1285      \n",
      "_________________________________________________________________\n",
      "multivariate_normal_tri_l_1  multiple                  0         \n",
      "=================================================================\n",
      "Total params: 43,028\n",
      "Trainable params: 43,028\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "decoder = get_decoder(latent_dim)\n",
    "decoder.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8998844",
   "metadata": {},
   "source": [
    "### Connect encoder to decoder "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "44663b3d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense_input (InputLayer)     [(None, 2)]               0         \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 256)               768       \n",
      "_________________________________________________________________\n",
      "dense_1 (Dense)              (None, 128)               32896     \n",
      "_________________________________________________________________\n",
      "dense_2 (Dense)              (None, 64)                8256      \n",
      "_________________________________________________________________\n",
      "dense_3 (Dense)              (None, 32)                2080      \n",
      "_________________________________________________________________\n",
      "dense_4 (Dense)              (None, 5)                 165       \n",
      "_________________________________________________________________\n",
      "multivariate_normal_tri_l (M multiple                  8         \n",
      "_________________________________________________________________\n",
      "sequential_1 (Sequential)    multiple                  43028     \n",
      "=================================================================\n",
      "Total params: 87,201\n",
      "Trainable params: 87,201\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "vae = Model(inputs=encoder.inputs, outputs=decoder(encoder.outputs))\n",
    "vae.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "193def59",
   "metadata": {},
   "source": [
    "### Specify the loss function "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f572dddd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# KL divergence is implicity incorporated to the loss function before\n",
    "# add reconstruction error to loss function \n",
    "def reconstruction_error(decoding_dist, x_true):\n",
    "    return -tf.reduce_mean(decoding_dist.log_prob(x_true))\n",
    "\n",
    "class custom_reconstruction_error(Loss):\n",
    "    def call(self, decoding_dist, x_true):\n",
    "        return -tf.reduce_mean(decoding_dist.log_prob(x_true))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b30984a3",
   "metadata": {},
   "source": [
    "### Selection process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "dfee364c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# of training samples: 224\n",
      "# of test samples: 224\n"
     ]
    }
   ],
   "source": [
    "print(f'# of training samples: {X_train.shape[0]}')\n",
    "print(f'# of test samples: {X_test.shape[0]}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "68e9f656",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train_prior = np.append(X_train.DGN.to_numpy().reshape(-1,1), scaler.inverse_transform(X_train[keepdims[1:]]), axis=1)\n",
    "X_test_prior = np.append(X_test.DGN.to_numpy().reshape(-1,1), scaler.inverse_transform(X_test[keepdims[1:]]), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e176949b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train_prior_visual = pd.DataFrame(X_train_prior, columns=keepdims)\n",
    "y_train.reset_index(drop=True, inplace=True)\n",
    "X_train_prior_visual['Risk1Yr'] = y_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a528b011",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train_prior_visual[keepdims[1:]] = X_train_prior_visual[keepdims[1:]].astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1c742030",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 55028 characters]\n",
      "text/plain": [
       "<Figure size 763.125x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = sns.catplot(x='DGN', y='PRE4', kind=\"violin\", hue='Risk1Yr', data=X_train_prior_visual, inner=None, legend_out=True, doge=True, palette='Set3', aspect=2, ci=95)\n",
    "sns.swarmplot(x=\"DGN\", y='PRE4', color=\"k\", size=3, data=X_train_prior_visual, ax=g.ax)\n",
    "sns.despine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b996f2ae",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61552 characters]\n",
      "text/plain": [
       "<Figure size 763.125x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = sns.catplot(x='DGN', y='PRE4', kind=\"violin\", hue='Risk1Yr', data=X_train_prior_visual, inner=None, doge=True, palette='Set3', legend_out=True, aspect=2, bw=0.2, ci=95)\n",
    "sns.swarmplot(x=\"DGN\", y='PRE4', color=\"k\", size=3, data=X_train_prior_visual, ax=g.ax)\n",
    "sns.despine()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce39bd4b",
   "metadata": {},
   "source": [
    "### Compile and fit the model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5aa3970d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train[keepdims[1:]] = X_train[keepdims[1:]].astype(np.float32)\n",
    "X_test[keepdims[1:]] = X_test[keepdims[1:]].astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ab6e83bc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train.drop('DGN',axis=1,inplace=True)\n",
    "X_test.drop('DGN',axis=1,inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "3215530a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Epoch 1/600 \n",
      "\t loss: 147.5679, val_loss: 140.1842\n",
      "\n",
      " Epoch 101/600 \n",
      "\t loss: 6.5048, val_loss: 6.7086\n",
      "\n",
      " Epoch 201/600 \n",
      "\t loss: 3.0677, val_loss: 3.7662\n",
      "\n",
      " Epoch 301/600 \n",
      "\t loss: 1.7498, val_loss: 2.5523\n",
      "\n",
      " Epoch 401/600 \n",
      "\t loss: 1.0859, val_loss: 0.9216\n",
      "\n",
      " Epoch 501/600 \n",
      "\t loss: 0.6835, val_loss: 0.4812\n"
     ]
    }
   ],
   "source": [
    "optimizer = Adam(learning_rate=3e-4)\n",
    "epochs = 600 \n",
    "epoch_callback = LambdaCallback(on_epoch_end=lambda epoch, logs: print('\\n Epoch {}/{}'.format(epoch+1, epochs, logs),\n",
    "                                                                       '\\n\\t ' + (': {:.4f}, '.join(logs.keys()) + ': {:.4f}').format(*logs.values()))\n",
    "                                if epoch % 100 == 0 else False\n",
    "                               )\n",
    "\n",
    "\n",
    "vae.compile(optimizer=optimizer, loss=reconstruction_error)\n",
    "history = vae.fit(X_train, \n",
    "                  validation_data=(X_test,),\n",
    "                  epochs=epochs, \n",
    "                  batch_size=32,\n",
    "                  verbose=0,\n",
    "                  shuffle=True,\n",
    "                  callbacks=[epoch_callback]\n",
    "                 )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77a16a56",
   "metadata": {},
   "source": [
    "### Plot training and validation losses "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "116abbd8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16296 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(history.history['loss'], label='training')\n",
    "plt.plot(history.history['val_loss'], label='validation')\n",
    "plt.legend(loc='best')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('KL Divergence')\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Loss function is ELBO maximization\n",
    "# ELBO maximization is equivalent to KL divergence minimization "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fadadad2",
   "metadata": {},
   "source": [
    "### Sample from the generative model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "e65f7315",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "X_train_sample = encoder(X_train.to_numpy()).sample()\n",
    "X_train_sample = pd.DataFrame(X_train_sample.numpy(), columns=X_test.columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "c5ebde09",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-1.531444</td>\n",
       "      <td>-1.589891</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.089678</td>\n",
       "      <td>-0.586031</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-1.267397</td>\n",
       "      <td>1.969045</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.239857</td>\n",
       "      <td>-3.134773</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-2.799463</td>\n",
       "      <td>-1.603101</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       PRE4      PRE5\n",
       "0 -1.531444 -1.589891\n",
       "1 -0.089678 -0.586031\n",
       "2 -1.267397  1.969045\n",
       "3  0.239857 -3.134773\n",
       "4 -2.799463 -1.603101"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e7bd3f85",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.903259</td>\n",
       "      <td>-13.371138</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3.167860</td>\n",
       "      <td>-2.181474</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2.134860</td>\n",
       "      <td>26.299038</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3.456901</td>\n",
       "      <td>-30.591393</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.791055</td>\n",
       "      <td>-13.518393</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       PRE4       PRE5\n",
       "0  1.903259 -13.371138\n",
       "1  3.167860  -2.181474\n",
       "2  2.134860  26.299038\n",
       "3  3.456901 -30.591393\n",
       "4  0.791055 -13.518393"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# propogate back from the pipeline\n",
    "# training data -- reverse standardization\n",
    "\n",
    "X_train_sample = pd.DataFrame(scaler.inverse_transform(X_train_sample.to_numpy()), columns=X_train.columns)\n",
    "X_train_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "39439ca5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>3.417876</td>\n",
       "      <td>-22.587229</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1.725919</td>\n",
       "      <td>22.992628</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3.199011</td>\n",
       "      <td>8.050227</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2.679654</td>\n",
       "      <td>-2.083026</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1.792891</td>\n",
       "      <td>-21.024277</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       PRE4       PRE5\n",
       "0  3.417876 -22.587229\n",
       "1  1.725919  22.992628\n",
       "2  3.199011   8.050227\n",
       "3  2.679654  -2.083026\n",
       "4  1.792891 -21.024277"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# propogate back from the pipeline\n",
    "# testing data -- reverse standardization\n",
    "\n",
    "X_test_sample = encoder(X_test.to_numpy()).sample()\n",
    "X_test_sample = pd.DataFrame(X_test_sample.numpy(), columns=X_test.columns)\n",
    "\n",
    "X_test_sample = pd.DataFrame(scaler.inverse_transform(X_test_sample.to_numpy()), columns=X_test.columns)\n",
    "X_test_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d57c8784",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "y_train.reset_index(drop=True,inplace=True)\n",
    "y_test.reset_index(drop=True,inplace=True) # for concatenation\n",
    "\n",
    "X_train_sample['Risk1Yr'] = y_train\n",
    "X_train_sample['DGN'] = X_train_prior[...,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c661a8d2",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 58108 characters]\n",
      "text/plain": [
       "<Figure size 763.125x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = sns.catplot(x='DGN', y='PRE4', kind=\"violin\", hue='Risk1Yr', data=X_train_sample, inner=None, legend_out=True, doge=True, palette='Set3', aspect=2, ci=95)\n",
    "sns.swarmplot(x=\"DGN\", y='PRE4', color=\"k\", size=3, data=X_train_sample, ax=g.ax)\n",
    "sns.despine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "979b9fce",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61552 characters]\n",
      "text/plain": [
       "<Figure size 763.125x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = sns.catplot(x='DGN', y='PRE4', kind=\"violin\", hue='Risk1Yr', data=X_train_prior_visual, inner=None, doge=True, palette='Set3', legend_out=True, aspect=2, bw=0.2, ci=95)\n",
    "f = sns.swarmplot(x=\"DGN\", y='PRE4', color=\"k\", size=3, data=X_train_prior_visual, ax=g.ax)\n",
    "sns.despine()\n",
    "f = f.get_figure()\n",
    "f.savefig(os.getcwd() + '/maf-results/p.jpeg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "2a20c845",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>4.826622</td>\n",
       "      <td>10.922666</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>4.016593</td>\n",
       "      <td>7.421674</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>4.803132</td>\n",
       "      <td>-11.967712</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5.686647</td>\n",
       "      <td>11.195478</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5.073198</td>\n",
       "      <td>24.740921</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       PRE4       PRE5\n",
       "0  4.826622  10.922666\n",
       "1  4.016593   7.421674\n",
       "2  4.803132 -11.967712\n",
       "3  5.686647  11.195478\n",
       "4  5.073198  24.740921"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# decode \n",
    "X_train_sample.drop('Risk1Yr',axis=1,inplace=True)\n",
    "X_train_sample.drop('DGN', axis=1, inplace=True)\n",
    "X_train_sample = pd.DataFrame(scaler.inverse_transform(decoder(X_train_sample.to_numpy()).sample().numpy()), columns=X_train.columns)\n",
    "X_train_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c0d87cf9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57160 characters]\n",
      "text/plain": [
       "<Figure size 763.125x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_train_sample['DGN'] = X_train_prior[...,0]\n",
    "X_train_sample['Risk1Yr'] = y_train\n",
    "\n",
    "g = sns.catplot(x='DGN', y='PRE4', kind=\"violin\", hue='Risk1Yr', data=X_train_sample, inner=None, doge=True, palette='Set3', legend_out=True, aspect=2, bw=0.2, ci=95)\n",
    "sns.swarmplot(x=\"DGN\", y='PRE4', color=\"k\", size=3, data=X_train_sample, ax=g.ax)\n",
    "sns.despine()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71993523",
   "metadata": {},
   "source": [
    "### Masked Autoregressive Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "8f847af5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train_sample.drop('Risk1Yr', axis=1, inplace=True) # drop target, was only intended for VAE sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "ce7944ce",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "loc = [X_train_sample[i].mean().astype('float32') for i in list(X_train.columns)]\n",
    "scale_diag = [X_train_sample[i].std().astype('float32') for i in list(X_train.columns)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "54436897",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.MultivariateNormalDiag 'MultivariateNormalDiag' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvn = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)\n",
    "mvn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "243e4115",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def masked_autoregressive_flow(hidden_units=[16,16], event_shape=[2], activation='relu'):\n",
    "    network = tfb.AutoregressiveNetwork(params=2, \n",
    "                                    hidden_units=hidden_units,\n",
    "                                    event_shape=event_shape,\n",
    "                                    activation=activation\n",
    "                                   )\n",
    "    return tfb.MaskedAutoregressiveFlow(shift_and_log_scale_fn=network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "8c6ff7a1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.TransformedDistribution 'masked_autoregressive_flowMultivariateNormalDiag' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainable_dist = tfd.TransformedDistribution(distribution=mvn,\n",
    "                                             bijector=masked_autoregressive_flow(\n",
    "                                             activation='sigmoid'))\n",
    "trainable_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "56ce9b96",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "n_samples = 300\n",
    "x = mvn.sample(sample_shape=n_samples)\n",
    "names = [mvn.name, trainable_dist.bijector.name]\n",
    "samples = [x, trainable_dist.bijector.forward(x)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "bded8729",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train_np = X_train_sample.to_numpy()\n",
    "X_test_np = X_test_sample.to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "b1e0c2b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train_np = X_train_np[:,0:2]\n",
    "X_test_np = X_test_np[:,0:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "da26e4e2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# standardize once again before feeding into network \n",
    "scaler.fit(X_train_np)\n",
    "X_train_np = scaler.transform(X_train_np)\n",
    "X_test_np = scaler.transform(X_test_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "0a671c81",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X_train = X_train_np.astype(np.float32)\n",
    "X_train = tf.data.Dataset.from_tensor_slices(X_train)\n",
    "X_train = X_train.batch(32)\n",
    "\n",
    "X_valid = X_test_np.astype(np.float32)\n",
    "X_valid = tf.data.Dataset.from_tensor_slices(X_valid)\n",
    "X_valid = X_valid.batch(32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "d9971446",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0...\n",
      "Epoch 100...\n",
      "Epoch 200...\n",
      "Epoch 300...\n",
      "Epoch 400...\n",
      "Epoch 500...\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 600\n",
    "opt = tf.keras.optimizers.Adam(3e-4)\n",
    "train_losses = []\n",
    "valid_losses = []\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    if epoch % 100 == 0:\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_dist.bijector.trainable_variables)\n",
    "            loss = -trainable_dist.log_prob(train_batch)\n",
    "        train_loss(loss)\n",
    "        grads = tape.gradient(loss, trainable_dist.bijector.trainable_variables)\n",
    "        opt.apply_gradients(zip(grads, trainable_dist.bijector.trainable_variables))\n",
    "    train_losses.append(train_loss.result().numpy())\n",
    "        \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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "8e670a4c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17548 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_losses = history.history['loss']\n",
    "valid_losses = history.history['val_loss']\n",
    "\n",
    "plt.plot(train_losses, label='train', c=color_list[0])\n",
    "plt.plot(valid_losses, label='valid', c=color_list[1])\n",
    "plt.legend()\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(\"Negative log likelihood\")\n",
    "plt.title(\"Training and validation loss curves\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "0e24f956",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "x = mvn.sample(sample_shape=n_samples)\n",
    "names = [mvn.name, trainable_dist.bijector.name]\n",
    "samples = [x, trainable_dist.bijector.forward(x)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "272037c9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "num_layers = 4\n",
    "flow_bijector = []\n",
    "\n",
    "# set trainable=False for once and for all initialization \n",
    "def init_once(x):\n",
    "    return tf.Variable(x, name='permutation', trainable=False)\n",
    "\n",
    "\n",
    "for i in range(num_layers):\n",
    "    flow_i = masked_autoregressive_flow(hidden_units=[256,256])\n",
    "    flow_bijector.append(flow_i) \n",
    "    flow_bijector.append(tfb.Permute([1,0]))\n",
    "# discard the last permute layer \n",
    "flow_bijector = tfb.Chain(list(reversed(flow_bijector[:-1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "43c7a1f0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.TransformedDistribution 'chain_of_masked_autoregressive_flow_of_permute_of_masked_autoregressive_flow_of_permute_of_masked_autoregressive_flow_of_permute_of_masked_autoregressive_flowMultivariateNormalDiag' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainable_dist = tfd.TransformedDistribution(distribution=mvn,\n",
    "                                            bijector=flow_bijector)\n",
    "trainable_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "fb2d1a9a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def make_samples():\n",
    "    x = mvn.sample(n_samples)\n",
    "    samples = [x]\n",
    "    names = [mvn.name]\n",
    "    for bijector in reversed(trainable_dist.bijector.bijectors):\n",
    "        x = bijector.forward(x)\n",
    "        samples.append(x)\n",
    "        names.append(bijector.name)\n",
    "    return names, samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "ef7ef2b6",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0...\n",
      "Epoch 100...\n",
      "Epoch 200...\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 300\n",
    "opt = tf.keras.optimizers.Adam(3e-4)\n",
    "train_losses = []\n",
    "valid_losses = []\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    if epoch % 100 == 0:\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_dist.bijector.trainable_variables)\n",
    "            loss = -trainable_dist.log_prob(train_batch)\n",
    "        train_loss(loss)\n",
    "        grads = tape.gradient(loss, trainable_dist.bijector.trainable_variables)\n",
    "        opt.apply_gradients(zip(grads, trainable_dist.bijector.trainable_variables))\n",
    "    train_losses.append(train_loss.result().numpy())\n",
    "        \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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "886ee70d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31428 characters]\n",
      "text/plain": [
       "<Figure size 864x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(12, 5))\n",
    "cutoff = 5\n",
    "train_losses = history.history['loss'][cutoff:]\n",
    "valid_losses = history.history['val_loss'][cutoff:]\n",
    "\n",
    "plt.plot(np.arange(cutoff, epochs), train_losses, label='training')\n",
    "plt.plot(np.arange(cutoff, epochs), valid_losses, label='validation')\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.ylabel(\"Negative log likelihood\")\n",
    "plt.title(\"Training and validation loss curves\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "4177a424",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "names, samples = make_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "b7617c92",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_training_data(samples, plot_color: str=color_list[0]):\n",
    "    f, arr = plt.subplots(1, 2, figsize=(20,5))\n",
    "    names = ['Data', 'Trainable']\n",
    "    samples = [tf.constant(X_train_np), samples[-1]]\n",
    "\n",
    "    for i in range(2):\n",
    "        res = samples[i]\n",
    "        X, Y = res[..., 0].numpy(), res[..., 1].numpy()\n",
    "        arr[i].scatter(X, Y, s=10, color=plot_color)\n",
    "        z = np.polyfit(X, Y, 1)\n",
    "        p = np.poly1d(z)\n",
    "        arr[i].plot(X, p(X), color=color_list[1])\n",
    "        arr[i].set_xlim([-2, 2])\n",
    "        arr[i].set_ylim([-2, 2])\n",
    "        arr[i].set_title(names[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "8b52f865",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27472 characters]\n",
      "text/plain": [
       "<Figure size 1440x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_training_data(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "ec7a463d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "samples_nf = pd.DataFrame(scaler.inverse_transform(samples[-1].numpy()),columns=keepdims[1:])\n",
    "samples_nf['Risk1Yr'] = y_train\n",
    "samples_nf['DGN'] = X_train_sample.DGN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "id": "ccaff6ba",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 59340 characters]\n",
      "text/plain": [
       "<Figure size 763.125x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = sns.catplot(x='DGN', y='PRE4', kind=\"violin\", hue='Risk1Yr', data=samples_nf, inner=None, doge=True, palette='Set3', legend_out=True, aspect=2, bw=0.2, ci=95)\n",
    "f = sns.swarmplot(x=\"DGN\", y='PRE4', color=\"k\", size=3, data=samples_nf, ax=g.ax)\n",
    "f = f.get_figure()\n",
    "sns.despine()\n",
    "f.savefig(os.getcwd() + '/maf-results/q.jpeg')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df0c53f6",
   "metadata": {},
   "source": [
    "### Measure KL Divergence "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "2be7e9a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.003868835\n",
      "-0.00078658777\n"
     ]
    }
   ],
   "source": [
    "# training \n",
    "prior_train = mvn.prob(X_train_np).numpy()\n",
    "learned_train = trainable_dist.prob(X_train_np).numpy()\n",
    "\n",
    "kl = tf.keras.metrics.KLDivergence()\n",
    "kl.update_state(prior_train, learned_train)\n",
    "print(kl.result().numpy())\n",
    "kl.reset_state()\n",
    "\n",
    "# testing\n",
    "prior_test = mvn.prob(X_test_np).numpy()\n",
    "learned_test = trainable_dist.prob(X_test_np).numpy()\n",
    "\n",
    "kl = tf.keras.metrics.KLDivergence()\n",
    "kl.update_state(prior_test, learned_test)\n",
    "print(kl.result().numpy())\n",
    "kl.reset_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd78efd5",
   "metadata": {},
   "source": [
    "### Measure Shannon Entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "9670b9db",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.14439699\n",
      "0.0005107625\n"
     ]
    }
   ],
   "source": [
    "# training\n",
    "cross_entropy = prior_train * np.log(learned_train)\n",
    "print(-np.ma.masked_invalid(cross_entropy).sum())\n",
    "\n",
    "# testing\n",
    "cross_entropy = prior_test * np.log(learned_test)\n",
    "print(-np.ma.masked_invalid(cross_entropy).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f92bd8d",
   "metadata": {},
   "source": [
    "### Measure Poisson "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "id": "62054131",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.09606825\n",
      "0.030584812\n"
     ]
    }
   ],
   "source": [
    "# training \n",
    "poisson = tf.keras.metrics.Poisson()\n",
    "poisson.update_state(prior_train, learned_train)\n",
    "print(poisson.result().numpy())\n",
    "poisson.reset_state()\n",
    "\n",
    "# testing \n",
    "poisson = tf.keras.metrics.Poisson()\n",
    "poisson.update_state(prior_test, learned_test)\n",
    "print(poisson.result().numpy())\n",
    "poisson.reset_state()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94b7cd11",
   "metadata": {},
   "source": [
    "### Measure MAE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "3734da33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.095406\n",
      "0.03058204\n"
     ]
    }
   ],
   "source": [
    "# training \n",
    "mae = tf.keras.losses.MeanAbsoluteError()\n",
    "print(mae(prior_train, learned_train).numpy())\n",
    "\n",
    "# testing\n",
    "print(mae(prior_test, learned_test).numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "427fcb58",
   "metadata": {},
   "source": [
    "### Performance Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "8535be8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# collect metrics\n",
    "def _collect(prior_train, prior_test, learned_train, learned_test, results, number_of_run):\n",
    "    \n",
    "    # KL training\n",
    "    results['Kullback-Leibler Divergence'][number_of_run][0] = np.ma.masked_invalid(prior_train * np.log(prior_train/learned_train)).sum()\n",
    "    # KL testing\n",
    "    results['Kullback-Leibler Divergence'][number_of_run][1] = np.ma.masked_invalid(prior_test * np.log(prior_test/learned_test)).sum()\n",
    "    \n",
    "    # Cross Entropy training\n",
    "    results['Cross Entropy'][number_of_run][0] = -np.ma.masked_invalid(prior_train * np.log(learned_train)).sum()\n",
    "    # Cross Entropy testing\n",
    "    results['Cross Entropy'][number_of_run][1] = -np.ma.masked_invalid(prior_test * np.log(learned_test)).sum()\n",
    "    \n",
    "    poisson = tf.keras.metrics.Poisson()\n",
    "    # Poisson training\n",
    "    poisson.update_state(prior_train, learned_train)\n",
    "    results['Poisson'][number_of_run][0]  = poisson.result().numpy()\n",
    "    poisson.reset_state()\n",
    "    # Poisson testing\n",
    "    poisson.update_state(prior_test, learned_test)\n",
    "    results['Poisson'][number_of_run][1]  = poisson.result().numpy()\n",
    "    poisson.reset_state()\n",
    "    \n",
    "    # MAE training \n",
    "    mae = tf.keras.losses.MeanAbsoluteError()\n",
    "    results['Mean Absolute Error'][number_of_run][0] = mae(prior_train, learned_train).numpy()\n",
    "    results['Mean Absolute Error'][number_of_run][1] = mae(prior_test, learned_test).numpy()\n",
    "    \n",
    "    return results \n",
    "\n",
    "# metrics to measure KL, Poisson, Cross Entropy, Mean Absolute Error\n",
    "def init_results():\n",
    "    results = {'Kullback-Leibler Divergence': [[None, None] for _ in range(NUMBER_OF_RUNS)], \n",
    "               'Cross Entropy': [[None, None] for _ in range(NUMBER_OF_RUNS)],\n",
    "               'Poisson': [[None, None] for _ in range(NUMBER_OF_RUNS)], \n",
    "               'Mean Absolute Error': [[None, None] for _ in range(NUMBER_OF_RUNS)]\n",
    "              }\n",
    "    return results\n",
    "\n",
    "NUMBER_OF_RUNS = 5\n",
    "results = init_results()\n",
    "\n",
    "# generate new observations and measure with the generative model \n",
    "def _results(X_train_np, X_test_np, NUMBER_OF_RUNS, results):\n",
    "    for num in range(NUMBER_OF_RUNS):\n",
    "        \n",
    "        X_gen_train = encoder(X_train_np).sample().numpy()\n",
    "        X_gen_test = encoder(X_test_np).sample().numpy()\n",
    "        \n",
    "        # convert arrays to probabilities for metric calculation, avoids nan or inf values\n",
    "        # X_gen_train = tf.exp(-X_gen_train) / tf.reduce_sum(tf.exp(-X_gen_train), axis=0)\n",
    "        # X_gen_test = tf.exp(-X_gen_test) / tf.reduce_sum(tf.exp(-X_gen_test), axis=0)\n",
    "        \n",
    "        # get distributions for train and test \n",
    "        prior_train = prior.prob(X_gen_train).numpy()\n",
    "        learned_train = trainable_dist.prob(X_gen_train).numpy()\n",
    "        \n",
    "        prior_test = prior.prob(X_gen_test).numpy()\n",
    "        learned_test = trainable_dist.prob(X_gen_test).numpy()\n",
    "        \n",
    "        results = _collect(prior_train, prior_test, learned_train, learned_test, results, num)\n",
    "        \n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "616f50ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0.99999946, 1.0000002 ], dtype=float32)>"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# convert output of encoder to probabilities \n",
    "trial_data = encoder(X_train_np).sample().numpy()\n",
    "tf.reduce_sum(tf.exp(-trial_data) / tf.reduce_sum(tf.exp(-trial_data), axis=0),axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "d43326bc",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Kullback-Leibler Divergence': [[31.453903, 32.703033],\n",
       "  [25.26685, 35.475212],\n",
       "  [25.705252, 22.91606],\n",
       "  [25.389862, 25.627352],\n",
       "  [24.527397, 27.624155]],\n",
       " 'Cross Entropy': [[57.719772, 57.106598],\n",
       "  [51.27077, 60.587334],\n",
       "  [51.205433, 46.404827],\n",
       "  [50.89733, 51.456905],\n",
       "  [51.333427, 52.093456]],\n",
       " 'Poisson': [[0.2381211, 0.2282642],\n",
       "  [0.23317902, 0.23885687],\n",
       "  [0.2213739, 0.20724127],\n",
       "  [0.22463407, 0.22597542],\n",
       "  [0.22960626, 0.21534447]],\n",
       " 'Mean Absolute Error': [[0.029114885, 0.031007694],\n",
       "  [0.03636083, 0.029514452],\n",
       "  [0.030275216, 0.028754879],\n",
       "  [0.03284185, 0.033900343],\n",
       "  [0.034867395, 0.030299732]]}"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# sanity check \n",
    "_results(X_train_np, X_test_np, NUMBER_OF_RUNS, results)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25851622",
   "metadata": {},
   "source": [
    "### [Pooled Estimate of Common Std. Deviation](https://sphweb.bumc.bu.edu/otlt/mph-modules/bs/bs704_confidence_intervals/bs704_confidence_intervals5.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "13fd827c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metric sampling in seconds: 4\n"
     ]
    }
   ],
   "source": [
    "results = init_results()\n",
    "start_time = time.time()\n",
    "stats = _results(X_train_np, X_test_np, NUMBER_OF_RUNS, results)\n",
    "print('Metric sampling in seconds: {sec}'.format(sec=round(time.time()-start_time)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "id": "9a68a94c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate stds and xbars\n",
    "stdxbars = {'Kullback-Leibler Divergence': [None, None],\n",
    "           'Cross Entropy': [None, None],\n",
    "            'Poisson': [None, None],\n",
    "            'Mean Absolute Error': [None, None]\n",
    "           }\n",
    "def _stdxbar(stdxbars, stats):\n",
    "    for k, v in stats.items(): \n",
    "        tensored_samples = tf.convert_to_tensor(v)\n",
    "        xbars = tf.reduce_mean(tensored_samples, axis=0).numpy()\n",
    "        xbar_train, xbar_test = xbars[0], xbars[1] \n",
    "        stds = tf.math.reduce_std(tensored_samples, axis=0).numpy()\n",
    "        std_train, std_test = stds[0], stds[1]\n",
    "        stdxbars[k] = [[xbar_train, std_train], [xbar_test, std_test]]\n",
    "    return stdxbars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "d6c242c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Kullback-Leibler Divergence': [[24.627493, 1.5918127],\n",
       "  [29.433971, 2.3994606]],\n",
       " 'Cross Entropy': [[50.440132, 1.5529034], [53.29887, 2.4224615]],\n",
       " 'Poisson': [[0.21973053, 0.00912408], [0.21816826, 0.008245382]],\n",
       " 'Mean Absolute Error': [[0.0320363, 0.0015654956],\n",
       "  [0.029442465, 0.0013707615]]}"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stdxbars = _stdxbar(results, stats)\n",
    "stdxbars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "bace02c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _pooled(stdxbars, n):\n",
    "    poolingstds = {'Kullback-Leibler Divergence': .0,\n",
    "                  'Cross Entropy': .0,\n",
    "                   'Poisson': .0,\n",
    "                   'Mean Absolute Error': .0\n",
    "                  }\n",
    "    for k, v in stdxbars.items():\n",
    "        trainstd = v[0][1]\n",
    "        teststd = v[1][1]\n",
    "        estimate = np.sqrt((((n-1)*np.square(trainstd))+((n-1)*np.square(teststd)))/(2*(n-1)))\n",
    "        \n",
    "        assertion = False\n",
    "        try:\n",
    "            assert .5 <= trainstd/teststd <= 2., '{metric}: One sample variance cannot be the double of the other!'.format(metric=k)\n",
    "        except AssertionError as e:\n",
    "            assertion = True\n",
    "            print(e)\n",
    "        \n",
    "        # discard metrics that don't pass sample variance check\n",
    "        if assertion == False:\n",
    "            poolingstds[k] = estimate\n",
    "        else:\n",
    "            del poolingstds[k]\n",
    "\n",
    "    return poolingstds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "1b4d04ff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Kullback-Leibler Divergence': 2.036084335079386,\n",
       " 'Cross Entropy': 2.0346779874550855,\n",
       " 'Poisson': 0.008695836910191171,\n",
       " 'Mean Absolute Error': 0.0014713537422005453}"
      ]
     },
     "execution_count": 129,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pooled_estimates = _pooled(stdxbars, n=NUMBER_OF_RUNS)\n",
    "pooled_estimates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "9fd2dd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate 95% CI\n",
    "def _CI(stdxbars, pooled_estimates, n):\n",
    "    zval95 = 1.96\n",
    "    conf_intervals = {'Kullback-Leibler Divergence': .0,\n",
    "                      'Cross Entropy': .0,\n",
    "                      'Poisson': .0,\n",
    "                      'Mean Absolute Error': .0\n",
    "                     }\n",
    "    trainxbar = stdxbars['Kullback-Leibler Divergence'][0][0]\n",
    "    testxbar = stdxbars['Kullback-Leibler Divergence'][1][0]\n",
    "    mean_diff = testxbar-trainxbar\n",
    "    estimate = pooled_estimates['Kullback-Leibler Divergence']\n",
    "    \n",
    "    upper_bound = np.round(mean_diff+(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    lower_bound = np.round(mean_diff-(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    conf_intervals['Kullback-Leibler Divergence'] = [lower_bound, upper_bound]\n",
    "    \n",
    "    trainxbar = stdxbars['Cross Entropy'][0][0]\n",
    "    testxbar = stdxbars['Cross Entropy'][1][0]\n",
    "    mean_diff = testxbar-trainxbar\n",
    "    estimate = pooled_estimates['Cross Entropy']\n",
    "\n",
    "    upper_bound = np.round(mean_diff+(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    lower_bound = np.round(mean_diff-(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    conf_intervals['Cross Entropy'] = [lower_bound, upper_bound]\n",
    "\n",
    "    trainxbar = stdxbars['Poisson'][0][0]\n",
    "    testxbar = stdxbars['Poisson'][1][0]\n",
    "    mean_diff = testxbar-trainxbar\n",
    "    estimate = pooled_estimates['Poisson']\n",
    "    \n",
    "    upper_bound = np.round(mean_diff+(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    lower_bound = np.round(mean_diff-(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    conf_intervals['Poisson'] = [lower_bound, upper_bound]\n",
    "\n",
    "    trainxbar = stdxbars['Mean Absolute Error'][0][0]\n",
    "    testxbar = stdxbars['Mean Absolute Error'][1][0]\n",
    "    mean_diff = testxbar-trainxbar\n",
    "    estimate = pooled_estimates['Mean Absolute Error']\n",
    "    \n",
    "    upper_bound = np.round(mean_diff+(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    lower_bound = np.round(mean_diff-(zval95*estimate*np.sqrt(2/n)),3)\n",
    "    conf_intervals['Mean Absolute Error'] = [lower_bound, upper_bound]\n",
    "    \n",
    "    return conf_intervals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "a13e3f52",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Kullback-Leibler Divergence': [2.283, 7.33],\n",
       " 'Cross Entropy': [0.337, 5.381],\n",
       " 'Poisson': [-0.012, 0.009],\n",
       " 'Mean Absolute Error': [-0.004, -0.001]}"
      ]
     },
     "execution_count": 131,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conf_intervals = _CI(stdxbars, pooled_estimates, n=NUMBER_OF_RUNS)\n",
    "conf_intervals"
   ]
  }
 ],
 "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
}
