{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d2e079f0",
   "metadata": {},
   "source": [
    "# Analysing and controlling our model zoos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b4a7985b",
   "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": "21eac799",
   "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 zoo_utils import load_hyperrep_edx\n",
    "from model_distances import compute_predictions, epsilon_delta_distance, distance_matrix\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": "382bfed0",
   "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": "code",
   "execution_count": 4,
   "id": "4910aef2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import StandardScaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "885a4411",
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPERIMENT_PATH = Path.cwd() / 'results' / 'zoo_analysis'\n",
    "\n",
    "if not EXPERIMENT_PATH.is_dir():\n",
    "    os.makedirs(EXPERIMENT_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36556984",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9280c33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "RERUN = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "676ff495",
   "metadata": {},
   "outputs": [],
   "source": [
    "ZOO_PATHS = {\n",
    "    'SVHN': Path('{YOUR_WORKING_PATH}/results/zoos/svhn_train_zoo_relu'),\n",
    "    'CIFAR-10': Path('{YOUR_WORKING_PATH}/results/zoos/cifar10_train_zoo_relu'),\n",
    "    'EuroSAT': Path('{YOUR_WORKING_PATH}/results/zoos/eurosat_train_zoo_relu')\n",
    "}\n",
    "\n",
    "TESTSETS = {\n",
    "    'SVHN': torch.load('{YOUR_WORKING_PATH}/results/data/svhn_train_zoo_relu/dataset.pt')['testset'],\n",
    "    'CIFAR-10': torch.load('{YOUR_WORKING_PATH}/results/data/cifar10_train_zoo_relu/dataset.pt')['testset'],\n",
    "    'EuroSAT': torch.load('{YOUR_WORKING_PATH}/results/data/eurosat_train_zoo_relu/dataset.pt')['testset']\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da843874",
   "metadata": {},
   "source": [
    "## Visualising zoo performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "705148fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Model zoo</th>\n",
       "      <th>Trial ID</th>\n",
       "      <th>Activation function</th>\n",
       "      <th>Initialisation</th>\n",
       "      <th>Learning rate</th>\n",
       "      <th>Weight decay</th>\n",
       "      <th>Seed</th>\n",
       "      <th>Epoch</th>\n",
       "      <th>Test accuracy</th>\n",
       "      <th>Generalisation gap</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>7597300394</td>\n",
       "      <td>relu</td>\n",
       "      <td>kaiming_uniform</td>\n",
       "      <td>0.00075</td>\n",
       "      <td>0.0005</td>\n",
       "      <td>6</td>\n",
       "      <td>0</td>\n",
       "      <td>0.071758</td>\n",
       "      <td>-0.001563</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>7597300394</td>\n",
       "      <td>relu</td>\n",
       "      <td>kaiming_uniform</td>\n",
       "      <td>0.00075</td>\n",
       "      <td>0.0005</td>\n",
       "      <td>6</td>\n",
       "      <td>1</td>\n",
       "      <td>0.775200</td>\n",
       "      <td>-0.182034</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>7597300394</td>\n",
       "      <td>relu</td>\n",
       "      <td>kaiming_uniform</td>\n",
       "      <td>0.00075</td>\n",
       "      <td>0.0005</td>\n",
       "      <td>6</td>\n",
       "      <td>2</td>\n",
       "      <td>0.829479</td>\n",
       "      <td>-0.016614</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>7597300394</td>\n",
       "      <td>relu</td>\n",
       "      <td>kaiming_uniform</td>\n",
       "      <td>0.00075</td>\n",
       "      <td>0.0005</td>\n",
       "      <td>6</td>\n",
       "      <td>3</td>\n",
       "      <td>0.846381</td>\n",
       "      <td>-0.003682</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>7597300394</td>\n",
       "      <td>relu</td>\n",
       "      <td>kaiming_uniform</td>\n",
       "      <td>0.00075</td>\n",
       "      <td>0.0005</td>\n",
       "      <td>6</td>\n",
       "      <td>4</td>\n",
       "      <td>0.848033</td>\n",
       "      <td>0.010410</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Model zoo    Trial ID Activation function   Initialisation  Learning rate  \\\n",
       "0      SVHN  7597300394                relu  kaiming_uniform        0.00075   \n",
       "1      SVHN  7597300394                relu  kaiming_uniform        0.00075   \n",
       "2      SVHN  7597300394                relu  kaiming_uniform        0.00075   \n",
       "3      SVHN  7597300394                relu  kaiming_uniform        0.00075   \n",
       "4      SVHN  7597300394                relu  kaiming_uniform        0.00075   \n",
       "\n",
       "   Weight decay  Seed  Epoch  Test accuracy  Generalisation gap  \n",
       "0        0.0005     6      0       0.071758           -0.001563  \n",
       "1        0.0005     6      1       0.775200           -0.182034  \n",
       "2        0.0005     6      2       0.829479           -0.016614  \n",
       "3        0.0005     6      3       0.846381           -0.003682  \n",
       "4        0.0005     6      4       0.848033            0.010410  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DF_ZOO_PATH = EXPERIMENT_PATH / 'df_zoo.csv'\n",
    "\n",
    "if not DF_ZOO_PATH.is_file() or RERUN:\n",
    "    df_zoo = pd.DataFrame(columns=[\n",
    "        'Model zoo',\n",
    "        'Trial ID',\n",
    "        'Activation function',\n",
    "        'Initialisation',\n",
    "        'Learning rate',\n",
    "        'Weight decay',\n",
    "        'Seed',\n",
    "        'Epoch',\n",
    "        'Test accuracy',\n",
    "        'Generalisation gap'\n",
    "    ])\n",
    "\n",
    "    for zoo_label, zoo_path in ZOO_PATHS.items():\n",
    "        for path in os.listdir(zoo_path):\n",
    "            if not re.match('NN_tune', path):\n",
    "                continue\n",
    "            path = zoo_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)\n",
    "\n",
    "            for edx in range(len(result)):\n",
    "                df_zoo.loc[len(df_zoo)] = (\n",
    "                    zoo_label,\n",
    "                    result.loc[edx]['trial_id'],\n",
    "                    config['model::nlin'],\n",
    "                    config['model::init_type'],\n",
    "                    config['optim::lr'],\n",
    "                    config['optim::wd'],\n",
    "                    config['seed'],\n",
    "                    result.loc[edx]['training_iteration'],\n",
    "                    result.loc[edx]['test_acc'],\n",
    "                    result.loc[edx]['train_acc'] - result.loc[edx]['test_acc']\n",
    "                )\n",
    "            \n",
    "    df_zoo.to_csv(DF_ZOO_PATH)\n",
    "\n",
    "df_zoo = pd.read_csv(DF_ZOO_PATH, index_col=0)\n",
    "df_zoo.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3bbcd7d9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 111792 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "zoo_label = 'SVHN'\n",
    "\n",
    "_, ax = plt.subplots(figsize=(12, 6))\n",
    "\n",
    "sns.boxplot(df_zoo[df_zoo['Model zoo'] == zoo_label], x='Epoch', y='Test accuracy', color=COLOURS[0], ax=ax).set(\n",
    "    title=f'Distribution of model performance by epoch in the {zoo_label} model zoo'\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "db5373d4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88592 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "zoo_label = 'CIFAR-10'\n",
    "\n",
    "_, ax = plt.subplots(figsize=(12, 6))\n",
    "\n",
    "sns.boxplot(df_zoo[df_zoo['Model zoo'] == zoo_label], x='Epoch', y='Test accuracy', color=COLOURS[0], ax=ax).set(\n",
    "    title=f'Distribution of model performance by epoch in the {zoo_label} model zoo'\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6fe3488c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45832 characters]",
      "text/plain": [
       "<Figure size 1200x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "zoo_label = 'EuroSAT'\n",
    "\n",
    "_, ax = plt.subplots(figsize=(12, 6))\n",
    "\n",
    "sns.boxplot(df_zoo[df_zoo['Model zoo'] == zoo_label], x='Epoch', y='Test accuracy', color=COLOURS[0], ax=ax).set(\n",
    "    title=f'Distribution of model performance by epoch in the {zoo_label} model zoo'\n",
    ")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b12319cc",
   "metadata": {},
   "source": [
    "### Compact figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "8f780a33",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 65316 characters]",
      "text/plain": [
       "<Figure size 1400x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(1, 3, sharey=False, figsize=(14, 3))\n",
    "\n",
    "for i, zoo_label in enumerate(ZOO_PATHS.keys()):\n",
    "    sns.boxplot(df_zoo[df_zoo['Model zoo'] == zoo_label], x='Epoch', y='Test accuracy', color=COLOURS[1], ax=axs[i]).set(\n",
    "        title=zoo_label, ylim=(0., 0.95)\n",
    "    )\n",
    "    \n",
    "    # Set the x-axis ticks to show only one every five ticks\n",
    "    ticks = axs[i].get_xticks()\n",
    "    axs[i].set_xticks(ticks[::5])\n",
    "\n",
    "plt.savefig('figures/model_zoos_acc.png', dpi=300, bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b96c0790",
   "metadata": {},
   "source": [
    "## Evaluating the eigendecomposition of model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3f7bac67",
   "metadata": {},
   "outputs": [],
   "source": [
    "def model_to_vector(model):\n",
    "    # Get the state dictionary of the model\n",
    "    state_dict = model.state_dict()\n",
    "    # Flatten the state dictionary into a one-dimensional vector\n",
    "    vector = np.concatenate([param.data.detach().cpu().numpy().flatten() for param in state_dict.values()])\n",
    "    return vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2c5ff277",
   "metadata": {},
   "outputs": [],
   "source": [
    "def vector_to_model(vector, device='cuda'):\n",
    "    # Create a new instance of LeNet5\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",
    "    # Get the state dictionary of the model\n",
    "    state_dict = model.state_dict()\n",
    "    # Unflatten the vector into a state dictionary\n",
    "    offset = 0\n",
    "    for key, param in state_dict.items():\n",
    "        size = np.prod(param.size())\n",
    "        state_dict[key] = torch.from_numpy(vector[offset:offset+size].reshape(param.size()))\n",
    "        offset += size\n",
    "    # Load the state dictionary into the model\n",
    "    model.load_state_dict(state_dict)\n",
    "    return model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ffc9caa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model, dataloader):\n",
    "    model.eval()  # Set the model to evaluation mode\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    \n",
    "    device = next(model.parameters()).device\n",
    "\n",
    "    with torch.no_grad():  # Disable gradient calculation to save memory\n",
    "        for data in dataloader:\n",
    "            images, labels = [tensor.to(device) for tensor in data]\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "8c339cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading reference checkpoint\n",
    "zoo_path = ZOO_PATHS['SVHN']\n",
    "\n",
    "for path in os.listdir(zoo_path):\n",
    "    if 'NN_tune' in path and \"_00000_0\" in path:\n",
    "        ref_path = zoo_path / path\n",
    "\n",
    "with open(os.path.join(ref_path, 'params.json')) as ifh:\n",
    "    reference_params = json.load(ifh)\n",
    "        \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\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c1a9e736",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Model zoo</th>\n",
       "      <th>k</th>\n",
       "      <th>Explained variance</th>\n",
       "      <th>Selected eigenvalues</th>\n",
       "      <th>Reconstruction error</th>\n",
       "      <th>Average accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10689</td>\n",
       "      <td>0.100304</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>763.911083</td>\n",
       "      <td>0.206005</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>2</td>\n",
       "      <td>0.137084</td>\n",
       "      <td>Top</td>\n",
       "      <td>749.162090</td>\n",
       "      <td>0.224545</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10757</td>\n",
       "      <td>0.200174</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>720.509158</td>\n",
       "      <td>0.228431</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>4</td>\n",
       "      <td>0.221497</td>\n",
       "      <td>Top</td>\n",
       "      <td>709.749607</td>\n",
       "      <td>0.217102</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10691</td>\n",
       "      <td>0.237388</td>\n",
       "      <td>Extremes</td>\n",
       "      <td>705.039311</td>\n",
       "      <td>0.306440</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Model zoo      k  Explained variance Selected eigenvalues  \\\n",
       "0      SVHN  10689            0.100304               Bottom   \n",
       "1      SVHN      2            0.137084                  Top   \n",
       "2      SVHN  10757            0.200174               Bottom   \n",
       "3      SVHN      4            0.221497                  Top   \n",
       "4      SVHN  10691            0.237388             Extremes   \n",
       "\n",
       "   Reconstruction error  Average accuracy  \n",
       "0            763.911083          0.206005  \n",
       "1            749.162090          0.224545  \n",
       "2            720.509158          0.228431  \n",
       "3            709.749607          0.217102  \n",
       "4            705.039311          0.306440  "
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DF_EIGEN_PATH = EXPERIMENT_PATH / 'df_eigen.csv'\n",
    "\n",
    "if not DF_EIGEN_PATH.is_file() or RERUN:\n",
    "    # Create dataframe\n",
    "    df_eigen = pd.DataFrame(columns=['Model zoo', 'k', 'Explained variance', 'Selected eigenvalues', 'Reconstruction error', 'Average accuracy'])\n",
    "    \n",
    "    # Running experiments for every zoo\n",
    "    for zoo_label, zoo_path in tqdm(ZOO_PATHS.items(), desc='Iterating through the 3 model zoos', unit='zoo', total=3):\n",
    "        zoo_ds = torch.load(zoo_path / 'dataset_test.pt')\n",
    "        zoo_dl = torch.utils.data.DataLoader(zoo_ds, batch_size=64, shuffle=False)\n",
    "        \n",
    "        dataloader = torch.utils.data.DataLoader(TESTSETS[zoo_label], batch_size=128, shuffle=False)\n",
    "    \n",
    "        X = list()\n",
    "\n",
    "        # Run experiments\n",
    "        for bdx, data in enumerate(zoo_dl):\n",
    "            x, _, p, _ = data\n",
    "\n",
    "            for idx in range(x.shape[0]):\n",
    "                # Get original model\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",
    "\n",
    "                X.append(model_to_vector(orig_model))\n",
    "\n",
    "        X = np.vstack(X)\n",
    "\n",
    "        # Step 1: Center and scale the data using StandardScaler\n",
    "        scaler = StandardScaler()\n",
    "        X_scaled = scaler.fit_transform(X)\n",
    "\n",
    "        # Step 2: Compute the covariance matrix\n",
    "        cov_matrix = np.cov(X_scaled, rowvar=False)\n",
    "\n",
    "        # Step 3: Compute eigenvalues and eigenvectors\n",
    "        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)\n",
    "\n",
    "        # Step 4: Sort eigenvalues and eigenvectors\n",
    "        sorted_indices = np.argsort(eigenvalues)\n",
    "        eigenvalues_sorted = eigenvalues[sorted_indices]\n",
    "        eigenvectors_sorted = eigenvectors[:, sorted_indices]\n",
    "\n",
    "        # Step 5: Compute explained variance\n",
    "        total_variance = np.sum(eigenvalues_sorted)\n",
    "        explained_variance = eigenvalues_sorted / total_variance\n",
    "\n",
    "        # Step 6: Select k for specific explained variance levels\n",
    "        explained_variance_levels = [x/10 for x in range(1, 10)]\n",
    "        k_values = []\n",
    "        for level in explained_variance_levels:\n",
    "            cumulative_variance = np.cumsum(explained_variance)\n",
    "            k = np.searchsorted(cumulative_variance, level) + 1\n",
    "            k_values.append(k)\n",
    "        k_values.append(X.shape[1]) # We also want the full approximation as a baseline\n",
    "\n",
    "        # Step 7: Repeat the process using the top eigenvalues\n",
    "        sorted_indices_top = sorted_indices[::-1]\n",
    "        eigenvalues_sorted_top = eigenvalues[sorted_indices_top]\n",
    "        eigenvectors_sorted_top = eigenvectors[:, sorted_indices_top]\n",
    "\n",
    "        explained_variance_top = eigenvalues_sorted_top / total_variance\n",
    "        k_values_top = []\n",
    "        for level in explained_variance_levels:\n",
    "            cumulative_variance_top = np.cumsum(explained_variance_top)\n",
    "            k_top = np.searchsorted(cumulative_variance_top, level) + 1\n",
    "            k_values_top.append(k_top)\n",
    "        k_values_top.append(X.shape[1])\n",
    "\n",
    "        explained_variance_levels.append(1.0) # We add the label for full reconstruction\n",
    "\n",
    "        # Step 8: Project and reconstruct data for each k value, compute accuracy\n",
    "        for expl_var_idx, explained_variance_target in tqdm(enumerate(explained_variance_levels), desc=f'{zoo_label}: computing reconstructed accuracy', unit='experiment', total=len(explained_variance_levels)):\n",
    "            # Bottom\n",
    "            k = k_values[expl_var_idx]\n",
    "            explained_variance_actual = np.sum(explained_variance[:k])\n",
    "            selected_eigenvectors = eigenvectors_sorted[:, :k]\n",
    "            X_projected = np.dot(X_scaled, selected_eigenvectors)\n",
    "            X_reconstructed = np.dot(X_projected, selected_eigenvectors.T)\n",
    "            X_reconstructed = scaler.inverse_transform(X_reconstructed)\n",
    "\n",
    "            l2_bot = np.linalg.norm(X - X_reconstructed)\n",
    "            acc_bot = list()\n",
    "\n",
    "            for i in range(X.shape[0]):\n",
    "                acc_bot.append(test_model(vector_to_model(X_reconstructed[i]), dataloader))\n",
    "\n",
    "            df_eigen.loc[len(df_eigen)] = (\n",
    "                zoo_label,\n",
    "                k,\n",
    "                explained_variance_actual,\n",
    "                'Bottom',\n",
    "                l2_bot,\n",
    "                np.mean(acc_bot)\n",
    "            )\n",
    "\n",
    "            # Top\n",
    "            k = k_values_top[expl_var_idx]\n",
    "            explained_variance_actual = np.sum(explained_variance_top[:k])\n",
    "            selected_eigenvectors_top = eigenvectors_sorted_top[:, :k]\n",
    "            X_projected_top = np.dot(X_scaled, selected_eigenvectors_top)\n",
    "            X_reconstructed_top = np.dot(X_projected_top, selected_eigenvectors_top.T)\n",
    "            X_reconstructed_top = scaler.inverse_transform(X_reconstructed_top)\n",
    "\n",
    "            l2_top = np.linalg.norm(X - X_reconstructed_top)\n",
    "            acc_top = list()\n",
    "\n",
    "            for i in range(X.shape[0]):\n",
    "                acc_top.append(test_model(vector_to_model(X_reconstructed_top[i]), dataloader))\n",
    "\n",
    "            df_eigen.loc[len(df_eigen)] = (\n",
    "                zoo_label,\n",
    "                k,\n",
    "                explained_variance_actual,\n",
    "                'Top',\n",
    "                l2_top,\n",
    "                np.mean(acc_top)\n",
    "            )\n",
    "            \n",
    "            # Extremes\n",
    "            if expl_var_idx == 9:\n",
    "                df_eigen.loc[len(df_eigen)] = (\n",
    "                    zoo_label,\n",
    "                    k,\n",
    "                    explained_variance_actual,\n",
    "                    'Extremes',\n",
    "                    l2_top,\n",
    "                    np.mean(acc_top)\n",
    "                )\n",
    "            elif expl_var_idx % 2 == 1:\n",
    "                k_bot = k_values[expl_var_idx // 2]\n",
    "                k_top = k_values_top[expl_var_idx // 2]\n",
    "                k = k_bot + k_top\n",
    "                \n",
    "                explained_variance_actual = np.sum(explained_variance[:k_bot]) + np.sum(explained_variance_top[:k_top])\n",
    "                \n",
    "                selected_eigenvectors = np.hstack((\n",
    "                    eigenvectors_sorted_top[:, :k_top],\n",
    "                    eigenvectors_sorted_top[:, -k_bot:]\n",
    "                ))\n",
    "                \n",
    "                X_projected = np.dot(X_scaled, selected_eigenvectors)\n",
    "                X_reconstructed = np.dot(X_projected, selected_eigenvectors.T)\n",
    "                X_reconstructed = scaler.inverse_transform(X_reconstructed)\n",
    "                \n",
    "                l2_ext = np.linalg.norm(X - X_reconstructed)\n",
    "                acc_ext = list()\n",
    "\n",
    "                for i in range(X.shape[0]):\n",
    "                    acc_ext.append(test_model(vector_to_model(X_reconstructed[i]), dataloader))\n",
    "\n",
    "                df_eigen.loc[len(df_eigen)] = (\n",
    "                    zoo_label,\n",
    "                    k,\n",
    "                    explained_variance_actual,\n",
    "                    'Extremes',\n",
    "                    l2_ext,\n",
    "                    np.mean(acc_ext)\n",
    "                )\n",
    "\n",
    "    # Save results\n",
    "    df_eigen.to_csv(DF_EIGEN_PATH)\n",
    "    \n",
    "df_eigen_small = pd.read_csv(DF_EIGEN_PATH, index_col=0)\n",
    "df_eigen_small.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c837a612",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 261240 characters]",
      "text/plain": [
       "<Figure size 1800x1200 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(2, 3, figsize=(18, 12))\n",
    "\n",
    "for z, zoo_label in enumerate(ZOO_PATHS.keys()):\n",
    "    sns.lineplot(df_eigen_small[df_eigen_small['Model zoo'] == zoo_label], x='Explained variance', y='Reconstruction error', hue='Selected eigenvalues', palette=COLOURS[:3], marker=\"o\", ax=axs[0][z])\n",
    "    sns.lineplot(df_eigen_small[df_eigen_small['Model zoo'] == zoo_label], x='Explained variance', y='Average accuracy', hue='Selected eigenvalues', palette=COLOURS[:3], marker=\"o\", ax=axs[1][z])\n",
    "\n",
    "    axs[0][z].set(title=f'{zoo_label}\\nReconstruction error (L2)', xlim=(0,1.05))\n",
    "    axs[1][z].set(title=f'{zoo_label}\\nReconstructed models average accuracy', xlim=(0,1.05))\n",
    "\n",
    "plt.suptitle('Comparing top vs bottom eigenvalue decomposition')\n",
    "plt.subplots_adjust(hspace=0.25)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "4d7e3872",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Model zoo</th>\n",
       "      <th>k</th>\n",
       "      <th>Explained variance</th>\n",
       "      <th>Selected eigenvalues</th>\n",
       "      <th>Reconstruction error</th>\n",
       "      <th>Average accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10074</td>\n",
       "      <td>0.100038</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1876.993939</td>\n",
       "      <td>0.185673</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>2</td>\n",
       "      <td>0.106798</td>\n",
       "      <td>Top</td>\n",
       "      <td>1873.017565</td>\n",
       "      <td>0.227617</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10426</td>\n",
       "      <td>0.200019</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1769.319689</td>\n",
       "      <td>0.187757</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>14</td>\n",
       "      <td>0.203883</td>\n",
       "      <td>Top</td>\n",
       "      <td>1766.370188</td>\n",
       "      <td>0.314731</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10076</td>\n",
       "      <td>0.206836</td>\n",
       "      <td>Extremes</td>\n",
       "      <td>1765.586348</td>\n",
       "      <td>0.233233</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  Model zoo      k  Explained variance Selected eigenvalues  \\\n",
       "0      SVHN  10074            0.100038               Bottom   \n",
       "1      SVHN      2            0.106798                  Top   \n",
       "2      SVHN  10426            0.200019               Bottom   \n",
       "3      SVHN     14            0.203883                  Top   \n",
       "4      SVHN  10076            0.206836             Extremes   \n",
       "\n",
       "   Reconstruction error  Average accuracy  \n",
       "0           1876.993939          0.185673  \n",
       "1           1873.017565          0.227617  \n",
       "2           1769.319689          0.187757  \n",
       "3           1766.370188          0.314731  \n",
       "4           1765.586348          0.233233  "
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DF_EIGEN_LARGE_PATH = EXPERIMENT_PATH / 'df_eigen_large.csv'\n",
    "\n",
    "if not DF_EIGEN_LARGE_PATH.is_file() or RERUN:\n",
    "    # Create dataframe\n",
    "    df_eigen = pd.DataFrame(columns=['Model zoo', 'k', 'Explained variance', 'Selected eigenvalues', 'Reconstruction error', 'Average accuracy'])\n",
    "    \n",
    "    # Running experiments for every zoo\n",
    "    for zoo_label, zoo_path in tqdm(ZOO_PATHS.items(), desc='Iterating through the 3 model zoos', unit='zoo', total=3):\n",
    "        zoo_ds = torch.load(zoo_path / 'dataset_train.pt')\n",
    "        zoo_dl = torch.utils.data.DataLoader(zoo_ds, batch_size=64, shuffle=False)\n",
    "        \n",
    "        dataloader = torch.utils.data.DataLoader(TESTSETS[zoo_label], batch_size=128, shuffle=False)\n",
    "    \n",
    "        X = list()\n",
    "\n",
    "        # Run experiments\n",
    "        for bdx, data in enumerate(zoo_dl):\n",
    "            x, _, p, _ = data\n",
    "\n",
    "            for idx in range(x.shape[0]):\n",
    "                # Get original model\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",
    "\n",
    "                X.append(model_to_vector(orig_model))\n",
    "\n",
    "        X = np.vstack(X)\n",
    "\n",
    "        # Step 1: Center and scale the data using StandardScaler\n",
    "        scaler = StandardScaler()\n",
    "        X_scaled = scaler.fit_transform(X)\n",
    "\n",
    "        # Step 2: Compute the covariance matrix\n",
    "        cov_matrix = np.cov(X_scaled, rowvar=False)\n",
    "\n",
    "        # Step 3: Compute eigenvalues and eigenvectors\n",
    "        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)\n",
    "\n",
    "        # Step 4: Sort eigenvalues and eigenvectors\n",
    "        sorted_indices = np.argsort(eigenvalues)\n",
    "        eigenvalues_sorted = eigenvalues[sorted_indices]\n",
    "        eigenvectors_sorted = eigenvectors[:, sorted_indices]\n",
    "\n",
    "        # Step 5: Compute explained variance\n",
    "        total_variance = np.sum(eigenvalues_sorted)\n",
    "        explained_variance = eigenvalues_sorted / total_variance\n",
    "\n",
    "        # Step 6: Select k for specific explained variance levels\n",
    "        explained_variance_levels = [x/10 for x in range(1, 10)]\n",
    "        k_values = []\n",
    "        for level in explained_variance_levels:\n",
    "            cumulative_variance = np.cumsum(explained_variance)\n",
    "            k = np.searchsorted(cumulative_variance, level) + 1\n",
    "            k_values.append(k)\n",
    "        k_values.append(X.shape[1]) # We also want the full approximation as a baseline\n",
    "\n",
    "        # Step 7: Repeat the process using the top eigenvalues\n",
    "        sorted_indices_top = sorted_indices[::-1]\n",
    "        eigenvalues_sorted_top = eigenvalues[sorted_indices_top]\n",
    "        eigenvectors_sorted_top = eigenvectors[:, sorted_indices_top]\n",
    "\n",
    "        explained_variance_top = eigenvalues_sorted_top / total_variance\n",
    "        k_values_top = []\n",
    "        for level in explained_variance_levels:\n",
    "            cumulative_variance_top = np.cumsum(explained_variance_top)\n",
    "            k_top = np.searchsorted(cumulative_variance_top, level) + 1\n",
    "            k_values_top.append(k_top)\n",
    "        k_values_top.append(X.shape[1])\n",
    "\n",
    "        explained_variance_levels.append(1.0) # We add the label for full reconstruction\n",
    "\n",
    "        # Step 8: Project and reconstruct data for each k value, compute accuracy\n",
    "        for expl_var_idx, explained_variance_target in tqdm(enumerate(explained_variance_levels), desc=f'{zoo_label}: computing reconstructed accuracy', unit='experiment', total=len(explained_variance_levels)):\n",
    "            # Bottom\n",
    "            k = k_values[expl_var_idx]\n",
    "            explained_variance_actual = np.sum(explained_variance[:k])\n",
    "            selected_eigenvectors = eigenvectors_sorted[:, :k]\n",
    "            X_projected = np.dot(X_scaled, selected_eigenvectors)\n",
    "            X_reconstructed = np.dot(X_projected, selected_eigenvectors.T)\n",
    "            X_reconstructed = scaler.inverse_transform(X_reconstructed)\n",
    "\n",
    "            l2_bot = np.linalg.norm(X - X_reconstructed)\n",
    "            acc_bot = list()\n",
    "\n",
    "            for i in range(X.shape[0]):\n",
    "                acc_bot.append(test_model(vector_to_model(X_reconstructed[i]), dataloader))\n",
    "\n",
    "            df_eigen.loc[len(df_eigen)] = (\n",
    "                zoo_label,\n",
    "                k,\n",
    "                explained_variance_actual,\n",
    "                'Bottom',\n",
    "                l2_bot,\n",
    "                np.mean(acc_bot)\n",
    "            )\n",
    "\n",
    "            # Top\n",
    "            k = k_values_top[expl_var_idx]\n",
    "            explained_variance_actual = np.sum(explained_variance_top[:k])\n",
    "            selected_eigenvectors_top = eigenvectors_sorted_top[:, :k]\n",
    "            X_projected_top = np.dot(X_scaled, selected_eigenvectors_top)\n",
    "            X_reconstructed_top = np.dot(X_projected_top, selected_eigenvectors_top.T)\n",
    "            X_reconstructed_top = scaler.inverse_transform(X_reconstructed_top)\n",
    "\n",
    "            l2_top = np.linalg.norm(X - X_reconstructed_top)\n",
    "            acc_top = list()\n",
    "\n",
    "            for i in range(X.shape[0]):\n",
    "                acc_top.append(test_model(vector_to_model(X_reconstructed_top[i]), dataloader))\n",
    "\n",
    "            df_eigen.loc[len(df_eigen)] = (\n",
    "                zoo_label,\n",
    "                k,\n",
    "                explained_variance_actual,\n",
    "                'Top',\n",
    "                l2_top,\n",
    "                np.mean(acc_top)\n",
    "            )\n",
    "            \n",
    "            # Extremes\n",
    "            if expl_var_idx == 9:\n",
    "                df_eigen.loc[len(df_eigen)] = (\n",
    "                    zoo_label,\n",
    "                    k,\n",
    "                    explained_variance_actual,\n",
    "                    'Extremes',\n",
    "                    l2_top,\n",
    "                    np.mean(acc_top)\n",
    "                )\n",
    "            elif expl_var_idx % 2 == 1:\n",
    "                k_bot = k_values[expl_var_idx // 2]\n",
    "                k_top = k_values_top[expl_var_idx // 2]\n",
    "                k = k_bot + k_top\n",
    "                \n",
    "                explained_variance_actual = np.sum(explained_variance[:k_bot]) + np.sum(explained_variance_top[:k_top])\n",
    "                \n",
    "                selected_eigenvectors = np.hstack((\n",
    "                    eigenvectors_sorted_top[:, :k_top],\n",
    "                    eigenvectors_sorted_top[:, -k_bot:]\n",
    "                ))\n",
    "                \n",
    "                X_projected = np.dot(X_scaled, selected_eigenvectors)\n",
    "                X_reconstructed = np.dot(X_projected, selected_eigenvectors.T)\n",
    "                X_reconstructed = scaler.inverse_transform(X_reconstructed)\n",
    "                \n",
    "                l2_ext = np.linalg.norm(X - X_reconstructed)\n",
    "                acc_ext = list()\n",
    "\n",
    "                for i in range(X.shape[0]):\n",
    "                    acc_ext.append(test_model(vector_to_model(X_reconstructed[i]), dataloader))\n",
    "\n",
    "                df_eigen.loc[len(df_eigen)] = (\n",
    "                    zoo_label,\n",
    "                    k,\n",
    "                    explained_variance_actual,\n",
    "                    'Extremes',\n",
    "                    l2_ext,\n",
    "                    np.mean(acc_ext)\n",
    "                )\n",
    "\n",
    "    # Save results\n",
    "    df_eigen.to_csv(DF_EIGEN_LARGE_PATH)\n",
    "    \n",
    "df_eigen_large = pd.read_csv(DF_EIGEN_LARGE_PATH, index_col=0)\n",
    "df_eigen_large.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "906c1e93",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Model zoo</th>\n",
       "      <th>k</th>\n",
       "      <th>Explained variance</th>\n",
       "      <th>Selected eigenvalues</th>\n",
       "      <th>Reconstruction error</th>\n",
       "      <th>Average accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10074</td>\n",
       "      <td>0.100038</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1876.993939</td>\n",
       "      <td>0.185673</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>2</td>\n",
       "      <td>0.106798</td>\n",
       "      <td>Top</td>\n",
       "      <td>1873.017565</td>\n",
       "      <td>0.227617</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10426</td>\n",
       "      <td>0.200019</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1769.319689</td>\n",
       "      <td>0.187757</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>14</td>\n",
       "      <td>0.203883</td>\n",
       "      <td>Top</td>\n",
       "      <td>1766.370188</td>\n",
       "      <td>0.314731</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10076</td>\n",
       "      <td>0.206836</td>\n",
       "      <td>Extremes</td>\n",
       "      <td>1765.586348</td>\n",
       "      <td>0.233233</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10623</td>\n",
       "      <td>0.300463</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1653.045657</td>\n",
       "      <td>0.209710</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>31</td>\n",
       "      <td>0.302037</td>\n",
       "      <td>Top</td>\n",
       "      <td>1653.172213</td>\n",
       "      <td>0.313871</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10716</td>\n",
       "      <td>0.401328</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1528.294671</td>\n",
       "      <td>0.229893</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>54</td>\n",
       "      <td>0.400561</td>\n",
       "      <td>Top</td>\n",
       "      <td>1531.713595</td>\n",
       "      <td>0.311439</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10440</td>\n",
       "      <td>0.403902</td>\n",
       "      <td>Extremes</td>\n",
       "      <td>1528.577742</td>\n",
       "      <td>0.350869</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10766</td>\n",
       "      <td>0.500124</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1397.576173</td>\n",
       "      <td>0.226096</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>88</td>\n",
       "      <td>0.502280</td>\n",
       "      <td>Top</td>\n",
       "      <td>1396.907272</td>\n",
       "      <td>0.301846</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10800</td>\n",
       "      <td>0.603201</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1246.336618</td>\n",
       "      <td>0.221846</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>138</td>\n",
       "      <td>0.600288</td>\n",
       "      <td>Top</td>\n",
       "      <td>1253.839344</td>\n",
       "      <td>0.309871</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10654</td>\n",
       "      <td>0.602500</td>\n",
       "      <td>Extremes</td>\n",
       "      <td>1245.605241</td>\n",
       "      <td>0.551239</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10823</td>\n",
       "      <td>0.702901</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>1077.681406</td>\n",
       "      <td>0.224772</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>231</td>\n",
       "      <td>0.700217</td>\n",
       "      <td>Top</td>\n",
       "      <td>1085.691830</td>\n",
       "      <td>0.317173</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10840</td>\n",
       "      <td>0.802892</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>875.968563</td>\n",
       "      <td>0.227462</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>428</td>\n",
       "      <td>0.800373</td>\n",
       "      <td>Top</td>\n",
       "      <td>884.283947</td>\n",
       "      <td>0.333168</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10770</td>\n",
       "      <td>0.801889</td>\n",
       "      <td>Extremes</td>\n",
       "      <td>876.256382</td>\n",
       "      <td>0.657859</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10852</td>\n",
       "      <td>0.910597</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>573.405058</td>\n",
       "      <td>0.284061</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>780</td>\n",
       "      <td>0.900168</td>\n",
       "      <td>Top</td>\n",
       "      <td>624.571467</td>\n",
       "      <td>0.421457</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10853</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Bottom</td>\n",
       "      <td>0.000071</td>\n",
       "      <td>0.735733</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10853</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Top</td>\n",
       "      <td>0.000071</td>\n",
       "      <td>0.735733</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>SVHN</td>\n",
       "      <td>10853</td>\n",
       "      <td>1.000000</td>\n",
       "      <td>Extremes</td>\n",
       "      <td>0.000071</td>\n",
       "      <td>0.735733</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   Model zoo      k  Explained variance Selected eigenvalues  \\\n",
       "0       SVHN  10074            0.100038               Bottom   \n",
       "1       SVHN      2            0.106798                  Top   \n",
       "2       SVHN  10426            0.200019               Bottom   \n",
       "3       SVHN     14            0.203883                  Top   \n",
       "4       SVHN  10076            0.206836             Extremes   \n",
       "5       SVHN  10623            0.300463               Bottom   \n",
       "6       SVHN     31            0.302037                  Top   \n",
       "7       SVHN  10716            0.401328               Bottom   \n",
       "8       SVHN     54            0.400561                  Top   \n",
       "9       SVHN  10440            0.403902             Extremes   \n",
       "10      SVHN  10766            0.500124               Bottom   \n",
       "11      SVHN     88            0.502280                  Top   \n",
       "12      SVHN  10800            0.603201               Bottom   \n",
       "13      SVHN    138            0.600288                  Top   \n",
       "14      SVHN  10654            0.602500             Extremes   \n",
       "15      SVHN  10823            0.702901               Bottom   \n",
       "16      SVHN    231            0.700217                  Top   \n",
       "17      SVHN  10840            0.802892               Bottom   \n",
       "18      SVHN    428            0.800373                  Top   \n",
       "19      SVHN  10770            0.801889             Extremes   \n",
       "20      SVHN  10852            0.910597               Bottom   \n",
       "21      SVHN    780            0.900168                  Top   \n",
       "22      SVHN  10853            1.000000               Bottom   \n",
       "23      SVHN  10853            1.000000                  Top   \n",
       "24      SVHN  10853            1.000000             Extremes   \n",
       "\n",
       "    Reconstruction error  Average accuracy  \n",
       "0            1876.993939          0.185673  \n",
       "1            1873.017565          0.227617  \n",
       "2            1769.319689          0.187757  \n",
       "3            1766.370188          0.314731  \n",
       "4            1765.586348          0.233233  \n",
       "5            1653.045657          0.209710  \n",
       "6            1653.172213          0.313871  \n",
       "7            1528.294671          0.229893  \n",
       "8            1531.713595          0.311439  \n",
       "9            1528.577742          0.350869  \n",
       "10           1397.576173          0.226096  \n",
       "11           1396.907272          0.301846  \n",
       "12           1246.336618          0.221846  \n",
       "13           1253.839344          0.309871  \n",
       "14           1245.605241          0.551239  \n",
       "15           1077.681406          0.224772  \n",
       "16           1085.691830          0.317173  \n",
       "17            875.968563          0.227462  \n",
       "18            884.283947          0.333168  \n",
       "19            876.256382          0.657859  \n",
       "20            573.405058          0.284061  \n",
       "21            624.571467          0.421457  \n",
       "22              0.000071          0.735733  \n",
       "23              0.000071          0.735733  \n",
       "24              0.000071          0.735733  "
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_eigen_large[df_eigen_large['Model zoo'] == 'SVHN']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "3eb899d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 265332 characters]",
      "text/plain": [
       "<Figure size 1800x1200 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(2, 3, figsize=(18, 12))\n",
    "\n",
    "for z, zoo_label in enumerate(ZOO_PATHS.keys()):\n",
    "    sns.lineplot(df_eigen_large[df_eigen_large['Model zoo'] == zoo_label], x='Explained variance', y='Reconstruction error', hue='Selected eigenvalues', palette=COLOURS[:3], marker=\"o\", ax=axs[0][z])\n",
    "    sns.lineplot(df_eigen_large[df_eigen_large['Model zoo'] == zoo_label], x='Explained variance', y='Average accuracy', hue='Selected eigenvalues', palette=COLOURS[:3], marker=\"o\", ax=axs[1][z])\n",
    "\n",
    "    axs[0][z].set(title=f'{zoo_label}\\nReconstruction error (L2)', xlim=(0,1.05))\n",
    "    axs[1][z].set(title=f'{zoo_label}\\nReconstructed models average accuracy', xlim=(0,1.05))\n",
    "\n",
    "plt.suptitle('Comparing top vs bottom eigenvalue decomposition (more models)')\n",
    "plt.subplots_adjust(hspace=0.25)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a545d34",
   "metadata": {},
   "source": [
    "#### Generating a Figure for the paper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "bf035b01",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 71936 characters]",
      "text/plain": [
       "<Figure size 1400x300 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_, axs = plt.subplots(1, 3, figsize=(14, 3))\n",
    "\n",
    "for z, zoo_label in enumerate(ZOO_PATHS.keys()):\n",
    "    sns.lineplot(df_eigen_large[df_eigen_large['Model zoo'] == zoo_label], x='Explained variance', y='Average accuracy', hue='Selected eigenvalues', palette=COLOURS[:3], marker=\"o\", ax=axs[z])\n",
    "    axs[z].set(title=f\"{zoo_label}\", xlim=(0,1.05), ylim=(0., 0.8))\n",
    "    if z != 1:\n",
    "        axs[z].get_legend().remove()\n",
    "    \n",
    "plt.savefig('figures/eigendecomposition_reconstruction.png', dpi=300, bbox_inches='tight')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
