{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a537cad2",
   "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": "0930b3f7",
   "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": "19d703eb",
   "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": "43caef8a",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1a6b258b",
   "metadata": {},
   "outputs": [],
   "source": [
    "RERUN = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "05c1a048",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENT_PATH = Path.cwd() / 'results' / 'hr_analysis_eurosat'\n",
    "\n",
    "if not EXPERIMENT_PATH.is_dir():\n",
    "    os.makedirs(EXPERIMENT_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d0820ef3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test set corresponding to the data used to train the model zoo\n",
    "TESTSET_DS = torch.load('{YOUR_WORKING_PATH}/results/data/eurosat_train_zoo_relu/dataset.pt')['testset']\n",
    "TESTSET_DL = torch.utils.data.DataLoader(TESTSET_DS, batch_size=64, shuffle=False)\n",
    "\n",
    "# Model zoo\n",
    "ZOO_PATH = Path('{YOUR_WORKING_PATH}/results/zoos/eurosat_train_zoo_relu')\n",
    "ZOO_DS = torch.load(ZOO_PATH / 'dataset_test.pt')\n",
    "ZOO_DL = torch.utils.data.DataLoader(ZOO_DS, batch_size=16, shuffle=False)\n",
    "\n",
    "ZOO_DS_TRAIN = torch.load(ZOO_PATH / 'dataset_train.pt')\n",
    "ZOO_DL_TRAIN = torch.utils.data.DataLoader(ZOO_DS_TRAIN, batch_size=16, shuffle=False)\n",
    "\n",
    "# Threshold used to define models that \"perform well\"; depends on the zoo\n",
    "MODEL_ACC_THR = 0.7\n",
    "\n",
    "# Path to the hyper-representation models\n",
    "HYPERREP_PATH = Path('{YOUR_WORKING_PATH}/results/hyperrepresentations/tune/eurosat_train_hyperrep')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5491e24",
   "metadata": {},
   "source": [
    "## Loading data\n",
    "\n",
    "### Loading HR models meta-data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "51e43a53",
   "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>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>Epoch</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>d28cf_00025</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>64</td>\n",
       "      <td>distillation</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.993533</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.993533</td>\n",
       "      <td>0.985235</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.985235</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.985235</td>\n",
       "      <td>0.883079</td>\n",
       "      <td>0.321712</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>d28cf_00025</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>64</td>\n",
       "      <td>distillation</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.012841</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.012841</td>\n",
       "      <td>0.027832</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.027832</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.027832</td>\n",
       "      <td>0.895981</td>\n",
       "      <td>0.315194</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>d28cf_00025</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>64</td>\n",
       "      <td>distillation</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2</td>\n",
       "      <td>...</td>\n",
       "      <td>0.011913</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.011913</td>\n",
       "      <td>0.016323</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.016323</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.016323</td>\n",
       "      <td>0.896555</td>\n",
       "      <td>0.314326</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>d28cf_00025</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>64</td>\n",
       "      <td>distillation</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>3</td>\n",
       "      <td>...</td>\n",
       "      <td>0.015708</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.015708</td>\n",
       "      <td>0.024025</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.024025</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.024025</td>\n",
       "      <td>0.896595</td>\n",
       "      <td>0.313355</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>d28cf_00025</td>\n",
       "      <td>0.0001</td>\n",
       "      <td>64</td>\n",
       "      <td>distillation</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>4</td>\n",
       "      <td>...</td>\n",
       "      <td>0.009568</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.009568</td>\n",
       "      <td>0.010813</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.010813</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.010813</td>\n",
       "      <td>0.896910</td>\n",
       "      <td>0.313311</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 22 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      Trial ID  Learning rate  Latent dimension Behaviour loss Query set  \\\n",
       "0  d28cf_00025         0.0001                64   distillation      data   \n",
       "1  d28cf_00025         0.0001                64   distillation      data   \n",
       "2  d28cf_00025         0.0001                64   distillation      data   \n",
       "3  d28cf_00025         0.0001                64   distillation      data   \n",
       "4  d28cf_00025         0.0001                64   distillation      data   \n",
       "\n",
       "  Query dataset  # of queries  Gamma  Beta  Epoch  ...  Train loss (recon)  \\\n",
       "0       EuroSAT           256    0.0   0.0      0  ...            0.993533   \n",
       "1       EuroSAT           256    0.0   0.0      1  ...            0.012841   \n",
       "2       EuroSAT           256    0.0   0.0      2  ...            0.011913   \n",
       "3       EuroSAT           256    0.0   0.0      3  ...            0.015708   \n",
       "4       EuroSAT           256    0.0   0.0      4  ...            0.009568   \n",
       "\n",
       "   Train loss (structure)  Train loss (behaviour)  Test loss  \\\n",
       "0                     0.0                0.993533   0.985235   \n",
       "1                     0.0                0.012841   0.027832   \n",
       "2                     0.0                0.011913   0.016323   \n",
       "3                     0.0                0.015708   0.024025   \n",
       "4                     0.0                0.009568   0.010813   \n",
       "\n",
       "   Test loss (contrast)  Test loss (recon)  Test loss (structure)  \\\n",
       "0                   0.0           0.985235                    0.0   \n",
       "1                   0.0           0.027832                    0.0   \n",
       "2                   0.0           0.016323                    0.0   \n",
       "3                   0.0           0.024025                    0.0   \n",
       "4                   0.0           0.010813                    0.0   \n",
       "\n",
       "   Test loss (behaviour)  Dstk: Test accuracy  Dstk: GGap  \n",
       "0               0.985235             0.883079    0.321712  \n",
       "1               0.027832             0.895981    0.315194  \n",
       "2               0.016323             0.896555    0.314326  \n",
       "3               0.024025             0.896595    0.313355  \n",
       "4               0.010813             0.896910    0.313311  \n",
       "\n",
       "[5 rows x 22 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MAX_EPOCH = 100\n",
    "\n",
    "df_hr = pd.DataFrame(columns=[\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 path in os.listdir(HYPERREP_PATH):\n",
    "    if not re.match('AE_trainable', path):\n",
    "        continue\n",
    "    path = HYPERREP_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",
    "            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": "markdown",
   "id": "0bb4e14b",
   "metadata": {},
   "source": [
    "### Loading checkpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "77b86358",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading hyper-reps\n",
    "def load_hr(trial_id):\n",
    "    for path in os.listdir(HYPERREP_PATH):\n",
    "        if trial_id in path:\n",
    "            model_path = HYPERREP_PATH / path\n",
    "    return load_hyperrep_edx(model_path, MAX_EPOCH)\n",
    "\n",
    "hyper_reps = {idx: load_hr(idx) for idx in df_hr['Trial ID'].unique()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aebcd6ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading reference config\n",
    "trial_id = df_hr.iloc[0]['Trial ID']\n",
    "for path in os.listdir(HYPERREP_PATH):\n",
    "    if trial_id in path:\n",
    "        model_path = HYPERREP_PATH / path\n",
    "        with open(HYPERREP_PATH / 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": "abb8a35a",
   "metadata": {},
   "source": [
    "## Analysing the training loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "048df463",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 247652 characters]",
      "text/plain": [
       "<Figure size 1200x1400 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axs = plt.subplots(3, 2, sharey=True, figsize=(12, 14))\n",
    "\n",
    "sns.lineplot(df_hr[(df_hr['Behaviour loss'] == 'l2') & (df_hr['Learning rate'] == 1e-5)], x='Epoch', y='Train loss', hue='Beta', style='Gamma', palette=COLOURS[:3], ax=axs[0][0]).set(\n",
    "    title='L2 Loss, lr = 1e-5'\n",
    ")\n",
    "\n",
    "sns.lineplot(df_hr[(df_hr['Behaviour loss'] == 'l2') & (df_hr['Learning rate'] == 1e-4)], x='Epoch', y='Train loss', hue='Beta', style='Gamma', palette=COLOURS[:3], ax=axs[0][1]).set(\n",
    "    title='L2 Loss, lr = 1e-4', ylim=(0., 1.5)\n",
    ")\n",
    "\n",
    "sns.lineplot(df_hr[(df_hr['Behaviour loss'] == 'distillation') & (df_hr['Learning rate'] == 1e-5)], x='Epoch', y='Train loss', hue='Beta', style='Gamma', palette=COLOURS[:3], ax=axs[1][0]).set(\n",
    "    title='Distillation Loss, lr = 1e-5'\n",
    ")\n",
    "\n",
    "sns.lineplot(df_hr[(df_hr['Behaviour loss'] == 'distillation') & (df_hr['Learning rate'] == 1e-4)], x='Epoch', y='Train loss', hue='Beta', style='Gamma', palette=COLOURS[:3], ax=axs[1][1]).set(\n",
    "    title='Distillation Loss, lr = 1e-4', ylim=(0., 1.5)\n",
    ")\n",
    "\n",
    "sns.lineplot(df_hr[(df_hr['Behaviour loss'] == 'cross_entropy') & (df_hr['Learning rate'] == 1e-5)], x='Epoch', y='Train loss', hue='Beta', style='Gamma', palette=COLOURS[:3], ax=axs[2][0]).set(\n",
    "    title='Cross-entropy Loss, lr = 1e-5'\n",
    ")\n",
    "\n",
    "sns.lineplot(df_hr[(df_hr['Behaviour loss'] == 'cross_entropy') & (df_hr['Learning rate'] == 1e-4)], x='Epoch', y='Train loss', hue='Beta', style='Gamma', palette=COLOURS[:3], ax=axs[2][1]).set(\n",
    "    title='Cross-entropy Loss, lr = 1e-4', ylim=(0., 1.5)\n",
    ")\n",
    "\n",
    "for ax in axs.flatten():\n",
    "    ax.legend(ncol=2)\n",
    "\n",
    "plt.suptitle('Training loss')\n",
    "plt.subplots_adjust(hspace=0.40)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a771087",
   "metadata": {},
   "source": [
    "### Extracting key models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "30d9da69",
   "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>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>Epoch</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>302</th>\n",
       "      <td>d28cf_00004</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.00</td>\n",
       "      <td>1.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.517407</td>\n",
       "      <td>0.517407</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.748472</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.748472</td>\n",
       "      <td>0.748472</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.935311</td>\n",
       "      <td>0.260450</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>403</th>\n",
       "      <td>d28cf_00010</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>1.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.893007</td>\n",
       "      <td>0.893007</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.233713</td>\n",
       "      <td>0.319473</td>\n",
       "      <td>1.281831</td>\n",
       "      <td>1.281831</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.949818</td>\n",
       "      <td>0.408585</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>807</th>\n",
       "      <td>d28cf_00009</td>\n",
       "      <td>0.00010</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>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.117895</td>\n",
       "      <td>0.617386</td>\n",
       "      <td>0.062396</td>\n",
       "      <td>0.141606</td>\n",
       "      <td>0.126279</td>\n",
       "      <td>0.142413</td>\n",
       "      <td>0.771686</td>\n",
       "      <td>0.072494</td>\n",
       "      <td>0.983024</td>\n",
       "      <td>0.623544</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>908</th>\n",
       "      <td>d28cf_00003</td>\n",
       "      <td>0.00010</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.1</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.115284</td>\n",
       "      <td>0.662928</td>\n",
       "      <td>0.054434</td>\n",
       "      <td>0.141328</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.141328</td>\n",
       "      <td>0.830902</td>\n",
       "      <td>0.064709</td>\n",
       "      <td>0.984058</td>\n",
       "      <td>0.616272</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1615</th>\n",
       "      <td>d28cf_00001</td>\n",
       "      <td>0.00010</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.534961</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.534961</td>\n",
       "      <td>0.533183</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.533183</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.533183</td>\n",
       "      <td>0.943571</td>\n",
       "      <td>0.449419</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1716</th>\n",
       "      <td>d28cf_00007</td>\n",
       "      <td>0.00010</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.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.478548</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.478548</td>\n",
       "      <td>0.457390</td>\n",
       "      <td>0.089005</td>\n",
       "      <td>0.476778</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.476778</td>\n",
       "      <td>0.963891</td>\n",
       "      <td>0.530346</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2120</th>\n",
       "      <td>d28cf_00008</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>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.149865</td>\n",
       "      <td>0.887436</td>\n",
       "      <td>0.067913</td>\n",
       "      <td>0.174474</td>\n",
       "      <td>0.284862</td>\n",
       "      <td>0.168664</td>\n",
       "      <td>1.076722</td>\n",
       "      <td>0.067769</td>\n",
       "      <td>0.969290</td>\n",
       "      <td>0.528638</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2221</th>\n",
       "      <td>d28cf_00002</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.00</td>\n",
       "      <td>0.1</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.153765</td>\n",
       "      <td>0.715594</td>\n",
       "      <td>0.091339</td>\n",
       "      <td>0.169478</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.169478</td>\n",
       "      <td>0.868098</td>\n",
       "      <td>0.091853</td>\n",
       "      <td>0.966495</td>\n",
       "      <td>0.484092</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2827</th>\n",
       "      <td>d28cf_00011</td>\n",
       "      <td>0.00010</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>1.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.569155</td>\n",
       "      <td>0.569155</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.866255</td>\n",
       "      <td>0.137753</td>\n",
       "      <td>0.904598</td>\n",
       "      <td>0.904598</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.957421</td>\n",
       "      <td>0.465414</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3029</th>\n",
       "      <td>d28cf_00005</td>\n",
       "      <td>0.00010</td>\n",
       "      <td>64</td>\n",
       "      <td>l2</td>\n",
       "      <td>data</td>\n",
       "      <td>EuroSAT</td>\n",
       "      <td>256</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.716043</td>\n",
       "      <td>0.716043</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.129177</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>1.129177</td>\n",
       "      <td>1.129177</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.953867</td>\n",
       "      <td>0.471478</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3332</th>\n",
       "      <td>d28cf_00006</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.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.477857</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.477857</td>\n",
       "      <td>0.464852</td>\n",
       "      <td>0.256462</td>\n",
       "      <td>0.475820</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.475820</td>\n",
       "      <td>0.949878</td>\n",
       "      <td>0.482086</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3433</th>\n",
       "      <td>d28cf_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.00</td>\n",
       "      <td>0.0</td>\n",
       "      <td>100</td>\n",
       "      <td>...</td>\n",
       "      <td>0.474465</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.474465</td>\n",
       "      <td>0.472818</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.472818</td>\n",
       "      <td>0.000000</td>\n",
       "      <td>0.472818</td>\n",
       "      <td>0.901377</td>\n",
       "      <td>0.296046</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>12 rows × 22 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         Trial ID  Learning rate  Latent dimension Behaviour loss Query set  \\\n",
       "302   d28cf_00004        0.00001                64             l2      data   \n",
       "403   d28cf_00010        0.00001                64             l2      data   \n",
       "807   d28cf_00009        0.00010                64             l2      data   \n",
       "908   d28cf_00003        0.00010                64             l2      data   \n",
       "1615  d28cf_00001        0.00010                64             l2      data   \n",
       "1716  d28cf_00007        0.00010                64             l2      data   \n",
       "2120  d28cf_00008        0.00001                64             l2      data   \n",
       "2221  d28cf_00002        0.00001                64             l2      data   \n",
       "2827  d28cf_00011        0.00010                64             l2      data   \n",
       "3029  d28cf_00005        0.00010                64             l2      data   \n",
       "3332  d28cf_00006        0.00001                64             l2      data   \n",
       "3433  d28cf_00000        0.00001                64             l2      data   \n",
       "\n",
       "     Query dataset  # of queries  Gamma  Beta  Epoch  ...  Train loss (recon)  \\\n",
       "302        EuroSAT           256   0.00   1.0    100  ...            0.517407   \n",
       "403        EuroSAT           256   0.05   1.0    100  ...            0.893007   \n",
       "807        EuroSAT           256   0.05   0.1    100  ...            0.117895   \n",
       "908        EuroSAT           256   0.00   0.1    100  ...            0.115284   \n",
       "1615       EuroSAT           256   0.00   0.0    100  ...            0.534961   \n",
       "1716       EuroSAT           256   0.05   0.0    100  ...            0.478548   \n",
       "2120       EuroSAT           256   0.05   0.1    100  ...            0.149865   \n",
       "2221       EuroSAT           256   0.00   0.1    100  ...            0.153765   \n",
       "2827       EuroSAT           256   0.05   1.0    100  ...            0.569155   \n",
       "3029       EuroSAT           256   0.00   1.0    100  ...            0.716043   \n",
       "3332       EuroSAT           256   0.05   0.0    100  ...            0.477857   \n",
       "3433       EuroSAT           256   0.00   0.0    100  ...            0.474465   \n",
       "\n",
       "      Train loss (structure)  Train loss (behaviour)  Test loss  \\\n",
       "302                 0.517407                0.000000   0.748472   \n",
       "403                 0.893007                0.000000   1.233713   \n",
       "807                 0.617386                0.062396   0.141606   \n",
       "908                 0.662928                0.054434   0.141328   \n",
       "1615                0.000000                0.534961   0.533183   \n",
       "1716                0.000000                0.478548   0.457390   \n",
       "2120                0.887436                0.067913   0.174474   \n",
       "2221                0.715594                0.091339   0.169478   \n",
       "2827                0.569155                0.000000   0.866255   \n",
       "3029                0.716043                0.000000   1.129177   \n",
       "3332                0.000000                0.477857   0.464852   \n",
       "3433                0.000000                0.474465   0.472818   \n",
       "\n",
       "      Test loss (contrast)  Test loss (recon)  Test loss (structure)  \\\n",
       "302               0.000000           0.748472               0.748472   \n",
       "403               0.319473           1.281831               1.281831   \n",
       "807               0.126279           0.142413               0.771686   \n",
       "908               0.000000           0.141328               0.830902   \n",
       "1615              0.000000           0.533183               0.000000   \n",
       "1716              0.089005           0.476778               0.000000   \n",
       "2120              0.284862           0.168664               1.076722   \n",
       "2221              0.000000           0.169478               0.868098   \n",
       "2827              0.137753           0.904598               0.904598   \n",
       "3029              0.000000           1.129177               1.129177   \n",
       "3332              0.256462           0.475820               0.000000   \n",
       "3433              0.000000           0.472818               0.000000   \n",
       "\n",
       "      Test loss (behaviour)  Dstk: Test accuracy  Dstk: GGap  \n",
       "302                0.000000             0.935311    0.260450  \n",
       "403                0.000000             0.949818    0.408585  \n",
       "807                0.072494             0.983024    0.623544  \n",
       "908                0.064709             0.984058    0.616272  \n",
       "1615               0.533183             0.943571    0.449419  \n",
       "1716               0.476778             0.963891    0.530346  \n",
       "2120               0.067769             0.969290    0.528638  \n",
       "2221               0.091853             0.966495    0.484092  \n",
       "2827               0.000000             0.957421    0.465414  \n",
       "3029               0.000000             0.953867    0.471478  \n",
       "3332               0.475820             0.949878    0.482086  \n",
       "3433               0.472818             0.901377    0.296046  \n",
       "\n",
       "[12 rows x 22 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_hr[(df_hr['Epoch'] == MAX_EPOCH) & (df_hr['Behaviour loss'] == 'l2')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4e1b214c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 73864 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "idx_dict = {\n",
    "    'd28cf_00011': 'Contrastive + Structure (Baseline)',\n",
    "    'd28cf_00000': 'Behavior',\n",
    "    'd28cf_00006': 'Contrastive + Behavior',\n",
    "    'd28cf_00002': 'Structure + Behavior',\n",
    "    'd28cf_00008': 'Contrastive + Structure + Behavior'\n",
    "}\n",
    "idx_ref_chk = list(idx_dict.keys())[0]\n",
    "f\"Reference checkpoint: {idx_ref_chk}\"\n",
    "\n",
    "\n",
    "_, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "    sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Train loss', color=COLOURS[i], label=label, ax=ax)\n",
    "    \n",
    "plt.legend()\n",
    "plt.suptitle('Training loss')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1de9f4f",
   "metadata": {},
   "source": [
    "## Analysing discriminative downstream tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b3dd3be8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 94144 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "    sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Dstk: Test accuracy', color=COLOURS[i], label=label, ax=ax)\n",
    "    \n",
    "ax.set(ylabel='$R^2$ score')\n",
    "    \n",
    "plt.legend()\n",
    "plt.suptitle('Downstream task: predicting test accuracy ($R^2$ score)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "aedca6ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 103836 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "    sns.lineplot(df_hr[df_hr['Trial ID'] == idx], x='Epoch', y='Dstk: GGap', color=COLOURS[i], label=label, ax=ax)\n",
    "    \n",
    "ax.set(ylabel='$R^2$ score')\n",
    "    \n",
    "plt.legend()\n",
    "plt.suptitle('Downstream task: predicting the generalization gap ($R^2$ score)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca39007d",
   "metadata": {},
   "source": [
    "## Analysing reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "10cf2249",
   "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": "markdown",
   "id": "d1df38e2",
   "metadata": {},
   "source": [
    "### Distribution of reconstructed accuracies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "099859c9",
   "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 index</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>0</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.240741</td>\n",
       "      <td>0.118272</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.273827</td>\n",
       "      <td>0.117037</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.241728</td>\n",
       "      <td>0.182963</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.235802</td>\n",
       "      <td>0.130617</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.795309</td>\n",
       "      <td>0.411852</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Model index          Hyper-representation model  Original accuracy  \\\n",
       "0            0  Contrastive + Structure (Baseline)           0.240741   \n",
       "1            1  Contrastive + Structure (Baseline)           0.273827   \n",
       "2            2  Contrastive + Structure (Baseline)           0.241728   \n",
       "3            3  Contrastive + Structure (Baseline)           0.235802   \n",
       "4            4  Contrastive + Structure (Baseline)           0.795309   \n",
       "\n",
       "   Reconstruction accuracy  \n",
       "0                 0.118272  \n",
       "1                 0.117037  \n",
       "2                 0.182963  \n",
       "3                 0.130617  \n",
       "4                 0.411852  "
      ]
     },
     "execution_count": 16,
     "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))[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 index', 'Hyper-representation model', 'Original accuracy', 'Reconstruction accuracy'])\n",
    "    df_x = pd.DataFrame(columns=['Model index', 'Hyper-representation model'] + X_COLS)\n",
    "    df_z = pd.DataFrame(columns=['Model index', 'Hyper-representation model'] + Z_COLS)\n",
    "\n",
    "    # Run experiments\n",
    "    for bdx, data in enumerate(tqdm(ZOO_DL)):\n",
    "        x, m, p, _ = data\n",
    "\n",
    "        for hr_idx, hr_label in idx_dict.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)\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)\n",
    "                \n",
    "                # Save results to DataFrame\n",
    "                mdx = bdx * x.shape[0] + idx\n",
    "                df_zoo_acc.loc[len(df_zoo_acc)] = (\n",
    "                    mdx,\n",
    "                    hr_label,\n",
    "                    orig_acc,\n",
    "                    recon_acc\n",
    "                )\n",
    "                df_x.loc[len(df_x)] = [mdx, hr_label] + x[idx].flatten().tolist()\n",
    "                df_z.loc[len(df_z)] = [mdx, 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": 17,
   "id": "15d6cf60",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 165112 characters]",
      "text/plain": [
       "<Figure size 800x2000 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(len(idx_dict), 1, figsize=(8, 4*len(idx_dict)))\n",
    "\n",
    "for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "    sns.histplot(df_zoo_acc[df_zoo_acc['Hyper-representation model'] == label], x='Original accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[0], alpha=0.5, label='Original', ax=axs[i])\n",
    "    sns.histplot(df_zoo_acc[df_zoo_acc['Hyper-representation model'] == label], x='Reconstruction accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[1], alpha=0.5, label='Reconstruction', ax=axs[i])\n",
    "    axs[i].set(title=label, xlabel='Test accuracy', ylabel='Probability', ylim=(0, 1))\n",
    "    axs[i].legend()\n",
    "    \n",
    "plt.subplots_adjust(hspace=0.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d0fafc8",
   "metadata": {},
   "source": [
    "### Pairwise $L^2$ distance between weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3d96534c",
   "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>Hyper-representation model</th>\n",
       "      <th>L2 (structure)</th>\n",
       "      <th>MSE (structure)</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>54.993130</td>\n",
       "      <td>0.112522</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>52.678055</td>\n",
       "      <td>0.103247</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>52.184288</td>\n",
       "      <td>0.101321</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>49.717228</td>\n",
       "      <td>0.091967</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>23.154524</td>\n",
       "      <td>0.019948</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           Hyper-representation model  L2 (structure)  MSE (structure)\n",
       "0  Contrastive + Structure (Baseline)       54.993130         0.112522\n",
       "1  Contrastive + Structure (Baseline)       52.678055         0.103247\n",
       "2  Contrastive + Structure (Baseline)       52.184288         0.101321\n",
       "3  Contrastive + Structure (Baseline)       49.717228         0.091967\n",
       "4  Contrastive + Structure (Baseline)       23.154524         0.019948"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DF_ZOO_L2WEIGHTS_PATH = EXPERIMENT_PATH / 'df_zoo_l2weights.csv'\n",
    "\n",
    "if not DF_ZOO_L2WEIGHTS_PATH.is_file() or RERUN:\n",
    "    # Create dataframe\n",
    "    df_zoo_l2weights = pd.DataFrame(columns=['Hyper-representation model', 'L2 (structure)', 'MSE (structure)'])\n",
    "\n",
    "    # Run experiments\n",
    "    for data in tqdm(ZOO_DL):\n",
    "        x, m, p, _ = data\n",
    "\n",
    "        for hr_idx, hr_label in idx_dict.items():\n",
    "            _, y, _ = hyper_reps[hr_idx].forward(x, p)\n",
    "            \n",
    "            for j in range(x.shape[0]):\n",
    "                df_zoo_l2weights.loc[len(df_zoo_l2weights)] = (\n",
    "                    hr_label,\n",
    "                    torch.sqrt(torch.sum((x[j] - y[j]) ** 2)).item(),\n",
    "                    torch.mean((x[j] - y[j]) ** 2).item()\n",
    "                )\n",
    "                \n",
    "    # Save results\n",
    "    df_zoo_l2weights.to_csv(DF_ZOO_L2WEIGHTS_PATH)\n",
    "    \n",
    "df_zoo_l2weights = pd.read_csv(DF_ZOO_L2WEIGHTS_PATH, index_col=0)\n",
    "df_zoo_l2weights.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5627d425",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 91284 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "    sns.kdeplot(df_zoo_l2weights[df_zoo_l2weights['Hyper-representation model'] == label], x='L2 (structure)', color=COLOURS[i], label=label, ax=ax)\n",
    "    \n",
    "plt.legend()\n",
    "plt.suptitle('Reconstruction structural error: $L^2$ distance between weights')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27a32819",
   "metadata": {},
   "source": [
    "### Pairwise $L^2$ distance between predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ca8fc380",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reconstruct_model(hr, x, p, i):\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 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "ce4e6e5a",
   "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>Hyper-representation model</th>\n",
       "      <th>L2 (behavior)</th>\n",
       "      <th>MSE (behavior)</th>\n",
       "      <th>Agreement</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>0.881178</td>\n",
       "      <td>0.119660</td>\n",
       "      <td>0.385926</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Behavior</td>\n",
       "      <td>0.641160</td>\n",
       "      <td>0.054532</td>\n",
       "      <td>0.383704</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Contrastive + Behavior</td>\n",
       "      <td>0.644318</td>\n",
       "      <td>0.054744</td>\n",
       "      <td>0.010123</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Structure + Behavior</td>\n",
       "      <td>0.617119</td>\n",
       "      <td>0.055881</td>\n",
       "      <td>0.378519</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Contrastive + Structure + Behavior</td>\n",
       "      <td>0.612932</td>\n",
       "      <td>0.054517</td>\n",
       "      <td>0.385185</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           Hyper-representation model  L2 (behavior)  MSE (behavior)  \\\n",
       "0  Contrastive + Structure (Baseline)       0.881178        0.119660   \n",
       "1                            Behavior       0.641160        0.054532   \n",
       "2              Contrastive + Behavior       0.644318        0.054744   \n",
       "3                Structure + Behavior       0.617119        0.055881   \n",
       "4  Contrastive + Structure + Behavior       0.612932        0.054517   \n",
       "\n",
       "   Agreement  \n",
       "0   0.385926  \n",
       "1   0.383704  \n",
       "2   0.010123  \n",
       "3   0.378519  \n",
       "4   0.385185  "
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DF_ZOO_L2PREDS_PATH = EXPERIMENT_PATH / 'df_zoo_l2preds.csv'\n",
    "\n",
    "if not DF_ZOO_L2PREDS_PATH.is_file() or RERUN:\n",
    "    # Create dataframe\n",
    "    df_zoo_l2preds = pd.DataFrame(columns=['Hyper-representation model', 'L2 (behavior)', 'MSE (behavior)', 'Agreement'])\n",
    "\n",
    "    # Run experiments\n",
    "    for data in tqdm(ZOO_DL):\n",
    "        x, m, p, _ = data\n",
    "\n",
    "        for i in range(x.shape[0]):\n",
    "            checkpoint = tokens_to_checkpoint(x[i], p[i], 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(checkpoint)\n",
    "\n",
    "            for hr_idx, hr_label in idx_dict.items():   \n",
    "                recon_model = reconstruct_model(hyper_reps[hr_idx], x, p, i)\n",
    "                \n",
    "                l2_dist, mse, agreement = compare_models_prediction(orig_model, recon_model, TESTSET_DL, n_batches=None)\n",
    "                \n",
    "                df_zoo_l2preds.loc[len(df_zoo_l2preds)] = (\n",
    "                    hr_label,\n",
    "                    l2_dist,\n",
    "                    mse,\n",
    "                    agreement\n",
    "                )\n",
    "                    \n",
    "                \n",
    "    # Save results\n",
    "    df_zoo_l2preds.to_csv(DF_ZOO_L2PREDS_PATH)\n",
    "    \n",
    "df_zoo_l2preds = pd.read_csv(DF_ZOO_L2PREDS_PATH, index_col=0)\n",
    "df_zoo_l2preds.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "6d42e37a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 122104 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(1 ,2, figsize=(12, 4))\n",
    "\n",
    "for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "    sns.kdeplot(df_zoo_l2preds[df_zoo_l2preds['Hyper-representation model'] == label], x='L2 (behavior)', color=COLOURS[i], label=label, ax=axs[0])\n",
    "    sns.kdeplot(df_zoo_l2preds[df_zoo_l2preds['Hyper-representation model'] == label], x='Agreement', color=COLOURS[i], label=label, ax=axs[1])\n",
    "    \n",
    "    \n",
    "axs[0].set(title='$L^2$ distance between predictions')\n",
    "axs[1].set(title='Model agreement')\n",
    "    \n",
    "plt.legend(loc=(-0.80, 0.60))\n",
    "plt.suptitle('Reconstruction performance')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99b8497e",
   "metadata": {},
   "source": [
    "## Analysing generative downstream tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "16951709",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_acc_from_z(hr, z, p, i, dataloader):\n",
    "    z = z.type(next(hr.parameters()).dtype)\n",
    "\n",
    "    y = hr.to(z.device).forward_decoder(z, p)\n",
    "    checkpoint = tokens_to_checkpoint(y[i], p[i], reference_checkpoint_cuda)\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "1a6c162a",
   "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>Original std</th>\n",
       "      <th>Sampled std</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Contrastive + Structure (Baseline)</th>\n",
       "      <td>0.294386</td>\n",
       "      <td>0.246665</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Behavior</th>\n",
       "      <td>0.187862</td>\n",
       "      <td>0.157544</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Contrastive + Behavior</th>\n",
       "      <td>0.095189</td>\n",
       "      <td>0.081251</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Structure + Behavior</th>\n",
       "      <td>0.069325</td>\n",
       "      <td>0.059606</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Contrastive + Structure + Behavior</th>\n",
       "      <td>0.057848</td>\n",
       "      <td>0.051122</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                    Original std  Sampled std\n",
       "Contrastive + Structure (Baseline)      0.294386     0.246665\n",
       "Behavior                                0.187862     0.157544\n",
       "Contrastive + Behavior                  0.095189     0.081251\n",
       "Structure + Behavior                    0.069325     0.059606\n",
       "Contrastive + Structure + Behavior      0.057848     0.051122"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Select only models with original accuracy above a threshold\n",
    "good_models_ids = df_zoo_acc[df_zoo_acc['Original accuracy'] >= MODEL_ACC_THR]['Model index'].unique().tolist()\n",
    "df_z_good_models = df_z[df_z['Model index'].isin(good_models_ids)]\n",
    "\n",
    "# Generate anchors and samples\n",
    "z_anchors = {hr_label: df_z_good_models[df_z_good_models['Hyper-representation model'] == hr_label][Z_COLS].values for hr_label in idx_dict.values()}\n",
    "z_samples = {k: generate_z_samples(z) for k, z in z_anchors.items()}\n",
    "\n",
    "# Print diagnostic statistics:\n",
    "df_diagnostic = pd.DataFrame(columns=['Original std', 'Sampled std'])\n",
    "for key in z_anchors.keys():\n",
    "    df_diagnostic.loc[key] = (\n",
    "        z_anchors[key].std(axis=0).mean(),\n",
    "        z_samples[key].std(axis=0).mean()\n",
    "    )\n",
    "display(df_diagnostic)\n",
    "\n",
    "# Reshape anchors and samples\n",
    "z_anchors = {k: z.reshape((-1, Z_DIM_1, Z_DIM_2)) for k, z in z_anchors.items()}\n",
    "z_samples = {k: z.reshape((-1, Z_DIM_1, Z_DIM_2)) for k, z in z_samples.items()}\n",
    "\n",
    "del df_z_good_models, df_diagnostic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "4e5a9533",
   "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>Hyper-representation model</th>\n",
       "      <th>Origin</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>Generated</td>\n",
       "      <td>0.365185</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>Generated</td>\n",
       "      <td>0.251605</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>Generated</td>\n",
       "      <td>0.212840</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>Generated</td>\n",
       "      <td>0.167160</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>Generated</td>\n",
       "      <td>0.420000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           Hyper-representation model     Origin  Accuracy\n",
       "0  Contrastive + Structure (Baseline)  Generated  0.365185\n",
       "1  Contrastive + Structure (Baseline)  Generated  0.251605\n",
       "2  Contrastive + Structure (Baseline)  Generated  0.212840\n",
       "3  Contrastive + Structure (Baseline)  Generated  0.167160\n",
       "4  Contrastive + Structure (Baseline)  Generated  0.420000"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DF_GENERATIVE_PATH = EXPERIMENT_PATH / 'df_generative.csv'\n",
    "\n",
    "if not DF_GENERATIVE_PATH.is_file() or RERUN:\n",
    "    # Create dataframe\n",
    "    df_generative = pd.DataFrame(columns=['Hyper-representation model', 'Origin', 'Accuracy'])\n",
    "    \n",
    "    p = next(iter(ZOO_DL))[2]\n",
    "\n",
    "    for hr_idx, hr_label in tqdm(idx_dict.items(), unit='hyperrepresentation model'):\n",
    "        z = torch.tensor(z_samples[hr_label]).to('cuda')\n",
    "        p_z = p[0].repeat(z.shape[0], 1, 1).to('cuda') \n",
    "        \n",
    "        for i in range(z.shape[0]):\n",
    "            recon_acc = compute_acc_from_z(hyper_reps[hr_idx], z, p_z, i, TESTSET_DL)\n",
    "            df_generative.loc[len(df_generative)] = (\n",
    "                hr_label,\n",
    "                'Generated',\n",
    "                recon_acc\n",
    "            )\n",
    "                \n",
    "    # Save results\n",
    "    df_generative.to_csv(DF_GENERATIVE_PATH)\n",
    "    \n",
    "df_generative = pd.read_csv(DF_GENERATIVE_PATH, index_col=0)\n",
    "df_generative.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6828e65b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 188764 characters]",
      "text/plain": [
       "<Figure size 800x2000 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "n_figures = len(df_generative['Hyper-representation model'].unique())\n",
    "\n",
    "fig, axs = plt.subplots(n_figures, figsize=(8, 4*n_figures))\n",
    "\n",
    "for i, (idx, label) in enumerate(idx_dict.items()):\n",
    "    sns.histplot(df_zoo_acc[(df_zoo_acc['Hyper-representation model'] == label) & (df_zoo_acc['Original accuracy'] >= MODEL_ACC_THR)], x='Original accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[0], alpha=0.5, label='Original', ax=axs[i])\n",
    "    sns.histplot(df_zoo_acc[(df_zoo_acc['Hyper-representation model'] == label) & (df_zoo_acc['Original accuracy'] >= MODEL_ACC_THR)], x='Reconstruction accuracy', stat='probability', binwidth=0.05, binrange=(0., 1.), kde=True, color=COLOURS[1], alpha=0.5, label='Reconstruction', ax=axs[i])\n",
    "    sns.histplot(\n",
    "        df_generative[(df_generative['Hyper-representation model'] == label) & (df_generative['Origin'] == 'Generated')],\n",
    "        x='Accuracy',\n",
    "        stat='probability',\n",
    "        binwidth=0.05,\n",
    "        binrange=(0., 1.),\n",
    "        kde=True,\n",
    "        color=COLOURS[2],\n",
    "        alpha=0.5,\n",
    "        label='Generated',\n",
    "        ax=axs[i]\n",
    "    )\n",
    "    axs[i].set(title=label, xlabel='Test accuracy', ylabel='Probability', ylim=(0, 1))\n",
    "    axs[i].legend()\n",
    "    \n",
    "plt.subplots_adjust(hspace=0.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "399d8362",
   "metadata": {},
   "source": [
    "### Verifying model diversity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "04036bc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def x_to_checkpoint(x, p, i):\n",
    "    checkpoint = tokens_to_checkpoint(x[i], p[i], reference_checkpoint_cuda)\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 model\n",
    "\n",
    "def z_to_checkpoint(hr, z, p, i):\n",
    "    z = z.type(next(hr.parameters()).dtype)\n",
    "\n",
    "    y = hr.to(z.device).forward_decoder(z, p)\n",
    "    checkpoint = tokens_to_checkpoint(y[i], p[i], reference_checkpoint_cuda)\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 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "404b8154",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e2311a37c3ab4145b094da909cd43caa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/5 [00:00<?, ?hyperrepresentation model/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "GEN_DIVERSITY_PATH = EXPERIMENT_PATH / 'gen_diversity.json'\n",
    "\n",
    "if not GEN_DIVERSITY_PATH.is_file() or RERUN:\n",
    "    gen_diversity = dict()\n",
    "    \n",
    "    p = next(iter(ZOO_DL))[2]\n",
    "    \n",
    "    # Original model zoo\n",
    "    df_x_good_models = df_x[df_x['Model index'].isin(good_models_ids)]\n",
    "    x = df_x_good_models[df_x_good_models['Hyper-representation model'] == list(idx_dict.values())[0]][X_COLS].values.reshape((-1, X_DIM_1, X_DIM_2))\n",
    "    x = torch.tensor(x).to('cuda')\n",
    "    p_x = p[0].repeat(x.shape[0], 1, 1).to('cuda')\n",
    "    original_models = [x_to_checkpoint(x, p_x, i) for i in range(x.shape[0])]\n",
    "    gen_diversity['Model zoo'] = dict()\n",
    "    \n",
    "    # L2 (weights)\n",
    "    dist_mat = distance_matrix(original_models, metric='l2')\n",
    "    distances = dist_mat[np.triu_indices(dist_mat.shape[0], k=1)]\n",
    "    gen_diversity['Model zoo']['l2_structure'] = distances.tolist()\n",
    "    \n",
    "    # L2 (predictions)\n",
    "    dist_mat = distance_matrix(original_models, metric='l2_preds', testset='data', ds_test=TESTSET_DS)\n",
    "    distances = dist_mat[np.triu_indices(dist_mat.shape[0], k=1)]\n",
    "    gen_diversity['Model zoo']['l2_behaviour'] = distances.tolist()\n",
    "    \n",
    "    del df_x_good_models, x, original_models\n",
    "\n",
    "    for hr_idx, hr_label in tqdm(idx_dict.items(), unit='hyperrepresentation model'):        \n",
    "        z = torch.tensor(z_samples[hr_label]).to('cuda')\n",
    "        p_z = p[0].repeat(z.shape[0], 1, 1).to('cuda')\n",
    "        generated_models = [z_to_checkpoint(hyper_reps[hr_idx], z, p_z, i) for i in range(z.shape[0])]\n",
    "        gen_diversity[hr_label] = dict()\n",
    "    \n",
    "        # L2 (weights)\n",
    "        dist_mat = distance_matrix(generated_models, metric='l2')\n",
    "        distances = dist_mat[np.triu_indices(dist_mat.shape[0], k=1)]\n",
    "        gen_diversity[hr_label]['l2_structure'] = distances.tolist()\n",
    "\n",
    "        # L2 (predictions)\n",
    "        dist_mat = distance_matrix(generated_models, metric='l2_preds', testset='data', ds_test=TESTSET_DS)\n",
    "        distances = dist_mat[np.triu_indices(dist_mat.shape[0], k=1)]\n",
    "        gen_diversity[hr_label]['l2_behaviour'] = distances.tolist()\n",
    "        \n",
    "        del z, generated_models\n",
    "                \n",
    "    # Save results\n",
    "    with open(GEN_DIVERSITY_PATH, 'w') as ofh:\n",
    "        json.dump(gen_diversity, ofh)\n",
    "    \n",
    "with open(GEN_DIVERSITY_PATH, 'r') as ifh:\n",
    "    gen_diversity = json.load(ifh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "b3389108",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 67736 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, ax = plt.subplots(figsize=(10, 6))\n",
    "\n",
    "sns.kdeplot(gen_diversity['Model zoo']['l2_structure'], color=COLOURS[0], label='Model zoo (anchors)', ax=ax)\n",
    "sns.kdeplot(gen_diversity['Contrastive + Structure (Baseline)']['l2_structure'], color=COLOURS[1], label='Contrastive + Structure (Baseline)', ax=ax)\n",
    "sns.kdeplot(gen_diversity['Contrastive + Structure + Behavior']['l2_structure'], color=COLOURS[2], label='Contrastive + Structure + Behavior', ax=ax)\n",
    "\n",
    "ax.set(xlabel='L2 distance - Structure')\n",
    "\n",
    "plt.legend()\n",
    "plt.suptitle('Distribution of the pairwise structural L2 distances')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "de29df74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 75580 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, ax = plt.subplots(figsize=(10, 6))\n",
    "\n",
    "sns.kdeplot(gen_diversity['Model zoo']['l2_behaviour'], color=COLOURS[0], label='Model zoo (anchors)', ax=ax)\n",
    "sns.kdeplot(gen_diversity['Contrastive + Structure (Baseline)']['l2_behaviour'], color=COLOURS[1], label='Contrastive + Structure (Baseline)', ax=ax)\n",
    "sns.kdeplot(gen_diversity['Contrastive + Structure + Behavior']['l2_behaviour'], color=COLOURS[2], label='Contrastive + Structure + Behavior', ax=ax)\n",
    "\n",
    "ax.set(xlabel='L2 distance - Behavior')\n",
    "\n",
    "plt.legend()\n",
    "plt.suptitle('Distribution of the pairwise behavioural L2 distances')\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
}
