{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from augerino import datasets, models\n",
    "import numpy as np\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import torch.nn.functional as F\n",
    "from augerino.models.e2_steerable import C8SteerableCNN, SmallE2\n",
    "from smallnet import smallnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_cuda = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpt_path = \"./saved-outputs/\"\n",
    "ncopies = 12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = smallnet(in_channels=1,num_targets=10)\n",
    "augerino = models.UniformAug()\n",
    "low_model = models.AugAveragedModel(net, augerino,ncopies=ncopies)\n",
    "\n",
    "net = smallnet(in_channels=1,num_targets=10)\n",
    "augerino = models.UniformAug()\n",
    "mid_model = models.AugAveragedModel(net, augerino,ncopies=ncopies)\n",
    "\n",
    "net = smallnet(in_channels=1,num_targets=10)\n",
    "augerino = models.UniformAug()\n",
    "high_model = models.AugAveragedModel(net, augerino,ncopies=ncopies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "low_model.load_state_dict(torch.load(checkpt_path + \"model0.01.pt\"))\n",
    "mid_model.load_state_dict(torch.load(checkpt_path + \"model0.05.pt\"))\n",
    "high_model.load_state_dict(torch.load(checkpt_path + \"model0.1.pt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = smallnet(in_channels=1,num_targets=10)\n",
    "augerino = models.UniformAug()\n",
    "init_model = models.AugAveragedModel(net, augerino,ncopies=ncopies)\n",
    "init_model.load_state_dict(torch.load(checkpt_path + \"model0.01_init.pt\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/greg_b/e2cnn/e2cnn/nn/modules/r2_conv/basisexpansion_singleblock.py:61: UserWarning: indexing with dtype torch.uint8 is now deprecated, please use a dtype torch.bool instead. (Triggered internally at  /pytorch/aten/src/ATen/native/IndexingUtils.h:20.)\n",
      "  sampled_basis = sampled_basis[mask, ...]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e2_model = SmallE2(channel_in=1, n_classes=10, rot_n=10)\n",
    "e2_model.load_state_dict(torch.load(\"./saved-outputs/e2_epoch200.pt\", map_location='cpu'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/greg_b/miniconda3/envs/tf-gpu/lib/python3.8/site-packages/torch/nn/functional.py:3447: UserWarning: Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.\n",
      "  warnings.warn(\"Default grid_sample and affine_grid behavior has changed \"\n",
      "/home/greg_b/miniconda3/envs/tf-gpu/lib/python3.8/site-packages/torch/nn/functional.py:3384: UserWarning: Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.\n",
      "  warnings.warn(\"Default grid_sample and affine_grid behavior has changed \"\n"
     ]
    }
   ],
   "source": [
    "dataset = datasets.RotMNIST(\"~/datasets/\", train=True)\n",
    "trainloader = DataLoader(dataset, batch_size=128)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "img, lab = next(iter(trainloader))\n",
    "init_aug = init_model.aug(img)\n",
    "end_aug = mid_model.aug(img)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's Mess Around With Rotating Inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = datasets.RotMNIST(\"~/datasets/\", train=False)\n",
    "testloader = DataLoader(dataset, batch_size=128)\n",
    "testimg, testlab =next(iter(trainloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([  2,   9,  20,  26,  53,  58,  60,  61,  64,  89,  92, 115, 127]),)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.where(testlab==4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/greg_b/miniconda3/envs/tf-gpu/lib/python3.8/site-packages/torch/nn/functional.py:3447: UserWarning: Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.\n",
      "  warnings.warn(\"Default grid_sample and affine_grid behavior has changed \"\n",
      "/home/greg_b/miniconda3/envs/tf-gpu/lib/python3.8/site-packages/torch/nn/functional.py:3384: UserWarning: Default grid_sample and affine_grid behavior has changed to align_corners=False since 1.3.0. Please specify align_corners=True if the old behavior is desired. See the documentation of grid_sample for details.\n",
      "  warnings.warn(\"Default grid_sample and affine_grid behavior has changed \"\n"
     ]
    }
   ],
   "source": [
    "n_ang = 50\n",
    "angles = torch.linspace(-np.pi, np.pi, n_ang)\n",
    "\n",
    "four = testimg[ind, ::].unsqueeze(0)\n",
    "batch_four = torch.cat(n_ang * [four])\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    # Build affine matrices for random translation of each image\n",
    "    affineMatrices = torch.zeros(n_ang,2,3)\n",
    "    affineMatrices[:,0,0] = angles.cos()\n",
    "    affineMatrices[:,1,1] = angles.cos()\n",
    "    affineMatrices[:,0,1] = angles.sin()\n",
    "    affineMatrices[:,1,0] = -angles.sin()\n",
    "\n",
    "    flowgrid = F.affine_grid(affineMatrices, size = batch_four.size())\n",
    "    \n",
    "    rot_four = F.grid_sample(batch_four, flowgrid)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "low_model.eval();\n",
    "high_model.eval();\n",
    "mid_model.eval();\n",
    "e2_model.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# low_model.train();\n",
    "# high_model.train();\n",
    "# mid_model.train();\n",
    "# e2_model.train();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    low_preds = low_model(rot_four)\n",
    "    high_preds = high_model(rot_four)\n",
    "    mid_preds = mid_model(rot_four)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    e2_preds = e2_model(rot_four)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "sftmx = torch.nn.Softmax(-1)\n",
    "low_probs = sftmx(low_preds)\n",
    "high_probs = sftmx(high_preds)\n",
    "mid_probs = sftmx(mid_preds)\n",
    "e2_probs = sftmx(e2_preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "tick_pts = [-np.pi, -np.pi/2, 0, np.pi/2, np.pi]\n",
    "tick_labs = [r\"-$\\pi$\", r'-$\\pi$/2', '0', r'$\\pi$/2', r'$\\pi$']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### distribution plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "softplus = torch.nn.Softplus()\n",
    "test_pts = torch.linspace(-2, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_density(test_pts, width):\n",
    "    dist = torch.distributions.Uniform(-width/2., width/2.)\n",
    "    dens = dist.log_prob(test_pts).exp().detach()\n",
    "    return dens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init_dens = get_density(test_pts, softplus(torch.tensor(-1.)))\n",
    "# high_dens = get_density(test_pts, high_logger[500:501]['width2'].values[0])\n",
    "# mid_dens = get_density(test_pts, mid_logger[500:501]['width2'].values[0])\n",
    "# low_dens = get_density(test_pts, low_logger[500:501]['width2'].values[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_dens = get_density(test_pts, softplus(torch.tensor(-1.)))\n",
    "width = min(softplus(high_model.aug.width[2]), 2*np.pi)\n",
    "high_dens = get_density(test_pts, width)\n",
    "width = min(softplus(mid_model.aug.width[2]), 2*np.pi)\n",
    "mid_dens = get_density(test_pts, width)\n",
    "width = min(softplus(low_model.aug.width[2]), 2*np.pi)\n",
    "low_dens = get_density(test_pts, width)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot it all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_img(fig, img, i, j, label):\n",
    "    ax = plt.Subplot(fig, inner[i,j])\n",
    "    ax.imshow(img, cmap=\"Greys\", interpolation='nearest')\n",
    "    ax.set_title(label, fontsize=ax_fs+4, y=-0.5)\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "    fig.add_subplot(ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 242976 characters]\n",
      "text/plain": [
       "<Figure size 4000x800 with 8 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.gridspec as gridspec\n",
    "sns.set_style(\"white\")\n",
    "fig = plt.figure(figsize=(20, 4), dpi=200)\n",
    "outer = gridspec.GridSpec(1, 3, wspace=0.2, hspace=0.2)\n",
    "alpha=0.5\n",
    "ax_fs=18\n",
    "lwd=4.\n",
    "## plot mario and iggy\n",
    "\n",
    "inner = gridspec.GridSpecFromSubplotSpec(2, 3,\n",
    "                subplot_spec=outer[0], wspace=0.1, hspace=0.25)\n",
    "ind=42\n",
    "plot_img(fig, img[ind, 0, ::].cpu().detach(),\n",
    "        0, 0, \"\")\n",
    "plot_img(fig, img[ind+1, 0, ::].cpu().detach(),\n",
    "        0, 1, \"\")\n",
    "plot_img(fig, img[ind+2, 0, ::].cpu().detach(),\n",
    "        0, 2, \"\")\n",
    "plot_img(fig, img[ind+3, 0, ::].cpu().detach(),\n",
    "        1, 0, \"\")\n",
    "plot_img(fig, img[ind+4, 0, ::].cpu().detach(),\n",
    "        1, 1, \"Data Samples\")\n",
    "plot_img(fig, img[ind+5, 0, ::].cpu().detach(),\n",
    "        1, 2, \"\")\n",
    "ax = plt.Subplot(fig, inner[1, 1])\n",
    "        \n",
    "## Plot Learned Distribution\n",
    "inner = gridspec.GridSpecFromSubplotSpec(1,1,\n",
    "                    subplot_spec=outer[1], wspace=0.1, hspace=0.1)\n",
    "ax = plt.Subplot(fig, inner[0])\n",
    "ax.plot(test_pts, low_dens.detach().cpu(), linewidth=lwd, label=\"No Reg.\",\n",
    "        alpha=alpha, linestyle=\"-\")\n",
    "ax.plot(test_pts, mid_dens.detach().cpu(), linewidth=lwd, label=\"Std. Reg.\",\n",
    "        alpha=alpha, linestyle=\"--\")\n",
    "ax.plot(test_pts, high_dens.detach().cpu(), linewidth=lwd, label=\"High Reg.\",\n",
    "        alpha=alpha, linestyle=\":\")\n",
    "ax.plot(test_pts, init_dens.detach(), linewidth=lwd, label=\"Initial\", color='gray', alpha=0.7)\n",
    "ax.legend(loc='lower left', fontsize=ax_fs-1, ncol=2,\n",
    "           bbox_to_anchor=(0.02, -0.5))\n",
    "ax.set_xlabel(\"Rotation\", fontsize=ax_fs)\n",
    "ax.set_ylabel(\"Probability\", fontsize=ax_fs)\n",
    "tick_pts = [-np.pi/2, -np.pi/4, 0, np.pi/4, np.pi/2]\n",
    "tick_labs = [r\"-$\\pi$/2\", r'-$\\pi$/4', '0', r'$\\pi$/4', r'$\\pi$/2']\n",
    "ax.set_xticks(tick_pts)\n",
    "ax.set_xticklabels(tick_labs)\n",
    "ax.set_xlim(-np.pi/2, np.pi/2)\n",
    "ax.tick_params(\"both\", labelsize=ax_fs-2)\n",
    "sns.despine(ax=ax)\n",
    "fig.add_subplot(ax)\n",
    "\n",
    "\n",
    "inner = gridspec.GridSpecFromSubplotSpec(1,1,\n",
    "                    subplot_spec=outer[2], wspace=0.1, hspace=0.1)\n",
    "ax = plt.Subplot(fig, inner[0])\n",
    "alpha=0.7\n",
    "num=4\n",
    "lwd=4\n",
    "ax.plot(angles, low_probs[:, num].detach().cpu(), linewidth=lwd, label=\"Low Reg.\",\n",
    "        alpha=alpha, linestyle=\"-\")\n",
    "ax.plot(angles, mid_probs[:, num].detach().cpu(), linewidth=lwd, label=\"Mid. Reg.\",\n",
    "        alpha=alpha, linestyle=\"-\")\n",
    "ax.plot(angles, high_probs[:, num].detach().cpu(), linewidth=lwd, label=\"High Reg.\",\n",
    "        alpha=alpha, linestyle=\"-\")\n",
    "ax.plot(angles, e2_probs[:, num].detach().cpu(), linewidth=3, label=\"E2\",\n",
    "        color='gray', alpha=0.6)\n",
    "ax.set_ylim(0.0, 1.001)\n",
    "ax.set_xlabel(\"Rotation of Input\", fontsize=ax_fs)\n",
    "ax.set_ylabel(\"Predicted Probability\", fontsize=ax_fs)\n",
    "tick_pts = [-np.pi, -np.pi/2, 0, np.pi/2, np.pi]\n",
    "tick_labs = [r\"-$\\pi$\", r'-$\\pi$/2', '0', r'$\\pi$/2', r'$\\pi$']\n",
    "ax.tick_params(\"both\", labelsize=ax_fs-2)\n",
    "ax.set_xticks(tick_pts)\n",
    "ax.set_xticklabels(tick_labs)\n",
    "\n",
    "# ax.set_yticks([0.98, 0.99, 1.])\n",
    "ax.set_xlim(-np.pi, np.pi)\n",
    "ax.legend(loc=\"lower left\", fontsize=ax_fs-1, ncol=2,\n",
    "           bbox_to_anchor=(0.02, -0.54))\n",
    "sns.despine(ax=ax)\n",
    "fig.add_subplot(ax)\n",
    "plt.savefig(\"./rotmnist_full.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
