mutrans_gene.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Varying the set of genes used in regression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook assumes you have run preprocessing and the `--vary-gene` experiment\n",
    "```sh\n",
    "make update\n",
    "python mutrans.py --vary-gene\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "matplotlib.rcParams[\"figure.dpi\"] = 200\n",
    "matplotlib.rcParams[\"axes.edgecolor\"] = \"gray\"\n",
    "matplotlib.rcParams[\"savefig.bbox\"] = \"tight\"\n",
    "matplotlib.rcParams[\"savefig.pad_inches\"] = 0.01\n",
    "matplotlib.rcParams['font.family'] = 'sans-serif'\n",
    "matplotlib.rcParams['font.sans-serif'] = ['Arial', 'Avenir', 'DejaVu Sans']\n",
    "matplotlib.rcParams.update({\n",
    "    # 'text.usetex': True,\n",
    "    'text.latex.preamble': r'\\usepackage{amsfonts}',\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = torch.load(\"results/mutrans.vary_gene.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\n\".join(result[()].keys()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_metric(metric, descending=False, ylabel=None, filenames=()):\n",
    "    plt.figure(figsize=(8,4))\n",
    "    empty = ((\"exclude\", ((\"gene\", \".*\"),)),)\n",
    "    include = {}\n",
    "    exclude = {}\n",
    "    for k, v in result.items():\n",
    "        if k == empty:\n",
    "            continue\n",
    "        k = dict(k)\n",
    "        if \"include\" in k:\n",
    "            gene = k[\"include\"][0][1]\n",
    "            include[gene] = v[metric]\n",
    "        if \"exclude\" in k:\n",
    "            gene = k[\"exclude\"][0][1]\n",
    "            exclude[gene] = v[metric]\n",
    "    assert set(include) == set(exclude)\n",
    "    rankby = [(include[g] - exclude[g], g) for g in include]\n",
    "    rankby.sort(reverse=descending)\n",
    "    genes = [g for _, g in rankby]\n",
    "    X = list(range(len(genes)))\n",
    "    plt.axhline(result[empty][metric], color=\"darkgreen\", linestyle=\":\", label=\"No genes\",\n",
    "                zorder=-2)\n",
    "    plt.plot(X, [include[g] for g in genes], \"k+\", color=\"darkred\", label=\"A single gene\")\n",
    "    plt.plot(X, [exclude[g] for g in genes], \"ko\", color=\"darkblue\",label=\"All but one gene\",\n",
    "             markerfacecolor=\"None\")\n",
    "    plt.plot(X, [include[g] for g in genes], \"wo\", markersize=8, zorder=-1)\n",
    "    plt.plot(X, [exclude[g] for g in genes], \"wo\", markersize=8, zorder=-1)\n",
    "    plt.axhline(result[()][metric], color=\"k\", lw=1, linestyle=\"--\", label=\"All genes\",\n",
    "                zorder=-2)\n",
    "    plt.legend(loc=\"best\")\n",
    "    plt.xticks(X, labels=[g[1:-1] for g in genes])\n",
    "    plt.ylabel(metric if ylabel is None else ylabel)\n",
    "    plt.xlabel(\"Gene\")\n",
    "    for f in filenames:\n",
    "        plt.savefig(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(\"loss\", ylabel=\"ELBO loss\",\n",
    "            filenames=[\"paper/vary_gene_loss.png\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(\"rate_scale\", ylabel=r\"E[$\\sigma_4$]\",\n",
    "            filenames=[\"paper/vary_gene_loss.png\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(\"KL\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(\"MAE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "plot_metric(\"RMSE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(\"naive KL\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(\"naive MAE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_metric(\"naive RMSE\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_elbo(descending=True, filenames=()):\n",
    "    plt.figure(figsize=(8,4))\n",
    "    empty = ((\"exclude\", ((\"gene\", \".*\"),)),)\n",
    "    include = {}\n",
    "    exclude = {}\n",
    "    for k, v in result.items():\n",
    "        if k == empty:\n",
    "            continue\n",
    "        k = dict(k)\n",
    "        if \"include\" in k:\n",
    "            gene = k[\"include\"][0][1]\n",
    "            include[gene] = -v[\"loss\"]\n",
    "        if \"exclude\" in k:\n",
    "            gene = k[\"exclude\"][0][1]\n",
    "            exclude[gene] = -v[\"loss\"]\n",
    "    assert set(include) == set(exclude)\n",
    "    rankby = [(include[g] - exclude[g], g) for g in include]\n",
    "    rankby.sort(reverse=descending)\n",
    "    genes = [g for _, g in rankby]\n",
    "    X = list(range(len(genes)))\n",
    "    plt.axhline(-result[()][\"loss\"], color=\"darkgreen\", linestyle=\":\", label=\"All genes\",\n",
    "                zorder=-2)\n",
    "    plt.plot(X, [include[g] for g in genes], \"ko\", color=\"darkred\", label=\"One gene\",\n",
    "            markeredgecolor=\"white\", markersize=7)\n",
    "    plt.axhline(-result[empty][\"loss\"], color=\"k\", linestyle=\"--\", lw=1, label=\"No genes\",\n",
    "                zorder=-2)\n",
    "    plt.legend(loc=\"best\")\n",
    "    plt.xticks(X, labels=[g[1:-1] for g in genes])\n",
    "    plt.ylabel(\"ELBO\")\n",
    "    plt.xlabel(\"Gene\")\n",
    "    plt.tight_layout()\n",
    "    for f in filenames:\n",
    "        plt.savefig(f)\n",
    "        \n",
    "plot_elbo(filenames=[\"paper/vary_gene_elbo.png\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## --vary-nsp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = torch.load(\"results/mutrans.vary_nsp.pt\")\n",
    "result.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_elbo(descending=True, filenames=()):\n",
    "    plt.figure(figsize=(8,4))\n",
    "    empty = ((\"exclude\", ((\"gene\", \".*\"),)),)\n",
    "    full = ((\"include\", ((\"gene\", \"^ORF1[ab]:\"),)),)\n",
    "    include = {}\n",
    "    for k, v in result.items():\n",
    "        k = dict(k)\n",
    "        if \"include\" not in k:\n",
    "            continue\n",
    "        k = dict(k[\"include\"])\n",
    "        if \"region\" not in k:\n",
    "            continue\n",
    "        gene, region = k[\"region\"]\n",
    "        include[gene, region] = -v[\"loss\"]\n",
    "    rankby = [(include[gr], gr[-1]) for gr in include]\n",
    "    rankby.sort(reverse=descending)\n",
    "    labels = [g for _, g in rankby]\n",
    "    elbos = [e for e, _ in rankby]\n",
    "    X = list(range(len(labels)))\n",
    "    plt.axhline(-result[full][\"loss\"], color=\"darkgreen\", linestyle=\":\", label=\"All of ORF1\",\n",
    "                zorder=-2)\n",
    "    plt.plot(X, elbos, \"ko\", color=\"darkred\", label=\"One nsp\",\n",
    "             markeredgecolor=\"white\", markersize=7)\n",
    "    plt.axhline(-result[empty][\"loss\"], color=\"k\", linestyle=\"--\", lw=1, label=\"No nsps\",\n",
    "                zorder=-2)\n",
    "    plt.legend(loc=\"best\")\n",
    "    plt.xticks(X, labels=labels)\n",
    "    plt.ylabel(\"ELBO\")\n",
    "    plt.xlabel(\"Nonstructural protein within ORF1\")\n",
    "    plt.tight_layout()\n",
    "    for f in filenames:\n",
    "        plt.savefig(f)\n",
    "        \n",
    "plot_elbo(filenames=[\"paper/vary_nsp_elbo.png\"])"
   ]
  },
  {
   "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.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}