notebook_templates/Other tables.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from concurrent.futures import ProcessPoolExecutor\n",
    "from pathlib import Path\n",
    "import pandas as pd\n",
    "from omegaconf import OmegaConf\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt \n",
    "from matplotlib.gridspec import GridSpec\n",
    "from tqdm.notebook import tqdm\n",
    "import bm\n",
    "os.chdir(Path(bm.__file__).parent.parent)\n",
    "from bm import train\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format='retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = train.main.dora.dir\n",
    "print(output_dir)\n",
    "eval_dir = output_dir / \"eval\" / \"signatures\"\n",
    "sigs_to_eval = [p.name for p in (output_dir / \"grids\" / \"nmi.main_table\").iterdir()]\n",
    "print(sigs_to_eval)\n",
    "assert output_dir.exists()\n",
    "assert eval_dir.exists()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Common functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data_from_sig(sig, level=\"segment\"):\n",
    "    \"\"\"\n",
    "    Load data from solver signature\n",
    "    - probs (torch tensor): probability on vocab [N, V]\n",
    "    - vocab (torch tensor): vocab of word hashes [V]\n",
    "    - words (torch tensor): the word hash for each sample [N]\n",
    "    - metadata (panda dataframe) of len [N] which contains for each sample:\n",
    "           'word_hashes', 'word_indices', 'seq_indices',\n",
    "           'word_strings', 'subject_id', 'recording_id'\n",
    "    \"\"\"\n",
    "    assert level in [\"word\", \"segment\"], \"level should be 'word' or 'segment'\"\n",
    "    probs = torch.load(eval_dir / sig / f\"probs_{level}.pth\") \n",
    "    vocab = torch.load(eval_dir / sig / f\"vocab_{level}.pth\") # vocab (hashes)\n",
    "    metadata = pd.read_csv(eval_dir / sig / \"metadata.csv\", index_col=0, low_memory=False) \n",
    "    \n",
    "    words = torch.tensor(metadata[f\"{level}_hashes\"].tolist()).long() # for each sample, the word (hashes)\n",
    "    assert probs.shape == (len(words), len(vocab)), (probs.shape, len(words), len(vocab))\n",
    "    assert len(words) == len(metadata)\n",
    "    metadata[\"idx\"] = range(len(words))\n",
    "    metadata[\"word_strings\"] = metadata[\"word_strings\"].str.lower()\n",
    "\n",
    "    return probs, vocab, words, metadata\n",
    "\n",
    "\n",
    "def get_accuracy_from_probs(probs, row_labels, col_labels, topk=10):\n",
    "    \"\"\"\n",
    "    probs: for each row, the probability distribution over a vocab\n",
    "    returns the topk accuracy that the topk best predicted labels\n",
    "    match the row_labels\n",
    "    Inputs:\n",
    "        probs: of shape [B, V] probability over vocab, each row sums to 1\n",
    "        row_labels: of shape [B] true word for each row\n",
    "        col_labels: [V] word that correspond to each column\n",
    "        topk: int\n",
    "    Returns: float scalar, topk accuracy\n",
    "    \"\"\"\n",
    "    assert len(row_labels) == len(probs)\n",
    "    assert len(col_labels) == probs.shape[1]\n",
    "\n",
    "    # Extract topk indices\n",
    "    idx = probs.topk(topk, dim=1).indices\n",
    "\n",
    "    # Get the corresponding topk labels\n",
    "    whs = col_labels[idx.view(-1)].reshape(idx.shape)\n",
    "\n",
    "    # 1 if the labels matches with the targets\n",
    "    correct = ((whs == row_labels[:, None]).any(1)).float()\n",
    "    assert len(correct) == len(row_labels)\n",
    "\n",
    "    # Average across samples\n",
    "    acc = correct.mean()\n",
    "\n",
    "    return acc.item()\n",
    "\n",
    "\n",
    "def eval_acc_one_sig(sig, topks=(1, 5, 10), level=\"word\", add_baselines=True):\n",
    "    \"\"\"\n",
    "    Return accuracy dataframe from one solver signature\n",
    "    level: whether to return `word` or `segment` level accuracy\n",
    "    \"\"\"\n",
    "    # Load data\n",
    "    probs, vocab, words, _ = load_data_from_sig(sig, level=level)\n",
    "#     if level == \"segment\":\n",
    "#         print(probs.shape)\n",
    "        \n",
    "    # Compute acc\n",
    "    acc_df = []\n",
    "    for topk in topks:\n",
    "        \n",
    "        # --- Acc ---\n",
    "        acc = get_accuracy_from_probs(probs, words, vocab, topk=topk)\n",
    "        \n",
    "        out = {\n",
    "            f\"acc\":acc,\n",
    "            \"topk\":topk,\n",
    "        }\n",
    "        \n",
    "        if add_baselines:\n",
    "        \n",
    "            # --- Baseline on vocab ---\n",
    "            # equivalent to : shuffle targets vocab (inf times)\n",
    "            # equivalent to : output uniform prob on vocab\n",
    "            # equivalent to : 1/vocab_len\n",
    "            rand_probs_vocab = torch.ones_like(probs) / len(vocab)\n",
    "            out[\"baseline_vocab\"] = get_accuracy_from_probs(rand_probs_vocab, words, vocab, topk=topk)\n",
    "\n",
    "            # --- Baseline on words ---\n",
    "            # equivalent to : shuffle word targets before aggregating on vocab (inf times)\n",
    "            # equivalent to : output uniform prob on samples\n",
    "            # equivalent to : each_word_freq\n",
    "            check_vocab, counts = torch.unique(words, return_counts=True)\n",
    "            import pdb\n",
    "#             assert (check_vocab == vocab).all()\n",
    "            rand_probs_words = torch.stack([counts/len(words)]*len(probs))\n",
    "            out[\"baseline\"] = get_accuracy_from_probs(rand_probs_words, words, vocab, topk=topk)\n",
    "\n",
    "            # Update\n",
    "            acc_df.append(out)\n",
    "    acc_df = pd.DataFrame(acc_df)\n",
    "    return acc_df\n",
    "\n",
    "def eval_acc(sigs, level=\"word\", add_baselines=True):\n",
    "    \"\"\"\n",
    "    Return accuracy dataframe for multiple sigs \n",
    "    level: whether to return word or segment level accuracy\n",
    "    \"\"\"\n",
    "    futures = []\n",
    "    acc = []\n",
    "    with ProcessPoolExecutor(20) as pool:\n",
    "        for sig in sigs:\n",
    "            future = pool.submit(eval_acc_one_sig, sig, level=level, add_baselines=add_baselines)\n",
    "            futures.append((sig, future))\n",
    "        for sig, future in tqdm(futures):\n",
    "            try:\n",
    "                acc_sig = future.result()\n",
    "            except Exception:\n",
    "                print(\"ERROR WITH\", sig)\n",
    "                raise\n",
    "                continue\n",
    "            acc_sig[\"sig\"] = sig\n",
    "            acc.append(acc_sig)\n",
    "    acc = pd.concat(acc)\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load meta dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select signatures\n",
    "valid_sigs = [sig for sig in sigs_to_eval if (eval_dir / sig / \"vocab_segment.pth\").is_file()]\n",
    "configs = [OmegaConf.load(eval_dir / sig / \"solver_config.yaml\") for sig in valid_sigs]\n",
    "for c, s in zip(configs, valid_sigs):\n",
    "    if not hasattr(c.dset, 'features'):\n",
    "        c.dset.features = c.dset.forcings\n",
    "        c.dset.features_params = c.dset.forcings_params\n",
    "        \n",
    "print(set(sigs_to_eval) - set(valid_sigs))\n",
    "run_df = pd.DataFrame({\n",
    "    \"sig\":valid_sigs,\n",
    "})\n",
    "run_df[\"dataset\"] = [\"-\".join(conf.dset.selections) for conf in configs]\n",
    "run_df[\"seed\"] = [conf.seed for conf in configs]\n",
    "run_df[\"forcings\"] = [\"-\".join(conf.dset.features) for conf in configs]\n",
    "run_df[\"loss\"] = [conf.optim.loss for conf in configs]\n",
    "run_df[\"is_random\"] = [conf.test.wer_random for conf in configs]\n",
    "run_df[\"max_scale\"] = [conf.norm.max_scale for conf in configs]\n",
    "run_df[\"n_mels\"] = [conf.dset.features_params.MelSpectrum.n_mels for conf in configs]\n",
    "run_df[\"deepmel\"] = [getattr(conf, 'feature_model_name', None) == 'deep_mel' for conf in configs]\n",
    "run_df[\"ft\"] = [conf.optim.epochs == 1 and not conf.test.wer_random for conf in configs]\n",
    "run_df[\"random\"] = [conf.test.wer_random for conf in configs]\n",
    "\n",
    "run_df[\"batch_size\"] = [conf.optim.batch_size for conf in configs]\n",
    "run_df[\"lr\"] = [conf.optim.lr for conf in configs]\n",
    "run_df[\"autorej\"] = [conf.dset.autoreject for conf in configs]\n",
    "run_df[\"n_rec\"] = [conf.dset.n_recordings for conf in configs]\n",
    "# run_df[\"ft\"] = [conf.optim.lr == 0 for conf in configs]\n",
    "run_df[\"dropout\"] = [conf.simpleconv.merger_dropout > 0 for conf in configs]\n",
    "run_df[\"gelu\"] = [bool(conf.simpleconv.gelu) for conf in configs]\n",
    "run_df[\"skip\"] = [bool(conf.simpleconv.skip) for conf in configs]\n",
    "run_df[\"initial\"] = [bool(conf.simpleconv.initial_linear) for conf in configs]\n",
    "run_df[\"complex\"] = [bool(conf.simpleconv.complex_out) for conf in configs]\n",
    "run_df[\"subject_lay\"] = [bool(conf.simpleconv.subject_layers) for conf in configs]\n",
    "run_df[\"subject_emb\"] = [bool(conf.simpleconv.subject_dim) for conf in configs]\n",
    "run_df[\"attention\"] = [bool(conf.simpleconv.merger) for conf in configs]\n",
    "run_df[\"glu\"] = [bool(conf.simpleconv.glu) for conf in configs]\n",
    "run_df[\"depth\"] = [conf.simpleconv.depth for conf in configs]\n",
    "run_df[\"offset_meg\"] = [conf.task.offset_meg_ms for conf in configs]\n",
    "run_df = run_df[run_df.loss == \"clip\"]\n",
    "# run_df = run_df[run_df.dataset == \"gwilliams2022\"]\n",
    "\n",
    "def get_name(row):\n",
    "    if row.ft:\n",
    "        return \"Trained with MSE\"\n",
    "    if row.random:\n",
    "        return \"Random model\"\n",
    "    if row.deepmel:\n",
    "        return \"Deep Mel\"\n",
    "    if row.forcings == 'MelSpectrum':\n",
    "        return 'MelSpectrum CLIP'\n",
    "    if not row.dropout:\n",
    "        return r\"\\wo spatial attention dropout\"\n",
    "    if not row.gelu:\n",
    "        return r\"\\wo GELU, \\w ReLU\"\n",
    "    if not row.skip:\n",
    "        return r\"\\wo skip connections\"\n",
    "    if not row.initial:\n",
    "        return r\"\\wo initial 1x1 conv.\"\n",
    "    if not row.complex:\n",
    "        return r\"\\wo final convs\"\n",
    "    if not row.attention:\n",
    "        return r\"\\wo spatial attention\"\n",
    "    if not row.glu:\n",
    "        return r\"\\wo non-residual GLU conv.\"\n",
    "    if not row.subject_lay:\n",
    "        if row.subject_emb:\n",
    "            return r\"\\w subj. embedding*\"\n",
    "        else:\n",
    "            return r\"\\wo subject-specific layer\"\n",
    "    if row.depth == 5:\n",
    "        return r\"less deep\"\n",
    "    if row.autorej:\n",
    "        return \"autoreject\"\n",
    "    if row.max_scale != 20:\n",
    "        if row.max_scale == 100:\n",
    "            return \"\\w clamp=100\"\n",
    "        else:\n",
    "            return \"\\wo clamping brain signal\"\n",
    "    return \"Our model\"\n",
    "    \n",
    "        \n",
    "run_df['name'] = run_df.apply(get_name, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Accuracy table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "acc_df = eval_acc(run_df[\"sig\"].values, level=\"segment\")\n",
    "# acc_df = pd.merge(acc_df, on=[\"sig\", \"topk\"], how=\"outer\")\n",
    "acc_df = pd.merge(acc_df, run_df, on=\"sig\", how=\"left\")\n",
    "def dset_order(name):\n",
    "    return name.map({\n",
    "       'audio_mous': 0, \n",
    "       'gwilliams2022': 1,\n",
    "       'broderick2019': 2,\n",
    "       'brennan2019': 3,\n",
    "    })\n",
    "acc_df = acc_df.sort_values([\"dataset\"], key=dset_order);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for name in run_df.name.unique():\n",
    "    print(name, ((run_df.name == name) & (run_df.dataset =='gwilliams2022')).sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Keys to set depends on the Table we want to generate\n",
    "keys = [\"dataset\", \"name\"]\n",
    "# keys = [\"dataset\", \"name\", \"lr\", \"batch_size\", \"offset_meg\", \"n_rec\"]\n",
    "# keys = [\"dataset\", \"name\", \"lr\", \"batch_size\"]\n",
    "# keys = [\"dataset\", \"name\", \"n_mels\"]\n",
    "# keys = [\"dataset\", \"offset_meg\"]\n",
    "acc_table = acc_df\n",
    "# acc_table = acc_df.query('batch_size == 256 & lr==3e-4 & n_rec != 16')\n",
    "# acc_table = acc_df.query('batch_size == 256 & lr==3e-4 & n_rec == 16')\n",
    "acc_table = acc_table.query(\"topk==10\").sort_values(keys).groupby(keys)[\"acc\"].agg([\"mean\", \"std\"])\n",
    "key = \"acc\"\n",
    "acc_table[\"str_acc\"] = (100 * acc_table[\"mean\"]).round(1).astype(str) + r\" PM \" + (100 * acc_table[\"std\"]).round(2).astype(str)\n",
    "print(acc_table)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert(x):\n",
    "    if isinstance(x, float):\n",
    "        print(x)\n",
    "    return float(x.split(\" \")[0])\n",
    "    \n",
    "toplot = acc_table.reset_index()\n",
    "index = list(keys)\n",
    "index.remove('dataset')\n",
    "# index.remove('n_mels')\n",
    "# index.remove('batch_size')\n",
    "toplot =  pd.pivot_table(toplot, values=[\"str_acc\"], columns=[\"dataset\"], index=index, aggfunc=\"first\")\n",
    "\n",
    "toplot[('str_acc', \"mean_dataset\")] = toplot.applymap(convert).mean(axis=1).round(1)\n",
    "# toplot.sort_values(index, ascending=False)\n",
    "# dsets = ['audio_mous', 'gwilliams2022', 'broderick2019', 'brennan2019'][:-1]\n",
    "dsets = ['audio_mous', 'gwilliams2022', 'broderick2019', 'brennan2019']\n",
    "# dsets = [\"gwilliams2022\"]\n",
    "# dsets = []\n",
    "extra = []\n",
    "\n",
    "toplot = toplot[[('str_acc', dset) \n",
    "                 for dset in dsets + ['mean_dataset'] +  extra]]\n",
    "# toplot = acc_table.reset_index()\n",
    "# index.remove('n_mels')\n",
    "# toplot =  pd.pivot_table(toplot, values=[\"str_acc\"], \n",
    "#                          columns=[\"n_mels\"], index=index, aggfunc=\"first\")\n",
    "toplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(toplot.to_latex(index=True).replace('PM', r'$\\pm$'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}