{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "42d2400a",
   "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": null,
   "id": "bcc6020d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import re\n",
    "\n",
    "from model_distances import distance_matrix\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": "6b010aea",
   "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": "370bf67e",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6a0cb5a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "RERUN = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a065ba3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENT_PATH = Path.cwd() / 'results' / 'ablation_queryset'\n",
    "\n",
    "if not EXPERIMENT_PATH.is_dir():\n",
    "    os.makedirs(EXPERIMENT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bd9383f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASETS = ['svhn', 'cifar10', 'eurosat']\n",
    "DATASET_LABELS = {\n",
    "    'svhn': 'SVHN',\n",
    "    'cifar10': 'CIFAR-10',\n",
    "    'eurosat': 'EuroSAT'\n",
    "}\n",
    "DS_ORDER = ['SVHN', 'CIFAR-10', 'EuroSAT']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4683468b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test set corresponding to the data used to train the model zoo\n",
    "TESTSET_DS = {ds: torch.load(f'{YOUR_WORKING_PATH}/results/data/{ds}_train_zoo_relu/dataset.pt')['testset'] for ds in DATASETS}\n",
    "TESTSET_DL = {ds: torch.utils.data.DataLoader(TESTSET_DS[ds], batch_size=64, shuffle=False) for ds in DATASETS}\n",
    "\n",
    "# Model zoo\n",
    "ZOO_PATH = {ds: Path(f'{YOUR_WORKING_PATH}/results/zoos/{ds}_train_zoo_relu') for ds in DATASETS}\n",
    "ZOO_DS = {ds: torch.load(ZOO_PATH[ds] / 'dataset_test.pt') for ds in DATASETS}\n",
    "ZOO_DL = {ds: torch.utils.data.DataLoader(ZOO_DS[ds], batch_size=16, shuffle=False) for ds in DATASETS}\n",
    "\n",
    "# Threshold used to define models that \"perform well\"; depends on the zoo\n",
    "MODEL_ACC_THR = {\n",
    "    'svhn': 0.8,\n",
    "    'cifar10': 0.6,\n",
    "    'eurosat': 0.7\n",
    "}\n",
    "\n",
    "# Path to the hyper-representation models\n",
    "HYPERREP_PATH = {ds: Path(f'{YOUR_WORKING_PATH}/results/hyperrepresentations/tune/{ds}_train_hyperrep_queryset') for ds in DATASETS}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd212da7",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "### Loading HR models meta-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ed29f0be",
   "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>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>5c618_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</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.971960</td>\n",
       "      <td>1.350016</td>\n",
       "      <td>0.929954</td>\n",
       "      <td>0.945631</td>\n",
       "      <td>1.008066</td>\n",
       "      <td>0.942345</td>\n",
       "      <td>1.146405</td>\n",
       "      <td>0.919672</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>5c618_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</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.204838</td>\n",
       "      <td>0.975027</td>\n",
       "      <td>0.119262</td>\n",
       "      <td>0.213107</td>\n",
       "      <td>0.817504</td>\n",
       "      <td>0.181297</td>\n",
       "      <td>0.732977</td>\n",
       "      <td>0.119999</td>\n",
       "      <td>0.533971</td>\n",
       "      <td>0.279194</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>5c618_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</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.201126</td>\n",
       "      <td>0.950988</td>\n",
       "      <td>0.117807</td>\n",
       "      <td>0.205383</td>\n",
       "      <td>0.634964</td>\n",
       "      <td>0.182774</td>\n",
       "      <td>0.756175</td>\n",
       "      <td>0.119063</td>\n",
       "      <td>0.531583</td>\n",
       "      <td>0.283261</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>5c618_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</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.200538</td>\n",
       "      <td>0.950121</td>\n",
       "      <td>0.117251</td>\n",
       "      <td>0.202396</td>\n",
       "      <td>0.571917</td>\n",
       "      <td>0.182947</td>\n",
       "      <td>0.753049</td>\n",
       "      <td>0.119603</td>\n",
       "      <td>0.530918</td>\n",
       "      <td>0.288594</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>5c618_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</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.196003</td>\n",
       "      <td>0.905218</td>\n",
       "      <td>0.117201</td>\n",
       "      <td>0.200439</td>\n",
       "      <td>0.629362</td>\n",
       "      <td>0.177865</td>\n",
       "      <td>0.712369</td>\n",
       "      <td>0.118475</td>\n",
       "      <td>0.537369</td>\n",
       "      <td>0.294035</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  5c618_00001        0.00001                64             l2   \n",
       "1      SVHN  5c618_00001        0.00001                64             l2   \n",
       "2      SVHN  5c618_00001        0.00001                64             l2   \n",
       "3      SVHN  5c618_00001        0.00001                64             l2   \n",
       "4      SVHN  5c618_00001        0.00001                64             l2   \n",
       "\n",
       "  Query set Query dataset  # of queries  Gamma  Beta  ...  Train loss (recon)  \\\n",
       "0    random          SVHN           256   0.05   0.1  ...            0.971960   \n",
       "1    random          SVHN           256   0.05   0.1  ...            0.204838   \n",
       "2    random          SVHN           256   0.05   0.1  ...            0.201126   \n",
       "3    random          SVHN           256   0.05   0.1  ...            0.200538   \n",
       "4    random          SVHN           256   0.05   0.1  ...            0.196003   \n",
       "\n",
       "   Train loss (structure)  Train loss (behaviour)  Test loss  \\\n",
       "0                1.350016                0.929954   0.945631   \n",
       "1                0.975027                0.119262   0.213107   \n",
       "2                0.950988                0.117807   0.205383   \n",
       "3                0.950121                0.117251   0.202396   \n",
       "4                0.905218                0.117201   0.200439   \n",
       "\n",
       "   Test loss (contrast)  Test loss (recon)  Test loss (structure)  \\\n",
       "0              1.008066           0.942345               1.146405   \n",
       "1              0.817504           0.181297               0.732977   \n",
       "2              0.634964           0.182774               0.756175   \n",
       "3              0.571917           0.182947               0.753049   \n",
       "4              0.629362           0.177865               0.712369   \n",
       "\n",
       "   Test loss (behaviour)  Dstk: Test accuracy  Dstk: GGap  \n",
       "0               0.919672             0.559680    0.283559  \n",
       "1               0.119999             0.533971    0.279194  \n",
       "2               0.119063             0.531583    0.283261  \n",
       "3               0.119603             0.530918    0.288594  \n",
       "4               0.118475             0.537369    0.294035  \n",
       "\n",
       "[5 rows x 23 columns]"
      ]
     },
     "execution_count": 8,
     "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 DATASETS:\n",
    "    for path in os.listdir(HYPERREP_PATH[ds]):\n",
    "        if not re.match('AE_trainable', path):\n",
    "            continue\n",
    "        path = HYPERREP_PATH[ds] / 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",
    "    '{YOUR_WORKING_PATH}/results/data/stl10_32/dataset.pt': 'STL-10_32'\n",
    "})\n",
    "df_hr.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c67a911",
   "metadata": {},
   "source": [
    "### Loading checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b83c6afd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading hyper-reps\n",
    "def load_hr(trial_id):\n",
    "    for ds in DATASETS:\n",
    "        for path in os.listdir(HYPERREP_PATH[ds]):\n",
    "            if trial_id in path:\n",
    "                return load_hyperrep_edx(HYPERREP_PATH[ds] / path, MAX_EPOCH)\n",
    "\n",
    "hyper_reps = {idx: load_hr(idx) for idx in df_hr['Trial ID'].unique()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "77f60054",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading reference config\n",
    "trial_id = df_hr.iloc[0]['Trial ID']\n",
    "for path in os.listdir(HYPERREP_PATH['svhn']):\n",
    "    if trial_id in path:\n",
    "        model_path = HYPERREP_PATH['svhn'] / path\n",
    "        with open(HYPERREP_PATH['svhn'] / path / 'params.json') as ifh:\n",
    "            ref_config = json.load(ifh)\n",
    "\n",
    "# Loading reference params\n",
    "ref_path = ref_config[\"training::distil_reference\"]\n",
    "with open(os.path.join(ref_path, 'params.json')) as ifh:\n",
    "    reference_params = json.load(ifh)\n",
    "    \n",
    "# Loading reference checkpoint\n",
    "ref_chkpth = None\n",
    "for chkpth in os.listdir(ref_path):\n",
    "    if 'checkpoint_' in chkpth:\n",
    "        if not ref_chkpth or int(chkpth.split('_')[-1]) > int(ref_chkpth.split('_')[-1]):\n",
    "            ref_chkpth = chkpth\n",
    "\n",
    "reference_checkpoint = torch.load(Path(ref_path) / str(ref_chkpth) / 'checkpoints', map_location=torch.device(\"cpu\"))\n",
    "reference_checkpoint_cuda = {k: v.to('cuda') for k, v in deepcopy(reference_checkpoint).items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca5bd807",
   "metadata": {},
   "source": [
    "### Extracting key models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6ba506ff",
   "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>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>100</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>5c618_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</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.171891</td>\n",
       "      <td>0.846770</td>\n",
       "      <td>0.096904</td>\n",
       "      <td>0.184744</td>\n",
       "      <td>0.302095</td>\n",
       "      <td>0.178568</td>\n",
       "      <td>0.681341</td>\n",
       "      <td>0.122704</td>\n",
       "      <td>0.740386</td>\n",
       "      <td>0.336324</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>201</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>5c618_00002</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>STL-10_32</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.290401</td>\n",
       "      <td>0.627302</td>\n",
       "      <td>0.252967</td>\n",
       "      <td>0.287471</td>\n",
       "      <td>0.289379</td>\n",
       "      <td>0.287370</td>\n",
       "      <td>0.505372</td>\n",
       "      <td>0.263148</td>\n",
       "      <td>0.748346</td>\n",
       "      <td>0.352681</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>302</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>5c618_00000</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</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.101269</td>\n",
       "      <td>0.655730</td>\n",
       "      <td>0.039662</td>\n",
       "      <td>0.099264</td>\n",
       "      <td>0.294052</td>\n",
       "      <td>0.089012</td>\n",
       "      <td>0.542366</td>\n",
       "      <td>0.038639</td>\n",
       "      <td>0.894404</td>\n",
       "      <td>0.397372</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>403</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>5c618_00003</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</td>\n",
       "      <td>STL-10_32</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.152279</td>\n",
       "      <td>0.654490</td>\n",
       "      <td>0.096478</td>\n",
       "      <td>0.168595</td>\n",
       "      <td>0.300918</td>\n",
       "      <td>0.161631</td>\n",
       "      <td>0.529810</td>\n",
       "      <td>0.120722</td>\n",
       "      <td>0.722982</td>\n",
       "      <td>0.346971</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>504</th>\n",
       "      <td>CIFAR-10</td>\n",
       "      <td>e7ff1_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</td>\n",
       "      <td>CIFAR-10</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.123115</td>\n",
       "      <td>0.843542</td>\n",
       "      <td>0.043067</td>\n",
       "      <td>0.152219</td>\n",
       "      <td>0.306361</td>\n",
       "      <td>0.144106</td>\n",
       "      <td>0.653344</td>\n",
       "      <td>0.087524</td>\n",
       "      <td>0.872274</td>\n",
       "      <td>0.683310</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>605</th>\n",
       "      <td>CIFAR-10</td>\n",
       "      <td>e7ff1_00003</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</td>\n",
       "      <td>STL-10_32</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.122027</td>\n",
       "      <td>0.800268</td>\n",
       "      <td>0.046667</td>\n",
       "      <td>0.154067</td>\n",
       "      <td>0.294455</td>\n",
       "      <td>0.146679</td>\n",
       "      <td>0.621572</td>\n",
       "      <td>0.093913</td>\n",
       "      <td>0.870390</td>\n",
       "      <td>0.691756</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>706</th>\n",
       "      <td>CIFAR-10</td>\n",
       "      <td>e7ff1_00002</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>STL-10_32</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.147304</td>\n",
       "      <td>0.654013</td>\n",
       "      <td>0.091003</td>\n",
       "      <td>0.142485</td>\n",
       "      <td>0.282516</td>\n",
       "      <td>0.135115</td>\n",
       "      <td>0.521082</td>\n",
       "      <td>0.092229</td>\n",
       "      <td>0.953338</td>\n",
       "      <td>0.792475</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>807</th>\n",
       "      <td>CIFAR-10</td>\n",
       "      <td>e7ff1_00000</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>CIFAR-10</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.134091</td>\n",
       "      <td>0.597939</td>\n",
       "      <td>0.082552</td>\n",
       "      <td>0.130553</td>\n",
       "      <td>0.265260</td>\n",
       "      <td>0.123463</td>\n",
       "      <td>0.481218</td>\n",
       "      <td>0.083713</td>\n",
       "      <td>0.944333</td>\n",
       "      <td>0.784994</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>908</th>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>4b0a2_00003</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</td>\n",
       "      <td>STL-10_32</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.098722</td>\n",
       "      <td>0.604228</td>\n",
       "      <td>0.042555</td>\n",
       "      <td>0.147919</td>\n",
       "      <td>0.350148</td>\n",
       "      <td>0.137276</td>\n",
       "      <td>0.746369</td>\n",
       "      <td>0.069599</td>\n",
       "      <td>0.935109</td>\n",
       "      <td>0.419745</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1009</th>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>4b0a2_00001</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>random</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.111631</td>\n",
       "      <td>0.759297</td>\n",
       "      <td>0.039668</td>\n",
       "      <td>0.163295</td>\n",
       "      <td>0.364696</td>\n",
       "      <td>0.152695</td>\n",
       "      <td>0.944515</td>\n",
       "      <td>0.064715</td>\n",
       "      <td>0.932646</td>\n",
       "      <td>0.420665</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1110</th>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>4b0a2_00000</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.142199</td>\n",
       "      <td>0.760975</td>\n",
       "      <td>0.073447</td>\n",
       "      <td>0.162952</td>\n",
       "      <td>0.280402</td>\n",
       "      <td>0.156770</td>\n",
       "      <td>0.922439</td>\n",
       "      <td>0.071696</td>\n",
       "      <td>0.970460</td>\n",
       "      <td>0.547347</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1211</th>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>4b0a2_00002</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>STL-10_32</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.229940</td>\n",
       "      <td>0.845695</td>\n",
       "      <td>0.161522</td>\n",
       "      <td>0.276827</td>\n",
       "      <td>0.286946</td>\n",
       "      <td>0.276294</td>\n",
       "      <td>1.037734</td>\n",
       "      <td>0.191690</td>\n",
       "      <td>0.954707</td>\n",
       "      <td>0.473230</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>12 rows × 23 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     Model zoo     Trial ID  Learning rate  Latent dimension Behaviour loss  \\\n",
       "100       SVHN  5c618_00001        0.00001                64             l2   \n",
       "201       SVHN  5c618_00002        0.00001                64             l2   \n",
       "302       SVHN  5c618_00000        0.00001                64             l2   \n",
       "403       SVHN  5c618_00003        0.00001                64             l2   \n",
       "504   CIFAR-10  e7ff1_00001        0.00001                64             l2   \n",
       "605   CIFAR-10  e7ff1_00003        0.00001                64             l2   \n",
       "706   CIFAR-10  e7ff1_00002        0.00001                64             l2   \n",
       "807   CIFAR-10  e7ff1_00000        0.00001                64             l2   \n",
       "908    EuroSAT  4b0a2_00003        0.00001                64             l2   \n",
       "1009   EuroSAT  4b0a2_00001        0.00001                64             l2   \n",
       "1110   EuroSAT  4b0a2_00000        0.00001                64             l2   \n",
       "1211   EuroSAT  4b0a2_00002        0.00001                64             l2   \n",
       "\n",
       "     Query set Query dataset  # of queries  Gamma  Beta  ...  \\\n",
       "100     random          SVHN           256   0.05   0.1  ...   \n",
       "201       data     STL-10_32           256   0.05   0.1  ...   \n",
       "302       data          SVHN           256   0.05   0.1  ...   \n",
       "403     random     STL-10_32           256   0.05   0.1  ...   \n",
       "504     random      CIFAR-10           256   0.05   0.1  ...   \n",
       "605     random     STL-10_32           256   0.05   0.1  ...   \n",
       "706       data     STL-10_32           256   0.05   0.1  ...   \n",
       "807       data      CIFAR-10           256   0.05   0.1  ...   \n",
       "908     random     STL-10_32           256   0.05   0.1  ...   \n",
       "1009    random       EuroSAT           256   0.05   0.1  ...   \n",
       "1110      data       EuroSAT           256   0.05   0.1  ...   \n",
       "1211      data     STL-10_32           256   0.05   0.1  ...   \n",
       "\n",
       "      Train loss (recon)  Train loss (structure)  Train loss (behaviour)  \\\n",
       "100             0.171891                0.846770                0.096904   \n",
       "201             0.290401                0.627302                0.252967   \n",
       "302             0.101269                0.655730                0.039662   \n",
       "403             0.152279                0.654490                0.096478   \n",
       "504             0.123115                0.843542                0.043067   \n",
       "605             0.122027                0.800268                0.046667   \n",
       "706             0.147304                0.654013                0.091003   \n",
       "807             0.134091                0.597939                0.082552   \n",
       "908             0.098722                0.604228                0.042555   \n",
       "1009            0.111631                0.759297                0.039668   \n",
       "1110            0.142199                0.760975                0.073447   \n",
       "1211            0.229940                0.845695                0.161522   \n",
       "\n",
       "      Test loss  Test loss (contrast)  Test loss (recon)  \\\n",
       "100    0.184744              0.302095           0.178568   \n",
       "201    0.287471              0.289379           0.287370   \n",
       "302    0.099264              0.294052           0.089012   \n",
       "403    0.168595              0.300918           0.161631   \n",
       "504    0.152219              0.306361           0.144106   \n",
       "605    0.154067              0.294455           0.146679   \n",
       "706    0.142485              0.282516           0.135115   \n",
       "807    0.130553              0.265260           0.123463   \n",
       "908    0.147919              0.350148           0.137276   \n",
       "1009   0.163295              0.364696           0.152695   \n",
       "1110   0.162952              0.280402           0.156770   \n",
       "1211   0.276827              0.286946           0.276294   \n",
       "\n",
       "      Test loss (structure)  Test loss (behaviour)  Dstk: Test accuracy  \\\n",
       "100                0.681341               0.122704             0.740386   \n",
       "201                0.505372               0.263148             0.748346   \n",
       "302                0.542366               0.038639             0.894404   \n",
       "403                0.529810               0.120722             0.722982   \n",
       "504                0.653344               0.087524             0.872274   \n",
       "605                0.621572               0.093913             0.870390   \n",
       "706                0.521082               0.092229             0.953338   \n",
       "807                0.481218               0.083713             0.944333   \n",
       "908                0.746369               0.069599             0.935109   \n",
       "1009               0.944515               0.064715             0.932646   \n",
       "1110               0.922439               0.071696             0.970460   \n",
       "1211               1.037734               0.191690             0.954707   \n",
       "\n",
       "      Dstk: GGap  \n",
       "100     0.336324  \n",
       "201     0.352681  \n",
       "302     0.397372  \n",
       "403     0.346971  \n",
       "504     0.683310  \n",
       "605     0.691756  \n",
       "706     0.792475  \n",
       "807     0.784994  \n",
       "908     0.419745  \n",
       "1009    0.420665  \n",
       "1110    0.547347  \n",
       "1211    0.473230  \n",
       "\n",
       "[12 rows x 23 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_hr[(df_hr['Epoch'] == MAX_EPOCH)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d1096e4b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'svhn': {'5c618_00000': 'SVHN',\n",
       "  '5c618_00002': 'STL-10_32',\n",
       "  '5c618_00001': 'Random'},\n",
       " 'cifar10': {'e7ff1_00000': 'CIFAR-10',\n",
       "  'e7ff1_00002': 'STL-10_32',\n",
       "  'e7ff1_00001': 'Random'},\n",
       " 'eurosat': {'4b0a2_00000': 'EuroSAT',\n",
       "  '4b0a2_00002': 'STL-10_32',\n",
       "  '4b0a2_00001': 'Random'}}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "IDX_DICT = {ds:\n",
    "    {\n",
    "        df_hr[(df_hr['Model zoo'] == DATASET_LABELS[ds]) & (df_hr['Query set'] == 'data') & (df_hr['Query dataset'] == DATASET_LABELS[ds])]['Trial ID'].iloc[0]: DATASET_LABELS[ds],\n",
    "        df_hr[(df_hr['Model zoo'] == DATASET_LABELS[ds]) & (df_hr['Query set'] == 'data') & (df_hr['Query dataset'] == 'STL-10_32')]['Trial ID'].iloc[0]: \"STL-10_32\",\n",
    "        df_hr[(df_hr['Model zoo'] == DATASET_LABELS[ds]) & (df_hr['Query set'] == 'random') & (df_hr['Query dataset'] == DATASET_LABELS[ds])]['Trial ID'].iloc[0]: \"Random\"\n",
    "    }\n",
    "    for ds in DATASETS\n",
    "}\n",
    "IDX_DICT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a887f1ba",
   "metadata": {},
   "source": [
    "## Analysing loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "306351ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 56000 characters]",
      "text/plain": [
       "<Figure size 1400x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(1, 3, figsize=(14, 3))\n",
    "\n",
    "for d, ds in enumerate(DATASETS):\n",
    "    ds_label = DATASET_LABELS[ds]\n",
    "    \n",
    "    for i, (idx, label) in enumerate(IDX_DICT[ds].items()):\n",
    "        sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Train loss', color=COLOURS[i], label=label, ax=axs[d])\n",
    "    \n",
    "    axs[d].set(title=f\"{ds_label}\", ylabel='Training loss', ylim=(0., 2.))\n",
    "\n",
    "plt.subplots_adjust(hspace=0.6, wspace=0.25)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc5f944d",
   "metadata": {},
   "source": [
    "## Analysing discriminative downstream tasks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7150104",
   "metadata": {},
   "source": [
    "### Test accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "06cbbe06",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88560 characters]",
      "text/plain": [
       "<Figure size 1400x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(1, 3, figsize=(14, 3))\n",
    "\n",
    "for d, ds in enumerate(DATASETS):\n",
    "    ds_label = DATASET_LABELS[ds]\n",
    "    \n",
    "    for i, (idx, label) in enumerate(IDX_DICT[ds].items()):\n",
    "        sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Dstk: Test accuracy', color=COLOURS[i], label=label, ax=axs[d])\n",
    "    \n",
    "    axs[d].set(title=f\"{ds_label}\", ylabel='$R^2$ score')\n",
    "\n",
    "plt.subplots_adjust(hspace=0.6, wspace=0.25)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2921f31",
   "metadata": {},
   "source": [
    "### Generalisation gap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "92eb91d3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 103020 characters]",
      "text/plain": [
       "<Figure size 1400x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(1, 3, figsize=(14, 3))\n",
    "\n",
    "for d, ds in enumerate(DATASETS):\n",
    "    ds_label = DATASET_LABELS[ds]\n",
    "    \n",
    "    for i, (idx, label) in enumerate(IDX_DICT[ds].items()):\n",
    "        sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Dstk: GGap', color=COLOURS[i], label=label, ax=axs[d])\n",
    "    \n",
    "    axs[d].set(title=f\"{ds_label}\", ylabel='$R^2$ score')\n",
    "\n",
    "plt.subplots_adjust(hspace=0.6, wspace=0.25)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "813724e9",
   "metadata": {},
   "source": [
    "### Combined"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b485c32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 179200 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",
    "    for i, (idx, label) in enumerate(IDX_DICT[ds].items()):\n",
    "        plot_label = 'Zoo trainset' if label == DATASET_LABELS[ds] else label\n",
    "        sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Dstk: Test accuracy', color=COLOURS[i], label=plot_label, ax=axs[0][d])\n",
    "        sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Dstk: GGap', color=COLOURS[i], label=plot_label, ax=axs[1][d])\n",
    "    \n",
    "    axs[0][d].set(title=f\"Predicting test accuracy: {ds_label}\", ylabel='$R^2$ score')\n",
    "    axs[0][d].legend().remove()\n",
    "    \n",
    "    axs[1][d].set(title=f\"Predicting generalization gap: {ds_label}\", ylabel='$R^2$ score')\n",
    "    axs[1][d].legend().remove()\n",
    "\n",
    "plt.legend(title='Query set', ncols=3, loc=(-1.40, -0.60))\n",
    "plt.subplots_adjust(hspace=0.4, wspace=0.25)\n",
    "plt.savefig('figures/queryset_discriminative.png', dpi=300, bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8be350ce",
   "metadata": {},
   "source": [
    "## Analysing reconstructive downstream tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bce0ae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_reconstruction_acc(hr, x, p, i, dataloader):\n",
    "    z, y, zp = hr.forward(x, p)\n",
    "    checkpoint = tokens_to_checkpoint(y[i], p[i], reference_checkpoint)\n",
    "    model = CNN3(\n",
    "        channels_in=reference_params['model::channels_in'],\n",
    "        nlin=reference_params['model::nlin'],\n",
    "        dropout=reference_params['model::dropout']\n",
    "    )\n",
    "    model.load_state_dict(checkpoint)\n",
    "    \n",
    "    return test_model(model, dataloader, device='cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a40cfaba",
   "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>Model zoo</th>\n",
       "      <th>Model index</th>\n",
       "      <th>Hyper-representation</th>\n",
       "      <th>Hyper-representation model</th>\n",
       "      <th>Original accuracy</th>\n",
       "      <th>Reconstruction accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>0</td>\n",
       "      <td>5c618_00000</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>0.743623</td>\n",
       "      <td>0.882529</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>1</td>\n",
       "      <td>5c618_00000</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>0.828596</td>\n",
       "      <td>0.896358</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>2</td>\n",
       "      <td>5c618_00000</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>0.856215</td>\n",
       "      <td>0.898510</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>3</td>\n",
       "      <td>5c618_00000</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>0.873848</td>\n",
       "      <td>0.899969</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>4</td>\n",
       "      <td>5c618_00000</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>0.887792</td>\n",
       "      <td>0.901659</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Model zoo  Model index Hyper-representation Hyper-representation model  \\\n",
       "0      SVHN            0          5c618_00000                       SVHN   \n",
       "1      SVHN            1          5c618_00000                       SVHN   \n",
       "2      SVHN            2          5c618_00000                       SVHN   \n",
       "3      SVHN            3          5c618_00000                       SVHN   \n",
       "4      SVHN            4          5c618_00000                       SVHN   \n",
       "\n",
       "   Original accuracy  Reconstruction accuracy  \n",
       "0           0.743623                 0.882529  \n",
       "1           0.828596                 0.896358  \n",
       "2           0.856215                 0.898510  \n",
       "3           0.873848                 0.899969  \n",
       "4           0.887792                 0.901659  "
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DF_ZOO_ACC_PATH = EXPERIMENT_PATH / 'df_zoo_acc.csv'\n",
    "DF_Z_PATH = EXPERIMENT_PATH / 'df_z_dict.csv'\n",
    "DF_X_PATH = EXPERIMENT_PATH / 'df_x_dict.csv'\n",
    "\n",
    "X_DIM_1, X_DIM_2 = next(iter(ZOO_DL['svhn']))[0].shape[1:3]\n",
    "X_COLS = [f'Dimension {i}' for i in range(X_DIM_1 * X_DIM_2)]\n",
    "\n",
    "Z_DIM_1 = X_DIM_1\n",
    "Z_DIM_2 = 64\n",
    "Z_COLS = [f'Dimension {i}' for i in range(Z_DIM_1 * Z_DIM_2)]\n",
    "\n",
    "if not DF_ZOO_ACC_PATH.is_file() or RERUN:\n",
    "    # Create dataframe\n",
    "    df_zoo_acc = pd.DataFrame(columns=['Model zoo', 'Model index', 'Hyper-representation', 'Hyper-representation model', 'Original accuracy', 'Reconstruction accuracy'])\n",
    "    df_x = pd.DataFrame(columns=['Model zoo', 'Model index', 'Hyper-representation', 'Hyper-representation model'] + X_COLS)\n",
    "    df_z = pd.DataFrame(columns=['Model zoo', 'Model index', 'Hyper-representation', 'Hyper-representation model'] + Z_COLS)\n",
    "\n",
    "    # Run experiments\n",
    "    for ds in DATASETS:\n",
    "        for bdx, data in enumerate(tqdm(ZOO_DL[ds])):\n",
    "            x, m, p, _ = data\n",
    "\n",
    "            for hr_idx, hr_label in IDX_DICT[ds].items():\n",
    "                z, y, _ = hyper_reps[hr_idx].forward(x, p)\n",
    "\n",
    "                for idx in range(x.shape[0]):\n",
    "                    # Original accuracy\n",
    "                    orig_checkpoint = tokens_to_checkpoint(x[idx], p[idx], reference_checkpoint)\n",
    "                    orig_model = CNN3(\n",
    "                        channels_in=reference_params['model::channels_in'],\n",
    "                        nlin=reference_params['model::nlin'],\n",
    "                        dropout=reference_params['model::dropout']\n",
    "                    )\n",
    "                    orig_model.load_state_dict(orig_checkpoint)\n",
    "                    orig_acc = test_model(orig_model, TESTSET_DL[ds])\n",
    "\n",
    "                    # Reconstruction accuracy\n",
    "                    recon_checkpoint = tokens_to_checkpoint(y[idx], p[idx], reference_checkpoint)\n",
    "                    recon_model = CNN3(\n",
    "                        channels_in=reference_params['model::channels_in'],\n",
    "                        nlin=reference_params['model::nlin'],\n",
    "                        dropout=reference_params['model::dropout']\n",
    "                    )\n",
    "                    recon_model.load_state_dict(recon_checkpoint)\n",
    "                    recon_acc = test_model(recon_model, TESTSET_DL[ds])\n",
    "\n",
    "                    # Save results to DataFrame\n",
    "                    mdx = bdx * x.shape[0] + idx\n",
    "                    df_zoo_acc.loc[len(df_zoo_acc)] = (\n",
    "                        DATASET_LABELS[ds],\n",
    "                        mdx,\n",
    "                        hr_idx,\n",
    "                        hr_label,\n",
    "                        orig_acc,\n",
    "                        recon_acc\n",
    "                    )\n",
    "                    df_x.loc[len(df_x)] = [DATASET_LABELS[ds], mdx, hr_idx, hr_label] + x[idx].flatten().tolist()\n",
    "                    df_z.loc[len(df_z)] = [DATASET_LABELS[ds], mdx, hr_idx, hr_label] + z[idx].flatten().tolist()\n",
    "\n",
    "    # Save results to disk\n",
    "    df_zoo_acc.to_csv(DF_ZOO_ACC_PATH)\n",
    "    df_x.to_csv(DF_X_PATH)\n",
    "    df_z.to_csv(DF_Z_PATH)\n",
    "    \n",
    "df_zoo_acc = pd.read_csv(DF_ZOO_ACC_PATH, index_col=0)\n",
    "df_x = pd.read_csv(DF_X_PATH, index_col=0)\n",
    "df_z = pd.read_csv(DF_Z_PATH, index_col=0)\n",
    "df_zoo_acc.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "932d993d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 178136 characters]",
      "text/plain": [
       "<Figure size 1600x1200 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(len(IDX_DICT['svhn']), 3, figsize=(16, 4*len(IDX_DICT['svhn'])))\n",
    "\n",
    "for d, ds in enumerate(DATASETS):\n",
    "    for i, (idx, label) in enumerate(IDX_DICT[ds].items()):\n",
    "        sns.histplot(df_zoo_acc[df_zoo_acc['Hyper-representation'] == idx], x='Original accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[0], alpha=0.5, label='Original', ax=axs[i][d])\n",
    "        sns.histplot(df_zoo_acc[df_zoo_acc['Hyper-representation'] == idx], x='Reconstruction accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[1], alpha=0.5, label='Reconstruction', ax=axs[i][d])\n",
    "        axs[i][d].set(title=f\"Model zoo {DATASET_LABELS[ds]}\\nQuery set {label}\", xlabel='Test accuracy', ylabel='Probability', ylim=(0, 1))\n",
    "        axs[i][d].legend()\n",
    "    \n",
    "plt.subplots_adjust(hspace=0.5, wspace=0.25)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "536d98c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 101120 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",
    "    for i, (idx, label) in enumerate(IDX_DICT[ds].items()):\n",
    "        if label == DATASET_LABELS[ds]:\n",
    "            continue\n",
    "        sns.histplot(df_zoo_acc[df_zoo_acc['Hyper-representation'] == idx], x='Original accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[0], alpha=0.5, label='Original', ax=axs[i-1][d])\n",
    "        sns.histplot(df_zoo_acc[df_zoo_acc['Hyper-representation'] == idx], x='Reconstruction accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[1], alpha=0.5, label='Reconstruction', ax=axs[i-1][d])\n",
    "        axs[i-1][d].set(title=f\"Model zoo {DATASET_LABELS[ds]}\\nQuery set {label}\", xlabel='Test accuracy', ylabel='Probability', ylim=(0, 1))\n",
    "        axs[i-1][d].legend().remove()\n",
    "\n",
    "# Get the legend handles and labels\n",
    "handles, labels = plt.gca().get_legend_handles_labels()\n",
    "\n",
    "# Remove extra handle and label\n",
    "handles = handles[:2]\n",
    "labels = labels[:2]\n",
    "\n",
    "plt.legend(handles, labels, loc=(-1.25, -0.50), ncols=2)\n",
    "plt.subplots_adjust(hspace=0.6, wspace=0.25)\n",
    "plt.savefig('figures/queryset_reconstructive.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
}
