{
 "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",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n"
   ]
  },
  {
   "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",
    "dataname = '2dgaussian_diffmag'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from functions.get_data import data_n_loaders\n",
    "# LAB_DIR = os.environ['USERDIR'] #change as needed\n",
    "# DATA_PATH = LAB_DIR+'/data'\n",
    "\n",
    "#load entire train data and transform parameters\n",
    "batch_size = 512 #doesnt matter since we disregard the dataloader\n",
    "_, _,train_data, _, mean, scale = data_n_loaders(dataname, batch_size, return_data=True, \\\n",
    "      data_path=DATA_PATH, standardise_data = True, return_transform=True)\n",
    "mean = mean.squeeze()\n",
    "train_data_raw = (train_data.data/scale) + mean\n",
    "\n",
    "\n",
    "#sample equal number of points from each concept\n",
    "torch.manual_seed(0)\n",
    "numpoints_perconcept = 1000\n",
    "num_concepts = 6\n",
    "data_subset_perconcept = []\n",
    "labels_subset_perconcept = []\n",
    "for i in range(num_concepts):\n",
    "    indices = train_data.labels == i\n",
    "    rand_indices_i = torch.randperm(indices.sum())[:numpoints_perconcept]\n",
    "    data_subset_perconcept.append(train_data_raw[indices][rand_indices_i])\n",
    "    labels_subset_perconcept.append(train_data.labels[indices][rand_indices_i])\n",
    "data_subset_raw = torch.cat(data_subset_perconcept, dim=0)\n",
    "data_subset = (data_subset_raw - mean) * scale #transform same as training data (to input to model)\n",
    "labels_subset = torch.cat(labels_subset_perconcept, dim=0)"
   ]
  },
  {
   "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",
    "            if ('gamreg' not in dir and 'topk' not in dir) or 'k15' in dir:\n",
    "                continue\n",
    "            if 'w128' not in dir:\n",
    "                continue\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",
    "kvals_all = []\n",
    "for exptname in exptnames:\n",
    "    if 'topk' in exptname:\n",
    "        kvals_all.append(int(exptname.split('_k')[1].split('_')[0]))\n",
    "    else:\n",
    "        kvals_all.append(0)\n",
    "gammavals_all = []\n",
    "for exptname in exptnames:\n",
    "    if 'gamreg' in exptname:\n",
    "        gammavals_all.append(float(exptname.split('_gamreg')[1].split('_')[0]))\n",
    "    else:\n",
    "        gammavals_all.append(0)\n",
    "max_iters = []\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_iters.append(max(epochs))\n",
    "\n",
    "# choose which models (hyperparams) to plot\n",
    "kvals_plot = [2*i for i in range(1, 11)]\n",
    "gammavals_relu_plot = [1.0, 0.5, 0.005, 0.001, 0.0005, 0.0001, 5e-05, 1e-05, 5e-06, 1e-06]\n",
    "gammavals_jumprelu_plot = [0.1, 0.05, 0.005, 0.00125, 0.001, 0.0005, 0.0001, 5e-05, 1e-05, 5e-06]\n",
    "gammavals_spade_plot = [0.5, 0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001, 5e-05, 1e-05]\n",
    "relu_indices2 = [i for i in range(len(exptnames)) if saetypes[i] == 'relu' and gammavals_all[i] in gammavals_relu_plot]\n",
    "jumprelu_indices2 = [i for i in range(len(exptnames)) if saetypes[i] == 'jumprelu' and gammavals_all[i] in gammavals_jumprelu_plot]\n",
    "topk_indices2 = [i for i in range(len(exptnames)) if saetypes[i] == 'topk_relu' and kvals_all[i] in kvals_plot]\n",
    "spade_indices2 = [i for i in range(len(exptnames)) if saetypes[i] == 'sparsemax_dist' and gammavals_all[i] in gammavals_spade_plot]\n",
    "numpoints_persae = len(relu_indices2)\n",
    "allindices2 = relu_indices2 + jumprelu_indices2 + topk_indices2 + spade_indices2\n",
    "exptnames = [exptnames[i] for i in allindices2]\n",
    "saetypes = [saetypes[i] for i in allindices2]\n",
    "kvals_all = [kvals_all[i] for i in allindices2]\n",
    "gammavals_all = [gammavals_all[i] for i in allindices2]\n",
    "max_iters = [max_iters[i] for i in allindices2]\n",
    "assert len(set(max_iters))==1 #max-iters same for all models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load trained models for chosen hyperparams\n",
    "from models import SAE\n",
    "width = 128 #of sae\n",
    "modelsall = []\n",
    "dimdata = 2\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(dimdata, 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_iters[i]}epochs.pt\")\n",
    "    modeli.load_state_dict(state_dict['model'])\n",
    "    modelsall.append(modeli)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#get latent representations (sparse codes) on data subset\n",
    "latent_reps_all = [] #latent reps on data subset for each model\n",
    "mse_all = []\n",
    "meannorm2_all = []\n",
    "datavar_all = []\n",
    "device = 'cuda'\n",
    "for i in range(len(exptnames)):\n",
    "    modeli = modelsall[i].to(device)\n",
    "    with torch.no_grad():\n",
    "        pts_recon, latent_reps = modeli(data_subset.to(device), return_hidden=True)\n",
    "        mse = ((pts_recon - data_subset.to(device))**2).sum(dim=-1).mean().item()\n",
    "        mean_norm2 = (data_subset**2).sum(dim=-1).mean().item()\n",
    "        mse_all.append(mse)\n",
    "        meannorm2_all.append(mean_norm2)\n",
    "        datavar = data_subset.var(dim=0).sum().item()\n",
    "        datavar_all.append(datavar)\n",
    "        latent_reps_all.append(latent_reps.detach().cpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saes_order = ['relu', 'jumprelu', 'topk_relu', 'sparsemax_dist']\n",
    "\n",
    "#get f1 score of each latent for each class\n",
    "eps = 1e-12\n",
    "f1_all = torch.zeros((len(exptnames), num_concepts, width))\n",
    "precision_all = torch.zeros((len(exptnames), num_concepts, width))\n",
    "recall_all = torch.zeros((len(exptnames), num_concepts, width))\n",
    "l0_all = torch.zeros((len(exptnames)))\n",
    "for i in range(len(exptnames)):\n",
    "    latents_i = latent_reps_all[i]\n",
    "    # threshold = torch.quantile(latents_i[latents_i>0.0].flatten(), 0.1) #10th percentile of data\n",
    "    threshold = 1e-6\n",
    "    l0_all[i] = (latents_i>threshold).sum(dim=-1).float().mean()\n",
    "    for c in range(num_concepts):\n",
    "        truepos = (latents_i[labels_subset==c,:]>threshold).sum(0).float()\n",
    "        falsepos = (latents_i[labels_subset!=c,:]>threshold).sum(0).float()\n",
    "        allpos = (labels_subset==c).sum().float()\n",
    "        pospred = (latents_i>threshold).sum(0).float()\n",
    "        precision = truepos/(pospred+eps)\n",
    "        recall = truepos/(allpos+eps)\n",
    "        f1 = 2*precision*recall/(precision+recall+eps)\n",
    "        f1_all[i,c, :] = f1\n",
    "        precision_all[i,c, :] = precision\n",
    "        recall_all[i,c, :] = recall\n",
    "monosemantic_latents = torch.argmax(f1_all, dim=-1)\n",
    "ordered_latents = torch.argsort(f1_all, dim=-1, descending=True)\n",
    "f1_sorted = torch.gather(f1_all, dim=-1, index=ordered_latents)\n",
    "precision_sorted = torch.gather(precision_all, dim=-1, index=ordered_latents)\n",
    "recall_sorted = torch.gather(recall_all, dim=-1, index=ordered_latents)\n",
    "\n",
    "#group gamma, k values by sae_type\n",
    "gammavals_saegrouped = []\n",
    "kvals_saegrouped = []\n",
    "for i in range(len(saes_order)):\n",
    "    gammavals_saegrouped.append(sorted([gammavals_all[j] for j in range(len(exptnames)) if saetypes[j] == saes_order[i]]))\n",
    "    kvals_saegrouped.append(sorted([kvals_all[j] for j in range(len(exptnames)) if saetypes[j] == saes_order[i]]))\n",
    "\n",
    "#rearrange f1 scores and L0 (sparsity) for plotting\n",
    "l0_lineplot = torch.zeros((len(saes_order), numpoints_persae))\n",
    "f1mean_lineplot = torch.zeros((num_concepts, len(saes_order), numpoints_persae))\n",
    "f1std_lineplot = torch.zeros((num_concepts, len(saes_order), numpoints_persae))\n",
    "shownum = 5\n",
    "for i in range(num_concepts):\n",
    "    for j in range(len(exptnames)):\n",
    "        colid = saes_order.index(saetypes[j])  \n",
    "        if saetypes[j] == 'topk_relu':\n",
    "            l0id = kvals_saegrouped[colid].index(kvals_all[j])\n",
    "        else:\n",
    "            l0id = gammavals_saegrouped[colid].index(gammavals_all[j])\n",
    "        l0_lineplot[colid, l0id] =  l0_all[j]\n",
    "        f1mean_lineplot[i, colid, l0id] = f1_sorted[j, i, :shownum].mean()\n",
    "        f1std_lineplot[i, colid, l0id] = f1_sorted[j, i, :shownum].std()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#get latent receptive fields\n",
    "xlims = [-2, 2]\n",
    "ylims = [-2, 2]\n",
    "numpoints_axis = 1000\n",
    "x, y = torch.meshgrid(torch.linspace(xlims[0], xlims[1], numpoints_axis), torch.linspace(ylims[0], ylims[1], numpoints_axis))\n",
    "Xmesh_raw = torch.stack([x, y], dim=2).view(-1, 2)\n",
    "Xmesh = scale*(Xmesh_raw - mean.reshape((1, 2))) #transform to pass to the model\n",
    "rfall = []\n",
    "heatall = []\n",
    "for i in range(len(exptnames)):\n",
    "    with torch.no_grad():\n",
    "        modeli = modelsall[i]\n",
    "        _, latents = modeli(Xmesh.to(device), return_hidden=True)\n",
    "        latents_s = latents.reshape((numpoints_axis, numpoints_axis, width)).T\n",
    "        latents_s = latents_s.detach().cpu()\n",
    "        rfall.append((latents_s>1e-6).float())\n",
    "        heatall.append(latents_s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# choose models to plot receptive fields for\n",
    "plot_concepts = [2,5] \n",
    "concept_colors = ['orange', 'purple']\n",
    "exptnames_plot = []\n",
    "plot_indices = []\n",
    "for i in range(len(saes_order)):\n",
    "    temp = f1mean_lineplot[plot_concepts, :,:].mean(dim=0)[i,:]\n",
    "    index_l0_highestmeanf1 = temp.argmax().item()\n",
    "    if saes_order[i] == 'topk_relu':\n",
    "        l0_plot_index = kvals_saegrouped[i][index_l0_highestmeanf1]\n",
    "        index_l0_saei = [j for j in range(len(exptnames)) if saetypes[j] == saes_order[i] and kvals_all[j] == l0_plot_index]\n",
    "    else:\n",
    "        l0_plot_index = gammavals_saegrouped[i][index_l0_highestmeanf1]\n",
    "        index_l0_saei = [j for j in range(len(exptnames)) if saetypes[j] == saes_order[i] and gammavals_all[j] == l0_plot_index]\n",
    "    assert len(index_l0_saei) == 1\n",
    "    plot_indices.append(index_l0_saei[0])\n",
    "    exptnames_plot.append(exptnames[index_l0_saei[0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "still801_relu_gamreg0.0001_w128_021725 done\n",
      "logic576_jumprelu_gamreg5e-05_w128_021725 done\n",
      "guava893_topk_relu_k18_w128_021725 done\n",
      "light893_spade_gamreg0.001_w128_021725 done\n",
      "Time for spectral clustering 40 data-data corrs = 18.93s\n"
     ]
    }
   ],
   "source": [
    "#data-data correlations, latent-latent correlations, spectral clusters\n",
    "Cdatacos = []\n",
    "Clatentcos = []\n",
    "concept_of_latents = torch.zeros((len(exptnames), width))\n",
    "for j in range(len(exptnames)):\n",
    "    for i in range(width):\n",
    "        concept_of_latents[j,i] = torch.argmax(f1_all[j, :, i])\n",
    "latents_order = torch.argsort(concept_of_latents, dim=-1)\n",
    "for i in range(len(exptnames)):    \n",
    "    latents = latent_reps_all[i][:,latents_order[i]]\n",
    "    deadlatents = (latents>1e-12).sum(dim=0)==0.0\n",
    "    latents = torch.cat([latents[:,~deadlatents], latents[:,deadlatents]], dim=1)\n",
    "    \n",
    "    #sparse code correlations for data pairs\n",
    "    Corr = latents @ latents.T\n",
    "    norm = torch.norm(latents, dim=1).unsqueeze(1)\n",
    "    Cdatacos.append( Corr/ (norm @ norm.T + 1e-12) )\n",
    "    \n",
    "    #sparse code correlations for latent pairs\n",
    "    Corr = latents.T @ latents\n",
    "    norm = torch.norm(latents, dim=0).unsqueeze(1)\n",
    "    Corr = Corr/ (norm @ norm.T + 1e-12)\n",
    "    Clatentcos.append(Corr)\n",
    "\n",
    "#compute stable rank of the correlation matrix\n",
    "stableranks = []\n",
    "for i in range(len(exptnames)):\n",
    "    if i in plot_indices:\n",
    "        sr = torch.trace(Cdatacos[i])/torch.linalg.matrix_norm(Cdatacos[i], ord=2)\n",
    "        stableranks.append(sr.item())\n",
    "    else:\n",
    "        stableranks.append(-1) #not computed since not plotted here\n",
    "\n",
    "\n",
    "#spectral clustering on data-data correlation matrix\n",
    "from sklearn.cluster import spectral_clustering as sc\n",
    "dataorder_all = []\n",
    "clusters_all = []\n",
    "clustersizes_all = []\n",
    "tic = time.perf_counter()\n",
    "for i in range(len(exptnames)):\n",
    "    if i in plot_indices:\n",
    "        n_clusters = math.ceil(stableranks[i])\n",
    "        labels = sc(Cdatacos[i].cpu().numpy(), n_clusters=n_clusters, random_state=0)\n",
    "        labels = torch.tensor(labels)\n",
    "        order = torch.argsort(labels)\n",
    "        # labels, order = spectral_clustering(Cdatacos[i], n_clusters=math.ceil(stableranks[i]), normalized=False, return_clusteredorder=True)\n",
    "        dataorder_all.append(order)\n",
    "        clusters_all.append(labels)\n",
    "        clustersizes = []\n",
    "        for j in range(n_clusters):\n",
    "            clustersizes.append((labels==j).sum())\n",
    "        clustersizes_all.append(clustersizes)\n",
    "        print(f\"{exptnames[i]} done\")\n",
    "    else: #not plotted so spectral clustering not computed\n",
    "        dataorder_all.append([])\n",
    "        clusters_all.append([])\n",
    "        clustersizes_all.append([])\n",
    "toc = time.perf_counter()\n",
    "print(f\"Time for spectral clustering {len(exptnames)} data-data corrs = {(toc-tic):.2f}s\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 702484 characters]",
      "text/plain": [
       "<Figure size 2520x1140 with 20 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#show only final receptive fields for all models, data in separate figures, each having 4X5 subplots (one per neuron)\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "plt.rcParams.update({'font.size': 13})\n",
    "row_heights = [1, 0.5, 0.5, 0.5, 0.5]\n",
    "numcols = len(exptnames_plot) \n",
    "numrows = 3\n",
    "# fig, axs = plt.subplots(numrows, numcols, figsize=(2.5*numcols, 2.5*(3)), sharey='row', constrained_layout=True, gridspec_kw={'height_ratios': row_heights})\n",
    "scaleval = 3\n",
    "\n",
    "#define figure layout and subplots\n",
    "plt.figure(figsize=(scaleval*8.4,scaleval*(3.8)))\n",
    "ax1 = plt.subplot2grid((4, 8), (0, 0), rowspan=2, colspan=2)\n",
    "ax2 = plt.subplot2grid((4, 8), (0, 2), rowspan=2, colspan=2)\n",
    "ax3 = plt.subplot2grid((4, 8), (0, 4), rowspan=2, colspan=2)\n",
    "ax4 = plt.subplot2grid((4, 8), (0, 6), rowspan=2, colspan=2)\n",
    "ax5 = plt.subplot2grid((4, 8), (2, 0), rowspan=1, colspan=1)\n",
    "ax6 = plt.subplot2grid((4, 8), (2, 1), rowspan=1, colspan=1)\n",
    "ax7 = plt.subplot2grid((4, 8), (2, 2), rowspan=1, colspan=1)\n",
    "ax8 = plt.subplot2grid((4, 8), (2,3), rowspan=1, colspan=1)\n",
    "ax9 = plt.subplot2grid((4, 8), (2, 4), rowspan=1, colspan=1)\n",
    "ax10 = plt.subplot2grid((4, 8), (2, 5), rowspan=1, colspan=1)\n",
    "ax11 = plt.subplot2grid((4, 8), (2, 6), rowspan=1, colspan=1)\n",
    "ax12 = plt.subplot2grid((4, 8), (2, 7), rowspan=1, colspan=1)\n",
    "ax13 = plt.subplot2grid((4, 8), (3, 0), rowspan=1, colspan=1)\n",
    "ax14 = plt.subplot2grid((4, 8), (3, 1), rowspan=1, colspan=1)\n",
    "ax15 = plt.subplot2grid((4, 8), (3, 2), rowspan=1, colspan=1)\n",
    "ax16 = plt.subplot2grid((4, 8), (3, 3), rowspan=1, colspan=1)\n",
    "ax17 = plt.subplot2grid((4, 8), (3, 4), rowspan=1, colspan=1)\n",
    "ax18 = plt.subplot2grid((4, 8), (3, 5), rowspan=1, colspan=1)\n",
    "ax19 = plt.subplot2grid((4, 8), (3, 6), rowspan=1, colspan=1)\n",
    "ax20 = plt.subplot2grid((4, 8), (3, 7), rowspan=1, colspan=1)\n",
    "axs = np.array([ [ax1, ax2, ax3, ax4], [ax5, ax7, ax9, ax11], [ax6, ax8, ax10, ax12], [ax13, ax15, ax17, ax19], [ax14, ax16, ax18, ax20] ])\n",
    "\n",
    "torch.manual_seed(0)\n",
    "concept_names = ['Concept '+str(plot_concepts[i]) for i in range(len(plot_concepts))]\n",
    "\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "# Define a custom colormap that interpolates from white to your desired color\n",
    "colormaps = []\n",
    "for k in range(len(plot_concepts)):\n",
    "    custom_cmap = LinearSegmentedColormap.from_list(\"custom\", [\"white\", concept_colors[k]])\n",
    "    colormaps.append(custom_cmap)\n",
    "indices = torch.randperm(len(data_subset))[:500]\n",
    "labels_forindices = labels_subset[indices]\n",
    "\n",
    "#make plots\n",
    "for j in range(len(exptnames_plot)):\n",
    "    for k in range(len(plot_concepts)):\n",
    "\n",
    "        #show datapoints\n",
    "        axs[k+1,j].scatter(data_subset_raw[indices,0], data_subset_raw[indices,1], \\\n",
    "            s=1, c='dimgray')\n",
    "        axs[k+1,j].scatter(data_subset_raw[indices,0][labels_forindices==plot_concepts[k]], data_subset_raw[indices,1][labels_forindices==plot_concepts[k]], \\\n",
    "            s=1, c='k') #active concept- whose receptive fields are shown\n",
    "            \n",
    "        #visualize receptive fields\n",
    "        heatjk = heatall[plot_indices[j]][monosemantic_latents[plot_indices[j],plot_concepts[k]].item(),:,:]\n",
    "        axs[k+1,j].imshow(heatjk, alpha=0.6,\n",
    "            extent=[xlims[0], xlims[1], ylims[0], ylims[1]], origin='lower', cmap=colormaps[k])\n",
    "        \n",
    "        #set axis labels, ticks, limits\n",
    "        axs[k+1,j].set_xlabel(r\"$d_1$\", fontsize=13)\n",
    "        if j==0 and k==0:\n",
    "            axs[k+1,j].set_ylabel(r\"$d_2$\", fontsize=13)\n",
    "        axs[k+1,j].set_xticks([])\n",
    "        axs[k+1,j].set_yticks([])\n",
    "        axs[k+1,j].set_xlim(*xlims)\n",
    "        axs[k+1,j].set_ylim(*ylims)\n",
    "        #draw boundaries around receptive fields\n",
    "        axs[k+1,j].contour(heatjk, levels=[torch.quantile(heatjk.flatten(), 0.1).item()], colors=concept_colors[k], \\\n",
    "            linewidths=2, extent=[xlims[0], xlims[1], ylims[0], ylims[1]])\n",
    "        # if j==len(exptnames)-1: #sparsemax\n",
    "        \n",
    "        #show encoder weights\n",
    "        protos = modelsall[plot_indices[j]].Ae.detach().cpu()\n",
    "        latentid = monosemantic_latents[plot_indices[j],plot_concepts[k]].item()\n",
    "        \n",
    "        #points at each encoder weight for SpaDE\n",
    "        if saetypes[plot_indices[j]] == 'sparsemax_dist':\n",
    "            protos = protos/scale + mean\n",
    "            axs[k+1,j].scatter(protos[latentid,0], protos[latentid,1], marker='d', color=concept_colors[k], s=40)\n",
    "        \n",
    "        #quivers for encoder weights (directions) for ReLU, TopK, JumpReLU\n",
    "        else:\n",
    "            protos = protos*scale\n",
    "            protos = protos/(1e-12+ protos.norm(dim=1, keepdim=True))\n",
    "            bdec = modelsall[plot_indices[j]].bd.detach().cpu().squeeze()\n",
    "            bdec = bdec/scale + mean\n",
    "            axs[k+1,j].quiver(bdec[0], bdec[1], protos[latentid,0], protos[latentid,1], scale=0.5, scale_units='xy', angles='xy', \\\n",
    "                color=concept_colors[k], width=0.015)\n",
    "\n",
    "for k in range(len(plot_concepts)):\n",
    "    for i in range(len(saes_order)):\n",
    "\n",
    "        #show F1-scores\n",
    "        axs[0, i].fill_between(l0_lineplot[i, :], \n",
    "                   f1mean_lineplot[plot_concepts[k], i, :] - f1std_lineplot[plot_concepts[k], i, :], \n",
    "                   f1mean_lineplot[plot_concepts[k], i, :] + f1std_lineplot[plot_concepts[k], i, :], \n",
    "                   color=concept_colors[k], alpha=0.1)\n",
    "        axs[0, i].plot(l0_lineplot[i, :], f1mean_lineplot[plot_concepts[k], i, :], 'o--', c=concept_colors[k], label=f'$C_{plot_concepts[k]}$')\n",
    "        axs[0, i].set_xlim(3.5, 20.5)\n",
    "        axs[0, i].set_ylim(0.25, 1.02)\n",
    "        # axs[0,i].set_xlabel(r\"$L_0$\")\n",
    "        # axs[0,i].set_ylabel(r\"$F_1$ score\")\n",
    "        axs[0,i].set_xticks([5, 10, 15, 20], labels=[5, 10, 15, 20], fontsize=18)\n",
    "        axs[0,i].set_xlabel(r\"$L_0$\", fontsize=20)\n",
    "        # axs[0,i].tick_params(axis='y', labelsize=10)\n",
    "        axs[0,i].set_yticks([0.4, 0.6, 0.8, 1.0], labels=[0.4, 0.6, 0.8, 1.0], fontsize=18)\n",
    "        if i>0:\n",
    "            axs[0,i].sharey(axs[0,0])\n",
    "            # axs[0,i].set_yticks([])\n",
    "        else:\n",
    "            axs[0,i].set_ylabel(r\"$F_1$ score\", fontsize=23, labelpad=10)\n",
    "        axs[0, 0].legend()\n",
    "\n",
    "# brightness of spectral clusters depends on the cluster size\n",
    "alphavals_specclus = []\n",
    "for i in range(len(saes_order)):\n",
    "    alphas_sci = []\n",
    "    for j in range(data_subset.size(0)):\n",
    "        sclabel_j = clusters_all[plot_indices[i]][j]\n",
    "        clussize_j = clustersizes_all[plot_indices[i]][sclabel_j]\n",
    "        max_clustersize = max(clustersizes_all[plot_indices[i]])\n",
    "        alphas_sci.append(clussize_j/max_clustersize)\n",
    "    alphavals_specclus.append(alphas_sci)\n",
    "\n",
    "for i in range(len(saes_order)):\n",
    "\n",
    "    #show sparse code correlations for data pairs\n",
    "    height, width = Cdatacos[plot_indices[i]].shape[:2]\n",
    "    axs[3, i].imshow(Cdatacos[plot_indices[i]], cmap='magma', vmin=0, vmax=1,\n",
    "                    extent=[0, width, height, 0])\n",
    "    axs[3, i].hlines([numpoints_perconcept*j for j in range(1, num_concepts+1)],\n",
    "                    0, width, color='#EFEFEF', linewidth=2)\n",
    "    axs[3, i].vlines([numpoints_perconcept*j for j in range(1, num_concepts+1)],\n",
    "                    0, height, color='#EFEFEF', linewidth=2)\n",
    "    axs[3,i].set_xticks(ticks=[500+1000*i for i in range(0, 6)], labels=[r\"$C_0$\", r\"$C_1$\", r\"$C_2$\", r\"$C_3$\", r\"$C_4$\", r\"$C_5$\"], fontsize=14)\n",
    "    if i==0:\n",
    "        axs[3,i].set_yticks(ticks=[500+1000*i for i in range(0, 6)], labels=[r\"$C_0$\", r\"$C_1$\", r\"$C_2$\", r\"$C_3$\", r\"$C_4$\", r\"$C_5$\"], fontsize=14)\n",
    "    else:\n",
    "        axs[3,i].set_yticks([])\n",
    "    axs[3,i].tick_params(axis='x', bottom=False)  # Remove x-axis ticks\n",
    "    axs[3,i].tick_params(axis='y', left=False)  # Remove y-axis tick\n",
    "\n",
    "    #show spectral clusters on above data-data correlation matrix\n",
    "    axs[4, i].scatter(data_subset_raw[indices,0], \\\n",
    "        data_subset_raw[indices,1], \\\n",
    "            c=clusters_all[plot_indices[i]][indices], \\\n",
    "                cmap='Set2', alpha=torch.tensor(alphavals_specclus[i])[indices], s=4)\n",
    "    axs[4,i].set_xticks([])\n",
    "    axs[4,i].set_yticks([])\n",
    "    if i==0:\n",
    "        axs[4,i].set_ylabel(r\"$d_2$\", fontsize=13)\n",
    "    axs[4,i].set_xlabel(r\"$d_1$\", fontsize=13)\n",
    "    axs[4,i].set_xlim(*xlims)\n",
    "    axs[4,i].set_ylim(*ylims)\n",
    "    \n",
    "    #names of concepts\n",
    "    color_clus = 'k'\n",
    "    for k in [1, 2, 4]:\n",
    "        axs[k,i].text(1.6, 0.6, r\"$C_0$\", fontsize=13, color=color_clus, ha='center', va='center')\n",
    "        axs[k,i].text(0.5, 1.0, r\"$C_1$\", fontsize=13, color=color_clus, ha='center', va='center')\n",
    "        axs[k,i].text(-1.5, 1.5, r\"$C_2$\", fontsize=13, color=color_clus, ha='center', va='center')\n",
    "        axs[k,i].text(-1.3, -0.2, r\"$C_3$\", fontsize=13, color=color_clus, ha='center', va='center')\n",
    "        axs[k,i].text(-0.2, -1.6, r\"$C_4$\", fontsize=13, color=color_clus, ha='center', va='center')\n",
    "        axs[k,i].text(0.8, -0.8, r\"$C_5$\", fontsize=13, color=color_clus, ha='center', va='center')\n",
    "    \n",
    "    #turn off top and right spines\n",
    "    for k in [0, 1, 2, 3, 4]:\n",
    "        axs[k,i].spines['top'].set_visible(False)\n",
    "        axs[k,i].spines['right'].set_visible(False)\n",
    "    \n",
    "for ax in axs.flatten():\n",
    "    ax.set_aspect('auto')\n",
    "\n",
    "# for i in range(axs.shape[0]):\n",
    "#     for j in range(axs.shape[1]):\n",
    "#         if i>0:\n",
    "#             axs[i,j].set_aspect('equal')\n",
    "#         else:\n",
    "#             axs[i,j].set_aspect('auto')\n",
    "    \n",
    "plt.tight_layout()\n",
    "# plt.subplots_adjust(wspace=0.5)\n",
    "# plt.show()\n",
    "# plt.savefig('./figs/sep_mainfigX.pdf', bbox_inches='tight', dpi=300)"
   ]
  },
  {
   "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
}
