{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 160,
   "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, BatchNormalization, LayerNormalization, Layer, Dense, ReLU\n",
    "from tensorflow.keras import Model\n",
    "\n",
    "from tensorflow.keras import Model\n",
    "from tensorflow.keras.callbacks import LambdaCallback\n",
    "from tensorflow.python.keras import backend\n",
    "from tensorflow.python.keras.engine import base_layer_utils\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": 161,
   "id": "e96e5477",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-2.0, 2.0)"
      ]
     },
     "execution_count": 161,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 34472 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": 162,
   "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": 179,
   "id": "7205510b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tfp.distributions.Sample 'SampleNormal' batch_shape=[] event_shape=[2] dtype=float32>"
      ]
     },
     "execution_count": 179,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvn = tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2])\n",
    "mvn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b69c3972",
   "metadata": {},
   "source": [
    "### RealNVP initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 190,
   "id": "157a5271",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NN(Layer):\n",
    "    def __init__(self, input_shape, n_hidden=[512, 512], activation=\"relu\", name=\"nn\"):\n",
    "        super(NN, self).__init__(name=\"nn\")\n",
    "        layer_list = []\n",
    "        for i, hidden in enumerate(n_hidden):\n",
    "            layer_list.append(Dense(hidden, activation=activation, name='dense_{}_1'.format(i)))\n",
    "            layer_list.append(Dense(hidden, activation=activation, name='dense_{}_2'.format(i)))\n",
    "        self.layer_list = layer_list\n",
    "        self.log_s_layer = Dense(input_shape, activation=\"tanh\", name='log_s')\n",
    "        self.t_layer = Dense(input_shape, name='t')\n",
    "\n",
    "    def call(self, x):\n",
    "        y = x\n",
    "        for layer in self.layer_list:\n",
    "            y = layer(y)\n",
    "        log_s = self.log_s_layer(y)\n",
    "        t = self.t_layer(y)\n",
    "        return log_s, t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 191,
   "id": "6fa5477a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "class RealNVP(tfb.Bijector):\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        input_shape,\n",
    "        n_hidden=[512, 512],\n",
    "        forward_min_event_ndims=1,\n",
    "        validate_args: bool = False,\n",
    "        name=\"real_nvp\",\n",
    "    ):\n",
    "        super(RealNVP, self).__init__(\n",
    "            validate_args=validate_args, forward_min_event_ndims=forward_min_event_ndims, name=name\n",
    "        )\n",
    "\n",
    "        assert input_shape[-1] % 2 == 0\n",
    "        self.input_shape = input_shape\n",
    "        nn_layer = NN(input_shape[-1] // 2, n_hidden)\n",
    "        nn_input_shape = input_shape.copy()\n",
    "        nn_input_shape[-1] = input_shape[-1] // 2\n",
    "        x = tf.keras.Input(nn_input_shape)\n",
    "        log_s, t = nn_layer(x)\n",
    "        self.nn = Model(x, [log_s, t], name=\"nn\")\n",
    "        \n",
    "    def _forward(self, x):\n",
    "        x_a, x_b = tf.split(x, 2, axis=-1)\n",
    "        y_b = x_b\n",
    "        log_s, t = self.nn(x_b)\n",
    "        s = tf.exp(log_s)\n",
    "        y_a = s * x_a + t\n",
    "        y = tf.concat([y_a, y_b], axis=-1)\n",
    "        return y\n",
    "\n",
    "    def _inverse(self, y):\n",
    "        y_a, y_b = tf.split(y, 2, axis=-1)\n",
    "        x_b = y_b\n",
    "        log_s, t = self.nn(y_b)\n",
    "        s = tf.exp(log_s)\n",
    "        x_a = (y_a - t) / s\n",
    "        x = tf.concat([x_a, x_b], axis=-1)\n",
    "        return x\n",
    "\n",
    "    def _forward_log_det_jacobian(self, x):\n",
    "        _, x_b = tf.split(x, 2, axis=-1)\n",
    "        log_s, t = self.nn(x_b)\n",
    "        return log_s"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2bb7fcd",
   "metadata": {},
   "source": [
    "### Transformed distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "f776f68d",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainable_dist = tfd.TransformedDistribution(distribution=mvn,\n",
    "                                             bijector=RealNVP(input_shape=[2], n_hidden=[16,16])\n",
    "                                            )\n",
    "# assure variables are tracked\n",
    "assert trainable_dist.trainable_variables != []"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "923ac546",
   "metadata": {},
   "source": [
    "### Initialize samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "id": "8a2aa1ea",
   "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": "markdown",
   "id": "84cfabcf",
   "metadata": {},
   "source": [
    "### Plot routine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 195,
   "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": 196,
   "id": "64f81179",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 31632 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": 197,
   "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": 198,
   "id": "f29033e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 19628 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": 201,
   "id": "fe88f36c",
   "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": 202,
   "id": "936bde0c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30952 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": 204,
   "id": "892d601a",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers = 8\n",
    "flow_bijector = []\n",
    "\n",
    "for i in range(num_layers):\n",
    "    flow_i = RealNVP(input_shape=[2], n_hidden=[512,512])\n",
    "    flow_bijector.append(flow_i)  \n",
    "    flow_bijector.append(tfb.Permute(permutation=[1,0]))\n",
    "flow_bijector = tfb.Chain(list(reversed(flow_bijector[:-1])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "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": 207,
   "id": "e73726cb",
   "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": 208,
   "id": "6a507799",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 164856 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": 209,
   "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": 210,
   "id": "4791a426",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 67996 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": 211,
   "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": 212,
   "id": "f8e8ea4d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39780 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": 213,
   "id": "7bc43dae",
   "metadata": {},
   "outputs": [],
   "source": [
    "names, samples = make_samples()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 214,
   "id": "41c676d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 107552 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": 215,
   "id": "ffe5a97d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69664 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
}
