{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:04.029098Z",
     "iopub.status.busy": "2020-09-27T04:02:04.028307Z",
     "iopub.status.idle": "2020-09-27T04:02:08.657879Z",
     "shell.execute_reply": "2020-09-27T04:02:08.656229Z"
    },
    "papermill": {
     "duration": 4.676923,
     "end_time": "2020-09-27T04:02:08.658353",
     "exception": false,
     "start_time": "2020-09-27T04:02:03.981430",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "import os\n",
    "import pickle\n",
    "import numpy as np\n",
    "import warnings\n",
    "import matplotlib.pylab as plt\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import imageio\n",
    "from tqdm import tqdm\n",
    "from src.experiment_tagging import get_model_path\n",
    "from src.visualization import cluster_and_visualize\n",
    "from src.lesion import perform_lesion_experiment, plot_all_damaged_clusters\n",
    "from src.pointers import DATA_PATHS\n",
    "from src.utils import suppress, all_logging_disabled\n",
    "\n",
    "sns.set(font_scale=1.3)\n",
    "sns.set_style('whitegrid')\n",
    "\n",
    "if not os.path.isdir('../results'):\n",
    "    os.mkdir('../results')\n",
    "\n",
    "if not os.path.isdir('../results/lucid_ims'):\n",
    "    os.mkdir('../results/lucid_ims')\n",
    "\n",
    "if not os.path.isdir('../results/lucid_ims'):\n",
    "    os.mkdir('../results/cluster_vis_corr')\n",
    "\n",
    "n_clust = 12\n",
    "n_shuffles = 20\n",
    "n_workers = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:08.726415Z",
     "iopub.status.busy": "2020-09-27T04:02:08.725486Z",
     "iopub.status.idle": "2020-09-27T04:02:08.760451Z",
     "shell.execute_reply": "2020-09-27T04:02:08.759018Z"
    },
    "papermill": {
     "duration": 0.072939,
     "end_time": "2020-09-27T04:02:08.760787",
     "exception": false,
     "start_time": "2020-09-27T04:02:08.687848",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Factors of clusterability plotting\n",
    "\n",
    "mlp_csv = '../results/clustering_factors_mlp.csv'\n",
    "mlp_factors = pd.read_csv(mlp_csv)\n",
    "mlp_factors = mlp_factors.loc[mlp_factors['Dataset']!='CIFAR-10']\n",
    "clust_grad_csv = '../results/clustering_factors_clust_grad.csv'\n",
    "clust_grad_factors = pd.read_csv(clust_grad_csv)\n",
    "cnn_csv = '../results/clustering_factors_cnn.csv'\n",
    "cnn_factors = pd.read_csv(cnn_csv)\n",
    "cnn_factors = cnn_factors.loc[cnn_factors['Dataset']!='CIFAR-10']\n",
    "cnn_vgg_csv = '../results/clustering_factors_cnn_vgg.csv'\n",
    "cnn_vgg_factors = pd.read_csv(cnn_vgg_csv)\n",
    "imagenet_csv = '../results/clustering_factors_imagenet.csv'\n",
    "imagenet_factors = pd.read_csv(imagenet_csv)\n",
    "\n",
    "sz = 150\n",
    "figh = 6\n",
    "figw = 8\n",
    "fa = 0.4\n",
    "color_map = 'cubehelix'\n",
    "\n",
    "def label_points(x, y, val, ax, offset=0.03):\n",
    "    a = pd.concat({'x': x, 'y': y, 'val': val}, axis=1)\n",
    "    for i, point in a.iterrows():\n",
    "        if isinstance(point['val'], str):\n",
    "            ax.text(point['x']+offset, point['y'], point['val'], size='small')\n",
    "        elif not np.isnan(point['val']):\n",
    "            ax.text(point['x']+offset, point['y'], str(round(point['val'], 2)), size='x-small')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:08.901787Z",
     "iopub.status.busy": "2020-09-27T04:02:08.852111Z",
     "iopub.status.idle": "2020-09-27T04:02:10.520645Z",
     "shell.execute_reply": "2020-09-27T04:02:10.521963Z"
    },
    "papermill": {
     "duration": 1.733389,
     "end_time": "2020-09-27T04:02:10.522380",
     "exception": false,
     "start_time": "2020-09-27T04:02:08.788991",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61692 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# MLP regularization, unpruned\n",
    "factor_types = ['Control', 'L1 Reg', 'L2 Reg', 'Dropout']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = mlp_factors.loc[mlp_factors['Dataset'].isin(dset_types)].loc[mlp_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim([8.2, 10.25])\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'MLP Clusterability, Regularization')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_mlp_regularization_unpruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:10.613731Z",
     "iopub.status.busy": "2020-09-27T04:02:10.605180Z",
     "iopub.status.idle": "2020-09-27T04:02:12.702649Z",
     "shell.execute_reply": "2020-09-27T04:02:12.703981Z"
    },
    "papermill": {
     "duration": 2.156205,
     "end_time": "2020-09-27T04:02:12.704387",
     "exception": false,
     "start_time": "2020-09-27T04:02:10.548182",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63464 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# MLP regularization, pruned\n",
    "factor_types = ['Control, Pruning', 'L1 Reg, Pruning', 'L2 Reg, Pruning', 'Dropout, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = mlp_factors.loc[mlp_factors['Dataset'].isin(dset_types)].loc[mlp_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'MLP Clusterability, Regularization with Pruning')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_mlp_regularization_pruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:12.799765Z",
     "iopub.status.busy": "2020-09-27T04:02:12.789832Z",
     "iopub.status.idle": "2020-09-27T04:02:14.885716Z",
     "shell.execute_reply": "2020-09-27T04:02:14.887015Z"
    },
    "papermill": {
     "duration": 2.153419,
     "end_time": "2020-09-27T04:02:14.887425",
     "exception": false,
     "start_time": "2020-09-27T04:02:12.734006",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69308 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# MLP halves\n",
    "factor_types = ['Halves Same', 'Halves Same, Pruning', 'Halves Diff', 'Halves Diff, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = mlp_factors.loc[mlp_factors['Dataset'].isin(dset_types)].loc[mlp_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim([8.2, 10.22])\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'MLP Clusterability, Halves Datasets')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_mlp_halves.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:14.987597Z",
     "iopub.status.busy": "2020-09-27T04:02:14.954764Z",
     "iopub.status.idle": "2020-09-27T04:02:16.877518Z",
     "shell.execute_reply": "2020-09-27T04:02:16.878021Z"
    },
    "papermill": {
     "duration": 1.95854,
     "end_time": "2020-09-27T04:02:16.878217",
     "exception": false,
     "start_time": "2020-09-27T04:02:14.919677",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 59512 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# MLP clusterable init\n",
    "factor_types = ['Control', 'Control, Pruning', 'Clusterable Init', 'Clusterable Init, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = mlp_factors.loc[mlp_factors['Dataset'].isin(dset_types)].loc[mlp_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(5.4, 10.4)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'MLP Clusterability, Clusterable Init')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_mlp_clusterable_init.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:17.056287Z",
     "iopub.status.busy": "2020-09-27T04:02:17.007149Z",
     "iopub.status.idle": "2020-09-27T04:02:18.409439Z",
     "shell.execute_reply": "2020-09-27T04:02:18.410783Z"
    },
    "papermill": {
     "duration": 1.500235,
     "end_time": "2020-09-27T04:02:18.411373",
     "exception": false,
     "start_time": "2020-09-27T04:02:16.911138",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62244 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# MLP clust grad\n",
    "factor_types = ['Control', 'Control, Pruning', 'Cluster Gradient', 'Cluster Gradient, Pruning']\n",
    "subset_factors = clust_grad_factors.loc[clust_grad_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(4.8, 10.5)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'MLP Clusterability, Cluster Gradient')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_clust_grad.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:18.517502Z",
     "iopub.status.busy": "2020-09-27T04:02:18.515569Z",
     "iopub.status.idle": "2020-09-27T04:02:19.736163Z",
     "shell.execute_reply": "2020-09-27T04:02:19.738070Z"
    },
    "papermill": {
     "duration": 1.289552,
     "end_time": "2020-09-27T04:02:19.738601",
     "exception": false,
     "start_time": "2020-09-27T04:02:18.449049",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38300 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# MLP Polynomials\n",
    "factor_types = ['Control', 'L1 Reg', 'L2 Reg']\n",
    "dset_types = ['Polynomials']\n",
    "subset_factors = mlp_factors.loc[mlp_factors['Dataset'].isin(dset_types)].loc[mlp_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(7.3, 10.1)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'MLP Clusterability, Polynomial Regression')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_mlp_poly.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:19.842842Z",
     "iopub.status.busy": "2020-09-27T04:02:19.841268Z",
     "iopub.status.idle": "2020-09-27T04:02:21.623707Z",
     "shell.execute_reply": "2020-09-27T04:02:21.625092Z"
    },
    "papermill": {
     "duration": 1.851183,
     "end_time": "2020-09-27T04:02:21.625488",
     "exception": false,
     "start_time": "2020-09-27T04:02:19.774305",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 52264 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN regularization, unpruned\n",
    "factor_types = ['Control', 'L1 Reg', 'L2 Reg', 'Dropout']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(7.7, 10.9)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, Regularization')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_regularization_unpruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47872 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN dropout, unpruned\n",
    "factor_types = ['Control', 'Dropout']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(10.58, 10.72)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, Dropout')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax, offset=0.002)\n",
    "savepath = '../results/clusterability_factors_cnn_dropout_unpruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45900 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN l1/l2 reg, unpruned\n",
    "factor_types = ['L1 Reg', 'L2 Reg']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(7.7, 10.8)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, L1/L2 Regularization')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_L1L2_unpruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:21.735976Z",
     "iopub.status.busy": "2020-09-27T04:02:21.708305Z",
     "iopub.status.idle": "2020-09-27T04:02:23.448079Z",
     "shell.execute_reply": "2020-09-27T04:02:23.449707Z"
    },
    "papermill": {
     "duration": 1.789793,
     "end_time": "2020-09-27T04:02:23.450123",
     "exception": false,
     "start_time": "2020-09-27T04:02:21.660330",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 63212 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN regularization, pruned\n",
    "factor_types = ['Control, Pruning', 'L1 Reg, Pruning', 'L2 Reg, Pruning', 'Dropout, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(7.25, 11.1)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, Regularization with Pruning')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_regularization_pruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49368 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN dropout, pruned\n",
    "factor_types = ['Control, Pruning','Dropout, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(9.85, 10.65)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, Dropout with Pruning')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax, offset=0.01)\n",
    "savepath = '../results/clusterability_factors_cnn_dropout_pruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 46200 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN l1/l2 reg, pruned\n",
    "factor_types = ['L1 Reg, Pruning', 'L2 Reg, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "# plt.xlim(7.25, 11.1)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, L1/L2 Reg with Pruning')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_L1L2_pruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:23.561492Z",
     "iopub.status.busy": "2020-09-27T04:02:23.536058Z",
     "iopub.status.idle": "2020-09-27T04:02:25.383649Z",
     "shell.execute_reply": "2020-09-27T04:02:25.384964Z"
    },
    "papermill": {
     "duration": 1.896303,
     "end_time": "2020-09-27T04:02:25.385351",
     "exception": false,
     "start_time": "2020-09-27T04:02:23.489048",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 70900 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN stacked datasets\n",
    "factor_types = ['Stacked Same', 'Stacked Same, Pruning', 'Stacked Diff', 'Stacked Diff, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(loc='upper left', framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(9.88, 10.85)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, Stacked Datasets')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_stacked.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:25.502800Z",
     "iopub.status.busy": "2020-09-27T04:02:25.494442Z",
     "iopub.status.idle": "2020-09-27T04:02:27.331217Z",
     "shell.execute_reply": "2020-09-27T04:02:27.332543Z"
    },
    "papermill": {
     "duration": 1.907363,
     "end_time": "2020-09-27T04:02:27.332923",
     "exception": false,
     "start_time": "2020-09-27T04:02:25.425560",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65288 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# CNN clusterable init\n",
    "factor_types = ['Control', 'Control, Pruning', 'Clusterable Init', 'Clusterable Init, Pruning']\n",
    "dset_types = ['MNIST', 'FASHION']\n",
    "subset_factors = cnn_factors.loc[cnn_factors['Dataset'].isin(dset_types)].loc[cnn_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    style='Dataset',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    style_order=dset_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(10.05, 10.8)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Small CNN Clusterability, Clusterable Init')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_clusterable_init.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:27.450947Z",
     "iopub.status.busy": "2020-09-27T04:02:27.449439Z",
     "iopub.status.idle": "2020-09-27T04:02:28.967004Z",
     "shell.execute_reply": "2020-09-27T04:02:28.968339Z"
    },
    "papermill": {
     "duration": 1.595932,
     "end_time": "2020-09-27T04:02:28.968730",
     "exception": false,
     "start_time": "2020-09-27T04:02:27.372798",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48472 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# VGG CNN regularization, unpruned\n",
    "factor_types = ['Control', 'L1 Reg', 'L2 Reg', 'Dropout', 'Dropout, L2 Reg']\n",
    "subset_factors = cnn_vgg_factors.loc[cnn_vgg_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'CNN-VGG Clusterability, Regularization, CIFAR-10')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_vgg_regularization_unpruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:29.094743Z",
     "iopub.status.busy": "2020-09-27T04:02:29.066032Z",
     "iopub.status.idle": "2020-09-27T04:02:30.417192Z",
     "shell.execute_reply": "2020-09-27T04:02:30.418501Z"
    },
    "papermill": {
     "duration": 1.407848,
     "end_time": "2020-09-27T04:02:30.418882",
     "exception": false,
     "start_time": "2020-09-27T04:02:29.011034",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54160 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# VGG CNN regularization, pruned\n",
    "factor_types = ['Control, Pruning', 'L1 Reg, Pruning', 'L2 Reg, Pruning', \n",
    "                'Dropout, Pruning', 'Dropout, L2 Reg, Pruning']\n",
    "subset_factors = cnn_vgg_factors.loc[cnn_vgg_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'CNN-VGG Clusterability, Regularization with Pruning, CIFAR-10')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_vgg_regularization_pruned.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:30.551697Z",
     "iopub.status.busy": "2020-09-27T04:02:30.546056Z",
     "iopub.status.idle": "2020-09-27T04:02:31.811655Z",
     "shell.execute_reply": "2020-09-27T04:02:31.813195Z"
    },
    "papermill": {
     "duration": 1.351143,
     "end_time": "2020-09-27T04:02:31.813735",
     "exception": false,
     "start_time": "2020-09-27T04:02:30.462592",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 46232 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# VGG CNN clusterable init\n",
    "factor_types = ['Control', 'Control, Pruning', 'Clusterable Init', 'Clusterable Init, Pruning']\n",
    "subset_factors = cnn_vgg_factors.loc[cnn_vgg_factors['Network'].isin(factor_types)]\n",
    "factors_data = subset_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'CNN-VGG Clusterability, Clusterable Init, CIFAR-10')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax)\n",
    "savepath = '../results/clusterability_factors_cnn_vgg_clusterable_init.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:31.989586Z",
     "iopub.status.busy": "2020-09-27T04:02:31.946520Z",
     "iopub.status.idle": "2020-09-27T04:02:33.247569Z",
     "shell.execute_reply": "2020-09-27T04:02:33.248901Z"
    },
    "papermill": {
     "duration": 1.382325,
     "end_time": "2020-09-27T04:02:33.249286",
     "exception": false,
     "start_time": "2020-09-27T04:02:31.866961",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 37500 characters]\n",
      "text/plain": [
       "<Figure size 576x432 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Imagenet\n",
    "factor_types = ['VGG-16', 'VGG-19', 'ResNet-18', 'ResNet-34', 'ResNet-50', 'Inception-V3']\n",
    "factors_data = imagenet_factors\n",
    "\n",
    "f, ax = plt.subplots(1, figsize=(figw, figh))\n",
    "g = sns.scatterplot(x='ncut',\n",
    "                    y='z_score',\n",
    "                    hue='Network',\n",
    "                    hue_order=factor_types,\n",
    "                    s=sz,\n",
    "                    palette=sns.color_palette(color_map, len(factor_types)),\n",
    "                    data=factors_data,\n",
    "                    ax=ax)\n",
    "g.legend(framealpha=fa)\n",
    "g.legend(loc='center left')\n",
    "for lh in g.legend_.legendHandles:\n",
    "    lh._sizes = [sz]\n",
    "plt.xlim(2.7, 8.3)\n",
    "plt.xlabel('N-cut')\n",
    "plt.ylabel('Z-score')\n",
    "plt.title(f'Imagenet Clusterability')\n",
    "label_points(factors_data['ncut'], factors_data['z_score'],\n",
    "             factors_data['percentile'], ax, offset=0.08)\n",
    "savepath = '../results/clusterability_factors_imagenet.png'\n",
    "plt.savefig(savepath, dpi=300)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:02:33.370445Z",
     "iopub.status.busy": "2020-09-27T04:02:33.369162Z",
     "iopub.status.idle": "2020-09-27T04:04:07.952163Z",
     "shell.execute_reply": "2020-09-27T04:04:07.950667Z"
    },
    "papermill": {
     "duration": 94.651602,
     "end_time": "2020-09-27T04:04:07.952499",
     "exception": false,
     "start_time": "2020-09-27T04:02:33.300897",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save cluster vis corr images\n",
    "\n",
    "vis_tags = ['MNIST+DROPOUT', 'HALVES-MNIST+DROPOUT', 'HALVES-SAME-MNIST+DROPOUT',\n",
    "            'HALVES-MNIST', 'HALVES-SAME-MNIST',\n",
    "            'FASHION+DROPOUT','CNN-STACKED-SAME-MNIST', 'CNN-STACKED-MNIST']\n",
    "is_unpruned = False\n",
    "\n",
    "with suppress(), all_logging_disabled():\n",
    "    for tag in tqdm(vis_tags):\n",
    "        model_path = get_model_path(tag)\n",
    "        side_len = 32 if 'VGG' in tag else 28\n",
    "        vis_results = cluster_and_visualize(model_path, model_path, side_len=side_len,\n",
    "                                            n_clusters=n_clust, filter_norm=1)\n",
    "        all_columns = []\n",
    "        for lyr_i, lyr in enumerate(vis_results[is_unpruned]):\n",
    "            for clust in vis_results[is_unpruned][lyr]:\n",
    "                ims = vis_results[is_unpruned][lyr][clust]['ims']\n",
    "                if ims[0].shape[0] > 1:\n",
    "                    all_vis = np.concatenate([np.squeeze(im) for im in ims], axis=1)\n",
    "                else:\n",
    "                    all_vis = np.concatenate([np.squeeze(im) for im in ims], axis=0)\n",
    "                all_vis -= np.min(all_vis)  # normalize to have values in in [0,1]\n",
    "                all_vis /= np.max(all_vis)  # normalize to have values in in [0,1]\n",
    "                if len(ims) > 1:\n",
    "                    all_columns.append(all_vis)\n",
    "\n",
    "        all_cols_vis = np.concatenate(all_columns, axis=-1)\n",
    "        if len(all_cols_vis.shape) > 2:\n",
    "            if all_cols_vis.shape[0] < 3:\n",
    "                zero_red = np.expand_dims(np.zeros((all_cols_vis.shape[1], all_cols_vis.shape[2])), axis=0)\n",
    "                all_cols_vis = np.concatenate([zero_red, all_cols_vis], axis=0)\n",
    "            white_line = np.ones((3, 1, all_cols_vis.shape[-1]))\n",
    "            all_cols_vis = np.concatenate([all_cols_vis[:, :side_len, :], white_line, \n",
    "                                           all_cols_vis[:, side_len:, :]], axis=1)\n",
    "            all_cols_vis = np.transpose(all_cols_vis, axes=(1, 2, 0))  # channels last\n",
    "        else:\n",
    "            white_line = np.ones((1, all_cols_vis.shape[-1]))\n",
    "            all_cols_vis = np.vstack([all_cols_vis[:side_len, :], white_line, all_cols_vis[side_len:, :]])\n",
    "        if np.max(all_cols_vis) <= 1.0:\n",
    "            all_cols_vis *= 255\n",
    "        savename = '../results/cluster_vis_corr_' + tag + '_unpruned=' + str(is_unpruned) + '.png'\n",
    "        savename = savename.replace('+', '_').replace('-', '_')\n",
    "        imageio.imwrite(savename, all_cols_vis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:04:08.097507Z",
     "iopub.status.busy": "2020-09-27T04:04:08.065371Z",
     "iopub.status.idle": "2020-09-27T04:04:19.305928Z",
     "shell.execute_reply": "2020-09-27T04:04:19.304549Z"
    },
    "papermill": {
     "duration": 11.293875,
     "end_time": "2020-09-27T04:04:19.306272",
     "exception": false,
     "start_time": "2020-09-27T04:04:08.012397",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Save some example lucid visualizations\n",
    "\n",
    "data_dir='../datasets/'\n",
    "image_dir = '../results/lucid_ims/'\n",
    "model_tags = ('MNIST+LUCID', 'CNN-MNIST+LUCID', 'CNN-VGG-CIFAR10+DROPOUT+L2REG')\n",
    "imagenet_nets = ('vgg16', 'vgg19', 'resnet50')\n",
    "\n",
    "with suppress(), all_logging_disabled():\n",
    "\n",
    "    for tag in model_tags:\n",
    "\n",
    "        for is_unpruned in [True, False]:\n",
    "\n",
    "            suff = '_unpruned_max_data.pkl' if is_unpruned else '_pruned_max_data.pkl'\n",
    "            pruning_str = '_unpruned' if is_unpruned else '_pruned'\n",
    "\n",
    "            with open(data_dir + tag + '0' + suff, 'rb') as f:\n",
    "                data = pickle.load(f)\n",
    "\n",
    "            max_images = data['max_images']\n",
    "            random_max_images = data['random_max_images']\n",
    "            sm_sizes = data['sm_sizes']\n",
    "            sm_layers = data['sm_layers']\n",
    "            sm_layer_sizes = data['sm_layer_sizes']\n",
    "            n_examples = len(sm_sizes)\n",
    "            n_random = int(len(random_max_images) / n_examples)\n",
    "\n",
    "            all_random = [np.hstack(random_max_images[i*n_random: (i+1)*n_random]) for i in range(n_examples)]\n",
    "\n",
    "            for i in range(n_examples):\n",
    "                savename = image_dir + tag + pruning_str + '_layer' + str(sm_layers[i]).replace('/', '') + \\\n",
    "                           '_layersize' + str(sm_layer_sizes[i]) + '_size' + str(sm_sizes[i]) + '.png'\n",
    "                random_savename = image_dir + tag + pruning_str + '_layer' + str(sm_layers[i]).replace('/', '') + \\\n",
    "                                  '_layersize' + str(sm_layer_sizes[i]) + '_size' + str(sm_sizes[i]) + '_random.png'\n",
    "                savename = savename.replace('+', '_').replace('-', '_')\n",
    "                random_savename = random_savename.replace('+', '_').replace('-', '_')\n",
    "                imageio.imwrite(savename, max_images[i])\n",
    "                imageio.imwrite(random_savename, all_random[i])\n",
    "\n",
    "    for net in imagenet_nets:\n",
    "\n",
    "        with open(data_dir + net + '_max_data.pkl', 'rb') as f:\n",
    "            data = pickle.load(f)\n",
    "\n",
    "        max_images = data['max_images']\n",
    "        random_max_images = data['random_max_images']\n",
    "        sm_sizes = data['sm_sizes']\n",
    "        sm_layers = data['sm_layers']\n",
    "        sm_layer_sizes = data['sm_layer_sizes']\n",
    "        n_examples = len(sm_sizes)\n",
    "        n_random = int(len(random_max_images) / n_examples)\n",
    "\n",
    "        all_random = [np.hstack(random_max_images[i*n_random: (i+1)*n_random]) for i in range(n_examples)]\n",
    "\n",
    "        for i in range(n_examples):\n",
    "            savename = image_dir + net + '_layer' + str(sm_layers[i]).replace('/', '') + \\\n",
    "                       '_layersize' + str(sm_layer_sizes[i]) + '_size' + str(sm_sizes[i]) + '.png'\n",
    "            random_savename = image_dir + net + '_layer' + str(sm_layers[i]).replace('/', '') + \\\n",
    "                              '_layersize' + str(sm_layer_sizes[i]) + '_size' + str(sm_sizes[i]) + '_random.png'\n",
    "            savename = savename.replace('+', '_').replace('-', '_')\n",
    "            random_savename = random_savename.replace('+', '_').replace('-', '_')\n",
    "            if np.max(max_images[i]) <= 1.0:\n",
    "                max_images[i] *= 255\n",
    "                all_random[i] *= 255\n",
    "            imageio.imwrite(savename, max_images[i])\n",
    "            imageio.imwrite(random_savename, all_random[i])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2020-09-27T04:04:19.410044Z",
     "iopub.status.busy": "2020-09-27T04:04:19.409286Z",
     "iopub.status.idle": "2020-09-27T04:09:53.963395Z",
     "shell.execute_reply": "2020-09-27T04:09:53.964401Z"
    },
    "papermill": {
     "duration": 334.606015,
     "end_time": "2020-09-27T04:09:53.964692",
     "exception": false,
     "start_time": "2020-09-27T04:04:19.358677",
     "status": "completed"
    },
    "pycharm": {
     "name": "#%%\n"
    },
    "scrolled": false,
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [05:31<00:00, 165.57s/it]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 393512 characters]\n",
      "text/plain": [
       "<Figure size 2160x576 with 10 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 254444 characters]\n",
      "text/plain": [
       "<Figure size 2160x864 with 15 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Save some example lesion visualizations\n",
    "\n",
    "n_clust = 4\n",
    "# os.environ['CUDA_VISIBLE_DEVICES'] = '3'\n",
    "lesion_models = (('fashion', 'FASHION'), ('fashion', 'CNN-FASHION'), \n",
    "                 #('poly', 'POLY')\n",
    "                )\n",
    "\n",
    "with warnings.catch_warnings():\n",
    "    warnings.simplefilter('ignore')\n",
    "\n",
    "    for dataset_name, model_name in tqdm(lesion_models):\n",
    "\n",
    "        if 'CNN' in model_name:\n",
    "            is_unpruned = True\n",
    "        else:\n",
    "            is_unpruned = False\n",
    "\n",
    "        (true_results,\n",
    "         all_random_results,\n",
    "         metadata,\n",
    "         evaluation) = perform_lesion_experiment('.' + DATA_PATHS[dataset_name],\n",
    "                                                 get_model_path(model_name), n_clusters=n_clust,\n",
    "                                                 n_shuffles=n_shuffles, unpruned=is_unpruned)\n",
    "\n",
    "        model_and_dataset = model_name if 'CNN' in model_name else 'MLP-' + model_name\n",
    "        axes = plot_all_damaged_clusters(true_results, all_random_results,\n",
    "                                         metadata, evaluation, title=f'{model_and_dataset}, k={n_clust}')\n",
    "        savename = '../results/lesion_' + model_name + '_unpruned=' + str(is_unpruned) + '.png'\n",
    "        plt.savefig(savename)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "papermill": {
     "duration": 0.054872,
     "end_time": "2020-09-27T04:09:54.086362",
     "exception": false,
     "start_time": "2020-09-27T04:09:54.031490",
     "status": "completed"
    },
    "tags": []
   },
   "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.5"
  },
  "papermill": {
   "duration": 473.437719,
   "end_time": "2020-09-27T04:09:56.360088",
   "environment_variables": {},
   "exception": null,
   "input_path": "./notebooks/plotting_main.ipynb",
   "output_path": "./notebooks/plotting_main.ipynb",
   "parameters": {},
   "start_time": "2020-09-27T04:02:02.922369",
   "version": "1.2.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
