notebooks/llcs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Learning Coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "if not \"AWS_ACCESS_KEY_ID\" in os.environ or not \"AWS_SECRET_ACCESS_KEY\" in os.environ:\n",
    "    raise Exception(\"AWS_ACCESS_KEY_ID or AWS_SECRET_ACCESS_KEY not found in environment variables. Please set them in .env file.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from pprint import pp\n",
    "from pathlib import Path\n",
    "from typing import Optional, Iterable, List, Tuple, Dict, Union, Callable\n",
    "import warnings\n",
    "\n",
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import torch \n",
    "from torch import nn\n",
    "import wandb\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "\n",
    "from infra.utils.iterables import filter_objs\n",
    "\n",
    "# del sys.modules['icl.analysis.utils']\n",
    "from icl.analysis.utils import wandb_runs_to_df\n",
    "from icl.constants import SWEEPS, FIGURES, ANALYSIS\n",
    "from icl.analysis.utils import get_unique_run\n",
    "\n",
    "\n",
    "api= wandb.Api()\n",
    "\n",
    "WANDB_ENTITY = os.getenv(\"WANDB_ENTITY\")\n",
    "\n",
    "from icl.constants import FIGURES, ANALYSIS, DEVICE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set-up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.analysis.utils import get_sweep_configs\n",
    "\n",
    "NUM_LAYERS = 2\n",
    "NUM_HEADS = 4\n",
    "LR = 0.003\n",
    "\n",
    "TRAINING_SWEEP_ID = \"8rxvu833\"\n",
    "TRAINING_SWEEP_FILENAME = \"training-runs/L2H4Minf.yaml\"\n",
    "\n",
    "filters = {\"task_config\": {\"num_layers\": NUM_LAYERS, \"num_heads\": NUM_HEADS}, \"optimizer_config\": {\"lr\": LR}}  # TODO: Where are the H=2 runs?\n",
    "configs = list(get_sweep_configs(SWEEPS / TRAINING_SWEEP_FILENAME, **filters))\n",
    "\n",
    "print(f\"Found {len(configs)} runs.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Figure out which checkpoints are available\n",
    "\n",
    "checkpointers = [config.checkpointer_config.factory() for config in tqdm(configs, desc=\"Reading checkpoints\")]\n",
    "\n",
    "for checkpointer in tqdm(checkpointers, desc=\"Loading checkpoints\"):\n",
    "    print(f\"Found {len(checkpointer.file_ids)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MS = [config.task_config.num_tasks for config in configs] # [1, 4, 64, 2**10, 2**20]\n",
    "STEPS = checkpointer.file_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Calibrating LLCs\n",
    "\n",
    "Let's look at finding the right hyperparameters for learning coefficient estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Temperature, Learning Rate, and Localization Strength"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FORCE_REEVAL = True\n",
    "\n",
    "EVAL_EPS_TEMP_GAMMA_SWEEP_NAME = \"eps-temp-gamma\"\n",
    "EVAL_EPS_TEMP_GAMMA_SWEEP_ID = \"20ygh3o1\"\n",
    "\n",
    "if not os.path.exists(ANALYSIS / EVAL_EPS_TEMP_GAMMA_SWEEP_NAME) or FORCE_REEVAL:\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        eval_eps_temp_gamma_sweep = api.sweep(f\"{WANDB_ENTITY}/icl-llc/{EVAL_EPS_TEMP_GAMMA_SWEEP_ID}\")\n",
    "        eval_eps_temp_gamma_runs = list(filter_objs([r for r in eval_eps_temp_gamma_sweep.runs], config=filters))\n",
    "        eval_eps_temp_gamma_df = wandb_runs_to_df(eval_eps_temp_gamma_runs)\n",
    "    \n",
    "    eval_eps_temp_gamma_df.to_pickle(ANALYSIS / EVAL_EPS_TEMP_GAMMA_SWEEP_NAME)\n",
    "\n",
    "else:\n",
    "    eval_eps_temp_gamma_df = pd.read_pickle(ANALYSIS / EVAL_EPS_TEMP_GAMMA_SWEEP_NAME)\n",
    "\n",
    "# columns_to_convert = ['loss/mean/mean', 'loss/std/mean', 'llc/mean/mean', 'llc/std/mean', 'wbic/mean/mean', 'wbic/std/mean']  # Replace with your actual column names\n",
    "\n",
    "# for col in columns_to_convert:\n",
    "#     eval_eps_temp_gamma_df[col] = pd.to_numeric(eval_eps_temp_gamma_df[col], errors='coerce')\n",
    "\n",
    "# eval_eps_temp_gamma_df.reset_index(inplace=True)\n",
    "eval_eps_temp_gamma_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_eps_temp_gamma_df[\"epsilon\"] = eval_eps_temp_gamma_df[\"sampler_config/noise_scale\"]\n",
    "eval_eps_temp_gamma_df[\"beta\"] = eval_eps_temp_gamma_df[\"sampler_config/gradient_scale\"] * 2 / (eval_eps_temp_gamma_df[\"sampler_config/noise_scale\"] *eval_eps_temp_gamma_df[\"task_config/num_tasks\"])\n",
    "eval_eps_temp_gamma_df[\"gamma\"] = eval_eps_temp_gamma_df[\"sampler_config/localization_scale\"] * 2 / eval_eps_temp_gamma_df[\"sampler_config/noise_scale\"]\n",
    "# eval_eps_temp_gamma_df['version'] = [f\"$t={row['step']}, LN={'T' if row['task_config/layer_norm'] else 'F'}$\" for i, row in eval_eps_temp_gamma_df.iterrows()]\n",
    "eval_eps_temp_gamma_df['version'] = [f\"$t={row['step']}$\" for i, row in eval_eps_temp_gamma_df.iterrows()]\n",
    "# eval_eps_temp_gamma_df = eval_eps_temp_gamma_df[eval_eps_temp_gamma_df[\"sampler_config/noise_scale\"] < 0.001]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_eps_temp_gamma_df.columns.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = eval_eps_temp_gamma_df.loc[eval_eps_temp_gamma_df[\"task_config/layer_norm\"] == True]\n",
    "# df = df.loc[(df[\"sampler_config/noise_scale\"] > 0.0001)]\n",
    "# df = df.loc[(df[\"sampler_config/gradient_scale\"] > 0.0001)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig = plt.figure(figsize=(HEIGHT * .5, WIDTH * .5))\n",
    "\n",
    "# Assuming eval_dataset_size_df is your DataFrame\n",
    "data = df.loc[df.step == df.step.max()]\n",
    "\n",
    "# Step 2: Create FacetGrid\n",
    "# Replace 'epsilon', 'gamma', 'temperature' with your actual column names\n",
    "g = sns.FacetGrid(data=data, col='sampler_config/noise_scale', row='sampler_config/gradient_scale', hue='sampler_config/localization_scale', palette='viridis', sharey=True)\n",
    "\n",
    "def plot_cell(x, y, data=None, **kwargs):\n",
    "    ordering = np.argsort(data[x])\n",
    "    plt.plot(data[x].values[ordering], data[y].values[ordering], '-', alpha=0.8, **kwargs)\n",
    "    # plt.fill_between(data[x].values[ordering], data[y].values[ordering] - data['loss/std/mean'].values[ordering], data[y].values[ordering] + data['loss/std/mean'].values[ordering], alpha=0.05, color=kwargs['color'])\n",
    "\n",
    "# Step 3: Plot the Data\n",
    "g = g.map_dataframe(plot_cell, '_step', 'loss/mean/mean').add_legend()\n",
    "\n",
    "# Adjust layout and aesthetics as needed\n",
    "g.set_axis_labels(r'$\\tau$', '$E[L_n(w)]$')\n",
    "g.set_titles(r'$\\tilde \\eta={row_name}$, $\\varepsilon={col_name}$')\n",
    "g.legend.set_title(r'$\\tilde \\gamma$')\n",
    "g.fig.subplots_adjust(top=0.9, right=0.95)\n",
    "g.fig.suptitle(fr'$E[L_n(w)]$ vs. $M$ for Various $\\epsilon$, $\\gamma$, and $\\beta$ ({version})')\n",
    "\n",
    "# log y \n",
    "g.set(yscale='log')\n",
    "\n",
    "g.fig.savefig(FIGURES / f\"llc-vs-eps-gamma-beta-1.pdf\", bbox_inches='tight')\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = df.loc[df['_step'] == df['_step'].max()]\n",
    "\n",
    "# Assuming eval_dataset_size_df is your DataFrame\n",
    "data.sort_values(by='sampler_config/gradient_scale', inplace=True)\n",
    "\n",
    "# Step 1: Filter Data\n",
    "max_draw_df = data[data['_step'] == data['_step'].max()].sort_values(by='task_config/num_tasks')\n",
    "\n",
    "# Step 2: Create FacetGrid\n",
    "# Replace 'epsilon', 'gamma', 'temperature' with your actual column names\n",
    "g = sns.FacetGrid(data=max_draw_df, col='sampler_config/gradient_scale', row='sampler_config/noise_scale', hue=\"sampler_config/localization_scale\", palette='viridis', sharey=False)\n",
    "\n",
    "def plot_cell(x, y, data=None, **kwargs):\n",
    "    plt.xscale('log', base=10)\n",
    "    ordering = np.argsort(data[x])\n",
    "    plt.plot(data[x].values[ordering], data[y].values[ordering], 'o-', alpha=1, **kwargs)\n",
    "    # plt.fill_between(data[x], data[y] - data['loss/std/mean'], data[y] + data['loss/std/mean'], alpha=0.5, color=kwargs['color'])\n",
    "\n",
    "\n",
    "# Step 3: Plot the Data\n",
    "g = g.map_dataframe(plot_cell, 'step', 'llc/mean/mean', data=data).add_legend()\n",
    "\n",
    "# Adjust layout and aesthetics as needed\n",
    "g.set_axis_labels(r'$\\tilde \\gamma$', '$E[L_n(w)]$')\n",
    "g.set_titles(r'$\\tilde\\eta={col_name}$, $\\varepsilon={row_name}$')\n",
    "g.legend.set_title(r'$\\tilde \\gamma$')\n",
    "g.fig.subplots_adjust(top=0.9, right=0.9)\n",
    "g.fig.suptitle(fr'$E[L_n(w)]$ vs. $M$ for Various $\\epsilon$, $\\gamma$, and $\\beta$')\n",
    "\n",
    "# Show the plot\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation Dataset Size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FORCE_REEVAL = True\n",
    "\n",
    "EVAL_DATASET_SIZE_SWEEP_NAME = \"eval-dataset-size\"\n",
    "EVAL_DATASET_SIZE_SWEEP_ID = \"2s8wbx80\"\n",
    "\n",
    "if not os.path.exists(ANALYSIS / EVAL_DATASET_SIZE_SWEEP_NAME) or FORCE_REEVAL:\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        eval_dataset_size_sweep = api.sweep(f\"{WANDB_ENTITY}/icl-llc/{EVAL_DATASET_SIZE_SWEEP_ID}\")\n",
    "        eval_dataset_size_runs = list(filter_objs([r for r in eval_dataset_size_sweep.runs], config=filters))\n",
    "        eval_dataset_size_df = wandb_runs_to_df(eval_dataset_size_runs)\n",
    "    \n",
    "    eval_dataset_size_df.to_pickle(ANALYSIS / EVAL_DATASET_SIZE_SWEEP_NAME)\n",
    "\n",
    "else:\n",
    "    eval_dataset_size_df = pd.read_pickle(ANALYSIS / EVAL_DATASET_SIZE_SWEEP_NAME)\n",
    "\n",
    "# columns_to_convert = ['loss/mean', 'loss/std', 'llc/mean', 'llc/std', 'wbic/mean', 'wbic/std']  # Replace with your actual column names\n",
    "\n",
    "# # Convert each column to numeric, coercing \"NaN\" strings to actual NaN values\n",
    "# for col in columns_to_convert:\n",
    "#     eval_dataset_size_df[col] = pd.to_numeric(eval_dataset_size_df[col], errors='coerce')\n",
    "\n",
    "eval_dataset_size_df.reset_index(inplace=True)\n",
    "\n",
    "\n",
    "eval_dataset_size_df['log_num_tasks'] = np.log2(eval_dataset_size_df['task_config/num_tasks'])\n",
    "# eval_dataset_size_df.dropna(inplace=True)\n",
    "eval_dataset_size_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_datasets_sizes = sorted(eval_dataset_size_df['sampler_config/eval_dataset_size'].unique())[:-1]\n",
    "eval_datasets_sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_dataset_size_df['llc/mean/mean'], eval_dataset_size_df['step'].unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FULL_WIDTH = 6.5\n",
    "\n",
    "FULL_HEIGHT = FULL_WIDTH / golden_ratio\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, figsize=(FULL_HEIGHT, FULL_WIDTH))\n",
    "\n",
    "ax = axes[0]\n",
    "\n",
    "final_step = eval_dataset_size_df['_step'].max()\n",
    "final_step_df = eval_dataset_size_df[eval_dataset_size_df['_step'] == final_step-1]\n",
    "\n",
    "colors = sns.color_palette('viridis', len(eval_dataset_size_df['step'].unique()))\n",
    "steps = sorted(eval_dataset_size_df['step'].unique())\n",
    "\n",
    "for i, step in enumerate(steps):\n",
    "    _df = final_step_df[final_step_df['step'] == step]\n",
    "    # _df = _df.dropna(subset=['llc/mean/mean'])\n",
    "    # _df.sort_values(by='sampler_config/eval_dataset_size', inplace=True)\n",
    "    # sns.lineplot(data=_df, x='sampler_config/eval_dataset_size', y='llc/mean/mean', ax=ax)\n",
    "\n",
    "    ax.plot(_df['sampler_config/eval_dataset_size'], _df['llc/mean/mean'], color=colors[i], alpha=0.5)\n",
    "    ax.fill_between(_df['sampler_config/eval_dataset_size'], _df['llc/mean/mean'] - _df['llc/std/mean'], _df['llc/mean/mean'] + _df['llc/std/mean'], alpha=0.05, color=colors[i])\n",
    "\n",
    "ax.legend(labels=[f\"$t={step}$\" if i else \"_\" for step in steps for i in [1, 0]], loc='upper right', bbox_to_anchor=(1.15, .85))        \n",
    "\n",
    "ax.set_xscale('log', base=2)\n",
    "ax.set_xlabel(r'$\\mu$')\n",
    "ax.set_ylabel(r'$\\hat\\lambda(w_t)$')\n",
    "\n",
    "ax.set_title(r'(a) $\\hat\\lambda(w_t)$ Dependence on SGLD Dataset Size for Selected Checkpoints')\n",
    "\n",
    "ax = axes[1]\n",
    "colors = sns.color_palette('viridis', len(eval_datasets_sizes) //2 + 1)\n",
    "\n",
    "for i, size in enumerate(eval_datasets_sizes[::2] + [eval_datasets_sizes[-1]]):\n",
    "    _df = final_step_df[final_step_df['sampler_config/eval_dataset_size'] == size]\n",
    "    _df.sort_values(by='step', inplace=True)\n",
    "    # _df = _df.dropna(subset=['llc/mean/mean'])\n",
    "    # _df.sort_values(by='step', inplace=True)\n",
    "    # sns.lineplot(data=_df, x='step', y='llc/mean/mean', ax=ax)\n",
    "\n",
    "    ax.plot(_df['step'], _df['llc/mean/mean'], color=colors[i], alpha=0.5)\n",
    "    ax.fill_between(_df['step'], _df['llc/mean/mean'] - _df['llc/std/mean'], _df['llc/mean/mean'] + _df['llc/std/mean'], alpha=0.05, color=colors[i])\n",
    "\n",
    "ax.legend(labels=[f\"$\\mu={size}$\" if i else \"_\" for size in eval_datasets_sizes for i in [1, 0]], loc='upper right', bbox_to_anchor=(1.15, .9))        \n",
    "ax.set_xscale('log')\n",
    "ax.set_xlabel(r'$t$')\n",
    "ax.set_ylabel(r'$\\hat\\lambda(w_t)$')\n",
    "ax.set_title(r'(b) $\\hat\\lambda(w_t)$ Dependence on SGLD Checkpoint for Selected Dataset Sizes')\n",
    "\n",
    "plt.tight_layout()\n",
    "\n",
    "# ax.plot(_df[''])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.figures.colors import PRIMARY, SECONDARY\n",
    "\n",
    "fig, axes = plt.subplots(2, len(eval_datasets_sizes) // 2, figsize=(2 * FULL_WIDTH, FULL_HEIGHT))\n",
    "\n",
    "fig.subplots_adjust(hspace=0.3, wspace=0.1)  # Adjust this value as needed\n",
    "\n",
    "LOSS_MEAN, LOSS_STD = 'llc/mean/mean', 'llc/mean/std'\n",
    "max_loss_over_time = eval_dataset_size_df[LOSS_MEAN].max() \n",
    "max_loss_at_end = eval_dataset_size_df[LOSS_MEAN].max() \n",
    "\n",
    "colors = sns.color_palette('viridis', len(steps))\n",
    "\n",
    "for i, (ax, eval_batch_size) in enumerate(zip(axes.flatten(), eval_datasets_sizes)):\n",
    "    _df = eval_dataset_size_df.loc[eval_dataset_size_df['sampler_config/eval_dataset_size'] == eval_batch_size].sort_values(by=['step', '_step'])\n",
    "\n",
    "    for j, step in enumerate(steps):\n",
    "        sns.lineplot(data=_df.loc[_df.step == step], x=\"_step\", y=LOSS_MEAN, color=colors[j], ax=ax, alpha=0.8)\n",
    "\n",
    "    for step, color in zip(steps, colors):\n",
    "        _step_df = _df.loc[_df['step'] == step]\n",
    "        # for chain in range(8):\n",
    "        #     ax.plot(_step_df['_step'], _step_df[f\"llc/mean/{chain}\"], alpha=0.5, color=color, linewidth=0.5)\n",
    "        ax.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.05)  \n",
    "        # ax.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.1)\n",
    "        \n",
    "    ax.set_title(f\"$\\mu = {eval_batch_size}$\")\n",
    "    if i < len(eval_datasets_sizes) // 2:\n",
    "        ax.set_xlabel('')\n",
    "    else:\n",
    "        ax.set_xlabel(r\"$\\tau$\")\n",
    "\n",
    "    # ax.legend().remove()\n",
    "    ax.set_ylim(0, max_loss_over_time * .7)\n",
    "    # ax.set_yscale('log')\n",
    "\n",
    "    # _final_df = _df.loc[_df['_step'] == _df['_step'].max()]\n",
    "\n",
    "    # ax2.plot(_final_df['task_config/num_tasks'], _final_df[LOSS_MEAN], alpha=1)\n",
    "    # ax2.fill_between(_final_df[\"task_config/num_tasks\"], _final_df[LOSS_MEAN] - _final_df[LOSS_STD], _final_df[LOSS_MEAN] + _final_df[LOSS_STD], alpha=0.5)\n",
    "\n",
    "    # ax2.set_title(fr\"$\\mu = {eval_batch_size}, t=500k$\")\n",
    "    # ax2.set_xlabel(r\"$M$\")\n",
    "    # ax2.set_xscale('log', base=2)\n",
    "\n",
    "    # ax2.set_ylim(0, max_loss_at_end * 1.25)\n",
    "\n",
    "    if i != 0 and i != len(eval_datasets_sizes) // 2:\n",
    "        ax.set_yticklabels([])\n",
    "        ax.set_ylabel(\"\")\n",
    "        # ax2.set_yticklabels([])\n",
    "        # ax2.set_ylabel(\"\")\n",
    "\n",
    "axes[0, 0].set_ylabel(r\"$\\hat\\lambda_t$\")\n",
    "axes[1, 0].set_ylabel(r\"$\\hat\\lambda_t$\")\n",
    "\n",
    "\n",
    "# Titles above rows (shift far to the right)\n",
    "# axes[0, 0].set_title(r\"(a) $\\hat\\lambda_t$ Dependence on $\\mu$ for SGLD\")\n",
    "# axes[1, 0].set_title(r\"(b) $\\hat\\lambda_t$ Dependence on $\\mu$ for SGLD-MC\")\n",
    "plt.tight_layout()\n",
    "\n",
    "ax= axes[-1, -1]       \n",
    "handles = [plt.Line2D([0, 0], [0, 0], color=colors[j]) for j in range(len(steps))]\n",
    "labels = [f\"$t={step}$\" for step in steps]\n",
    "ax.legend(handles=handles, labels=labels, title=r\"Step $t$\", loc='upper left', bbox_to_anchor=(1.1,1.75))\n",
    "\n",
    "fig.set_facecolor('white')\n",
    "fig.savefig(FIGURES / \"llc-fixed-dataset.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot heatmap of llc/mean/mean vs. step and eval_dataset_size\n",
    "\n",
    "# Reset index\n",
    "final_step_df.reset_index(inplace=True)\n",
    "\n",
    "# Create pivot table\n",
    "eval_dataset_size_pivot = final_step_df.pivot(index='step', columns='sampler_config/eval_dataset_size', values='llc/mean/mean')\n",
    "\n",
    "# Create heatmap\n",
    "fig, ax = plt.subplots(figsize=(FULL_WIDTH, FULL_HEIGHT))\n",
    "\n",
    "sns.heatmap(eval_dataset_size_pivot, ax=ax, cmap='viridis', cbar_kws={'label': r'$\\hat\\lambda(w_t)$'})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FULL_WIDTH = 6.5\n",
    "FULL_HEIGHT = FULL_WIDTH * golden_ratio\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(FULL_WIDTH, FULL_HEIGHT / 2))\n",
    "\n",
    "colors = sns.color_palette('viridis', len(eval_datasets_sizes))\n",
    "\n",
    "for i, size in enumerate(eval_datasets_sizes):\n",
    "    _df = final_step_df[final_step_df['sampler_config/eval_dataset_size'] == size]\n",
    "    _df.sort_values(by='step', inplace=True)\n",
    "    # _df = _df.dropna(subset=['llc/mean/mean'])\n",
    "    # _df.sort_values(by='step', inplace=True)\n",
    "    # sns.lineplot(data=_df, x='step', y='llc/mean/mean', ax=ax)\n",
    "\n",
    "    ax.plot(_df['step'], _df['llc/mean/mean'], color=colors[i], alpha=0.25)\n",
    "    ax.fill_between(_df['step'], _df['llc/mean/mean'] - _df['llc/std/mean'], _df['llc/mean/mean'] + _df['llc/std/mean'], alpha=0.025, color=colors[i])\n",
    "\n",
    "ax.legend(labels=[f\"$\\mu={size}$\" if i else \"_\" for size in eval_datasets_sizes for i in [1, 0]], loc='upper right', bbox_to_anchor=(1.25, 1.1))        \n",
    "ax.set_xscale('log')\n",
    "ax.set_xlabel(r'$t$')\n",
    "ax.set_ylabel(r'$\\hat\\lambda(w_t)$')\n",
    "ax.set_title(r'(a) $\\hat\\lambda(w_t)$ Dependence on Dataset Size')\n",
    "\n",
    "\n",
    "# ax.plot(_df[''])\n",
    "fig.savefig(FIGURES / \"llc-vs-dataset-size.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, len(eval_datasets_sizes), figsize=(FULL_HEIGHT, FULL_WIDTH))\n",
    "\n",
    "for i, (ax, size) in enumerate(zip(axes.flatten(), eval_datasets_sizes)):\n",
    "    _df = final_step_df[final_step_df['sampler_config/eval_dataset_size'] == size]\n",
    "    _df.sort_values(by='step', inplace=True)\n",
    "    # _df = _df.dropna(subset=['llc/mean/mean'])\n",
    "    # _df.sort_values(by='step', inplace=True)\n",
    "    # sns.lineplot(data=_df, x='step', y='llc/mean/mean', ax=ax)\n",
    "\n",
    "    ax.plot(_df['step'], _df['llc/mean/mean'], color=colors[i], alpha=0.5)\n",
    "    ax.fill_between(_df['step'], _df['llc/mean/mean'] - _df['llc/std/mean'], _df['llc/mean/mean'] + _df['llc/std/mean'], alpha=0.05, color=colors[i])\n",
    "\n",
    "ax.legend(labels=[f\"$\\mu={size}$\" if i else \"_\" for size in eval_datasets_sizes for i in [1, 0]], loc='upper right', bbox_to_anchor=(1.15, .9))        \n",
    "ax.set_xscale('log')\n",
    "ax.set_xlabel(r'$t$')\n",
    "ax.set_ylabel(r'$\\hat\\lambda(w_t)$')\n",
    "ax.set_title(r'(b) $\\hat\\lambda(w_t)$ Dependence on Dataset Size')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, len(eval_datasets_sizes), figsize=(5*len(eval_datasets_sizes), 10))\n",
    "\n",
    "fig.subplots_adjust(hspace=0.3, wspace=0.1)  # Adjust this value as needed\n",
    "\n",
    "max_loss_over_time = eval_dataset_size_df['llc/mean/mean'].max()\n",
    "max_loss_at_end = eval_dataset_size_df[eval_dataset_size_df['_step'] == eval_dataset_size_df['_step'].max()]['llc/mean/mean'].max()\n",
    "\n",
    "for i, (ax1, ax2, eval_dataset_size) in enumerate(zip(axes[0, :], axes[1, :], eval_datasets_sizes)):\n",
    "    _df = eval_dataset_size_df.loc[eval_dataset_size_df['sampler_config/eval_dataset_size'] == eval_dataset_size]\n",
    "\n",
    "    sns.lineplot(data=_df, x=\"_step\", y=\"llc/mean/mean\", hue=\"log_num_tasks\", palette=\"viridis\", ax=ax1, alpha=0.5)\n",
    "\n",
    "    ax1.set_title(f\"$\\mu = {eval_dataset_size}$\")\n",
    "    ax1.set_xlabel(r\"$\\tau$\")\n",
    "    ax1.legend().remove()\n",
    "    ax1.set_ylim(0, max_loss_over_time * 1.1)\n",
    "\n",
    "    _final_df = _df.loc[_df['_step'] == _df['_step'].max()]\n",
    "\n",
    "    ax2.plot(_final_df['task_config/num_tasks'], _final_df['llc/mean/mean'], alpha=1)\n",
    "    ax2.fill_between(_final_df[\"task_config/num_tasks\"], _final_df[\"llc/mean/mean\"] - _final_df[\"llc/mean/std\"], _final_df[\"llc/mean/mean\"] + _final_df[\"llc/mean/std\"], alpha=0.5)\n",
    "\n",
    "    ax2.set_title(fr\"$\\mu = {eval_dataset_size}, \\tau = T_\\mathrm{{SGLD}}$\")\n",
    "    ax2.set_xlabel(r\"$M$\")\n",
    "    ax2.set_xscale('log', base=2)\n",
    "\n",
    "    # ax2.set_ylim(0, max_loss_at_end * 1.25)\n",
    "\n",
    "    if i > 0:\n",
    "        ax1.set_yticklabels([])\n",
    "        ax1.set_ylabel(\"\")\n",
    "        ax2.set_yticklabels([])\n",
    "        ax2.set_ylabel(\"\")\n",
    "\n",
    "axes[0, 0].set_ylabel(r\"$E[L_\\mu(w_\\tau)]$\")\n",
    "axes[1, 0].set_ylabel(r\"$E[L_\\mu(w_\\tau)]$\")\n",
    "\n",
    "# Add a colorbar for the hues on the left side of the top row\n",
    "sm = plt.cm.ScalarMappable(cmap=\"viridis\", norm=plt.Normalize(vmin=_df[\"log_num_tasks\"].min(), vmax=_df[\"log_num_tasks\"].max()))\n",
    "sm._A = []  # Fake up the array of the scalar mappable\n",
    "cb_ax = fig.add_axes([0.92, 0.3, 0.0125, 0.35])  # Adjust these values as needed for positioning and size\n",
    "cbar = fig.colorbar(sm, cax=cb_ax)\n",
    "cbar.set_ticklabels([f\"${{{int(tick)}}}$\" for tick in cbar.get_ticks()])\n",
    "cbar.set_label('$\\log_2 M$')  # Update this label as needed\n",
    "\n",
    "fig.set_facecolor('white')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixed Evaluation Datasets\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FORCE_REEVAL = True\n",
    "\n",
    "EVAL_FIXED_BATCH_SWEEP_NAME = \"fixed-batch\"\n",
    "EVAL_FIXED_BATCH_SWEEP_ID = \"z54jluit\"\n",
    "\n",
    "if not os.path.exists(ANALYSIS / EVAL_FIXED_BATCH_SWEEP_NAME) or FORCE_REEVAL:\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        eval_fixed_batch_sweep = api.sweep(f\"{WANDB_ENTITY}/icl-llc/{EVAL_FIXED_BATCH_SWEEP_ID}\")\n",
    "        eval_fixed_batch_runs = list(filter_objs([r for r in eval_fixed_batch_sweep.runs], config=filters))\n",
    "        eval_fixed_batch_df = wandb_runs_to_df(eval_fixed_batch_runs)\n",
    "    \n",
    "    eval_fixed_batch_df.to_pickle(ANALYSIS / EVAL_FIXED_BATCH_SWEEP_NAME)\n",
    "\n",
    "else:\n",
    "    eval_fixed_batch_df = pd.read_pickle(ANALYSIS / EVAL_FIXED_BATCH_SWEEP_NAME)\n",
    "\n",
    "# columns_to_convert = ['loss/mean', 'loss/std', 'llc/mean', 'llc/std', 'wbic/mean', 'wbic/std']  # Replace with your actual column names\n",
    "\n",
    "# for col in columns_to_convert:\n",
    "#     eval_fixed_batch_df[col] = pd.to_numeric(eval_fixed_batch_df[col], errors='coerce')\n",
    "\n",
    "    \n",
    "eval_fixed_batch_df.reset_index(inplace=True)\n",
    "eval_fixed_batch_df['log_num_tasks'] = np.log2(eval_fixed_batch_df['task_config/num_tasks'])\n",
    "eval_fixed_batch_df = eval_fixed_batch_df.loc[eval_fixed_batch_df['task_config/layer_norm'] == True]\n",
    "eval_fixed_batch_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_fixed_batch_df.columns.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_batch_sizes = np.array([2 ** i for i in range(10, 14)])\n",
    "validation_batch_sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.figures.colors import PRIMARY, SECONDARY\n",
    "\n",
    "fig, axes = plt.subplots(2, len(validation_batch_sizes), figsize=(2 * FULL_WIDTH, FULL_HEIGHT))\n",
    "\n",
    "fig.subplots_adjust(hspace=0.3, wspace=0.1)  # Adjust this value as needed\n",
    "\n",
    "LOSS_MEAN, LOSS_STD = 'llc/mean/mean', 'llc/mean/std'\n",
    "max_loss_over_time = eval_fixed_batch_df[LOSS_MEAN].max() \n",
    "max_loss_at_end = eval_fixed_batch_df[eval_fixed_batch_df['_step'] == eval_fixed_batch_df['_step'].max()][LOSS_MEAN].max() \n",
    "\n",
    "colors = sns.color_palette('viridis', len(steps))\n",
    "\n",
    "for axes_row, variant in zip(axes, reversed(eval_fixed_batch_df['sampler_config/eval_method'].unique())):\n",
    "    row_df = eval_fixed_batch_df.loc[eval_fixed_batch_df['sampler_config/eval_method'] == variant]\n",
    "    for i, (ax, eval_batch_size) in enumerate(zip(axes_row, validation_batch_sizes)):\n",
    "        _df = row_df.loc[row_df['sampler_config/grad_batch_size'] == eval_batch_size].sort_values(by=['step', '_step'])\n",
    "\n",
    "        for j, step in enumerate(steps):\n",
    "            sns.lineplot(data=_df.loc[_df.step == step], x=\"_step\", y=LOSS_MEAN, color=colors[j], ax=ax, alpha=0.8)\n",
    "\n",
    "        for step, color in zip(steps, colors):\n",
    "            _step_df = _df.loc[_df['step'] == step]\n",
    "            # for chain in range(8):\n",
    "            #     ax.plot(_step_df['_step'], _step_df[f\"llc/mean/{chain}\"], alpha=0.5, color=color, linewidth=0.5)\n",
    "            ax.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.05)  \n",
    "            # ax.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.1)\n",
    "            \n",
    "        if variant == \"grad-minibatch\":\n",
    "            ax.set_title(f\"$\\mu = {eval_batch_size}$\")\n",
    "            ax.set_xlabel('')\n",
    "        else:\n",
    "            ax.set_xlabel(r\"$\\tau$\")\n",
    "\n",
    "        # ax.legend().remove()\n",
    "        ax.set_ylim(0, max_loss_over_time * .8)\n",
    "        # ax.set_yscale('log')\n",
    "\n",
    "        # _final_df = _df.loc[_df['_step'] == _df['_step'].max()]\n",
    "\n",
    "        # ax2.plot(_final_df['task_config/num_tasks'], _final_df[LOSS_MEAN], alpha=1)\n",
    "        # ax2.fill_between(_final_df[\"task_config/num_tasks\"], _final_df[LOSS_MEAN] - _final_df[LOSS_STD], _final_df[LOSS_MEAN] + _final_df[LOSS_STD], alpha=0.5)\n",
    "\n",
    "        # ax2.set_title(fr\"$\\mu = {eval_batch_size}, t=500k$\")\n",
    "        # ax2.set_xlabel(r\"$M$\")\n",
    "        # ax2.set_xscale('log', base=2)\n",
    "\n",
    "        # ax2.set_ylim(0, max_loss_at_end * 1.25)\n",
    "\n",
    "        if i > 0:\n",
    "            ax.set_yticklabels([])\n",
    "            ax.set_ylabel(\"\")\n",
    "            # ax2.set_yticklabels([])\n",
    "            # ax2.set_ylabel(\"\")\n",
    "\n",
    "axes[0, 0].set_ylabel(r\"$\\hat\\lambda_t$\")\n",
    "axes[1, 0].set_ylabel(r\"$\\hat\\lambda_t$\")\n",
    "\n",
    "# Titles above rows (shift far to the right)\n",
    "# axes[0, 0].set_title(r\"(a) $\\hat\\lambda_t$ Dependence on $\\mu$ for SGLD\")\n",
    "# axes[1, 0].set_title(r\"(b) $\\hat\\lambda_t$ Dependence on $\\mu$ for SGLD-MC\")\n",
    "plt.tight_layout()\n",
    "\n",
    "ax= axes[-1, -1]       \n",
    "handles = [plt.Line2D([0, 0], [0, 0], color=colors[j]) for j in range(len(steps))]\n",
    "labels = [f\"$t={step}$\" for step in steps]\n",
    "ax.legend(handles=handles, labels=labels, title=r\"Step $t$\", loc='upper left', bbox_to_anchor=(1.1,1.75))\n",
    "\n",
    "fig.set_facecolor('white')\n",
    "fig.savefig(FIGURES / \"llc-vs-batch-size.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Minibatch Evaluation Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FORCE_REEVAL = True\n",
    "\n",
    "EVAL_MINIBATCH_SWEEP_NAME = \"mini-batch\"\n",
    "EVAL_MINIBATCH_SWEEP_ID = \"0kbtugcd\"\n",
    "\n",
    "if not os.path.exists(ANALYSIS / EVAL_MINIBATCH_SWEEP_NAME) or FORCE_REEVAL:\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        eval_minibatch_sweep = api.sweep(f\"{WANDB_ENTITY}/icl-llc/{EVAL_MINIBATCH_SWEEP_ID}\")\n",
    "        eval_minibatch_runs = list(filter_objs([r for r in eval_minibatch_sweep.runs], config=filters))\n",
    "        eval_minibatch_df = wandb_runs_to_df(eval_minibatch_runs)\n",
    "    \n",
    "    eval_minibatch_df.to_pickle(ANALYSIS / EVAL_MINIBATCH_SWEEP_NAME)\n",
    "\n",
    "else:\n",
    "    eval_minibatch_df = pd.read_pickle(ANALYSIS / EVAL_MINIBATCH_SWEEP_NAME)\n",
    "\n",
    "# columns_to_convert = ['loss/mean', 'loss/std', 'llc/mean', 'llc/std', 'wbic/mean', 'wbic/std']  # Replace with your actual column names\n",
    "\n",
    "# for col in columns_to_convert:\n",
    "#    eval_minibatch_df[col] = pd.to_numeric(eval_minibatch_df[col], errors='coerce')\n",
    "\n",
    "\n",
    "eval_minibatch_df.reset_index(inplace=True)\n",
    "eval_minibatch_df['log_num_tasks'] = np.log2(eval_minibatch_df['task_config/num_tasks'])\n",
    "eval_minibatch_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_minibatch_df.columns.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_batch_sizes = [256, 512, 1024, 2048, 4096, 8192]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_minibatch_df= eval_minibatch_df.loc[eval_minibatch_df['task_config/layer_norm'] == True]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.figures.colors import PRIMARY, SECONDARY\n",
    "\n",
    "fig, axes = plt.subplots(2, len(grad_batch_sizes), figsize=(12, 5))\n",
    "\n",
    "fig.subplots_adjust(hspace=0.3, wspace=0.1)  # Adjust this value as needed\n",
    "\n",
    "LOSS_MEAN, LOSS_STD = 'llc/mean/mean', 'llc/mean/std'\n",
    "max_loss_over_time = eval_minibatch_df[LOSS_MEAN].max() \n",
    "max_loss_at_end = eval_minibatch_df[eval_minibatch_df['_step'] == eval_minibatch_df['_step'].max()][LOSS_MEAN].max() \n",
    "\n",
    "for r, (row_axes, batch_type) in enumerate(zip(axes, ['grad', 'new'])):\n",
    "    for i, (ax1, grad_batch) in enumerate(zip(row_axes, grad_batch_sizes)):\n",
    "        _df = eval_minibatch_df.loc[(eval_minibatch_df['sampler_config/grad_batch_size'] == grad_batch) & (eval_minibatch_df['sampler_config/eval_method'] == f'{batch_type}-minibatch')].sort_values(by=['step', '_step'])\n",
    "\n",
    "        sns.lineplot(data=_df, x=\"_step\", y=LOSS_MEAN, hue=\"step\", palette=\"tab10\", ax=ax1, alpha=1)\n",
    "\n",
    "        for step, color in zip(_df['step'].unique(), [PRIMARY, SECONDARY]):\n",
    "            _step_df = _df.loc[_df['step'] == step]\n",
    "            # for chain in range(8):\n",
    "            #     ax1.plot(_step_df['_step'], _step_df[f\"llc/mean/{chain}\"], alpha=0.5, color=color, linewidth=0.5)\n",
    "            ax1.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.1)  \n",
    "            # ax1.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.1)\n",
    "            \n",
    "\n",
    "        ax1.set_title(f\"$m = {grad_batch}$\")\n",
    "        ax1.set_xlabel(r\"$\\tau$\")\n",
    "\n",
    "        if i < 5:\n",
    "            ax1.legend().remove()\n",
    "        else:\n",
    "            ax1.legend(['10.1k', '_', '500k'], title=r\"Step, $t$\", loc='upper left', bbox_to_anchor=(1.05, 1))\n",
    "            # ax1.legend(['1', '_', '3.2k', '_', '27.1k', '_', '121k', '_', '500k'], title=r\"Step, $t$\", loc='upper left', bbox_to_anchor=(1.05, 1))\n",
    "\n",
    "        ax1.set_ylim(0, max_loss_over_time * .8)\n",
    "        # ax1.set_yscale('log')\n",
    "\n",
    "        # _final_df = _df.loc[_df['_step'] == _df['_step'].max()]\n",
    "\n",
    "        # ax2.plot(_final_df['task_config/num_tasks'], _final_df[LOSS_MEAN], alpha=1)\n",
    "        # ax2.fill_between(_final_df[\"task_config/num_tasks\"], _final_df[LOSS_MEAN] - _final_df[LOSS_STD], _final_df[LOSS_MEAN] + _final_df[LOSS_STD], alpha=0.5)\n",
    "\n",
    "        # ax2.set_title(fr\"$\\mu = {eval_batch_size}, t=500k$\")\n",
    "        # ax2.set_xlabel(r\"$M$\")\n",
    "        # ax2.set_xscale('log', base=2)\n",
    "\n",
    "        # ax2.set_ylim(0, max_loss_at_end * 1.25)\n",
    "\n",
    "        if i > 0:\n",
    "            ax1.set_yticklabels([])\n",
    "            ax1.set_ylabel(\"\")\n",
    "            # ax2.set_yticklabels([])\n",
    "            # ax2.set_ylabel(\"\")\n",
    "\n",
    "axes[0, 0].set_ylabel(\"Same minibatch\\n\" r\"$\\hat\\lambda_t$\")\n",
    "axes[1, 0].set_ylabel(\"New minibatch\\n\" r\"$\\hat\\lambda_t$\")\n",
    "\n",
    "fig.set_facecolor('white')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation Loss Function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FORCE_REEVAL = True\n",
    "\n",
    "EVAL_LOSS_FN_SWEEP_NAME = \"loss-fn\"\n",
    "EVAL_LOSS_FN_SWEEP_ID = \"skcddlzg\"\n",
    "\n",
    "if not os.path.exists(ANALYSIS / EVAL_LOSS_FN_SWEEP_NAME) or FORCE_REEVAL:\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        eval_loss_fn_sweep = api.sweep(f\"{WANDB_ENTITY}/icl-llc/{EVAL_LOSS_FN_SWEEP_ID}\")\n",
    "        eval_loss_fn_runs = list(filter_objs([r for r in eval_loss_fn_sweep.runs], config=filters))\n",
    "        eval_loss_fn_df = wandb_runs_to_df(eval_loss_fn_runs)\n",
    "    \n",
    "    eval_loss_fn_df.to_pickle(ANALYSIS / EVAL_LOSS_FN_SWEEP_NAME)\n",
    "\n",
    "else:\n",
    "    eval_loss_fn_df = pd.read_pickle(ANALYSIS / EVAL_LOSS_FN_SWEEP_NAME)\n",
    "\n",
    "# columns_to_convert = ['loss/mean', 'loss/std', 'llc/mean', 'llc/std', 'wbic/mean', 'wbic/std']  # Replace with your actual column names\n",
    "\n",
    "# for col in columns_to_convert:\n",
    "#     eval_loss_fn_df[col] = pd.to_numeric(eval_loss_fn_df[col], errors='coerce')\n",
    "\n",
    "\n",
    "eval_loss_fn_df.reset_index(inplace=True)\n",
    "eval_loss_fn_df['log_num_tasks'] = np.log2(eval_loss_fn_df['task_config/num_tasks'])\n",
    "eval_loss_fn_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_loss_fn_df = eval_loss_fn_df.loc[eval_loss_fn_df['task_config/layer_norm'] == True]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fns = ['mse', \"subsequence-mse\"]\n",
    "focus_steps = [10204, 499999]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.figures.colors import PRIMARY, SECONDARY\n",
    "\n",
    "fig, axes = plt.subplots(1, len(loss_fns), figsize=(12, 2.5))\n",
    "\n",
    "fig.subplots_adjust(hspace=0.3, wspace=0.1)  # Adjust this value as needed\n",
    "\n",
    "LOSS_MEAN, LOSS_STD = 'llc/mean/mean', 'llc/mean/std'\n",
    "max_loss_over_time = eval_loss_fn_df[LOSS_MEAN].max() \n",
    "max_loss_at_end = eval_loss_fn_df[eval_loss_fn_df['_step'] == eval_loss_fn_df['_step'].max()][LOSS_MEAN].max() \n",
    "\n",
    "colors = sns.color_palette('viridis', len(steps))\n",
    "\n",
    "\n",
    "for i, (ax1, loss_fn) in enumerate(zip(axes, loss_fns)):\n",
    "    _df = eval_loss_fn_df.loc[eval_loss_fn_df['sampler_config/eval_loss_fn'] == loss_fn].sort_values(by=['step', '_step'])\n",
    "\n",
    "\n",
    "    for step, color in zip(_df['step'].unique(), colors):\n",
    "        _step_df = _df.loc[_df['step'] == step]\n",
    "        # for chain in range(8):\n",
    "        #     ax1.plot(_step_df['_step'], _step_df[f\"llc/mean/{chain}\"], alpha=0.5, color=color, linewidth=0.5)\n",
    "        sns.lineplot(data=_step_df, x=\"_step\", y=LOSS_MEAN, color=color, ax=ax1, alpha=0.8)\n",
    "\n",
    "    for step, color in zip(_df['step'].unique(), colors):\n",
    "        _step_df = _df.loc[_df['step'] == step]\n",
    "\n",
    "        ax1.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.05)  \n",
    "        # ax1.fill_between(_step_df['_step'], _step_df[LOSS_MEAN] - _step_df[LOSS_STD], _step_df[LOSS_MEAN] + _step_df[LOSS_STD], alpha=0.1)\n",
    "        \n",
    "\n",
    "    ax1.set_title({ \"mse\": \"(a) MSE\", \"subsequence-mse\": \"(b) Subsequence MSE\" }[loss_fn])\n",
    "    ax1.set_xlabel(r\"$\\tau$\")\n",
    "\n",
    "    if i < 1:\n",
    "        ax1.legend().remove()\n",
    "    else:\n",
    "        ax1.legend([l if i else \"_\" for l in labels for i in [1, 0]], title=r\"Step $t$\", loc='upper left', bbox_to_anchor=(1.05, 1))\n",
    "\n",
    "    ax1.set_ylim(0, max_loss_over_time * .8)\n",
    "    # ax1.set_yscale('log')\n",
    "\n",
    "    # _final_df = _df.loc[_df['_step'] == _df['_step'].max()]\n",
    "\n",
    "    # ax2.plot(_final_df['task_config/num_tasks'], _final_df[LOSS_MEAN], alpha=1)\n",
    "    # ax2.fill_between(_final_df[\"task_config/num_tasks\"], _final_df[LOSS_MEAN] - _final_df[LOSS_STD], _final_df[LOSS_MEAN] + _final_df[LOSS_STD], alpha=0.5)\n",
    "\n",
    "    # ax2.set_title(fr\"$\\mu = {eval_batch_size}, t=500k$\")\n",
    "    # ax2.set_xlabel(r\"$M$\")\n",
    "    # ax2.set_xscale('log', base=2)\n",
    "\n",
    "    # ax2.set_ylim(0, max_loss_at_end * 1.25)\n",
    "\n",
    "    if i > 0:\n",
    "        ax1.set_yticklabels([])\n",
    "        ax1.set_ylabel(\"\")\n",
    "        # ax2.set_yticklabels([])\n",
    "        # ax2.set_ylabel(\"\")\n",
    "\n",
    "axes[0].set_ylabel(r\"$\\hat\\lambda_t$\")\n",
    "\n",
    "\n",
    "fig.set_facecolor('white')\n",
    "fig.savefig(FIGURES / \"llc-vs-loss-fn.pdf\", bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Role of seed\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seeded_run = api.run(\"{WANDB_ENTITY}/icl/tptexwlo\")\n",
    "unseeded_run = api.run(\"{WANDB_ENTITY}/icl/aegfv4rj\")\n",
    "seeds_df = wandb_runs_to_df([seeded_run, unseeded_run])\n",
    "print(seeds_df['sampler_config/init_seed'].unique())\n",
    "seeds_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.normal(0, 1, (1_000_000, 4)).norm(dim=-1).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(FULL_WIDTH, FULL_HEIGHT * 3 / 5))\n",
    "\n",
    "cmap = sns.color_palette(\"viridis\", 8)\n",
    "\n",
    "for seed in [None, 0]:\n",
    "    if seed is None:\n",
    "        _df = seeds_df.loc[seeds_df['sampler_config/init_seed'].isna()]\n",
    "    else:\n",
    "        _df = seeds_df.loc[seeds_df['sampler_config/init_seed'] == seed]\n",
    "\n",
    "    print(seed, _df.values.shape)\n",
    "    alpha = 0.75 if seed is None else .5\n",
    "    linewidth = 0.5 if seed is None else 1.5\n",
    "\n",
    "    for t in range(8):\n",
    "        color = cmap[t]\n",
    "        sns.lineplot(data=_df, x=\"_step\", y=f\"llc/mean/{t}\", alpha=alpha, ax=ax, linewidth=linewidth, color=color)\n",
    "\n",
    "ax.set_ylabel('$\\hat\\lambda_t$')\n",
    "ax.set_xlabel(r\"Step $t$\")\n",
    "ax.set_xlim(100, 500000)\n",
    "ax.set_xscale('log')\n",
    "# handles = [plt.Line2D([0], [0], color=cmap[t], linewidth=[1, 3][s], alpha=[1, 0.25][s]) for s in range(2) for t in range(8)]\n",
    "# labels = [rf\"$\\hat \\lambda_t^{{({t+1})}}$\" for s in range(2) for t in range(8)]\n",
    "\n",
    "handles = [plt.Line2D([0], [0], color=cmap[t], linewidth=1) for t in range(8)]\n",
    "labels = [rf\"$\\hat \\lambda_t^{{({t+1})}}$\" for t in range(8)]\n",
    "handles += [plt.Line2D([0], [0], color='black', linewidth=1), plt.Line2D([0], [0], color='black', linewidth=3, alpha=0.5)]\n",
    "labels += [r\"Diff. seed\", r\"Same seed\"]\n",
    "\n",
    "ax.legend(handles, labels, bbox_to_anchor=(1.05, 1.05), title=\"Token\", loc='upper left', ncol=1, fontsize='xx-small', title_fontsize='xx-small')\n",
    "\n",
    "# twinax = ax.twinx()\n",
    "# twinax.legend(handles, labels, bbox_to_anchor=(1.05, 1), title=\"Seed\", loc='center left', ncol=1, fontsize='xx-small', title_fontsize='xx-small')\n",
    "fig.savefig(FIGURES / \"llc-vs-seed.pdf\", bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLCs over time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from infra.utils.iterables import flatten_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp(list(df.columns))\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llc_chain_columns = [f'llc-chain/{i}' for i in range(25)]\n",
    "df[llc_chain_columns] = df[llc_chain_columns].replace(\"NaN\", np.nan)\n",
    "\n",
    "\n",
    "# Calculate the average of non-NaN values in llc-chain columns\n",
    "# and the fraction of NaN values\n",
    "llc_chain_values = df[llc_chain_columns]\n",
    "mean_llc_chain = llc_chain_values.mean(axis=1, skipna=True)\n",
    "frac_nan = llc_chain_values.isna().mean(axis=1)\n",
    "\n",
    "df[\"llc/mean-fixed\"] = mean_llc_chain\n",
    "df[\"llc/frac-nan\"] = frac_nan\n",
    "df[\"log_num_tasks\"] = np.log(df[\"task_config/num_tasks\"])\n",
    "\n",
    "mean_llc_chain, frac_nan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.lineplot(data=df, x=\"_step\", y=\"llc/frac-nan\", hue=\"task_config/num_tasks\")\n",
    "plt.xscale(\"log\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import patches\n",
    "\n",
    "\n",
    "TRANSITIONS = [\n",
    "    (100, 800, 'A1'),\n",
    "    (800, 10_000, 'A2'),\n",
    "    (10_000, 28_000, 'B1'),\n",
    "    (28_000, 280_000, 'B2'),\n",
    "]\n",
    "\n",
    "INIT_X = TRANSITIONS[0][0]\n",
    "FINAL_X = TRANSITIONS[-1][1]\n",
    "\n",
    "def plot_transitions(axes, **kwargs):\n",
    "    from icl.figures.colors import plot_transitions as _plot_transitions\n",
    "    return _plot_transitions(axes, TRANSITIONS, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = checkpointers[0].file_ids\n",
    "steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "from icl.analysis.evals import RegressionEvaluator\n",
    "from icl.regression.experiments.activations_analysis import iter_models\n",
    "from icl.train import Run\n",
    "\n",
    "evals = []\n",
    "functional_metrics = []\n",
    "gradient_norms = []\n",
    "\n",
    "B = 8192\n",
    "K = 8\n",
    "D = 4\n",
    "OOD_MULTIPLIER = 3\n",
    "\n",
    "def eval_loss(yhats, ys):\n",
    "    return ((yhats - ys) ** 2).mean(dim=0)[:, 0]\n",
    "\n",
    "def apply_transformations(ws, xs):\n",
    "    return xs @ ws.view(B, D, 1)\n",
    "\n",
    "for log2_M, config in tqdm(enumerate(configs)):\n",
    "    run = Run(config)\n",
    "    run.evaluator = RegressionEvaluator(\n",
    "        pretrain_dist=run.pretrain_dist,\n",
    "        true_dist=run.true_dist,\n",
    "        max_examples=config.task_config.max_examples,\n",
    "        eval_batch_size=8192,\n",
    "        seed=config.task_config.true_seed,\n",
    "    )\n",
    "    pretrain_dist_noiseless = run.config.task_config.pretrain_dist_factory().to(\n",
    "        DEVICE\n",
    "    )\n",
    "    noise_std = pretrain_dist_noiseless.std\n",
    "    pretrain_dist_noiseless.std = 0.\n",
    "\n",
    "    ws = pretrain_dist_noiseless.task_distribution.sample_tasks(B) # -> B D \n",
    "    wpriors = pretrain_dist_noiseless.task_distribution.tasks.mean(dim=0) # -> D\n",
    "    wpriors = wpriors.repeat(B, 1) # -> B D\n",
    "\n",
    "    xs = torch.normal(\n",
    "        mean=0.,\n",
    "        std=1.,\n",
    "        size=(B, K, D,),\n",
    "        device=DEVICE\n",
    "    )\n",
    "    ood_xs = OOD_MULTIPLIER * xs\n",
    "\n",
    "    errors = torch.normal(\n",
    "        mean=0.,\n",
    "        std=noise_std,\n",
    "        size=(B, K, 1,),\n",
    "        device=DEVICE,\n",
    "    )\n",
    "\n",
    "    ys_without_noise = apply_transformations(ws, xs)\n",
    "    ood_ys_without_noise= OOD_MULTIPLIER * ys_without_noise\n",
    "\n",
    "    ys = ys_without_noise + errors\n",
    "    ood_ys = ood_ys_without_noise + errors\n",
    "\n",
    "    yhats_prior = apply_transformations(wpriors, xs)\n",
    "    yhats_zero = torch.zeros_like(ys)\n",
    "    # ood_yhats_prior = apply_transformations(wpriors, ood_xs)\n",
    "    \n",
    "    for step, model in zip(steps, iter_models(run.model, run.checkpointer)):\n",
    "        yhats = model(xs, ys)\n",
    "        ood_yhats = model(xs, ood_ys)\n",
    "        # yhats_without_noise = model(xs, ys_without_noise)\n",
    "\n",
    "        losses = eval_loss(yhats, ys)\n",
    "        # losses_without_noise = eval_loss(yhats_without_noise, ys_without_noise)\n",
    "        losses_prior = eval_loss(yhats, yhats_prior)\n",
    "        losses_zero = eval_loss(yhats, yhats_zero)\n",
    "\n",
    "        ood_losses = eval_loss(ood_yhats, ood_ys)\n",
    "        # ood_losses_midpoint = eval_loss(ood_yhats, ood_yhats_prior)\n",
    "        # ood_losses_without_noise = eval_loss(yhats, ood_ys_without_noise)\n",
    "\n",
    "        loss = losses.mean()\n",
    "        loss.backward()\n",
    "\n",
    "        for n, p in model.named_parameters():\n",
    "            if p.grad is None:\n",
    "                continue\n",
    "\n",
    "            grad_sq_mean = (p.grad ** 2).mean().item()\n",
    "            grad_sq_std = (p.grad ** 2).std().item()\n",
    "\n",
    "            gradient_norms.append({\n",
    "                \"m\": log2_M,\n",
    "                \"M\": 2 ** log2_M,    \n",
    "                \"step\": step,\n",
    "                \"layer\": n,\n",
    "                \"grad/norm\": grad_sq_mean ** 0.5,\n",
    "                \"grad_sq/mean\": grad_sq_mean,\n",
    "                \"grad_sq/std\": grad_sq_std,\n",
    "                \"numel\": p.numel(),\n",
    "                \"loss\": loss.item(),\n",
    "            })          \n",
    "\n",
    "            p.grad = None \n",
    "\n",
    "        for token in range(8):\n",
    "            functional_metrics.append({\n",
    "                \"m\": log2_M,\n",
    "                \"M\": 2 ** log2_M,\n",
    "                \"step\": step,\n",
    "                \"loss\": losses[token].item(),\n",
    "                \"ood_loss\": ood_losses[token].item(),\n",
    "                # \"loss_without_noise\": losses_without_noise[i],\n",
    "                # \"ood_loss_without_noise\": ood_losses_without_noise[i],\n",
    "                \"loss_prior\": losses_prior[token].item(),\n",
    "                \"loss_zero\": losses_zero[token].item(),\n",
    "                # \"ood_loss_midpoint\": ood_losses_midpoint[i],\n",
    "                \"token\": token\n",
    "            })\n",
    "\n",
    "        evals.append({\n",
    "            \"m\": log2_M,\n",
    "            \"M\": 2 ** log2_M,\n",
    "            \"step\": step,\n",
    "            \"weight_norm\": (sum([(p ** 2).sum() for p in model.parameters()]) ** 0.5).item(),\n",
    "            **run.evaluator(model),\n",
    "        })\n",
    "\n",
    "\n",
    "evals = pd.DataFrame(evals)\n",
    "evals.to_csv(ANALYSIS / \"small-model-evals.csv\", index=False)\n",
    "functional_metrics = pd.DataFrame(functional_metrics)\n",
    "functional_metrics.to_csv(ANALYSIS / \"small-model-functional-metrics.csv\", index=False)\n",
    "gradient_norms = pd.DataFrame(gradient_norms)\n",
    "gradient_norms.to_csv(ANALYSIS / \"small-model-gradient-norms.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evals = pd.read_csv(ANALYSIS / \"small-model-evals.csv\")\n",
    "functional_metrics = pd.read_csv(ANALYSIS / \"small-model-functional-metrics.csv\")\n",
    "gradient_norms = pd.read_csv(ANALYSIS / \"small-model-gradient-norms.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evals.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.figures.derivatives import d_dt, d_dlogt, dlog_dlogt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for log2_M in range(20):\n",
    "    # Filter the DataFrame and compute the derivatives\n",
    "    mse_values = evals.loc[evals.m == log2_M, \"pretrain/mse\"].values\n",
    "    llc_values = df.loc[df[\"task_config/num_tasks\"] == int(2 ** log2_M), \"llc/mean-fixed\"].values\n",
    "    weightnorm_values = evals.loc[evals.m == log2_M, \"weight_norm\"].values\n",
    "    \n",
    "    # Compute the derivatives using your d_dlogt function\n",
    "    dloss_dlogt_values = d_dlogt(steps, mse_values)\n",
    "    dllc_dlogt_values = d_dlogt(steps, llc_values)\n",
    "    dweightnorm_dlogt_values = d_dlogt(steps, weightnorm_values)\n",
    "\n",
    "    # Assign the computed derivatives back to the original DataFrame\n",
    "    evals.loc[evals.m == log2_M, \"dloss_dlogt\"] = dloss_dlogt_values\n",
    "    evals.loc[evals.m == log2_M, \"dllc_dlogt\"] = dllc_dlogt_values\n",
    "    evals.loc[evals.m == log2_M, \"dweightnorm_dlogt\"] = dweightnorm_dlogt_values\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evals.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import colors, lines, patches\n",
    "\n",
    "def get_reduced_viridis_palette(num_colors, ratio=3 / 3.5):\n",
    "    return sns.color_palette(\"viridis\", int(num_colors // ratio))[:num_colors]\n",
    "\n",
    "LINE_PALETTE = get_reduced_viridis_palette(21-5)\n",
    "# num_palette_steps = int((21 * 3.75) // 3)\n",
    "# LINE_PALETTE = [sns.color_palette(\"coolwarm\", num_palette_steps)[i] for i in [*range(10), *range(num_palette_steps - 10, num_palette_steps)]]\n",
    "\n",
    "print(LINE_PALETTE)\n",
    "# \"viridis\"\n",
    "ALPHA=0.75\n",
    "\n",
    "fig, axes = plt.subplots(2, 2, figsize=(20, 6))\n",
    "\n",
    "ax = axes[0, 0]\n",
    "\n",
    "filtered_evals = evals.loc[(evals.step != 20408) & (evals['m'] > 5)]\n",
    "\n",
    "sns.lineplot(data=filtered_evals, x=\"step\", y=\"pretrain/mse\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$L_\\mathrm{val}$\")\n",
    "ax.legend().remove()\n",
    "\n",
    "ax = axes[1, 0]\n",
    "\n",
    "sns.lineplot(data=filtered_evals, x=\"step\", y=\"dloss_dlogt\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\delta L_\\mathrm{val}/\\delta \\log t$\")\n",
    "ax.legend().remove()\n",
    "ax.set_ylim(-2.5, 2.5)\n",
    "\n",
    "ax = axes[0, 1]\n",
    "\n",
    "sns.lineplot(data=df.loc[(df._step != 20408) & (df.log_num_tasks > 5)], x=\"_step\", y=\"llc/mean-fixed\", hue=\"log_num_tasks\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\hat\\lambda$\")\n",
    "ax.legend().remove()\n",
    "\n",
    "ax = axes[1, 1]\n",
    "\n",
    "sns.lineplot(data=filtered_evals, x=\"step\", y=\"dllc_dlogt\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\delta \\hat\\lambda/\\delta \\log t$\")\n",
    "ax.set_ylim(-500, 500)\n",
    "ax.legend().remove()\n",
    "\n",
    "# ax = axes[0, 2]\n",
    "\n",
    "# sns.lineplot(data=filtered_evals, x=\"step\", y=\"weight_norm\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "# ax.set_xscale(\"log\")\n",
    "# ax.set_xlim(INIT_X, FINAL_X)\n",
    "# ax.set_xlabel(\"Step, $t$\")\n",
    "# ax.set_ylabel(\"$|w_t|$\")\n",
    "# ax.legend().remove()\n",
    "\n",
    "# ax.set_ylim(20, 800)\n",
    "\n",
    "# ax = axes[1, 2]\n",
    "\n",
    "# sns.lineplot(data=filtered_evals, x=\"step\", y=\"dweightnorm_dlogt\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "# ax.set_xscale(\"log\")\n",
    "# ax.set_xlim(INIT_X, FINAL_X)\n",
    "# ax.set_xlabel(\"Step, $t$\")\n",
    "# ax.set_ylabel(\"$\\delta|w_t|/\\delta\\log t$\")\n",
    "# ax.legend().remove()\n",
    "\n",
    "_patches = plot_transitions(axes)\n",
    "\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "\n",
    "# Create custom handles\n",
    "# handles = [lines.Line2D([], [], color=custom_colors[i], marker='o', linestyle='', label=custom_labels[i]) for i in range(len(custom_labels))]\n",
    "\n",
    "# Add the custom handles to the existing ones\n",
    "# handles.extend(custom_handles)\n",
    "\n",
    "# Now, you can create the legend with the updated handles and custom labels\n",
    "# axes[0, 2].legend(handles=handles, title=\"$\\log_2 M$\", loc='center right', bbox_to_anchor=(0.9, .5))\n",
    "\n",
    "# plt.tight_layout(rect=[0, 0, 0.9, 1])  # Adjust layout to make room for colorbar\n",
    "# plt.tight_layout()\n",
    "\n",
    "cbar_ax = fig.add_axes([0.93, 0.125, 0.02, 0.33])  # Adjust as necessary for position and size\n",
    "custom_cmap = colors.LinearSegmentedColormap.from_list(\"custom_cmap\", LINE_PALETTE)\n",
    "\n",
    "sm = plt.cm.ScalarMappable(cmap=custom_cmap, norm=plt.Normalize(vmin=0+5, vmax=20), )\n",
    "sm._A = []  # Dummy array for the ScalarMappable. \n",
    "cbar = fig.colorbar(sm, cax=cbar_ax)\n",
    "\n",
    "tick_positions =  [5, 10, 15, 20]  # Positions for each color\n",
    "tick_labels = map(str, tick_positions)  # Labels for each color\n",
    "cbar.set_ticks(tick_positions)\n",
    "cbar.set_ticklabels(tick_labels)\n",
    "cbar.set_label(\"$\\log_2 M$\")\n",
    "\n",
    "\n",
    "stages_legend_ax = fig.add_axes([0.945, 0.68, 0.02, 0.25])  # Adjust as necessary for position and size\n",
    "stages_legend_ax.axis('off')\n",
    "stages_legend_ax.legend(handles=_patches, title=\"Stage\", loc='upper center', bbox_to_anchor=(0, .5))\n",
    "\n",
    "fig.set_facecolor('white')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import colors, lines, patches\n",
    "\n",
    "def get_reduced_viridis_palette(num_colors, ratio=3 / 3.5):\n",
    "    return sns.color_palette(\"viridis\", int(num_colors // ratio))[:num_colors]\n",
    "\n",
    "LINE_PALETTE = get_reduced_viridis_palette(5)\n",
    "# num_palette_steps = int((21 * 3.75) // 3)\n",
    "# LINE_PALETTE = [sns.color_palette(\"coolwarm\", num_palette_steps)[i] for i in [*range(10), *range(num_palette_steps - 10, num_palette_steps)]]\n",
    "\n",
    "print(LINE_PALETTE)\n",
    "# \"viridis\"\n",
    "ALPHA=0.75\n",
    "\n",
    "fig, axes = plt.subplots(2, 2, figsize=(20, 6))\n",
    "\n",
    "ax = axes[0, 0]\n",
    "\n",
    "filtered_evals = evals.loc[(evals.step != 20408) & (evals['m'] <= 5)]\n",
    "\n",
    "sns.lineplot(data=filtered_evals, x=\"step\", y=\"pretrain/mse\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$L_\\mathrm{val}$\")\n",
    "ax.legend().remove()\n",
    "\n",
    "ax = axes[1, 0]\n",
    "\n",
    "sns.lineplot(data=filtered_evals, x=\"step\", y=\"dloss_dlogt\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\delta L_\\mathrm{val}/\\delta \\log t$\")\n",
    "ax.legend().remove()\n",
    "ax.set_ylim(-2.5, 2.5)\n",
    "\n",
    "ax = axes[0, 1]\n",
    "\n",
    "sns.lineplot(data=df.loc[(df._step != 20408) & (df.log_num_tasks <= 5)], x=\"_step\", y=\"llc/mean-fixed\", hue=\"log_num_tasks\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\hat\\lambda$\")\n",
    "ax.legend().remove()\n",
    "\n",
    "ax = axes[1, 1]\n",
    "\n",
    "sns.lineplot(data=filtered_evals, x=\"step\", y=\"dllc_dlogt\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\delta \\hat\\lambda/\\delta \\log t$\")\n",
    "ax.set_ylim(-500, 500)\n",
    "ax.legend().remove()\n",
    "\n",
    "# ax = axes[0, 2]\n",
    "\n",
    "# sns.lineplot(data=filtered_evals, x=\"step\", y=\"weight_norm\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "# ax.set_xscale(\"log\")\n",
    "# ax.set_xlim(INIT_X, FINAL_X)\n",
    "# ax.set_xlabel(\"Step, $t$\")\n",
    "# ax.set_ylabel(\"$|w_t|$\")\n",
    "# ax.legend().remove()\n",
    "\n",
    "# ax.set_ylim(20, 800)\n",
    "\n",
    "# ax = axes[1, 2]\n",
    "\n",
    "# sns.lineplot(data=filtered_evals, x=\"step\", y=\"dweightnorm_dlogt\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "# ax.set_xscale(\"log\")\n",
    "# ax.set_xlim(INIT_X, FINAL_X)\n",
    "# ax.set_xlabel(\"Step, $t$\")\n",
    "# ax.set_ylabel(\"$\\delta|w_t|/\\delta\\log t$\")\n",
    "# ax.legend().remove()\n",
    "\n",
    "_patches = plot_transitions(axes)\n",
    "\n",
    "\n",
    "handles, labels = ax.get_legend_handles_labels()\n",
    "\n",
    "# Create custom handles\n",
    "# handles = [lines.Line2D([], [], color=custom_colors[i], marker='o', linestyle='', label=custom_labels[i]) for i in range(len(custom_labels))]\n",
    "\n",
    "# Add the custom handles to the existing ones\n",
    "# handles.extend(custom_handles)\n",
    "\n",
    "# Now, you can create the legend with the updated handles and custom labels\n",
    "# axes[0, 2].legend(handles=handles, title=\"$\\log_2 M$\", loc='center right', bbox_to_anchor=(0.9, .5))\n",
    "\n",
    "# plt.tight_layout(rect=[0, 0, 0.9, 1])  # Adjust layout to make room for colorbar\n",
    "# plt.tight_layout()\n",
    "\n",
    "cbar_ax = fig.add_axes([0.93, 0.125, 0.02, 0.33])  # Adjust as necessary for position and size\n",
    "custom_cmap = colors.LinearSegmentedColormap.from_list(\"custom_cmap\", LINE_PALETTE)\n",
    "\n",
    "sm = plt.cm.ScalarMappable(cmap=custom_cmap, norm=plt.Normalize(vmin=0, vmax=5), )\n",
    "sm._A = []  # Dummy array for the ScalarMappable. \n",
    "cbar = fig.colorbar(sm, cax=cbar_ax)\n",
    "\n",
    "tick_positions =  [0, 1, 2, 3, 4, 5]  # Positions for each color\n",
    "tick_labels = map(str, tick_positions)  # Labels for each color\n",
    "cbar.set_ticks(tick_positions)\n",
    "cbar.set_ticklabels(tick_labels)\n",
    "cbar.set_label(\"$\\log_2 M$\")\n",
    "\n",
    "\n",
    "stages_legend_ax = fig.add_axes([0.945, 0.68, 0.02, 0.25])  # Adjust as necessary for position and size\n",
    "stages_legend_ax.axis('off')\n",
    "stages_legend_ax.legend(handles=_patches, title=\"Stage\", loc='upper center', bbox_to_anchor=(0, .5))\n",
    "\n",
    "fig.set_facecolor('white')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 4, figsize=(20, 5))\n",
    "\n",
    "# Average across token colum\n",
    "functional_metrics_averages = functional_metrics.groupby([\"m\", \"step\"]).mean().reset_index()\n",
    "\n",
    "ax = axes[0]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"loss\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "\n",
    "ax = axes[1]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"loss_prior\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_ylim(0.01, 10)\n",
    "\n",
    "ax = axes[2]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"loss_zero\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_ylim(0.01, 10)\n",
    "\n",
    "ax = axes[3]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"ood_loss\", hue=\"m\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "\n",
    "for ax in axes.flatten():\n",
    "    ax.set_yscale('log')\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_xlim(100, 400_000)\n",
    "    ax.set_xlabel(\"Step, $t$\")\n",
    "    ax.legend().remove()\n",
    "\n",
    "add_milestones(axes) #, alpha=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LINE_PALETTE = get_reduced_viridis_palette(8)\n",
    "ALPHA = 1\n",
    "\n",
    "fig, axes = plt.subplots(1, 4, figsize=(20, 5))\n",
    "\n",
    "# Average across token colum\n",
    "functional_metrics_m20 = functional_metrics.loc[functional_metrics.m == 20]\n",
    "\n",
    "ax = axes[0]\n",
    "sns.lineplot(data=functional_metrics_m20.loc[functional_metrics_m20.step != 20408], x=\"step\", y=\"loss\", hue=\"token\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "\n",
    "ax = axes[1]\n",
    "sns.lineplot(data=functional_metrics_m20.loc[functional_metrics_m20.step != 20408], x=\"step\", y=\"loss_prior\", hue=\"token\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_ylim(0.01, 10)\n",
    "\n",
    "ax = axes[2]\n",
    "sns.lineplot(data=functional_metrics_m20.loc[functional_metrics_m20.step != 20408], x=\"step\", y=\"loss_zero\", hue=\"token\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "ax.set_ylim(0.01, 10)\n",
    "\n",
    "ax = axes[3]\n",
    "sns.lineplot(data=functional_metrics_m20.loc[functional_metrics_m20.step != 20408], x=\"step\", y=\"ood_loss\", hue=\"token\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "\n",
    "for ax in axes.flatten():\n",
    "    ax.set_yscale('log')\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_xlim(100, 400_000)\n",
    "    ax.set_xlabel(\"Step, $t$\")\n",
    "    ax.legend().remove()\n",
    "\n",
    "add_milestones(axes) #, alpha=0.8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import colors as mcolors\n",
    "\n",
    "# LINE_PALETTE=\"viridis\"\n",
    "# ALPHA=1\n",
    "\n",
    "fig, axes = plt.subplots(1, 5, figsize=(20, 4))\n",
    "\n",
    "for token, log2_M in enumerate([0, 1, 3, 6, 20]):\n",
    "    # Average across token colum\n",
    "    ax = axes[token]\n",
    "    functional_metrics_specific = functional_metrics.loc[functional_metrics.m == log2_M]\n",
    "    sns.lineplot(data=functional_metrics_specific, x=\"step\", y=\"loss\", hue=\"token\", palette=LINE_PALETTE, alpha=ALPHA, ax=ax)\n",
    "    ax.set_title(f\"$M = 2^{{{log2_M}}}$\")\n",
    "\n",
    "\n",
    "for ax in axes.flatten():\n",
    "    ax.set_yscale('log')\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_xlim(INIT_X, FINAL_X)\n",
    "    ax.set_xlabel(\"Step, $t$\")\n",
    "    ax.set_ylabel(\"$L_\\mathrm{val}$\")\n",
    "    ax.legend().remove()\n",
    "\n",
    "add_milestones(axes, alpha=0.2) #, alpha=0.25)\n",
    "\n",
    "cbar_ax = fig.add_axes([0.93, 0.15, 0.02, 0.7])  # Adjust as necessary for position and size\n",
    "\n",
    "custom_cmap = mcolors.LinearSegmentedColormap.from_list(\"custom_cmap\", LINE_PALETTE)\n",
    "\n",
    "sm = plt.cm.ScalarMappable(cmap=custom_cmap, norm=plt.Normalize(vmin=1, vmax=8), )\n",
    "sm._A = []  # Dummy array for the ScalarMappable. \n",
    "cbar = fig.colorbar(sm, cax=cbar_ax)\n",
    "\n",
    "tick_positions = range(1, len(LINE_PALETTE)+1)  # Positions for each color\n",
    "tick_labels = [f\"${i}$\" for i in range(1, len(LINE_PALETTE) + 1)] # Replace with your labels\n",
    "cbar.set_ticks(tick_positions)\n",
    "cbar.set_ticklabels(tick_labels)\n",
    "cbar.set_label(\"$k$\")\n",
    "\n",
    "plt.tight_layout(rect=[0, 0, 0.91, 1])  # Adjust layout to make room for colorbar\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import colors as mcolors\n",
    "\n",
    "# LINE_PALETTE=\"viridis\"\n",
    "# ALPHA=1\n",
    "\n",
    "fig, axes = plt.subplots(1, 5, figsize=(20, 4))\n",
    "\n",
    "for token, log2_M in enumerate([0, 1, 3, 6, 20]):\n",
    "    # Average across token colum\n",
    "    ax = axes[token]\n",
    "    functional_metrics_specific = functional_metrics.loc[functional_metrics.m == log2_M]\n",
    "\n",
    "    icl_score = functional_metrics_specific.loc[functional_metrics_specific.token == 7, \"loss\"].values - functional_metrics_specific.loc[functional_metrics_specific.token == 4, \"loss\"].values\n",
    "    sns.lineplot(x=steps, y=icl_score, alpha=ALPHA, ax=ax)\n",
    "    ax.set_title(f\"$M = 2^{{{log2_M}}}$\")\n",
    "\n",
    "\n",
    "for ax in axes.flatten():\n",
    "    # ax.set_yscale('symlog')\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_xlim(INIT_X, FINAL_X)\n",
    "    ax.set_xlabel(\"Step, $t$\")\n",
    "    ax.set_ylabel(\"$L_\\mathrm{val}$\")\n",
    "    ax.legend().remove()\n",
    "\n",
    "add_milestones(axes, alpha=0.2) #, alpha=0.25)\n",
    "\n",
    "cbar_ax = fig.add_axes([0.93, 0.15, 0.02, 0.7])  # Adjust as necessary for position and size\n",
    "\n",
    "custom_cmap = mcolors.LinearSegmentedColormap.from_list(\"custom_cmap\", LINE_PALETTE)\n",
    "\n",
    "sm = plt.cm.ScalarMappable(cmap=custom_cmap, norm=plt.Normalize(vmin=1, vmax=8), )\n",
    "sm._A = []  # Dummy array for the ScalarMappable. \n",
    "cbar = fig.colorbar(sm, cax=cbar_ax)\n",
    "\n",
    "tick_positions = range(1, len(LINE_PALETTE)+1)  # Positions for each color\n",
    "tick_labels = [f\"${i}$\" for i in range(1, len(LINE_PALETTE) + 1)] # Replace with your labels\n",
    "cbar.set_ticks(tick_positions)\n",
    "cbar.set_ticklabels(tick_labels)\n",
    "cbar.set_label(\"$k$\")\n",
    "\n",
    "plt.tight_layout(rect=[0, 0, 0.91, 1])  # Adjust layout to make room for colorbar\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wpriors_over_m = []\n",
    "wprior_norms_over_m = []\n",
    "\n",
    "for log2_M, config in tqdm(enumerate(configs)):\n",
    "    run = Run(config)\n",
    "    run.evaluator = RegressionEvaluator(\n",
    "        pretrain_dist=run.pretrain_dist,\n",
    "        true_dist=run.true_dist,\n",
    "        max_examples=config.task_config.max_examples,\n",
    "        eval_batch_size=8192,\n",
    "        seed=config.task_config.true_seed,\n",
    "    )\n",
    "    pretrain_dist_noiseless = run.config.task_config.pretrain_dist_factory().to(\n",
    "        DEVICE\n",
    "    )\n",
    "    noise_std = pretrain_dist_noiseless.std\n",
    "    pretrain_dist_noiseless.std = 0.\n",
    "\n",
    "    ws = pretrain_dist_noiseless.task_distribution.sample_tasks(B) # -> B D \n",
    "    wpriors = pretrain_dist_noiseless.task_distribution.tasks.mean(dim=0) # -> D\n",
    "    wpriors_over_m.append(wpriors)\n",
    "    wprior_norms_over_m.append(wpriors.norm())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.arange(21), [w.item() for w in wprior_norms_over_m])\n",
    "plt.yscale(\"log\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 4, figsize=(20, 5))\n",
    "\n",
    "# Average across token colum\n",
    "for log2_M in range(21):\n",
    "    functional_metrics_averages.loc[functional_metrics_averages.m == log2_M, \"loss_zero_norm\"] = functional_metrics_averages.loc[functional_metrics_averages.m == log2_M, \"loss_zero\"] / wprior_norms_over_m[log2_M].item()\n",
    "    functional_metrics_averages.loc[functional_metrics_averages.m == log2_M, \"loss_prior_norm\"] = functional_metrics_averages.loc[functional_metrics_averages.m == log2_M, \"loss_prior\"] / wprior_norms_over_m[log2_M].item()\n",
    "\n",
    "ax = axes[0]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"loss\", hue=\"m\", palette=\"viridis\", alpha=0.5, ax=ax)\n",
    "\n",
    "ax = axes[1]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"loss_prior_norm\", hue=\"m\", palette=\"viridis\", alpha=0.5, ax=ax)\n",
    "ax.set_ylim(0.01, 1000)\n",
    "\n",
    "ax = axes[2]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"loss_zero_norm\", hue=\"m\", palette=\"viridis\", alpha=0.5, ax=ax)\n",
    "ax.set_ylim(0.01, 1000)\n",
    "\n",
    "ax = axes[3]\n",
    "sns.lineplot(data=functional_metrics_averages.loc[functional_metrics_averages.step != 20408], x=\"step\", y=\"ood_loss\", hue=\"m\", palette=\"viridis\", alpha=0.5, ax=ax)\n",
    "\n",
    "for ax in axes.flatten():\n",
    "    ax.set_yscale('log')\n",
    "    ax.set_xscale(\"log\")\n",
    "    ax.set_xlim(100, 400_000)\n",
    "    ax.set_xlabel(\"Step, $t$\")\n",
    "    ax.legend().remove()\n",
    "\n",
    "add_milestones(axes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m20_batch_losses_np = [m20_batch_losses[I:I + 100] for I in range(0, len(m20_batch_losses), 100)]\n",
    "m20_batch_losses_np = np.array(m20_batch_losses_np)\n",
    "m20_batch_losses_np.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m20_batch_losses_cumsum = np.cumsum(m20_batch_losses_np, axis=1)\n",
    "m20_batch_losses_cumavg = m20_batch_losses_cumsum / np.arange(1, 101)\n",
    "\n",
    "m20_batch_losses_df = pd.DataFrame([{\"step\": step, \"loss\": loss, \"b\": b} for step, losses in zip(steps, m20_batch_losses_cumavg) for b, loss in enumerate(losses)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from mpl_toolkits.axes_grid1.inset_locator import inset_axes\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(20, 5))\n",
    "\n",
    "\n",
    "ax = axes[0]\n",
    "# First plot\n",
    "m20_functional_metrics = functional_metrics_averages.loc[functional_metrics_averages.m == 20]\n",
    "sns.lineplot(data=m20_batch_losses_df, x=\"step\", y=\"loss\", hue=\"b\", palette=get_reduced_viridis_palette(100), alpha=0.5, ax=ax, legend=None)\n",
    "sns.lineplot(data=m20_functional_metrics, x=\"step\", y=\"loss\", color=sns.color_palette('bright')[1], alpha=1, ax=ax, linewidth=2, label=\"$L_\\mathrm{val}$\")\n",
    "ax.set_yscale('log')\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\mathrm{Loss}$\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.legend(loc=\"lower left\")\n",
    "\n",
    "ax=axes[1]\n",
    "\n",
    "for b in range(100):\n",
    "    m20_batch_loss_slope = d_dlogt(steps, m20_batch_losses_cumavg[:, b])\n",
    "    m20_batch_losses_df.loc[m20_batch_losses_df.b == b, \"slope\"] = m20_batch_loss_slope\n",
    "\n",
    "m20_val_loss_slopes = d_dlogt(steps, m20_functional_metrics.loss.values)\n",
    "sns.lineplot(data=m20_batch_losses_df, x=\"step\", y=\"slope\", hue=\"b\", palette=get_reduced_viridis_palette(100), alpha=0.5, ax=ax, legend=None)\n",
    "sns.lineplot(data=m20_functional_metrics, x=\"step\", y=m20_val_loss_slopes, color=sns.color_palette('bright')[1], alpha=0.75, ax=ax, linewidth=2, label=\"$\\delta L_\\mathrm{val}/\\delta \\log t$\")\n",
    "# ax.set_yscale('log')\n",
    "ax.set_xscale(\"log\")\n",
    "ax.set_xlabel(\"Step, $t$\")\n",
    "ax.set_ylabel(\"$\\delta \\mathrm{Loss}/\\delta \\log t$\")\n",
    "ax.set_xlim(INIT_X, FINAL_X)\n",
    "ax.set_ylim(-2, 1)\n",
    "ax.legend(loc=\"lower left\")\n",
    "\n",
    "# Create an inset for the second plot\n",
    "\n",
    "# ax_inset = inset_axes(ax, width=\"30%\", height=\"30%\", loc='upper right')\n",
    "\n",
    "# # Second plot (inset)\n",
    "# sns.lineplot(data=m20_batch_losses_df, x=\"step\", y=\"loss\", hue=\"b\", palette=get_reduced_viridis_palette(100), alpha=0.5, ax=ax_inset, legend=None)\n",
    "# sns.lineplot(data=m20_functional_metrics, x=\"step\", y=\"loss\", color=sns.color_palette('deep')[3], alpha=1, ax=ax_inset)\n",
    "# ax_inset.set_xlim(10_000, 250_000)\n",
    "# ax_inset.set_ylim(1.75, 2.25)\n",
    "# ax_inset.set_xscale(\"log\")\n",
    "# ax_inset.set_yscale('log')\n",
    "\n",
    "add_milestones(axes)\n",
    "\n",
    "# plt.show()\n",
    "\n",
    "cbar_ax = fig.add_axes([0.93, 0.15, 0.02, 0.7])  # Adjust as necessary for position and size\n",
    "\n",
    "custom_cmap = mcolors.LinearSegmentedColormap.from_list(\"custom_cmap\", get_reduced_viridis_palette(100))\n",
    "\n",
    "sm = plt.cm.ScalarMappable(cmap=custom_cmap, norm=plt.Normalize(vmin=1, vmax=101), )\n",
    "sm._A = []  # Dummy array for the ScalarMappable. \n",
    "cbar = fig.colorbar(sm, cax=cbar_ax)\n",
    "\n",
    "tick_positions = list(range(1, 101, 10)) + [100] # Positions for each color\n",
    "tick_labels = [\"1\"] +  [f\"${i}$\" for i in range(10, 101, 10)]  # Replace with your labels\n",
    "cbar.set_ticks(tick_positions)\n",
    "cbar.set_ticklabels(tick_labels)\n",
    "cbar.set_label(\"$b$\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fitting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_milestone_indices(steps, milestones):\n",
    "    milestone_indices = []\n",
    "    for step in steps:\n",
    "        # Find the index of the milestone that the current step falls into\n",
    "        index = next((i for i, milestone in enumerate(milestones) if milestone[0] <= step < milestone[1]), None)\n",
    "        milestone_indices.append(index if index is not None else 'Out of defined milestones')\n",
    "    return milestone_indices\n",
    "\n",
    "milestones_of_steps = get_milestone_indices(steps, TRANSITIONS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "run = Run(configs[0])\n",
    "sum(p.numel() for p in run.model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.analysis.baselines import dmmse_predictor, ridge_predictor\n",
    "from icl.regression.tasks import TaskDistribution, DiscreteTaskDistribution, RegressionSequenceDistribution\n",
    "\n",
    "class DMMSE(nn.Module):\n",
    "    def __init__(self, dist: RegressionSequenceDistribution, noise_variance: float, learn_prior: bool = False):\n",
    "        super().__init__()\n",
    "\n",
    "        self.prior = dist.task_distribution\n",
    "        self.noise_variance = nn.Parameter(torch.tensor(noise_variance))\n",
    "\n",
    "        if learn_prior:\n",
    "            self.tasks = nn.Parameter(self.prior.tasks)\n",
    "            self.prior.tasks = self.tasks\n",
    "    \n",
    "    def forward(self, xs, ys):\n",
    "        return dmmse_predictor(xs, ys, self.prior, self.noise_variance)\n",
    "\n",
    "\n",
    "class Ridge(nn.Module):\n",
    "    def __init__(self, noise_variance: float):\n",
    "        super().__init__()\n",
    "\n",
    "        self.noise_variance = nn.Parameter(torch.tensor(noise_variance))\n",
    "\n",
    "\n",
    "    def forward(self, xs, ys):\n",
    "        return ridge_predictor(xs, ys, self.noise_variance)\n",
    "\n",
    "\n",
    "def fit_baseline_predictor(baseline: nn.Module, model: nn.Module, dist: RegressionSequenceDistribution, num_steps: int=1000, lr: float=0.0001, device: str = \"cpu\", batch_size=128, num_examples=8, verbose=True):\n",
    "    optimizer = torch.optim.Adam(baseline.parameters(), lr=lr)\n",
    "    criterion = nn.MSELoss()\n",
    "\n",
    "    # We're fitting just a single parameter (the noise variance)\n",
    "\n",
    "    if verbose:\n",
    "        losses = []\n",
    "        sigmas = []\n",
    "\n",
    "    for step in tqdm(range(num_steps), desc=\"Fitting...\"):\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Get a batch of data\n",
    "        xs, ys = dist.get_batch(batch_size=batch_size, num_examples=num_examples)\n",
    "        xs = xs.to(device)\n",
    "        ys = ys.to(device)\n",
    "\n",
    "        # Get the predictions of the reference model\n",
    "        with torch.no_grad():\n",
    "            yhats = model(xs, ys)\n",
    "\n",
    "        # Get the predictions of the baseline\n",
    "        baseline_preds = baseline(xs, ys)\n",
    "\n",
    "        # Update the baseline to be closer to the reference model\n",
    "        loss = criterion(baseline_preds, yhats)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if verbose:\n",
    "            losses.append(loss.item())\n",
    "            sigmas.append(baseline.noise_variance.item())\n",
    "\n",
    "    if verbose:\n",
    "        fig, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "        plt.suptitle(f\"Baseline fitting {baseline.__class__.__name__} on {dist.task_distribution.__class__.__name__}\")\n",
    "        axes[0].plot(losses)\n",
    "        axes[0].set_title(\"Loss\")\n",
    "        axes[1].plot(sigmas)\n",
    "        axes[1].set_title(\"Noise variance\")\n",
    "        plt.show()\n",
    "\n",
    "    return baseline\n",
    "\n",
    "def eval_delta_predictor(baseline: nn.Module, model: nn.Module, xs, ys, device: str = \"cpu\"):\n",
    "    baseline_preds = baseline(xs, ys)\n",
    "    preds = model(xs, ys)\n",
    "\n",
    "    return nn.MSELoss()(baseline_preds, preds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from infra.utils.iterables import flatten_dict\n",
    "from icl.train import Run\n",
    "import random\n",
    "\n",
    "fit_baseline_results = []\n",
    "\n",
    "lr = 0.01\n",
    "num_steps = 2_00\n",
    "LEARN_PRIOR = False\n",
    "\n",
    "for config, checkpointer in zip(configs, checkpointers):\n",
    "    run = Run(config)\n",
    "    run.model.load_state_dict(checkpointer[-1][\"model\"])\n",
    "\n",
    "    print(\"Evaluating\", run.config.to_slug())\n",
    "\n",
    "    batch_size = run.config.batch_size\n",
    "    num_examples = run.config.task_config.max_examples\n",
    "    \n",
    "    noise = run.config.task_config.noise_variance\n",
    "\n",
    "    learned_dmmse_pretrain = DMMSE(run.pretrain_dist, noise_variance=noise, learn_prior=LEARN_PRIOR)\n",
    "    learned_ridge_pretrain = Ridge(noise_variance=noise)\n",
    "    learned_ridge_true = Ridge(noise_variance=noise)\n",
    "\n",
    "    init_learned_dmmse_pretrain_delta = eval_delta_predictor(learned_dmmse_pretrain, run.model, run.evaluator.pretrain_xs, run.evaluator.pretrain_ys, device=DEVICE)\n",
    "    init_learned_ridge_pretrain_delta = eval_delta_predictor(learned_ridge_pretrain, run.model, run.evaluator.pretrain_xs, run.evaluator.pretrain_ys, device=DEVICE)\n",
    "    init_learned_ridge_true_delta = eval_delta_predictor(learned_ridge_true, run.model, run.evaluator.true_xs, run.evaluator.true_ys, device=DEVICE)\n",
    "    \n",
    "    fit_baseline_predictor(learned_dmmse_pretrain, run.model, run.pretrain_dist, num_steps=num_steps, lr=lr, device=DEVICE, batch_size=batch_size, num_examples=num_examples)\n",
    "    fit_baseline_predictor(learned_ridge_pretrain, run.model, run.pretrain_dist, num_steps=num_steps, lr=lr, device=DEVICE, batch_size=batch_size, num_examples=num_examples)\n",
    "    fit_baseline_predictor(learned_ridge_true, run.model, run.true_dist, num_steps=num_steps, lr=lr, device=DEVICE, batch_size=batch_size, num_examples=num_examples)\n",
    "\n",
    "    learned_dmmse_pretrain_delta = eval_delta_predictor(learned_dmmse_pretrain, run.model, run.evaluator.pretrain_xs, run.evaluator.pretrain_ys, device=DEVICE)\n",
    "    learned_ridge_pretrain_delta = eval_delta_predictor(learned_ridge_pretrain, run.model, run.evaluator.pretrain_xs, run.evaluator.pretrain_ys, device=DEVICE)\n",
    "    learned_ridge_true_delta = eval_delta_predictor(learned_ridge_true, run.model, run.evaluator.true_xs, run.evaluator.true_ys, device=DEVICE)\n",
    "\n",
    "    fit_baseline_results.append({\n",
    "        \"step\": checkpointer.file_ids[-1],\n",
    "        \"config\": run.config.to_slug(),\n",
    "        \"learned_dmmse_pretrain/init_delta\": init_learned_dmmse_pretrain_delta.item(),\n",
    "        \"learned_ridge_pretrain/init_delta\": init_learned_ridge_pretrain_delta.item(),\n",
    "        \"learned_ridge_true/init_delta\": init_learned_ridge_true_delta.item(),\n",
    "        \"learned_dmmse_pretrain/delta\": learned_dmmse_pretrain_delta.item(),\n",
    "        \"learned_ridge_pretrain/delta\": learned_ridge_pretrain_delta.item(),\n",
    "        \"learned_ridge_true/delta\": learned_ridge_true_delta.item(),\n",
    "        \"learned_dmmse_pretrain/delta_delta\": learned_dmmse_pretrain_delta.item() - init_learned_dmmse_pretrain_delta.item(),\n",
    "        \"learned_ridge_pretrain/delta_delta\": learned_ridge_pretrain_delta.item() - init_learned_ridge_pretrain_delta.item(),        \n",
    "        \"learned_ridge_true/delta_delta\": learned_ridge_true_delta.item() - init_learned_ridge_true_delta.item(),\n",
    "        \"learned_dmmse_pretrain/noise_variance\": learned_dmmse_pretrain.noise_variance.item(),\n",
    "        \"learned_ridge_pretrain/noise_variance\": learned_ridge_pretrain.noise_variance.item(),\n",
    "        \"learned_ridge_true/noise_variance\": learned_ridge_true.noise_variance.item(),\n",
    "        **flatten_dict(run.config.task_config.model_dump(), flatten_lists=True)\n",
    "    })\n",
    "\n",
    "    pp(fit_baseline_results[-1])\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "baseline_fits_df = pd.DataFrame(fit_baseline_results)\n",
    "\n",
    "# Create 2x3 subplots\n",
    "fig, axes = plt.subplots(3, 2, figsize=(12, 18))\n",
    "\n",
    "plt.suptitle(\n",
    "    \"L2-H4-K8-D4-err0.125-dmlp64-dembed64-seeds0-1-2-3-n128000000-lr0.01-B256-T500000@t=499999\"\n",
    ")\n",
    "\n",
    "# Define the labels for rows\n",
    "row_labels = ['dmmse_pretrain', 'ridge_pretrain', 'ridge_true']\n",
    "\n",
    "# Loop through the rows\n",
    "for token, row_label in enumerate(row_labels):\n",
    "    \n",
    "    # First column: init_delta and delta\n",
    "    ax1 = axes[token, 0]\n",
    "    baseline_fits_df.plot(x='num_tasks', y=f'learned_{row_label}/init_delta', ax=ax1, label=f'{row_label} init_delta')\n",
    "    baseline_fits_df.plot(x='num_tasks', y=f'learned_{row_label}/delta', ax=ax1, label=f'{row_label} delta')\n",
    "    ax1.set_title(f\"{row_label} init_delta and delta\")\n",
    "    ax1.set_xlabel('num_tasks')\n",
    "    ax1.set_ylabel('Value')\n",
    "    \n",
    "    # Second column: noise_variance\n",
    "    ax2 = axes[token, 1]\n",
    "    baseline_fits_df.plot(x='num_tasks', y=f'learned_{row_label}/noise_variance', ax=ax2, label=f'{row_label} noise_variance')\n",
    "    ax2.axhline(y=0.125, color='r', linestyle='-', label='True noise_variance')\n",
    "\n",
    "    ax2.legend()\n",
    "\n",
    "    ax2.set_title(f\"{row_label} noise_variance\")\n",
    "    ax2.set_xlabel('num_tasks')\n",
    "    ax2.set_ylabel('Noise Variance')\n",
    "\n",
    "for ax in axes.flat:\n",
    "    ax.set_xscale(\"log\")\n",
    "\n",
    "plt.tight_layout(rect=[0.1, 0.1, 1, 1])\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PCA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def iter_enumerated_models(model, checkpointer, verbose=False):\n",
    "    for file_id in tqdm(checkpointer.file_ids, desc=\"Iterating over checkpoints\", disable=not verbose):\n",
    "        model.load_state_dict(checkpointer.load_file(file_id)[\"model\"])\n",
    "        yield file_id, model\n",
    "\n",
    "def iter_models(model, checkpointer, verbose=False):\n",
    "    for file_id in tqdm(checkpointer.file_ids, desc=\"Iterating over checkpoints\", disable=not verbose):\n",
    "        model.load_state_dict(checkpointer.load_file(file_id)[\"model\"])\n",
    "        yield model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, Iterable, Tuple\n",
    "from sklearn.decomposition import PCA\n",
    "from collections import defaultdict\n",
    "from icl.analysis.hooks import hook\n",
    "import numpy as np\n",
    "from icl.analysis.utils import map_evals_over_checkpoints, get_unique_run\n",
    "from icl.train import Run\n",
    "from infra.utils.tensors import convert_tensor, ReturnTensor\n",
    "\n",
    "\n",
    "def extract_activations_over_checkpoints(models: Iterable[nn.Module], xs, ys, *paths, return_type: ReturnTensor=\"np\"):\n",
    "    def eval_activations(model):\n",
    "        hooked_model = hook(model, *paths)\n",
    "        return {k: convert_tensor(v, return_type) for k, v in hooked_model.run_with_cache(xs, ys)[1].items() if k in paths and v is not None}\n",
    "    \n",
    "    for model in models:\n",
    "        yield eval_activations(model)\n",
    "\n",
    "\n",
    "def get_vectorized_activations_trace(models: Iterable[nn.Module], xs, ys, *paths):\n",
    "    evals: Dict[str, list] = defaultdict(list)\n",
    "    \n",
    "    for activations in extract_activations_over_checkpoints(models, xs, ys, *paths):\n",
    "        for path, activation in activations.items():\n",
    "            evals[path].append(activation)\n",
    "\n",
    "    return {\n",
    "        k: np.array(v).reshape(len(v), -1) for k, v in evals.items()\n",
    "    }\n",
    "\n",
    "\n",
    "def get_pca_activations_trace(models: Iterable[nn.Module], xs, ys, *paths, num_components=3) -> Dict[str, Tuple[PCA, np.ndarray]]:\n",
    "    results = {}\n",
    "\n",
    "    for path, activations in get_vectorized_activations_trace(models, xs, ys, *paths).items():\n",
    "        pca = PCA(n_components=num_components)\n",
    "        activations_reduced = pca.fit_transform(activations)\n",
    "        results[path] = pca, activations_reduced\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# demo = Run(configs[2])\n",
    "# demo_models = iter_models(demo.model, demo.checkpointer, verbose=True)\n",
    "\n",
    "# demo_logits_pca_3, demo_logits_reduced_3  = get_pca_activations_trace(\n",
    "#     demo_models, \n",
    "#     demo.evaluator.pretrain_xs, \n",
    "#     demo.evaluator.pretrain_ys, \n",
    "#     \"token_sequence_transformer\",\n",
    "#     num_components=3\n",
    "# )['token_sequence_transformer']\n",
    "\n",
    "# steps = demo.checkpointer.file_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "def plot_sample_evolution(steps, samples, title=\"Sample Evolution in 2D Plane\", num_points_to_label=10, save: Optional[str] = None, ax: Optional = None, connect_dots=False):\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(figsize=(15, 8))\n",
    "    \n",
    "    # Main plot\n",
    "    sc = ax.scatter(samples[:, 0], samples[:, 1], c=steps, cmap='viridis', s=50, alpha=0.6)\n",
    "\n",
    "    if connect_dots:\n",
    "        ax.plot(samples[:, 0], samples[:, 1], c='black', alpha=0.2)\n",
    "\n",
    "    plt.colorbar(sc, ax=ax, label='Steps')\n",
    "    \n",
    "    # Label some points\n",
    "    total_samples = len(samples)\n",
    "    step = total_samples // num_points_to_label\n",
    "    for i in range(0, total_samples, step):\n",
    "        sample_step = steps[i]\n",
    "        ax.text(samples[i, 0], samples[i, 1], str(sample_step), fontsize=12, ha='right', va='bottom')\n",
    "        \n",
    "    ax.set_xlabel('Feature 1')\n",
    "    ax.set_ylabel('Feature 2')\n",
    "    ax.set_title(title)\n",
    "    \n",
    "    # Inset for explained variance at the bottom right corner with slight transparency\n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "\n",
    "        plt.savefig(save)\n",
    "\n",
    "\n",
    "def plot_explained_variance(pca, title=\"Explained Variance\", ax: Optional[plt.Axes] = None):\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(figsize=(15, 8))\n",
    "\n",
    "    ax.bar(range(len(pca.explained_variance_ratio_)), pca.explained_variance_ratio_)\n",
    "\n",
    "    for i, ratio in enumerate(pca.explained_variance_ratio_):\n",
    "        ax.text(i, ratio, f\"{ratio:.2f}\", fontsize=12, ha='center', va='bottom')\n",
    "\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel('Component')\n",
    "    ax.set_ylabel('Variance')\n",
    "\n",
    "\n",
    "def plot_sample_evolution_with_inset(steps, samples, pca, title=\"Sample Evolution in 2D Plane\", num_points_to_label=10, save: Optional[str] = None, ax: Optional = None, connect_dots=False):\n",
    "    if ax is None:\n",
    "        fig, ax = plt.subplots(figsize=(15, 8))\n",
    "    \n",
    "    plot_sample_evolution(steps, samples, title=title, num_points_to_label=num_points_to_label, ax=ax, connect_dots=connect_dots)\n",
    "\n",
    "    axins = ax.inset_axes([0.7, 0.05, 0.25, 0.25])  # x, y, width, height\n",
    "    axins.patch.set_alpha(0.5)\n",
    "    plot_explained_variance(pca, ax=axins)\n",
    "    \n",
    "    # Inset for explained variance at the bottom right corner with slight transparency\n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "\n",
    "        plt.savefig(save)\n",
    "\n",
    "    \n",
    "def plot_multiple_slices(steps, samples, pca, title=\"Sample Evolution in 2D Plane\", num_points_to_label=10, save: Optional[str] = None, ax: Optional = None, connect_dots=False):\n",
    "    num_pca_components = samples.shape[-1]\n",
    "    num_rows = num_pca_components - 1\n",
    "    fig, ax = plt.subplots(num_rows, num_rows, figsize=(20, 20))\n",
    "\n",
    "    fig.suptitle(title)\n",
    "\n",
    "    for i in range(num_pca_components):\n",
    "        for j in range(i):\n",
    "            sc = ax[i-1, j].scatter(samples[:, i], samples[:, j], c=steps, cmap='viridis', s=50, alpha=0.6)\n",
    "            ax[i-1, j].set_xlabel(f'Feature {i}')\n",
    "            ax[i-1, j].set_ylabel(f'Feature {j}')\n",
    "            ax[i-1, j].set_title(f'Feature {i} vs Feature {j}')\n",
    "\n",
    "            if connect_dots:\n",
    "                ax[i-1, j].plot(samples[:, i], samples[:, j], c='black', alpha=0.2)\n",
    "\n",
    "            # Label some points\n",
    "            total_samples = len(samples)\n",
    "            step = total_samples // num_points_to_label\n",
    "            for k in range(0, total_samples, step):\n",
    "                sample_step = steps[k]\n",
    "                ax[i-1, j].text(samples[k, i], samples[k, j], str(sample_step), fontsize=12, ha='right', va='bottom')\n",
    "\n",
    "        for j in range(i + 1, num_rows):\n",
    "            ax[i, j].axis('off')\n",
    "\n",
    "\n",
    "    ax[0, -1].axis('on')\n",
    "    plot_explained_variance(pca, ax=ax[0, -1])\n",
    "\n",
    "    plt.colorbar(sc, ax=ax[0, -1], label='Steps')\n",
    "    plt.tight_layout()\n",
    "\n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "\n",
    "        plt.savefig(save)\n",
    "\n",
    "# plot_multiple_slices(steps, demo_logits_reduced_3, demo_logits_pca_3, title=demo.config.to_latex(), connect_dots=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for config, checkpointer in zip(configs, checkpointers):\n",
    "    run = Run(config)\n",
    "    _steps = checkpointer.file_ids\n",
    "\n",
    "    _pca, _logits_reduced = get_pca_activations_trace(\n",
    "        iter_models(run.model, run.checkpointer, verbose=True), \n",
    "        run.evaluator.pretrain_xs, \n",
    "        run.evaluator.pretrain_ys, \n",
    "        \"token_sequence_transformer\",\n",
    "        num_components=3\n",
    "    )['token_sequence_transformer']\n",
    "    \n",
    "    plot_multiple_slices(\n",
    "        _steps, \n",
    "        _logits_reduced, \n",
    "        _pca, \n",
    "        connect_dots=True, \n",
    "        title=config.to_latex(), \n",
    "        save=FIGURES / (\"pca3-logits-\" + config.to_slug(delimiter=\"-\") + \".png\")\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Attention Entropies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Union\n",
    "from torchtyping import TensorType\n",
    "from infra.utils.iterables import map_nested\n",
    "\n",
    "from infra.utils.iterables import flatten_dict\n",
    "\n",
    "from icl.train import Run\n",
    "\n",
    "def compute_attention_entropies(attn: TensorType[\"B\", \"H\", \"2K\", \"2K\"]):\n",
    "    \"\"\"\n",
    "    Computes the entropy of each token in each head, averaged across the batch, \n",
    "    then averages this over heads. \n",
    "\n",
    "    \"\"\"\n",
    "    \n",
    "    # Threshold attention weights to avoid log(0)\n",
    "    log_attention = torch.where(attn > 0, torch.log(attn), torch.tensor(0.0).to(attn.device))\n",
    "    entropy_per_token = - torch.sum(attn * log_attention, dim=-1).mean(dim=0).squeeze(-1) # TensorType[\"H\", \"2K\"]\n",
    "\n",
    "    num_heads, num_tokens = entropy_per_token.shape\n",
    "\n",
    "    entropy_per_head = entropy_per_token.mean(dim=-1) # TensorType[\"H\"]\n",
    "    entropy = entropy_per_head.mean() # TensorType[]    \n",
    "    \n",
    "    # Each token computes entropy over a variable context length, so we normalize by the maximum possible entropy\n",
    "    # for a token with a fixed context length.\n",
    "\n",
    "    max_entropy_per_token = torch.log2(torch.arange(1, num_tokens + 1).to(attn.device)) # TensorType[\"H\", \"2K\"]\n",
    "    max_entropy_per_token[0] = 1. # Special case for the first token to avoid dividing by 0\n",
    "\n",
    "    entropy_per_token_normalized = entropy_per_token / max_entropy_per_token\n",
    "    entropy_per_head_normalized = entropy_per_token_normalized.mean(dim=-1) # TensorType[\"H\"]\n",
    "    entropy_normalized = entropy_per_head_normalized.mean() # TensorType[]    \n",
    "\n",
    "    results: Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]] = {\"mean\": entropy, \"mean_normalized\": entropy_normalized}\n",
    "\n",
    "    for i in range(num_heads):\n",
    "        head_results = {\"mean\": entropy_per_head[i], \"mean_normalized\": entropy_per_head_normalized[i]}\n",
    "\n",
    "        for j in range(num_tokens):\n",
    "            head_results[f\"token_{j}\"] = entropy_per_token[i, j]\n",
    "            head_results[f\"token_{j}_normalized\"] = entropy_per_token_normalized[i, j]\n",
    "\n",
    "        results[f\"head_{i}\"] = head_results\n",
    "\n",
    "    return map_nested(lambda x: convert_tensor(x, \"np\"), results)\n",
    "\n",
    "\n",
    "def get_attention_entropies_trace(\n",
    "    steps: List[int],\n",
    "    models: Iterable[nn.Module],\n",
    "    xs: torch.Tensor,\n",
    "    ys: torch.Tensor,\n",
    "    **paths,\n",
    "):\n",
    "    results = defaultdict(list)\n",
    "    reverse_paths = {v: k for k, v in paths.items()}\n",
    "\n",
    "    for activations in extract_activations_over_checkpoints(models, xs, ys, *paths.values(), return_type=\"pt\"):\n",
    "        for k, v in activations.items():\n",
    "            path = reverse_paths[k]\n",
    "            results[path].append(compute_attention_entropies(v))\n",
    "\n",
    "    values = []\n",
    "\n",
    "    for i in range(len(steps)):\n",
    "        value = {}\n",
    "\n",
    "        for block in results.keys():\n",
    "            value[block] = results[block][i]\n",
    "        \n",
    "        value[\"step\"] = steps[i]\n",
    "        values.append(flatten_dict(value, flatten_lists=True))\n",
    "\n",
    "    return pd.DataFrame(values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_attention_patterns(df: pd.DataFrame, num_blocks: int, num_heads: int, num_tokens: int, title=\"\", save: Optional[str] = None, normalized=False, figsize=(20, 25), logx=False, logy=False):\n",
    "    fig = plt.figure(figsize=figsize)\n",
    "    plt.suptitle(title)\n",
    "\n",
    "    num_cols = num_blocks * 2\n",
    "    num_rows = 1 + 1 + num_heads\n",
    "\n",
    "    suffix = \"\" if not normalized else \"_normalized\"\n",
    "    suffix_title = \"\" if not normalized else \" (Normalized)\"\n",
    "\n",
    "    fig.set_facecolor('white')\n",
    "\n",
    "    # Create subplot for mean entropy of first two blocks\n",
    "    ax0 = plt.subplot2grid((num_rows, num_cols), (0, 0), colspan=num_cols)\n",
    "    block_cmap = sns.color_palette(\"viridis\", num_blocks)\n",
    "\n",
    "    for b in range(num_blocks):\n",
    "        ax0.plot(df.step, df[f\"block_{b}/mean{suffix}\"], label=f\"block_{b}\", color=block_cmap[b])\n",
    "\n",
    "    ax0.set_title(\"Blocks\")\n",
    "    ax0.set_xlabel(\"Step\")\n",
    "    ax0.set_ylabel(f\"Entropy{suffix_title}\")\n",
    "    ax0.legend()\n",
    "\n",
    "    # Create subplots for each block, showing entropy in different heads\n",
    "    ax1 = [plt.subplot2grid((num_rows, num_cols), (1, i*2), colspan=2) for i in range(num_blocks)]\n",
    "    head_cmap = sns.color_palette(\"viridis\", num_heads)\n",
    "    \n",
    "    for b in range(num_blocks):\n",
    "        ax1[b].set_title(f\"Block {b}\")\n",
    "        ax1[b].set_xlabel(\"Step\")\n",
    "        ax1[b].set_ylabel(f\"Entropy{suffix_title}\")\n",
    "        for h in range(num_heads):\n",
    "            series = df[f\"block_{b}/head_{h}/mean{suffix}\"]\n",
    "            ax1[b].plot(df.step, series, label=f\"Head {h}\", color=head_cmap[h])\n",
    "\n",
    "    ax1[0].legend()\n",
    "\n",
    "    # Create subplots for each head in each block, detailing entropy for each token\n",
    "    ax2 = [plt.subplot2grid((num_rows, num_cols), (i//(num_cols) + 2, i%(num_cols))) for i in range(num_heads * num_blocks * 2)]\n",
    "    ax_idx = 0\n",
    "    token_cmap = sns.color_palette(\"viridis\", num_tokens)\n",
    "\n",
    "\n",
    "    for h in range(num_heads):\n",
    "        for b in range(num_blocks):\n",
    "            for x_or_y in (1, 0):\n",
    "                ax2[ax_idx].set_title(f\"Block {b} Head {h}\")\n",
    "                ax2[ax_idx].set_xlabel(\"Step\")\n",
    "                ax2[ax_idx].set_ylabel(f\"Entropy{suffix_title}\")\n",
    "\n",
    "                for t in range(1-int(x_or_y), num_tokens, 2):\n",
    "                    series = df[f\"block_{b}/head_{h}/token_{t}{suffix}\"]\n",
    "                    ax2[ax_idx].plot(df.step, series, label=f\"Token {t}\", color=token_cmap[t])\n",
    "                    \n",
    "                ax_idx += 1\n",
    "\n",
    "    ax2[0].legend()\n",
    "    ax2[1].legend()\n",
    "\n",
    "    for ax in [ax0, *ax1, *ax2]:\n",
    "        if logx:\n",
    "            ax.set_xscale(\"log\")\n",
    "        if logy:\n",
    "            ax.set_yscale(\"log\")\n",
    "\n",
    "    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
    "\n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "        plt.savefig(save)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# demo = Run(configs[2])\n",
    "\n",
    "# num_blocks = demo.config.task_config.num_layers\n",
    "# num_heads = demo.config.task_config.num_heads\n",
    "# num_tokens = demo.config.task_config.max_examples * 2\n",
    "\n",
    "# df = get_attention_entropies_trace(\n",
    "#     demo.checkpointer.file_ids,\n",
    "#     iter_models(demo.model, demo.checkpointer, verbose=True), \n",
    "#     demo.evaluator.pretrain_xs, \n",
    "#     demo.evaluator.pretrain_ys, \n",
    "#     **{f\"block_{b}\": f\"token_sequence_transformer.blocks.{b}.attention.attention_softmax\" for b in range(num_blocks)}\n",
    "# )\n",
    "\n",
    "# demo_attn_entropy_slug = \"attn-S-\" + demo.config.to_slug(delimiter=\"-\")\n",
    "\n",
    "for normalized in (True, False):\n",
    "    plot_attention_patterns(\n",
    "        subdf, \n",
    "        num_blocks=num_blocks, \n",
    "        num_heads=num_heads, \n",
    "        num_tokens=num_tokens, \n",
    "        title=demo.config.to_latex(), \n",
    "        save=FIGURES / (demo_attn_entropy_slug + \".png\"),\n",
    "        figsize=(25, 25),\n",
    "        normalized=normalized\n",
    "    )\n",
    "\n",
    "# df.to_csv(ANALYSIS / (demo_attn_entropy_slug + \".csv\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for config, checkpointer in zip(configs, checkpointers):\n",
    "    run = Run(config)\n",
    "    \n",
    "    num_blocks = run.config.task_config.num_layers\n",
    "    num_heads = run.config.task_config.num_heads\n",
    "    num_tokens = run.config.task_config.max_examples * 2\n",
    "\n",
    "    subdf = get_attention_entropies_trace(\n",
    "        checkpointer.file_ids,\n",
    "        iter_models(run.model, checkpointer, verbose=True), \n",
    "        run.evaluator.pretrain_xs, \n",
    "        run.evaluator.pretrain_ys, \n",
    "        **{f\"block_{b}\": f\"token_sequence_transformer.blocks.{b}.attention.attention_softmax\" for b in range(num_blocks)}\n",
    "    )\n",
    "    \n",
    "    slug = \"attn-S-\" + run.config.to_slug(delimiter=\"-\")\n",
    "\n",
    "    plot_attention_patterns(\n",
    "        subdf, \n",
    "        num_blocks=num_blocks, \n",
    "        num_heads=num_heads, \n",
    "        num_tokens=num_tokens, \n",
    "        title=run.config.to_latex(), \n",
    "        save=FIGURES / (slug + \".png\"),\n",
    "        figsize=(25, 25),\n",
    "        normalized=True\n",
    "    )\n",
    "\n",
    "    subdf.to_csv(ANALYSIS / (slug + \".csv\"))\n",
    "\n",
    "# os.system('say \"Your program has finished.\"')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.train import Run\n",
    "demo = Run(configs[2])\n",
    "attn_weights = demo.model.token_sequence_transformer.blocks[0].attention.attention.weight\n",
    "attn_weights.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "numel_per_layer = attn_weights.numel()\n",
    "\n",
    "def num_params_to_gb(num: int):\n",
    "    return f\"{num * (32 / 8) / (10 ** 9):.2f} Gb\"\n",
    "\n",
    "for num_blocks in [2, 4, 8]:\n",
    "    for num_heads in [2, 4]:\n",
    "        numel_per_head = numel_per_layer // num_heads\n",
    "\n",
    "        within_head_cov_size = (numel_per_head ** 2)  * num_heads * num_blocks\n",
    "        between_head_cov_size = (numel_per_head ** 2) * num_heads * num_heads * (num_blocks-1)\n",
    "\n",
    "        full_cov_size = (numel_per_layer * num_blocks) ** 2\n",
    "\n",
    "        reduction = full_cov_size - within_head_cov_size - between_head_cov_size\n",
    "\n",
    "        print(f\"\\nL{num_blocks}H{num_heads}\")\n",
    "        print(\"Full:\", f\"{full_cov_size:,} ({num_params_to_gb(full_cov_size)})\")\n",
    "        print(\"Within heads:\", f\"{within_head_cov_size:,} ({num_params_to_gb(within_head_cov_size)})\")\n",
    "        print(\"Between heads:\", f\"{between_head_cov_size:,} ({num_params_to_gb(between_head_cov_size)})\")\n",
    "        print(\"Reduction:\", f\"-{reduction:,} (-{reduction/full_cov_size * 100:.2f}%)\")\n",
    "\n",
    "# attn_weights.numel(), f\"{(32 // 8 * (attn_weights.numel() * 2 ) ** 2):,}\", attn_weights.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_attn_weights(W: torch.Tensor, num_heads: int, embed_dim: int, head_size: int):\n",
    "    W_split = W.view((embed_dim, num_heads, head_size * 3))\n",
    "    \n",
    "    for h in range(num_heads):\n",
    "        yield tuple(W_split[:, h, i*head_size:(i+1)*head_size] for i in range(3))\n",
    "\n",
    "\n",
    "def plot_attn_weights(W: torch.Tensor, num_heads: int, embed_dim: int, head_size: int, subtitles=(\"$W_Q^{(h)}$\", \"$W_K^{(h)}$\", \"$W_V^{(h)}$\"), title=\"\", save: Optional[str] = None):\n",
    "    heads = list(split_attn_weights(W, num_heads, embed_dim, head_size))\n",
    "\n",
    "    fig, axs = plt.subplots(num_heads, 3, figsize=(25, 10))\n",
    "    plt.suptitle(title)\n",
    "\n",
    "    for h, head in enumerate(heads):\n",
    "        axs[h, 0].set_ylabel(f\"Head {h}\\nHead Size\")\n",
    "\n",
    "        for i, mat in enumerate(head):\n",
    "            axs[h, i].matshow(mat.detach().cpu().numpy().T, cmap='viridis') \n",
    "\n",
    "    for i, subtitle in enumerate(subtitles):\n",
    "        axs[0, i].set_title(subtitle)\n",
    "        axs[-1, i].set_xlabel(\"Embedding Dimension\")\n",
    "\n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "        plt.savefig(save)\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_attn_head_weights(head: torch.Tensor, embed_dim, head_size: int, title=\"\", subtitles=(\"$W_Q$\", \"$W_K$\", \"$W_V$\"), save: Optional[str] = None):\n",
    "    head_Ex3c = head.view((embed_dim, head_size * 3))\n",
    "    q, k, v = tuple(head_Ex3c[:, i*head_size:(i+1)*head_size].detach().cpu().numpy() for i in range(3))\n",
    "\n",
    "    fig, ax = plt.subplots(1, 3, figsize=(30, 3.5))\n",
    "    plt.suptitle(title)\n",
    "\n",
    "    for i, (mat, subtitle) in enumerate(zip((q, k, v), subtitles)):\n",
    "        ax[i].set_title(subtitle)\n",
    "        ax[i].matshow(mat.T, cmap='viridis')\n",
    "        ax[i].set_xlabel(\"Embedding Dimension\")\n",
    "        ax[i].set_ylabel(\"Head Size\")\n",
    "\n",
    "    plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n",
    "\n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "        plt.savefig(save)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_attn_eigencomponents(evecs, evals, slug: Optional[str] = None):\n",
    "    for i in range(1, 1 + len(evals)):\n",
    "        attn0, attn1 = evecs[:evecs.shape[0]//2, -i], evecs[evecs.shape[0]//2:, -i]\n",
    "\n",
    "        for layer, attn in enumerate((attn0, attn1)):\n",
    "            plot_attn_weights(\n",
    "                torch.Tensor(attn), \n",
    "                num_heads=4,\n",
    "                embed_dim=64, \n",
    "                head_size=16, \n",
    "                title=f\"Eigenvector {i-1} of covariance matrix within attention layer 0 ($\\lambda_{i-1}={evals[-i]}$)\",\n",
    "                subtitles=(f\"$u_{{Q,{i-1}}}^{{({layer})}}$\", f\"$u_{{K,{i-1}}}^{{({layer})}}$\", f\"$u_{{V,{i-1}}}^{{({layer})}}$\"),\n",
    "                save=(FIGURES / (f\"cov-attn{layer}-evec{i-1}-\" + slug + \".png\") if slug else None)\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_attn_weights(attn0, 4, 64, 16, title=\"Attention layer 0\")\n",
    "\n",
    "num_heads = 4\n",
    "attn0_view = attn0.view((64, num_heads, 16 * 3))\n",
    "heads = [attn0_view[:, h, :] for h in range(num_heads)]\n",
    "full_head_size = 16 * 3 * 64\n",
    "pseudo_cov = heads[0].reshape((full_head_size, 1)) * heads[1].reshape((1, full_head_size)) \n",
    "head_evals, head_evecs = eigsh(pseudo_cov.detach().cpu().numpy(), k=3, which=\"LM\")\n",
    "del pseudo_cov\n",
    "\n",
    "\n",
    "print(head_evals)\n",
    "plot_attn_head_weights(\n",
    "    torch.Tensor(head_evecs[:, -1]), \n",
    "    64, \n",
    "    16, \n",
    "    title=\"Principal eigenvalue of covariance matrix within head 1\",\n",
    "    subtitles=(\"$u_{Q,1}^{(1)}$\", \"$u_{K,1}^{(1)}$\", \"$u_{V,1}^{(1)}$\")   \n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from icl.analysis.sample import make_slt_evals\n",
    "\n",
    "def generate_slt_observables(\n",
    "    steps: List[int],\n",
    "    models: Iterable[nn.Module],\n",
    "    xs: torch.Tensor,\n",
    "    ys: torch.Tensor,\n",
    "    **kwargs\n",
    "):\n",
    "    trainset = torch.utils.data.TensorDataset(xs, ys)\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=len(xs))\n",
    "    slt_evals = make_slt_evals(\n",
    "        dataset=trainset,\n",
    "        loader=trainloader,\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "    for step, model in zip(steps, models):\n",
    "        yield step, slt_evals(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(checkpointer.file_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_learning_coeff_over_time(steps, lcs, lc_stds, title=\"\", save: Optional[str] = None):\n",
    "\n",
    "    # Initialize the figure\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(18, 12))\n",
    "    ax.set_title(title)\n",
    "\n",
    "    # Plot mean values as a line\n",
    "    ax.plot(steps, lcs, 'o-', linewidth=2)\n",
    "    \n",
    "    # Add shaded area for error\n",
    "    ax.fill_between(steps, lcs - lc_stds, lcs + lc_stds, color='gray', alpha=0.4)\n",
    "\n",
    "    # Labels and scales\n",
    "    ax.set_xlabel(\"Steps\")\n",
    "    ax.set_ylabel(r\"$\\hat\\lambda$\")\n",
    "    \n",
    "    fig.tight_layout()\n",
    "\n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "        plt.savefig(save)\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def plot_cov_evals_over_time(steps, *eval_traces, title=\"\", save: Optional[str] = None):\n",
    "\n",
    "    # Initialize the figure\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(18, 12))\n",
    "    ax.set_title(title)\n",
    "\n",
    "    # Plot mean values as a line\n",
    "    for i, eval_trace in enumerate(eval_traces):\n",
    "        ax.plot(steps, eval_trace, 'o-', label=f\"Eigenvalue {i}\", linewidth=2)\n",
    "    \n",
    "    # Labels and scales\n",
    "    ax.set_xlabel(\"Steps\")\n",
    "    ax.set_ylabel(r\"$\\hat\\lambda$\")\n",
    "    \n",
    "    # Show legend\n",
    "    ax.legend()\n",
    "\n",
    "    fig.tight_layout()\n",
    "    \n",
    "    if save:\n",
    "        parent_dir = os.path.dirname(save)\n",
    "        if not os.path.exists(parent_dir):\n",
    "            os.makedirs(parent_dir)\n",
    "        plt.savefig(save)\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "from scipy.sparse.linalg import eigsh\n",
    "\n",
    "\n",
    "\n",
    "for log2_M in MS:\n",
    "    wandb.init(entity=WANDB_ENTITY, project=\"icl\")\n",
    "\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "\n",
    "    log2_m = int(np.log2(log2_M))\n",
    "    config, checkpointer = configs[log2_m], checkpointers[log2_m]\n",
    "    run = Run(config)\n",
    "\n",
    "    xs, ys = run.evaluator.pretrain_xs, run.evaluator.pretrain_ys\n",
    "    trainset = torch.utils.data.TensorDataset(xs, ys)\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=len(xs))\n",
    "    observables_over_time = []\n",
    "    \n",
    "    slt_evals = make_slt_evals(\n",
    "        dataset=trainset,\n",
    "        loader=trainloader,\n",
    "        cores=1,\n",
    "        lr=1e-5,\n",
    "        num_draws=100,\n",
    "        elasticity=1.,\n",
    "        num_chains=20,\n",
    "        device=\"cuda\",\n",
    "        covariance_paths=[\n",
    "            f\"token_sequence_transformer.blocks.{b}.attention.attention\"\n",
    "            for b in range(run.config.task_config.num_layers)\n",
    "        ],\n",
    "    )\n",
    "\n",
    "    slug = run.config.to_slug(delimiter=\"-\")\n",
    "\n",
    "    min_step = -1\n",
    "\n",
    "    if os.path.exists(ANALYSIS / f\"cov-tmp-{slug}.pt\"):\n",
    "        min_step, observables = torch.load(ANALYSIS / f\"cov-tmp-{slug}.pt\")\n",
    "        print(f\"Loaded observables from previous step {min_step} from {ANALYSIS / f'cov-tmp-{slug}.pt'}\")\n",
    "    \n",
    "    for step in STEPS:\n",
    "        # if step > min_step:\n",
    "        run.model.load_state_dict(checkpointer.load_file(step)[\"model\"])\n",
    "        observables = slt_evals(run.model)\n",
    "        torch.save((step, observables), ANALYSIS / f\"cov-tmp-{slug}.pt\")\n",
    "\n",
    "        cov = observables.pop(\"covariance\")\n",
    "        evals, evecs = eigsh(cov, k=K, which='LM')\n",
    "\n",
    "        for token in range(1, 1+K):\n",
    "            observables[f\"cov-eval/{token-1}\"] = evals[-token]\n",
    "\n",
    "        observables_over_time.append(observables)\n",
    "        del cov\n",
    "        pp(observables)\n",
    "        plot_attn_eigencomponents(evecs, evals, slug=slug + f\"@t={step}\")\n",
    "\n",
    "    plot_cov_evals_over_time(\n",
    "        STEPS,\n",
    "        *[[o[f\"cov-eval/{k}\"] for o in observables_over_time] for k in K],\n",
    "        title=run.config.to_latex(),\n",
    "        save=FIGURES / f\"cov-eval-of-t-{slug}.png\"\n",
    "    )\n",
    "\n",
    "    plot_learning_coeff_over_time(\n",
    "        STEPS,\n",
    "        [o[\"mean\"] for o in observables_over_time],\n",
    "        [o[\"std\"] for o in observables_over_time],\n",
    "        title=run.config.to_latex(),\n",
    "        save=FIGURES / f\"lc-of-t-{slug}.png\"\n",
    "    )\n",
    "\n",
    "    observables_df = pd.DataFrame(observables_over_time)\n",
    "    observables_df.to_csv(ANALYSIS / f\"cov/cov-{slug}.csv\")\n",
    "    os.remove(ANALYSIS / f\"cov-tmp-{slug}.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(observables_over_time[0].keys())\n",
    "plot_cov_evals_over_time(\n",
    "    STEPS,\n",
    "    [o[\"cov-eval/0\"] for o in observables_over_time],\n",
    "    [o[\"cov-eval/1\"] for o in observables_over_time],\n",
    "    title=run.config.to_latex(),\n",
    "    save=FIGURES / f\"cov-eval-of-t-{slug}.png\"\n",
    ")\n",
    "\n",
    "plot_learning_coeff_over_time(\n",
    "    STEPS,\n",
    "    np.array([o[\"mean\"] for o in observables_over_time]),\n",
    "    np.array([o[\"std\"] for o in observables_over_time]),\n",
    "    title=run.config.to_latex(),\n",
    "    save=FIGURES / f\"lc-of-t-{slug}.png\"\n",
    ")\n",
    "\n",
    "observables_df = pd.DataFrame(observables_over_time)\n",
    "observables_df.to_csv(ANALYSIS / f\"cov/cov-{slug}.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()\n",
    "del run, trainset, trainloader\n",
    "del slt_evals\n",
    "import gc\n",
    "gc.collect()\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "covariances = observables.pop(\"covariances\")\n",
    "evals, evecs = eigsh(covariances, k=3, which='LM')\n",
    "\n",
    "for token, (eval, evec) in enumerate(zip(evals, evecs.T)):\n",
    "    slug = f\"cov-u{token}\" + run.config.to_slug(delimiter=\"-\") + f\"@t={step}\"\n",
    "    attn0, attn1 = evec.split(64 * 16 * 3)\n",
    "    \n",
    "\n",
    "# TODO: Need to rename the new files otherwise you can't tell easily tell what step they come from.\n",
    "\n",
    "\n",
    "os.system('say \"Your program has finished.\"')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "from scipy.sparse.linalg import eigsh\n",
    "\n",
    "\n",
    "for config, checkpointer in zip(configs, checkpointers):\n",
    "    run = Run(config)\n",
    "    observables_over_time = []\n",
    "    \n",
    "    for step, observables in generate_slt_observables(\n",
    "        checkpointer.file_ids,\n",
    "        iter_models(run.model, checkpointer, verbose=True), \n",
    "        run.evaluator.pretrain_xs, \n",
    "        run.evaluator.pretrain_ys, \n",
    "        cores=4,\n",
    "        lr=1e-5,\n",
    "        num_draws=100,\n",
    "        elasticity=1.,\n",
    "        num_chains=20,\n",
    "        device=\"cuda\",\n",
    "        covariance_paths=[\n",
    "            f\"token_sequence_transformer.blocks.{b}.attention.attention\"\n",
    "            for b in range(run.config.task_config.num_layers)\n",
    "        ]\n",
    "    ):\n",
    "        # wandb.log(observables, step=step)\n",
    "        observables[\"step\"] = step\n",
    "        covariances = observables.pop(\"covariances\")\n",
    "\n",
    "        # I only want the two largest eigenvalues in evals and evecs\n",
    "        covariances = np.linalg.eigvalsh(covariances)\n",
    "        evals, evecs = eigsh(covariances, k=3, which='LM')\n",
    "        \n",
    "        observables_over_time.append(observables)\n",
    "        print(yaml.dump({\n",
    "            **observables,\n",
    "            \"covariances\": covariances.shape\n",
    "        }))\n",
    "\n",
    "        raise NotImplementedError(\"TODO: Save covariances\")\n",
    "\n",
    "    subdf = pd.DataFrame(observables_over_time)\n",
    "    slug = \"slt-\" + run.config.to_slug(delimiter=\"-\")\n",
    "    subdf.to_csv(ANALYSIS / (slug + \".csv\"))\n",
    "\n",
    "# wandb.finish()\n",
    "os.system('say \"Your program has finished.\"')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import shutil\n",
    "from icl.config import RegressionConfig\n",
    "\n",
    "import os\n",
    "from PIL import Image\n",
    "\n",
    "def gather_images_side_by_side(folder, save: Optional[str] = None, delete: bool = True):\n",
    "    \"\"\"\n",
    "    Assumes folder contains folders that contain pngs. \n",
    "    \"\"\"\n",
    "    folder = Path(folder)\n",
    "    folder_paths = folder.glob(\"*\")\n",
    "\n",
    "    # Create a dictionary to store images by filename\n",
    "    images_by_filename = {}\n",
    "\n",
    "    if save:\n",
    "        save = Path(save)\n",
    "\n",
    "        if not os.path.exists(save):\n",
    "            os.makedirs(save)\n",
    "\n",
    "    # Load images from each folder and organize them by filename\n",
    "    for folder_path in folder_paths:\n",
    "        filenames = [f for f in os.listdir(folder_path) if f.endswith('.png')] \n",
    "        for filename in filenames:\n",
    "            img = Image.open(os.path.join(folder_path, filename))\n",
    "            if filename in images_by_filename:\n",
    "                images_by_filename[filename].append(img)\n",
    "            else:\n",
    "                images_by_filename[filename] = [img]\n",
    "\n",
    "    # Create comparison images for each unique filename\n",
    "    for filename, image_list in images_by_filename.items():\n",
    "        # Calculate the width and height of the result image\n",
    "        width = sum(img.width for img in image_list)\n",
    "        height = max(img.height for img in image_list)\n",
    "\n",
    "        # Create a new image for the comparison\n",
    "        result_image = Image.new('RGB', (width, height))\n",
    "\n",
    "        # Paste images side by side\n",
    "        x_offset = 0\n",
    "        for img in image_list:\n",
    "            result_image.paste(img, (x_offset, 0))\n",
    "            x_offset += img.width\n",
    "\n",
    "        # Display or save the result image\n",
    "        if save: \n",
    "            result_image.save(save / filename)  # You can replace this with result_image.save() to save the comparison images\n",
    "\n",
    "    if delete:\n",
    "        # Delete the temporary folder\n",
    "        shutil.rmtree(folder)\n",
    "\n",
    "\n",
    "def plot_activations(config: RegressionConfig, activations: Dict[str, torch.Tensor], save: Optional[str] = None):\n",
    "    B = 1\n",
    "    E = config.task_config.embed_size\n",
    "    T = 2 * config.task_config.max_examples\n",
    "    H = config.task_config.num_heads\n",
    "\n",
    "    def optionally_rotate(x, name):\n",
    "        if len(x.shape) != 2:\n",
    "            raise ValueError(\"Tensor should have two dimensions.\")\n",
    "\n",
    "        if x.shape[0] > x.shape[1]:\n",
    "            return x.T, f\"{name}.T\"\n",
    "        \n",
    "        return x, name \n",
    "\n",
    "    def separate_attention(qkv: TensorType[\"B\", \"T\", \"C\"], num_heads: int, batch_size: int, head_size: int, num_tokens: int):\n",
    "        return (qkv   \n",
    "            .view(batch_size, num_tokens, num_heads, 3*head_size)\n",
    "            .transpose(-2, -3)     \n",
    "            .split(head_size, dim=-1)\n",
    "        )\n",
    "\n",
    "    if save:\n",
    "        save = Path(save)\n",
    "\n",
    "        if not os.path.exists(save):\n",
    "            os.makedirs(save)\n",
    "\n",
    "    for location, v in activations.items():\n",
    "        activation_slice = v[0]\n",
    "\n",
    "        if location.endswith(\"attention.attention\"):\n",
    "            q, k, v = separate_attention(v, num_heads=H, batch_size=B, head_size=E//H, num_tokens=T)\n",
    "            qk = q @ k.transpose(-2, -1)\n",
    "            q, k, qk, v = q[0], k[0], v[0], qk[0]\n",
    "            \n",
    "            fig, axs = plt.subplots(H, 4, figsize=(15, 15))\n",
    "\n",
    "            for j, (name, x) in enumerate(zip([\"Q\", \"K\", \"QK\", \"V\"], [q, k, qk, v])):\n",
    "                for h in range(H):\n",
    "                    ax = axs[h, j]\n",
    "                    im = ax.matshow(x[h].detach().to(\"cpu\").numpy())\n",
    "                    ax.set_title(f\"{h}.{name}\")\n",
    "                    # fig.colorbar(im, ax=ax)\n",
    "\n",
    "            plt.suptitle(location)\n",
    "            plt.tight_layout()\n",
    "\n",
    "            if save:\n",
    "                plt.savefig(save / (location + \".png\"))\n",
    "                plt.close(fig)\n",
    "\n",
    "                del fig\n",
    "                del axs\n",
    "\n",
    "        elif len(activation_slice.shape) == 2:\n",
    "            fig = plt.figure()\n",
    "\n",
    "            x, location = optionally_rotate(activation_slice, location)\n",
    "            plt.matshow(x.detach().to(\"cpu\").numpy())\n",
    "            plt.title(f\"{location}\")\n",
    "            # fig.colorbar(im)\n",
    "\n",
    "            if save:\n",
    "                plt.savefig(save / (location + \".png\"))\n",
    "                plt.close(fig)\n",
    "\n",
    "                del fig\n",
    "\n",
    "\n",
    "        elif len(activation_slice.shape) == 3:  # [heads, xs, ys]\n",
    "            heads, xs, ys = activation_slice.shape\n",
    "            fig, axs = plt.subplots(1, heads, figsize=(15, 15))\n",
    "\n",
    "            for j in range(heads):\n",
    "                ax = axs[j]\n",
    "                x, name = optionally_rotate(activation_slice[j], str(j))\n",
    "                im = ax.matshow(x.detach().to(\"cpu\").numpy())\n",
    "                ax.set_title(f\"{name}\")\n",
    "                # fig.colorbar(im, ax=ax)\n",
    "            \n",
    "            plt.suptitle(f\"{location}.#\")\n",
    "            plt.tight_layout()\n",
    "\n",
    "            if save:\n",
    "                plt.savefig(save / (location + \".png\"))\n",
    "                plt.close(fig)\n",
    "\n",
    "            del fig\n",
    "            del axs\n",
    "\n",
    "        else:\n",
    "            raise ValueError(\"Unsupported number of dimensions.\")\n",
    "\n",
    "\n",
    "def compare_activations(config: ICLConfig, model, x: TensorType[\"B\", \"D\"], y: TensorType[\"B\", 1], save: Optional[str] = None, names: Optional[List[str]] = None):\n",
    "    B = len(x)\n",
    "    hooked_model = hook(model)\n",
    "\n",
    "    activations = {}\n",
    "    output, activations_ = hooked_model.run_with_cache(x, y)\n",
    "    activations[\"x\"] = x\n",
    "    activations[\"y\"] = y\n",
    "    activations[\"output\"] = output\n",
    "    activations.update(activations_)\n",
    "\n",
    "    def activations_per_sample(activations, index, keep_batch_dim=False):\n",
    "        if keep_batch_dim:\n",
    "            print({k: type(v) for k, v in activations.items()})\n",
    "            return {k: v[index].unsqueeze(0) for k, v in activations.items() if v is not None}\n",
    "        \n",
    "        return {k: v[index] for k, v in activations.items() if v is not None}\n",
    "\n",
    "    tmp_folder = Path(\"tmp\")\n",
    "\n",
    "    names = names or list(map(str, range(B)))\n",
    "\n",
    "    for (name, b) in zip(names, range(B)):\n",
    "        activations_b = activations_per_sample(activations, b, keep_batch_dim=True)\n",
    "        plot_activations(config, activations_b, save=tmp_folder / str(name))\n",
    "\n",
    "    gather_images_side_by_side(tmp_folder, save=save, delete=True)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demo = Run.create_and_restore(configs[2])\n",
    "compare_activations(demo.config, demo.model, demo.evaluator.pretrain_xs[:3], demo.evaluator.pretrain_ys[:3], save=FIGURES / \"demo\", names=[\"$x_0$\", \"$x_1$\", \"$x_2$\"])\n",
    "# gather_images_side_by_side(\"tmp\", save=FIGURES/\"demo\", delete=True)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot a few samples for each model at the end of training\n",
    "from icl.train import Run\n",
    "\n",
    "NUM_SAMPLES = 4\n",
    "\n",
    "for config, checkpointer in zip(configs, checkpointers):\n",
    "    run = Run.create_and_restore(config)\n",
    "    \n",
    "    sample_names = [f\"$x_{i}$\" for i in range(NUM_SAMPLES)]\n",
    "    slug = \"activations-\" + run.config.to_slug(delimiter=\"-\")\n",
    "\n",
    "    compare_activations(\n",
    "        run.config, \n",
    "        run.model, \n",
    "        run.evaluator.pretrain_xs[:NUM_SAMPLES], \n",
    "        run.evaluator.pretrain_ys[:NUM_SAMPLES], \n",
    "        save=FIGURES / slug, \n",
    "        names=sample_names\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(checkpointer.file_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot a few samples for a subset of models over training\n",
    "\n",
    "MS = [1, 4, 64, 2**10, 2**20]\n",
    "STEPS = [0, 1_805, 3_084, 15_381, 26_279, 100_262, 153_061, 193_877, 255_102, 306_122, 408_163]\n",
    "\n",
    "for log2_M in MS:\n",
    "    log2_m = int(np.log2(log2_M))\n",
    "    config, checkpointer = configs[log2_m], checkpointers[log2_m]\n",
    "    run = Run(config)\n",
    "\n",
    "    for step in STEPS:\n",
    "        run.model.load_state_dict(checkpointer.load_file(step)[\"model\"])\n",
    "\n",
    "        sample_names = [f\"$x_{i}$\" for i in range(NUM_SAMPLES)]\n",
    "        slug = \"activations-\" + run.config.to_slug(delimiter=\"-\") + f\"@t={step}\"\n",
    "\n",
    "        # TODO: Need to rename the new files otherwise you can't tell easily tell what step they come from.\n",
    "        compare_activations(\n",
    "            run.config, \n",
    "            run.model, \n",
    "            run.evaluator.pretrain_xs[:NUM_SAMPLES], \n",
    "            run.evaluator.pretrain_ys[:NUM_SAMPLES], \n",
    "            save=FIGURES / slug, \n",
    "            names=sample_names\n",
    "        )\n",
    "        \n",
    "        os.system('say \"Your program has finished.\"')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLC hyperparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "import wandb\n",
    "from icl.config import get_config\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from infra.utils.iterables import flatten_dict\n",
    "\n",
    "api = wandb.Api()\n",
    "sweep = api.sweep(f\"{WANDB_ENTITY}/icl-llc/ebu13rjw\")  # L4H4\n",
    "\n",
    "def wandb_run_to_df(run):\n",
    "    history_df = run.history()\n",
    "    config_dict = get_config(**run.config).model_dump()\n",
    "    config_dict[\"analysis_config\"] = run.config[\"analysis_config\"]\n",
    "\n",
    "    del config_dict[\"logger_config\"]\n",
    "    del config_dict[\"checkpointer_config\"]\n",
    "\n",
    "    config_dict_flat = flatten_dict(config_dict, flatten_lists=True)\n",
    "    \n",
    "    for k, v in config_dict_flat.items():\n",
    "        if isinstance(v, tuple):\n",
    "            # Repeat the tuple for the entire length of the DataFrame\n",
    "            v = [v] * len(history_df)\n",
    "            \n",
    "        history_df[k] = v\n",
    "\n",
    "    return history_df\n",
    "\n",
    "\n",
    "def wandb_runs_to_df(runs):\n",
    "    return pd.concat([wandb_run_to_df(run) for run in tqdm(runs, desc=\"Converting runs to dfs\")])\n",
    "\n",
    "\n",
    "subdf = wandb_runs_to_df(sweep.runs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers = 2\n",
    "num_heads = 4\n",
    "# df.to_csv(\"../analysis/L4H4-llc-grid-search.csv\") \n",
    "subdf = pd.read_csv(f\"../analysis/L{num_layers}H{num_heads}-llc-grid-search.csv\")\n",
    "# df = pd.read_csv(\"../analysis/L2H4-llc-grid-search.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# del df\n",
    "subdf.columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "\n",
    "# Get unique values for lrs, elasticitys, and num_tasks\n",
    "num_chains = 25\n",
    "unique_lrs = subdf['analysis_config/lr'].unique()\n",
    "unique_elasticities = subdf['analysis_config/elasticity'].unique()\n",
    "unique_num_tasks = subdf['task_config/num_tasks'].unique()\n",
    "\n",
    "show_std = False\n",
    "\n",
    "# Sort for visual consistency\n",
    "unique_lrs.sort()\n",
    "unique_lrs = unique_lrs[:-1]\n",
    "unique_elasticities.sort()\n",
    "unique_num_tasks.sort()\n",
    "\n",
    "prefix = \"\" # \"thresholded-\" # \"\"\n",
    "Prefix = \"\" # \"Thresholded \" # \"\"\n",
    "\n",
    "# Initialize colormap\n",
    "cmap = plt.cm.viridis\n",
    "\n",
    "# Create subplots\n",
    "fig, axes = plt.subplots(len(unique_lrs), len(unique_elasticities), figsize=(15, 15))\n",
    "fig.set_facecolor('white')\n",
    "fig.suptitle(f\"{Prefix}$\\hat\\lambda$ hyperparameter sweep ($n_\\mathrm{{chains}}={num_chains}$)\\n$L={num_layers}, H={num_heads}, t=500k$\")\n",
    "\n",
    "# Loop through the grid\n",
    "for token, lr in enumerate(unique_lrs):\n",
    "    for j, elasticity in enumerate(unique_elasticities):\n",
    "        ax = axes[token, j]\n",
    "\n",
    "        # Filter DataFrame for specific lr and elasticity\n",
    "        filtered_df = subdf[(subdf['analysis_config/lr'] == lr) & (subdf['analysis_config/elasticity'] == elasticity)]\n",
    "\n",
    "        for log_num_tasks in unique_num_tasks:\n",
    "            task_specific_df = filtered_df[filtered_df['task_config/num_tasks'] == log_num_tasks]\n",
    "\n",
    "            # Sort by 'num_draws' for plotting\n",
    "            task_specific_df = task_specific_df.sort_values('_step')\n",
    "\n",
    "            # Calculate color based on log2(num_tasks)\n",
    "            color = cmap(np.log2(log_num_tasks) / np.log2(max(unique_num_tasks)))\n",
    "\n",
    "            # Plot using Seaborn for better aesthetics\n",
    "            filtered_data = task_specific_df[(task_specific_df[f'{prefix}llc/mean'] != \"NaN\") & (task_specific_df[f'{prefix}llc/mean'] <1_000)]\n",
    "            sns.lineplot(x='_step', y=f'{prefix}llc/mean', data=filtered_data, ax=ax, label=f'_M={log_num_tasks}', color=color)\n",
    "            \n",
    "            if show_std:\n",
    "                mean_val = task_specific_data[f'{prefix}llc/mean']\n",
    "                std_val = task_specific_data[f'{prefix}llc/std']\n",
    "\n",
    "                if not np.isnan(mean_val) and not np.isnan(std_val) and np.isfinite(mean_val) and np.isfinite(std_val):\n",
    "                    lower = mean_val - std_val\n",
    "                    upper = mean_val + std_val\n",
    "                else:\n",
    "                    lower = np.nan\n",
    "                    upper = np.nan\n",
    "\n",
    "                ax.fill_between(task_specific_data['_step'], lower, upper, color=color, alpha=0.1)\n",
    "\n",
    "        ax.set_title(f\"$\\epsilon={lr}, \\gamma={elasticity}$\")\n",
    "        ax.set_xlabel(r\"$t_\\mathrm{SGLD}$\")\n",
    "        ax.set_ylabel(r\"$\\hat\\lambda$\")\n",
    "\n",
    "# plt.legend()\n",
    "plt.tight_layout()\n",
    "\n",
    "# Plot a color bar to the right of the grid\n",
    "norm = Normalize(vmin=0, vmax=20)\n",
    "cbar = fig.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=axes)\n",
    "# cbar.ax.set_clim(0, 20)\n",
    "cbar.ax.set_ylabel(r\"$\\log_2(M)$\")\n",
    "cbar.locator = MaxNLocator(integer=True)\n",
    "cbar.update_ticks()\n",
    "\n",
    "\n",
    "if show_std:\n",
    "    plt.savefig(f\"../figures/llc-grid-over-t-L{num_layers}_H{num_heads}.png\")\n",
    "else:\n",
    "    plt.savefig(f\"../figures/llc-grid-over-t-L{num_layers}_H{num_heads}.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "\n",
    "for M in range(0, 21, 3):\n",
    "    upper_M = min(21, M + 5)\n",
    "\n",
    "    data = subdf[(subdf['task_config/num_tasks'] >= 2**M) & (subdf['task_config/num_tasks'] < 2**upper_M)]\n",
    "\n",
    "    # Get unique values for lrs, elasticitys, and num_tasks\n",
    "    num_chains = 25\n",
    "    unique_lrs = data['analysis_config/lr'].unique()\n",
    "    unique_elasticities = data['analysis_config/elasticity'].unique()\n",
    "    unique_num_tasks = data['task_config/num_tasks'].unique()\n",
    "\n",
    "    show_std = True\n",
    "\n",
    "    # Sort for visual consistency\n",
    "    unique_lrs.sort()\n",
    "    unique_lrs = unique_lrs[:-1]\n",
    "    unique_elasticities.sort()\n",
    "    unique_num_tasks.sort()\n",
    "\n",
    "    # Initialize colormap\n",
    "    cmap = plt.cm.viridis\n",
    "\n",
    "    # Create subplots\n",
    "    fig, axes = plt.subplots(len(unique_lrs), len(unique_elasticities), figsize=(15, 15))\n",
    "    fig.set_facecolor('white')\n",
    "\n",
    "    prefix = \"thresholded-\" # \"\"\n",
    "    Prefix = \"Thresholded \" # \"\"\n",
    "\n",
    "    fig.suptitle(f\"{Prefix}$\\hat\\lambda$ hyperparameter sweep ($n_\\mathrm{{chains}}={num_chains}$)\\n$L=2, H=4$\\n$M \\in [{2**M}, {2**upper_M}, t=500k)$\")\n",
    "\n",
    "    # Loop through the grid\n",
    "    for token, lr in enumerate(unique_lrs):\n",
    "        for j, elasticity in enumerate(unique_elasticities):\n",
    "            ax = axes[token, j]\n",
    "\n",
    "            # Filter DataFrame for specific lr and elasticity\n",
    "            filtered_data = data[(data['analysis_config/lr'] == lr) & (data['analysis_config/elasticity'] == elasticity)]\n",
    "\n",
    "            for log_num_tasks in unique_num_tasks:\n",
    "                task_specific_data = filtered_data[filtered_data['task_config/num_tasks'] == log_num_tasks]\n",
    "\n",
    "                # Sort by 'num_draws' for plotting\n",
    "                task_specific_data = task_specific_data.sort_values('_step')\n",
    "\n",
    "                # Calculate color based on log2(num_tasks)\n",
    "                color = cmap((np.log2(log_num_tasks)-M)/5)\n",
    "\n",
    "                # Plot using Seaborn for better aesthetics\n",
    "                more_filtered_data = task_specific_data.loc[(task_specific_data[f'{prefix}llc/mean'] != \"NaN\") & (task_specific_data[f'{prefix}llc/std'] != \"NaN\")]\n",
    "                sns.lineplot(x='_step', y=f'{prefix}llc/mean', data=more_filtered_data, ax=ax, label=f'_M={log_num_tasks}', color=color)\n",
    "\n",
    "                if show_std:\n",
    "                    # Print types of each cell in more_filtered_data\n",
    "                    steps = more_filtered_data['_step'].to_numpy()\n",
    "                    means = more_filtered_data[f\"{prefix}llc/mean\"].to_numpy()\n",
    "                    stds = more_filtered_data[f\"{prefix}llc/std\"].to_numpy()\n",
    "                    means = pd.to_numeric(means, errors='coerce')\n",
    "                    stds = pd.to_numeric(stds, errors='coerce')\n",
    "\n",
    "                    ax.fill_between(steps, means-stds, means+stds, color=color, alpha=0.2)\n",
    "                    \n",
    "            ax.set_title(f\"$\\epsilon={lr}, \\gamma={elasticity}$\")\n",
    "            ax.set_xlabel(r\"$t_\\mathrm{SGLD}$\")\n",
    "            ax.set_ylabel(r\"$\\hat\\lambda$\")\n",
    "\n",
    "    # plt.legend()\n",
    "    plt.tight_layout()\n",
    "\n",
    "    # Plot a color bar to the right of the grid\n",
    "    norm = Normalize(vmin=M, vmax=upper_M)\n",
    "    cbar = fig.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=axes)\n",
    "    # cbar.ax.set_clim(0, 20)\n",
    "    cbar.ax.set_ylabel(r\"$\\log_2(M)$\")\n",
    "    cbar.locator = MaxNLocator(integer=True)\n",
    "    cbar.update_ticks()\n",
    "\n",
    "\n",
    "    plt.savefig(f\"../figures/llc-grid-search-M{M}-{upper_M}.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's plot this as a function of M on the x axis\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "\n",
    "# Get unique values for lrs, elasticitys, and num_tasks\n",
    "num_chains = 25\n",
    "unique_lrs = subdf['analysis_config/lr'].unique()\n",
    "unique_elasticities = subdf['analysis_config/elasticity'].unique()\n",
    "unique_num_tasks = subdf['task_config/num_tasks'].unique()\n",
    "\n",
    "num_layers = 4\n",
    "num_heads = 4\n",
    "\n",
    "show_std = True\n",
    "\n",
    "# Sort for visual consistency\n",
    "unique_lrs.sort()\n",
    "# unique_lrs = np.array([lr for lr in unique_lrs if lr <= 0.0001])\n",
    "unique_lrs = unique_lrs[:-2]\n",
    "unique_elasticities.sort()\n",
    "unique_num_tasks.sort()\n",
    "\n",
    "unique_num_tasks = np.array([2**m for m in range(0, 21)])\n",
    "\n",
    "prefix = \"\" # \"thresholded-\" \n",
    "Prefix = \"\" # \"Thresholded \" \n",
    "\n",
    "# Initialize colormap\n",
    "cmap = plt.cm.viridis\n",
    "\n",
    "# Create subplots\n",
    "fig, axes = plt.subplots(len(unique_lrs), len(unique_elasticities), figsize=(15, 15))\n",
    "fig.set_facecolor('white')\n",
    "fig.suptitle(f\"{Prefix}$\\hat\\lambda$ hyperparameter sweep ($n_\\mathrm{{chains}}={num_chains}$)\\n$L={num_layers}, H={num_heads}, t=500k$\")\n",
    "\n",
    "steps = np.array([9, 29, 99, 299, 999])\n",
    "\n",
    "# Loop through the grid\n",
    "for token, lr in enumerate(unique_lrs):\n",
    "    for j, elasticity in enumerate(unique_elasticities):\n",
    "        ax = axes[token, j]\n",
    "\n",
    "        # Filter DataFrame for specific lr and elasticity\n",
    "        filtered_df = subdf[(subdf['analysis_config/lr'] == lr) & (subdf['analysis_config/elasticity'] == elasticity)]\n",
    "\n",
    "        for step in steps:\n",
    "            # Find the closest step to the desired step for each num_tasks \n",
    "            # Problem is wandb sometimes drops a log.\n",
    "            closest_step_df = filtered_df.groupby('task_config/num_tasks').apply(lambda x: x.iloc[(x['_step']-step).abs().argsort()[:1]]).reset_index(drop=True)\n",
    "\n",
    "            # Sort by 'num_draws' for plotting\n",
    "            closest_step_df = closest_step_df.sort_values('task_config/num_tasks')\n",
    "\n",
    "            # Calculate color based on log2(num_tasks)\n",
    "            color = cmap(step / 999)\n",
    "\n",
    "            # Plot using Seaborn for better aesthetics\n",
    "            filtered_data = closest_step_df[closest_step_df[f'{prefix}llc/mean'] != \"NaN\"]\n",
    "            log_num_tasks = filtered_data['task_config/num_tasks'].to_numpy()\n",
    "            log_num_tasks = pd.to_numeric(log_num_tasks, errors='coerce')\n",
    "            means = filtered_data[f\"{prefix}llc/mean\"].to_numpy()\n",
    "            means = pd.to_numeric(means, errors='coerce')\n",
    "\n",
    "            ax.plot(log_num_tasks, means, label=f'_step={step}', color=color)\n",
    "\n",
    "            if show_std:\n",
    "                stds = filtered_data[f\"{prefix}llc/std\"].to_numpy()\n",
    "                stds = pd.to_numeric(stds, errors='coerce')\n",
    "                ax.fill_between(log_num_tasks, means-stds, means+stds, color=color, alpha=0.2)\n",
    "                \n",
    "        ax.set_title(f\"$\\epsilon={lr}, \\gamma={elasticity}$\")\n",
    "        ax.set_xlabel(r\"$M$\")\n",
    "        ax.set_xscale(\"log\")\n",
    "        ax.set_xticks([2**m for m in range(0, 21, 4)], [f\"$2^{{{m}}}$\" for m in range(0, 21, 4)])\n",
    "        ax.set_ylabel(r\"$\\hat\\lambda$\")\n",
    "\n",
    "# plt.legend()\n",
    "plt.tight_layout()\n",
    "\n",
    "# Plot a color bar to the right of the grid\n",
    "norm = Normalize(vmin=0, vmax=1000)\n",
    "cbar = fig.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=axes)\n",
    "# cbar.ax.set_clim(0, 20)\n",
    "cbar.ax.set_ylabel(r\"$t_\\mathrm{SGLD}$\")\n",
    "cbar.locator = MaxNLocator(integer=True)\n",
    "cbar.update_ticks()\n",
    "\n",
    "\n",
    "if show_std:\n",
    "    plt.savefig(\"../figures/llc-grid-search-std.png\")\n",
    "else:\n",
    "    plt.savefig(\"../figures/llc-grid-search.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's plot this as a function of M on the x axis\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "\n",
    "# Get unique values for lrs, elasticitys, and num_tasks\n",
    "num_chains = 25\n",
    "unique_lrs = subdf['analysis_config/lr'].unique()\n",
    "unique_elasticities = subdf['analysis_config/elasticity'].unique()\n",
    "unique_num_tasks = subdf['task_config/num_tasks'].unique()\n",
    "\n",
    "num_layers = 2\n",
    "num_heads = 4\n",
    "\n",
    "show_std = True\n",
    "\n",
    "# Sort for visual consistency\n",
    "unique_lrs.sort()\n",
    "unique_lrs = unique_lrs[:-1]\n",
    "# unique_lrs = np.array([lr for lr in unique_lrs if lr <= 0.0001])\n",
    "unique_elasticities.sort()\n",
    "unique_num_tasks.sort()\n",
    "\n",
    "unique_num_tasks = np.array([2**m for m in range(0, 21)])\n",
    "\n",
    "prefix = \"\" # \"thresholded-\" \n",
    "Prefix = \"\" # \"Thresholded \" \n",
    "\n",
    "# Initialize colormap\n",
    "cmap = plt.cm.viridis\n",
    "\n",
    "# Create subplots\n",
    "fig, axes = plt.subplots(len(unique_lrs), len(unique_elasticities), figsize=(15, 15))\n",
    "fig.set_facecolor('white')\n",
    "fig.suptitle(f\"{Prefix}$\\hat\\lambda$ hyperparameter sweep ($n_\\mathrm{{chains}}={num_chains}$)\\n$L={num_layers}, H={num_heads}, t=500k$\")\n",
    "\n",
    "# Loop through the grid\n",
    "for token, lr in enumerate(unique_lrs):\n",
    "    for j, elasticity in enumerate(unique_elasticities):\n",
    "        ax = axes[token, j]\n",
    "\n",
    "        # Filter DataFrame for specific lr and elasticity\n",
    "        filtered_df = subdf[(subdf['analysis_config/lr'] == lr) & (subdf['analysis_config/elasticity'] == elasticity)]\n",
    "\n",
    "        # Get the last step for each num_tasks\n",
    "        last_step_df = filtered_df.groupby('task_config/num_tasks').last().reset_index()\n",
    "\n",
    "        # Calculate color based on log2(num_tasks)\n",
    "        color = sns.color_palette()[0]\n",
    "\n",
    "        # Plot using Seaborn for better aesthetics\n",
    "        filtered_data = last_step_df[last_step_df[f'{prefix}llc/mean'] != \"NaN\"]\n",
    "        log_num_tasks = filtered_data['task_config/num_tasks'].to_numpy()\n",
    "        log_num_tasks = pd.to_numeric(log_num_tasks, errors='coerce')\n",
    "        means = filtered_data[f\"{prefix}llc/mean\"].to_numpy()\n",
    "        means = pd.to_numeric(means, errors='coerce')\n",
    "\n",
    "        ax.plot(log_num_tasks, means, label=f'_step={step}', color=color)\n",
    "\n",
    "        if show_std:\n",
    "            stds = filtered_data[f\"{prefix}llc/std\"].to_numpy()\n",
    "            stds = pd.to_numeric(stds, errors='coerce')\n",
    "            ax.fill_between(log_num_tasks, means-stds, means+stds, color=color, alpha=0.2)\n",
    "                \n",
    "        ax.set_title(f\"$\\epsilon={lr}, \\gamma={elasticity}, t_\\mathrm{{SGLD}}=1000$\")\n",
    "        ax.set_xlabel(r\"$M$\")\n",
    "        ax.set_xscale(\"log\")\n",
    "        ax.set_xticks([2**m for m in range(0, 21, 4)], [f\"$2^{{{m}}}$\" for m in range(0, 21, 4)])\n",
    "        ax.set_ylabel(r\"$\\hat\\lambda$\")\n",
    "\n",
    "# plt.legend()\n",
    "plt.tight_layout()\n",
    "\n",
    "# Plot a color bar to the right of the grid\n",
    "# norm = Normalize(vmin=0, vmax=1000)\n",
    "# cbar = fig.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=axes)\n",
    "# cbar.ax.set_clim(0, 20)\n",
    "# cbar.ax.set_ylabel(r\"$t_\\mathrm{SGLD}$\")\n",
    "# cbar.locator = MaxNLocator(integer=True)\n",
    "# cbar.update_ticks()\n",
    "\n",
    "\n",
    "if show_std:\n",
    "    plt.savefig(\"../figures/llc-grid-search-std.png\")\n",
    "else:\n",
    "    plt.savefig(\"../figures/llc-grid-search.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "import wandb\n",
    "from icl.config import get_config\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from infra.utils.iterables import flatten_dict\n",
    "\n",
    "api = wandb.Api()\n",
    "sweep = api.sweep(f\"{WANDB_ENTITY}/icl-llc/eli1wlds\")  # L4H4\n",
    "\n",
    "def wandb_run_to_df(run):\n",
    "    history_df = run.history()\n",
    "    config_dict = get_config(**run.config).model_dump()\n",
    "    config_dict[\"analysis_config\"] = run.config[\"analysis_config\"]\n",
    "\n",
    "    del config_dict[\"logger_config\"]\n",
    "    del config_dict[\"checkpointer_config\"]\n",
    "\n",
    "    config_dict_flat = flatten_dict(config_dict, flatten_lists=True)\n",
    "    \n",
    "    for k, v in config_dict_flat.items():\n",
    "        if isinstance(v, tuple):\n",
    "            # Repeat the tuple for the entire length of the DataFrame\n",
    "            v = [v] * len(history_df)\n",
    "            \n",
    "        history_df[k] = v\n",
    "\n",
    "    return history_df\n",
    "\n",
    "\n",
    "def wandb_runs_to_df(runs):\n",
    "    return pd.concat([wandb_run_to_df(run) for run in tqdm(runs, desc=\"Converting runs to dfs\")])\n",
    "\n",
    "\n",
    "subdf = wandb_runs_to_df(sweep.runs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_layers = 2\n",
    "num_heads = 4\n",
    "\n",
    "name = f\"../analysis/L{num_layers}H{num_heads}-llc-grid-search-batches.csv\"\n",
    "subdf.to_csv(name) \n",
    "# subdf = pd.read_csv(name)\n",
    "# df = pd.read_csv(\"../analysis/L2H4-llc-grid-search.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "\n",
    "def plot_llc_estimation_hyperparams_sweep(observations_df, y=\"llc/mean\", row=\"analysis_config/lr\", col=\"analysis_config/batch_size\"):\n",
    "    fig = plt.figure(figsize=(15, 6))\n",
    "\n",
    "    # Get rid of the NaNs\n",
    "    observations_df = observations_df[observations_df[y] != \"NaN\"]\n",
    "    observations_df = observations_df.rename(columns={\"task_config/num_tasks\": \"# Tasks\",\n",
    "                                                      \"analysis_config/lr\": \"Learning Rate\",\n",
    "                                                      \"analysis_config/batch_size\": \"Batch Size\"})\n",
    "\n",
    "    g = sns.FacetGrid(observations_df, col=\"Batch Size\", row=\"Learning Rate\", sharey=False)\n",
    "    g.map_dataframe(sns.lineplot, x=\"# Tasks\", y=y)\n",
    "    g.add_legend()\n",
    "    g.set(xscale=\"log\", yscale=\"log\")\n",
    "    # g.set(xscale=\"log\", yscale=\"linear\")\n",
    "\n",
    "    plt.suptitle(\"Covariance estimation hyperparameter sweep\")\n",
    "    g.fig.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "plot_llc_estimation_hyperparams_sweep(subdf)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.9.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}