notebooks/PartialFlows.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison of Partial and Loc-scale Flows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "import numpy as np\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import edward2 as ed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../mdnf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import aux\n",
    "import flows\n",
    "import flows_edward2 as fed\n",
    "import flows_edward2_made as made\n",
    "from flows_factorized import DiscreteFactorizedFlowPartial\n",
    "\n",
    "from flows_transformations import CopiableMADELocScale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# can be run as a script with args in format KEY=VAL,KEY=[STRVAL],...\n",
    "args = aux.parse_script_args() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARTIAL_FLOWS = bool(args.get(\"PARTIAL_FLOWS\", 1)) # 1 = use partial flows / 0 = use loc-scale flows\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Results output file: PARTIAL_5_1_10_12411.csv\n"
     ]
    }
   ],
   "source": [
    "SEED = args.get(\"SEED\", 12411)\n",
    "K = args.get(\"K\", 5)\n",
    "NSAMPLES = args.get(\"NSAMPLES\", 1000)\n",
    "NITER = args.get(\"NITER\", 10000)\n",
    "t = args.get(\"TEMP\", 1.0)\n",
    "B = args.get(\"B\", 10)\n",
    "\n",
    "OUT = args.get(\"OUT\", \"PARTIAL_%s_%s_%s_%s.csv\" % (K, int(PARTIAL_FLOWS), B, SEED))\n",
    "print(\"Results output file: %s\" % OUT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Target distribution (sorted)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(SEED)\n",
    "tf.random.set_seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.07, 0.13, 0.2 , 0.27, 0.33]], dtype=float32)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 'true' data generating distribution -> ordered probs\n",
    "probs = np.array(np.arange(1,K+1), dtype=float)\n",
    "probs /= np.sum(probs)\n",
    "\n",
    "target = tfp.distributions.OneHotCategorical(probs = [probs])\n",
    "np.round(target.probs, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Base distibution (shuffled)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Base probs: [[0.2  0.07 0.33 0.27 0.13]]\n"
     ]
    }
   ],
   "source": [
    "source_probs = np.array(probs[-1: :-1])\n",
    "np.random.shuffle(source_probs)\n",
    "base = tfp.distributions.OneHotCategorical(probs = [source_probs])\n",
    "print(\"Base probs: %s\" % np.round(base.probs, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Flows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(SEED)\n",
    "tf.random.set_seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "if PARTIAL_FLOWS:\n",
    "    # Partial flows\n",
    "    category_pairs = [[i,i+1] for u in range(K-1,0,-1) for i in range(0,u)]    \n",
    "    layers = [DiscreteFactorizedFlowPartial(1,K, pair, temperature=t) for pair in category_pairs]\n",
    "else:\n",
    "    # Loc-scale Tran's flows\n",
    "    layers = []\n",
    "    for _ in range(B):\n",
    "        made_class = CopiableMADELocScale\n",
    "        network = made_class(K*2, hidden_dims=[], hidden_order=\"left-to-right\") \n",
    "        flow = ed.layers.DiscreteAutoregressiveFlow(network, temperature=t)\n",
    "        layers.append(flow)\n",
    "    #layers = [DiscreteFactorizedFlowLocScale(1, K, temperature=t) for _ in range(B)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 10 layers: [<flows_factorized.DiscreteFactorizedFlowPartial object at 0x7f4242277d10>, <flows_factorized.Discre...\n",
      "Initial probs: tf.Tensor([[0.069 0.186 0.329 0.262 0.154]], shape=(1, 5), dtype=float32)\n"
     ]
    }
   ],
   "source": [
    "print(\"Using %s layers: %s...\" % (len(layers), str(layers)[:100]))\n",
    "flow = flows.DiscreteFlow(layers=layers)\n",
    "\n",
    "# initial output distribution of a flow\n",
    "outprobs = tf.reduce_mean(flow(tf.cast(base.sample(NSAMPLES), 'float32')), 0)\n",
    "print(\"Initial probs: %s\" % outprobs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.0s] iter=0 loss=1.624 recovered=[[0.07 0.2  0.33 0.26 0.14]]\n",
      "[0.1s] iter=1 loss=1.672 recovered=[[0.2  0.08 0.32 0.13 0.27]]\n",
      "[0.2s] iter=2 loss=1.564 recovered=[[0.06 0.2  0.33 0.14 0.27]]\n",
      "[0.3s] iter=3 loss=1.605 recovered=[[0.06 0.21 0.33 0.13 0.27]]\n",
      "[0.4s] iter=4 loss=1.616 recovered=[[0.07 0.31 0.21 0.14 0.27]]\n",
      "[0.5s] iter=5 loss=1.627 recovered=[[0.06 0.34 0.18 0.15 0.26]]\n",
      "[0.6s] iter=6 loss=1.663 recovered=[[0.06 0.34 0.2  0.28 0.13]]\n",
      "[0.7s] iter=7 loss=1.690 recovered=[[0.06 0.33 0.2  0.29 0.12]]\n",
      "[0.8s] iter=8 loss=1.682 recovered=[[0.06 0.34 0.19 0.27 0.14]]\n",
      "[0.9s] iter=9 loss=1.732 recovered=[[0.32 0.06 0.22 0.16 0.25]]\n",
      "[0.9s] iter=10 loss=1.688 recovered=[[0.22 0.07 0.29 0.15 0.27]]\n",
      "[1.8s] iter=20 loss=1.659 recovered=[[0.15 0.19 0.05 0.32 0.28]]\n",
      "[2.8s] iter=30 loss=1.527 recovered=[[0.07 0.2  0.14 0.31 0.27]]\n",
      "[3.6s] iter=40 loss=1.536 recovered=[[0.08 0.2  0.12 0.34 0.26]]\n",
      "[4.5s] iter=50 loss=1.817 recovered=[[0.33 0.21 0.07 0.12 0.27]]\n",
      "[5.3s] iter=60 loss=1.797 recovered=[[0.35 0.19 0.07 0.14 0.25]]\n",
      "[6.1s] iter=70 loss=1.717 recovered=[[0.21 0.24 0.07 0.12 0.36]]\n",
      "[6.9s] iter=80 loss=1.781 recovered=[[0.14 0.27 0.32 0.08 0.2 ]]\n",
      "[7.7s] iter=90 loss=1.513 recovered=[[0.07 0.18 0.15 0.28 0.32]]\n",
      "[8.6s] iter=100 loss=1.526 recovered=[[0.07 0.2  0.15 0.31 0.27]]\n",
      "[9.5s] iter=110 loss=1.589 recovered=[[0.06 0.18 0.33 0.14 0.28]]\n",
      "[10.3s] iter=120 loss=1.903 recovered=[[0.29 0.33 0.12 0.06 0.21]]\n",
      "[11.2s] iter=130 loss=1.561 recovered=[[0.14 0.06 0.27 0.21 0.32]]\n",
      "[12.0s] iter=140 loss=1.791 recovered=[[0.22 0.31 0.12 0.07 0.27]]\n",
      "[12.8s] iter=150 loss=1.674 recovered=[[0.06 0.37 0.24 0.12 0.2 ]]\n",
      "[13.6s] iter=160 loss=1.474 recovered=[[0.07 0.14 0.19 0.34 0.26]]\n",
      "[14.4s] iter=170 loss=1.830 recovered=[[0.33 0.16 0.22 0.07 0.23]]\n",
      "[15.3s] iter=180 loss=1.575 recovered=[[0.07 0.25 0.14 0.2  0.34]]\n",
      "[15.4s] iter=182 loss=1.458 recovered=[[0.07 0.15 0.18 0.29 0.32]]\n",
      "Recovered in 183 iterations\n"
     ]
    }
   ],
   "source": [
    "optimizer = tf.keras.optimizers.Adam(lr=0.1)\n",
    "is_sorted = False\n",
    "start = time.time()\n",
    "for i in range(NITER):\n",
    "    with tf.GradientTape() as tape:         \n",
    "        # our 'features'\n",
    "        target_samples = tf.cast(target.sample(NSAMPLES), 'float32') # cast to the right type\n",
    "\n",
    "        # move samples to the space where we know how to evaluate probabilities\n",
    "        reversed_target_samples = flow.reverse(target_samples)\n",
    "        \n",
    "        # evaluate log-probs of the samples (output shape=batch x N)\n",
    "        # (i.e., log_probs = base.log_prob(reversed_target_samples) )\n",
    "        probs = tf.reduce_sum(reversed_target_samples*base.probs, -1)\n",
    "        log_probs = tf.math.log(probs+1e-12)    \n",
    "        \n",
    "        # independent variables -> we just sum up log-probs \n",
    "        # to get joint log prob of a N-dim sample\n",
    "        log_probs = tf.reduce_sum(log_probs, -1) \n",
    "                \n",
    "        # loss = minus average log-likelihood\n",
    "        loss = -tf.reduce_mean(log_probs) \n",
    "\n",
    "        outprobs = tf.reduce_mean(flow(tf.cast(base.sample(NSAMPLES), 'float32')), 0)\n",
    "        \n",
    "        is_sorted = (sorted(outprobs.numpy()[0])==outprobs.numpy()[0]).all()\n",
    "\n",
    "        if i%10==0 or i<10 or is_sorted:        \n",
    "            print(\"[%.1fs] iter=%i loss=%.3f recovered=%s\" % \\\n",
    "                  (time.time()-start, i, loss, np.round(outprobs, 2)))\n",
    "        \n",
    "        if is_sorted:\n",
    "            print(\"Recovered in %i iterations\" % (i+1))\n",
    "            break\n",
    "                                  \n",
    "            \n",
    "    gradients = tape.gradient(loss, flow.trainable_variables)        \n",
    "    optimizer.apply_gradients(zip(gradients, flow.trainable_variables))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Store results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = open(OUT, \"w\")\n",
    "fmt = (\"%s,\"*8+\"%s\\n\")\n",
    "f.write((fmt % (K,NSAMPLES,NITER,SEED,t,B,int(PARTIAL_FLOWS),time.time()-start,i)))\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}