{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 240,
   "id": "45e05b72",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nasbench201.search_model_darts_proj import TinyNetworkDartsProj\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from nasbench201.cell_operations import SearchSpaceNames\n",
    "import nasbench201.utils as ig_utils\n",
    "import torch.utils\n",
    "import torchvision.datasets as dset\n",
    "import numpy as np\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 241,
   "id": "eaa02532",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "torch.cuda.manual_seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 242,
   "id": "29976057",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AGRS():\n",
    "    def __init__(self):\n",
    "        self.data = '../data'\n",
    "        self.dataset = 'cifar10'\n",
    "        self.train_portion = 0.5\n",
    "        self.batch_size = 64\n",
    "        self.init_channels=16\n",
    "        self.layers = 8\n",
    "        self.learning_rate = 0.025\n",
    "        self.learning_rate_min = 0.001\n",
    "        self.momentum = 0.9\n",
    "        self.nesterov = False\n",
    "        self.weight_decay = 3e-4\n",
    "        self.grad_clip = 5\n",
    "        self.cutout = False\n",
    "args = AGRS()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 243,
   "id": "3725b779",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Jocab_Score(ori_model, input, target, weights=None):\n",
    "    model = copy.deepcopy(ori_model)\n",
    "    model.eval()\n",
    "    model.proj_weights = weights\n",
    "    num_edge, num_op = model.num_edge, model.num_op\n",
    "    for i in range(num_edge):\n",
    "        model.candidate_flags[i] = False\n",
    "    batch_size = input.shape[0]\n",
    "    model.K = torch.zeros(batch_size, batch_size).cuda()\n",
    "    model.K_list = {}\n",
    "    def counting_forward_hook(module, inp, out):\n",
    "        if isinstance(inp, tuple):\n",
    "            inp = inp[0]\n",
    "        inp = inp.view(inp.size(0), -1)\n",
    "        x = (inp > 0).float()\n",
    "        K = x @ x.t()\n",
    "        if x.cpu().numpy().sum() == 0:\n",
    "            model.K = model.K\n",
    "        else:\n",
    "            K2 = (1.-x) @ (1.-x.t())\n",
    "            model.K = model.K + K + K2\n",
    "        model.K_list[module.name]=K\n",
    "        #print(module)\n",
    "        \n",
    "\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, nn.ReLU):\n",
    "            module.name = name\n",
    "            module.register_forward_hook(counting_forward_hook)\n",
    "    \n",
    "    input = input.cuda()\n",
    "    model(input)\n",
    "    K = model.K.cpu().numpy()\n",
    "    score = hooklogdet(model.K.cpu().numpy())\n",
    "    #print(model.K_list)\n",
    "    K_list = model.K_list\n",
    "    del model\n",
    "    del input\n",
    "    return score, K,K_list\n",
    "\n",
    "def hooklogdet(K, labels=None):\n",
    "    s, ld = np.linalg.slogdet(K)\n",
    "    return ld"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 244,
   "id": "ae134d08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "train_transform, valid_transform = ig_utils._data_transforms_cifar10(args)\n",
    "train_data = dset.CIFAR10(root=args.data, train=True, download=True, transform=train_transform)\n",
    "valid_data = dset.CIFAR10(root=args.data, train=False, download=True, transform=valid_transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 245,
   "id": "cd9923d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_train = len(train_data)\n",
    "indices = list(range(num_train))\n",
    "split = 64\n",
    "\n",
    "train_queue = torch.utils.data.DataLoader(\n",
    "    train_data, batch_size=args.batch_size,\n",
    "    sampler=torch.utils.data.sampler.SubsetRandomSampler(indices[:split]),\n",
    "    pin_memory=True)\n",
    "input, target = next(iter(train_queue))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1529,
   "id": "e08b4613",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(2)\n",
    "torch.manual_seed(2)\n",
    "torch.cuda.manual_seed(2)\n",
    "from scipy.stats import rankdata\n",
    "input, target = next(iter(train_queue))\n",
    "LAYER=8\n",
    "OPN=4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1530,
   "id": "58c0ad9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nasbench201.cell_operations import OPS\n",
    "class TinyNetwork(nn.Module):\n",
    "    def __init__(self, C, N, num_classes, criterion, affine=False, track_running_stats=True, stem_channels=3):\n",
    "        super(TinyNetwork, self).__init__()\n",
    "        self.stem = nn.Sequential(\n",
    "                nn.Conv2d(stem_channels, C, kernel_size=3, padding=1, bias=False),\n",
    "                nn.BatchNorm2d(C))\n",
    "        op_names=['skip_connect','nor_conv_1x1', 'nor_conv_3x3', 'avg_pool_3x3']\n",
    "        self.N=N\n",
    "        self.edges = nn.ModuleDict()\n",
    "        for i in range(N):\n",
    "            self.edges[str(i)]=nn.ModuleList([OPS[op_name](C, C, 1, affine, track_running_stats) for op_name in op_names])\n",
    "        \n",
    "        self.lastact    = nn.Sequential(nn.BatchNorm2d(C), nn.ReLU(inplace=True))\n",
    "        self.global_pooling = nn.AdaptiveAvgPool2d(1)\n",
    "        self.classifier = nn.Linear(C, num_classes)\n",
    "        self.cos = nn.CosineSimilarity(dim=1, eps=1e-6)\n",
    "        self.weights=[[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4]]\n",
    "        self.weights = np.array(self.weights)\n",
    "        self.weights=torch.from_numpy(self.weights)\n",
    "    \n",
    "    def forward(self, inputs):\n",
    "        weights = self.weights\n",
    "        sum_value=[]\n",
    "        feature = self.stem(inputs)\n",
    "        for i in range(self.N):\n",
    "            feature=sum(op(feature, block_input=True)*w if w==0 else op(feature) * w for op, w in zip(self.edges[str(i)], weights[i]))\n",
    "#             with torch.no_grad():\n",
    "#                 print(self.calc_k(feature))\n",
    "#                 print(torch.mean(torch.abs(feature)))\n",
    "#                 print(torch.count_nonzero((feature>0).float()))\n",
    "                \n",
    "        out = self.lastact(feature)\n",
    "        out = self.global_pooling( out )\n",
    "        out = out.view(out.size(0), -1)\n",
    "        logits = self.classifier(out)\n",
    "        #print(sum_value)\n",
    "        #print('model end')\n",
    "        return logits\n",
    "    \n",
    "    def calc_k(self, inp):\n",
    "        inp = inp.view(inp.size(0), -1)\n",
    "        x = (inp > 0).float()\n",
    "        K = x @ x.t()\n",
    "        if x.cpu().numpy().sum() == 0:\n",
    "            return 0\n",
    "        else:\n",
    "            K2 = (1.-x) @ (1.-x.t())\n",
    "            K =  K + K2\n",
    "            return hooklogdet(K.cpu().numpy())\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1533,
   "id": "b84ded93",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = TinyNetwork(C=16, N=LAYER, num_classes=10, criterion=nn.CrossEntropyLoss())\n",
    "#model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1534,
   "id": "baddff20",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hongkaiw/anaconda2/envs/mct/lib/python3.7/site-packages/torch/tensor.py:593: RuntimeWarning: Iterating over a tensor might cause the trace to be incorrect. Passing a tensor of different shape won't change the number of iterations executed (and might lead to errors or silently give incorrect results).\n",
      "  'incorrect results).', category=RuntimeWarning)\n",
      "/home/hongkaiw/.local/lib/python3.7/site-packages/ipykernel_launcher.py:27: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n"
     ]
    }
   ],
   "source": [
    "import torch.onnx \n",
    "model.eval()\n",
    "  \n",
    "torch.onnx.export(model,         # model being run \n",
    "     input,       # model input (or a tuple for multiple inputs) \n",
    "     \"toy.onnx\",       # where to save the model  \n",
    "     export_params=False,  # store the trained parameter weights inside the model file \n",
    "     opset_version=10,    # the ONNX version to export the model to \n",
    "     do_constant_folding=True,  # whether to execute constant folding for optimization \n",
    "     input_names = ['modelInput'],   # the model's input names \n",
    "     output_names = ['modelOutput'], # the model's output names \n",
    "     dynamic_axes={'modelInput' : {1 : 'batch_size'},    # variable length axes \n",
    "                            'modelOutput' : {0 : 'batch_size'}}) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1522,
   "id": "78f85e24",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "def Jocab_Score(ori_model, input, target, weights=None):\n",
    "    model = copy.deepcopy(ori_model)\n",
    "    model.eval()\n",
    "    model.proj_weights = weights\n",
    "    batch_size = input.shape[0]\n",
    "    model.K = torch.zeros(batch_size, batch_size).cuda()\n",
    "    model.K_list = {}\n",
    "    model.count = 0\n",
    "    def counting_forward_hook(module, inp, out):\n",
    "        if isinstance(inp, tuple):\n",
    "            inp = inp[0]\n",
    "        inp = inp.view(inp.size(0), -1)\n",
    "        #with torch.no_grad():\n",
    "            #print(torch.sum((inp > 0).float()), torch.count_nonzero(inp))\n",
    "        x = (inp > 0).float()\n",
    "        K = x @ x.t()\n",
    "        if x.cpu().numpy().sum() == 0:\n",
    "            model.K = model.K\n",
    "        else:\n",
    "            K2 = (1.-x) @ (1.-x.t())\n",
    "            model.K = model.K + K + K2\n",
    "        model.K_list[module.name]=K\n",
    "        #print(module)\n",
    "        \n",
    "\n",
    "    for name, module in model.named_modules():\n",
    "        if isinstance(module, nn.ReLU):\n",
    "        #if 'ReLU' in str(type(module)):\n",
    "            module.name = name\n",
    "            #print(module)\n",
    "            model.count+=1\n",
    "            module.register_forward_hook(counting_forward_hook)\n",
    "    \n",
    "    input = input.cuda()\n",
    "    model(input, weights)\n",
    "    K = model.K.cpu().numpy()\n",
    "    score = hooklogdet(model.K.cpu().numpy())\n",
    "    #print(model.K_list)\n",
    "    K_list = model.K_list\n",
    "    #print(model.count)\n",
    "    del model\n",
    "    del input\n",
    "    return score, K,K_list\n",
    "\n",
    "def hooklogdet(K, labels=None):\n",
    "    s, ld = np.linalg.slogdet(K)\n",
    "    return ld"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1523,
   "id": "42ae5a33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "0\n",
      "1\n",
      "1\n",
      "1\n",
      "1\n",
      "1\n",
      "2\n",
      "[724.5167, 724.62964, 721.42676, 726.57513, 721.924, 724.0039, 724.2308, 724.328, 724.6446, 723.4378, 723.59174, 726.2936, 726.9928, 722.4523, 723.66644, 727.96545, 727.3341, 722.5211, 723.89703, 727.9818, 726.8876, 723.4647, 724.414, 727.43134, 727.0108, 724.00287, 724.3993, 727.34503, 727.5785, 724.18317, 724.01965, 727.66486]\n",
      "[2.625 1.5   1.875 4.   ]\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "weights=[[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4],[1/4, 1/4, 1/4, 1/4]]\n",
    "\n",
    "\n",
    "pt_score = []\n",
    "avg_skip_rank=np.array([0.0,0.0,0.0,0.0])\n",
    "count_skip = 0\n",
    "crit,K,K_list = Jocab_Score(model, input,  target, weights)\n",
    "for l in range(LAYER):\n",
    "    op_s = []\n",
    "    for o in range(OPN):\n",
    "        w = copy.deepcopy(weights)\n",
    "        w[l][o]=0\n",
    "        crit,K,K_list = Jocab_Score(model, input,  target, w)\n",
    "        pt_score.append(crit)\n",
    "        op_s.append(crit)\n",
    "    avg_skip_rank +=(rankdata(op_s))\n",
    "    select=np.argmin(op_s)\n",
    "    print(select)\n",
    "    if select ==0:\n",
    "        count_skip+=1\n",
    "print(pt_score)\n",
    "print(avg_skip_rank/LAYER)\n",
    "print(count_skip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1467,
   "id": "3133fd42",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "[557.9228, 597.0493, 597.9796, 550.9689]\n",
      "[3. 2. 1. 4.]\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "disc_score = []\n",
    "avg_skip_rank=np.array([0.0,0.0,0.0,0.0])\n",
    "count_skip=0\n",
    "for l in range(LAYER):\n",
    "    op_s = []\n",
    "    for o in range(OPN):\n",
    "        w = copy.deepcopy(weights)\n",
    "        w[l]=np.zeros_like(w[l])\n",
    "        w[l][o]=1\n",
    "        #w[l][0]=1\n",
    "        crit,K,K_list = Jocab_Score(model, input,  target, w)\n",
    "       #print(w)\n",
    "        op_s.append(crit)\n",
    "        disc_score.append(crit)\n",
    "    #print([5-x for x in rankdata(op_s)])\n",
    "    avg_skip_rank +=(5-rankdata(op_s))\n",
    "    select=np.argmax(op_s)\n",
    "    print(select)\n",
    "    if select ==0:\n",
    "        count_skip+=1\n",
    "print(disc_score)\n",
    "print(avg_skip_rank/LAYER)\n",
    "print(count_skip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1468,
   "id": "e091015c",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = copy.deepcopy(weights)\n",
    "arch=[1,1,1,1,1,1,1,1]\n",
    "for i in range(len(arch)):    \n",
    "    w[i]=np.zeros_like(w[i])\n",
    "    w[i][arch[i]]=1\n",
    "crit,K,K_list = Jocab_Score(model, input,  target, w)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1469,
   "id": "05ffd6a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7116abcb5edb4d99822e2a5f26b3cfb1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[557.9228, 597.0493, 597.9796, 550.9689]\n"
     ]
    }
   ],
   "source": [
    "from itertools import combinations_with_replacement,permutations,product\n",
    "from tqdm.notebook import tqdm\n",
    "final_score = []\n",
    "archs=[]\n",
    "archs =list(product([0,1,2,3], repeat=LAYER)) \n",
    "                 \n",
    "archs = [list(x) for x in archs]\n",
    "#print(archs)\n",
    "for i in tqdm(range(len(archs))):\n",
    "    arch = archs[i]\n",
    "    w = copy.deepcopy(weights)\n",
    "    for i in range(len(arch)):\n",
    "        w[i]=np.zeros_like(w[i])\n",
    "        w[i][arch[i]]=1\n",
    "    crit,K,K_list = Jocab_Score(model, input,  target, w)\n",
    "    final_score.append(crit)\n",
    "print(final_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1470,
   "id": "24aab655",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>arch</th>\n",
       "      <th>naswot</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[0]</td>\n",
       "      <td>557.922791</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[1]</td>\n",
       "      <td>597.049316</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[2]</td>\n",
       "      <td>597.979614</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[3]</td>\n",
       "      <td>550.968872</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  arch      naswot\n",
       "0  [0]  557.922791\n",
       "1  [1]  597.049316\n",
       "2  [2]  597.979614\n",
       "3  [3]  550.968872"
      ]
     },
     "execution_count": 1470,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.DataFrame(list(zip(archs, final_score)),columns =['arch', 'naswot'])\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1471,
   "id": "2da3ceab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[557.9227905273438, 597.04931640625, 597.9796142578125, 550.9688720703125]\n"
     ]
    }
   ],
   "source": [
    "index=0\n",
    "best_nwot=[]\n",
    "for l in range(LAYER):\n",
    "    for o in range(OPN): \n",
    "        max_nwot=max(df[df.apply(lambda x: x['arch'][l]==o, axis=1)]['naswot'])\n",
    "        best_nwot.append(max_nwot)\n",
    "print(best_nwot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1472,
   "id": "a03d51ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\n",
      "[3. 2. 1. 4.]\n"
     ]
    }
   ],
   "source": [
    "avg_rank=np.array([0.0,0.0,0.0,0.0])\n",
    "\n",
    "for i in range(LAYER):\n",
    "    #print(np.argmax(best_nwot[i*4:(i+1)*4]))\n",
    "    avg_rank +=(5-rankdata(best_nwot[i*4:(i+1)*4]))\n",
    "    #print((5-rankdata(best_nwot[i*4:(i+1)*4])))\n",
    "print(avg_rank/LAYER)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1473,
   "id": "5357b16a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SpearmanrResult(correlation=1.0, pvalue=0.0)"
      ]
     },
     "execution_count": 1473,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from scipy import stats\n",
    "stats.spearmanr([x*-1 for x in pt_score], best_nwot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1474,
   "id": "4ebc7c45",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[557.9228, 597.0493, 597.9796, 550.9689]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "SpearmanrResult(correlation=1.0, pvalue=0.0)"
      ]
     },
     "execution_count": 1474,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(disc_score)\n",
    "stats.spearmanr(disc_score, best_nwot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1485,
   "id": "02a823f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "cor_dic={}\n",
    "cor_dic['zc_pt(nwot)']=[1.0,0.80,0.82,0.83,0.79,0.69,0.66,0.65]\n",
    "cor_dic['disc_zc(nwot)']=[1.0,0.85,0.89,0.71,0.63,0.35,0.28,0.07]\n",
    "#cor_dic['zc_pt(nwot)_w/o_skip']  =[1.0,0.5217,0.5533,0.6655,0.7019,0.5058,0.5801,0.637]\n",
    "#cor_dic['disc_zc(nwot)_w/o_skip']=[1.0,0.4638,0.5788,0.7332,0.7419,0.5711,0.6141,0.6845]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1486,
   "id": "66c58b61",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "import statistics as stat\n",
    "import itertools\n",
    "marker = itertools.cycle(('^', 'x', 'o', 's', '*', '+', '1')) \n",
    "color = itertools.cycle(('b', 'c', 'r', 'g', 'y', 'm', 'k')) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1488,
   "id": "89d82376",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49016 characters]\n",
      "text/plain": [
       "<Figure size 936x936 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pathlib\n",
    "import seaborn as sns\n",
    "x = np.array([1, 2, 3, 4, 5, 6,7,8])\n",
    "#plt.figure(figsize=(8.5,5.5))\n",
    "idx=0\n",
    "cor_df = pd.DataFrame.from_dict(cor_dic, orient='index')\n",
    "# Set up the matplotlib figure\n",
    "f, ax = plt.subplots(figsize=(13, 13))\n",
    "sns.set(font_scale=1.8)\n",
    "# Generate a custom diverging colormap\n",
    "cmap = sns.diverging_palette(200, 55, as_cmap=True)\n",
    "g = sns.heatmap(cor_df, cmap=cmap, center=0,\n",
    "         square=True, linewidths=.5, cbar_kws={\"shrink\": .267}, annot=True,xticklabels=x)\n",
    "plt.xlabel('#N Layers')\n",
    "plt.ylabel('Spearman-$\\\\rho$')\n",
    "#plt.xticks(x)\n",
    "#plt.legend(bbox_to_anchor=(0.55, 0.65), prop={'size': 13})\n",
    "#plt.grid()\n",
    "plt.savefig(pathlib.Path('op_correl_layer_increase_toy').with_suffix('.pdf'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 999,
   "id": "7876c4d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "rank_dic={}\n",
    "rank_dic['zc_pt(nwot)']  =[4,4,4,4,2.6,2.67,2.71,2.75]\n",
    "rank_dic['disc_zc(nwot)']=[4,4,4,2.5,2.2,2.33,2.28,2.25]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1087,
   "id": "afca5839",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 44208 characters]\n",
      "text/plain": [
       "<Figure size 936x936 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pathlib\n",
    "x = np.array([1, 2, 3, 4, 5, 6,7,8])\n",
    "idx=0\n",
    "cor_df = pd.DataFrame.from_dict(rank_dic, orient='index')\n",
    "# Set up the matplotlib figure\n",
    "f, ax = plt.subplots(figsize=(13, 13))\n",
    "#sns.set(font_scale=1.8)\n",
    "# Generate a custom diverging colormap\n",
    "cmap = sns.diverging_palette(200, 55, as_cmap=True)\n",
    "g = sns.heatmap(cor_df, cmap=cmap, center=0,\n",
    "         square=True, linewidths=.5, cbar_kws={\"shrink\": .267}, annot=True,xticklabels=x)\n",
    "plt.xlabel('#N Layers in Toy model')\n",
    "plt.ylabel('Everage Rank for Skip')\n",
    "#plt.legend(bbox_to_anchor=(0.55, 0.65), prop={'size': 13})\n",
    "#plt.grid()\n",
    "plt.savefig(pathlib.Path('skip_layer_increase_toy').with_suffix('.pdf'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1489,
   "id": "bcc59139",
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_dic={}\n",
    " #[3.         1.66666667 1.33333333 4.        ]             \n",
    "pt_dic['skip']  =     [3,3.00,3.00,3.00,3.00,2.83,2.71,2.75]# 00000011\n",
    "pt_dic['skip_select']=[0,0,0,0,0,0,1,1] \n",
    "pt_dic['conv_1x1']=   [2,1.50,1.67,1.75,1.60,1.83,1.71,1.75]#01112222\n",
    "pt_dic['conv_1x1_select']=[0,1,1,1,2,2,2,2]\n",
    "pt_dic['conv_3x3']=   [1,1.50,1.33,1.25,1.40,1.33,1.57,1.5]#21233445\n",
    "pt_dic['conv_3x3_select']=[2,1,2,3,3,4,4,5]\n",
    "pt_dic['avg_pooling']=[4,4.00,4.00,4.00,4.00,4.00,4.00,4.00]#00000000\n",
    "pt_dic['avg_pooling_select']=[0,0,0,0,0,0,0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1490,
   "id": "30f0c50c",
   "metadata": {},
   "outputs": [],
   "source": [
    "disc_dic={}\n",
    "#[3.         1.33333333 1.66666667 4.        ]\n",
    "disc_dic['skip']  =     [3,3.00,3.00,2.75,2.40,2.17,2.14,1.75]#00001225\n",
    "disc_dic['skip_select']=[0,0,0,0,1,2,2,5] \n",
    "disc_dic['conv_1x1']=   [2,2.00,1.33,1.50,2.00,2.00,2.57,2.75]#00221321\n",
    "disc_dic['conv_1x1_select']=[0,0,2,2,1,3,2,1]\n",
    "disc_dic['conv_3x3']=   [1,1.00,1.67,1.75,1.80,2.17,2.00,2.25]#22123132\n",
    "disc_dic['conv_3x3_select']=[2,2,1,2,3,1,3,2]\n",
    "disc_dic['avg_pooling']=[4,4.00,4.00,4.00,3.80,3.67,3.29,3.25]#00000000\n",
    "disc_dic['avg_pooling_select']=[0,0,0,0,0,0,0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1491,
   "id": "acbdfcc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "best_dic={}\n",
    "\n",
    "best_dic['skip']  =     [3,3.00,3.00,3.00,3.00,3.00,3.00,3.00]#00000000\n",
    "best_dic['skip_select']=[0,0,0,0,0,0,0,0] \n",
    "best_dic['conv_1x1']=   [2,1.50,1.67,2.00,1.80,1.50,1.42,1.5]#01101344\n",
    "best_dic['conv_1x1_select']=[0,1,1,0,1,3,4,4]\n",
    "best_dic['conv_3x3']=   [1,1.50,1.33,1.00,1.20,1.50,1.57,1.5]#21244334\n",
    "best_dic['conv_3x3_select']=[2,1,2,4,4,3,3,4]\n",
    "best_dic['avg_pooling']=[4,4.00,4.00,4.00,4.00,4.00,4.00,4.00]#00000000\n",
    "best_dic['avg_pooling_select']=[0,0,0,0,0,0,0,0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1496,
   "id": "5c0200f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62164 characters]\n",
      "text/plain": [
       "<Figure size 936x936 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pathlib\n",
    "x = np.array([1, 2, 3, 4, 5, 6,7,8])\n",
    "idx=0\n",
    "pt_df = pd.DataFrame.from_dict(pt_dic, orient='index')\n",
    "# Set up the matplotlib figure\n",
    "f, ax = plt.subplots(figsize=(13, 13))\n",
    "#sns.set(font_scale=1.8)\n",
    "# Generate a custom diverging colormap\n",
    "cmap = sns.diverging_palette(200, 55, as_cmap=True)\n",
    "g = sns.heatmap(pt_df, cmap=cmap, center=0,\n",
    "         square=True, linewidths=.5, cbar_kws={\"shrink\": .267}, annot=True,xticklabels=x)\n",
    "plt.xlabel('#N Layers')\n",
    "plt.ylabel('')\n",
    "#plt.legend(bbox_to_anchor=(0.55, 0.65), prop={'size': 13})\n",
    "#plt.grid()\n",
    "plt.savefig(pathlib.Path('op_select_pt').with_suffix('.pdf'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1497,
   "id": "c96d7468",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 69384 characters]\n",
      "text/plain": [
       "<Figure size 936x936 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pathlib\n",
    "x = np.array([1, 2, 3, 4, 5, 6,7,8])\n",
    "idx=0\n",
    "disc_df = pd.DataFrame.from_dict(disc_dic, orient='index')\n",
    "# Set up the matplotlib figure\n",
    "f, ax = plt.subplots(figsize=(13, 13))\n",
    "#sns.set(font_scale=1.8)\n",
    "# Generate a custom diverging colormap\n",
    "cmap = sns.diverging_palette(200, 55, as_cmap=True)\n",
    "g = sns.heatmap(disc_df, cmap=cmap, center=0,\n",
    "         square=True, linewidths=.5, cbar_kws={\"shrink\": .267}, annot=True,xticklabels=x)\n",
    "plt.xlabel('#N Layers')\n",
    "plt.ylabel('')\n",
    "#plt.legend(bbox_to_anchor=(0.55, 0.65), prop={'size': 13})\n",
    "#plt.grid()\n",
    "plt.savefig(pathlib.Path('op_select_disc').with_suffix('.pdf'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1498,
   "id": "b9a92916",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57848 characters]\n",
      "text/plain": [
       "<Figure size 936x936 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import pathlib\n",
    "x = np.array([1, 2, 3, 4, 5, 6,7,8])\n",
    "idx=0\n",
    "best_df = pd.DataFrame.from_dict(best_dic, orient='index')\n",
    "# Set up the matplotlib figure\n",
    "f, ax = plt.subplots(figsize=(13, 13))\n",
    "#sns.set(font_scale=1.8)\n",
    "# Generate a custom diverging colormap\n",
    "cmap = sns.diverging_palette(200, 55, as_cmap=True)\n",
    "g = sns.heatmap(best_df, cmap=cmap, center=0,\n",
    "         square=True, linewidths=.5, cbar_kws={\"shrink\": .267}, annot=True,xticklabels=x)\n",
    "plt.xlabel('#N Layers')\n",
    "plt.ylabel('')\n",
    "#plt.legend(bbox_to_anchor=(0.55, 0.65), prop={'size': 13})\n",
    "#plt.grid()\n",
    "plt.savefig(pathlib.Path('op_select_best').with_suffix('.pdf'), bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ef09034",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mct",
   "language": "python",
   "name": "mct"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
