mutrans_prediction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predictive accuracy of mutrans model on new lineages\n",
    "\n",
    "This notebook assumes you have run\n",
    "```sh\n",
    "make update  # downloads and preprocesses data\n",
    "python mutrans.py --vary-leaves=50  # or some largish number\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from pyrocov import pangolin\n",
    "from pyrocov.util import pearson_correlation\n",
    "\n",
    "matplotlib.rcParams[\"figure.dpi\"] = 200\n",
    "matplotlib.rcParams[\"axes.edgecolor\"] = \"gray\"\n",
    "matplotlib.rcParams[\"savefig.bbox\"] = \"tight\"\n",
    "matplotlib.rcParams['font.family'] = 'sans-serif'\n",
    "matplotlib.rcParams['font.sans-serif'] = ['Arial', 'Avenir', 'DejaVu Sans']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torch.load(\"results/mutrans.data.single.None.pt\", map_location=\"cpu\")\n",
    "print(dataset.keys())\n",
    "locals().update(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lineage_id = {name: i for i, name in enumerate(lineage_id_inv)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = torch.load(\"results/mutrans.pt\")\n",
    "for key in results:\n",
    "    print(key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_fit = list(results.values())[0]\n",
    "print(best_fit.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loo = torch.load(\"results/mutrans.vary_leaves.pt\", map_location=\"cpu\")\n",
    "print(len(loo))\n",
    "print(list(loo)[0])\n",
    "print(list(loo.values())[0].keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(mutations))\n",
    "print(best_fit[\"median\"][\"rate_loc\"].shape)\n",
    "print(list(loo.values())[0][\"median\"][\"rate_loc\"].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loo_rate_loc = {\n",
    "    k[-1][-1][-1][-1][-1].replace(\"$\", \"\").replace(\"^\", \"\"): v[\"median\"][\"rate_loc\"]\n",
    "    for k, v in loo.items()\n",
    "}\n",
    "print(\" \".join(loo_rate_loc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_rate_loc = best_fit[\"median\"][\"rate_loc\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_prediction(filenames=[], debug=False, use_who=True):\n",
    "    X1, Y1, X2, Y2, labels, debug_labels = [], [], [], [], [], []\n",
    "    who = {vs[0]: k for k, vs in pangolin.WHO_ALIASES.items()}\n",
    "    ancestors = set(lineage_id)\n",
    "    for child, rate_loc in loo_rate_loc.items():\n",
    "        parent = pangolin.compress(\n",
    "            pangolin.get_most_recent_ancestor(\n",
    "                pangolin.decompress(child), ancestors\n",
    "            )\n",
    "        )\n",
    "        c = lineage_id[child]\n",
    "        p = lineage_id[parent]\n",
    "        truth = best_rate_loc[c].item()\n",
    "        baseline = rate_loc[p].item()\n",
    "        guess = rate_loc[c].item()\n",
    "        X1.append(truth)\n",
    "        Y1.append(guess)\n",
    "        X2.append(truth - baseline)\n",
    "        Y2.append(guess - baseline)\n",
    "        labels.append(who.get(child))\n",
    "        debug_labels.append(child)\n",
    "    fig, axes = plt.subplots(1, 2, figsize=(7, 3.5))\n",
    "    for ax, X, Y in zip(axes, [X1, X2], [Y1, Y2]):\n",
    "        X = np.array(X)\n",
    "        Y = np.array(Y)\n",
    "        ax.scatter(X, Y, 40, lw=0, alpha=1, color=\"white\", zorder=-5)\n",
    "        ax.scatter(X, Y, 20, lw=0, alpha=0.5, color=\"darkred\")\n",
    "        lb = min(min(X), min(Y))\n",
    "        ub = max(max(X), max(Y))\n",
    "        d = ub - lb\n",
    "        lb -= 0.03 * d\n",
    "        ub += 0.05 * d\n",
    "        ax.plot([lb, ub], [lb, ub], \"k--\", alpha=0.2, zorder=-10)\n",
    "        ax.set_xlim(lb, ub)\n",
    "        ax.set_ylim(lb, ub)\n",
    "        ax.text(lb + 0.06 * d, ub - 0.1 * d, f\"ρ = {pearson_correlation(X, Y):0.2g}\",\n",
    "                backgroundcolor=\"white\")\n",
    "        for x, y, label, debug_label in zip(X, Y, labels, debug_labels):\n",
    "            if label is not None:\n",
    "                pad = 0.01\n",
    "                ax.plot([x], [y], \"ko\", mfc=\"#c77\", c=\"black\", ms=4, mew=0.5)\n",
    "                ax.text(x, y + pad, label if use_who else debug_label,\n",
    "                        va=\"bottom\", ha=\"center\", fontsize=6)\n",
    "            if debug and abs(x - y) > 0.2:\n",
    "                ax.text(x, y - pad, debug_label, va=\"top\", ha=\"center\", fontsize=5)\n",
    "    axes[0].set_xlabel(\"full estimate\")\n",
    "    axes[0].set_ylabel(\"LOO estimate\")\n",
    "    axes[1].set_xlabel(\"full estimate − baseline\")\n",
    "    axes[1].set_ylabel(\"LOO estimate − baseline\")\n",
    "    plt.tight_layout()\n",
    "    for f in filenames:\n",
    "        plt.savefig(f)\n",
    "plot_prediction(debug=True)\n",
    "plot_prediction(use_who=False, filenames=[\"paper/lineage_prediction.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
}