{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "812e1cf0",
   "metadata": {},
   "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/preprocessingv1.0.ipynb'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9f0b013a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \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 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",
    "\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": {},
   "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": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prior event shape: 16\n",
      "# of Gaussions: 16\n",
      "Covariance matrix: MultivariateNormalDiag\n"
     ]
    }
   ],
   "source": [
    "latent_dim = len(X_train.columns)\n",
    "input_shape = len(X_train.columns)\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": {},
   "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": {},
   "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": {},
   "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": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense (Dense)                (None, 256)               4352      \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, 152)               5016      \n",
      "_________________________________________________________________\n",
      "multivariate_normal_tri_l (M multiple                  512       \n",
      "=================================================================\n",
      "Total params: 53,112\n",
      "Trainable params: 53,112\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": {},
   "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": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_1\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense_5 (Dense)              (None, 5)                 85        \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, 152)               39064     \n",
      "_________________________________________________________________\n",
      "multivariate_normal_tri_l_1  multiple                  0         \n",
      "=================================================================\n",
      "Total params: 80,877\n",
      "Trainable params: 80,877\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": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "dense_input (InputLayer)     [(None, 16)]              0         \n",
      "_________________________________________________________________\n",
      "dense (Dense)                (None, 256)               4352      \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, 152)               5016      \n",
      "_________________________________________________________________\n",
      "multivariate_normal_tri_l (M multiple                  512       \n",
      "_________________________________________________________________\n",
      "sequential_1 (Sequential)    multiple                  80877     \n",
      "=================================================================\n",
      "Total params: 133,989\n",
      "Trainable params: 133,989\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": {},
   "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": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# of training samples: 300\n",
      "# of test samples: 148\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": "e176949b",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train_prior_visual = pd.DataFrame(scaler.inverse_transform(X_train), columns=X_train.columns).astype('int64')\n",
    "y_train.reset_index(drop=True, inplace=True)\n",
    "X_train_prior_visual['Risk1Yr'] = y_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1c742030",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 70632 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": 16,
   "id": "b996f2ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 80252 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": 17,
   "id": "3215530a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Epoch 1/1000 \n",
      "\t loss: 412.1550, val_loss: 406.0995\n",
      "\n",
      " Epoch 101/1000 \n",
      "\t loss: 34.0448, val_loss: 35.2328\n",
      "\n",
      " Epoch 201/1000 \n",
      "\t loss: 13.7975, val_loss: 14.0460\n",
      "\n",
      " Epoch 301/1000 \n",
      "\t loss: 6.4840, val_loss: 5.9566\n",
      "\n",
      " Epoch 401/1000 \n",
      "\t loss: 3.3904, val_loss: 3.6168\n",
      "\n",
      " Epoch 501/1000 \n",
      "\t loss: 2.1311, val_loss: 2.2261\n",
      "\n",
      " Epoch 601/1000 \n",
      "\t loss: 1.4892, val_loss: 1.8160\n",
      "\n",
      " Epoch 701/1000 \n",
      "\t loss: 0.9465, val_loss: 0.9693\n",
      "\n",
      " Epoch 801/1000 \n",
      "\t loss: 0.6962, val_loss: 1.5155\n",
      "\n",
      " Epoch 901/1000 \n",
      "\t loss: 0.7830, val_loss: 0.6719\n"
     ]
    }
   ],
   "source": [
    "optimizer = Adam(learning_rate=3e-4)\n",
    "epochs = 1000 \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": 18,
   "id": "116abbd8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 14748 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": 126,
   "id": "e65f7315",
   "metadata": {
    "scrolled": true
   },
   "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": 127,
   "id": "c5ebde09",
   "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>DGN</th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "      <th>PRE6</th>\n",
       "      <th>PRE7</th>\n",
       "      <th>PRE8</th>\n",
       "      <th>PRE9</th>\n",
       "      <th>PRE10</th>\n",
       "      <th>PRE11</th>\n",
       "      <th>PRE14</th>\n",
       "      <th>PRE17</th>\n",
       "      <th>PRE19</th>\n",
       "      <th>PRE25</th>\n",
       "      <th>PRE30</th>\n",
       "      <th>PRE32</th>\n",
       "      <th>AGE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2.407982</td>\n",
       "      <td>-0.831916</td>\n",
       "      <td>2.002110</td>\n",
       "      <td>-0.153428</td>\n",
       "      <td>0.954857</td>\n",
       "      <td>0.083955</td>\n",
       "      <td>0.436395</td>\n",
       "      <td>-2.078765</td>\n",
       "      <td>-1.488115</td>\n",
       "      <td>-0.231488</td>\n",
       "      <td>-0.061668</td>\n",
       "      <td>-0.949218</td>\n",
       "      <td>1.709186</td>\n",
       "      <td>-0.830009</td>\n",
       "      <td>2.592370</td>\n",
       "      <td>0.038623</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2.054684</td>\n",
       "      <td>-2.579462</td>\n",
       "      <td>0.267117</td>\n",
       "      <td>-2.559943</td>\n",
       "      <td>0.820520</td>\n",
       "      <td>1.518479</td>\n",
       "      <td>-0.525705</td>\n",
       "      <td>-0.696743</td>\n",
       "      <td>0.276630</td>\n",
       "      <td>0.158306</td>\n",
       "      <td>-1.480168</td>\n",
       "      <td>0.939017</td>\n",
       "      <td>-1.205487</td>\n",
       "      <td>2.143132</td>\n",
       "      <td>0.010564</td>\n",
       "      <td>0.333124</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-1.431643</td>\n",
       "      <td>-3.029876</td>\n",
       "      <td>0.148751</td>\n",
       "      <td>0.273040</td>\n",
       "      <td>0.946856</td>\n",
       "      <td>-2.778357</td>\n",
       "      <td>-1.209445</td>\n",
       "      <td>-0.415556</td>\n",
       "      <td>1.428541</td>\n",
       "      <td>1.260927</td>\n",
       "      <td>-0.324407</td>\n",
       "      <td>0.132441</td>\n",
       "      <td>-0.229503</td>\n",
       "      <td>0.225607</td>\n",
       "      <td>0.491041</td>\n",
       "      <td>-0.748245</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.281104</td>\n",
       "      <td>-0.509333</td>\n",
       "      <td>1.997721</td>\n",
       "      <td>-1.335895</td>\n",
       "      <td>0.376429</td>\n",
       "      <td>0.630013</td>\n",
       "      <td>-1.233507</td>\n",
       "      <td>-0.079674</td>\n",
       "      <td>-1.723972</td>\n",
       "      <td>-2.168686</td>\n",
       "      <td>2.634992</td>\n",
       "      <td>0.437542</td>\n",
       "      <td>0.865715</td>\n",
       "      <td>-0.459556</td>\n",
       "      <td>0.774192</td>\n",
       "      <td>0.524889</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.376555</td>\n",
       "      <td>-1.034092</td>\n",
       "      <td>0.058366</td>\n",
       "      <td>1.709675</td>\n",
       "      <td>0.435779</td>\n",
       "      <td>-0.455206</td>\n",
       "      <td>-2.279676</td>\n",
       "      <td>0.541935</td>\n",
       "      <td>0.410181</td>\n",
       "      <td>0.310724</td>\n",
       "      <td>-0.937057</td>\n",
       "      <td>1.826935</td>\n",
       "      <td>0.648489</td>\n",
       "      <td>-0.013052</td>\n",
       "      <td>0.119601</td>\n",
       "      <td>-1.192695</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "        DGN      PRE4      PRE5      PRE6      PRE7      PRE8      PRE9  \\\n",
       "0  2.407982 -0.831916  2.002110 -0.153428  0.954857  0.083955  0.436395   \n",
       "1  2.054684 -2.579462  0.267117 -2.559943  0.820520  1.518479 -0.525705   \n",
       "2 -1.431643 -3.029876  0.148751  0.273040  0.946856 -2.778357 -1.209445   \n",
       "3 -0.281104 -0.509333  1.997721 -1.335895  0.376429  0.630013 -1.233507   \n",
       "4  0.376555 -1.034092  0.058366  1.709675  0.435779 -0.455206 -2.279676   \n",
       "\n",
       "      PRE10     PRE11     PRE14     PRE17     PRE19     PRE25     PRE30  \\\n",
       "0 -2.078765 -1.488115 -0.231488 -0.061668 -0.949218  1.709186 -0.830009   \n",
       "1 -0.696743  0.276630  0.158306 -1.480168  0.939017 -1.205487  2.143132   \n",
       "2 -0.415556  1.428541  1.260927 -0.324407  0.132441 -0.229503  0.225607   \n",
       "3 -0.079674 -1.723972 -2.168686  2.634992  0.437542  0.865715 -0.459556   \n",
       "4  0.541935  0.410181  0.310724 -0.937057  1.826935  0.648489 -0.013052   \n",
       "\n",
       "      PRE32       AGE  \n",
       "0  2.592370  0.038623  \n",
       "1  0.010564  0.333124  \n",
       "2  0.491041 -0.748245  \n",
       "3  0.774192  0.524889  \n",
       "4  0.119601 -1.192695  "
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_train_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "id": "e7bd3f85",
   "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>DGN</th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "      <th>PRE6</th>\n",
       "      <th>PRE7</th>\n",
       "      <th>PRE8</th>\n",
       "      <th>PRE9</th>\n",
       "      <th>PRE10</th>\n",
       "      <th>PRE11</th>\n",
       "      <th>PRE14</th>\n",
       "      <th>PRE17</th>\n",
       "      <th>PRE19</th>\n",
       "      <th>PRE25</th>\n",
       "      <th>PRE30</th>\n",
       "      <th>PRE32</th>\n",
       "      <th>AGE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>27</td>\n",
       "      <td>96</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>24</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>-19</td>\n",
       "      <td>51</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>27</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>-32</td>\n",
       "      <td>48</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>18</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>35</td>\n",
       "      <td>96</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>28</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>21</td>\n",
       "      <td>46</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>14</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   DGN  PRE4  PRE5  PRE6  PRE7  PRE8  PRE9  PRE10  PRE11  PRE14  PRE17  PRE19  \\\n",
       "0    2    27    96     0     0     0     0      0      0      0      0      0   \n",
       "1    1   -19    51     0     0     0     0      0      0      0      0      0   \n",
       "2    0   -32    48     0     0     0     0      0      0      1      0      0   \n",
       "3    0    35    96     0     0     0     0      0      0      0      0      0   \n",
       "4    1    21    46     1     0     0     0      0      0      0      0      0   \n",
       "\n",
       "   PRE25  PRE30  PRE32  AGE  \n",
       "0      0      0      0   24  \n",
       "1      0      1      0   27  \n",
       "2      0      0      0   18  \n",
       "3      0      0      0   28  \n",
       "4      0      0      0   14  "
      ]
     },
     "execution_count": 128,
     "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 = X_train_sample.astype('int64')\n",
    "X_train_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "id": "39439ca5",
   "metadata": {},
   "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>DGN</th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "      <th>PRE6</th>\n",
       "      <th>PRE7</th>\n",
       "      <th>PRE8</th>\n",
       "      <th>PRE9</th>\n",
       "      <th>PRE10</th>\n",
       "      <th>PRE11</th>\n",
       "      <th>PRE14</th>\n",
       "      <th>PRE17</th>\n",
       "      <th>PRE19</th>\n",
       "      <th>PRE25</th>\n",
       "      <th>PRE30</th>\n",
       "      <th>PRE32</th>\n",
       "      <th>AGE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>20</td>\n",
       "      <td>30</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>51</td>\n",
       "      <td>68</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>18</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>48</td>\n",
       "      <td>41</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>15</td>\n",
       "      <td>58</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>-1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>19</td>\n",
       "      <td>27</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>38</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   DGN  PRE4  PRE5  PRE6  PRE7  PRE8  PRE9  PRE10  PRE11  PRE14  PRE17  PRE19  \\\n",
       "0    0    20    30     0     0     0     0      0      0      0      0      0   \n",
       "1    0    51    68     0     0     0     0      0      0      0      0      0   \n",
       "2    1    48    41     0     0     0     0      0      0      0      0      0   \n",
       "3    0    15    58     0     0     0     0      1      0     -1      0      0   \n",
       "4    1    19    27     0     0     0     0      0      0      0      0      0   \n",
       "\n",
       "   PRE25  PRE30  PRE32  AGE  \n",
       "0      0      1      0   17  \n",
       "1      0      0      0   18  \n",
       "2      0      1      0   19  \n",
       "3      0      1      0   14  \n",
       "4      0      0      0   38  "
      ]
     },
     "execution_count": 129,
     "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_train.columns)\n",
    "\n",
    "X_test_sample = pd.DataFrame(scaler.inverse_transform(X_test_sample.to_numpy()), columns=X_test.columns)\n",
    "X_test_sample = X_test_sample.astype('int64')\n",
    "X_test_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "id": "d57c8784",
   "metadata": {},
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "c661a8d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57000 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": 132,
   "id": "979b9fce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61308 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, 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": "code",
   "execution_count": 133,
   "id": "2a20c845",
   "metadata": {},
   "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>DGN</th>\n",
       "      <th>PRE4</th>\n",
       "      <th>PRE5</th>\n",
       "      <th>PRE6</th>\n",
       "      <th>PRE7</th>\n",
       "      <th>PRE8</th>\n",
       "      <th>PRE9</th>\n",
       "      <th>PRE10</th>\n",
       "      <th>PRE11</th>\n",
       "      <th>PRE14</th>\n",
       "      <th>PRE17</th>\n",
       "      <th>PRE19</th>\n",
       "      <th>PRE25</th>\n",
       "      <th>PRE30</th>\n",
       "      <th>PRE32</th>\n",
       "      <th>AGE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>69</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>26</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>46</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>9</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>41</td>\n",
       "      <td>24</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>49</td>\n",
       "      <td>42</td>\n",
       "      <td>4</td>\n",
       "      <td>-2</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>-6</td>\n",
       "      <td>-1</td>\n",
       "      <td>7</td>\n",
       "      <td>-1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>5</td>\n",
       "      <td>0</td>\n",
       "      <td>-47</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>36</td>\n",
       "      <td>27</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>11</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   DGN  PRE4  PRE5  PRE6  PRE7  PRE8  PRE9  PRE10  PRE11  PRE14  PRE17  PRE19  \\\n",
       "0    1     7    69     1     0     2     0      1      0      4      0      0   \n",
       "1    1    46     0     0     0     0     0      1      0      0      0      0   \n",
       "2    1    41    24     0     0     0     0      0      0      1      0      0   \n",
       "3    1    49    42     4    -2     0     0     -6     -1      7     -1      0   \n",
       "4    1    36    27     0     0     1     0      2      0      0      1      0   \n",
       "\n",
       "   PRE25  PRE30  PRE32  AGE  \n",
       "0      0      0      0   26  \n",
       "1      0      0      0    9  \n",
       "2      0      0      0   11  \n",
       "3      0      5      0  -47  \n",
       "4      0      0      0   11  "
      ]
     },
     "execution_count": 133,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# decode \n",
    "X_train_sample.drop('Risk1Yr',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 = X_train_sample.astype('int64')\n",
    "X_train_sample.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "c0d87cf9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61368 characters]\n",
      "text/plain": [
       "<Figure size 763.125x360 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "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": 135,
   "id": "8f847af5",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train_sample.drop('Risk1Yr', axis=1, inplace=True) # drop target, was only intended for VAE sampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "ce7944ce",
   "metadata": {},
   "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": 137,
   "id": "54436897",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.MultivariateNormalDiag 'MultivariateNormalDiag' batch_shape=[] event_shape=[16] dtype=float32>"
      ]
     },
     "execution_count": 137,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvn = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)\n",
    "mvn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "243e4115",
   "metadata": {},
   "outputs": [],
   "source": [
    "def masked_autoregressive_flow(hidden_units=[16,16], event_shape=[16], 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": 139,
   "id": "8c6ff7a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.TransformedDistribution 'masked_autoregressive_flowMultivariateNormalDiag' batch_shape=[] event_shape=[16] dtype=float32>"
      ]
     },
     "execution_count": 139,
     "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": 140,
   "id": "56ce9b96",
   "metadata": {},
   "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": 141,
   "id": "bded8729",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train_np = X_train_sample.to_numpy()\n",
    "X_test_np = X_test_sample.to_numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "id": "da26e4e2",
   "metadata": {},
   "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": 146,
   "id": "0a671c81",
   "metadata": {},
   "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(128)\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(128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "id": "d9971446",
   "metadata": {},
   "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": 148,
   "id": "8e670a4c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15508 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": 149,
   "id": "0e24f956",
   "metadata": {},
   "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": 162,
   "id": "272037c9",
   "metadata": {},
   "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(init_once(x=np.random.permutation(16).astype('int'))))\n",
    "# discard the last permute layer \n",
    "flow_bijector = tfb.Chain(list(reversed(flow_bijector[:-1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "43c7a1f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainable_dist = tfd.TransformedDistribution(distribution=mvn,\n",
    "                                            bijector=flow_bijector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "fb2d1a9a",
   "metadata": {},
   "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\n",
    "\n",
    "names, samples = make_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "ef7ef2b6",
   "metadata": {},
   "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": 166,
   "id": "886ee70d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15840 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')\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "4177a424",
   "metadata": {},
   "outputs": [],
   "source": [
    "names, samples = make_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "id": "ec7a463d",
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_nf = pd.DataFrame(scaler.inverse_transform(samples[0]).astype(np.int64),columns=X_test.columns)\n",
    "samples_nf['Risk1Yr'] = y_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "id": "ccaff6ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62964 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",
    "sns.swarmplot(x=\"DGN\", y='PRE4', color=\"k\", size=3, data=samples_nf, ax=g.ax)\n",
    "sns.despine()"
   ]
  }
 ],
 "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
}
