{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1f9b9a54",
   "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": "efa4494c",
   "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 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": "c65f70fc",
   "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": "9694f2d3",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f932efb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENT_PATH = Path.cwd() / 'results' / 'ablation_compute'\n",
    "\n",
    "if not EXPERIMENT_PATH.is_dir():\n",
    "    os.makedirs(EXPERIMENT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "100ba77a",
   "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']\n",
    "\n",
    "HR_ORDER = ['Contrastive + Structure (Baseline)', 'Contrastive + Structure + Behavior']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3326a919",
   "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') for ds in DATASETS}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2523b4f",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "### Loading HR models meta-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "784426aa",
   "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>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": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MAX_EPOCH = {\n",
    "    'Contrastive + Structure (Baseline)': 100,\n",
    "    'Contrastive + Structure + Behavior': 50\n",
    "}\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": "47c96b73",
   "metadata": {},
   "source": [
    "### Loading checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "10ce8007",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading hyper-reps\n",
    "def load_hr(trial_id, edx):\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, edx)\n",
    "\n",
    "# Load at epoch 50 if \\beta < 0.9, i.e. the behavioral loss is used\n",
    "hyper_reps = {idx: load_hr(idx, edx=50 if df_hr[df_hr['Trial ID'] == idx]['Beta'].iloc[0] < 0.9 else 100) for idx in df_hr['Trial ID'].unique()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6833f054",
   "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": "23e5a636",
   "metadata": {},
   "source": [
    "### Extracting key models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "550ebb40",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'svhn': {'50f57_00011': 'Contrastive + Structure (Baseline)',\n",
       "  '50f57_00008': 'Contrastive + Structure + Behavior'},\n",
       " 'cifar10': {'0d5ff_00011': 'Contrastive + Structure (Baseline)',\n",
       "  '0d5ff_00008': 'Contrastive + Structure + Behavior'},\n",
       " 'eurosat': {'d28cf_00011': 'Contrastive + Structure (Baseline)',\n",
       "  'd28cf_00008': 'Contrastive + Structure + Behavior'}}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "IDX_DICT = {ds:\n",
    "    {\n",
    "        df_hr[(df_hr['Model zoo'] == DATASET_LABELS[ds]) & (df_hr['Learning rate'] == 1e-4) & (df_hr['Gamma'] == 0.05) & (df_hr['Beta'] == 1.0) & (df_hr['Behaviour loss'] == 'l2')]['Trial ID'].iloc[0]: 'Contrastive + Structure (Baseline)',\n",
    "        df_hr[(df_hr['Model zoo'] == DATASET_LABELS[ds]) & (df_hr['Learning rate'] == 1e-5) & (df_hr['Gamma'] == 0.05) & (df_hr['Beta'] == 0.1) & (df_hr['Behaviour loss'] == 'l2')]['Trial ID'].iloc[0]: 'Contrastive + Structure + Behavior'\n",
    "    }\n",
    "    for ds in DATASETS\n",
    "}\n",
    "IDX_DICT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88b92b6a",
   "metadata": {},
   "source": [
    "## Analysing loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "98a7f4ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61716 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": "b172c65e",
   "metadata": {},
   "source": [
    "## Analysing reconstructive downstream tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f10f2d70",
   "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": 13,
   "id": "c013ca98",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9fc17454c21345f1a56c29910689735b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "71cec3dc9c3440c684aa3c179cf1608b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3316be556748419284325079ce96865c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/45 [00:00<?, ?it/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>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>50f57_00011</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.743623</td>\n",
       "      <td>0.195874</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>1</td>\n",
       "      <td>50f57_00011</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.828596</td>\n",
       "      <td>0.195913</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>2</td>\n",
       "      <td>50f57_00011</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.856215</td>\n",
       "      <td>0.295444</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>3</td>\n",
       "      <td>50f57_00011</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.873848</td>\n",
       "      <td>0.429779</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>4</td>\n",
       "      <td>50f57_00011</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.887792</td>\n",
       "      <td>0.581400</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Model zoo  Model index Hyper-representation  \\\n",
       "0      SVHN            0          50f57_00011   \n",
       "1      SVHN            1          50f57_00011   \n",
       "2      SVHN            2          50f57_00011   \n",
       "3      SVHN            3          50f57_00011   \n",
       "4      SVHN            4          50f57_00011   \n",
       "\n",
       "           Hyper-representation model  Original accuracy  \\\n",
       "0  Contrastive + Structure (Baseline)           0.743623   \n",
       "1  Contrastive + Structure (Baseline)           0.828596   \n",
       "2  Contrastive + Structure (Baseline)           0.856215   \n",
       "3  Contrastive + Structure (Baseline)           0.873848   \n",
       "4  Contrastive + Structure (Baseline)           0.887792   \n",
       "\n",
       "   Reconstruction accuracy  \n",
       "0                 0.195874  \n",
       "1                 0.195913  \n",
       "2                 0.295444  \n",
       "3                 0.429779  \n",
       "4                 0.581400  "
      ]
     },
     "execution_count": 13,
     "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": 15,
   "id": "45021443",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 96680 characters]",
      "text/plain": [
       "<Figure size 1400x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(len(IDX_DICT['svhn']), 3, figsize=(14, 3*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\"{DATASET_LABELS[ds]}\\n{label}\", xlabel='Test accuracy', ylabel='Probability', ylim=(0, 1))\n",
    "        axs[i][d].legend().remove()\n",
    "    \n",
    "plt.legend(loc=(-1.25, -0.50), ncols=2)\n",
    "plt.subplots_adjust(hspace=0.6, wspace=0.25)\n",
    "plt.savefig('figures/ablation_compute_reconstruction.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
}
