{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5394ec2c",
   "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",
    "import tensorflow as tf \n",
    "import tensorflow_probability as tfp \n",
    "from tensorflow.keras.layers import Input\n",
    "from tensorflow.keras import Model\n",
    "from tensorflow.keras.callbacks import LambdaCallback\n",
    "\n",
    "tfd = tfp.distributions\n",
    "tfb = tfp.bijectors\n",
    "\n",
    "color_list = ['#bcad', '#dacb']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "021893ba",
   "metadata": {},
   "source": [
    "### Define the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e96e5477",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-2.0, 2.0)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 34480 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_samples = 1000\n",
    "split_index = int(n_samples*.75) # 75, 25 split\n",
    "noisy_moons = make_moons(n_samples=n_samples, noise=.05)\n",
    "X, y = noisy_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": "code",
   "execution_count": 3,
   "id": "567da5b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, y_train = X_data[:split_index].astype(np.float32), y[:split_index].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_data[split_index:].astype(np.float32), y[split_index:].astype(np.float32)\n",
    "X_valid = tf.data.Dataset.from_tensor_slices(X_valid)\n",
    "X_valid = X_valid.batch(128)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7989c77",
   "metadata": {},
   "source": [
    "### Define the base distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7205510b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.Sample 'SampleNormal' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvn = tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2])\n",
    "mvn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91c03525",
   "metadata": {},
   "source": [
    "### MAF initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f3f9ffa1",
   "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": "markdown",
   "id": "d2bb7fcd",
   "metadata": {},
   "source": [
    "### Transformed distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f776f68d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.TransformedDistribution 'invert_masked_autoregressive_flowSampleNormal' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 6,
     "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": "markdown",
   "id": "923ac546",
   "metadata": {},
   "source": [
    "### Initialize samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8a2aa1ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = mvn.sample(sample_shape=(n_samples, 2))\n",
    "names = [mvn.name, trainable_dist.bijector.name]\n",
    "samples = [x, trainable_dist.bijector.forward(x)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84cfabcf",
   "metadata": {},
   "source": [
    "### Plot routine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6071af1c",
   "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": "markdown",
   "id": "129b234d",
   "metadata": {},
   "source": [
    "### Prior training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "64f81179",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40024 characters]\n",
      "text/plain": [
       "<Figure size 576x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "_plot(samples)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6260147",
   "metadata": {},
   "source": [
    "### Train the distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "63463772",
   "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": 11,
   "id": "f29033e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20508 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": "markdown",
   "id": "7be786e2",
   "metadata": {},
   "source": [
    "### After training -- w/o stacking & permutation & batch normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fe88f36c",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = mvn.sample(sample_shape=(n_samples,2))\n",
    "names = [mvn.name, trainable_dist.bijector.name]\n",
    "samples = [x, trainable_dist.bijector.forward(x)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "936bde0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31164 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": "markdown",
   "id": "dec4652e",
   "metadata": {},
   "source": [
    "### Complex model -- apply bijectors in reverse order"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "892d601a",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers = 8\n",
    "flow_bijector = []\n",
    "\n",
    "for i in range(num_layers):\n",
    "    flow_i = tfb.Invert(masked_autoregressive_flow(hidden_units=[512,512]))\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": 19,
   "id": "174b9850",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainable_dist = tfd.TransformedDistribution(distribution=mvn,\n",
    "                                            bijector=flow_bijector)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "015eb594",
   "metadata": {},
   "source": [
    "### Generate samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "e73726cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_samples():\n",
    "    x = mvn.sample((n_samples, 2))\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": 21,
   "id": "6a507799",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 185720 characters]\n",
      "text/plain": [
       "<Figure size 2304x576 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "_plot(samples, rows=2, plot_color=color_list[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "17d3488a",
   "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": "markdown",
   "id": "11c7b853",
   "metadata": {},
   "source": [
    "### Visualization -- before training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "4791a426",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88256 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": 24,
   "id": "c2adde55",
   "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": 25,
   "id": "f8e8ea4d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 27568 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": 26,
   "id": "7bc43dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "names, samples = make_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "41c676d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 91320 characters]\n",
      "text/plain": [
       "<Figure size 2304x576 with 16 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "_plot(samples, rows=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ffe5a97d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 85920 characters]\n",
      "text/plain": [
       "<Figure size 1440x360 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "visualize_training_data(samples)"
   ]
  }
 ],
 "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
}
