{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d98765e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import warnings\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm\n",
    "from copy import deepcopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "116617c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "{YOUR_WORKING_PATH}/zoo_utils.py:6: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import re\n",
    "\n",
    "from utils import test_model, compare_models_prediction, generate_z_samples\n",
    "from zoo_utils import load_hyperrep_edx\n",
    "\n",
    "from shrp.models.def_net import CNN, CNN3\n",
    "from shrp.models.def_AE_trainable import get_transformations\n",
    "from shrp.datasets.dataset_auxiliaries import tokens_to_recipe, tokens_to_checkpoint\n",
    "from shrp.models.def_loss import DistillationLoss, ReconDistillationLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0d465bf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "sns.set_theme(context='notebook', style='whitegrid',\n",
    "              rc={\n",
    "                  'axes.spines.right': False,\n",
    "                  'axes.spines.top': False,\n",
    "                  'xtick.bottom': True,\n",
    "                  'ytick.left': True,\n",
    "                  'font.family': 'Fira Sans'\n",
    "              })\n",
    "\n",
    "COLOURS = [\n",
    "    '#960018',  # Carmine Red\n",
    "    '#FED85D',  # Mustard Gold\n",
    "    '#434384',  # Marian Blue\n",
    "    '#54B674',  # Emerald Green\n",
    "    '#D183C9',  # French Mauve\n",
    "    '#318CE7',  # French Blue\n",
    "    '#D08000'   # Fulvous Orange\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a45fdd7",
   "metadata": {},
   "source": [
    "## Loading data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "755588f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASETS = ['svhn', 'cifar10', 'eurosat']\n",
    "DATASET_LABELS = {\n",
    "    'svhn': 'SVHN',\n",
    "    'cifar10': 'CIFAR-10',\n",
    "    'eurosat': 'EuroSAT'\n",
    "}\n",
    "\n",
    "DS_ORDER = ['SVHN', 'CIFAR-10', 'EuroSAT']\n",
    "HR_ORDER = ['Contrastive + Structure (Baseline)', 'Contrastive + Structure + Behavior']\n",
    "\n",
    "HYPERREP_PATH = Path('{YOUR_WORKING_PATH}/results/hyperrepresentations/tune')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e0f27f84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ea53564d50a84775be3c635d2cbf844a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading data:   0%|          | 0/3 [00:00<?, ?dataset/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "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>Model zoo</th>\n",
       "      <th>Trial ID</th>\n",
       "      <th>Learning rate</th>\n",
       "      <th>Latent dimension</th>\n",
       "      <th>Behaviour loss</th>\n",
       "      <th>Query set</th>\n",
       "      <th>Query dataset</th>\n",
       "      <th># of queries</th>\n",
       "      <th>Gamma</th>\n",
       "      <th>Beta</th>\n",
       "      <th>...</th>\n",
       "      <th>Train loss (recon)</th>\n",
       "      <th>Train loss (structure)</th>\n",
       "      <th>Train loss (behaviour)</th>\n",
       "      <th>Test loss</th>\n",
       "      <th>Test loss (contrast)</th>\n",
       "      <th>Test loss (recon)</th>\n",
       "      <th>Test loss (structure)</th>\n",
       "      <th>Test loss (behaviour)</th>\n",
       "      <th>Dstk: Test accuracy</th>\n",
       "      <th>Dstk: GGap</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>1.344194</td>\n",
       "      <td>0.930796</td>\n",
       "      <td>1.390127</td>\n",
       "      <td>1.059503</td>\n",
       "      <td>1.045855</td>\n",
       "      <td>1.060222</td>\n",
       "      <td>0.791913</td>\n",
       "      <td>1.090034</td>\n",
       "      <td>0.559680</td>\n",
       "      <td>0.283559</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.275532</td>\n",
       "      <td>0.770928</td>\n",
       "      <td>0.220488</td>\n",
       "      <td>0.257198</td>\n",
       "      <td>0.887651</td>\n",
       "      <td>0.224016</td>\n",
       "      <td>0.591360</td>\n",
       "      <td>0.183200</td>\n",
       "      <td>0.554503</td>\n",
       "      <td>0.276677</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.266038</td>\n",
       "      <td>0.677502</td>\n",
       "      <td>0.220320</td>\n",
       "      <td>0.247892</td>\n",
       "      <td>0.708432</td>\n",
       "      <td>0.223653</td>\n",
       "      <td>0.544859</td>\n",
       "      <td>0.187963</td>\n",
       "      <td>0.550643</td>\n",
       "      <td>0.276275</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.264123</td>\n",
       "      <td>0.664841</td>\n",
       "      <td>0.219599</td>\n",
       "      <td>0.239949</td>\n",
       "      <td>0.652948</td>\n",
       "      <td>0.218212</td>\n",
       "      <td>0.527920</td>\n",
       "      <td>0.183800</td>\n",
       "      <td>0.551713</td>\n",
       "      <td>0.279411</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.261248</td>\n",
       "      <td>0.628665</td>\n",
       "      <td>0.220424</td>\n",
       "      <td>0.242219</td>\n",
       "      <td>0.741064</td>\n",
       "      <td>0.215964</td>\n",
       "      <td>0.495595</td>\n",
       "      <td>0.184894</td>\n",
       "      <td>0.538348</td>\n",
       "      <td>0.282729</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 23 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "  Model zoo     Trial ID  Learning rate  Latent dimension Behaviour loss  \\\n",
       "0      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "1      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "2      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "3      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "4      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "\n",
       "  Query set Query dataset  # of queries  Gamma  Beta  ...  Train loss (recon)  \\\n",
       "0      data          SVHN           256   0.05   0.1  ...            1.344194   \n",
       "1      data          SVHN           256   0.05   0.1  ...            0.275532   \n",
       "2      data          SVHN           256   0.05   0.1  ...            0.266038   \n",
       "3      data          SVHN           256   0.05   0.1  ...            0.264123   \n",
       "4      data          SVHN           256   0.05   0.1  ...            0.261248   \n",
       "\n",
       "   Train loss (structure)  Train loss (behaviour)  Test loss  \\\n",
       "0                0.930796                1.390127   1.059503   \n",
       "1                0.770928                0.220488   0.257198   \n",
       "2                0.677502                0.220320   0.247892   \n",
       "3                0.664841                0.219599   0.239949   \n",
       "4                0.628665                0.220424   0.242219   \n",
       "\n",
       "   Test loss (contrast)  Test loss (recon)  Test loss (structure)  \\\n",
       "0              1.045855           1.060222               0.791913   \n",
       "1              0.887651           0.224016               0.591360   \n",
       "2              0.708432           0.223653               0.544859   \n",
       "3              0.652948           0.218212               0.527920   \n",
       "4              0.741064           0.215964               0.495595   \n",
       "\n",
       "   Test loss (behaviour)  Dstk: Test accuracy  Dstk: GGap  \n",
       "0               1.090034             0.559680    0.283559  \n",
       "1               0.183200             0.554503    0.276677  \n",
       "2               0.187963             0.550643    0.276275  \n",
       "3               0.183800             0.551713    0.279411  \n",
       "4               0.184894             0.538348    0.282729  \n",
       "\n",
       "[5 rows x 23 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MAX_EPOCH = 100\n",
    "\n",
    "df_hr = pd.DataFrame(columns=[\n",
    "    'Model zoo',\n",
    "    'Trial ID',\n",
    "    'Learning rate',\n",
    "    'Latent dimension',\n",
    "    'Behaviour loss',\n",
    "    'Query set',\n",
    "    'Query dataset',\n",
    "    '# of queries',\n",
    "    'Gamma',\n",
    "    'Beta',\n",
    "    'Epoch',\n",
    "    'Train loss',\n",
    "    'Train loss (contrast)', \n",
    "    'Train loss (recon)',\n",
    "    'Train loss (structure)', \n",
    "    'Train loss (behaviour)',\n",
    "    'Test loss',\n",
    "    'Test loss (contrast)', \n",
    "    'Test loss (recon)',\n",
    "    'Test loss (structure)', \n",
    "    'Test loss (behaviour)',\n",
    "    'Dstk: Test accuracy',\n",
    "    'Dstk: GGap'\n",
    "])\n",
    "\n",
    "for ds in tqdm(DATASETS, desc='Loading data', unit='dataset'):\n",
    "    hr_path = HYPERREP_PATH / f\"{ds}_train_hyperrep\"\n",
    "    for path in os.listdir(hr_path):\n",
    "        if not re.match('AE_trainable', path):\n",
    "            continue\n",
    "        path = hr_path / path\n",
    "\n",
    "        with open(path / 'params.json') as ifh_config, open(path / 'result.json') as ifh_result:\n",
    "            config = json.load(ifh_config)\n",
    "            result = pd.read_json(ifh_result, lines=True, dtype={'trial_id': str})\n",
    "\n",
    "        for edx in range(len(result)):\n",
    "            df_hr.loc[len(df_hr)] = (\n",
    "                DATASET_LABELS[ds],\n",
    "                result.loc[edx]['trial_id'],\n",
    "                config['optim::lr'],\n",
    "                config['ae:lat_dim'],\n",
    "                config['training::loss_distillation'],\n",
    "                config['training::queryset_distillation'],\n",
    "                config['training::queryset_dump'],\n",
    "                config['training::n_queries_distillation'],\n",
    "                config['training::gamma'],\n",
    "                config['training::beta'],\n",
    "                result.loc[edx]['training_iteration'],\n",
    "                result.loc[edx]['loss/loss_train'],\n",
    "                result.loc[edx]['loss/loss_contrast_train'],\n",
    "                result.loc[edx]['loss/loss_recon_train'],\n",
    "                result.loc[edx]['loss/loss_structure_train'],\n",
    "                result.loc[edx]['loss/loss_behaviour_train'],\n",
    "                result.loc[edx]['loss/loss_test'],\n",
    "                result.loc[edx]['loss/loss_contrast_test'],\n",
    "                result.loc[edx]['loss/loss_recon_test'],\n",
    "                result.loc[edx]['loss/loss_structure_test'],\n",
    "                result.loc[edx]['loss/loss_behaviour_test'],\n",
    "                result.loc[edx]['dstk/test_acc_test'],\n",
    "                result.loc[edx]['dstk/ggap_test']\n",
    "            )\n",
    "\n",
    "df_hr['Query dataset'] = df_hr['Query dataset'].map({\n",
    "    '{YOUR_WORKING_PATH}/results/data/svhn_train_zoo_relu/dataset.pt': 'SVHN',\n",
    "    '{YOUR_WORKING_PATH}/results/data/cifar10_train_zoo_relu/dataset.pt': 'CIFAR-10',\n",
    "    '{YOUR_WORKING_PATH}/results/data/eurosat_train_zoo_relu/dataset.pt': 'EuroSAT',\n",
    "    '{YOUR_WORKING_PATH}/results/data/tinyimagenet_32/dataset.pt': 'TinyImagenet-32'\n",
    "})\n",
    "df_hr.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "407a49eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx_dict = {\n",
    "    '00011': 'Contrastive + Structure (Baseline)',\n",
    "    '00000': 'Behavior',\n",
    "    '00006': 'Contrastive + Behavior',\n",
    "    '00002': 'Structure + Behavior',\n",
    "    '00008': 'Contrastive + Structure + Behavior'\n",
    "}\n",
    "\n",
    "df_hr['Hyper-representation'] = df_hr['Trial ID'].str.slice(start=-5).map(idx_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b0e5bf",
   "metadata": {},
   "source": [
    "## Analysing data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "829bdbd0",
   "metadata": {},
   "source": [
    "### Predicting test accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b637ec6e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{lrrr}\n",
      "Model zoo & SVHN & CIFAR-10 & EuroSAT \\\\\n",
      "Hyper-representation &  &  &  \\\\\n",
      "Contrastive + Structure (Baseline) & 0.742 & 0.890 & 0.957 \\\\\n",
      "Behavior & 0.538 & 0.771 & 0.901 \\\\\n",
      "Contrastive + Behavior & 0.752 & 0.893 & 0.950 \\\\\n",
      "Structure + Behavior & 0.887 & 0.939 & 0.966 \\\\\n",
      "Contrastive + Structure + Behavior & 0.886 & 0.947 & 0.969 \\\\\n",
      "Difference & 0.144 & 0.056 & 0.012 \\\\\n",
      "\\end{tabular}\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "</style>\n",
       "<table id=\"T_3f512\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >Model zoo</th>\n",
       "      <th id=\"T_3f512_level0_col0\" class=\"col_heading level0 col0\" >SVHN</th>\n",
       "      <th id=\"T_3f512_level0_col1\" class=\"col_heading level0 col1\" >CIFAR-10</th>\n",
       "      <th id=\"T_3f512_level0_col2\" class=\"col_heading level0 col2\" >EuroSAT</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >Hyper-representation</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "      <th class=\"blank col1\" >&nbsp;</th>\n",
       "      <th class=\"blank col2\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_3f512_level0_row0\" class=\"row_heading level0 row0\" >Contrastive + Structure (Baseline)</th>\n",
       "      <td id=\"T_3f512_row0_col0\" class=\"data row0 col0\" >0.742</td>\n",
       "      <td id=\"T_3f512_row0_col1\" class=\"data row0 col1\" >0.890</td>\n",
       "      <td id=\"T_3f512_row0_col2\" class=\"data row0 col2\" >0.957</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_3f512_level0_row1\" class=\"row_heading level0 row1\" >Behavior</th>\n",
       "      <td id=\"T_3f512_row1_col0\" class=\"data row1 col0\" >0.538</td>\n",
       "      <td id=\"T_3f512_row1_col1\" class=\"data row1 col1\" >0.771</td>\n",
       "      <td id=\"T_3f512_row1_col2\" class=\"data row1 col2\" >0.901</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_3f512_level0_row2\" class=\"row_heading level0 row2\" >Contrastive + Behavior</th>\n",
       "      <td id=\"T_3f512_row2_col0\" class=\"data row2 col0\" >0.752</td>\n",
       "      <td id=\"T_3f512_row2_col1\" class=\"data row2 col1\" >0.893</td>\n",
       "      <td id=\"T_3f512_row2_col2\" class=\"data row2 col2\" >0.950</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_3f512_level0_row3\" class=\"row_heading level0 row3\" >Structure + Behavior</th>\n",
       "      <td id=\"T_3f512_row3_col0\" class=\"data row3 col0\" >0.887</td>\n",
       "      <td id=\"T_3f512_row3_col1\" class=\"data row3 col1\" >0.939</td>\n",
       "      <td id=\"T_3f512_row3_col2\" class=\"data row3 col2\" >0.966</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_3f512_level0_row4\" class=\"row_heading level0 row4\" >Contrastive + Structure + Behavior</th>\n",
       "      <td id=\"T_3f512_row4_col0\" class=\"data row4 col0\" >0.886</td>\n",
       "      <td id=\"T_3f512_row4_col1\" class=\"data row4 col1\" >0.947</td>\n",
       "      <td id=\"T_3f512_row4_col2\" class=\"data row4 col2\" >0.969</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_3f512_level0_row5\" class=\"row_heading level0 row5\" >Difference</th>\n",
       "      <td id=\"T_3f512_row5_col0\" class=\"data row5 col0\" >0.144</td>\n",
       "      <td id=\"T_3f512_row5_col1\" class=\"data row5 col1\" >0.056</td>\n",
       "      <td id=\"T_3f512_row5_col2\" class=\"data row5 col2\" >0.012</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7fcdccbf7370>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_dstk = df_hr[df_hr['Epoch'] == MAX_EPOCH]\n",
    "df_dstk = df_dstk.groupby(['Hyper-representation', 'Model zoo'])['Dstk: Test accuracy'].mean().reset_index().pivot(index='Hyper-representation', columns='Model zoo', values='Dstk: Test accuracy')\n",
    "df_dstk = df_dstk.loc[HR_ORDER, DS_ORDER]\n",
    "df_dstk.loc['Difference'] = df_dstk.loc['Contrastive + Structure + Behavior'] - df_dstk.loc['Contrastive + Structure (Baseline)']\n",
    "df_dstk = df_dstk.style.format(precision=3)\n",
    "print(df_dstk.to_latex())\n",
    "df_dstk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49dc11bc",
   "metadata": {},
   "source": [
    "### Predicting generalization gap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3815eb0a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{lrrr}\n",
      "Model zoo & SVHN & CIFAR-10 & EuroSAT \\\\\n",
      "Hyper-representation &  &  &  \\\\\n",
      "Contrastive + Structure (Baseline) & 0.347 & 0.700 & 0.465 \\\\\n",
      "Behavior & 0.286 & 0.576 & 0.296 \\\\\n",
      "Contrastive + Behavior & 0.337 & 0.710 & 0.482 \\\\\n",
      "Structure + Behavior & 0.378 & 0.785 & 0.484 \\\\\n",
      "Contrastive + Structure + Behavior & 0.368 & 0.785 & 0.529 \\\\\n",
      "Difference & 0.021 & 0.085 & 0.063 \\\\\n",
      "\\end{tabular}\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<style type=\"text/css\">\n",
       "</style>\n",
       "<table id=\"T_fa291\">\n",
       "  <thead>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >Model zoo</th>\n",
       "      <th id=\"T_fa291_level0_col0\" class=\"col_heading level0 col0\" >SVHN</th>\n",
       "      <th id=\"T_fa291_level0_col1\" class=\"col_heading level0 col1\" >CIFAR-10</th>\n",
       "      <th id=\"T_fa291_level0_col2\" class=\"col_heading level0 col2\" >EuroSAT</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th class=\"index_name level0\" >Hyper-representation</th>\n",
       "      <th class=\"blank col0\" >&nbsp;</th>\n",
       "      <th class=\"blank col1\" >&nbsp;</th>\n",
       "      <th class=\"blank col2\" >&nbsp;</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th id=\"T_fa291_level0_row0\" class=\"row_heading level0 row0\" >Contrastive + Structure (Baseline)</th>\n",
       "      <td id=\"T_fa291_row0_col0\" class=\"data row0 col0\" >0.347</td>\n",
       "      <td id=\"T_fa291_row0_col1\" class=\"data row0 col1\" >0.700</td>\n",
       "      <td id=\"T_fa291_row0_col2\" class=\"data row0 col2\" >0.465</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_fa291_level0_row1\" class=\"row_heading level0 row1\" >Behavior</th>\n",
       "      <td id=\"T_fa291_row1_col0\" class=\"data row1 col0\" >0.286</td>\n",
       "      <td id=\"T_fa291_row1_col1\" class=\"data row1 col1\" >0.576</td>\n",
       "      <td id=\"T_fa291_row1_col2\" class=\"data row1 col2\" >0.296</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_fa291_level0_row2\" class=\"row_heading level0 row2\" >Contrastive + Behavior</th>\n",
       "      <td id=\"T_fa291_row2_col0\" class=\"data row2 col0\" >0.337</td>\n",
       "      <td id=\"T_fa291_row2_col1\" class=\"data row2 col1\" >0.710</td>\n",
       "      <td id=\"T_fa291_row2_col2\" class=\"data row2 col2\" >0.482</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_fa291_level0_row3\" class=\"row_heading level0 row3\" >Structure + Behavior</th>\n",
       "      <td id=\"T_fa291_row3_col0\" class=\"data row3 col0\" >0.378</td>\n",
       "      <td id=\"T_fa291_row3_col1\" class=\"data row3 col1\" >0.785</td>\n",
       "      <td id=\"T_fa291_row3_col2\" class=\"data row3 col2\" >0.484</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_fa291_level0_row4\" class=\"row_heading level0 row4\" >Contrastive + Structure + Behavior</th>\n",
       "      <td id=\"T_fa291_row4_col0\" class=\"data row4 col0\" >0.368</td>\n",
       "      <td id=\"T_fa291_row4_col1\" class=\"data row4 col1\" >0.785</td>\n",
       "      <td id=\"T_fa291_row4_col2\" class=\"data row4 col2\" >0.529</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th id=\"T_fa291_level0_row5\" class=\"row_heading level0 row5\" >Difference</th>\n",
       "      <td id=\"T_fa291_row5_col0\" class=\"data row5 col0\" >0.021</td>\n",
       "      <td id=\"T_fa291_row5_col1\" class=\"data row5 col1\" >0.085</td>\n",
       "      <td id=\"T_fa291_row5_col2\" class=\"data row5 col2\" >0.063</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n"
      ],
      "text/plain": [
       "<pandas.io.formats.style.Styler at 0x7fcdcd59e530>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_dstk = df_hr[df_hr['Epoch'] == MAX_EPOCH]\n",
    "df_dstk = df_dstk.groupby(['Hyper-representation', 'Model zoo'])['Dstk: GGap'].mean().reset_index().pivot(index='Hyper-representation', columns='Model zoo', values='Dstk: GGap')\n",
    "df_dstk = df_dstk.loc[HR_ORDER, DS_ORDER]\n",
    "df_dstk.loc['Difference'] = df_dstk.loc['Contrastive + Structure + Behavior'] - df_dstk.loc['Contrastive + Structure (Baseline)']\n",
    "df_dstk = df_dstk.style.format(precision=3)\n",
    "print(df_dstk.to_latex())\n",
    "df_dstk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa5fff5a",
   "metadata": {},
   "source": [
    "### Plotting discriminative downstream tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ab036cd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 258952 characters]",
      "text/plain": [
       "<Figure size 1400x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(2, 3, figsize=(14, 6))\n",
    "\n",
    "for d, ds in enumerate(DATASETS):\n",
    "    ds_label = DATASET_LABELS[ds]\n",
    "    \n",
    "    df_hr_plot = df_hr[df_hr['Model zoo'] == ds_label]\n",
    "    \n",
    "    for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "        sns.lineplot(df_hr_plot[df_hr_plot['Hyper-representation'] == label], x='Epoch', y='Dstk: Test accuracy', color=COLOURS[i], label=label, ax=axs[0][d])\n",
    "        sns.lineplot(df_hr_plot[df_hr_plot['Hyper-representation'] == label], x='Epoch', y='Dstk: GGap', color=COLOURS[i], label=label, ax=axs[1][d])\n",
    "    \n",
    "    axs[0][d].set(title=f\"{ds_label}\\nPredicting test accuracy ($R^2$)\", ylabel='$R^2$ score')\n",
    "    axs[0][d].get_legend().remove()\n",
    "    \n",
    "    axs[1][d].set(title=f\"{ds_label}\\nPredicting generalization gap ($R^2$)\", ylabel='$R^2$ score')\n",
    "    axs[1][d].get_legend().remove()\n",
    "\n",
    "plt.legend(loc=(-2.1, -0.60), ncols=3)\n",
    "plt.subplots_adjust(hspace=0.6, wspace=0.25)\n",
    "plt.savefig('figures/discriminative_dstk.png', dpi=300, bbox_inches='tight')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
