{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import utils\n",
    "from typing import Any\n",
    "import torch as t\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import itertools\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "# from functools import partial\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# parameters for top-k distribution\n",
    "rng = np.random.RandomState(1234)\n",
    "n = 10\n",
    "k = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0.47143516 -1.19097569  1.43270697 -0.3126519  -0.72058873  0.88716294\n",
      "  0.85958841 -0.6365235   0.01569637 -2.24268495]\n"
     ]
    }
   ],
   "source": [
    "# generate random parameters for the distribution: these are the initial weights\n",
    "theta = rng.randn(n)\n",
    "print(theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of possible states: 10\n",
      "[[1. 0. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 1. 0. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 1. 0. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 1. 0. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 1. 0. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 1. 0. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0. 1. 0.]\n",
      " [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]]\n"
     ]
    }
   ],
   "source": [
    "# create all possible_states:\n",
    "combs = list(itertools.combinations(range(n), k))\n",
    "n_states = len(combs)\n",
    "assert n_states == np.math.factorial(n)/(np.math.factorial(k)*np.math.factorial(n-k))\n",
    "print('Number of possible states:', n_states)\n",
    "mat_x = np.zeros((len(combs), n))\n",
    "for i in range(n_states):\n",
    "    mat_x[i, combs[i]] = 1.\n",
    "\n",
    "print(mat_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# put all of this in pytorch\n",
    "theta_t = t.from_numpy(theta).float().requires_grad_(True)\n",
    "states_t = t.from_numpy(mat_x).float()\n",
    "\n",
    "def tow_t(_theta):\n",
    "    return states_t @ _theta\n",
    "\n",
    "def Z_t(_theta):\n",
    "    return t.log(t.sum(t.exp(tow_t(_theta))))\n",
    "\n",
    "def pmf_t(_theta):\n",
    "    return t.exp(tow_t(_theta) - Z_t(_theta))\n",
    "\n",
    "def sample_state_from_pdf(_theta):\n",
    "    _pmft = pmf_t(_theta)\n",
    "    indx_ch = rng.choice(list(range(n_states)), p=_pmft.detach().numpy())\n",
    "    return indx_ch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1., dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "print(t.sum(pmf_t(theta)))  # so far so good"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#groundtruth distribution\n",
    "b_t = t.abs(t.from_numpy(rng.randn(n)).float())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1.1500, 0.9919, 0.9533, 2.0213, 0.3341, 0.0021, 0.4055, 0.2891, 1.3212,\n",
      "        1.5469])\n"
     ]
    }
   ],
   "source": [
    "print(b_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.7361908\n"
     ]
    }
   ],
   "source": [
    "sorted_bt = np.sort(b_t.detach().numpy())\n",
    "min_value_of_exp = np.sum((sorted_bt[:5])**2) + np.sum((sorted_bt[5:] - 1)**2)\n",
    "print(min_value_of_exp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(10.8149)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def objective(index):\n",
    "    return t.sum((states_t[index] - b_t)**2)\n",
    "\n",
    "objective(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1.1500, 0.9919, 0.9533, 2.0213, 0.3341, 0.0021, 0.4055, 0.2891, 1.3212,\n",
       "        1.5469])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(pmf_t(theta_t))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# writing explicitly the expectation of this objective summing over\n",
    "# all possible states:\n",
    "def expectation_t(_theta):\n",
    "    _pmf = pmf_t(_theta)\n",
    "    _p_values = t.stack([_pmf[i] * objective(i) for i in range(n_states)])\n",
    "    return t.sum(_p_values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(11.2864, grad_fn=<SumBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expectation_t(theta_t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "#Ground truth gradient\n",
    "exact_gradient = t.autograd.grad(expectation_t(theta_t), theta_t)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Essentially we are now solving explicitly\n",
    "$\\min_{\\theta} \\mathbb{E}_{z\\sim p(z, \\theta)} b^\\intercal z$\n",
    "where $p(z, \\theta)$ is top-k distribution.\n",
    "\n",
    "With full optimization we simply write $\\mathbb{E}_{z\\sim p(z, \\theta)} b^\\intercal z= \\sum_{i=1}^{N} p(z_i, \\theta) b^\\intercal z_i $\n",
    "summing over all possible states, where $N=\\binom{n}{k}$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def return_grad(strategy, reinitialize=True):\n",
    "    global theta_t\n",
    "    if reinitialize:\n",
    "        theta_t = t.from_numpy(theta).float().requires_grad_(True)\n",
    "        \n",
    "    # redefine objective with given strategy\n",
    "    def objective_(_theta):\n",
    "        sample = strategy(_theta)\n",
    "#         print(\"Sample from obj_\", sample)\n",
    "#         import pdb; pdb.set_trace()\n",
    "#         return t.sum((sample - b_t)**2)\n",
    "        if len(sample.shape) == 2:\n",
    "            return ((sample - b_t)**2).mean(dim=0).sum()\n",
    "        else:\n",
    "            return t.sum((sample - b_t)**2)\n",
    "        \n",
    "    obj = objective_(theta_t)\n",
    "    obj.backward()\n",
    "    return theta_t.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "\n",
    "def sample_gumbel(shape, eps=1e-20):\n",
    "    U = torch.rand(shape)\n",
    "    return -Variable(torch.log(-torch.log(U + eps) + eps))\n",
    "\n",
    "def gumbel_softmax_sample(logits, temperature):\n",
    "#     import pdb; pdb.set_trace()\n",
    "    y = logits + sample_gumbel(logits.size())\n",
    "    return F.softmax(y / temperature, dim=-1)\n",
    "\n",
    "def gumbel_softmax(logits, temperature=1.0):\n",
    "    \"\"\"\n",
    "    input: [*, n_class]\n",
    "    return: [*, n_class] an one-hot vector\n",
    "    \"\"\"\n",
    "    y = gumbel_softmax_sample(logits, temperature)\n",
    "    shape = y.size()\n",
    "    _, ind = y.max(dim=-1)\n",
    "    y_hard = torch.zeros_like(y).view(-1, shape[-1])\n",
    "    y_hard.scatter_(1, ind.view(-1, 1), 1)\n",
    "    y_hard = y_hard.view(*shape)\n",
    "    return (y_hard - y).detach() + y\n",
    "\n",
    "grads = []\n",
    "for i in range(10000):\n",
    "    grads += [return_grad(gumbel_softmax)]\n",
    "STGS_grads = torch.stack(grads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2149684/1349740020.py:13: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  y_perturbed = F.softmax(logits + sample_gumbel(logits.size()))\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "\n",
    "def SIMPLE(logits):\n",
    "    \"\"\"\n",
    "    input: [*, n_class]\n",
    "    return: [*, n_class] an one-hot vector\n",
    "    \"\"\"\n",
    "    y = F.softmax(logits, dim=-1)\n",
    "    y_perturbed = F.softmax(logits + sample_gumbel(logits.size()))\n",
    "    shape = y.size()\n",
    "    _, ind = y_perturbed.max(dim=-1)\n",
    "    y_hard = torch.zeros_like(y_perturbed).view(-1, shape[-1])\n",
    "    y_hard.scatter_(1, ind.view(-1, 1), 1)\n",
    "    y_hard = y_hard.view(*shape)\n",
    "    return (y_hard - y).detach() + y\n",
    "\n",
    "grads = []\n",
    "for i in range(10000):\n",
    "    grads += [return_grad(SIMPLE)]\n",
    "SIMPLE_grads = torch.stack(grads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "exact_gradient = exact_gradient[0].expand(10000, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11536 characters]\n",
      "text/plain": [
       "<Figure size 288x192.24 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import math\n",
    "import seaborn as sns\n",
    "sns.set_theme(style=\"white\")\n",
    "\n",
    "sanae_colors = [sns.cubehelix_palette(as_cmap=False)[i] for i in range(6)]\n",
    "\n",
    "gradients = {'Exact': exact_gradient, 'ST Gumbel Softmax': STGS_grads, 'SIMPLE':SIMPLE_grads}\n",
    "x = ['Exact', 'ST Gumbel Softmax', 'SIMPLE']\n",
    "x_axis = np.arange(len(x))\n",
    "\n",
    "bias = [1.0 - F.cosine_similarity(exact_gradient.mean(axis=0), gradients[estimator].mean(axis=0), dim=0) for estimator in x]\n",
    "\n",
    "fig, ax1 = plt.subplots(figsize=(4, 2.67));\n",
    "lns1 = ax1.bar(x=[0, 3, 6], height=bias, alpha=0.75, edgecolor=\"k\", lw=1.5, color=sanae_colors[1], label='Bias')\n",
    "ax1.grid(axis=\"y\")\n",
    "plt.xticks([0.5, 3.5, 6.5], x, fontsize=10)\n",
    "\n",
    "\n",
    "ax2=plt.twinx()\n",
    "\n",
    "variance = []\n",
    "for estimator in x:\n",
    "    mu = gradients[estimator].mean(axis=0)\n",
    "    variance += [(1.0 - F.cosine_similarity(gradients[estimator], mu)).var()]\n",
    "\n",
    "lns3 = ax2.bar(x=[1, 4, 7], height=variance, alpha=0.75, edgecolor=\"k\", lw=1.5, color=sanae_colors[4], label='Variance')\n",
    "\n",
    "ax2.set_ylim([0.0, 0.14])\n",
    "plt.legend([lns1, lns3], ['Bias', 'Variance'],loc=1)\n",
    "\n",
    "plt.savefig('bias_variance_SIMPLE_STGS.pdf',bbox_inches='tight')\n",
    "plt.show()"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
