examples/comparison_plan_rad_aff.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Flow training and visualization\n",
    "In this notebook, we train normalizing flows to fit predefined prior distributions, testing their expressivity. The plots are generated to visualize the learned distributions for given layers $K$, and the training loss is plotted to compare the expressivity of different flows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "# Import required packages\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "import normflows as nf\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from tqdm import tqdm\n",
    "\n",
    "print(\"PyTorch version: %s\" % torch.__version__)\n",
    "dev = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(\"Using device: %s\" % dev)\n",
    "\n",
    "#z shape is (batch_size, num_samples, dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prior target distribution visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "priors = []\n",
    "priors.append(nf.distributions.TwoModes(2.0, 0.2))\n",
    "priors.append(nf.distributions.Sinusoidal(0.4, 4))\n",
    "priors.append(nf.distributions.Sinusoidal_gap(0.4, 4))\n",
    "priors.append(nf.distributions.Sinusoidal_split(0.4, 4))\n",
    "priors.append(nf.distributions.Smiley(0.15))\n",
    "\n",
    "\n",
    "# Plot prior distributions\n",
    "grid_size = 200\n",
    "grid_length = 4.0\n",
    "grid_shape = ([-grid_length, grid_length], [-grid_length, grid_length])\n",
    "\n",
    "space_mesh = torch.linspace(-grid_length, grid_length, grid_size)\n",
    "xx, yy = torch.meshgrid(space_mesh, space_mesh)\n",
    "z = torch.cat([xx.unsqueeze(2), yy.unsqueeze(2)], 2)\n",
    "z = z.reshape(-1, 2)\n",
    "\n",
    "K_arr = [2, 8, 32]\n",
    "max_iter = 30000\n",
    "batch_size = 512\n",
    "num_samples = 256\n",
    "save_iter = 1000\n",
    "\n",
    "for k in range(len(priors)):\n",
    "    log_prob = priors[k].log_prob(z)\n",
    "    prob = torch.exp(log_prob)\n",
    "\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.pcolormesh(xx, yy, prob.reshape(grid_size, grid_size))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Flow training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "flow_types = (\"Planar\", \"Radial\", \"NICE\", \"RealNVP\")\n",
    "max_iter = 20000\n",
    "batch_size = 1024\n",
    "plot_batches = 10 ** 2\n",
    "plot_samples = 10 ** 4\n",
    "save_iter = 50\n",
    "\n",
    "for name in flow_types:\n",
    "    K_arr = [2, 8, 32]\n",
    "    for K in K_arr:\n",
    "        print(\"Flow type {} with K = {}\".format(name, K))\n",
    "        for k in range(len(priors)):\n",
    "            if k == 0 or k == 4:\n",
    "                anneal_iter = 10000\n",
    "            else: # turn annealing off when fitting to sinusoidal distributions\n",
    "                anneal_iter = 1\n",
    "        \n",
    "            flows = []\n",
    "            b = torch.tensor([0,1])\n",
    "            for i in range(K):\n",
    "                if name == \"Planar\":\n",
    "                    flows += [nf.flows.Planar((2,))]\n",
    "                elif name == \"Radial\":\n",
    "                    flows += [nf.flows.Radial((2,))]\n",
    "                elif name == \"NICE\":\n",
    "                    flows += [nf.flows.MaskedAffineFlow(b, nf.nets.MLP([2, 16, 16, 2], init_zeros=True))]\n",
    "                elif name == \"RealNVP\":\n",
    "                    flows += [nf.flows.MaskedAffineFlow(b, nf.nets.MLP([2, 16, 16, 2], init_zeros=True), \n",
    "                                                        nf.nets.MLP([2, 16, 16, 2], init_zeros=True))]\n",
    "                b = 1-b # parity alternation for mask\n",
    "\n",
    "            q0 = nf.distributions.DiagGaussian(2)\n",
    "            nfm = nf.NormalizingFlow(p=priors[k], q0=q0, flows=flows)\n",
    "            nfm.to(dev) # Move model on GPU if available\n",
    "    \n",
    "            # Train model\n",
    "            loss_hist = np.array([])\n",
    "            log_q_hist = np.array([])\n",
    "            log_p_hist = np.array([])\n",
    "            x = torch.zeros(batch_size, device=dev)\n",
    "\n",
    "            optimizer = torch.optim.Adam(nfm.parameters(), lr=1e-3, weight_decay=1e-3)\n",
    "            for it in tqdm(range(max_iter)):\n",
    "                optimizer.zero_grad()\n",
    "                loss = nfm.reverse_kld(batch_size, np.min([1.0, 0.01 + it / anneal_iter]))\n",
    "                if ~(torch.isnan(loss) | torch.isinf(loss)):\n",
    "                    loss.backward()\n",
    "                    optimizer.step()\n",
    "            \n",
    "                if (it + 1) % save_iter == 0:\n",
    "                    loss_hist = np.append(loss_hist, loss.cpu().data.numpy())\n",
    "\n",
    "            # Plot learned posterior distribution\n",
    "            z_np = np.zeros((0, 2))\n",
    "            for i in range(plot_batches):\n",
    "                z, _ = nfm.sample(plot_samples)\n",
    "                z_np = np.concatenate((z_np, z.cpu().data.numpy()))\n",
    "            plt.figure(figsize=(10, 10))\n",
    "            plt.hist2d(z_np[:, 0], z_np[:, 1], (grid_size, grid_size), grid_shape)\n",
    "            plt.show()\n",
    "            np.save(\"{}-K={}-k={}\".format(name,K,k), (z_np, loss.cpu().data.numpy()))\n",
    "    \n",
    "            # Plot training history\n",
    "            plt.figure(figsize=(10, 10))\n",
    "            plt.plot(loss_hist, label='loss')\n",
    "            plt.legend()\n",
    "            plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Expressivity plot of flows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(14, 10))\n",
    "K_arr = [2, 8, 32]\n",
    "nrows=5\n",
    "ncols=7\n",
    "axes = [ fig.add_subplot(nrows, ncols, r * ncols + c + 1) for r in range(0, nrows) for c in range(0, ncols) ]\n",
    "\n",
    "for ax in axes:\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "\n",
    "grid_size = 100\n",
    "grid_length = 4.0\n",
    "grid_shape = ([-grid_length, grid_length], [-grid_length, grid_length])\n",
    "\n",
    "space_mesh = torch.linspace(-grid_length, grid_length, grid_size)\n",
    "xx, yy = torch.meshgrid(space_mesh, space_mesh)\n",
    "z = torch.cat([xx.unsqueeze(2), yy.unsqueeze(2)], 2)\n",
    "z = z.reshape(-1, 2)\n",
    "axes[0].annotate('Target', xy=(0.5, 1.10), xytext=(0.5, 1.20), xycoords='axes fraction', \n",
    "            fontsize=24, ha='center', va='bottom',\n",
    "            arrowprops=dict(arrowstyle='-[, widthB=1.5, lengthB=0.2', lw=2.0))\n",
    "for k in range(5):\n",
    "    axes[k*ncols].set_ylabel('{}'.format(k+1), rotation=0, fontsize=20, labelpad=15)\n",
    "    log_prob = priors[k].log_prob(z)\n",
    "    prob = torch.exp(log_prob)\n",
    "    axes[k*ncols + 0].pcolormesh(xx, yy, prob.reshape(grid_size, grid_size))\n",
    "\n",
    "\n",
    "for l in range(len(K_arr)):\n",
    "    K = K_arr[l]\n",
    "    if l == 1:\n",
    "        axes[0*ncols + l+1].annotate('Planar flows', xy=(0.5, 1.10), xytext=(0.5, 1.20), xycoords='axes fraction', \n",
    "            fontsize=24, ha='center', va='bottom',\n",
    "            arrowprops=dict(arrowstyle='-[, widthB=6.0, lengthB=0.2', lw=2.0))\n",
    "    axes[4*ncols + l+1].set_xlabel('K = {}'.format(K), fontsize=20)\n",
    "    for k in range(5):\n",
    "        z_np, _ = np.load(\"Planar-K={}-k={}.npy\".format(K,k), allow_pickle=True)\n",
    "        axes[k*ncols + l+1].hist2d(z_np[:, 0], z_np[:, 1], (grid_size, grid_size), grid_shape)\n",
    "        \n",
    "for l in range(len(K_arr)):\n",
    "    K = K_arr[l]\n",
    "    if l == 1:\n",
    "        axes[0*ncols + l+1+len(K_arr)].annotate('Radial flows', xy=(0.5, 1.10), xytext=(0.5, 1.20), xycoords='axes fraction', \n",
    "            fontsize=24, ha='center', va='bottom',\n",
    "            arrowprops=dict(arrowstyle='-[, widthB=6.0, lengthB=0.2', lw=2.0))\n",
    "    axes[4*ncols + l+1+len(K_arr)].set_xlabel('K = {}'.format(K), fontsize=20)\n",
    "    for k in range(5):\n",
    "        z_np, _ = np.load(\"Radial-K={}-k={}.npy\".format(K,k), allow_pickle=True)\n",
    "        axes[k*ncols + l+1+len(K_arr)].hist2d(z_np[:, 0], z_np[:, 1], (grid_size, grid_size), grid_shape)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.02, wspace=0.02)\n",
    "\n",
    "for l in range(1,4):\n",
    "    for k in range(5):\n",
    "        pos1 = axes[k*ncols + l].get_position() # get the original position \n",
    "        pos2 = [pos1.x0 + 0.01, pos1.y0,  pos1.width, pos1.height] \n",
    "        axes[k*ncols + l].set_position(pos2) # set a new position\n",
    "        \n",
    "for l in range(4,7):\n",
    "    for k in range(5):\n",
    "        pos1 = axes[k*ncols + l].get_position() # get the original position \n",
    "        pos2 = [pos1.x0 + 0.02, pos1.y0,  pos1.width, pos1.height] \n",
    "        axes[k*ncols + l].set_position(pos2) # set a new position\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison of Planar, Radial,  and Affine on given prior distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from itertools import repeat\n",
    "\n",
    "k_arr = [0, 2, 4]\n",
    "fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))\n",
    "markers = ['s', 'o', 'v', 'P', 'd']\n",
    "\n",
    "for k in range(len(k_arr)):\n",
    "    loss = [[] for i in repeat(None, len(flow_types))]\n",
    "    for intt, name in enumerate(flow_types):\n",
    "        for K in K_arr:\n",
    "            _, loss_v = np.load(\"{}-K={}-k={}.npy\".format(name,K,k), allow_pickle=True)\n",
    "            loss[intt].append(loss_v)\n",
    "        axes[k].plot(K_arr, loss[intt], marker=markers[intt], label=name)\n",
    "    axes[k].set_title('Target {}'.format(k_arr[k]+1), fontsize=16)\n",
    "    axes[k].set_xlabel('Flow length', fontsize=12)\n",
    "    axes[k].set_ylabel('Variational bound (nats)', fontsize=12)\n",
    "    axes[k].legend()\n",
    "    axes[k].grid('major')\n",
    "\n",
    "fig.tight_layout(pad=2.0)"
   ]
  }
 ],
 "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.11"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}