{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "import torch\n",
    "assert torch.cuda.is_available()\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "from tqdm.auto import tqdm\n",
    "import json\n",
    "from train import eval\n",
    "from Utils import load\n",
    "plt.rcParams[\"font.family\"] = \"Segoe UI\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pruners = np.array(\n",
    "    [\n",
    "        # 'cs',\n",
    "        # 'grasp',\n",
    "        # 'sf',\n",
    "        # 'mica_cs_shuffle',\n",
    "        # 'mica_grasp_shuffle',\n",
    "        # 'mica_sf_shuffle',\n",
    "        'mica_erk_shuffle',\n",
    "        'mica_igq_shuffle',\n",
    "        # 'mica_cs_rand',\n",
    "        # 'mica_grasp_rand',\n",
    "        # 'mica_sf_rand',\n",
    "        'mica_erk_rand',\n",
    "        'mica_igq_rand'\n",
    "    ],\n",
    "    dtype=object\n",
    "    )\n",
    "\n",
    "colors = {\n",
    "    'cs'         : plt.cm.tab10(9),\n",
    "    'grasp'      : plt.cm.tab10(8),\n",
    "    'sf'         : plt.cm.tab10(6),\n",
    "    'mica_cs'    : plt.cm.tab10(9),\n",
    "    'mica_grasp' : plt.cm.tab10(8),\n",
    "    'mica_sf'    : plt.cm.tab10(6),\n",
    "    'mica_erk'   : plt.cm.tab10(4),\n",
    "    'mica_igq'   : plt.cm.tab10(3),\n",
    "}\n",
    "\n",
    "markers = {\n",
    "    'cs'         : \"^\", \n",
    "    'grasp'      : \"h\", \n",
    "    'sf'         : \"H\", \n",
    "    'mica_cs'    : \"^\", \n",
    "    'mica_grasp' : \"h\", \n",
    "    'mica_sf'    : \"H\",\n",
    "    'mica_erk'   : \"p\", \n",
    "    'mica_igq'   : \"*\"\n",
    "}\n",
    "\n",
    "pruner_names = {\n",
    "    'cs'                 : 'RbPI-SNIP',\n",
    "    'grasp'              : 'RbPI-GraSP',\n",
    "    'sf'                 : 'RbPI-SynFlow',\n",
    "    'mica_cs_shuffle'    : 'RPI-SNIP',\n",
    "    'mica_grasp_shuffle' : 'RPI-GraSP',\n",
    "    'mica_sf_shuffle'    : 'RPI-SynFlow',\n",
    "    'mica_erk_shuffle'   : 'RPI-ERK',\n",
    "    'mica_igq_shuffle'   : 'RPI-IGQ',\n",
    "    'mica_cs_rand'       : '$\\\\bf{MiCA}$-$\\\\bf{SNIP}$',\n",
    "    'mica_grasp_rand'    : '$\\\\bf{MiCA}$-$\\\\bf{GraSP}$',\n",
    "    'mica_sf_rand'       : '$\\\\bf{MiCA}$-$\\\\bf{SynFlow}$',\n",
    "    'mica_erk_rand'      : '$\\\\bf{MiCA}$-$\\\\bf{ERK}$',\n",
    "    'mica_igq_rand'      : '$\\\\bf{MiCA}$-$\\\\bf{IGQ}$',\n",
    "    }\n",
    "\n",
    "pepochs = {\n",
    "    'cs'         : '1',\n",
    "    'grasp'      : '1',\n",
    "    'sf'         : '100',\n",
    "    'mica_cs'    : '1',\n",
    "    'mica_grasp' : '1',\n",
    "    'mica_sf'    : '1',\n",
    "    'mica_erk'   : '1',\n",
    "    'mica_igq'   : '1',\n",
    "}\n",
    "\n",
    "for pruner in pruners:\n",
    "    if '_shuffle' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_shuffle', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_shuffle', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_shuffle', '')]\n",
    "    elif '_rand' in pruner:\n",
    "        colors[pruner]  = colors[pruner.replace('_rand', '')]\n",
    "        markers[pruner]  = markers[pruner.replace('_rand', '')]\n",
    "        pepochs[pruner] = pepochs[pruner.replace('_rand', '')]\n",
    "\n",
    "runs = [0, 1, 2]\n",
    "\n",
    "max_run_num = len(runs)\n",
    "\n",
    "print(pruners)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_experimentid(pruner, ratio, pepochs, run_num, expid):\n",
    "    if '_shuffle' in pruner:\n",
    "        base_dir = f'../Results/data/singleshot/{expid}_shuffle'\n",
    "        p_name = pruner.replace('_shuffle', '')\n",
    "    elif '_min' in pruner:\n",
    "        base_dir = f'../Results/data/singleshot/{expid}_min'\n",
    "        p_name = pruner.replace('_min', '')\n",
    "    elif '_max' in pruner:\n",
    "        base_dir = f'../Results/data/singleshot/{expid}_max'\n",
    "        p_name = pruner.replace('_max', '')\n",
    "    elif '_rand' in pruner:\n",
    "        base_dir = f'../Results/data/singleshot/{expid}_rand'\n",
    "        p_name = pruner.replace('_rand', '')\n",
    "    else:\n",
    "        base_dir = f'../Results/data/singleshot/{expid}'\n",
    "        p_name = pruner\n",
    "    b_dir  = f'{base_dir}'\n",
    "    experimentid =  f'{b_dir}/' \\\n",
    "                    f'{p_name}-{ratio}-{pepochs}/' \\\n",
    "                    f'run_{run_num}/' \n",
    "    return experimentid\n",
    "\n",
    "def zero_to_nan(values):\n",
    "    return [float('nan') if x==0 else x for x in values]\n",
    "\n",
    "def get_ordered_list(l, order):\n",
    "    out = []\n",
    "    for i in order:\n",
    "        out.append(l[i])\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expid = 'cifar10_lottery_resnet20'\n",
    "ratios = [\n",
    "    '0.5', '1.0', '1.5', '2.0', \n",
    "    '2.5', '3.0', '3.5', '4.0', \n",
    "    ]\n",
    "\n",
    "resnet_edges = {}\n",
    "resnet_effective_edges = {}\n",
    "for l, pruner in enumerate(pruners):\n",
    "    resnet_edges[pruner]           = np.zeros((len(ratios)))\n",
    "    resnet_effective_edges[pruner] = np.zeros((len(ratios)))\n",
    "    for i, ratio in enumerate(tqdm(ratios, desc=pruner)): \n",
    "        edges = []\n",
    "        effective_edges = []\n",
    "        for r, run_num in enumerate(runs):\n",
    "            experimentid      = get_experimentid(pruner, ratio, pepochs[pruner], run_num, expid)\n",
    "            model_file        = experimentid + 'model.pt'\n",
    "            train_result_file = experimentid + 'post-train.pkl'\n",
    "            compression_file  = experimentid + 'compression.pkl'\n",
    "            train_result      = pd.read_pickle(train_result_file)\n",
    "            compression       = pd.read_pickle(compression_file)\n",
    "            prunable          = compression['prunable']==True\n",
    "            weight            = compression['param']=='weight'\n",
    "            modules           = compression[np.logical_and(prunable, weight)]['module']\n",
    "            densities         = np.array(compression[np.logical_and(prunable, weight)]['density'])\n",
    "            shape             = compression[np.logical_and(prunable,weight)]['shape']\n",
    "            e                 = np.array([np.prod(s) for s in shape]) * (densities)\n",
    "            edges.append(np.sum(e))\n",
    "\n",
    "            with open(experimentid + 'args.json') as f:\n",
    "                args = json.load(f)\n",
    "            input_shape, num_classes = load.dimension(args['dataset']) \n",
    "            loss = nn.CrossEntropyLoss()\n",
    "            if i == 0 and r == 0:\n",
    "                test_loader, _, _ = load.dataloader(\n",
    "                    args['dataset'], args['test_batch_size'], args['test_batch_size'],\n",
    "                    False, args['workers'], 0, 0)\n",
    "                dataset = args['dataset']\n",
    "                test_batch_size = args['test_batch_size']\n",
    "                workers = args['workers']\n",
    "                (data, _) = next(iter(test_loader))\n",
    "                input_dim = list(data[0,:].shape)\n",
    "                input = torch.ones([1] + input_dim).to('cuda')\n",
    "            assert args['seed'] == run_num * 1000 # adhoc\n",
    "            assert args['run_number'] == str(run_num)\n",
    "            assert dataset == args['dataset']\n",
    "            assert test_batch_size == args['test_batch_size']\n",
    "            assert workers == args['workers']\n",
    "\n",
    "            model = load.model(args['model'], args['model_class'])(\n",
    "                input_shape, num_classes, args['dense_classifier'], args['pretrained']).to('cuda')\n",
    "            model.load_state_dict(torch.load(model_file))\n",
    "\n",
    "            print(f'Forward :')\n",
    "            model.set_effective_sparsity(is_forward=True, x=input)\n",
    "            print(f'Backward :')\n",
    "            model.set_effective_sparsity(is_forward=False, x=None)\n",
    "\n",
    "            effective_densities = []\n",
    "            set_modules = set(modules.unique())\n",
    "            for n, m in model.named_modules():\n",
    "                if hasattr(m, 'weight_mask') and n in set_modules:\n",
    "                    effective_densities.append((m.weight_mask.sum() / m.weight_mask.numel()).to('cpu'))\n",
    "            effective_densities = np.array(effective_densities)\n",
    "            e = np.array([np.prod(s) for s in shape]) * (effective_densities)\n",
    "            effective_edges.append(np.sum(e))\n",
    "\n",
    "        resnet_edges[pruner][i]           = np.mean(edges)    \n",
    "        resnet_effective_edges[pruner][i] = np.mean(effective_edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expid = 'cifar10_lottery_vgg16_bn'\n",
    "ratios = [\n",
    "    '0.5', '1.0', '1.5', '2.0', \n",
    "    '2.5', '3.0', '3.5', '4.0', \n",
    "    '4.5', '5.0', '5.5', '6.0', \n",
    "    ]\n",
    "\n",
    "vgg_edges = {}\n",
    "vgg_effective_edges = {}\n",
    "for l, pruner in enumerate(pruners):\n",
    "    vgg_edges[pruner]           = np.zeros((len(ratios)))\n",
    "    vgg_effective_edges[pruner] = np.zeros((len(ratios)))\n",
    "    for i, ratio in enumerate(tqdm(ratios, desc=pruner)): \n",
    "        edges = []\n",
    "        effective_edges = []\n",
    "        for r, run_num in enumerate(runs):\n",
    "            experimentid      = get_experimentid(pruner, ratio, pepochs[pruner], run_num, expid)\n",
    "            model_file        = experimentid + 'model.pt'\n",
    "            train_result_file = experimentid + 'post-train.pkl'\n",
    "            compression_file  = experimentid + 'compression.pkl'\n",
    "            train_result      = pd.read_pickle(train_result_file)\n",
    "            compression       = pd.read_pickle(compression_file)\n",
    "            prunable          = compression['prunable']==True\n",
    "            weight            = compression['param']=='weight'\n",
    "            modules           = compression[np.logical_and(prunable, weight)]['module']\n",
    "            densities         = np.array(compression[np.logical_and(prunable, weight)]['density'])\n",
    "            shape             = compression[np.logical_and(prunable,weight)]['shape']\n",
    "            e                 = np.array([np.prod(s) for s in shape]) * (densities)\n",
    "            edges.append(np.sum(e))\n",
    "\n",
    "            with open(experimentid + 'args.json') as f:\n",
    "                args = json.load(f)\n",
    "            input_shape, num_classes = load.dimension(args['dataset']) \n",
    "            loss = nn.CrossEntropyLoss()\n",
    "            if i == 0 and r == 0:\n",
    "                test_loader, _, _ = load.dataloader(\n",
    "                    args['dataset'], args['test_batch_size'], args['test_batch_size'],\n",
    "                    False, args['workers'], 0, 0)\n",
    "                dataset = args['dataset']\n",
    "                test_batch_size = args['test_batch_size']\n",
    "                workers = args['workers']\n",
    "                (data, _) = next(iter(test_loader))\n",
    "                input_dim = list(data[0,:].shape)\n",
    "                input = torch.ones([1] + input_dim).to('cuda')\n",
    "            assert args['seed'] == run_num * 1000 # adhoc\n",
    "            assert args['run_number'] == str(run_num)\n",
    "            assert dataset == args['dataset']\n",
    "            assert test_batch_size == args['test_batch_size']\n",
    "            assert workers == args['workers']\n",
    "\n",
    "            model = load.model(args['model'], args['model_class'])(\n",
    "                input_shape, num_classes, args['dense_classifier'], args['pretrained']).to('cuda')\n",
    "            model.load_state_dict(torch.load(model_file))\n",
    "\n",
    "            print(f'Forward :')\n",
    "            model.set_effective_sparsity(is_forward=True, x=input)\n",
    "            print(f'Backward :')\n",
    "            model.set_effective_sparsity(is_forward=False, x=None)\n",
    "\n",
    "            effective_densities = []\n",
    "            set_modules = set(modules.unique())\n",
    "            for n, m in model.named_modules():\n",
    "                if hasattr(m, 'weight_mask') and n in set_modules:\n",
    "                    effective_densities.append((m.weight_mask.sum() / m.weight_mask.numel()).to('cpu'))\n",
    "            effective_densities = np.array(effective_densities)\n",
    "            e = np.array([np.prod(s) for s in shape]) * (effective_densities)\n",
    "            effective_edges.append(np.sum(e))\n",
    "\n",
    "        vgg_edges[pruner][i]           = np.mean(edges)    \n",
    "        vgg_effective_edges[pruner][i] = np.mean(effective_edges)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "lw=4\n",
    "markersize=30\n",
    "fontsize=35\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0, Total\n",
      "1, $\\bf{MiCA}$-$\\bf{ERK}$\n",
      "2, RPI-ERK\n",
      "3, $\\bf{MiCA}$-$\\bf{IGQ}$\n",
      "4, RPI-IGQ\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 108592 characters]",
      "text/plain": [
       "<Figure size 2304x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_pruners = [\n",
    "    'mica_erk_rand',\n",
    "    'mica_erk_shuffle',\n",
    "    'mica_igq_rand',\n",
    "    'mica_igq_shuffle',\n",
    "    ]\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8*4, 8))\n",
    "\n",
    "bar_width = 0.12\n",
    "for i, arch in enumerate(['vgg', 'resnet']):\n",
    "    if arch == 'vgg':\n",
    "        effective_edges = vgg_effective_edges\n",
    "        edges = vgg_edges\n",
    "        xticks = np.arange(12) - bar_width * 2\n",
    "    elif arch == 'resnet':\n",
    "        effective_edges = resnet_effective_edges\n",
    "        edges = resnet_edges \n",
    "        xticks = np.arange(8) - bar_width * 2\n",
    "    ax = axes[i]\n",
    "\n",
    "    for l, pruner in enumerate(plot_pruners):\n",
    "        if l == 0:\n",
    "            ax.bar(xticks, edges[pruner], color='gray', \n",
    "                    label='Total', hatch='x',\n",
    "                    width=bar_width, ec=\"k\",\n",
    "                    align='center')\n",
    "            xticks += bar_width\n",
    "\n",
    "        if '_shuffle' in pruner:\n",
    "            ax.bar(xticks, effective_edges[pruner], color=colors[pruner], \n",
    "                    label=pruner_names[pruner], hatch='/', \n",
    "                    width=bar_width, ec=\"k\", \n",
    "                    align='center')\n",
    "        elif 'mica' in pruner:\n",
    "            ax.bar(xticks, effective_edges[pruner], color=colors[pruner], \n",
    "                    label=pruner_names[pruner],\n",
    "                    width=bar_width, ec=\"k\",\n",
    "                    align='center')\n",
    "        xticks += bar_width\n",
    "\n",
    "    ax.set_ylim(1, 10**6.7)\n",
    "    ax.set_yscale('log')\n",
    "    ax.set_xticks(np.arange(1, 12, 2))\n",
    "    ax.set_xticklabels([fr'$10^{i}$' for i in np.arange(1, 7)])\n",
    "    if arch == 'vgg':\n",
    "        ax.set_xlim(-.5, 11.5)\n",
    "    else:\n",
    "        ax.set_xlim(-.5, 9.5)\n",
    "    # ax.set_yticks(np.arange(10, 100, 20))\n",
    "    ax.minorticks_off()\n",
    "    ax.set_axisbelow(True)\n",
    "    \n",
    "    ax.tick_params(\n",
    "        which       = \"both\", bottom    = True,  top   = False,\n",
    "        labelbottom = True,   left      = True,  right = False,\n",
    "        labelleft   = True,   \n",
    "        direction   = 'out',  length    = 7,     width = 1.0,    \n",
    "        pad         = 8,      labelsize = fontsize\n",
    "        )\n",
    "    ax.set_title('VGG-16' if arch == 'vgg' else 'ResNet-20', size=fontsize)\n",
    "    ax.set_xlabel('Pre-Defined Compression Ratio', size=fontsize)\n",
    "    if i == 0:\n",
    "        ax.set_ylabel('Number of Functional Edges', size=fontsize)\n",
    "    ax.grid()\n",
    "    \n",
    "\n",
    "    # boundary and zoom box\n",
    "    for dir in [\"top\", \"bottom\", \"right\", \"left\"]:\n",
    "        ax.spines[dir].set_linewidth(3)\n",
    "    ax.spines[\"right\"].set_linestyle((0,(3,3)))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.2)\n",
    "fig.tight_layout()\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "for nl, label in enumerate(labels):\n",
    "    print(f'{nl}, {label}')\n",
    "order = [1,3,2,4,0]\n",
    "ordered_handles = get_ordered_list(handles, order)\n",
    "ordered_labels = get_ordered_list(labels, order)\n",
    "plt.legend(\n",
    "    handles=ordered_handles,labels=ordered_labels, \n",
    "    fontsize=fontsize, loc='upper center', \n",
    "    bbox_to_anchor=(-.1, -.25), ncol=3)\n",
    "os.makedirs('Plots', exist_ok=True)\n",
    "plt.savefig('./Plots/cifar10_sparsity_dist.eps', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_sparsity_dist.svg', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_sparsity_dist.pdf', bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vgg\n",
      "[4.65347289e+06 1.47155702e+06 4.65345008e+05 1.47157998e+05\n",
      " 4.65379995e+04 1.47140003e+04 4.65599994e+03 1.46699998e+03\n",
      " 4.67000006e+02 1.46000000e+02 4.70000013e+01 2.10000005e+01]\n",
      "vgg\n",
      "[4.65347289e+06 1.47155702e+06 4.65345008e+05 1.47157998e+05\n",
      " 4.65379995e+04 1.47140003e+04 4.65599994e+03 1.46699998e+03\n",
      " 4.67000006e+02 1.46000000e+02 4.70000013e+01 2.10000005e+01]\n",
      "vgg\n",
      "[4.65347109e+06 1.47154898e+06 4.65342005e+05 1.47148999e+05\n",
      " 4.65280003e+04 1.47070002e+04 4.64800009e+03 1.46200004e+03\n",
      " 4.61000008e+02 1.40000005e+02 4.20000012e+01 0.00000000e+00]\n",
      "vgg\n",
      "[4.65347109e+06 1.47154898e+06 4.65342005e+05 1.47148999e+05\n",
      " 4.65280003e+04 1.47070002e+04 4.64800009e+03 1.46200004e+03\n",
      " 4.61000008e+02 1.40000005e+02 4.20000012e+01 0.00000000e+00]\n",
      "resnet\n",
      "[8.56670008e+04 2.70910000e+04 8.56099994e+03 2.71499998e+03\n",
      " 8.60000002e+02 2.66999999e+02 8.20000015e+01 2.70000005e+01]\n",
      "resnet\n",
      "[8.56670008e+04 2.70910000e+04 8.56099994e+03 2.71499998e+03\n",
      " 8.60000002e+02 2.66999999e+02 8.20000015e+01 2.70000005e+01]\n",
      "resnet\n",
      "[8.56550008e+04 2.70829998e+04 8.55500006e+03 2.69599996e+03\n",
      " 8.43000024e+02 2.64000007e+02 6.60000018e+01 2.20000002e+01]\n",
      "resnet\n",
      "[8.56550008e+04 2.70829998e+04 8.55500006e+03 2.69599996e+03\n",
      " 8.43000024e+02 2.64000007e+02 6.60000018e+01 2.20000002e+01]\n",
      "0, $\\bf{MiCA}$-$\\bf{ERK}$\n",
      "1, RPI-ERK\n",
      "2, $\\bf{MiCA}$-$\\bf{IGQ}$\n",
      "3, RPI-IGQ\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<ipython-input-12-64c9ba7c04cc>:44: RuntimeWarning: invalid value encountered in divide\n",
      "  ax.bar(xticks, np.array(effective_edges[pruner])/np.array(edges[pruner]), color=colors[pruner],\n",
      "<ipython-input-12-64c9ba7c04cc>:39: RuntimeWarning: invalid value encountered in divide\n",
      "  ax.bar(xticks, np.array(effective_edges[pruner])/np.array(edges[pruner]), color=colors[pruner],\n",
      "The PostScript backend does not support transparency; partially transparent artists will be rendered opaque.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 89348 characters]",
      "text/plain": [
       "<Figure size 2304x576 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_pruners = [\n",
    "    'mica_erk_rand',\n",
    "    'mica_erk_shuffle',\n",
    "    'mica_igq_rand',\n",
    "    'mica_igq_shuffle',\n",
    "    ]\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(8*4, 8))\n",
    "\n",
    "\n",
    "for i, arch in enumerate(['vgg', 'resnet']):\n",
    "    if arch == 'vgg':\n",
    "        xmin = -.5\n",
    "        xmax = 11.5\n",
    "        bar_width = (xmax-xmin)/100\n",
    "        effective_edges = vgg_effective_edges\n",
    "        edges = vgg_edges\n",
    "        xticks = np.arange(12) - bar_width * 2\n",
    "    elif arch == 'resnet':\n",
    "        xmin = -.5\n",
    "        xmax = 8.5\n",
    "        bar_width = (xmax-xmin)/100\n",
    "        effective_edges = resnet_effective_edges\n",
    "        edges = resnet_edges \n",
    "        xticks = np.arange(8) - bar_width * 2\n",
    "    \n",
    "    ax = axes[i]\n",
    "\n",
    "    for l, pruner in enumerate(plot_pruners):\n",
    "        # if l == 0:\n",
    "        #     ax.bar(xticks, edges[pruner], color='gray', \n",
    "        #             label='Total', hatch='x',\n",
    "        #             width=bar_width, ec=\"k\",\n",
    "        #             align='center')\n",
    "        #     xticks += bar_width\n",
    "        print(arch)\n",
    "        print(edges[pruner])\n",
    "        if '_shuffle' in pruner:\n",
    "            ax.bar(xticks, np.array(effective_edges[pruner])/np.array(edges[pruner]), color=colors[pruner], \n",
    "                    label=pruner_names[pruner], hatch='/', \n",
    "                    width=bar_width, ec=\"k\", \n",
    "                    align='edge')\n",
    "        elif 'mica' in pruner:\n",
    "            ax.bar(xticks, np.array(effective_edges[pruner])/np.array(edges[pruner]), color=colors[pruner], \n",
    "                    label=pruner_names[pruner],\n",
    "                    width=bar_width, ec=\"k\",\n",
    "                    align='edge')\n",
    "        xticks += bar_width\n",
    "\n",
    "    # ax.set_ylim(1, 10**6.7)\n",
    "    # ax.set_yscale('log')\n",
    "    ax.set_ylim(0, 1)\n",
    "    ax.set_xticks(np.arange(1, 12, 2))\n",
    "    ax.set_xticklabels([fr'$10^{i}$' for i in np.arange(1, 7)])\n",
    "    if arch == 'vgg':\n",
    "        ax.set_xlim(-.5, 11.5)\n",
    "    else:\n",
    "        ax.set_xlim(-.5, 8.5)\n",
    "    # ax.set_yticks(np.arange(10, 100, 20))\n",
    "    ax.minorticks_off()\n",
    "    ax.set_axisbelow(True)\n",
    "    \n",
    "    ax.tick_params(\n",
    "        which       = \"both\", bottom    = True,  top   = False,\n",
    "        labelbottom = True,   left      = True,  right = False,\n",
    "        labelleft   = True if arch == 'vgg' else False,   \n",
    "        direction   = 'out',  length    = 7,     width = 1.0,    \n",
    "        pad         = 8,      labelsize = fontsize\n",
    "        )\n",
    "    ax.set_title('VGG-16' if arch == 'vgg' else 'ResNet-20', size=fontsize)\n",
    "    ax.set_xlabel('Pre-Defined Compression Ratio', size=fontsize)\n",
    "    if i == 0:\n",
    "        ax.set_ylabel('Ratio of Functional Edges', size=fontsize)\n",
    "    ax.grid()\n",
    "    \n",
    "\n",
    "    # boundary and zoom box\n",
    "    for dir in [\"top\", \"bottom\", \"right\", \"left\"]:\n",
    "        ax.spines[dir].set_linewidth(3)\n",
    "    ax.spines[\"right\"].set_linestyle((0,(3,3)))\n",
    "    ax.spines['right'].set_visible(False)\n",
    "    ax.spines['top'].set_visible(False)\n",
    "\n",
    "fig.subplots_adjust(hspace=0.2)\n",
    "fig.tight_layout()\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "for nl, label in enumerate(labels):\n",
    "    print(f'{nl}, {label}')\n",
    "order = [0,2,1,3]\n",
    "ordered_handles = get_ordered_list(handles, order)\n",
    "ordered_labels = get_ordered_list(labels, order)\n",
    "plt.legend(\n",
    "    handles=ordered_handles,labels=ordered_labels, \n",
    "    fontsize=fontsize, loc='upper center', \n",
    "    bbox_to_anchor=(-.1, -.25), ncol=2)\n",
    "os.makedirs('Plots', exist_ok=True)\n",
    "plt.savefig('./Plots/cifar10_sparsity_dist.eps', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_sparsity_dist.svg', bbox_inches=\"tight\")\n",
    "plt.savefig('./Plots/cifar10_sparsity_dist.pdf', bbox_inches=\"tight\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "connected_edge",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
