{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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 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",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "tfd = tfp.distributions\n",
    "tfb = tfp.bijectors\n",
    "tfpl = tfp.layers\n",
    "\n",
    "color_list = ['#bcad', '#dacb']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad5ec2a8",
   "metadata": {},
   "source": [
    "### Define the dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6a51cc0f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-2.0, 2.0)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 33904 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "np.random.seed(42)\n",
    "\n",
    "n_samples = 1000\n",
    "niosy_moons = make_moons(n_samples=n_samples, noise=.05)\n",
    "X, y = niosy_moons\n",
    "X_data = StandardScaler().fit_transform(X)\n",
    "xlim, ylim = [-2, 2], [-2, 2]\n",
    "\n",
    "y_label = y.astype(np.bool)\n",
    "X_train, Y_train = X_data[...,0], X_data[...,1]\n",
    "plt.scatter(X_train[y_label], Y_train[y_label], s=10, color=color_list[0])\n",
    "plt.scatter(X_train[y_label == False], Y_train[y_label == False], s=10, color=color_list[1])\n",
    "plt.legend(['label: 1', 'label: 0'])\n",
    "plt.xlim(xlim)\n",
    "plt.ylim(ylim)"
   ]
  },
  {
   "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: 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=latent_dim) \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": {},
   "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=(2,), 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)                 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": {},
   "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": {},
   "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": "2a055ad4",
   "metadata": {},
   "source": [
    "### Define training and validation sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "dfee364c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# of training samples: 750\n",
      "# of test samples: 250\n"
     ]
    }
   ],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=.25, \n",
    "                                                    shuffle=True, stratify=y,\n",
    "                                                    random_state=42)\n",
    "print(f'# of training samples: {X_train.shape[0]}')\n",
    "print(f'# of test samples: {X_test.shape[0]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b30984a3",
   "metadata": {},
   "source": [
    "### Visualize validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "29025384",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16624 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "X_valid, Y_valid = X_test[...,0], X_test[...,1]\n",
    "plt.scatter(X_valid[y_test == True], Y_valid[y_test == True], s=10, color=color_list[0])\n",
    "plt.scatter(X_valid[y_test == False], Y_valid[y_test == False], s=10, color=color_list[1])\n",
    "plt.legend(['label: 1', 'label: 0'])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce39bd4b",
   "metadata": {},
   "source": [
    "### Compile and fit the model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3215530a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " Epoch 1/100 \n",
      "\t loss: 23.1562, val_loss: 14.0041\n",
      "\n",
      " Epoch 51/100 \n",
      "\t loss: 0.1843, val_loss: 0.1066\n"
     ]
    }
   ],
   "source": [
    "optimizer = Adam(learning_rate=3e-4)\n",
    "epochs = 100 \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 % 50 == 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": 16,
   "id": "116abbd8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15932 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 proportional to KL divergence minimization "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fadadad2",
   "metadata": {},
   "source": [
    "### Sample from the generative model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e65f7315",
   "metadata": {},
   "outputs": [],
   "source": [
    "my_sample = np.concatenate((X_valid, Y_valid)).reshape(2, -1).transpose()\n",
    "sampled = encoder(my_sample).sample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "46a37036",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11204 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(my_sample[...,0], my_sample[..., 1], label='sampled', s=1)\n",
    "plt.show()\n",
    "\n",
    "## learns the distribution very well"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71993523",
   "metadata": {},
   "source": [
    "### Inverse Masked Autoregressive Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ce7944ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "loc = [my_sample[..., 0].mean().astype('float32'), my_sample[..., 1].mean().astype('float32')]\n",
    "scale_diag = [my_sample[..., 0].std().astype('float32'), my_sample[..., 1].std().astype('float32')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "54436897",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.MultivariateNormalDiag 'MultivariateNormalDiag' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvn = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)\n",
    "mvn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "243e4115",
   "metadata": {},
   "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": 22,
   "id": "8c6ff7a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.TransformedDistribution 'invert_masked_autoregressive_flowMultivariateNormalDiag' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainable_dist = tfd.TransformedDistribution(distribution=mvn,\n",
    "                                             bijector=tfb.Invert(masked_autoregressive_flow(\n",
    "                                             activation='sigmoid')))\n",
    "trainable_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "56ce9b96",
   "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": 24,
   "id": "a2ee96bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _plot(results, rows=1, legend=False, plot_color: str=color_list[0]):\n",
    "    cols = int(len(results) / rows)\n",
    "    f, arr = plt.subplots(rows, cols, figsize=(4 * cols, 4 * rows))\n",
    "    i = 0\n",
    "    for r in range(rows):\n",
    "        for c in range(cols):\n",
    "            res = results[i]\n",
    "            X, Y = res[..., 0].numpy(), res[..., 1].numpy()\n",
    "            if rows == 1:\n",
    "                p = arr[c]\n",
    "            else:\n",
    "                p = arr[r, c]\n",
    "            p.scatter(X, Y, s=10, color=plot_color)\n",
    "            p.set_xlim([-5, 5])\n",
    "            p.set_ylim([-5, 5])\n",
    "            p.set_title(names[i])\n",
    "            \n",
    "            i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "1fb98514",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19568 characters]\n",
      "text/plain": [
       "<Figure size 576x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "_plot(samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "dd219342",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, y_train = X_train.astype(np.float32), y_train.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, y_valid = X_test.astype(np.float32), y_test.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": 27,
   "id": "0a671c81",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0...\n",
      "Epoch 50...\n",
      "Epoch 100...\n",
      "Epoch 150...\n",
      "Epoch 200...\n",
      "Epoch 250...\n",
      "Epoch 300...\n",
      "Epoch 350...\n",
      "Epoch 400...\n",
      "Epoch 450...\n",
      "Epoch 500...\n",
      "Epoch 550...\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 600\n",
    "opt = tf.keras.optimizers.Adam()\n",
    "train_losses = []\n",
    "valid_losses = []\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    if epoch % 50 == 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": 28,
   "id": "8e670a4c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16508 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "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": 29,
   "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": 30,
   "id": "e290de78",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19468 characters]\n",
      "text/plain": [
       "<Figure size 576x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "_plot(samples, rows=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "272037c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers = 4\n",
    "flow_bijector = []\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": 32,
   "id": "43c7a1f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainable_dist = tfd.TransformedDistribution(distribution=mvn,\n",
    "                                            bijector=flow_bijector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "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": 34,
   "id": "afbbeed0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62616 characters]\n",
      "text/plain": [
       "<Figure size 1152x576 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "_plot(samples, rows=2, plot_color=color_list[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "afad9e11",
   "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_data), 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",
    "        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": 36,
   "id": "c88164f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 68632 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": 38,
   "id": "ef7ef2b6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0...\n",
      "Epoch 200...\n",
      "Epoch 400...\n"
     ]
    }
   ],
   "source": [
    "num_epochs = 600\n",
    "opt = tf.keras.optimizers.Adam()\n",
    "train_losses = []\n",
    "valid_losses = []\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    if epoch % 200 == 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": 39,
   "id": "886ee70d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 17516 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": 40,
   "id": "4177a424",
   "metadata": {},
   "outputs": [],
   "source": [
    "names, samples = make_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "c0ef6b68",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50124 characters]\n",
      "text/plain": [
       "<Figure size 1152x576 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "_plot(samples, rows=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "e5add6a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_training_data_updated(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_data), 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",
    "        arr[i].set_title(names[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "a8ec530d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 75204 characters]\n",
      "text/plain": [
       "<Figure size 1440x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_training_data_updated(samples, plot_color=color_list[1])"
   ]
  }
 ],
 "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
}
