{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analysis of trained SAEs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load packages\n",
    "\n",
    "import math, time\n",
    "import torch\n",
    "import os\n",
    "import sys\n",
    "sys.path.append('../..')\n",
    "from functions.get_data import data_n_loaders\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SET THE PATH OF DATA and name of dataset\n",
    "DATA_PATH = './data/' #change this to your data path\n",
    "dataset = f'128dgaussian_intrinsicdim'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#choose models to visualize\n",
    "\n",
    "#get exptnames from folder names within saved_models\n",
    "exptnames = []\n",
    "for root, dirs, files in os.walk('saved_models'):\n",
    "    for dir in dirs:\n",
    "        if 'checkpoints_' in dir:\n",
    "            exptnames.append(dir.split('checkpoints_')[1])\n",
    "\n",
    "#sort exptnames to have all with same sae_type (name after the first _ within exptnames) together; after saename, the float following must be used for sorting\n",
    "order = {'relu': 0, 'jumprelu': 1, 'topk': 2, 'spade': 3}\n",
    "exptnames = sorted(exptnames, key=lambda x: (order[x.split('_')[1]]))\n",
    "\n",
    "#get parameters from experiment names\n",
    "saetypes = []\n",
    "for exptname in exptnames:\n",
    "    if 'spade' in exptname:\n",
    "        saetypes.append('sparsemax_dist')\n",
    "    elif 'topk_relu' in exptname:\n",
    "        saetypes.append('topk_relu')\n",
    "    else:\n",
    "        saetypes.append(exptname.split('_')[1])\n",
    "max_epochs = []\n",
    "for exptname in exptnames:\n",
    "    epochs = []\n",
    "    for root, dirs, files in os.walk('saved_models/checkpoints_'+exptname):\n",
    "        for file in files:\n",
    "            if 'model_' in file:\n",
    "                epochs.append(int(file.split('_')[1].split('.pt')[0].split('epochs')[0]))\n",
    "    max_epochs.append(max(epochs))\n",
    "kvals_all = torch.zeros((len(exptnames),), dtype=int)\n",
    "for i in range(len(exptnames)):\n",
    "    if saetypes[i] == 'topk_relu':\n",
    "        kvals_all[i] = int(exptnames[i].split('_k')[1].split('_')[0])\n",
    "    else:\n",
    "        kvals_all[i] = 0\n",
    "kvals_all = kvals_all.tolist()\n",
    "gammavals_all = []\n",
    "for i in range(len(exptnames)):\n",
    "    if saetypes[i] == 'topk_relu':\n",
    "        gammavals_all.append(0.0)\n",
    "    else:\n",
    "        gamma = float(exptnames[i].split('_gamreg')[1].split('_')[0])\n",
    "        gammavals_all.append(gamma)\n",
    "\n",
    "# choose which models (hyperparams) to plot\n",
    "gammavals_relu = [0.001, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0]\n",
    "gammavals_jumprelu = [0.001, 0.01, 0.025, 0.05, 0.075, 0.1, 1.0]\n",
    "gammavals_spade = [0.001, 0.01, 0.1, 1.0, 3.0, 5.0, 10.0]\n",
    "kvals_topk = [4, 8, 16, 32, 64, 128, 256]\n",
    "indices_relux = [i for i in range(len(exptnames)) if saetypes[i] == 'relu' and gammavals_all[i] in gammavals_relu]\n",
    "indices_jumprelux = [i for i in range(len(exptnames)) if saetypes[i] == 'jumprelu' and gammavals_all[i] in gammavals_jumprelu]\n",
    "indices_topkx = [i for i in range(len(exptnames)) if saetypes[i] == 'topk_relu' and kvals_all[i] in kvals_topk]\n",
    "indices_spadex = [i for i in range(len(exptnames)) if saetypes[i] == 'sparsemax_dist' and gammavals_all[i] in gammavals_spade]\n",
    "indices_all = indices_relux + indices_jumprelux + indices_topkx + indices_spadex\n",
    "exptnames = [exptnames[i] for i in indices_all]\n",
    "max_epochs = [max_epochs[indices_all[i]] for i in range(len(indices_all))]\n",
    "saetypes = [saetypes[i] for i in indices_all]\n",
    "kvals_all = [kvals_all[i] for i in indices_all]\n",
    "gammavals_all = [gammavals_all[i] for i in indices_all]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load trained models for chosen hyperparams\n",
    "from models import SAE\n",
    "dim = 128 #of points in dataset\n",
    "import numpy as np\n",
    "device = 'cuda'\n",
    "modelsall = []\n",
    "width = 512 #standard in this expt set\n",
    "for i in range(len(exptnames)):\n",
    "    exptnamei = exptnames[i]\n",
    "    normalize_decoder = False if saetypes[i] == 'sparsemax_dist' else True \n",
    "    modeli = SAE(dim, width, sae_type=saetypes[i],kval_topk=kvals_all[i], normalize_decoder=normalize_decoder)\n",
    "    state_dict = torch.load(f\"./saved_models/checkpoints_{exptnamei}/model_{max_epochs[i]}epochs.pt\")\n",
    "    modeli.load_state_dict(state_dict['model'])\n",
    "    modelsall.append(modeli.to(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load data subset\n",
    "\n",
    "# LAB_DIR = os.environ['USERDIR']\n",
    "# DATA_PATH = LAB_DIR+'/data'\n",
    "\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "#load entire train data and transform parameters\n",
    "torch.manual_seed(0)\n",
    "numpoints_perconcept = 1000\n",
    "num_concepts = 5\n",
    "batchsize = 2048\n",
    "meannorm2_concepts_all = []\n",
    "datavars_perconcept = torch.zeros((num_concepts,))\n",
    "_,_,train_data, _ = data_n_loaders(dataset, batchsize, return_data=True, \\\n",
    "                                    data_path=DATA_PATH, standardise_data = True)\n",
    "\n",
    "#sample equal number of points from each class\n",
    "intrinsic_dims = train_data.truefeatures['intrinsic_dims']\n",
    "data_subset_perconcept = []\n",
    "labels_subset_perconcept = []\n",
    "for j in range(num_concepts):\n",
    "    indices = train_data.labels == j\n",
    "    rand_indices_i = torch.randperm(indices.sum())[:numpoints_perconcept]\n",
    "    data_subset_perconcept.append(train_data.data[indices][rand_indices_i])\n",
    "    labels_subset_perconcept.append(train_data.labels[indices][rand_indices_i])\n",
    "data_subset = torch.cat(data_subset_perconcept, dim=0)\n",
    "labels_subset = torch.cat(labels_subset_perconcept, dim=0)\n",
    "meannorm2_concepts = torch.stack([torch.mean(torch.norm(data_subset[labels_subset==j,:], dim=1)**2) for j in range(num_concepts)])\n",
    "meannorm2_concepts_all.append(meannorm2_concepts)\n",
    "for k in range(num_concepts):\n",
    "    datavars_perconcept[k] = torch.var(data_subset[labels_subset==k,:], dim=0).sum()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#get latent representations (sparse codes) on data subset\n",
    "latent_reps_all = []\n",
    "Cdatacos = []\n",
    "stableranks = []\n",
    "l0_perconcept_all = []\n",
    "l0_all = []\n",
    "mse_perconcept_all = torch.zeros((len(exptnames),num_concepts))\n",
    "mse_all = torch.zeros((len(exptnames),))\n",
    "nmse_all = torch.zeros((len(exptnames),))\n",
    "mse_center = torch.zeros((len(exptnames), num_concepts))\n",
    "for i in range(len(exptnames)):\n",
    "    modeli = modelsall[i]\n",
    "    with torch.no_grad():\n",
    "        \n",
    "        preds, latent_reps = modeli(data_subset.to(device), return_hidden=True)\n",
    "        mse_persample = torch.sum((preds - data_subset.to(device))**2, dim=-1)\n",
    "        mse_perconcept_all[i,:] = torch.stack([torch.mean(mse_persample[labels_subset==j]) for j in range(num_concepts)])\n",
    "        mse_all[i] = torch.mean(mse_persample)\n",
    "        \n",
    "        #mse wrt mean of each concept\n",
    "        for k in range(num_concepts):\n",
    "            center_k = data_subset[labels_subset==k,:].mean(dim=0, keepdim=True).to(device)\n",
    "            mse_center[i,k] = (preds[labels_subset==k,:]-center_k).pow(2).sum(dim=-1).mean()\n",
    "            \n",
    "        latentsflat = latent_reps.flatten()\n",
    "        active = latentsflat>1e-12\n",
    "        if active.sum()>0:\n",
    "            # threshold = torch.quantile(latentsflat[active], 0.01)\n",
    "            threshold = 1e-6\n",
    "            activelatents = (latent_reps>threshold).float().sum(dim=1)\n",
    "            # activelatentdist_all.append(activelatents)\n",
    "            l0_perconcept = torch.stack([torch.mean(activelatents[labels_subset==j]) for j in range(num_concepts)])\n",
    "            l0_perconcept_all.append(l0_perconcept)\n",
    "            l0_all.append(torch.mean(activelatents))\n",
    "        else:\n",
    "            # activelatentdist_all.append(torch.zeros((latent_reps_all[i].shape[0],)))\n",
    "            l0_perconcept_all.append(torch.zeros((num_concepts,)))\n",
    "            l0_all.append(torch.tensor([0.0]))\n",
    "\n",
    "        #sparse code correlations for data pairs\n",
    "        latent_reps_all.append(latent_reps)\n",
    "        Cdata = latent_reps @ latent_reps.T\n",
    "        \n",
    "        #sparse code correlations for latent pairs\n",
    "        norms = torch.norm(latent_reps, dim=1).unsqueeze(1)\n",
    "        Cdatacos.append(Cdata/(norms @ norms.T + 1e-12))\n",
    "\n",
    "l0_perconcept_all = [l0.cpu().numpy() for l0 in l0_perconcept_all]\n",
    "l0_perconcept_all = torch.tensor(l0_perconcept_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "#sort latents acc to avg concept act, get latent-latent corrs\n",
    "\n",
    "deadlatents_all = []\n",
    "for i in range(len(exptnames)):\n",
    "    if (latent_reps_all[i]>1e-12).sum()==0:\n",
    "        deadlatents_all.append(torch.ones((latent_reps_all[i].shape[-1],), dtype=bool))\n",
    "    else:\n",
    "        threshold = torch.quantile(latent_reps_all[i][latent_reps_all[i]>0.0].flatten(), 0.1)\n",
    "        # threshold = 1e-12\n",
    "        deadlatents = torch.all(latent_reps_all[i] <= threshold, dim=0)\n",
    "        deadlatents_all.append(deadlatents.cpu())\n",
    "\n",
    "#identify concept which latent activates most for\n",
    "classavg_reps = [torch.zeros(num_concepts, latent_reps_all[0].shape[-1]) for i in range(len(exptnames))]\n",
    "for i in range(len(exptnames)):\n",
    "    for c in range(num_concepts):\n",
    "        classavg_reps[i][c] = torch.mean(latent_reps_all[i][labels_subset==c], dim=0)    \n",
    "order_by_concept_all = []\n",
    "concepts_of_latents_all = []\n",
    "for i in range(len(exptnames)):\n",
    "    concepts_of_latents = torch.argmax(classavg_reps[i][:, ~deadlatents_all[i]], dim=0) #gives indices of concepts with max value in each latent dimension\n",
    "    order_by_concept = torch.argsort(concepts_of_latents) #sorts latents by concept\n",
    "    order_by_concept_all.append(order_by_concept)\n",
    "    concepts_of_latents_all.append(concepts_of_latents)\n",
    "\n",
    "\n",
    "#sort latent representations by concept\n",
    "latentreps_conceptsorted_all = []\n",
    "for i in range(len(exptnames)):\n",
    "    activelatents = ~deadlatents_all[i]\n",
    "    latentsactive = latent_reps_all[i][:, activelatents]\n",
    "    latentsactive_conceptsorted = latentsactive[:, order_by_concept_all[i]]\n",
    "    latentsdead = latent_reps_all[i][:, ~activelatents]\n",
    "    latents_sortedanddead = torch.cat([latentsactive_conceptsorted, latentsdead], dim=1)\n",
    "    latentreps_conceptsorted_all.append(latents_sortedanddead) #reorder latents by concepts\n",
    "\n",
    "\n",
    "#track co-occurrence of latents\n",
    "Clatent = []\n",
    "for i in range(len(exptnames)):\n",
    "    Clatent.append(latentreps_conceptsorted_all[i].T @ latentreps_conceptsorted_all[i])\n",
    "\n",
    "Clatent_cos = []\n",
    "Clatent_cos_active = []\n",
    "for i in range(len(exptnames)):\n",
    "    norms = torch.norm(latentreps_conceptsorted_all[i], dim=0).unsqueeze(1)\n",
    "    active_i = (norms>1e-12).squeeze()\n",
    "    eps = 1e-6\n",
    "    C = Clatent[i]/(eps+ (norms @ norms.T))\n",
    "    Clatent_cos.append(C)\n",
    "    C  = C[active_i,:]\n",
    "    C = C[:, active_i]\n",
    "    Clatent_cos_active.append(C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#group by sae_type\n",
    "indices_spade = [i for i in range(len(exptnames)) if saetypes[i] == 'sparsemax_dist']\n",
    "indices_topk = [i for i in range(len(exptnames)) if saetypes[i] == 'topk_relu']\n",
    "indices_relu = [i for i in range(len(exptnames)) if saetypes[i] == 'relu']\n",
    "indices_jumprelu = [i for i in range(len(exptnames)) if saetypes[i] == 'jumprelu']\n",
    "indices_allsaes = [indices_relu, indices_jumprelu, indices_topk, indices_spade]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 136464 characters]",
      "text/plain": [
       "<Figure size 2000x400 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot variance explained vs sparsity for each concept\n",
    "saeorder = ['relu', 'jumprelu', 'topk_relu', 'sparsemax_dist']\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "fig, ax = plt.subplots(1, num_concepts, figsize=(4*num_concepts, 4), sharey=True)\n",
    "for j in range(num_concepts):\n",
    "    \n",
    "    count = 0\n",
    "    for ind in [indices_relu, indices_jumprelu, indices_topk, indices_spade]:\n",
    "        l0vals = l0_perconcept_all[ind,j]\n",
    "        l0vals_sorted, order = torch.sort(l0vals)\n",
    "        label = saeorder[count]\n",
    "        if label=='sparsemax_dist':\n",
    "            label = 'spade'\n",
    "        elif label=='topk_relu':\n",
    "            label = 'topk'\n",
    "        ax[j].plot(l0vals_sorted, (mse_perconcept_all[ind,j][order].cpu()/datavars_perconcept[j]), \\\n",
    "            c=colors_saetypes[saeorder[count]], alpha=1.0, marker='o', linestyle='solid', label=label)\n",
    "        count+=1\n",
    "    ax[j].set_yscale('log')\n",
    "    ax[j].set_xscale('log')\n",
    "    ax[j].legend()\n",
    "    ax[j].set_xlabel(r'Sparsity ($L_0$)')\n",
    "    if j==0:\n",
    "        ax[j].set_ylabel(r'Variance Explained')\n",
    "    ax[j].set_title(\"Dim= \"+str(train_data.truefeatures['intrinsic_dims'][j].item()))\n",
    "    ax[j].vlines(intrinsic_dims[j], 0, 1, linestyles='dashed', colors='gray')\n",
    "\n",
    "    min_l0 = min([l0_perconcept_all[i][j].min().cpu() for i in range(len(exptnames))])\n",
    "    max_l0 = max([l0_perconcept_all[i][j].max().cpu() for i in range(len(exptnames))])\n",
    "    ax[j].hlines(1-0.8, min_l0, max_l0, colors='pink', linestyles='dashed')\n",
    "\n",
    "\n",
    "fig.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 112076 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot normalized mse vs sparsity for all concepts for each SAE type\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "\n",
    "cmap = plt.get_cmap('rainbow')\n",
    "norm = plt.Normalize(vmin=0, vmax=1)\n",
    "fig, ax = plt.subplots(1, len(set(saetypes)), figsize=(len(set(saetypes))*5, 5), sharey=True)\n",
    "colors_conceptsx = [cmap(norm(intrinsic_dims[i]/max(intrinsic_dims))) for i in range(len(intrinsic_dims))]\n",
    "for i in range(len(exptnames)):\n",
    "    plotid = saeorder.index(saetypes[i])    \n",
    "    # colors = colors_concepts\n",
    "    colors = cmap(norm(intrinsic_dims/max(intrinsic_dims)))\n",
    "    \n",
    "    ax[plotid].plot(l0_perconcept_all[i,:], mse_perconcept_all[i, :]/datavars_perconcept, color='gray', linestyle='solid', alpha=0.7)\n",
    "    # ax[plotid].plot(intrinsic_dims, intrinsic_dims, color='gray', linestyle='dashed', alpha=0.1)\n",
    "    ax[plotid].scatter(l0_perconcept_all[i,:], mse_perconcept_all[i, :]/datavars_perconcept, \\\n",
    "        color=colors, marker='^', alpha=1.0)\n",
    "\n",
    "    ax[plotid].set_xlabel(r'Concept Sparsity')\n",
    "    if plotid==0:\n",
    "        ax[plotid].set_ylabel(r'normalized mse')\n",
    "        # Create custom legend only for the 5 colors_concepts and label them with their intrinsic dims\n",
    "        handles = [plt.Line2D([0], [0], marker='^', color='w', markerfacecolor=colors_conceptsx[i], markersize=10) for i in range(num_concepts)]\n",
    "        labels = [f'Dim {intrinsic_dims[i].item()}' for i in range(num_concepts)]\n",
    "        ax[plotid].legend(handles, labels)\n",
    "    # ax[plotid].set_ylabel(r'Sparsity $L_0$')\n",
    "    # ax.set_yscale('log')\n",
    "    \n",
    "    # ax[plotid].set_xticks(ticks=intrinsic_dims, labels=intrinsic_dims.tolist())\n",
    "    ax[plotid].hlines(1-0.8, l0_perconcept_all[indices_allsaes[plotid],:].min(), \\\n",
    "        l0_perconcept_all[indices_allsaes[plotid],:].max(),\\\n",
    "         colors='pink', linestyles='dashed')\n",
    "    ax[plotid].set_yscale('log')\n",
    "    ax[plotid].set_title(saetypes[i])\n",
    "    # ax[plotid].set_xlim(min(intrinsic_dims).item()-1, max(intrinsic_dims).item()+5)\n",
    "\n",
    "#for each axis, set xticks to be powers of 2 between its xmin and xmax\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "# Loop over each axis\n",
    "for countax, axis in enumerate(ax):\n",
    "    \n",
    "    minx = l0_perconcept_all[indices_allsaes[countax],:].min() * 0.9\n",
    "    maxx = l0_perconcept_all[indices_allsaes[countax],:].max() * 1.1\n",
    "    \n",
    "    axis.set_xlim(minx, maxx)\n",
    "    axis.set_xscale('log')\n",
    "    \n",
    "    # Convert to scalars before applying log2\n",
    "    min_pow = int(math.floor(math.log2(float(minx))))\n",
    "    max_pow = int(math.ceil(math.log2(float(maxx))))\n",
    "    \n",
    "    # Generate powers of 2 within range\n",
    "    xticks = [2**i for i in range(min_pow, max_pow + 1)]\n",
    "    \n",
    "    # Set xticks explicitly\n",
    "    axis.set_xticks(xticks)\n",
    "    axis.set_xticklabels([2**i for i in range(min_pow, max_pow + 1)])  # LaTeX-style labels\n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 822264 characters]",
      "text/plain": [
       "<Figure size 1200x1200 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.patches as patches\n",
    "\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "\n",
    "cmap = plt.get_cmap('rainbow')\n",
    "norm = plt.Normalize(vmin=0, vmax=1)\n",
    "fig, axes = plt.subplots(len(set(saetypes)), 4, figsize=(4*3, 3*len(set(saetypes))))\n",
    "\n",
    "ax = axes[:,1]\n",
    "colors_conceptsx = [cmap(norm(intrinsic_dims[i]/max(intrinsic_dims))) for i in range(len(intrinsic_dims))]\n",
    "for i in range(len(exptnames)):\n",
    "    plotid = saeorder.index(saetypes[i])    \n",
    "    # colors = colors_concepts\n",
    "    colors = cmap(norm(intrinsic_dims/max(intrinsic_dims)))\n",
    "    \n",
    "    ax[plotid].plot(l0_perconcept_all[i,:], mse_perconcept_all[i, :]/datavars_perconcept, color='gray', linestyle='solid', alpha=0.7)\n",
    "    # ax[plotid].plot(intrinsic_dims, intrinsic_dims, color='gray', linestyle='dashed', alpha=0.1)\n",
    "    ax[plotid].scatter(l0_perconcept_all[i,:], mse_perconcept_all[i, :]/datavars_perconcept, \\\n",
    "        color=colors, marker='^', alpha=0.8)\n",
    "\n",
    "    # ax[plotid].set_xlabel(r'Concept $L_0$')\n",
    "    if plotid==0:\n",
    "        # ax[plotid].set_ylabel(r'normalized mse')\n",
    "        # Create custom legend only for the 5 colors_concepts and label them with their intrinsic dims\n",
    "        handles = [plt.Line2D([0], [0], marker='^', color='w', markerfacecolor=colors_conceptsx[i], markersize=10) for i in range(num_concepts)]\n",
    "        labels = [f'$d=${intrinsic_dims[i].item()}' for i in range(num_concepts)]\n",
    "        ax[plotid].legend(handles, labels)\n",
    "    # ax[plotid].set_ylabel(r'Sparsity $L_0$')\n",
    "    # ax.set_yscale('log')\n",
    "    \n",
    "    # ax[plotid].set_xticks(ticks=intrinsic_dims, labels=intrinsic_dims.tolist())\n",
    "    ax[plotid].hlines(1-0.8, l0_perconcept_all[indices_allsaes[plotid],:].min(), \\\n",
    "        l0_perconcept_all[indices_allsaes[plotid],:].max(),\\\n",
    "         colors='pink', linestyles='dashed')\n",
    "    ax[plotid].set_yscale('log')\n",
    "    # ax[plotid].set_title(saetypes[i])\n",
    "    # ax[plotid].set_xlim(min(intrinsic_dims).item()-1, max(intrinsic_dims).item()+5)\n",
    "\n",
    "#for each axis, set xticks to be powers of 2 between its xmin and xmax\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "# Loop over each axis\n",
    "for countax, axis in enumerate(ax):\n",
    "    \n",
    "    minx = l0_perconcept_all[indices_allsaes[countax],:].min() - 0.1\n",
    "    maxx = l0_perconcept_all[indices_allsaes[countax],:].max() * 1.01\n",
    "    \n",
    "    axis.set_xlim(minx, maxx)\n",
    "    axis.set_xscale('log')\n",
    "    \n",
    "    # Convert to scalars before applying log2\n",
    "    min_pow = int(math.floor(math.log2(float(minx))))\n",
    "    max_pow = int(math.ceil(math.log2(float(maxx))))\n",
    "    \n",
    "    # Generate powers of 2 within range\n",
    "    xticks = [2**i for i in range(min_pow, max_pow + 1)]\n",
    "    \n",
    "    # Set xticks explicitly\n",
    "    axis.set_xticks(xticks)\n",
    "    # labelsz = [f\"$2^{i}$\" if i >= 0 else f\"$2^{{-{abs(i)}}}$\" for i in range(min_pow, max_pow + 1)]\n",
    "    labelsz = [None] * (max_pow - min_pow + 1)\n",
    "    for i in range(min_pow, max_pow+1):\n",
    "        if i < 0:\n",
    "            labelsz[i-min_pow] = f\"$2^{{-{abs(i)}}}$\"\n",
    "        elif i>=10:\n",
    "            labelsz[i-min_pow] = ''\n",
    "        else:\n",
    "            labelsz[i-min_pow] = f\"$2^{i}$\"\n",
    "    axis.set_xticklabels(labelsz)\n",
    "\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "\n",
    "\n",
    "ax = axes[:,0]\n",
    "cmap = plt.get_cmap('plasma')\n",
    "norm = plt.Normalize(vmin=0, vmax=1)\n",
    "for i in range(len(exptnames)):\n",
    "    plotid = saeorder.index(saetypes[i])\n",
    "    if saetypes[i] == 'sparsemax_dist':\n",
    "        label = 'spade'\n",
    "    elif saetypes[i] == 'topk_relu':\n",
    "        label = 'topk'\n",
    "    else:\n",
    "        label = saetypes[i]\n",
    "    # varex = mse_center[i, :].cpu() / datavars_perconcept\n",
    "    varex = mse_perconcept_all[i, :].cpu() / datavars_perconcept\n",
    "    colors = cmap(norm(varex))\n",
    "    \n",
    "    ax[plotid].plot(intrinsic_dims, l0_perconcept_all[i, :], color='gray', linestyle='solid', alpha=0.7)\n",
    "    ax[plotid].plot(intrinsic_dims, intrinsic_dims, color='gray', linestyle='dashed', alpha=0.1)\n",
    "    ax[plotid].scatter(intrinsic_dims, l0_perconcept_all[i, :], color=colors, marker='s', alpha=0.8)\n",
    "\n",
    "    # ax[plotid].set_xlabel(f'Concept dim $d$')\n",
    "    ax[plotid].set_xscale('log')\n",
    "    ax[plotid].set_xticks(ticks=intrinsic_dims, labels=intrinsic_dims.tolist())\n",
    "    \n",
    "    ax[plotid].set_yscale('log')\n",
    "    # ax[plotid].set_title(label)\n",
    "\n",
    "ax = axes[:, 2:]\n",
    "for i in range(len(exptnames)):\n",
    "    if saetypes[i]!='topk_relu':\n",
    "        if saetypes[i]=='relu':\n",
    "            gammvals = gammavals_relu\n",
    "        elif saetypes[i]=='jumprelu':\n",
    "            gammvals = gammavals_jumprelu\n",
    "        elif saetypes[i]=='sparsemax_dist':\n",
    "            gammvals = gammavals_spade\n",
    "        gami = float(exptnames[i].split('gamreg')[1].split('_')[0])\n",
    "        loc = gammvals.index(gami)\n",
    "    else:\n",
    "        kvalues_order = [256, 128, 64, 32, 16, 8, 4]\n",
    "        kval_i = int(exptnames[i].split('topk_relu_k')[1].split('_')[0])\n",
    "        loc = kvalues_order.index(kval_i)\n",
    "    concepts_now = concepts_of_latents_all[i][order_by_concept_all[i]]\n",
    "    numlatents_perconc = [torch.sum(concepts_now==j).item() for j in range(num_concepts)]\n",
    "    latentconcboundarylocs = []\n",
    "    for jn in range(len(numlatents_perconc)):\n",
    "        latentconcboundarylocs.append(sum(numlatents_perconc[:jn]))\n",
    "    latentconclabellocs = [(latentconcboundarylocs[i]+latentconcboundarylocs[i+1])/2 for i in range(len(latentconcboundarylocs)-1)]\n",
    "    if loc in [2, 4]:\n",
    "        loc = [2,4].index(loc)\n",
    "        rowid = saeorder.index(saetypes[i])\n",
    "        height, width = Clatent_cos_active[i].shape\n",
    "        ax[rowid, loc].imshow(Clatent_cos_active[i].cpu(), cmap='magma', vmin=0, vmax=1, extent=[0, width, height, 0])\n",
    "        ax[rowid, loc].vlines(latentconcboundarylocs[1:], 0, Clatent_cos_active[i].shape[0], colors='#EFEFEF', linestyles='solid',linewidth=2)\n",
    "        ax[rowid, loc].hlines(latentconcboundarylocs[1:], 0, Clatent_cos_active[i].shape[1], colors='#EFEFEF', linestyles='solid',linewidth=2)\n",
    "        # if saetypes[i]=='topk_relu':\n",
    "        #     ax[rowid, loc].set_title(f\"k={kval_i}\")\n",
    "        # else:\n",
    "        #     ax[rowid, loc].set_title(f\"$\\gamma$={gami}\")\n",
    "        # ax[rowid, loc].set_xticks([0, Clatent_cos_active[i].shape[0]])\n",
    "        # ax[rowid, loc].set_yticks([0, Clatent_cos_active[i].shape[1]])\n",
    "        ax[rowid, loc].set_xticks([])\n",
    "        ax[rowid, loc].set_yticks([])\n",
    "\n",
    "#add width between subplots\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.subplots_adjust(hspace=0.5)\n",
    "# plt.savefig('./figs/intrinsicdimfig2.pdf', bbox_inches='tight', dpi=300)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 554900 characters]",
      "text/plain": [
       "<Figure size 1720x900 with 16 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#make main figure on heterogeneity\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "\n",
    "cmap = plt.get_cmap('rainbow')\n",
    "norm = plt.Normalize(vmin=0, vmax=1)\n",
    "# fig, axes = plt.subplots(len(set(saetypes)), 4, figsize=(4*3, 3*len(set(saetypes))))\n",
    "# scaleval = 3\n",
    "# fig = plt.figure(figsize=(scaleval*8.2, scaleval*4.6))\n",
    "\n",
    "# ax1 = plt.subplot2grid((5, 8), (0, 0), rowspan=2, colspan=2)\n",
    "# ax2 = plt.subplot2grid((5, 8), (0, 2), rowspan=2, colspan=2)\n",
    "# ax3 = plt.subplot2grid((5, 8), (0, 4), rowspan=2, colspan=2)\n",
    "# ax4 = plt.subplot2grid((5, 8), (0, 6), rowspan=2, colspan=2)\n",
    "\n",
    "# ax5 = plt.subplot2grid((5, 8), (2, 0), rowspan=2, colspan=2)\n",
    "# ax6 = plt.subplot2grid((5, 8), (2, 2), rowspan=2, colspan=2)\n",
    "# ax7 = plt.subplot2grid((5, 8), (2, 4), rowspan=2, colspan=2)\n",
    "# ax8 = plt.subplot2grid((5, 8), (2, 6), rowspan=2, colspan=2)\n",
    "\n",
    "# ax9 = plt.subplot2grid((5, 8), (4, 0), rowspan=1, colspan=1)\n",
    "# ax10 = plt.subplot2grid((5, 8), (4, 1), rowspan=1, colspan=1)\n",
    "# ax11 = plt.subplot2grid((5, 8), (4, 2), rowspan=1, colspan=1)\n",
    "# ax12 = plt.subplot2grid((5, 8), (4, 3), rowspan=1, colspan=1)\n",
    "# ax13 = plt.subplot2grid((5, 8), (4, 4), rowspan=1, colspan=1)\n",
    "# ax14 = plt.subplot2grid((5, 8), (4, 5), rowspan=1, colspan=1)\n",
    "# ax15 = plt.subplot2grid((5, 8), (4, 6), rowspan=1, colspan=1)\n",
    "# ax16 = plt.subplot2grid((5, 8), (4, 7), rowspan=1, colspan=1)\n",
    "\n",
    "import matplotlib.gridspec as gridspec\n",
    "\n",
    "# Scaling factor for figure size\n",
    "scaleval = 2  # Adjust this as needed\n",
    "fig = plt.figure(figsize=(scaleval * 8.6, scaleval * 4.5))\n",
    "\n",
    "# Define the GridSpec layout (5 rows, 8 columns)\n",
    "gs = gridspec.GridSpec(5, 8, figure=fig, height_ratios=[0.8, 0.8, 0.8, 0.8, 1.4])\n",
    "\n",
    "# Define subplots using GridSpec, mirroring the subplot2grid layout\n",
    "ax1 = fig.add_subplot(gs[0:2, 0:2])\n",
    "ax2 = fig.add_subplot(gs[0:2, 2:4], sharey=ax1)\n",
    "ax3 = fig.add_subplot(gs[0:2, 4:6], sharey=ax1)\n",
    "ax4 = fig.add_subplot(gs[0:2, 6:8], sharey=ax1)\n",
    "\n",
    "ax5 = fig.add_subplot(gs[2:4, 0:2])\n",
    "ax6 = fig.add_subplot(gs[2:4, 2:4], sharey=ax5)\n",
    "ax7 = fig.add_subplot(gs[2:4, 4:6], sharey=ax5)\n",
    "ax8 = fig.add_subplot(gs[2:4, 6:8], sharey=ax5)\n",
    "\n",
    "ax9 = fig.add_subplot(gs[4, 0])\n",
    "ax10 = fig.add_subplot(gs[4, 1])\n",
    "ax11 = fig.add_subplot(gs[4, 2])\n",
    "ax12 = fig.add_subplot(gs[4, 3])\n",
    "ax13 = fig.add_subplot(gs[4, 4])\n",
    "ax14 = fig.add_subplot(gs[4, 5])\n",
    "ax15 = fig.add_subplot(gs[4, 6])\n",
    "ax16 = fig.add_subplot(gs[4, 7])\n",
    "\n",
    "axes = np.array([[ax1, ax2, ax3, ax4], [ax5, ax6, ax7, ax8], [ax9, ax11, ax13, ax15], [ax10, ax12, ax14, ax16]])\n",
    "axes = axes.T\n",
    "\n",
    "#plot normalized mse vs sparsity for all concepts for each SAE type\n",
    "ax = axes[:,1]\n",
    "colors_conceptsx = [cmap(norm(intrinsic_dims[i]/max(intrinsic_dims))) for i in range(len(intrinsic_dims))]\n",
    "for i in range(len(exptnames)):\n",
    "    plotid = saeorder.index(saetypes[i])    \n",
    "    # colors = colors_concepts\n",
    "    colors = cmap(norm(intrinsic_dims/max(intrinsic_dims)))\n",
    "    \n",
    "    ax[plotid].plot(l0_perconcept_all[i,:], mse_perconcept_all[i, :]/datavars_perconcept, color='gray', linestyle='solid', alpha=0.7)\n",
    "    # ax[plotid].plot(intrinsic_dims, intrinsic_dims, color='gray', linestyle='dashed', alpha=0.1)\n",
    "    ax[plotid].scatter(l0_perconcept_all[i,:], mse_perconcept_all[i, :]/datavars_perconcept, \\\n",
    "        color=colors, marker='^', alpha=0.8)\n",
    "    ax[plotid].set_xlabel(r\"$L_0$\", fontsize=14)\n",
    "    mseticks = [1e-3, 1e-2, 1e-1, 1e0]\n",
    "    ax[plotid].set_yticks(mseticks, mseticks, fontsize=13)\n",
    "    if plotid==0:\n",
    "        ax[plotid].set_ylabel(r\"Normalized MSE\", fontsize=16, labelpad=10)\n",
    "\n",
    "    # ax[plotid].set_xlabel(r'Concept $L_0$')\n",
    "    if plotid==0:\n",
    "        # ax[plotid].set_ylabel(r'normalized mse')\n",
    "        # Create custom legend only for the 5 colors_concepts and label them with their intrinsic dims\n",
    "        handles = [plt.Line2D([0], [0], marker='^', color='w', markerfacecolor=colors_conceptsx[i], markersize=10) for i in range(num_concepts)]\n",
    "        labels = [f'$d=${intrinsic_dims[i].item()}' for i in range(num_concepts)]\n",
    "        ax[plotid].legend(handles, labels)\n",
    "    # ax[plotid].set_ylabel(r'Sparsity $L_0$')\n",
    "    # ax.set_yscale('log')\n",
    "    \n",
    "    # ax[plotid].set_xticks(ticks=intrinsic_dims, labels=intrinsic_dims.tolist())\n",
    "    ax[plotid].hlines(1-0.8, l0_perconcept_all[indices_allsaes[plotid],:].min(), \\\n",
    "        l0_perconcept_all[indices_allsaes[plotid],:].max(),\\\n",
    "         colors='pink', linestyles='dashed')\n",
    "    ax[plotid].set_yscale('log')\n",
    "    # ax[plotid].set_title(saetypes[i])\n",
    "    # ax[plotid].set_xlim(min(intrinsic_dims).item()-1, max(intrinsic_dims).item()+5)\n",
    "\n",
    "#for each axis, set xticks to be powers of 2 between its xmin and xmax\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "# Loop over each axis\n",
    "for countax, axis in enumerate(ax):\n",
    "    \n",
    "    minx = l0_perconcept_all[indices_allsaes[countax],:].min() - 0.1\n",
    "    maxx = l0_perconcept_all[indices_allsaes[countax],:].max() * 1.01\n",
    "    \n",
    "    axis.set_xlim(minx, maxx)\n",
    "    axis.set_xscale('log')\n",
    "    \n",
    "    # Convert to scalars before applying log2\n",
    "    min_pow = int(math.floor(math.log2(float(minx))))\n",
    "    max_pow = int(math.ceil(math.log2(float(maxx))))\n",
    "    \n",
    "    # Generate powers of 2 within range\n",
    "    xticks = [2**i for i in range(min_pow, max_pow + 1)]\n",
    "    \n",
    "    # Set xticks explicitly\n",
    "    axis.set_xticks(xticks, labels=xticks, fontsize=13)\n",
    "    # labelsz = [f\"$2^{i}$\" if i >= 0 else f\"$2^{{-{abs(i)}}}$\" for i in range(min_pow, max_pow + 1)]\n",
    "    labelsz = [None] * (max_pow - min_pow + 1)\n",
    "    for i in range(min_pow, max_pow+1):\n",
    "        if i < 0:\n",
    "            labelsz[i-min_pow] = f\"$2^{{-{abs(i)}}}$\"\n",
    "        elif i>=10:\n",
    "            labelsz[i-min_pow] = ''\n",
    "        else:\n",
    "            labelsz[i-min_pow] = f\"$2^{i}$\"\n",
    "    axis.set_xticklabels(labelsz)\n",
    "    \n",
    "\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "\n",
    "#plot sparsity vs intrinsic dimension of each concept\n",
    "ax = axes[:,0]\n",
    "cmap = plt.get_cmap('plasma')\n",
    "norm = plt.Normalize(vmin=0, vmax=1)\n",
    "for i in range(len(exptnames)):\n",
    "    plotid = saeorder.index(saetypes[i])\n",
    "    if saetypes[i] == 'sparsemax_dist':\n",
    "        label = 'spade'\n",
    "    elif saetypes[i] == 'topk_relu':\n",
    "        label = 'topk'\n",
    "    else:\n",
    "        label = saetypes[i]\n",
    "    # varex = mse_center[i, :].cpu() / datavars_perconcept\n",
    "    varex = mse_perconcept_all[i, :].cpu() / datavars_perconcept\n",
    "    colors = cmap(norm(varex))\n",
    "    \n",
    "    ax[plotid].plot(intrinsic_dims, l0_perconcept_all[i, :], color='gray', linestyle='solid', alpha=0.7)\n",
    "    ax[plotid].plot(intrinsic_dims, intrinsic_dims, color='gray', linestyle='dashed', alpha=0.1)\n",
    "    ax[plotid].scatter(intrinsic_dims, l0_perconcept_all[i, :], color=colors, marker='s', alpha=0.8)\n",
    "\n",
    "    # ax[plotid].set_xlabel(f'Concept dim $d$')\n",
    "    ax[plotid].set_xscale('log')\n",
    "    ax[plotid].set_xticks(ticks=intrinsic_dims, labels=intrinsic_dims.tolist(), fontsize=13)\n",
    "    l0ticks = [1e0, 1e1, 1e2]\n",
    "    ax[plotid].set_yticks(l0ticks, l0ticks, fontsize=13)\n",
    "    ax[plotid].set_xlabel(r\"$d$\", fontsize=14)\n",
    "    if plotid==0:\n",
    "        ax[plotid].set_ylabel(r\"$L_0$\", fontsize=16, labelpad=10)\n",
    "    \n",
    "    ax[plotid].set_yscale('log')\n",
    "    \n",
    "    # ax[plotid].set_title(label)\n",
    "\n",
    "#plot latent-latent correlations for active latents\n",
    "ax = axes[:, 2:]\n",
    "for i in range(len(exptnames)):\n",
    "    if saetypes[i]!='topk_relu':\n",
    "        if saetypes[i]=='relu':\n",
    "            gammvals = gammavals_relu\n",
    "        elif saetypes[i]=='jumprelu':\n",
    "            gammvals = gammavals_jumprelu\n",
    "        elif saetypes[i]=='sparsemax_dist':\n",
    "            gammvals = gammavals_spade\n",
    "        gami = float(exptnames[i].split('gamreg')[1].split('_')[0])\n",
    "        loc = gammvals.index(gami)\n",
    "    else:\n",
    "        kvalues_order = [256, 128, 64, 32, 16, 8, 4]\n",
    "        kval_i = int(exptnames[i].split('topk_relu_k')[1].split('_')[0])\n",
    "        loc = kvalues_order.index(kval_i)\n",
    "    concepts_now = concepts_of_latents_all[i][order_by_concept_all[i]]\n",
    "    numlatents_perconc = [torch.sum(concepts_now==j).item() for j in range(num_concepts)]\n",
    "    latentconcboundarylocs = []\n",
    "    for jn in range(len(numlatents_perconc)):\n",
    "        latentconcboundarylocs.append(sum(numlatents_perconc[:jn]))\n",
    "    latentconclabellocs = [(latentconcboundarylocs[i]+latentconcboundarylocs[i+1])/2 for i in range(len(latentconcboundarylocs)-1)]\n",
    "    if loc in [2, 4]:\n",
    "        loc = [4, 2].index(loc)\n",
    "        rowid = saeorder.index(saetypes[i])\n",
    "        height, width = Clatent_cos_active[i].shape\n",
    "        ax[rowid, loc].imshow(Clatent_cos_active[i].cpu(), cmap='magma', vmin=0, vmax=1, extent=[0, width, height, 0])\n",
    "        ax[rowid, loc].vlines(latentconcboundarylocs[1:], 0, Clatent_cos_active[i].shape[0], colors='#EFEFEF', linestyles='solid',linewidth=1)\n",
    "        ax[rowid, loc].hlines(latentconcboundarylocs[1:], 0, Clatent_cos_active[i].shape[1], colors='#EFEFEF', linestyles='solid',linewidth=1)\n",
    "        ax[rowid, loc].set_title(f\"$L_0$={l0_all[i].item():.0f}\", fontsize=13)\n",
    "        # if saetypes[i]=='topk_relu':\n",
    "        #     ax[rowid, loc].set_title(f\"k={kval_i}\")\n",
    "        # else:\n",
    "        #     ax[rowid, loc].set_title(f\"$\\gamma$={gami}\")\n",
    "        # ax[rowid, loc].set_xticks([0, Clatent_cos_active[i].shape[0]])\n",
    "        # ax[rowid, loc].set_yticks([0, Clatent_cos_active[i].shape[1]])\n",
    "        ax[rowid, loc].set_xticks([])\n",
    "        ax[rowid, loc].set_yticks([])\n",
    "\n",
    "#add width between subplots\n",
    "for axx in axes.flatten():\n",
    "    axx.set_aspect('auto')\n",
    "    axx.spines['right'].set_visible(False)\n",
    "    axx.spines['top'].set_visible(False)\n",
    "\n",
    "for ax in [ax2, ax3, ax4, ax6, ax7, ax8]:\n",
    "    ax.tick_params(labelleft=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "# fig.subplots_adjust(hspace=0.5)\n",
    "# plt.savefig('./figs/hetero_mainfigX.pdf', bbox_inches='tight', dpi=300)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 117864 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot sparsity vs intrinsic dimension for each SAE type\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "\n",
    "cmap = plt.get_cmap('rainbow')\n",
    "norm = plt.Normalize(vmin=0, vmax=1)\n",
    "fig, ax = plt.subplots(1, len(set(saetypes)), figsize=(len(set(saetypes))*5, 5), sharey=True)\n",
    "\n",
    "for i in range(len(exptnames)):\n",
    "    plotid = saeorder.index(saetypes[i])\n",
    "    if saetypes[i] == 'sparsemax_dist':\n",
    "        label = 'spade'\n",
    "    elif saetypes[i] == 'topk_relu':\n",
    "        label = 'topk'\n",
    "    else:\n",
    "        label = saetypes[i]\n",
    "    # varex = mse_center[i, :].cpu() / datavars_perconcept\n",
    "    varex = mse_perconcept_all[i, :].cpu() / datavars_perconcept\n",
    "    colors = cmap(norm(varex))\n",
    "    \n",
    "    ax[plotid].plot(intrinsic_dims, l0_perconcept_all[i, :], color='gray', linestyle='solid', alpha=0.7)\n",
    "    ax[plotid].plot(intrinsic_dims, intrinsic_dims, color='gray', linestyle='dashed', alpha=0.1)\n",
    "    ax[plotid].scatter(intrinsic_dims, l0_perconcept_all[i, :], color=colors, marker='^', alpha=1.0)\n",
    "\n",
    "    ax[plotid].set_xlabel(r'Intrinsic Dim')\n",
    "    if plotid==0:\n",
    "        ax[plotid].set_ylabel(r'Sparsity $L_0$')\n",
    "    # ax[plotid].set_ylabel(r'Sparsity $L_0$')\n",
    "    # ax.set_yscale('log')\n",
    "    ax[plotid].set_xscale('log')\n",
    "    ax[plotid].set_xticks(ticks=intrinsic_dims, labels=intrinsic_dims.tolist())\n",
    "    \n",
    "    ax[plotid].set_yscale('log')\n",
    "    ax[plotid].set_title(label)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 148988 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot normalized mse vs intrinsic dimension of concepts for each SAE type\n",
    "colors_saetypes = {'relu': '#377eb8', 'jumprelu':'#984ea3' , 'topk_relu': '#4daf4a', 'sparsemax_dist': '#ff7f00'}\n",
    "\n",
    "cmap = plt.get_cmap('rainbow')\n",
    "norm = plt.Normalize(vmin=0, vmax=1)\n",
    "fig, ax = plt.subplots(1, len(set(saetypes)), figsize=(len(set(saetypes))*5, 5), sharey=True)\n",
    "\n",
    "for i in range(len(exptnames)):\n",
    "    plotid = saeorder.index(saetypes[i])\n",
    "    if saetypes[i] == 'sparsemax_dist':\n",
    "        label = 'spade'\n",
    "    elif saetypes[i] == 'topk_relu':\n",
    "        label = 'topk'\n",
    "    else:\n",
    "        label = saetypes[i]\n",
    "    # varex = mse_center[i, :].cpu() / datavars_perconcept\n",
    "    # varex = mse_perconcept_all[i, :].cpu() / datavars_perconcept\n",
    "    varex = l0_perconcept_all[i,:]/l0_perconcept_all.max()\n",
    "    colors = cmap(norm(varex))\n",
    "    \n",
    "    ax[plotid].plot(intrinsic_dims, mse_perconcept_all[i, :]/datavars_perconcept, color='gray', linestyle='solid', alpha=0.7)\n",
    "    # ax[plotid].plot(intrinsic_dims, intrinsic_dims, color='gray', linestyle='dashed', alpha=0.1)\n",
    "    ax[plotid].scatter(intrinsic_dims, mse_perconcept_all[i, :]/datavars_perconcept, color=colors, marker='^', alpha=1.0)\n",
    "\n",
    "    ax[plotid].set_xlabel(r'Intrinsic Dim')\n",
    "    if plotid==0:\n",
    "        ax[plotid].set_ylabel(r'normalized mse')\n",
    "    # ax[plotid].set_ylabel(r'Sparsity $L_0$')\n",
    "    # ax.set_yscale('log')\n",
    "    ax[plotid].set_xscale('log')\n",
    "    ax[plotid].set_xticks(ticks=intrinsic_dims, labels=intrinsic_dims.tolist())\n",
    "    ax[plotid].hlines(1-0.8, min(intrinsic_dims).item(), max(intrinsic_dims).item(),\\\n",
    "         colors='pink', linestyles='dashed')\n",
    "    ax[plotid].set_yscale('log')\n",
    "    ax[plotid].set_title(label)\n",
    "    ax[plotid].set_xlim(min(intrinsic_dims).item()-1, max(intrinsic_dims).item()+5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1007076 characters]",
      "text/plain": [
       "<Figure size 1400x800 with 28 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#plot correlations between pairs of latents over all concepts\n",
    "kvalues_order = [256, 128, 64, 32, 16, 8, 4]\n",
    "fig, ax = plt.subplots(len(set(saetypes)), 7, figsize=(14, 8))\n",
    "for i in range(len(exptnames)):\n",
    "    if saetypes[i]!='topk_relu':\n",
    "        if saetypes[i]=='relu':\n",
    "            gammvals = gammavals_relu\n",
    "        elif saetypes[i]=='jumprelu':\n",
    "            gammvals = gammavals_jumprelu\n",
    "        elif saetypes[i]=='sparsemax_dist':\n",
    "            gammvals = gammavals_spade\n",
    "        gami = float(exptnames[i].split('gamreg')[1].split('_')[0])\n",
    "        loc = gammvals.index(gami)\n",
    "    else:\n",
    "        kvalues_oder = [256, 128, 64, 32, 16, 8, 4]\n",
    "        kval_i = int(exptnames[i].split('topk_relu_k')[1].split('_')[0])\n",
    "        loc = kvalues_order.index(kval_i)\n",
    "    \n",
    "    rowid = saeorder.index(saetypes[i])\n",
    "    ax[rowid, loc].imshow(Clatent_cos_active[i].cpu(), cmap='viridis', vmin=0, vmax=1)\n",
    "    # ax[rowid, loc].set_title(exptnames[i])\n",
    "    if loc==0:\n",
    "        ax[rowid, 0].set_ylabel(saetypes[i])\n",
    "    # if rowid==0:\n",
    "    if saetypes[i]=='topk_relu':\n",
    "        ax[rowid, loc].set_title(f\"k={kval_i}\")\n",
    "    else:\n",
    "        ax[rowid, loc].set_title(f\"gamreg={gami}\")\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
