{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d98765e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import warnings\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm\n",
    "from copy import deepcopy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "116617c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "{YOUR_WORKING_PATH}/zoo_utils.py:6: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import re\n",
    "\n",
    "from utils import test_model, compare_models_prediction, generate_z_samples\n",
    "from zoo_utils import load_hyperrep_edx\n",
    "\n",
    "from shrp.models.def_net import CNN, CNN3\n",
    "from shrp.models.def_AE_trainable import get_transformations\n",
    "from shrp.datasets.dataset_auxiliaries import tokens_to_recipe, tokens_to_checkpoint\n",
    "from shrp.models.def_loss import DistillationLoss, ReconDistillationLoss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0d465bf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "sns.set_theme(context='notebook', style='whitegrid',\n",
    "              rc={\n",
    "                  'axes.spines.right': False,\n",
    "                  'axes.spines.top': False,\n",
    "                  'xtick.bottom': True,\n",
    "                  'ytick.left': True,\n",
    "                  'font.family': 'Fira Sans'\n",
    "              })\n",
    "\n",
    "COLOURS = [\n",
    "    '#960018',  # Carmine Red\n",
    "    '#FED85D',  # Mustard Gold\n",
    "    '#434384',  # Marian Blue\n",
    "    '#54B674',  # Emerald Green\n",
    "    '#D183C9',  # French Mauve\n",
    "    '#318CE7',  # French Blue\n",
    "    '#D08000'   # Fulvous Orange\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a45fdd7",
   "metadata": {},
   "source": [
    "## Loading data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "755588f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASETS = ['svhn', 'cifar10', 'eurosat']\n",
    "DATASET_LABELS = {\n",
    "    'svhn': 'SVHN',\n",
    "    'cifar10': 'CIFAR-10',\n",
    "    'eurosat': 'EuroSAT'\n",
    "}\n",
    "\n",
    "DS_ORDER = ['SVHN', 'CIFAR-10', 'EuroSAT']\n",
    "HR_ORDER = ['Contrastive + Structure (Baseline)', 'Behavior', 'Contrastive + Behavior', 'Structure + Behavior', 'Contrastive + Structure + Behavior']\n",
    "\n",
    "HYPERREP_PATH = Path('{YOUR_WORKING_PATH}/results/hyperrepresentations/tune')\n",
    "\n",
    "EXPERIMENT_PATH = Path.cwd() / 'results' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e0f27f84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "189b6e9e26cc4cf6962917424f17535f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading data:   0%|          | 0/3 [00:00<?, ?dataset/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Model zoo</th>\n",
       "      <th>Trial ID</th>\n",
       "      <th>Learning rate</th>\n",
       "      <th>Latent dimension</th>\n",
       "      <th>Behaviour loss</th>\n",
       "      <th>Query set</th>\n",
       "      <th>Query dataset</th>\n",
       "      <th># of queries</th>\n",
       "      <th>Gamma</th>\n",
       "      <th>Beta</th>\n",
       "      <th>...</th>\n",
       "      <th>Train loss (recon)</th>\n",
       "      <th>Train loss (structure)</th>\n",
       "      <th>Train loss (behaviour)</th>\n",
       "      <th>Test loss</th>\n",
       "      <th>Test loss (contrast)</th>\n",
       "      <th>Test loss (recon)</th>\n",
       "      <th>Test loss (structure)</th>\n",
       "      <th>Test loss (behaviour)</th>\n",
       "      <th>Dstk: Test accuracy</th>\n",
       "      <th>Dstk: GGap</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>1.344194</td>\n",
       "      <td>0.930796</td>\n",
       "      <td>1.390127</td>\n",
       "      <td>1.059503</td>\n",
       "      <td>1.045855</td>\n",
       "      <td>1.060222</td>\n",
       "      <td>0.791913</td>\n",
       "      <td>1.090034</td>\n",
       "      <td>0.559680</td>\n",
       "      <td>0.283559</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.275532</td>\n",
       "      <td>0.770928</td>\n",
       "      <td>0.220488</td>\n",
       "      <td>0.257198</td>\n",
       "      <td>0.887651</td>\n",
       "      <td>0.224016</td>\n",
       "      <td>0.591360</td>\n",
       "      <td>0.183200</td>\n",
       "      <td>0.554503</td>\n",
       "      <td>0.276677</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.266038</td>\n",
       "      <td>0.677502</td>\n",
       "      <td>0.220320</td>\n",
       "      <td>0.247892</td>\n",
       "      <td>0.708432</td>\n",
       "      <td>0.223653</td>\n",
       "      <td>0.544859</td>\n",
       "      <td>0.187963</td>\n",
       "      <td>0.550643</td>\n",
       "      <td>0.276275</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.264123</td>\n",
       "      <td>0.664841</td>\n",
       "      <td>0.219599</td>\n",
       "      <td>0.239949</td>\n",
       "      <td>0.652948</td>\n",
       "      <td>0.218212</td>\n",
       "      <td>0.527920</td>\n",
       "      <td>0.183800</td>\n",
       "      <td>0.551713</td>\n",
       "      <td>0.279411</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>50f57_00020</td>\n",
       "      <td>0.00001</td>\n",
       "      <td>64</td>\n",
       "      <td>cross_entropy</td>\n",
       "      <td>data</td>\n",
       "      <td>SVHN</td>\n",
       "      <td>256</td>\n",
       "      <td>0.05</td>\n",
       "      <td>0.1</td>\n",
       "      <td>...</td>\n",
       "      <td>0.261248</td>\n",
       "      <td>0.628665</td>\n",
       "      <td>0.220424</td>\n",
       "      <td>0.242219</td>\n",
       "      <td>0.741064</td>\n",
       "      <td>0.215964</td>\n",
       "      <td>0.495595</td>\n",
       "      <td>0.184894</td>\n",
       "      <td>0.538348</td>\n",
       "      <td>0.282729</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 23 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "  Model zoo     Trial ID  Learning rate  Latent dimension Behaviour loss  \\\n",
       "0      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "1      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "2      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "3      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "4      SVHN  50f57_00020        0.00001                64  cross_entropy   \n",
       "\n",
       "  Query set Query dataset  # of queries  Gamma  Beta  ...  Train loss (recon)  \\\n",
       "0      data          SVHN           256   0.05   0.1  ...            1.344194   \n",
       "1      data          SVHN           256   0.05   0.1  ...            0.275532   \n",
       "2      data          SVHN           256   0.05   0.1  ...            0.266038   \n",
       "3      data          SVHN           256   0.05   0.1  ...            0.264123   \n",
       "4      data          SVHN           256   0.05   0.1  ...            0.261248   \n",
       "\n",
       "   Train loss (structure)  Train loss (behaviour)  Test loss  \\\n",
       "0                0.930796                1.390127   1.059503   \n",
       "1                0.770928                0.220488   0.257198   \n",
       "2                0.677502                0.220320   0.247892   \n",
       "3                0.664841                0.219599   0.239949   \n",
       "4                0.628665                0.220424   0.242219   \n",
       "\n",
       "   Test loss (contrast)  Test loss (recon)  Test loss (structure)  \\\n",
       "0              1.045855           1.060222               0.791913   \n",
       "1              0.887651           0.224016               0.591360   \n",
       "2              0.708432           0.223653               0.544859   \n",
       "3              0.652948           0.218212               0.527920   \n",
       "4              0.741064           0.215964               0.495595   \n",
       "\n",
       "   Test loss (behaviour)  Dstk: Test accuracy  Dstk: GGap  \n",
       "0               1.090034             0.559680    0.283559  \n",
       "1               0.183200             0.554503    0.276677  \n",
       "2               0.187963             0.550643    0.276275  \n",
       "3               0.183800             0.551713    0.279411  \n",
       "4               0.184894             0.538348    0.282729  \n",
       "\n",
       "[5 rows x 23 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "MAX_EPOCH = 100\n",
    "\n",
    "df_hr = pd.DataFrame(columns=[\n",
    "    'Model zoo',\n",
    "    'Trial ID',\n",
    "    'Learning rate',\n",
    "    'Latent dimension',\n",
    "    'Behaviour loss',\n",
    "    'Query set',\n",
    "    'Query dataset',\n",
    "    '# of queries',\n",
    "    'Gamma',\n",
    "    'Beta',\n",
    "    'Epoch',\n",
    "    'Train loss',\n",
    "    'Train loss (contrast)', \n",
    "    'Train loss (recon)',\n",
    "    'Train loss (structure)', \n",
    "    'Train loss (behaviour)',\n",
    "    'Test loss',\n",
    "    'Test loss (contrast)', \n",
    "    'Test loss (recon)',\n",
    "    'Test loss (structure)', \n",
    "    'Test loss (behaviour)',\n",
    "    'Dstk: Test accuracy',\n",
    "    'Dstk: GGap'\n",
    "])\n",
    "\n",
    "for ds in tqdm(DATASETS, desc='Loading data', unit='dataset'):\n",
    "    hr_path = HYPERREP_PATH / f\"{ds}_train_hyperrep\"\n",
    "    for path in os.listdir(hr_path):\n",
    "        if not re.match('AE_trainable', path):\n",
    "            continue\n",
    "        path = hr_path / path\n",
    "\n",
    "        with open(path / 'params.json') as ifh_config, open(path / 'result.json') as ifh_result:\n",
    "            config = json.load(ifh_config)\n",
    "            result = pd.read_json(ifh_result, lines=True, dtype={'trial_id': str})\n",
    "\n",
    "        for edx in range(len(result)):\n",
    "            df_hr.loc[len(df_hr)] = (\n",
    "                DATASET_LABELS[ds],\n",
    "                result.loc[edx]['trial_id'],\n",
    "                config['optim::lr'],\n",
    "                config['ae:lat_dim'],\n",
    "                config['training::loss_distillation'],\n",
    "                config['training::queryset_distillation'],\n",
    "                config['training::queryset_dump'],\n",
    "                config['training::n_queries_distillation'],\n",
    "                config['training::gamma'],\n",
    "                config['training::beta'],\n",
    "                result.loc[edx]['training_iteration'],\n",
    "                result.loc[edx]['loss/loss_train'],\n",
    "                result.loc[edx]['loss/loss_contrast_train'],\n",
    "                result.loc[edx]['loss/loss_recon_train'],\n",
    "                result.loc[edx]['loss/loss_structure_train'],\n",
    "                result.loc[edx]['loss/loss_behaviour_train'],\n",
    "                result.loc[edx]['loss/loss_test'],\n",
    "                result.loc[edx]['loss/loss_contrast_test'],\n",
    "                result.loc[edx]['loss/loss_recon_test'],\n",
    "                result.loc[edx]['loss/loss_structure_test'],\n",
    "                result.loc[edx]['loss/loss_behaviour_test'],\n",
    "                result.loc[edx]['dstk/test_acc_test'],\n",
    "                result.loc[edx]['dstk/ggap_test']\n",
    "            )\n",
    "\n",
    "df_hr['Query dataset'] = df_hr['Query dataset'].map({\n",
    "    '{YOUR_WORKING_PATH}/results/data/svhn_train_zoo_relu/dataset.pt': 'SVHN',\n",
    "    '{YOUR_WORKING_PATH}/results/data/cifar10_train_zoo_relu/dataset.pt': 'CIFAR-10',\n",
    "    '{YOUR_WORKING_PATH}/results/data/eurosat_train_zoo_relu/dataset.pt': 'EuroSAT',\n",
    "    '{YOUR_WORKING_PATH}/results/data/tinyimagenet_32/dataset.pt': 'TinyImagenet-32'\n",
    "})\n",
    "df_hr.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "407a49eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "idx_dict = {\n",
    "    '00011': 'Contrastive + Structure (Baseline)',\n",
    "    '00000': 'Behavior',\n",
    "    '00006': 'Contrastive + Behavior',\n",
    "    '00002': 'Structure + Behavior',\n",
    "    '00008': 'Contrastive + Structure + Behavior'\n",
    "}\n",
    "\n",
    "df_hr['Hyper-representation'] = df_hr['Trial ID'].str.slice(start=-5).map(idx_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c04ee8e1",
   "metadata": {},
   "source": [
    "### Loading $L^2_{weights}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a33468f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f348d32a7a045ef9e2d9857c3ff98f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading data:   0%|          | 0/3 [00:00<?, ?dataset/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Hyper-representation model</th>\n",
       "      <th>L2 (structure)</th>\n",
       "      <th>MSE (structure)</th>\n",
       "      <th>Model zoo</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>36.630211</td>\n",
       "      <td>0.049923</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>32.174301</td>\n",
       "      <td>0.038516</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>29.809624</td>\n",
       "      <td>0.033062</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>28.346432</td>\n",
       "      <td>0.029896</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>25.570820</td>\n",
       "      <td>0.024328</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           Hyper-representation model  L2 (structure)  MSE (structure)  \\\n",
       "0  Contrastive + Structure (Baseline)       36.630211         0.049923   \n",
       "1  Contrastive + Structure (Baseline)       32.174301         0.038516   \n",
       "2  Contrastive + Structure (Baseline)       29.809624         0.033062   \n",
       "3  Contrastive + Structure (Baseline)       28.346432         0.029896   \n",
       "4  Contrastive + Structure (Baseline)       25.570820         0.024328   \n",
       "\n",
       "  Model zoo  \n",
       "0      SVHN  \n",
       "1      SVHN  \n",
       "2      SVHN  \n",
       "3      SVHN  \n",
       "4      SVHN  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l2_weights = list()\n",
    "\n",
    "for ds in tqdm(DATASETS, desc='Loading data', unit='dataset'):\n",
    "    df_l2_weights_path = EXPERIMENT_PATH / f'hr_analysis_{ds}' / 'df_zoo_l2weights.csv'\n",
    "    df_l2_weights = pd.read_csv(df_l2_weights_path, index_col=0)\n",
    "    \n",
    "    df_l2_weights['Model zoo'] = DATASET_LABELS[ds]\n",
    "    \n",
    "    l2_weights.append(df_l2_weights)\n",
    "    \n",
    "l2_weights = pd.concat(l2_weights)\n",
    "l2_weights.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56efc173",
   "metadata": {},
   "source": [
    "### Loading $L^2_{preds}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4b19f338",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbf17903abbb4f38adfc1fae2fbc6a88",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading data:   0%|          | 0/3 [00:00<?, ?dataset/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Hyper-representation model</th>\n",
       "      <th>L2 (behavior)</th>\n",
       "      <th>MSE (behavior)</th>\n",
       "      <th>Agreement</th>\n",
       "      <th>Model zoo</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Contrastive + Structure (Baseline)</td>\n",
       "      <td>1.006281</td>\n",
       "      <td>0.121435</td>\n",
       "      <td>0.217386</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Behavior</td>\n",
       "      <td>0.671536</td>\n",
       "      <td>0.049624</td>\n",
       "      <td>0.217386</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Contrastive + Behavior</td>\n",
       "      <td>0.671696</td>\n",
       "      <td>0.049618</td>\n",
       "      <td>0.217386</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Structure + Behavior</td>\n",
       "      <td>0.235445</td>\n",
       "      <td>0.010870</td>\n",
       "      <td>0.780040</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Contrastive + Structure + Behavior</td>\n",
       "      <td>0.241480</td>\n",
       "      <td>0.011683</td>\n",
       "      <td>0.781000</td>\n",
       "      <td>SVHN</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           Hyper-representation model  L2 (behavior)  MSE (behavior)  \\\n",
       "0  Contrastive + Structure (Baseline)       1.006281        0.121435   \n",
       "1                            Behavior       0.671536        0.049624   \n",
       "2              Contrastive + Behavior       0.671696        0.049618   \n",
       "3                Structure + Behavior       0.235445        0.010870   \n",
       "4  Contrastive + Structure + Behavior       0.241480        0.011683   \n",
       "\n",
       "   Agreement Model zoo  \n",
       "0   0.217386      SVHN  \n",
       "1   0.217386      SVHN  \n",
       "2   0.217386      SVHN  \n",
       "3   0.780040      SVHN  \n",
       "4   0.781000      SVHN  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l2_preds = list()\n",
    "\n",
    "for ds in tqdm(DATASETS, desc='Loading data', unit='dataset'):\n",
    "    df_l2_preds_path = EXPERIMENT_PATH / f'hr_analysis_{ds}' / 'df_zoo_l2preds.csv'\n",
    "    df_l2_preds = pd.read_csv(df_l2_preds_path, index_col=0)\n",
    "    \n",
    "    df_l2_preds['Model zoo'] = DATASET_LABELS[ds]\n",
    "    \n",
    "    l2_preds.append(df_l2_preds)\n",
    "    \n",
    "l2_preds = pd.concat(l2_preds)\n",
    "l2_preds.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b0e5bf",
   "metadata": {},
   "source": [
    "## Analysing data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "829bdbd0",
   "metadata": {},
   "source": [
    "### Analysing structural distance ($L_2$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "79ebb7c5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{llll}\n",
      "\\toprule\n",
      "Model zoo &        SVHN &    CIFAR-10 &     EuroSAT \\\\\n",
      "Hyper-representation model         &             &             &             \\\\\n",
      "\\midrule\n",
      "Contrastive + Structure (Baseline) &  30.7 (±10) &  31.3 (±13) &  41.4 (±22) \\\\\n",
      "Behavior                           &   57.8 (±6) &   59.5 (±9) &  67.5 (±17) \\\\\n",
      "Contrastive + Behavior             &   55.7 (±7) &   60.7 (±9) &  67.4 (±17) \\\\\n",
      "Structure + Behavior               &  26.9 (±10) &  29.6 (±13) &  45.5 (±21) \\\\\n",
      "Contrastive + Structure + Behavior &  27.1 (±10) &  30.0 (±13) &  43.6 (±21) \\\\\n",
      "\\bottomrule\n",
      "\\end{tabular}\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_28150/2934977399.py:5: FutureWarning: In future versions `DataFrame.to_latex` is expected to utilise the base implementation of `Styler.to_latex` for formatting and rendering. The arguments signature may therefore change. It is recommended instead to use `DataFrame.style.to_latex` which also contains additional functionality.\n",
      "  print(df_dist_struct.to_latex())\n"
     ]
    },
    {
     "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>Model zoo</th>\n",
       "      <th>SVHN</th>\n",
       "      <th>CIFAR-10</th>\n",
       "      <th>EuroSAT</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Hyper-representation model</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Contrastive + Structure (Baseline)</th>\n",
       "      <td>30.7 (±10)</td>\n",
       "      <td>31.3 (±13)</td>\n",
       "      <td>41.4 (±22)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Behavior</th>\n",
       "      <td>57.8 (±6)</td>\n",
       "      <td>59.5 (±9)</td>\n",
       "      <td>67.5 (±17)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Contrastive + Behavior</th>\n",
       "      <td>55.7 (±7)</td>\n",
       "      <td>60.7 (±9)</td>\n",
       "      <td>67.4 (±17)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Structure + Behavior</th>\n",
       "      <td>26.9 (±10)</td>\n",
       "      <td>29.6 (±13)</td>\n",
       "      <td>45.5 (±21)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Contrastive + Structure + Behavior</th>\n",
       "      <td>27.1 (±10)</td>\n",
       "      <td>30.0 (±13)</td>\n",
       "      <td>43.6 (±21)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Model zoo                                 SVHN    CIFAR-10     EuroSAT\n",
       "Hyper-representation model                                            \n",
       "Contrastive + Structure (Baseline)  30.7 (±10)  31.3 (±13)  41.4 (±22)\n",
       "Behavior                             57.8 (±6)   59.5 (±9)  67.5 (±17)\n",
       "Contrastive + Behavior               55.7 (±7)   60.7 (±9)  67.4 (±17)\n",
       "Structure + Behavior                26.9 (±10)  29.6 (±13)  45.5 (±21)\n",
       "Contrastive + Structure + Behavior  27.1 (±10)  30.0 (±13)  43.6 (±21)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_dist_struct = l2_weights.groupby(['Hyper-representation model', 'Model zoo'])['L2 (structure)'].agg(['mean', 'std']).reset_index()\n",
    "df_dist_struct['print'] = df_dist_struct['mean'].map(lambda x: f\"{x:.1f}\") + df_dist_struct['std'].map(lambda x: f\" (±{x:.0f})\")\n",
    "df_dist_struct = df_dist_struct.pivot(index='Hyper-representation model', columns='Model zoo', values='print')\n",
    "df_dist_struct = df_dist_struct.loc[HR_ORDER, DS_ORDER]\n",
    "print(df_dist_struct.to_latex())\n",
    "df_dist_struct"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49dc11bc",
   "metadata": {},
   "source": [
    "### Analysing behavioral distance (agreement)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "395b45e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\begin{tabular}{llll}\n",
      "\\toprule\n",
      "Model zoo &          SVHN &      CIFAR-10 &       EuroSAT \\\\\n",
      "Hyper-representation model         &               &               &               \\\\\n",
      "\\midrule\n",
      "Contrastive + Structure (Baseline) &  50.0\\% (±20\\%) &  31.9\\% (±16\\%) &  35.9\\% (±24\\%) \\\\\n",
      "Behavior                           &  26.6\\% (±22\\%) &   10.3\\% (±8\\%) &  19.2\\% (±25\\%) \\\\\n",
      "Contrastive + Behavior             &  26.6\\% (±22\\%) &  11.8\\% (±14\\%) &  18.5\\% (±25\\%) \\\\\n",
      "Structure + Behavior               &  86.1\\% (±11\\%) &  59.4\\% (±20\\%) &  66.5\\% (±25\\%) \\\\\n",
      "Contrastive + Structure + Behavior &   87.0\\% (±9\\%) &  59.6\\% (±19\\%) &  65.8\\% (±25\\%) \\\\\n",
      "\\bottomrule\n",
      "\\end{tabular}\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_28150/1650146689.py:5: FutureWarning: In future versions `DataFrame.to_latex` is expected to utilise the base implementation of `Styler.to_latex` for formatting and rendering. The arguments signature may therefore change. It is recommended instead to use `DataFrame.style.to_latex` which also contains additional functionality.\n",
      "  print(df_dist_struct.to_latex())\n"
     ]
    },
    {
     "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>Model zoo</th>\n",
       "      <th>SVHN</th>\n",
       "      <th>CIFAR-10</th>\n",
       "      <th>EuroSAT</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Hyper-representation model</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Contrastive + Structure (Baseline)</th>\n",
       "      <td>50.0% (±20%)</td>\n",
       "      <td>31.9% (±16%)</td>\n",
       "      <td>35.9% (±24%)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Behavior</th>\n",
       "      <td>26.6% (±22%)</td>\n",
       "      <td>10.3% (±8%)</td>\n",
       "      <td>19.2% (±25%)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Contrastive + Behavior</th>\n",
       "      <td>26.6% (±22%)</td>\n",
       "      <td>11.8% (±14%)</td>\n",
       "      <td>18.5% (±25%)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Structure + Behavior</th>\n",
       "      <td>86.1% (±11%)</td>\n",
       "      <td>59.4% (±20%)</td>\n",
       "      <td>66.5% (±25%)</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Contrastive + Structure + Behavior</th>\n",
       "      <td>87.0% (±9%)</td>\n",
       "      <td>59.6% (±19%)</td>\n",
       "      <td>65.8% (±25%)</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "Model zoo                                   SVHN      CIFAR-10       EuroSAT\n",
       "Hyper-representation model                                                  \n",
       "Contrastive + Structure (Baseline)  50.0% (±20%)  31.9% (±16%)  35.9% (±24%)\n",
       "Behavior                            26.6% (±22%)   10.3% (±8%)  19.2% (±25%)\n",
       "Contrastive + Behavior              26.6% (±22%)  11.8% (±14%)  18.5% (±25%)\n",
       "Structure + Behavior                86.1% (±11%)  59.4% (±20%)  66.5% (±25%)\n",
       "Contrastive + Structure + Behavior   87.0% (±9%)  59.6% (±19%)  65.8% (±25%)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_dist_struct = l2_preds.groupby(['Hyper-representation model', 'Model zoo'])['Agreement'].agg(['mean', 'std']).reset_index()\n",
    "df_dist_struct['print'] = df_dist_struct['mean'].map(lambda x: f\"{x:.1%}\") + df_dist_struct['std'].map(lambda x: f\" (±{x:.0%})\")\n",
    "df_dist_struct = df_dist_struct.pivot(index='Hyper-representation model', columns='Model zoo', values='print')\n",
    "df_dist_struct = df_dist_struct.loc[HR_ORDER, DS_ORDER]\n",
    "print(df_dist_struct.to_latex())\n",
    "df_dist_struct"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d157bbf",
   "metadata": {},
   "source": [
    "### Plotting figures"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59919522",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 168044 characters]",
      "text/plain": [
       "<Figure size 1400x600 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "hr_to_plot = ['Contrastive + Structure (Baseline)', 'Contrastive + Behavior', 'Contrastive + Structure + Behavior']\n",
    "\n",
    "_, axs = plt.subplots(2, 3, figsize=(14, 6))\n",
    "\n",
    "for d, ds in enumerate(DATASETS):\n",
    "    ds_label = DATASET_LABELS[ds]\n",
    "    df_s = l2_weights[l2_weights['Model zoo'] == ds_label]\n",
    "    df_b = l2_preds[l2_preds['Model zoo'] == ds_label]\n",
    "\n",
    "    for i, label in enumerate(hr_to_plot):\n",
    "        sns.kdeplot(df_s[df_s['Hyper-representation model'] == label], x='L2 (structure)', color=COLOURS[i], label=label, ax=axs[0][d])\n",
    "    axs[0][d].set(title=f\"{ds_label}\\nStructural distance ($L^2$)\", xlabel=\"$L^2$ between weights\")\n",
    "    \n",
    "    for i, label in enumerate(hr_to_plot):\n",
    "        sns.histplot(df_b[df_b['Hyper-representation model'] == label], x='Agreement', binwidth=0.1, binrange=(0., 1.), kde=True, color=COLOURS[i], label=label, ax=axs[1][d])\n",
    "    axs[1][d].set(title=f\"{ds_label}\\nBehavioral similarity (Agreement)\", xlabel=\"Model agreement\")\n",
    "    \n",
    "plt.legend(loc=(-2.1, -0.50), ncols=3)\n",
    "plt.subplots_adjust(hspace=0.6, wspace=0.25)\n",
    "plt.savefig('figures/reconstructive_dstk.png', dpi=300, bbox_inches='tight')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
