mini_mutrans.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Mini version of mutrans.py model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import logging\n",
    "import math\n",
    "import re\n",
    "from collections import defaultdict\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "import pyro.distributions as dist\n",
    "import pyro\n",
    "from pyro.infer import SVI, Trace_ELBO\n",
    "from pyro.optim import ClippedAdam\n",
    "from pyrocov import mutrans, pangolin, stats\n",
    "from pyrocov.util import pretty_print\n",
    "\n",
    "logging.basicConfig(format=\"%(message)s\", level=logging.INFO)\n",
    "matplotlib.rcParams[\"figure.dpi\"] = 200\n",
    "matplotlib.rcParams[\"axes.edgecolor\"] = \"gray\"\n",
    "matplotlib.rcParams['font.family'] = 'sans-serif'\n",
    "matplotlib.rcParams['font.sans-serif'] = ['Arial', 'Avenir', 'DejaVu Sans']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def load_data_subset(*args, **kwargs):\n",
    "    filename = \"results/mutrans.data.single.pt\"\n",
    "    if os.path.exists(filename):\n",
    "        dataset = torch.load(filename)\n",
    "    else:\n",
    "        dataset = mutrans.load_gisaid_data()\n",
    "        torch.save(dataset, filename)\n",
    "    dataset = mutrans.subset_gisaid_data(dataset, *args, **kwargs)\n",
    "    dataset.update(mutrans.load_jhu_data(dataset))\n",
    "    return dataset\n",
    "\n",
    "queries = [\"England\", \"USA / California\", \"Brazil\", \"India\"]\n",
    "dataset = load_data_subset(\n",
    "    queries,\n",
    "    max_strains=100,\n",
    ")\n",
    "locals().update(dataset)\n",
    "\n",
    "values, indices = weekly_strains.sum([0,-1]).sort(0, descending=True)\n",
    "location_id_inv = list(location_id)\n",
    "for v, i in zip(values.tolist(), indices.tolist()):\n",
    "    print(\"{: >7d} {}\".format(int(v), location_id_inv[i]))\n",
    "print()\n",
    "for k, v in sorted(dataset.items()):\n",
    "    if isinstance(v, torch.Tensor):\n",
    "        print(f\"{k} \\t{type(v).__name__} of shape {tuple(v.shape)}\")\n",
    "    else:\n",
    "        print(f\"{k} \\t{type(v).__name__} of size {len(v)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "fit = mutrans.fit_svi(\n",
    "    dataset,\n",
    "    cond_data={},\n",
    "    model_type=\"sparse-poisson\",\n",
    "    guide_type=\"gaussian\",\n",
    "    forecast_steps=6,\n",
    "    learning_rate=0.02,\n",
    "    learning_rate_decay=1,\n",
    "    num_steps=2001,\n",
    "    clip_norm=10.0,\n",
    "    rank=10,\n",
    "    log_every=100,\n",
    "    seed=20210319,\n",
    ")\n",
    "mutrans.log_stats(dataset, fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def plusminus(mean, std):\n",
    "    p95 = 1.96 * std\n",
    "    return torch.stack([mean - p95, mean, mean + p95])\n",
    "\n",
    "def plot_forecast(fit, queries=None, num_strains=10):\n",
    "    if queries is None:\n",
    "        queries = list(location_id)\n",
    "    elif isinstance(queries, str):\n",
    "        queries = [queries]\n",
    "    fig, axes = plt.subplots(len(queries), figsize=(8, 1.5 + 2 * len(queries)), sharex=True)\n",
    "    if not isinstance(axes, (list, np.ndarray)):\n",
    "        axes = [axes]\n",
    "    dates = matplotlib.dates.date2num(mutrans.date_range(len(fit[\"mean\"][\"probs\"])))\n",
    "    forecast_steps = len(fit[\"mean\"][\"probs\"]) - len(weekly_cases)\n",
    "    assert forecast_steps >= 0\n",
    "    probs = plusminus(fit[\"mean\"][\"probs\"], fit[\"std\"][\"probs\"])  # [3, T, P, S]\n",
    "    padding = 1 + weekly_cases[-1:].expand(forecast_steps, -1)\n",
    "    weekly_cases_ = torch.cat([weekly_cases, padding], 0)\n",
    "    predicted = probs * weekly_cases_[..., None]\n",
    "    ids = torch.tensor([i for name, i in location_id.items()\n",
    "                        if any(q in name for q in queries)])\n",
    "    strain_ids = weekly_strains[:, ids].sum([0, 1]).sort(-1, descending=True).indices\n",
    "    strain_ids = strain_ids[:num_strains]\n",
    "    colors = [f\"C{i}\" for i in range(10)] + [\"black\"] * 90\n",
    "    assert len(colors) >= num_strains\n",
    "    light = \"#bbbbbb\"\n",
    "    for row, (query, ax) in enumerate(zip(queries, axes)):\n",
    "        ids = torch.tensor([i for name, i in location_id.items() if query in name])\n",
    "        print(f\"{query} matched {len(ids)} regions\")\n",
    "        counts = weekly_cases[:, ids].sum(1)\n",
    "        counts /= counts.max()\n",
    "        ax.plot(dates[:len(counts)], counts, \"k-\", color=light, lw=0.8, zorder=-20)\n",
    "        counts = weekly_strains[:, ids].sum([1, 2])\n",
    "        counts /= counts.max()\n",
    "        ax.plot(dates[:len(counts)], counts, \"k--\", color=light, lw=1, zorder=-20)\n",
    "        pred = predicted.index_select(-2, ids).sum(-2)\n",
    "        pred /= pred[1].sum(-1, True).clamp_(min=1e-8)\n",
    "        obs = weekly_strains[:, ids].sum(1)\n",
    "        obs /= obs.sum(-1, True).clamp_(min=1e-9)\n",
    "        for s, color in zip(strain_ids, colors):\n",
    "            lb, mean, ub = pred[..., s]\n",
    "            ax.fill_between(dates, lb, ub, color=color, alpha=0.2, zorder=-10)\n",
    "            ax.plot(dates, mean, color=color, lw=1, zorder=-9)\n",
    "            strain = lineage_id_inv[s]\n",
    "            ax.plot(dates[:len(obs)], obs[:, s], color=color, lw=0, marker='o', markersize=3,\n",
    "                    label=strain if row == 0 else None)\n",
    "        ax.set_ylim(0, 1)\n",
    "        ax.set_yticks(())\n",
    "        ax.set_ylabel(query.replace(\" / \", \"\\n\"))\n",
    "        ax.set_xlim(dates.min(), dates.max())\n",
    "        if row == 0:\n",
    "            ax.legend(loc=\"upper left\", fontsize=8)\n",
    "        elif row == 1:\n",
    "            ax.plot([], \"k--\", color=light, lw=1, label=\"relative #samples\")\n",
    "            ax.plot([], \"k-\", color=light, lw=0.8, label=\"relative #cases\")\n",
    "            ax.plot([], lw=0, marker='o', markersize=3, color='gray',\n",
    "                    label=\"observed portion\")\n",
    "            ax.fill_between([], [], [], color='gray', label=\"predicted portion\")\n",
    "            ax.legend(loc=\"upper left\", fontsize=8)\n",
    "    ax.xaxis.set_major_locator(matplotlib.dates.MonthLocator())\n",
    "    ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter(\"%b %Y\"))\n",
    "    plt.xticks(rotation=90)\n",
    "    plt.subplots_adjust(hspace=0)\n",
    "plot_forecast(fit, queries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def plot_prevalence(fit):\n",
    "    rate = fit[\"median\"][\"rate\"]\n",
    "    if rate.dim() == 2:\n",
    "        rate = rate.mean(0)\n",
    "    rate = rate - rate.median()\n",
    "    R = rate.exp()\n",
    "    probs = fit[\"mean\"][\"probs\"].mean(0)\n",
    "    cases = torch.einsum(\"ps,p->s\", probs, weekly_cases[-2])\n",
    "    cases = cases / mutrans.TIMESTEP\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.scatter(R, cases, lw=0, alpha=0)\n",
    "    for name, i in lineage_id.items():\n",
    "        plt.text(R[i], cases[i], name, fontsize=8, alpha=0.8,\n",
    "                 horizontalalignment=\"center\", verticalalignment=\"center\",)\n",
    "    plt.ylabel(\"confirmed cases / day\")\n",
    "    plt.yscale(\"symlog\")\n",
    "    lb10 = math.floor(10 * R.min().item())\n",
    "    ub10 = math.ceil(10 * R.max().item())\n",
    "    plt.xlabel(\"relative reproduction number $R_{strain} / R_{median}$\")\n",
    "    plt.title(f\"Transmissibility of {len(lineage_id)} PANGO lineages\")\n",
    "    plt.tight_layout()\n",
    "\n",
    "plot_prevalence(fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_volcano(mean, std, linthresh=2, top_k=50):\n",
    "    xs = mean * 0.01\n",
    "    ys = mean.abs() / std\n",
    "    assert len(xs) == len(mutations)\n",
    "    y0, y1 = float(ys.min()), float(ys.max())\n",
    "    x0, x1 = float(xs.min()), float(xs.max())\n",
    "    ys, idx = ys.sort(0, descending=True)\n",
    "    xs = xs[idx]\n",
    "    pos = (0 < xs) & (xs < math.inf)\n",
    "    neg = (-math.inf < xs) & (xs < 0)\n",
    "    ys_pos, ys_neg = ys[pos], ys[neg]\n",
    "    xs_pos, xs_neg = xs[pos], xs[neg]\n",
    "    idx_pos, idx_neg = idx[pos], idx[neg]\n",
    "    N = top_k\n",
    "\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    plt.title(f\"Increased transmissibility of {len(mutations)} mutations\")\n",
    "    plt.plot(xs[pos][:N], ys[pos][:N], 'k.', lw=0, markersize=2, zorder=10)\n",
    "    plt.plot(xs[pos][N:], ys[pos][N:], 'k.', lw=0, markersize=2, zorder=10, color=\"#aaa\")\n",
    "    plt.plot(xs[neg], ys[neg], 'k.', lw=0, markersize=2, zorder=10, color=\"#aaa\")\n",
    "    plt.xlabel(\"effect size (R / R_A = exp(μ))\")\n",
    "    plt.ylabel(\"statistical significance (|μ|/σ)\")\n",
    "    lpad = 0\n",
    "    rpad = 0.33 if any(\",\" in mutations[i] for i in idx_pos[:N].tolist()) else 0.18\n",
    "    plt.xlim(x0 - (x1 - x0) * lpad, x1 + (x1 - x0) * rpad)\n",
    "    plt.ylim(0, None)\n",
    "    plt.yscale(\"symlog\", linthresh=linthresh)\n",
    "    yticks = [y for y in [0, 1, 2, 5, 10, 20, 50, 100, 200, 500, 1000, 2000] if y < y1]\n",
    "    plt.yticks(yticks, list(map(str, yticks)))\n",
    "    p95 = dist.Normal(0, 1).icdf(torch.tensor(0.95)).item()\n",
    "    plt.plot([x0, x1], [p95, p95], 'k--', alpha=0.2)\n",
    "    comment_pos = 0.5\n",
    "    plt.text((1 - comment_pos) * x0 + comment_pos * x1, p95 * 0.95,\n",
    "             \"95% probabililty\\nof correct sign(μ)\",\n",
    "             fontsize=7, horizontalalignment=\"center\", verticalalignment=\"top\",\n",
    "             alpha=0.8, zorder=100)\n",
    "        \n",
    "    colors = {\"N\": \"blue\", \"S\": \"red\", \"M\": \"purple\", \"ORF3a\": \"darkgreen\"}\n",
    "    ax = plt.gca()\n",
    "    t = (ax.transScale + ax.transLimits).inverted()\n",
    "    for i in range(N):\n",
    "        x = x1\n",
    "        _, y = t.transform((0, 1 - (i + 1) / (N + 1)))\n",
    "        name = mutations[int(idx_pos[i])]\n",
    "        plt.plot([x, xs_pos[i]], [y, ys_pos[i]], color='gray', lw=0.2)\n",
    "        plt.text(x, y, \" \" + name, color=colors.get(name.split(\":\")[0], \"gray\"),\n",
    "                 fontsize=8, verticalalignment=\"center\", horizontalalignment=\"left\")\n",
    "    for spine in ax.spines.values():\n",
    "        spine.set_linewidth(0.5)\n",
    "    plt.tight_layout()\n",
    "\n",
    "plot_volcano(fit[\"mean\"][\"coef\"], fit[\"std\"][\"coef\"], linthresh=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "def plot_fit(fit):\n",
    "    num_nonzero = int(torch.count_nonzero(weekly_strains))\n",
    "    median = fit.get(\"median\", fit.get(\"mean\", {}))\n",
    "    plt.figure(figsize=(8, 7))\n",
    "    time = np.arange(1, 1 + len(fit[\"losses\"]))\n",
    "    plt.plot(fit[\"losses\"], \"k--\", label=\"loss\")\n",
    "    locs = []\n",
    "    grads = []\n",
    "    for name, series in fit[\"series\"].items():\n",
    "        rankby = -torch.tensor(series).log1p().mean().item()\n",
    "        if \"Guide.\" in name:\n",
    "            name = name.split(\"Guide.\")[-1].replace(\"$$$\", \".\")\n",
    "            grads.append((name, series, rankby))\n",
    "        elif \"_centered\" in name or name == \"local_time\":\n",
    "            grads.append((name, series, rankby))\n",
    "        elif name != \"loss\":\n",
    "            locs.append((name, series, rankby))\n",
    "    locs.sort(key=lambda x: x[-1])\n",
    "    grads.sort(key=lambda x: x[-1])\n",
    "    for name, series, _ in locs:\n",
    "        plt.plot(time, series, label=name)\n",
    "    for name, series, _ in locs:\n",
    "        plt.plot(time, series, color=\"white\", lw=3, alpha=0.3, zorder=-1)\n",
    "    for name, series, _ in grads:\n",
    "        plt.plot(time, series, lw=1, alpha=0.3, label=name, zorder=-2)\n",
    "    plt.yscale(\"log\")\n",
    "    plt.xscale(\"log\")\n",
    "    plt.xlim(1, len(fit[\"losses\"]))\n",
    "    plt.legend(loc=\"upper left\", fontsize=8)\n",
    "    plt.xlabel(\"SVI step (duration = {:0.1f} minutes)\".format(fit[\"walltime\"]/60))\n",
    "    loss = np.median(fit[\"losses\"][-201:]) / num_nonzero\n",
    "    plt.title(\" \".join([f\"L={loss:0.6g}\"] + [\n",
    "        \"{}={:0.3g}\".format(\n",
    "            \"\".join(p[0] for p in k.split(\"_\")).upper(), v\n",
    "        )\n",
    "        for k, v in median.items()\n",
    "        if v.numel() == 1\n",
    "    ]))\n",
    "plot_fit(fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "centered = defaultdict(list)\n",
    "for k, v in fit[\"params\"].items():\n",
    "    if \"_centered\" in k:\n",
    "        centered[k.split(\"_centered\")[0]].append(v)\n",
    "for k, v in sorted(centered.items()):\n",
    "    v = torch.distributions.transform_to(dist.constraints.unit_interval)(sum(v))\n",
    "    print(f\"{k}: {v.min().item():0.3g} {v.mean().item():0.3g}, {v.max().item():0.3g}\")"
   ]
  },
  {
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}