{
 "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",
    "%matplotlib inline\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "SEEDS = [1234, 1067, 9198, 9453, 6240, 7044, 1755, 7898, 3131, 1277]\n",
    "\n",
    "for SEED in SEEDS:\n",
    "    random.seed(SEED)\n",
    "    torch.manual_seed(SEED)\n",
    "    np.random.seed(SEED)\n",
    "    rng = np.random.RandomState(SEED)\n",
    "    torch.use_deterministic_algorithms(True)\n",
    "\n",
    "    n = 10\n",
    "    k = 5\n",
    "\n",
    "    # Generate random parameters for the distribution: these are the initial weights\n",
    "    theta = rng.randn(n)\n",
    "    print(theta)\n",
    "\n",
    "    # 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",
    "    \n",
    "    mat_x = np.zeros((len(combs), n))\n",
    "    for i in range(n_states):\n",
    "        mat_x[i, combs[i]] = 1.\n",
    "\n",
    "\n",
    "    # Create pytorch tensors from numpy array\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\n",
    "\n",
    "    assert(t.sum(pmf_t(theta)) == 1)  # so far so good\n",
    "\n",
    "    # Groundtruth weights\n",
    "    b_t = t.abs(t.from_numpy(rng.randn(n)).float())\n",
    "    print(b_t)\n",
    "\n",
    "    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)\n",
    "    \n",
    "    # Objective function to minimize\n",
    "    def objective(index):\n",
    "        return t.sum((states_t[index] - b_t)**2)\n",
    "\n",
    "\n",
    "    # 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)\n",
    "\n",
    "\n",
    "    # Ground truth gradient\n",
    "    exact_gradient = t.autograd.grad(expectation_t(theta_t), theta_t)\n",
    "\n",
    "    # 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)} \n",
    "    # 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",
    "\n",
    "\n",
    "    # Generic function that uses a given strategy\n",
    "    # and returns the estimated gradient\n",
    "    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",
    "            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\n",
    "\n",
    "    #1. Let's try the Straight through estimator\n",
    "    def ste_grad(grad_out):\n",
    "        return grad_out\n",
    "\n",
    "    def sample(_theta):\n",
    "        _sampled_index = sample_state_from_pdf(_theta)\n",
    "        return states_t[_sampled_index]\n",
    "\n",
    "    def sample_many(_theta):\n",
    "        _pmft = pmf_t(_theta)\n",
    "        indices = torch.multinomial(_pmft, num_samples=5000, replacement=True)\n",
    "        return states_t[indices]\n",
    "\n",
    "\n",
    "    # define top-k sample dist with STE gradient\n",
    "    class TopKTrueSamplingSTEGrad(t.autograd.Function):\n",
    "\n",
    "        @staticmethod\n",
    "        def forward(ctx, _theta):\n",
    "            return sample(_theta)\n",
    "\n",
    "        @staticmethod\n",
    "        def backward(ctx, grad_outputs):\n",
    "            return ste_grad(grad_outputs)\n",
    "\n",
    "    grads = []\n",
    "    for i in range(10000):\n",
    "        grads += [return_grad(TopKTrueSamplingSTEGrad.apply)]\n",
    "    STE = torch.stack(grads)\n",
    "\n",
    "    # 2.  let's try I-MLE with faithful samples\n",
    "    def imle_forward(ctx, _theta, _lambda):\n",
    "        ctx._lambda = _lambda\n",
    "        ctx._theta = _theta\n",
    "        with torch.no_grad():\n",
    "            ctx._fw = sample(_theta)\n",
    "        return ctx._fw\n",
    "\n",
    "    def imle_backward(ctx, grad_out):\n",
    "        theta_prime = ctx._theta - ctx._lambda*grad_out #q\n",
    "        sample_prime = sample(theta_prime)\n",
    "        return ctx._fw - sample_prime # The gradient is the difference of sample(p) and sample(q)\n",
    "\n",
    "    class TopKTrueSamplingIMLEGradWithImplicitQ(t.autograd.Function):\n",
    "\n",
    "        @staticmethod\n",
    "        def forward(ctx, _theta, _lambda):\n",
    "            return imle_forward(ctx, _theta, _lambda)\n",
    "\n",
    "        @staticmethod\n",
    "        def backward(ctx, grad_outputs):\n",
    "\n",
    "            grad = imle_backward(ctx, grad_outputs)\n",
    "            return grad, None\n",
    "\n",
    "    imle_ts_ap = TopKTrueSamplingIMLEGradWithImplicitQ.apply\n",
    "    imle_ts_strat = lambda _th: imle_ts_ap(_th, torch.tensor(2.5))\n",
    "\n",
    "    grads = []\n",
    "    for i in range(10000):\n",
    "        grads += [return_grad(imle_ts_strat)]\n",
    "    ES_IMLE = torch.stack(grads)\n",
    "    \n",
    "\n",
    "    def logaddexp(x, y):\n",
    "        with torch.no_grad():\n",
    "            m = torch.max(x, y)\n",
    "            m = m.masked_fill_(torch.isneginf(m), 0.)\n",
    "\n",
    "        z = (x - m).exp_() + (y - m).exp_()\n",
    "        mask = z == 0\n",
    "        z = z.masked_fill_(mask, 1.).log_().add_(m)\n",
    "        z = z.masked_fill_(mask, -float('inf'))\n",
    "\n",
    "        return z\n",
    "\n",
    "    def log1mexp(x):\n",
    "        assert(torch.all(x >= 0))\n",
    "        return torch.where(x < 0.6931471805599453094, torch.log(-torch.expm1(-x)), torch.log1p(-torch.exp(-x)))\n",
    "\n",
    "\n",
    "    def prob_k(probs, k, log_space=True):\n",
    "        \"\"\"\n",
    "        probs: a tensor of shape: (batch_size, num_vars)\n",
    "        where probs[:, i] corresponds to the batch probabilities\n",
    "        of Bernoulli variable Xi\n",
    "\n",
    "        if log_space=True, we expect log_probabilities as input\n",
    "        \"\"\"\n",
    "        batch_size = probs.size(0)\n",
    "        n = probs.size(1)\n",
    "        # a[:, i, j] = Pr(Sum(X1, ... , Xi) = j)\n",
    "        # Note: a[:, 0, 0] corresponds to the Pr\n",
    "        # that an empty sequence summing up to -1\n",
    "        # which is always 0\n",
    "        a = torch.zeros((batch_size, n+1, k+2), requires_grad=True)\n",
    "        if log_space:\n",
    "            a = torch.log(a)\n",
    "        # The probability of an empty sequence\n",
    "        # summing to 0 is 1\n",
    "        a[:, 0, 1] = 0 if log_space else 1\n",
    "        for i in range(1, n+1):\n",
    "\n",
    "            # To get a sequence of length i with only\n",
    "            # j true variables, I either take a seq.\n",
    "            # with k true variables, and set variable\n",
    "            # i to false, or take a sequence with k-1\n",
    "            # true variables and set variable i to true\n",
    "            if log_space:\n",
    "                a[:, i, 1:] = logaddexp(a[:, i-1, :-1] + probs[:, i-1:i],  a[:, i-1, 1:] + log1mexp(-probs[:, i-1:i].detach()))\n",
    "            else:\n",
    "                a[:, i, 1:] =  a[:, i-1, :-1].clone() * probs[:, i-1:i] + a[:, i-1, 1:].clone() * (1 - probs[:, i-1:i]).detach()\n",
    "\n",
    "\n",
    "        return a[:, n, k+1:k+2]\n",
    "\n",
    "    #4.  let's try I-MLE with inexact samples and exact marginals\n",
    "    _k_gamma = 5.0\n",
    "    _tau_gamma = 1.0\n",
    "\n",
    "    def sog_th1(s=10):\n",
    "        return (_tau_gamma/_k_gamma)*( np.sum([rng.gamma(1.0/_k_gamma, _k_gamma/(i+1.0)) for i in range(s)] ) - np.log(s) )\n",
    "\n",
    "    def map(_theta):\n",
    "        arg_sort = t.argsort(_theta)[k:]\n",
    "        _x = t.zeros(_theta.size())\n",
    "        _x[arg_sort] = 1.\n",
    "        return _x\n",
    "\n",
    "    def perturb_and_map(ctx, _theta):\n",
    "        if hasattr(ctx, 'eps'):\n",
    "            eps = ctx.eps\n",
    "        else:\n",
    "            eps = t.tensor([sog_th1() for _ in range(n)])\n",
    "            try: ctx.eps = eps\n",
    "            except AttributeError: print('Problems with ctx')\n",
    "        theta_prime = _theta + eps\n",
    "        return map(theta_prime)\n",
    "\n",
    "\n",
    "    def logsigmoid(x):\n",
    "        return -F.softplus(-x)\n",
    "\n",
    "    def imle_forward(ctx, _theta, _lambda):\n",
    "        ctx._lambda = _lambda\n",
    "        ctx._theta = _theta\n",
    "        with torch.no_grad():\n",
    "            ctx._fw = perturb_and_map(ctx, _theta)\n",
    "        return ctx._fw\n",
    "\n",
    "    def imle_backward(ctx, grad_out):\n",
    "        with torch.enable_grad():\n",
    "            theta_prime = ctx._theta - ctx._lambda*grad_out #q\n",
    "            log_p = logsigmoid(ctx._theta).unsqueeze(0)\n",
    "            log_q = logsigmoid(theta_prime).unsqueeze(0)\n",
    "            a_p = prob_k(log_p, k)\n",
    "            a_q = prob_k(log_q, k)\n",
    "            mar_p = torch.autograd.grad(a_p, log_p)\n",
    "            mar_q = torch.autograd.grad(a_q, log_q)\n",
    "        return mar_p[0] - mar_q[0]\n",
    "\n",
    "\n",
    "    class TopKTrueSamplingIMLEGradWithImplicitQ(t.autograd.Function):\n",
    "\n",
    "        @staticmethod\n",
    "        def forward(ctx, _theta, _lambda):\n",
    "            return imle_forward(ctx, _theta, _lambda)\n",
    "\n",
    "        @staticmethod\n",
    "        def backward(ctx, grad_outputs):\n",
    "\n",
    "            grad = imle_backward(ctx, grad_outputs)\n",
    "            return grad, None\n",
    "\n",
    "    imle_ts_ap = TopKTrueSamplingIMLEGradWithImplicitQ.apply\n",
    "    imle_ts_strat = lambda _th: imle_ts_ap(_th, torch.tensor(2.5))\n",
    "\n",
    "    grads = []\n",
    "    for i in range(10000):\n",
    "        grads += [return_grad(imle_ts_strat)]\n",
    "    EM_IMLE = torch.stack(grads)\n",
    "\n",
    "    #5.  let's try I-MLE with inexact samples and inexact marginals\n",
    "    _k_gamma = 5.0\n",
    "    _tau_gamma = 1.0\n",
    "\n",
    "    def sog_th1(s=10):\n",
    "        return (_tau_gamma/_k_gamma)*( np.sum([rng.gamma(1.0/_k_gamma, _k_gamma/(i+1.0)) for i in range(s)] ) - np.log(s) )\n",
    "\n",
    "    def map(_theta):\n",
    "        arg_sort = t.argsort(_theta)[k:]\n",
    "        _x = t.zeros(_theta.size())\n",
    "        _x[arg_sort] = 1.\n",
    "        return _x\n",
    "\n",
    "    def perturb_and_map(ctx, _theta):\n",
    "        if hasattr(ctx, 'eps'):\n",
    "            eps = ctx.eps\n",
    "        else:\n",
    "            eps = t.tensor([sog_th1() for _ in range(n)])\n",
    "            try: ctx.eps = eps\n",
    "            except AttributeError: print('Problems with ctx')\n",
    "        theta_prime = _theta + eps\n",
    "        return map(theta_prime)\n",
    "\n",
    "\n",
    "    def logsigmoid(x):\n",
    "        return -F.softplus(-x)\n",
    "\n",
    "    def imle_forward(ctx, _theta, _lambda):\n",
    "        ctx._lambda = _lambda\n",
    "        ctx._theta = _theta\n",
    "        with torch.no_grad():\n",
    "            ctx._fw = perturb_and_map(ctx, _theta)\n",
    "        return ctx._fw\n",
    "\n",
    "    def imle_backward(ctx, grad_out):\n",
    "        theta_prime = ctx._theta - ctx._lambda*grad_out #q\n",
    "        sample_prime = perturb_and_map(ctx, theta_prime)\n",
    "        return ctx._fw - sample_prime # The gradient is the difference of sample(p) and sample(q)\n",
    "    \n",
    "\n",
    "    class TopKTrueSamplingIMLEGradWithImplicitQ(t.autograd.Function):\n",
    "\n",
    "        @staticmethod\n",
    "        def forward(ctx, _theta, _lambda):\n",
    "            return imle_forward(ctx, _theta, _lambda)\n",
    "\n",
    "        @staticmethod\n",
    "        def backward(ctx, grad_outputs):\n",
    "\n",
    "            grad = imle_backward(ctx, grad_outputs)\n",
    "            return grad, None\n",
    "\n",
    "    imle_ts_ap = TopKTrueSamplingIMLEGradWithImplicitQ.apply\n",
    "    imle_ts_strat = lambda _th: imle_ts_ap(_th, torch.tensor(2.5))\n",
    "\n",
    "    grads = []\n",
    "    for i in range(10000):\n",
    "        grads += [return_grad(imle_ts_strat)]\n",
    "    gumbel_IMLE = torch.stack(grads)\n",
    "\n",
    "    # Softsub\n",
    "    import numpy as np\n",
    "\n",
    "    EPSILON = 1e-07\n",
    "\n",
    "    random.seed(SEED)\n",
    "    torch.manual_seed(SEED)\n",
    "    np.random.seed(SEED)\n",
    "    rng = np.random.RandomState(SEED)\n",
    "    torch.use_deterministic_algorithms(True)\n",
    "\n",
    "\n",
    "    def gumbel_keys(w):\n",
    "        uniform = torch.rand(w.shape)\n",
    "        z = -torch.log(-torch.log(uniform))\n",
    "        w = w + z\n",
    "        return w\n",
    "\n",
    "\n",
    "    def continuous_topk(w, k, t, separate=False):\n",
    "        khot_list = []\n",
    "        onehot_approx = torch.zeros_like(w, dtype=torch.float32)\n",
    "        for i in range(k):\n",
    "            khot_mask = torch.maximum(1.0 - onehot_approx, torch.full_like(onehot_approx, EPSILON))\n",
    "            w += torch.log(khot_mask)\n",
    "            onehot_approx = torch.nn.Softmax(dim=-1)(w/t)\n",
    "            khot_list.append(onehot_approx)\n",
    "        if separate:\n",
    "            return khot_list\n",
    "        else:\n",
    "            return torch.sum(torch.stack(khot_list), dim=0)\n",
    "\n",
    "\n",
    "    def sample_subset(w, k, t=2.0):\n",
    "        '''\n",
    "        Args:\n",
    "            w (Tensor): Float Tensor of weights for each element. In gumbel mode\n",
    "                these are interpreted as log probabilities\n",
    "            k (int): number of elements in the subset sample\n",
    "            t (float): temperature of the softmax\n",
    "        '''\n",
    "        w = gumbel_keys(w)\n",
    "        return continuous_topk(w, k, t)\n",
    "\n",
    "    grads = []\n",
    "    for i in range(10000):\n",
    "        theta_t = t.from_numpy(theta).float().requires_grad_(True)\n",
    "        obj = t.sum((sample_subset(logsigmoid(theta_t), k) - b_t)**2)\n",
    "        obj.backward()\n",
    "        grads += [theta_t.grad]\n",
    "    softsub = torch.stack(grads)\n",
    "\n",
    "    grads = []\n",
    "    for i in range(10000):\n",
    "        theta_t = t.from_numpy(theta).float().requires_grad_(True)\n",
    "        s = sample(theta_t)\n",
    "        obj = t.sum((s - b_t)**2)\n",
    "        logprob = pmf_t(theta_t)[(states_t == s).all(dim=1).nonzero().squeeze()].log()\n",
    "        logprob.backward()\n",
    "        grads += [theta_t.grad*obj]\n",
    "\n",
    "    SFE = torch.stack(grads)\n",
    "\n",
    "    #3.  let's try SIMPLE\n",
    "    import functorch\n",
    "    import random\n",
    "\n",
    "    random.seed(SEED)\n",
    "    torch.manual_seed(SEED)\n",
    "    np.random.seed(SEED)\n",
    "    rng = np.random.RandomState(SEED)\n",
    "    torch.use_deterministic_algorithms(True)\n",
    "\n",
    "    def logsigmoid(x):\n",
    "        return -F.softplus(-x)\n",
    "\n",
    "    def imle_forward(ctx, _theta, _lambda):\n",
    "        ctx._lambda = _lambda\n",
    "        ctx._theta = _theta\n",
    "        with torch.no_grad():\n",
    "            ctx._fw = sample(_theta)\n",
    "        return ctx._fw\n",
    "\n",
    "    def log_pr_exactly_k(theta, complement_theta):\n",
    "        # Note: this is written with an explicit dependence on `complement_theta = log(1 - prob)` so\n",
    "        # that we can compute marginals by differentiating wrt `theta` while holding its complement fixed.\n",
    "        probs = torch.exp(theta)\n",
    "        complement_probs = torch.exp(complement_theta)\n",
    "        a = torch.zeros([n + 1, k + 2])\n",
    "        a[0, 1] = 1\n",
    "\n",
    "        for i in range(1, n + 1):\n",
    "        a [i, 1:] = a[i-1, :-1].clone() * probs[i-1] + a[i-1, 1:].clone() * complement_probs[i - 1]\n",
    "        return torch.log(a[n, k + 1])\n",
    "\n",
    "    def marginals(theta):\n",
    "        log_p = torch.log(torch.sigmoid(theta))\n",
    "        log_p_complement = torch.log1p(-torch.exp(log_p))\n",
    "        res = functorch.grad(log_pr_exactly_k, argnums=0)(log_p, log_p_complement)\n",
    "        return res\n",
    "\n",
    "    def imle_backward(ctx, grad_out):\n",
    "        return functorch.jvp(marginals, tuple([ctx._theta]), tuple([grad_out]))[1]\n",
    "\n",
    "    class TopKTrueSamplingIMLEGradWithImplicitQ(t.autograd.Function):\n",
    "\n",
    "        @staticmethod\n",
    "        def forward(ctx, _theta, _lambda):\n",
    "            return imle_forward(ctx, _theta, _lambda)\n",
    "\n",
    "        @staticmethod\n",
    "        def backward(ctx, grad_outputs):\n",
    "\n",
    "            grad = imle_backward(ctx, grad_outputs)\n",
    "            return grad, None\n",
    "\n",
    "    imle_ts_ap = TopKTrueSamplingIMLEGradWithImplicitQ.apply\n",
    "    imle_ts_strat = lambda _th: imle_ts_ap(_th, torch.tensor(0.001))\n",
    "\n",
    "    grads = []\n",
    "    for i in range(10000):\n",
    "        grads += [return_grad(imle_ts_strat)]\n",
    "    E_IMLE_jvp = torch.stack(grads)\n",
    "\n",
    "    exact_gradient = exact_gradient[0].expand(10000, 10)\n",
    "\n",
    "    gradients = {'Exact': exact_gradient, 'STE': STE, 'SoftSub':softsub,'IMLE': gumbel_IMLE, \n",
    "                 'SFE': SFE,  'SIMPLE-F': ES_IMLE, 'SIMPLE-B': EM_IMLE, 'SIMPLE': E_IMLE_jvp}\n",
    "\n",
    "    torch.save(gradients, f'gradients_{SEED}.pt')\n",
    "    \n",
    "    x = ['Exact', 'SoftSub', 'IMLE', 'SFE', 'SIMPLE-F', 'SIMPLE-B', 'SIMPLE']\n",
    "    x_axis = np.arange(len(x))\n",
    "\n",
    "    # plot the bias-variance of the estimator\n",
    "    y = [1.0 - F.cosine_similarity(exact_gradient.mean(axis=0), gradients[estimator].mean(axis=0), dim=0) for estimator in x]\n",
    "    plt.bar(x_axis-(0.2/2), y, color='blue', width=0.2, label='Bias')\n",
    "\n",
    "    y = []\n",
    "    for estimator in x:\n",
    "        mu = gradients[estimator].mean(axis=0)\n",
    "        y += [F.cosine_similarity(gradients[estimator], mu).var()]\n",
    "    plt.bar(x_axis+(0.2/2), y, color='red', width=0.2, label='Variance')\n",
    "\n",
    "    plt.legend()\n",
    "    ax2 = plt.twinx()\n",
    "    ax2.set_ylim([0.0, 0.14])\n",
    "    plt.xticks(x_axis, x)\n",
    "    plt.savefig('bias_variance.png',bbox_inches='tight')\n",
    "    plt.show()\n",
    "\n",
    "    # Plot the average error of the estimators\n",
    "    y = [1.0 - F.cosine_similarity(exact_gradient, gradients[estimator], dim=1).mean() for estimator in x]\n",
    "    errors = [F.cosine_similarity(exact_gradient, gradients[estimator], dim=1).std() for estimator in x]\n",
    "    plt.errorbar(x, y, yerr=errors, capsize=3, fmt='o', color='blue')\n",
    "    plt.savefig('error.png', bbox_inches='tight')\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1234, 1067, 9198, 9453, 6240, 7044, 1755, 7898, 3131, 1277\n",
    "gradients_0 = torch.load('gradients_1234.pt')\n",
    "gradients_1 = torch.load('gradients_1067.pt')\n",
    "gradients_2 = torch.load('gradients_9198.pt')\n",
    "gradients_3 = torch.load('gradients_9453.pt')\n",
    "gradients_4 = torch.load('gradients_6240.pt')\n",
    "gradients_5 = torch.load('gradients_7044.pt')\n",
    "gradients_6 = torch.load('gradients_1755.pt')\n",
    "gradients_7 = torch.load('gradients_7898.pt')\n",
    "gradients_8 = torch.load('gradients_3131.pt')\n",
    "gradients_9 = torch.load('gradients_1277.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 380,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12408 characters]\n",
      "text/plain": [
       "<Figure size 432x192.24 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 10500 characters]\n",
      "text/plain": [
       "<Figure size 432x192.24 with 1 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",
    "x = ['Exact', 'SoftSub', 'IMLE', 'SFE', 'SIMPLE-F', 'SIMPLE-B', 'SIMPLE']\n",
    "\n",
    "x_axis = np.arange(len(x))\n",
    "\n",
    "# plot the bias-variance of the estimator\n",
    "all_grads = [gradients_0, gradients_1, gradients_2, gradients_3, gradients_4, gradients_5, gradients_6, gradients_7, gradients_8, gradients_9]\n",
    "\n",
    "biases = []\n",
    "for gradients in all_grads:\n",
    "    bias = [1.0 - F.cosine_similarity(gradients['Exact'].mean(axis=0), gradients[estimator].mean(axis=0), dim=0) for estimator in x]\n",
    "    biases.append(torch.tensor(bias))\n",
    "biases = torch.stack(biases)\n",
    "\n",
    "fig, ax1 = plt.subplots(figsize=(6, 2.67));\n",
    "lns1 = ax1.bar(x=[0, 3, 6, 9, 12, 15, 18], height=biases.mean(dim=0), alpha=0.75, edgecolor=\"k\", lw=1.5, color=sanae_colors[1], label='Bias')\n",
    "ax1.errorbar(x=[0, 3, 6, 9, 12, 15, 18] ,y=biases.mean(dim=0), yerr=biases.std(dim=0), fmt='none', c='k', alpha=1., zorder=3, mew=1.5, capsize=3)\n",
    "ax1.grid(axis=\"y\")\n",
    "plt.xticks([0.5, 3.5, 6.5, 9.5, 12.5, 15.5, 18.5],\n",
    "           x, fontsize=10)\n",
    "\n",
    "\n",
    "ax2=plt.twinx()\n",
    "\n",
    "variances = []\n",
    "for gradients in all_grads:\n",
    "    variance = []\n",
    "    for estimator in x:\n",
    "        mu = gradients[estimator].mean(axis=0)\n",
    "        variance += [F.cosine_similarity(gradients[estimator], mu).var()]\n",
    "    variances.append(torch.tensor(variance))\n",
    "variances = torch.stack(variances)\n",
    "lns3 = ax2.bar(x=[1, 4, 7, 10, 13, 16, 19], height=variances.mean(dim=0), alpha=0.75, edgecolor=\"k\", lw=1.5, color=sanae_colors[4], label='Variance')\n",
    "ax2.errorbar(x=[1, 4, 7, 10, 13, 16, 19] ,y=variances.mean(dim=0), yerr=variances.std(dim=0), fmt='none', c='k', alpha=1., zorder=3, mew=1.5, capsize=3)\n",
    "\n",
    "ax2.set_ylim([0.0, 0.168])\n",
    "plt.legend([lns1, lns3], ['Bias', 'Variance'],loc=1)\n",
    "\n",
    "plt.savefig('bias_variance.pdf',bbox_inches='tight')\n",
    "plt.show()\n",
    "\n",
    "# Plot the average error of the estimators\n",
    "errors = []\n",
    "ys = []\n",
    "for gradients in all_grads:\n",
    "    y = [1.0 - F.cosine_similarity(gradients['Exact'], gradients[estimator], dim=1).mean() for estimator in x]\n",
    "    y_ = [(1.0 - F.cosine_similarity(gradients['Exact'], gradients[estimator], dim=1)).std() for estimator in x]\n",
    "    ys.append(torch.tensor(y_))\n",
    "    errors.append(torch.tensor(y))\n",
    "errors = torch.stack(errors)\n",
    "ys = torch.stack(ys)\n",
    "\n",
    "plt.figure(figsize=(6, 2.67))\n",
    "# fig, ax1 = plt.subplots();\n",
    "plt.errorbar(x=[0, 2, 4, 6, 8, 10, 12] ,y=errors.mean(dim=0), yerr=ys.mean(dim=0), fmt='none', c='k', alpha=1.,mew=1.5, capsize=3)\n",
    "ax1=plt.scatter([0, 2, 4, 6, 8, 10, 12], errors.mean(dim=0), color=sanae_colors[4], alpha=0.75, edgecolor=\"k\", lw=1.5, s=75)\n",
    "# plt.xlim(-0.3, 12.3)\n",
    "plt.grid(axis=\"y\")\n",
    "plt.xticks([0, 2, 4, 6, 8, 10, 12],x, fontsize=10)\n",
    "plt.tick_params(axis='y', which='minor', bottom=False)\n",
    "plt.savefig('errors.pdf',bbox_inches='tight')"
   ]
  }
 ],
 "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.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
